fix: pack scalar tmem softmax as fp16

hw/dpi/float_dpi.cpp

@@ -12,6 +12,7 @@
 // limitations under the License.
 
 #include <stdio.h>
+#include <stdint.h>
 #include <math.h>
 #include <unordered_map>
 #include <vector>
@@ -38,6 +39,10 @@ extern "C" {
   void dpi_fdiv(bool enable, int dst_fmt, int64_t a, int64_t b, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags);
   void dpi_fsqrt(bool enable, int dst_fmt, int64_t a, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags);
   void dpi_fexp(bool enable, int dst_fmt, int64_t a, int64_t* result, svBitVecVal* fflags);
+  int dpi_f32_max(bool enable, int a_bits, int b_bits);
+  int dpi_softmax_exp_acc(bool enable, int score_bits, int max_bits, int accum_bits);
+  int dpi_softmax_prob_to_f16x2(bool enable, int score_bits, int max_bits, int denom_bits);
+  int dpi_softmax_prob_to_fp8e4m3x4(bool enable, int score_bits, int max_bits, int denom_bits);
 
   void dpi_ftoi(bool enable, int dst_fmt, int src_fmt, int64_t a, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags);
   void dpi_ftou(bool enable, int dst_fmt, int src_fmt, int64_t a, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags);
@@ -100,6 +105,73 @@ inline uint32_t float_to_bits(float value) {
   return bits.u;
 }
 
+int dpi_f32_max(bool enable, int a_bits, int b_bits) {
+  if (!enable)
+    return 0;
+  float a = bits_to_float(static_cast<uint32_t>(a_bits));
+  float b = bits_to_float(static_cast<uint32_t>(b_bits));
+  return static_cast<int>(float_to_bits(a > b ? a : b));
+}
+
+int dpi_softmax_exp_acc(bool enable, int score_bits, int max_bits, int accum_bits) {
+  if (!enable)
+    return 0;
+  float score = bits_to_float(static_cast<uint32_t>(score_bits));
+  float row_max = bits_to_float(static_cast<uint32_t>(max_bits));
+  float accum = bits_to_float(static_cast<uint32_t>(accum_bits));
+  return static_cast<int>(float_to_bits(accum + expf(score - row_max)));
+}
+
+int dpi_softmax_prob_to_f16x2(bool enable, int score_bits, int max_bits, int denom_bits) {
+  if (!enable)
+    return 0;
+  float score = bits_to_float(static_cast<uint32_t>(score_bits));
+  float row_max = bits_to_float(static_cast<uint32_t>(max_bits));
+  float denom = bits_to_float(static_cast<uint32_t>(denom_bits));
+  float prob = denom == 0.0f ? 0.0f : expf(score - row_max) / denom;
+  half_float::half half_prob(prob);
+  uint16_t half_bits;
+  static_assert(sizeof(half_bits) == sizeof(half_prob), "unexpected half size");
+  __builtin_memcpy(&half_bits, &half_prob, sizeof(half_bits));
+  return static_cast<int>(half_bits | (uint32_t(half_bits) << 16));
+}
+
+static float fp8e4m3_positive_to_float(uint8_t bits) {
+  const int exp = (bits >> 3) & 0xf;
+  const int frac = bits & 0x7;
+  if (exp == 0) {
+    return frac == 0 ? 0.0f : ldexpf(static_cast<float>(frac) / 8.0f, -6);
+  }
+  return ldexpf(1.0f + static_cast<float>(frac) / 8.0f, exp - 7);
+}
+
+static uint8_t float_to_fp8e4m3_positive(float value) {
+  if (!(value > 0.0f))
+    return 0;
+  uint8_t best = 0;
+  float best_diff = fabsf(value);
+  for (uint32_t bits = 1; bits < 0x80; ++bits) {
+    const float candidate = fp8e4m3_positive_to_float(static_cast<uint8_t>(bits));
+    const float diff = fabsf(candidate - value);
+    if (diff < best_diff || (diff == best_diff && ((bits & 1u) == 0u))) {
+      best = static_cast<uint8_t>(bits);
+      best_diff = diff;
+    }
+  }
+  return best;
+}
+
+int dpi_softmax_prob_to_fp8e4m3x4(bool enable, int score_bits, int max_bits, int denom_bits) {
+  if (!enable)
+    return 0;
+  float score = bits_to_float(static_cast<uint32_t>(score_bits));
+  float row_max = bits_to_float(static_cast<uint32_t>(max_bits));
+  float denom = bits_to_float(static_cast<uint32_t>(denom_bits));
+  float prob = denom == 0.0f ? 0.0f : expf(score - row_max) / denom;
+  uint32_t fp8 = float_to_fp8e4m3_positive(prob);
+  return static_cast<int>(fp8 | (fp8 << 8) | (fp8 << 16) | (fp8 << 24));
+}
+
 void dpi_fadd(bool enable, int dst_fmt, int64_t a, int64_t b, const svBitVecVal* frm, int64_t* result, svBitVecVal* fflags) {
   if (!enable) 
     return;

hw/dpi/float_dpi.vh

@@ -27,6 +27,10 @@ import "DPI-C" function void dpi_fnmsub(input logic enable, input int dst_fmt, i
 import "DPI-C" function void dpi_fdiv(input logic enable, input int dst_fmt, input longint a, input longint b, input bit[2:0] frm, output longint result, output bit[4:0] fflags);
 import "DPI-C" function void dpi_fsqrt(input logic enable, input int dst_fmt, input longint a, input bit[2:0] frm, output longint result, output bit[4:0] fflags);
 import "DPI-C" function void dpi_fexp(input logic enable, input int dst_fmt, input longint a, output longint result, output bit[4:0] fflags);
+import "DPI-C" function int dpi_f32_max(input logic enable, input int a_bits, input int b_bits);
+import "DPI-C" function int dpi_softmax_exp_acc(input logic enable, input int score_bits, input int max_bits, input int accum_bits);
+import "DPI-C" function int dpi_softmax_prob_to_f16x2(input logic enable, input int score_bits, input int max_bits, input int denom_bits);
+import "DPI-C" function int dpi_softmax_prob_to_fp8e4m3x4(input logic enable, input int score_bits, input int max_bits, input int denom_bits);
 
 import "DPI-C" function void dpi_ftoi(input logic enable, input int dst_fmt, input int src_fmt, input longint a, input bit[2:0] frm, output longint result, output bit[4:0] fflags);
 import "DPI-C" function void dpi_ftou(input logic enable, input int dst_fmt, input int src_fmt, input longint a, input bit[2:0] frm, output longint result, output bit[4:0] fflags);

hw/rtl/VX_define.vh

@@ -210,6 +210,7 @@
 `define INST_LSU_SD          4'b1011 // new for RV64I SD
 `define INST_LSU_TMEM_LD     4'b1100
 `define INST_LSU_TMEM_ST     4'b1101
+`define INST_LSU_TMEM_SOFTMAX 4'b1110
 `define INST_LSU_FENCE       4'b1111
 `define INST_LSU_BITS        4
 `define INST_LSU_FMT(op)     op[2:0]

hw/rtl/core/VX_decode.sv

@@ -550,6 +550,16 @@ module VX_decode  #(
                         end
                     end
                 `endif
+                    3'h3: begin
+                        if (func7 == 7'h30) begin
+                            ex_type = `EX_LSU;
+                            op_type = `INST_LSU_TMEM_SOFTMAX;
+                            use_rd = 1;
+                            `USED_IREG (rd);
+                            `USED_IREG (rs1);
+                            `USED_IREG (rs2);
+                        end
+                    end
                     default:;
                 endcase
             end

hw/rtl/core/VX_execute.sv

@@ -12,6 +12,9 @@
 // limitations under the License.
 
 `include "VX_define.vh"
+`ifdef SV_DPI
+`include "float_dpi.vh"
+`endif
 
 module VX_execute import VX_gpu_pkg::*; #(
     parameter CORE_ID = 0,
@@ -304,6 +307,7 @@ module VX_execute import VX_gpu_pkg::*; #(
     wire [`ISSUE_WIDTH-1:0] scalar_tmem_dispatch;
     wire [`ISSUE_WIDTH-1:0] scalar_tmem_ld_dispatch;
     wire [`ISSUE_WIDTH-1:0] scalar_tmem_st_dispatch;
+    wire [`ISSUE_WIDTH-1:0] scalar_tmem_softmax_dispatch;
     wire [`ISSUE_WIDTH-1:0] scalar_tmem_dispatch_ready;
     wire [`ISSUE_WIDTH-1:0] scalar_tmem_dispatch_fire;
     wire [`ISSUE_WIDTH-1:0] scalar_tmem_commit_ready;
@@ -314,7 +318,11 @@ module VX_execute import VX_gpu_pkg::*; #(
                                          && (lsu_dispatch_if[i].data.op_type == `INST_LSU_TMEM_LD);
         assign scalar_tmem_st_dispatch[i] = lsu_dispatch_if[i].valid
                                          && (lsu_dispatch_if[i].data.op_type == `INST_LSU_TMEM_ST);
-        assign scalar_tmem_dispatch[i] = scalar_tmem_ld_dispatch[i] || scalar_tmem_st_dispatch[i];
+        assign scalar_tmem_softmax_dispatch[i] = lsu_dispatch_if[i].valid
+                                             && (lsu_dispatch_if[i].data.op_type == `INST_LSU_TMEM_SOFTMAX);
+        assign scalar_tmem_dispatch[i] = scalar_tmem_ld_dispatch[i]
+                                      || scalar_tmem_st_dispatch[i]
+                                      || scalar_tmem_softmax_dispatch[i];
 
         assign scalar_mem_lsu_dispatch_if[i].valid = lsu_dispatch_if[i].valid && !scalar_tmem_dispatch[i];
         assign scalar_mem_lsu_dispatch_if[i].data = lsu_dispatch_if[i].data;
@@ -347,6 +355,16 @@ module VX_execute import VX_gpu_pkg::*; #(
     reg                     scalar_tmem_load_pending;
     reg                     scalar_tmem_store_pending;
     reg                     scalar_tmem_commit_pending;
+    reg                     scalar_tmem_softmax_active;
+    reg                     scalar_tmem_softmax_read_pending;
+    reg [1:0]               scalar_tmem_softmax_stage;
+    reg [5:0]               scalar_tmem_softmax_index;
+    reg [5:0]               scalar_tmem_softmax_read_index_r;
+    reg [8:0]               scalar_tmem_softmax_score_base_r;
+    reg [8:0]               scalar_tmem_softmax_p_base_r;
+    reg [31:0]              scalar_tmem_softmax_scores_r [0:31][0:`NUM_THREADS-1];
+    reg [31:0]              scalar_tmem_softmax_row_max_r [0:`NUM_THREADS-1];
+    reg [31:0]              scalar_tmem_softmax_denom_r [0:`NUM_THREADS-1];
     reg  [`UUID_WIDTH-1:0]  scalar_tmem_uuid_r;
     reg  [`NW_WIDTH-1:0]    scalar_tmem_wid_r;
     reg  [`NUM_THREADS-1:0] scalar_tmem_tmask_r;
@@ -361,6 +379,51 @@ module VX_execute import VX_gpu_pkg::*; #(
     wire [`ISSUE_WIDTH-1:0][`NUM_THREADS-1:0] scalar_tmem_req_tmask;
     wire [`ISSUE_WIDTH-1:0][`XLEN-1:0] scalar_tmem_req_pc;
     wire [`ISSUE_WIDTH-1:0][`NR_BITS-1:0] scalar_tmem_req_rd;
+    wire [`ISSUE_WIDTH-1:0][8:0] scalar_tmem_req_softmax_p_addr;
+
+    localparam [1:0] SCALAR_TMEM_SOFTMAX_MAX   = 2'd0;
+    localparam [1:0] SCALAR_TMEM_SOFTMAX_DENOM = 2'd1;
+    localparam [1:0] SCALAR_TMEM_SOFTMAX_WRITE = 2'd2;
+    localparam [5:0] SCALAR_TMEM_SOFTMAX_ROW_LAST  = 6'd31;
+    localparam [5:0] SCALAR_TMEM_SOFTMAX_TILE_LAST = 6'd63;
+
+    function automatic [31:0] scalar_tmem_f32_max;
+        input [31:0] a_bits;
+        input [31:0] b_bits;
+        begin
+        `ifdef SV_DPI
+            scalar_tmem_f32_max = 32'(dpi_f32_max(1'b1, int'(a_bits), int'(b_bits)));
+        `else
+            scalar_tmem_f32_max = a_bits;
+        `endif
+        end
+    endfunction
+
+    function automatic [31:0] scalar_tmem_softmax_exp_acc;
+        input [31:0] score_bits;
+        input [31:0] max_bits;
+        input [31:0] accum_bits;
+        begin
+        `ifdef SV_DPI
+            scalar_tmem_softmax_exp_acc = 32'(dpi_softmax_exp_acc(1'b1, int'(score_bits), int'(max_bits), int'(accum_bits)));
+        `else
+            scalar_tmem_softmax_exp_acc = accum_bits;
+        `endif
+        end
+    endfunction
+
+    function automatic [31:0] scalar_tmem_softmax_prob_pack;
+        input [31:0] score_bits;
+        input [31:0] max_bits;
+        input [31:0] denom_bits;
+        begin
+        `ifdef SV_DPI
+            scalar_tmem_softmax_prob_pack = 32'(dpi_softmax_prob_to_f16x2(1'b1, int'(score_bits), int'(max_bits), int'(denom_bits)));
+        `else
+            scalar_tmem_softmax_prob_pack = '0;
+        `endif
+        end
+    endfunction
 
     assign scalar_tmem_pending = |scalar_tmem_dispatch;
 
@@ -375,11 +438,30 @@ module VX_execute import VX_gpu_pkg::*; #(
     end
 
     wire scalar_tmem_grant_valid = scalar_tmem_pending && !scalar_tmem_load_pending
-                                && !scalar_tmem_store_pending && !scalar_tmem_commit_pending;
+                                && !scalar_tmem_store_pending && !scalar_tmem_commit_pending
+                                && !scalar_tmem_softmax_active;
     wire scalar_tmem_grant_is_load = |(scalar_tmem_grant & scalar_tmem_ld_dispatch);
     wire scalar_tmem_grant_is_store = |(scalar_tmem_grant & scalar_tmem_st_dispatch);
-    wire scalar_tmem_req_ready = scalar_tmem_grant_is_load ? scalar_tmem_rready : scalar_tmem_wready;
+    wire scalar_tmem_grant_is_softmax = |(scalar_tmem_grant & scalar_tmem_softmax_dispatch);
+    wire scalar_tmem_req_ready = scalar_tmem_grant_is_softmax ? 1'b1 :
+                                 scalar_tmem_grant_is_load ? scalar_tmem_rready : scalar_tmem_wready;
     wire scalar_tmem_req_fire = scalar_tmem_grant_valid && scalar_tmem_req_ready;
+    wire scalar_tmem_softmax_read_valid = scalar_tmem_softmax_active
+                                       && !scalar_tmem_softmax_read_pending
+                                       && (scalar_tmem_softmax_stage != SCALAR_TMEM_SOFTMAX_WRITE);
+    wire scalar_tmem_softmax_write_valid = scalar_tmem_softmax_active
+                                        && (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_WRITE);
+    wire scalar_tmem_softmax_read_issue = scalar_tmem_softmax_read_valid && scalar_tmem_rready;
+    wire scalar_tmem_softmax_write_issue = scalar_tmem_softmax_write_valid && scalar_tmem_wready;
+    wire [8:0] scalar_tmem_softmax_read_addr = scalar_tmem_softmax_score_base_r
+                                             + 9'(scalar_tmem_softmax_index[4:0]);
+    wire [8:0] scalar_tmem_softmax_write_addr = scalar_tmem_softmax_p_base_r
+                                              + 9'(scalar_tmem_softmax_index);
+    wire [`NUM_THREADS*`XLEN-1:0] scalar_tmem_softmax_wdata;
+    wire scalar_tmem_softmax_commit = scalar_tmem_commit_pending
+                                    && !scalar_tmem_load_pending
+                                    && !scalar_tmem_store_pending;
+    wire [`XLEN-1:0] scalar_tmem_softmax_token = `XLEN'(scalar_tmem_softmax_p_base_r);
 
     for (genvar i = 0; i < `ISSUE_WIDTH; ++i) begin : g_scalar_tmem_ready
         assign scalar_tmem_dispatch_ready[i] = scalar_tmem_grant_valid && scalar_tmem_grant[i] && scalar_tmem_req_ready;
@@ -390,15 +472,19 @@ module VX_execute import VX_gpu_pkg::*; #(
         assign scalar_tmem_req_tmask[i] = lsu_dispatch_if[i].data.tmask;
         assign scalar_tmem_req_pc[i] = lsu_dispatch_if[i].data.PC;
         assign scalar_tmem_req_rd[i] = lsu_dispatch_if[i].data.rd;
+        assign scalar_tmem_req_softmax_p_addr[i] = lsu_dispatch_if[i].data.rs2_data[0][8:0];
 
         assign scalar_tmem_commit_if[i].valid = scalar_tmem_commit_pending && scalar_tmem_grant_r[i];
         assign scalar_tmem_commit_if[i].data.uuid = scalar_tmem_uuid_r;
         assign scalar_tmem_commit_if[i].data.wid = scalar_tmem_wid_r;
         assign scalar_tmem_commit_if[i].data.tmask = scalar_tmem_tmask_r;
         assign scalar_tmem_commit_if[i].data.PC = scalar_tmem_pc_r;
-        assign scalar_tmem_commit_if[i].data.wb = scalar_tmem_load_pending;
+        assign scalar_tmem_commit_if[i].data.wb = scalar_tmem_load_pending || scalar_tmem_softmax_commit;
-        assign scalar_tmem_commit_if[i].data.rd = scalar_tmem_load_pending ? scalar_tmem_rd_r : '0;
+        assign scalar_tmem_commit_if[i].data.rd = (scalar_tmem_load_pending || scalar_tmem_softmax_commit) ? scalar_tmem_rd_r : '0;
-        assign scalar_tmem_commit_if[i].data.data = scalar_tmem_rdata_valid_r ? scalar_tmem_rdata_r : scalar_tmem_rdata;
+        assign scalar_tmem_commit_if[i].data.data =
+            scalar_tmem_load_pending
+                ? (scalar_tmem_rdata_valid_r ? scalar_tmem_rdata_r : scalar_tmem_rdata)
+                : {`NUM_THREADS{scalar_tmem_softmax_token}};
         assign scalar_tmem_commit_if[i].data.tensor = 1'b0;
         assign scalar_tmem_commit_if[i].data.pid = '0;
         assign scalar_tmem_commit_if[i].data.sop = 1'b1;
@@ -417,9 +503,49 @@ module VX_execute import VX_gpu_pkg::*; #(
         end
     endfunction
 
+    for (genvar lane = 0; lane < `NUM_THREADS; ++lane) begin : g_scalar_tmem_softmax_wdata
+        assign scalar_tmem_softmax_wdata[lane * `XLEN +: `XLEN] =
+            scalar_tmem_softmax_prob_pack(
+                scalar_tmem_softmax_scores_r[scalar_tmem_softmax_index[4:0]][lane],
+                scalar_tmem_softmax_row_max_r[lane],
+                scalar_tmem_softmax_denom_r[lane]);
+    end
+
+`ifdef DBG_TRACE_CORE_PIPELINE_VCS
+    always @(posedge clk) begin
+        if (!reset && ($time > `TRACE_STARTTIME) && (CORE_ID == 0)) begin
+            if (scalar_tmem_req_fire && scalar_tmem_grant_is_softmax) begin
+                `TRACE(1, ("%d: core%0d-scalar-tmem-softmax-start: grant=%b wid=%0d PC=0x%0h rd=%0d score=%0d prob=%0d tmask=%b (#%0d)\n",
+                    $time, CORE_ID, scalar_tmem_grant, scalar_tmem_req_wid[0], scalar_tmem_req_pc[0],
+                    scalar_tmem_req_rd[0], scalar_tmem_req_addr[0], scalar_tmem_req_softmax_p_addr[0],
+                    scalar_tmem_req_tmask[0], scalar_tmem_req_uuid[0]));
+            end
+            if (scalar_tmem_softmax_read_issue && (scalar_tmem_softmax_index == 6'd0 || scalar_tmem_softmax_index == SCALAR_TMEM_SOFTMAX_ROW_LAST)) begin
+                `TRACE(1, ("%d: core%0d-scalar-tmem-softmax-read: stage=%0d index=%0d addr=%0d rready=%b PC=0x%0h rd=%0d\n",
+                    $time, CORE_ID, scalar_tmem_softmax_stage, scalar_tmem_softmax_index,
+                    scalar_tmem_softmax_read_addr, scalar_tmem_rready, scalar_tmem_pc_r, scalar_tmem_rd_r));
+            end
+            if (scalar_tmem_softmax_write_issue && (scalar_tmem_softmax_index == 6'd0 || scalar_tmem_softmax_index == SCALAR_TMEM_SOFTMAX_TILE_LAST)) begin
+                `TRACE(1, ("%d: core%0d-scalar-tmem-softmax-write: index=%0d addr=%0d wready=%b PC=0x%0h rd=%0d\n",
+                    $time, CORE_ID, scalar_tmem_softmax_index, scalar_tmem_softmax_write_addr,
+                    scalar_tmem_wready, scalar_tmem_pc_r, scalar_tmem_rd_r));
+            end
+            if (scalar_tmem_commit_pending && scalar_tmem_pc_r == 32'h80000318) begin
+                `TRACE(1, ("%d: core%0d-scalar-tmem-softmax-commit-pending: grant_r=%b ready=%b valid0=%b wb0=%0d rd=%0d PC=0x%0h token=0x%0h\n",
+                    $time, CORE_ID, scalar_tmem_grant_r, scalar_tmem_commit_ready,
+                    scalar_tmem_commit_if[0].valid, scalar_tmem_commit_if[0].data.wb,
+                    scalar_tmem_commit_if[0].data.rd, scalar_tmem_commit_if[0].data.PC,
+                    scalar_tmem_softmax_token));
+            end
+        end
+    end
+`endif
+
     always @(*) begin
-        scalar_tmem_ren = scalar_tmem_grant_valid && scalar_tmem_grant_is_load;
+        scalar_tmem_ren = (scalar_tmem_grant_valid && scalar_tmem_grant_is_load)
-        scalar_tmem_wen = scalar_tmem_grant_valid && scalar_tmem_grant_is_store;
+                       || scalar_tmem_softmax_read_valid;
+        scalar_tmem_wen = (scalar_tmem_grant_valid && scalar_tmem_grant_is_store)
+                       || scalar_tmem_softmax_write_valid;
         scalar_tmem_raddr = '0;
         scalar_tmem_waddr = '0;
         scalar_tmem_wdata = '0;
@@ -433,6 +559,14 @@ module VX_execute import VX_gpu_pkg::*; #(
                                  & {(`NUM_THREADS * (`XLEN / 8)){scalar_tmem_grant_is_store}};
             end
         end
+        if (scalar_tmem_softmax_read_valid) begin
+            scalar_tmem_raddr = scalar_tmem_softmax_read_addr;
+        end
+        if (scalar_tmem_softmax_write_valid) begin
+            scalar_tmem_waddr = scalar_tmem_softmax_write_addr;
+            scalar_tmem_wdata = scalar_tmem_softmax_wdata;
+            scalar_tmem_mask = scalar_tmem_expand_tmask(scalar_tmem_tmask_r);
+        end
     end
 
     always @(posedge clk) begin
@@ -441,6 +575,13 @@ module VX_execute import VX_gpu_pkg::*; #(
             scalar_tmem_load_pending <= 1'b0;
             scalar_tmem_store_pending <= 1'b0;
             scalar_tmem_commit_pending <= 1'b0;
+            scalar_tmem_softmax_active <= 1'b0;
+            scalar_tmem_softmax_read_pending <= 1'b0;
+            scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_MAX;
+            scalar_tmem_softmax_index <= '0;
+            scalar_tmem_softmax_read_index_r <= '0;
+            scalar_tmem_softmax_score_base_r <= '0;
+            scalar_tmem_softmax_p_base_r <= '0;
             scalar_tmem_uuid_r <= '0;
             scalar_tmem_wid_r <= '0;
             scalar_tmem_tmask_r <= '0;
@@ -453,7 +594,7 @@ module VX_execute import VX_gpu_pkg::*; #(
                 scalar_tmem_grant_r <= scalar_tmem_grant;
                 scalar_tmem_load_pending <= scalar_tmem_grant_is_load;
                 scalar_tmem_store_pending <= scalar_tmem_grant_is_store;
-                scalar_tmem_commit_pending <= 1'b1;
+                scalar_tmem_commit_pending <= !scalar_tmem_grant_is_softmax;
                 scalar_tmem_rdata_valid_r <= 1'b0;
                 for (integer i = 0; i < `ISSUE_WIDTH; ++i) begin
                     if (scalar_tmem_grant[i]) begin
@@ -462,6 +603,19 @@ module VX_execute import VX_gpu_pkg::*; #(
                         scalar_tmem_tmask_r <= scalar_tmem_req_tmask[i];
                         scalar_tmem_pc_r <= scalar_tmem_req_pc[i];
                         scalar_tmem_rd_r <= scalar_tmem_req_rd[i];
+                        if (scalar_tmem_grant_is_softmax) begin
+                            scalar_tmem_softmax_active <= 1'b1;
+                            scalar_tmem_softmax_read_pending <= 1'b0;
+                            scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_MAX;
+                            scalar_tmem_softmax_index <= '0;
+                            scalar_tmem_softmax_read_index_r <= '0;
+                            scalar_tmem_softmax_score_base_r <= scalar_tmem_req_addr[i];
+                            scalar_tmem_softmax_p_base_r <= scalar_tmem_req_softmax_p_addr[i];
+                            for (integer lane = 0; lane < `NUM_THREADS; ++lane) begin
+                                scalar_tmem_softmax_row_max_r[lane] <= '0;
+                                scalar_tmem_softmax_denom_r[lane] <= '0;
+                            end
+                        end
                     end
                 end
             end else if (scalar_tmem_load_pending && scalar_tmem_commit_pending && !scalar_tmem_rdata_valid_r) begin
@@ -469,6 +623,58 @@ module VX_execute import VX_gpu_pkg::*; #(
                 scalar_tmem_rdata_valid_r <= 1'b1;
             end
 
+            if (scalar_tmem_softmax_read_issue) begin
+                scalar_tmem_softmax_read_pending <= 1'b1;
+                scalar_tmem_softmax_read_index_r <= scalar_tmem_softmax_index;
+            end else if (scalar_tmem_softmax_read_pending) begin
+                scalar_tmem_softmax_read_pending <= 1'b0;
+                for (integer lane = 0; lane < `NUM_THREADS; ++lane) begin
+                    scalar_tmem_softmax_scores_r[scalar_tmem_softmax_read_index_r[4:0]][lane]
+                        <= scalar_tmem_rdata[lane * `XLEN +: `XLEN];
+                    if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_MAX) begin
+                        scalar_tmem_softmax_row_max_r[lane] <=
+                            (scalar_tmem_softmax_read_index_r == '0)
+                                ? scalar_tmem_rdata[lane * `XLEN +: `XLEN]
+                                : scalar_tmem_f32_max(scalar_tmem_softmax_row_max_r[lane],
+                                                      scalar_tmem_rdata[lane * `XLEN +: `XLEN]);
+                    end else if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_DENOM) begin
+                        scalar_tmem_softmax_denom_r[lane] <=
+                            scalar_tmem_softmax_exp_acc(
+                                scalar_tmem_rdata[lane * `XLEN +: `XLEN],
+                                scalar_tmem_softmax_row_max_r[lane],
+                                scalar_tmem_softmax_denom_r[lane]);
+                    end
+                end
+
+                if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_MAX) begin
+                    if (scalar_tmem_softmax_read_index_r == SCALAR_TMEM_SOFTMAX_ROW_LAST) begin
+                        scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_DENOM;
+                        scalar_tmem_softmax_index <= '0;
+                        for (integer lane = 0; lane < `NUM_THREADS; ++lane) begin
+                            scalar_tmem_softmax_denom_r[lane] <= 32'h00000000;
+                        end
+                    end else begin
+                        scalar_tmem_softmax_index <= scalar_tmem_softmax_read_index_r + 6'd1;
+                    end
+                end else if (scalar_tmem_softmax_stage == SCALAR_TMEM_SOFTMAX_DENOM) begin
+                    if (scalar_tmem_softmax_read_index_r == SCALAR_TMEM_SOFTMAX_ROW_LAST) begin
+                        scalar_tmem_softmax_stage <= SCALAR_TMEM_SOFTMAX_WRITE;
+                        scalar_tmem_softmax_index <= '0;
+                    end else begin
+                        scalar_tmem_softmax_index <= scalar_tmem_softmax_read_index_r + 6'd1;
+                    end
+                end
+            end
+
+            if (scalar_tmem_softmax_write_issue) begin
+                if (scalar_tmem_softmax_index == SCALAR_TMEM_SOFTMAX_TILE_LAST) begin
+                    scalar_tmem_softmax_active <= 1'b0;
+                    scalar_tmem_commit_pending <= 1'b1;
+                end else begin
+                    scalar_tmem_softmax_index <= scalar_tmem_softmax_index + 6'd1;
+                end
+            end
+
             if (scalar_tmem_commit_pending && (|(scalar_tmem_grant_r & scalar_tmem_commit_ready))) begin
                 scalar_tmem_grant_r <= '0;
                 scalar_tmem_load_pending <= 1'b0;