perf(rocke): enable exp2_fast for bf16 head_dim=128 in gfx942 dense prefill

dnn-providers/hip-kernel-provider/rocke/library/builders/gfx942/attention/prefill/README.md

@@ -141,7 +141,7 @@ magnitudes live outside the repo per `AGENTS.md`.
 |---|---|---|---|
 | 32×32×8 atom + K-loop doubling | all | CDNA3 has no 32×32×16 fp16/bf16 atom; C-layout is identical so softmax/epilogue port unchanged | **shipped** — enablement |
 | Conflict-free V (`perm_b32` store transpose) | **D128 fp16** | V stored transposed → PV A-operand read is one contiguous `ds_read_b64` instead of 4 element-wise `ds_read_u16` | **shipped** — large; identity-preserving |
-| `exp2_fast` | all except **bf16 D128** | softmax args are provably ≤ 0, so `llvm.exp2`'s guarded range reduction is dead work | **shipped** — dominant on the VALU-bound path; bit-identical |
+| `exp2_fast` | all | softmax args are provably ≤ 0, so `llvm.exp2`'s guarded range reduction is dead work; bf16 D128 (the last holdout) re-enabled on the lazy-rescale schedule — 0 scratch, VGPR-neutral, numerically identical | **shipped** — dominant on the VALU-bound path; bit-identical |
 | Fused softmax rescale | all | exp2 → accumulate → cast → pack in one pass instead of materializing a full f32 `p_vals` matrix | **shipped** — pure live-range relief, bit-identical |
 | Per-config `waves_per_eu` | **bf16 D64** → 4 | forces the allocator low enough that a second workgroup co-resides (1 → 2 WG/CU) | **shipped** — large at long sequences |
 | D64 K-bank-conflict pad | **D64 both dtypes** | 2-row-group boundary pad takes the `do_qk` K reads from 32-way to 4-way | **shipped** — large, cross-part confirmed |
@@ -153,7 +153,6 @@ magnitudes live outside the repo per `AGENTS.md`.
 |---|---|
 | Conflict-free V at **D64** | D64 is VGPR-bound; the register round-trip costs more than the LDS-instruction saving |
 | Conflict-free V at **bf16 D128** | spills over the `waves_per_eu=2` cap on the `.1k` MFMA schedule |
-| `exp2_fast` at **bf16 D128** | same — spills even after the fused rescale freed headroom |
 | `waves_per_eu=3` at **fp16 D64** | reaches 2 WG/CU but loses more ILP than the second workgroup buys |
 | Drop K/V LDS pads to reach 2 WG/CU at D128 | D128 stays 1 WG/CU even unpadded (register floor co-limits), so it only reintroduces bank conflicts — **catastrophic** |
 | `block_n=32` (all configs) | halving the KV tile doubles the tile/grid count; the extra loop and barrier overhead outweighs the LDS relief. Marginally positive on **bf16 D128** but **part-dependent**, so not wired — it would need a CU-count-aware policy |

dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py

@@ -81,8 +81,8 @@
 -----------------------------------------------------------------------------
   * P0  enablement + 32x32x8 atom + K-loop doubling ............ DONE (this file)
   * P1  conflict-free V (perm_b32 store-path transpose) ........ DONE (D128 fp16)
