2 Commits

Author SHA1 Message Date
Zhongdi LUO
eb1b2ea330 feat: add pipelined scalar tmem softmax unit 2026-07-12 01:59:32 +00:00
Zhongdi LUO
2bfc6c4bde fix: pack scalar tmem softmax as fp16 2026-07-10 13:01:58 +00:00

View File

@@ -161,50 +161,29 @@ module VX_tmem_softmax_unit #(
end end
endfunction endfunction
// Convert a positive unsigned Q15 probability to FP8 E4M3 using function automatic [15:0] q15_to_fp16;
// round-to-nearest-even. Softmax probabilities are in [0, 1], so the
// positive subnormal and exponent ranges through 1.0 are sufficient.
function automatic [7:0] q15_to_fp8e4m3;
input [15:0] value; input [15:0] value;
integer msb; integer msb;
integer i; integer i;
integer shift; reg [15:0] normalized;
reg [3:0] exponent;
reg [4:0] rounded;
reg [15:0] remainder;
reg [15:0] halfway;
begin begin
if (value < 16'd512) begin msb = -1;
// E4M3 subnormals and the minimum normal are spaced by 2^-9, for (i = 0; i < 16; ++i) begin
// which is 64 in unsigned Q15. if (value[i]) msb = i;
rounded = value[10:6]; end
remainder = value & 16'h003f; if (msb < 0) begin
if ((remainder > 16'd32) q15_to_fp16 = 16'd0;
|| ((remainder == 16'd32) && rounded[0]))
rounded = rounded + 1'b1;
q15_to_fp8e4m3 = {3'b000, rounded};
end else begin end else begin
msb = 0; if (msb >= 10)
for (i = 0; i < 16; ++i) begin normalized = value >> (msb - 10);
if (value[i]) msb = i;
end
exponent = 4'(msb - 8);
shift = integer'(exponent) + 5;
rounded = 5'(value >> shift);
remainder = value & ((16'd1 << shift) - 1'b1);
halfway = 16'd1 << (shift - 1);
if ((remainder > halfway)
|| ((remainder == halfway) && rounded[0]))
rounded = rounded + 1'b1;
if (rounded == 5'd16)
q15_to_fp8e4m3 = {1'b0, exponent + 1'b1, 3'b000};
else else
q15_to_fp8e4m3 = {1'b0, exponent, rounded[2:0]}; normalized = value << (10 - msb);
q15_to_fp16 = {1'b0, msb[4:0], normalized[9:0]};
end end
end end
endfunction endfunction
function automatic [7:0] normalize_to_fp8e4m3; function automatic [15:0] normalize_to_fp16;
input [23:0] exp_value; input [23:0] exp_value;
input [31:0] reciprocal; input [31:0] reciprocal;
reg [55:0] product; reg [55:0] product;
@@ -212,10 +191,10 @@ module VX_tmem_softmax_unit #(
begin begin
product = exp_value * reciprocal; product = exp_value * reciprocal;
probability = (product + 56'd524288) >> 20; probability = (product + 56'd524288) >> 20;
if (probability > 56'd32768) if (probability > 36'd32768)
normalize_to_fp8e4m3 = 8'h38; normalize_to_fp16 = 16'h3c00;
else else
normalize_to_fp8e4m3 = q15_to_fp8e4m3(probability[15:0]); normalize_to_fp16 = q15_to_fp16(probability[15:0]);
end end
endfunction endfunction
@@ -368,7 +347,7 @@ module VX_tmem_softmax_unit #(
end else if (!out_valid || out_ready) begin end else if (!out_valid || out_ready) begin
for (lane = 0; lane < NUM_LANES; ++lane) begin for (lane = 0; lane < NUM_LANES; ++lane) begin
out_data[lane] <= { out_data[lane] <= {
4{normalize_to_fp8e4m3( 2{normalize_to_fp16(
exp_values[norm_index[4:0]][lane], reciprocals[lane])} exp_values[norm_index[4:0]][lane], reciprocals[lane])}
}; };
end end