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

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,25 @@ 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: an ``llvm-readelf --notes`` read of the compiled code object
+    (``.private_segment_fixed_size``/``.vgpr_count``) on the shipped bf16-D128 kernel --
+    both the default and persistent grids -- reports 0 scratch and a lower VGPR count
+    with exp2_fast than without, 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 +1450,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

@@ -849,18 +849,21 @@ 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,
+        ),  # exp2_fast enabled: lazy-rescale schedule no longer spills
     ],
 )
 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 -- both softmax args (m_i-m_new and
+    s-m_new) are <= 0, so exp2_fast is numerically safe universally. bf16 D128 was
+    previously held out for a register spill on the fused-rescale schedule (measured
+    175->256 VGPR / 22 spill over the waves-per-eu=2 cap); the current lazy-rescale
+    schedule does not spill (rocprofv3: lower VGPR, 0 scratch), so the holdout is
+    removed. This pins the exact enabled set so a future edit that flips an arm has to
+    update this matrix on purpose.
     """
     assert _use_exp2_fast(head_size, dtype) is expected