From d4de8c1faf26d8aee23118568cab87f94eb3e6fe Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Fri, 14 Aug 2026 17:36:31 -0400 Subject: [PATCH 1/3] perf(rocke): enable exp2_fast for bf16 head_dim=128 in gfx942 dense prefill The bf16-D128 exp2_fast holdout was justified by a pre-lazy-rescale spill over the waves-per-eu=2 cap. rocprofv3 on the current lazy-rescale kernel shows the opposite: VGPR drops and scratch stays zero, output numerically identical to plain exp2. Remove the stale holdout in _use_exp2_fast (return True), update the in-file status/policy comments that referenced the old exception, and re-bless the four bf16-D128 golden cases across both llvm flavors; fp16 and D64 are byte-identical. --- .../library/kernels/gfx942/attention_dense.py | 42 ++++++++++--------- .../attention_dense_gfx942_ir_sha256.json | 32 +++++++------- 2 files changed, 38 insertions(+), 36 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py index ea627888647f..6edef0c04a0e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py +++ b/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)) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json index 64b8afa5c3e7..b449f709a1f8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/attention_dense_gfx942_ir_sha256.json +++ b/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, From 8aa1b284fc25c65956d0174f66dd9e2b609e31f7 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Sat, 15 Aug 2026 00:48:24 -0400 Subject: [PATCH 2/3] test(rocke): update exp2_fast gate matrix for bf16 D128 enablement _use_exp2_fast now returns True for all configs (the bf16 D128 spill holdout is stale on the lazy-rescale schedule), so the pinned matrix in test_exp2_fast_gate_matches_the_spill_measured_matrix must expect True for (128, bf16). Flip the case and refresh the docstring; the parametrized case fails on the pre-change gate and passes after, guarding the flip. --- .../tests/test_attention_dense_gfx942.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py index 90067d1058a5..11ebcde7e7ec 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py @@ -849,18 +849,17 @@ 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 From e775fe95452d5c87b549608cdcf034d2362b4db6 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Sat, 15 Aug 2026 00:58:10 -0400 Subject: [PATCH 3/3] chore(rocke): format exp2_fast gate case + cite code-object read for no-spill - black reformats the flipped (128, bf16, True) parametrize case - _use_exp2_fast docstring now cites the llvm-readelf --notes code-object read (.private_segment_fixed_size / .vgpr_count) for the bf16-D128 no-spill claim, matching the PR body: the lazy-rescale build reports 0 scratch on both the default and persistent grids (VGPR reduced, not increased, vs plain exp2). --- .../rocke/library/kernels/gfx942/attention_dense.py | 7 ++++--- .../rocke/library/tests/test_attention_dense_gfx942.py | 6 +++++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py index 6edef0c04a0e..eb6038377a53 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py @@ -526,9 +526,10 @@ def _use_exp2_fast(head_size: int, dtype: str) -> bool: 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. + 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 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py index 11ebcde7e7ec..fbe095e432a9 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py @@ -849,7 +849,11 @@ def _walk_op_names(op): (64, "fp16", True), (128, "fp16", True), (64, "bf16", True), # fused rescale gave the headroom (P2) - (128, "bf16", True), # exp2_fast enabled: lazy-rescale schedule no longer spills + ( + 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):