fix: update fp16 softmax case25 score expectation

kernels/blackwell_fp8_e4m3/test_case_01_gemm/Makefile

@@ -4,7 +4,7 @@ VX_SRCS = kernel.cpp
 VX_INCLUDES = fp8_common.hpp
 OPTS ?= -n1
 
-include ../common.mk
+include ../../common.mk
 
-args.bin input.a.bin input.b.bin input.c.bin: ../wu_arch_cases/zero.bin
+args.bin input.a.bin input.b.bin input.c.bin: ../../wu_arch_cases/zero.bin
 	cp $< $@

kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/Makefile

@@ -0,0 +1,13 @@
+PROJECT = blackwell_fp8_e4m3_softmax_vector
+
+VX_SRCS = kernel.cpp
+VX_INCLUDES = fp8_common.hpp ../../wu_arch_cases/common_wu_blackwell_fa.h ../../wu_arch_cases/common_wu_min.h
+OPTS ?= -n1
+
+VORTEX_KN_PATH ?= $(realpath ../../../lib)
+GEMMINI_SW_PATH ?= $(realpath ../../../lib/gemmini)
+
+include ../../common.mk
+
+args.bin input.a.bin input.b.bin input.c.bin: ../../wu_arch_cases/zero.bin
+	cp $< $@

kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/README.md

@@ -0,0 +1,20 @@
+# test_case_02_scalar_tmem_softmax_vector
+
+FP8 E4M3 version of the scalar TMEM softmax vector handoff test.
+
+The test keeps the FP16-oriented `wu_arch_cases/case25_scalar_tmem_softmax_vector`
+unchanged. This case uses FP8 E4M3 TMEM A and SMEM B layout throughout:
+
+- first BWGMMA computes uniform FP32 scores: `Q = 1.0`, `K = 1.0`, `C = 0.0`,
+  expected score `32.0`;
+- scalar TMEM softmax consumes the FP32 C score tile and writes FP8 E4M3
+  probabilities into TMEM A;
+- uniform softmax probability is `1/32`, exactly representable as E4M3 `0x10`;
+- second BWGMMA consumes that FP8 P tile with `V = 1.0`, expected FP32 output
+  `1.0`.
+
+Build:
+
+```bash
+make -C /home/lzd/wu/wuarch/virgo-kernels/kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector
+```

kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/fp8_common.hpp

@@ -0,0 +1,41 @@
+#ifndef BLACKWELL_FP8_E4M3_COMMON_HPP
+#define BLACKWELL_FP8_E4M3_COMMON_HPP
+
+#include <stdint.h>
+#include <vx_intrinsics.h>
+
+#define WU_FP8_E4M3_ZERO 0x00u
+#define WU_FP8_E4M3_ONE_OVER_32 0x10u
+#define WU_FP8_E4M3_HALF 0x30u
+#define WU_FP8_E4M3_ONE 0x38u
+#define WU_FP8_E4M3_TWO 0x40u
+
+#define WU_FP8_PACK4(a, b, c, d)                                             \
+  ((((uint32_t)(a) & 0xffu) << 0) | (((uint32_t)(b) & 0xffu) << 8) |          \
+   (((uint32_t)(c) & 0xffu) << 16) | (((uint32_t)(d) & 0xffu) << 24))
+
+#define WU_FP8_REP2(x) x, x
+#define WU_FP8_REP4(x) WU_FP8_REP2(x), WU_FP8_REP2(x)
+
+#define WU_FP8_E4M3_ONE_PACKED                                               \
+  WU_FP8_PACK4(WU_FP8_E4M3_ONE, WU_FP8_E4M3_ONE, WU_FP8_E4M3_ONE,            \
+               WU_FP8_E4M3_ONE)
+#define WU_FP8_E4M3_ONE_OVER_32_PACKED                                       \
+  WU_FP8_PACK4(WU_FP8_E4M3_ONE_OVER_32, WU_FP8_E4M3_ONE_OVER_32,             \
+               WU_FP8_E4M3_ONE_OVER_32, WU_FP8_E4M3_ONE_OVER_32)
+
+static inline void wu_bwgmma_fp8(uint32_t addr_tmem_c, uint32_t addr_tmem_a,
+                                 uint32_t addr_smem_b) {
+  asm volatile(".insn r %0, 0, 0, %1, %2, %3"
+               :
+               : "i"(RISCV_CUSTOM3), "r"(addr_tmem_c), "r"(addr_tmem_a),
+                 "r"(addr_smem_b)
+               : "memory");
+}
+
+static inline void wu_bwgmma_wait() {
+  asm volatile(".insn r %0, 1, 0, x0, x0, x0" :: "i"(RISCV_CUSTOM3)
+               : "memory");
+}
+
+#endif

