File size: 2,186 Bytes
2166d7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | // Shared constants and lookup tables for the FlashAttention engine.
// Numeric formats are defined bit-exactly by software/fa_int_model.py.
package fa_pkg;
localparam int D_MAX = 16; // max head dimension
localparam int BR = 4; // query rows per tile
localparam int BC = 4; // keys per tile (part of the numeric definition)
localparam int NEG_INF = -(1 << 30); // initial running max
localparam int FA_TRAP = -2; // SUBLEQ 'a' operand that selects the FA opcode
// exp(-i/16 in log2 units) table, Q0.16: round(65536 * 2^(-i/16)), i = 0..16
function automatic logic [16:0] exp_lut(input logic [4:0] i);
case (i)
5'd0: exp_lut = 17'd65536;
5'd1: exp_lut = 17'd62757;
5'd2: exp_lut = 17'd60097;
5'd3: exp_lut = 17'd57549;
5'd4: exp_lut = 17'd55109;
5'd5: exp_lut = 17'd52773;
5'd6: exp_lut = 17'd50535;
5'd7: exp_lut = 17'd48393;
5'd8: exp_lut = 17'd46341;
5'd9: exp_lut = 17'd44376;
5'd10: exp_lut = 17'd42495;
5'd11: exp_lut = 17'd40693;
5'd12: exp_lut = 17'd38968;
5'd13: exp_lut = 17'd37316;
5'd14: exp_lut = 17'd35734;
5'd15: exp_lut = 17'd34219;
default: exp_lut = 17'd32768;
endcase
endfunction
// round(256 / sqrt(d)), d = 1..16
function automatic logic [8:0] scale_q(input logic [4:0] d);
case (d)
5'd1: scale_q = 9'd256;
5'd2: scale_q = 9'd181;
5'd3: scale_q = 9'd148;
5'd4: scale_q = 9'd128;
5'd5: scale_q = 9'd114;
5'd6: scale_q = 9'd105;
5'd7: scale_q = 9'd97;
5'd8: scale_q = 9'd91;
5'd9: scale_q = 9'd85;
5'd10: scale_q = 9'd81;
5'd11: scale_q = 9'd77;
5'd12: scale_q = 9'd74;
5'd13: scale_q = 9'd71;
5'd14: scale_q = 9'd68;
5'd15: scale_q = 9'd66;
default: scale_q = 9'd64;
endcase
endfunction
endpackage : fa_pkg
|