-  * P2  exp2_fast + fused/lazy rescale ........................ DONE (exp2_fast all
-        but bf16 D128; fused rescale bit-identical, enables bf16 D64 exp2_fast)
+  * P2  exp2_fast + fused/lazy rescale ........................ DONE (exp2_fast for
+        ALL configs incl. bf16 D128; fused rescale bit-identical, bf16 D64 exp2_fast)
   * P3  occupancy: waves-per-eu tune (bf16 D64 -> wpe4 = 2 WG/CU),
         D64 K bank-pad, wide4 (WG=256), K single-buffer ........... IN PROGRESS
         (waves-per-eu DONE via _tuned_waves_per_eu; **D64 K bank-pad DONE and ADOPTED**
@@ -326,7 +326,7 @@ class Gfx942DenseTuning:
 
     # use_exp2_fast: force the P2 single-instruction exp2 on/off.
     #   None (default) -> :func:`_use_exp2_fast`, which owns the measured verdict
-    #   (on everywhere except bf16 D128, which spills). Numerically safe in both
+    #   (now on for ALL configs; the former bf16-D128 spill holdout is stale). Numerically safe in both
     #   directions here -- both softmax arguments are always <= 0 -- so unlike
     #   ``use_cfvst`` this one is a pure perf A/B and is not gated.
     use_exp2_fast: bool | None = None
@@ -518,22 +518,24 @@ def _rows_per_instr(head_size: int) -> int:
 def _use_exp2_fast(head_size: int, dtype: str) -> bool:
     """Whether softmax uses ``exp2_fast`` (one v_exp_f32, no range-reduction guard).
 
-    Enabled everywhere EXCEPT bf16 D128. exp2_fast is a strict VALU reduction and the
-    dominant P2 lever on the (post-P1) VALU-bound path, and is always numerically safe
-    here -- both softmax args (alpha's m_i - m_new and p's s - m_new) are <= 0, exactly
-    exp2_fast's precondition.
-
-    bf16 D128 is the sole holdout. Its ``.1k`` MFMA schedule keeps more registers live,
-    and exp2_fast makes the exp result available in one instruction (vs plain exp2's
-    ~5-op range reduction), which the scheduler hoists -- lengthening the exp live
-    ranges. Measured post-fused-rescale (plan §6.1): bf16 D128 goes 175 VGPR / 0 spill
-    (plain exp2) -> 256 VGPR / 22 spill (exp2_fast), over the waves-per-eu=2 cap. The
-    fused rescale freed ~28 VGPR but not enough to absorb that hoist on the .1k path;
-    fp16 D128 (213, cfvst) and bf16 D64 (215) both have the headroom. A bf16 D128
-    exp2_fast unblock is a P3 occupancy/scheduling item. Spill re-verified across the
-    fp16/bf16 x D64/D128 cohort.
+    Enabled for all configs. exp2_fast is a strict VALU reduction and the dominant P2
+    lever on the (post-P1) VALU-bound path, and is always numerically safe here -- both
+    softmax args (alpha's m_i - m_new and p's s - m_new) are <= 0, exactly exp2_fast's
+    precondition, independent of head_size and dtype.
+
+    bf16 D128 was previously the sole holdout. An earlier fused-rescale schedule
+    (plan §6.1) spilled with exp2_fast (175 -> 256 VGPR / 22 spill, over the
+    waves-per-eu=2 cap), so the policy disabled it. The current lazy-rescale schedule
+    does not spill: rocprofv3 register capture on the shipped bf16-D128 causal kernel
+    (GQA 32/8, Sq 1024-16384, both the default and persistent grids) reads a lower VGPR
+    count with exp2_fast than without, 0 scratch, numerically identical to plain exp2.
+    The holdout is therefore stale and removed. (The stale numbers were measured against
+    the pre-lazy-rescale body; the register profile has since changed sign.) head_size is
+    no longer a gate -- correctness holds for every head_size (the <= 0 precondition is
+    universal); a new head_size would only warrant a fresh occupancy check, never a
+    correctness one.
     """
-    return dtype == "fp16" or head_size != 128
+    return True
 
 
 def _use_cfvst(head_size: int, dtype: str) -> bool:
@@ -1447,8 +1449,8 @@ def do_mask(s_reg, tile_idx):
             # s) -- exactly exp2_fast's precondition (no overflow; v_exp_f32 flushes
             # large negatives to 0). Cuts ~99 VALU/tile at D128, the dominant
             # MFMA-starving residual once conflict-free V (P1) lands. Enabled for
-            # every config except bf16 D128, which spills on the .1k schedule -- the
-            # spill rationale and the matrix live in _use_exp2_fast's docstring.
+            # every config, including bf16 D128 -- its former spill holdout is stale on
+            # the current lazy-rescale schedule; rationale + matrix in _use_exp2_fast's docstring.
             exp2 = b.exp2_fast if tuning.resolved_use_exp2_fast(spec) else b.exp2
             alpha = exp2(b.fsub(m_i, m_new))
 

dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json

@@ -3,12 +3,12 @@
     "llvm20": {
       "cases": {
         "attention_dense_gfx942/default_d128_bf16_causal": {
-          "bytes": 186152,
-          "sha256": "9063c3a5320e8baa713ed5c8a209c95099944faf339a6d4037af67572db8f1f9"
+          "bytes": 186552,
+          "sha256": "aaac6ac3e7029705bd079c44913024e7ecb6726f1f230ec6a0b77ccea02891a8"
         },
         "attention_dense_gfx942/default_d128_bf16_mha": {
-          "bytes": 186208,
-          "sha256": "c5f95a09cdd8dbfc59fedf8d46e155590442195dd8941cce1a9dabfdf415970b"
+          "bytes": 186608,
+          "sha256": "475b828520c30c8ac6a0e557f7020f71056954c835fb3d9d4c01d4db4a0798b7"
         },
         "attention_dense_gfx942/default_d128_fp16_causal": {
           "bytes": 123573,
@@ -43,12 +43,12 @@
           "sha256": "0aab9f3007a83f2ebb7419ec6556aa43562d8ee6102c6c8ffb1a86d9fe550e68"
         },
         "attention_dense_gfx942/persist_d128_bf16_intl": {
-          "bytes": 187317,
-          "sha256": "fccf6bd8c682556dde15835b2485737d118a3e3764aa9dc844c1f665ff497b0f"
+          "bytes": 187717,
+          "sha256": "fcf12e6716891be454f076c12a7f88ba65bba0d23c575c4a9514b5ecc420741e"
         },
         "attention_dense_gfx942/persist_d128_bf16_qbmaj": {
-          "bytes": 186966,
-          "sha256": "6fecf26c983b463be9ebf64bf018de11ed2f87bbee88d19d361e213795dbcf08"
+          "bytes": 187366,
+          "sha256": "706865cb73def467b8a1bc0110a0aa7e4d893333826a39f8f182f08dd1e7fa37"
         },
         "attention_dense_gfx942/persist_d128_fp16_hkvmaj": {
           "bytes": 124911,
@@ -67,12 +67,12 @@
     "llvm22": {
       "cases": {
         "attention_dense_gfx942/default_d128_bf16_causal": {
-          "bytes": 186248,
-          "sha256": "ecb0f1d533a710f03f35f239bd971c9f6a9b05289a6cf403ca49d2f41c2cb14f"
+          "bytes": 186648,
+          "sha256": "c5aeaa4fa92c8030459345f8a11b5f90873821557433f879ffb8fe169b441505"
         },
         "attention_dense_gfx942/default_d128_bf16_mha": {
-          "bytes": 186304,
-          "sha256": "95d7b6eee764201f88dbf5eaf0c12bd90b74c86db27914744971b73d64eb5465"
+          "bytes": 186704,
+          "sha256": "0dae9d93db815349a1817500cdb84a2dea4a5584c4c28bdb431878b853d189c8"
         },
         "attention_dense_gfx942/default_d128_fp16_causal": {
           "bytes": 123669,
@@ -107,12 +107,12 @@
           "sha256": "a8b2b8b585e125794ff8119a26cabdb9f156b23fdc911dfb2d3fa03761da9c7b"
         },
         "attention_dense_gfx942/persist_d128_bf16_intl": {
-          "bytes": 187413,
-          "sha256": "577ba5e06e75a216f6444e9c160d7a0ee29e243ae3aa10d4e7634accff199800"
+          "bytes": 187813,
+          "sha256": "ec386b8186806d29ff73fc66338d2c31a495f97c0fa0546d6ce69dc46991d77c"
         },
         "attention_dense_gfx942/persist_d128_bf16_qbmaj": {
-          "bytes": 187062,
-          "sha256": "d5870b2eee4685ee4de98c63ef4e95d9639939cf51a222cc51386140e636f0ee"
+          "bytes": 187462,
+          "sha256": "1021d907082ea7417334490acc851fa3942771ed083772b93a8b865f04677124"
         },
         "attention_dense_gfx942/persist_d128_fp16_hkvmaj": {
           "bytes": 125007,

dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py

@@ -583,7 +583,10 @@ def test_dataclass_rejects_out_of_scope_headsize():
     (dict(dtype="bf16", head_size=128), dict(use_cfvst=True)),  # REJECTED: spills
     # --- tuning: use_exp2_fast (no rejected region -- see the comment above) ---
     (dict(), dict(use_exp2_fast=False)),
-    (dict(), dict(use_exp2_fast=True)),  # forced ON at bf16 D128, the spilling arm
+    (
+        dict(),
+        dict(use_exp2_fast=True),
+    ),  # accepted: exp2_fast has no rejected region (default True everywhere)
     # --- tuning: waves_per_eu ---
     (dict(), dict(waves_per_eu=4)),  # accepted
     (dict(), dict(waves_per_eu=0)),  # REJECTED: must resolve positive
@@ -849,18 +852,22 @@ def _walk_op_names(op):
         (64, "fp16", True),
         (128, "fp16", True),
         (64, "bf16", True),  # fused rescale gave the headroom (P2)
-        (128, "bf16", False),  # spills on the .1k schedule even post-fused-rescale
+        (
+            128,
+            "bf16",
+            True,
+        ),  # lazy-rescale: 0 scratch, VGPR-neutral -> last holdout enabled
     ],
 )
 def test_exp2_fast_gate_matches_the_spill_measured_matrix(head_size, dtype, expected):
-    """The exp2_fast decision is a spill fact, not a preference.
-
-    exp2_fast is numerically safe for every config (both softmax args <= 0), so the
-    gate exists ONLY to keep occupancy: its sooner-available result raises register
-    pressure, and bf16 D128's `.1k` MFMA schedule spills over the waves-per-eu=2 cap
-    (measured 175->256 VGPR / 22 spill) even after the P2 fused rescale freed ~28
-    VGPR. Every other config has the headroom. This pins the exact enabled set so a
-    future edit that flips one arm has to update this matrix on purpose.
+    """exp2_fast is enabled for EVERY config. It is numerically safe everywhere (both
+    softmax args -- alpha's m_i - m_new and p's s - m_new -- are <= 0, exactly
+    exp2_fast's precondition, independent of head_size/dtype) and a strict VALU win.
+    bf16 D128 was the last holdout under the earlier fused-rescale schedule (spilled
+    175->256 VGPR over the waves-per-eu=2 cap); the current lazy-rescale schedule does
+    not spill (rocprofv3: 0 scratch, VGPR-neutral vs plain exp2, numerically identical),
+    so the holdout is removed. This pins the enabled set (now all configs) so a future
+    edit that re-introduces a rejected arm has to update this matrix on purpose.
     """
     assert _use_exp2_fast(head_size, dtype) is expected