From 4ab18ab6b1f8db79b9c327d4265b1a400192179f Mon Sep 17 00:00:00 2001
From: Zhongdi LUO <zluo@172.26.10.1>
Date: Sun, 12 Jul 2026 02:00:14 +0000
Subject: [PATCH] feat: pipeline fp8 scalar softmax and blackwell mma issue

---
 src/main/resources/vsrc/vortex                |  2 +-
 .../radiance/core/TensorCoreBlackwell.scala   | 38 +++++++++++--------
 src/main/scala/radiance/tile/VortexCore.scala |  1 +
 .../radiance/TensorCoreBlackwellTest.scala    |  4 ++
 4 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/src/main/resources/vsrc/vortex b/src/main/resources/vsrc/vortex
index dff1107..331c4de 160000
--- a/src/main/resources/vsrc/vortex
+++ b/src/main/resources/vsrc/vortex
@@ -1 +1 @@
-Subproject commit dff1107bf5e8b4e38e2acb9c623b1286ac662ec4
+Subproject commit 331c4decaa5347ec0bb4259d8882532583f9e4b1
diff --git a/src/main/scala/radiance/core/TensorCoreBlackwell.scala b/src/main/scala/radiance/core/TensorCoreBlackwell.scala
index 6c80682..f595956 100644
--- a/src/main/scala/radiance/core/TensorCoreBlackwell.scala
+++ b/src/main/scala/radiance/core/TensorCoreBlackwell.scala
@@ -149,7 +149,8 @@ class TensorCoreBlackwell(
   val bIndexReg = RegInit(0.U(log2Ceil(numBFragsPerGroup).W))
   val mGroupReg = RegInit(0.U(log2Ceil(numMGroups).W))
   val substepReg = RegInit(0.U(1.W))
-  val elemReg = RegInit(0.U(log2Ceil(numLanes).W))
+  val issueElemReg = RegInit(0.U(log2Ceil(numLanes).W))
+  val retireElemReg = RegInit(0.U(log2Ceil(numLanes).W))
   val waitCounter = RegInit(0.U(3.W))
 
   val aBuf = Reg(Vec(numAFragsPerSet, UInt(memWidth.W)))
@@ -225,14 +226,14 @@ class TensorCoreBlackwell(
     inputType = TensorInputType.FP8E4M3
   ))
 
-  val elemM = if (numLanes == 4) elemReg(0, 0) else elemReg(1, 0)
-  val elemN = if (numLanes == 4) elemReg(1) else elemReg(2)
+  val elemM = if (numLanes == 4) issueElemReg(0, 0) else issueElemReg(1, 0)
+  val elemN = if (numLanes == 4) issueElemReg(1) else issueElemReg(2)
   dpu.io.in.valid := dpuInValid
   for (k <- 0 until 8) {
     dpu.io.in.bits.a(k) := TensorCoreBlackwellFP8Packing.selectA(operandA, k, elemM, numLanes)
     dpu.io.in.bits.b(k) := TensorCoreBlackwellFP8Packing.selectB(operandB, k, elemN)
   }
-  dpu.io.in.bits.c := cWords(elemReg)
+  dpu.io.in.bits.c := cWords(issueElemReg)
   dpu.io.stall := false.B
   val dpuValid = dpu.io.out.valid
 
@@ -253,7 +254,8 @@ class TensorCoreBlackwell(
     bIndexReg := 0.U
     mGroupReg := 0.U
     substepReg := 0.U
-    elemReg := 0.U
+    issueElemReg := 0.U
+    retireElemReg := 0.U
     switch(io.initiate.bits.op) {
       is(Ops.bwgmma) { state := State.bwLoadAReq }
       is(Ops.tcgen05Cp) { state := State.cpRead }
@@ -322,24 +324,28 @@ class TensorCoreBlackwell(
 
   when(state === State.bwReadCResp) {
     cDataReg := io.tmemC.cRdata
-    elemReg  := 0.U
+    issueElemReg := 0.U
+    retireElemReg := 0.U
     state    := State.bwCompute
   }
 
   when(state === State.bwCompute) {
     dpuInValid := true.B
-    state := State.bwDpuResp
+    when(issueElemReg === (numLanes - 1).U) {
+      state := State.bwDpuResp
+    }.otherwise {
+      issueElemReg := issueElemReg + 1.U
+    }
   }
 
-  when(state === State.bwDpuResp) {
-    when(dpuValid) {
-      mmaDataReg(elemReg) := dpu.io.out.bits.data
-      when(elemReg === (numLanes - 1).U) {
-        state := State.bwWriteCReq
-      }.otherwise {
-        elemReg := elemReg + 1.U
-        state := State.bwCompute
-      }
+  when(dpuValid) {
+    assert(state === State.bwCompute || state === State.bwDpuResp,
+      "BWGMMA DPU response arrived outside the compute states")
+    mmaDataReg(retireElemReg) := dpu.io.out.bits.data
+    when(retireElemReg === (numLanes - 1).U) {
+      state := State.bwWriteCReq
+    }.otherwise {
+      retireElemReg := retireElemReg + 1.U
     }
   }
 
diff --git a/src/main/scala/radiance/tile/VortexCore.scala b/src/main/scala/radiance/tile/VortexCore.scala
index 1b06a88..dfef1bb 100644
--- a/src/main/scala/radiance/tile/VortexCore.scala
+++ b/src/main/scala/radiance/tile/VortexCore.scala
@@ -213,6 +213,7 @@ class Vortex(tile: RadianceTile)(implicit p: Parameters)
   addResource("/vsrc/vortex/hw/rtl/core/VX_dispatch.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_dispatch_unit.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_dispatch_unit_sane.sv")
+  addResource("/vsrc/vortex/hw/rtl/core/VX_tmem_softmax_unit.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_execute.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_fetch.sv")
   addResource("/vsrc/vortex/hw/rtl/core/VX_gather_unit.sv")
diff --git a/src/test/scala/radiance/TensorCoreBlackwellTest.scala b/src/test/scala/radiance/TensorCoreBlackwellTest.scala
index f85b0de..c872ccd 100644
--- a/src/test/scala/radiance/TensorCoreBlackwellTest.scala
+++ b/src/test/scala/radiance/TensorCoreBlackwellTest.scala
@@ -283,6 +283,7 @@ class TensorCoreBlackwellTest extends AnyFlatSpec with ChiselScalatestTester {
 
       var pendingB = Option.empty[(BigInt, BigInt)]
       var sawWriteback = false
+      var cycles = 0
 
       for (_ <- 0 until 20000 if !sawWriteback) {
         // Drive TMEM reads/writes
@@ -306,11 +307,14 @@ class TensorCoreBlackwellTest extends AnyFlatSpec with ChiselScalatestTester {
           } else None
 
           c.clock.step()
+          cycles += 1
           pendingB = nextB
         }
       }
 
       assert(sawWriteback, "BWGMMA did not complete")
+      assert(cycles < 5000,
+        s"BWGMMA took $cycles cycles; fragment elements are not issuing back-to-back")
       c.io.writeback.bits.wid.expect(1.U)
       // Verify all 32 C frags in TMEM
       for (i <- 0 until 32) {