1 Commits

Author SHA1 Message Date
Zhongdi LUO
331c4decaa feat: add pipelined fp8 scalar tmem softmax unit 2026-07-12 02:07:07 +00:00

View File

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