kernels/blackwell_fp8_e4m3/test_case_02_scalar_tmem_softmax_vector/kernel.cpp

@@ -0,0 +1,218 @@
+#define WU_CASE_WAIT_SPIN 16384u
+
+#include "fp8_common.hpp"
+#include "../../wu_arch_cases/common_wu_blackwell_fa.h"
+
+#define WU_FP8_SOFTMAX_SCORE_READY_BASE 0xc500u
+#define WU_FP8_SOFTMAX_P_READY_BASE 0xc600u
+#define WU_FP8_SOFTMAX_DONE_BASE 0xc700u
+
+#define WU_FP8_SOFTMAX_FP32_ZERO 0x00000000u
+#define WU_FP8_SOFTMAX_FP32_ONE 0x3f800000u
+#define WU_FP8_SOFTMAX_FP32_THIRTY_TWO 0x42000000u
+
+extern "C" {
+volatile uint32_t g_fp8_softmax_q_frag[4] __attribute__((aligned(16))) = {
+    WU_FP8_REP4(WU_FP8_E4M3_ONE_PACKED)};
+volatile uint32_t g_fp8_softmax_zero_frag[4] __attribute__((aligned(16))) = {
+    WU_FP8_SOFTMAX_FP32_ZERO, WU_FP8_SOFTMAX_FP32_ZERO,
+    WU_FP8_SOFTMAX_FP32_ZERO, WU_FP8_SOFTMAX_FP32_ZERO};
+volatile uint32_t g_fp8_softmax_score_ready __attribute__((aligned(16)));
+volatile uint32_t g_fp8_softmax_p_ready __attribute__((aligned(16)));
+volatile uint32_t g_fp8_softmax_done __attribute__((aligned(16)));
+volatile uint32_t g_fp8_softmax_score_bits[NUM_THREADS]
+    __attribute__((aligned(16)));
+volatile uint32_t g_fp8_softmax_p_bits[NUM_THREADS]
+    __attribute__((aligned(16)));
+volatile uint32_t g_fp8_softmax_out[WU_BW_OUT_WORDS]
+    __attribute__((aligned(16)));
+}
+
+#undef WU_FP8_REP2
+#undef WU_FP8_REP4
+
+static inline int wu_fp8_softmax_wait_status(volatile uint32_t *status,
+                                             uint32_t expected) {
+  for (uint32_t spin = 0; spin < WU_CASE_WAIT_SPIN; ++spin) {
+    if (*status == expected) {
+      return 0;
+    }
+  }
+  return 1;
+}
+
+extern "C" void __attribute__((naked, noinline, used))
+tensor_fp8_softmax_worker() {
+  asm volatile(
+      "csrr x5, %[csr_wid]\n\t"
+      "addi x1, x5, -%[num_scalar_warps]\n\t"
+      "slli x1, x1, 11\n\t"
+      "addi x2, x1, %[c_offset]\n\t"
+      "la x3, g_fp8_softmax_q_frag\n\t"
+      "li x7, 0\n\t"
+      "1:\n\t"
+      "add x4, x1, x7\n\t"
+      ".insn r %[custom3], 2, 0, x0, x4, x3\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 1b\n\t"
+      "la x3, g_fp8_softmax_zero_frag\n\t"
+      "li x7, 0\n\t"
+      "2:\n\t"
+      "add x4, x2, x7\n\t"
+      ".insn r %[custom3], 2, 0, x0, x4, x3\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 2b\n\t"
+      ".insn r %[custom3], 3, 0, x0, x0, x0\n\t"
+      "li x4, %[smem_base]\n\t"
+      ".insn r %[custom3], 0, 0, x2, x1, x4\n\t"
+      ".insn r %[custom3], 1, 0, x0, x0, x0\n\t"
+      "la x7, g_fp8_softmax_score_ready\n\t"
+      "li x6, %[score_ready]\n\t"
+      "sw x6, 0(x7)\n\t"
+      "la x7, g_fp8_softmax_p_ready\n\t"
+      "li x4, %[p_ready]\n\t"
+      "3:\n\t"
+      "lw x6, 0(x7)\n\t"
+      "bne x6, x4, 3b\n\t"
+      "la x3, g_fp8_softmax_zero_frag\n\t"
+      "li x7, 0\n\t"
+      "4:\n\t"
+      "add x4, x2, x7\n\t"
+      ".insn r %[custom3], 2, 0, x0, x4, x3\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 4b\n\t"
+      ".insn r %[custom3], 3, 0, x0, x0, x0\n\t"
+      "li x4, %[smem_base]\n\t"
+      ".insn r %[custom3], 0, 0, x2, x1, x4\n\t"
+      ".insn r %[custom3], 1, 0, x0, x0, x0\n\t"
+      "la x3, g_fp8_softmax_out\n\t"
+      "li x7, 0\n\t"
+      "5:\n\t"
+      "add x4, x2, x7\n\t"
+      "add x6, x3, x7\n\t"
+      ".insn r %[custom3], 6, 0, x0, x4, x6\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 5b\n\t"
+      ".insn r %[custom3], 3, 0, x0, x0, x0\n\t"
+      "la x7, g_fp8_softmax_done\n\t"
+      "li x6, %[done]\n\t"
+      "sw x6, 0(x7)\n\t"
+      ".insn r %[custom0], 0, 0, x0, x0, x0\n\t"
+      "6: j 6b\n\t"
+      :
+      : [csr_wid] "i"(VX_CSR_WARP_ID), [custom0] "i"(RISCV_CUSTOM0),
+        [custom3] "i"(RISCV_CUSTOM3),
+        [num_scalar_warps] "i"(NUM_SCALAR_WARPS),
+        [c_offset] "i"(WU_BW_TMEM_C_BYTE_OFFSET),
+        [tile_bytes] "i"(WU_BW_TMEM_TILE_BYTES),
+        [smem_base] "i"(WU_BW_DEV_SMEM_START_ADDR),
+        [score_ready] "i"(WU_FP8_SOFTMAX_SCORE_READY_BASE),
+        [p_ready] "i"(WU_FP8_SOFTMAX_P_READY_BASE),
+        [done] "i"(WU_FP8_SOFTMAX_DONE_BASE)
+      : "memory");
+}
+
+extern "C" int wu_main() {
+  if (vx_core_id() != 0 || vx_warp_id() != 0) {
+    return 0;
+  }
+
+  const uint32_t tid = wu_tid();
+  if (tid == 0) {
+    wu_case_reset();
+    g_fp8_softmax_score_ready = 0;
+    g_fp8_softmax_p_ready = 0;
+    g_fp8_softmax_done = 0;
+    for (uint32_t i = 0; i < NUM_THREADS; ++i) {
+      g_fp8_softmax_score_bits[i] = 0;
+      g_fp8_softmax_p_bits[i] = 0;
+    }
+    for (uint32_t i = 0; i < WU_BW_OUT_WORDS; ++i) {
+      g_fp8_softmax_out[i] = 0;
+    }
+    wu_bw_fill_smem_tile(
+        reinterpret_cast<volatile uint32_t *>(WU_BW_DEV_SMEM_START_ADDR),
+        WU_FP8_E4M3_ONE_PACKED);
+    vx_spawn_tensor(1u << NUM_SCALAR_WARPS, tensor_fp8_softmax_worker);
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (wu_fp8_softmax_wait_status(&g_fp8_softmax_score_ready,
+                                 WU_FP8_SOFTMAX_SCORE_READY_BASE) != 0) {
+    if (tid == 0) {
+      g_case_mem[1] = 0x81u;
+    }
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  const uint32_t c_frag =
+      wu_bw_tmem_c_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES;
+  const uint32_t observed_score = wu_bw_scalar_tmem_ld(c_frag);
+  g_fp8_softmax_score_bits[tid] = observed_score;
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (tid == 0 && g_case_mem[1] == 0 &&
+      observed_score != WU_FP8_SOFTMAX_FP32_THIRTY_TWO) {
+    g_aux[0] = observed_score;
+    g_case_mem[1] = 0x82u;
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  uint32_t softmax_token = 0;
+  const uint32_t p_frag =
+      wu_bw_tmem_a_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES;
+  if (g_case_mem[1] == 0) {
+    vx_tmc(wu_bw_all_lanes_mask());
+    softmax_token = wu_bw_scalar_tmem_softmax(c_frag, p_frag);
+    vx_tmc_one();
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  const uint32_t observed_p = wu_bw_scalar_tmem_ld(p_frag);
+  g_fp8_softmax_p_bits[tid] = observed_p;
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (tid == 0 && g_case_mem[1] == 0) {
+    if (softmax_token != p_frag) {
+      g_aux[0] = softmax_token;
+      g_aux[1] = p_frag;
+      g_case_mem[1] = 0x85u;
+    } else if (observed_p != WU_FP8_E4M3_ONE_OVER_32_PACKED) {
+      g_aux[0] = observed_p;
+      g_aux[1] = WU_FP8_E4M3_ONE_OVER_32_PACKED;
+      g_case_mem[1] = 0x86u;
+    } else {
+      g_fp8_softmax_p_ready = WU_FP8_SOFTMAX_P_READY_BASE;
+    }
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (tid == 0) {
+    if (g_case_mem[1] == 0 &&
+        wu_fp8_softmax_wait_status(&g_fp8_softmax_done,
+                                   WU_FP8_SOFTMAX_DONE_BASE) != 0) {
+      g_case_mem[1] = 0x83u;
+    }
+    if (g_case_mem[1] == 0) {
+      volatile uint32_t bad_actual = 0;
+      const uint32_t bad =
+          wu_bw_verify_constant(g_fp8_softmax_out, WU_BW_OUT_WORDS,
+                                WU_FP8_SOFTMAX_FP32_ONE, &bad_actual);
+      if (bad != WU_BW_OUT_WORDS) {
+        g_aux[0] = bad;
+        g_aux[1] = bad_actual;
+        g_case_mem[1] = 0x84u;
+      }
+    }
+    if (g_case_mem[1] != 0) {
+      wu_case_fail(g_case_mem[1]);
+      return 1;
+    }
+    wu_case_pass();
+  }
+  return 0;
+}

kernels/wu_arch_cases/Makefile

@@ -25,7 +25,8 @@ CASES := \
 	case21_moe_gating \
 	case22_gemm_silu \
 	case23_softmax_only \
-	case24_flash_sw_pipeline
+	case24_flash_sw_pipeline \
+	case25_scalar_tmem_softmax_vector
 
 SMOKE_CASES := \
 	case00_boot_scalar \

kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/Makefile

@@ -0,0 +1,3 @@
+PROJECT = case25_scalar_tmem_softmax_vector
+
+include ../case.mk

kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/README.md

@@ -0,0 +1,11 @@
+# case25_scalar_tmem_softmax_vector
+
+Validates the scalar-subcore vector softmax TMEM path.
+
+The tensor warp writes one score tile into its TMEM C partition, scalar warp 0
+executes `wu_bw_scalar_tmem_softmax()` to read the C tile and write packed fp16
+probabilities into TMEM A.  The softmax instruction returns the destination
+fragment as a completion token only after the P tile has been written; scalar
+warp 0 uses that token before publishing `p_ready`, then the tensor warp
+consumes that P tile through BWGMMA.  The case uses constant scores so the
+expected probability is uniform `1/32` and the final PV output is fp32 `1.0`.

kernels/wu_arch_cases/case25_scalar_tmem_softmax_vector/kernel.cpp

@@ -0,0 +1,198 @@
+#define WU_CASE_WAIT_SPIN 16384u
+
+#include "../common_wu_blackwell_fa.h"
+
+#define WU_CASE25_SCORE_READY_BASE 0xb500u
+#define WU_CASE25_P_READY_BASE 0xb600u
+#define WU_CASE25_DONE_BASE 0xb700u
+#define WU_CASE25_FP32_ZERO 0x00000000u
+#define WU_CASE25_FP32_THIRTY_TWO 0x42000000u
+
+extern "C" {
+volatile uint32_t g_case25_q_row[4] __attribute__((aligned(16))) = {
+    WU_BW_FP16_ONE_PACKED, WU_BW_FP16_ONE_PACKED, WU_BW_FP16_ONE_PACKED,
+    WU_BW_FP16_ONE_PACKED};
+volatile uint32_t g_case25_zero_row[4] __attribute__((aligned(16))) = {
+    WU_CASE25_FP32_ZERO, WU_CASE25_FP32_ZERO, WU_CASE25_FP32_ZERO,
+    WU_CASE25_FP32_ZERO};
+volatile uint32_t g_case25_score_ready __attribute__((aligned(16)));
+volatile uint32_t g_case25_p_ready __attribute__((aligned(16)));
+volatile uint32_t g_case25_done __attribute__((aligned(16)));
+volatile uint32_t g_case25_score_bits[NUM_THREADS] __attribute__((aligned(16)));
+volatile uint32_t g_case25_out[WU_BW_OUT_WORDS] __attribute__((aligned(16)));
+}
+
+static inline int wu_case25_wait_status(volatile uint32_t *status,
+                                        uint32_t expected) {
+  for (uint32_t spin = 0; spin < WU_CASE_WAIT_SPIN; ++spin) {
+    if (*status == expected) {
+      return 0;
+    }
+  }
+  return 1;
+}
+
+extern "C" void __attribute__((naked, noinline, used)) tensor_case25_worker() {
+  asm volatile(
+      "csrr x5, %[csr_wid]\n\t"
+      "addi x1, x5, -%[num_scalar_warps]\n\t"
+      "slli x1, x1, 11\n\t"
+      "addi x2, x1, %[c_offset]\n\t"
+      "la x3, g_case25_q_row\n\t"
+      "li x7, 0\n\t"
+      "1:\n\t"
+      "add x4, x1, x7\n\t"
+      ".insn r %[custom3], 2, 0, x0, x4, x3\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 1b\n\t"
+      "la x3, g_case25_zero_row\n\t"
+      "li x7, 0\n\t"
+      "2:\n\t"
+      "add x4, x2, x7\n\t"
+      ".insn r %[custom3], 2, 0, x0, x4, x3\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 2b\n\t"
+      ".insn r %[custom3], 3, 0, x0, x0, x0\n\t"
+      "li x4, %[smem_base]\n\t"
+      ".insn r %[custom3], 0, 0, x2, x1, x4\n\t"
+      ".insn r %[custom3], 1, 0, x0, x0, x0\n\t"
+      "la x7, g_case25_score_ready\n\t"
+      "li x6, %[score_ready]\n\t"
+      "sw x6, 0(x7)\n\t"
+      "la x7, g_case25_p_ready\n\t"
+      "li x4, %[p_ready]\n\t"
+      "3:\n\t"
+      "lw x6, 0(x7)\n\t"
+      "bne x6, x4, 3b\n\t"
+      "la x3, g_case25_zero_row\n\t"
+      "li x7, 0\n\t"
+      "4:\n\t"
+      "add x4, x2, x7\n\t"
+      ".insn r %[custom3], 2, 0, x0, x4, x3\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 4b\n\t"
+      ".insn r %[custom3], 3, 0, x0, x0, x0\n\t"
+      "li x4, %[smem_base]\n\t"
+      ".insn r %[custom3], 0, 0, x2, x1, x4\n\t"
+      ".insn r %[custom3], 1, 0, x0, x0, x0\n\t"
+      "la x3, g_case25_out\n\t"
+      "li x7, 0\n\t"
+      "5:\n\t"
+      "add x4, x2, x7\n\t"
+      "add x6, x3, x7\n\t"
+      ".insn r %[custom3], 6, 0, x0, x4, x6\n\t"
+      "addi x7, x7, 16\n\t"
+      "li x4, %[tile_bytes]\n\t"
+      "blt x7, x4, 5b\n\t"
+      ".insn r %[custom3], 3, 0, x0, x0, x0\n\t"
+      "la x7, g_case25_done\n\t"
+      "li x6, %[done]\n\t"
+      "sw x6, 0(x7)\n\t"
+      ".insn r %[custom0], 0, 0, x0, x0, x0\n\t"
+      "6: j 6b\n\t"
+      :
+      : [csr_wid] "i"(VX_CSR_WARP_ID), [custom0] "i"(RISCV_CUSTOM0),
+        [custom3] "i"(RISCV_CUSTOM3),
+        [num_scalar_warps] "i"(NUM_SCALAR_WARPS),
+        [c_offset] "i"(WU_BW_TMEM_C_BYTE_OFFSET),
+        [tile_bytes] "i"(WU_BW_TMEM_TILE_BYTES),
+        [smem_base] "i"(WU_BW_DEV_SMEM_START_ADDR),
+        [score_ready] "i"(WU_CASE25_SCORE_READY_BASE),
+        [p_ready] "i"(WU_CASE25_P_READY_BASE),
+        [done] "i"(WU_CASE25_DONE_BASE)
+      : "memory");
+}
+
+extern "C" int wu_main() {
+  if (vx_core_id() != 0 || vx_warp_id() != 0) {
+    return 0;
+  }
+
+  const uint32_t tid = wu_tid();
+  if (tid == 0) {
+    wu_case_reset();
+    g_case25_score_ready = 0;
+    g_case25_p_ready = 0;
+    g_case25_done = 0;
+    for (uint32_t i = 0; i < NUM_THREADS; ++i) {
+      g_case25_score_bits[i] = 0;
+    }
+    for (uint32_t i = 0; i < WU_BW_OUT_WORDS; ++i) {
+      g_case25_out[i] = 0;
+    }
+    wu_bw_fill_smem_tile(
+        reinterpret_cast<volatile uint32_t *>(WU_BW_DEV_SMEM_START_ADDR),
+        WU_BW_FP16_ONE_PACKED);
+    vx_spawn_tensor(1u << NUM_SCALAR_WARPS, tensor_case25_worker);
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (wu_case25_wait_status(&g_case25_score_ready,
+                            WU_CASE25_SCORE_READY_BASE) != 0) {
+    if (tid == 0) {
+      g_case_mem[1] = 0x91u;
+    }
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  const uint32_t c_frag =
+      wu_bw_tmem_c_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES;
+  const uint32_t observed = wu_bw_scalar_tmem_ld(c_frag);
+  g_case25_score_bits[tid] = observed;
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (tid == 0 && g_case_mem[1] == 0 &&
+      observed != WU_CASE25_FP32_THIRTY_TWO) {
+    g_aux[0] = observed;
+    g_case_mem[1] = 0x92u;
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  uint32_t softmax_token = 0;
+  const uint32_t p_frag =
+      wu_bw_tmem_a_byte_base(0) / WU_BW_TMEM_FRAGMENT_BYTES;
+  if (g_case_mem[1] == 0) {
+    vx_tmc(wu_bw_all_lanes_mask());
+    softmax_token = wu_bw_scalar_tmem_softmax(c_frag, p_frag);
+    vx_tmc_one();
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (tid == 0 && g_case_mem[1] == 0) {
+    if (softmax_token != p_frag) {
+      g_aux[0] = softmax_token;
+      g_aux[1] = p_frag;
+      g_case_mem[1] = 0x95u;
+    } else {
+      g_case25_p_ready = WU_CASE25_P_READY_BASE;
+    }
+  }
+  asm volatile("fence rw, rw" ::: "memory");
+
+  if (tid == 0) {
+    if (g_case_mem[1] == 0 &&
+        wu_case25_wait_status(&g_case25_done, WU_CASE25_DONE_BASE) != 0) {
+      g_case_mem[1] = 0x93u;
+    }
+    if (g_case_mem[1] == 0) {
+      volatile uint32_t bad_actual = 0;
+      const uint32_t bad =
+          wu_bw_verify_constant(g_case25_out, WU_BW_OUT_WORDS,
+                                WU_BW_FP32_ONE, &bad_actual);
+      if (bad != WU_BW_OUT_WORDS) {
+        g_aux[0] = bad;
+        g_aux[1] = bad_actual;
+        g_case_mem[1] = 0x94u;
+      }
+    }
+    if (g_case_mem[1] != 0) {
+      wu_case_fail(g_case_mem[1]);
+      return 1;
+    }
+    wu_case_pass();
+  }
+  return 0;
+}

kernels/wu_arch_cases/common_wu_blackwell_fa.h

@@ -50,6 +50,17 @@ static inline void wu_bw_scalar_tmem_st(uint32_t frag_addr, uint32_t value) {
                : "memory");
 }
 
+static inline uint32_t wu_bw_scalar_tmem_softmax(uint32_t score_frag_addr,
+                                                 uint32_t p_frag_addr) {
+  uint32_t token;
+  asm volatile(".insn r %[custom1], 3, 0x30, %[token], %[score], %[prob]"
+               : [token] "=r"(token)
+               : [custom1] "i"(RISCV_CUSTOM1), [score] "r"(score_frag_addr),
+                 [prob] "r"(p_frag_addr)
+               : "memory");
+  return token;
+}
+
 static inline uint32_t wu_bw_tmem_a_byte_base(uint32_t tensor_block) {
   return tensor_block * WU_BW_TMEM_WARP_BYTES;
 }