From 1c4e5d8de537d01342a014b29469d920061d8d49 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Fri, 3 Jul 2026 18:02:38 -0400 Subject: [PATCH 01/23] feat(rocke/attention): deterministic LDS-budget resolver for 2D prefill The register-PV 2D-tiled prefill path stages only K/V tiles in LDS; at head_dim=256 the default double-buffered K tile is 204800 B, over gfx950's 163840 B cap, so the kernel fails to compile ("local memory (204800) exceeds limit (163840)"). Add _resolve_lds_budget, hooked on _tiled_spec_from_problem's main return: it computes the exact register-PV LDS footprint and, when over the arch cap, deterministically shrinks the geometry (single-buffer K, then T=64) until it fits -- a strict no-op otherwise, so all currently-compiling configs are byte-identical. Scoped to gfx950 register-PV; mechanism reads the cap dynamically. GPU-verified on MI350X (bf16, GQA 16/2): D256 prefill compiles (139264 <= 163840), runs, and matches the fp32 paged reference (max_abs 1.5625e-02) at Sq 4096 and 8192. --- .../kernels/common/attention_unified.py | 105 +++++++++++++++++- .../library/tests/test_lds_budget_resolver.py | 87 +++++++++++++++ 2 files changed, 191 insertions(+), 1 deletion(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 900046945c07..717f5e7b6f57 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -1814,6 +1814,108 @@ def _enable_i64_kv_addr(problem: UnifiedAttentionProblem) -> bool: return cache_bytes > 0x8000_0000 +# --- LDS-budget resolver ---------------------------------------------------- +# The 2D-tiled register-PV path stages the K/V tiles plus a small epilogue Acc_lds +# buffer in LDS (Q, P^T and the running O accumulator are register-resident, so +# Acc_lds is the only auxiliary term). At head_dim=256 the default double-buffered +# K tile makes that footprint 204800 B, over gfx950's 163840 B cap, so the kernel +# fails to compile ("local memory (204800) exceeds limit (163840)"). Whether a +# geometry fits is deterministic (closed-form footprint vs the arch cap), so +# rather than hard-fail we deterministically shrink the geometry until it fits, +# using only reductions empirically confirmed to compile: (1) single-buffer K, +# (2) T=64. The resolver is a strict NO-OP whenever the footprint already fits or +# the path is not register-PV, so every currently-compiling config is byte-identical. +# Compute-dtype (bf16/fp16) width in bytes. The fp8 lever changes only the K/V +# *cache* storage width, never the output accumulator, so the epilogue Acc_lds +# staging buffer is always 16-bit. +_ACC_LDS_ELEM_BYTES = 2 + + +def _kv_lds_elem_bytes(spec) -> int: + if getattr(spec, "kv_storage_dtype", None) in ("fp8", "bf8", "e4m3", "e5m2"): + return 1 + return 2 # bf16 / fp16 K/V in LDS + + +def _acc_lds_bytes(spec) -> int: + """Epilogue Acc_lds staging buffer -- ``smem_alloc([BLOCK_M, OUT_STRIPE_COLS])`` + of the compute dtype (see ``gfx950/attention_tiled_2d.py`` ~L1473/1248). In the + register-PV path this is the *only* LDS beyond the K/V tiles (Q, P^T and the + running softmax stats are register-resident), so it is the full auxiliary term. + Scales with ``num_warps``, ``block_m_per_warp`` and ``head_size``.""" + out_stripe_cols = 32 if spec.head_size <= 64 else spec.head_size + block_m = spec.num_warps * spec.block_m_per_warp + return block_m * out_stripe_cols * _ACC_LDS_ELEM_BYTES + + +def _lds_bytes_regpv(spec) -> int: + """Exact LDS footprint of the register-PV 2D path: the K/V tiles plus the + epilogue Acc_lds. Verified against comgr for the D256 T=128 K-double geometry: + ``(2+1)*128*256*2 (K/V) + 16*256*2 (Acc) = 196608 + 8192 = 204800``.""" + wb = _kv_lds_elem_bytes(spec) + kbuf = 1 if spec.use_k_single_buffer else 2 + vbuf = 2 if spec.use_v_double_buffer else 1 + kv_bytes = (kbuf + vbuf) * spec.tile_size * spec.head_size * wb + return kv_bytes + _acc_lds_bytes(spec) + + +def _lds_capacity_bytes() -> int: + # Lazy import: keep ``common/`` arch-neutral (see module top). + from rocke.core.arch.target import ArchTarget + + return ArchTarget.from_gfx(_resolve_attention_arch()).lds_capacity_bytes + + +def _ldsfix_single_k(spec): + if spec.use_k_single_buffer: + return None + # K-single needs Q to fit the lone K slot: block_m <= tile_size. + if spec.num_warps * spec.block_m_per_warp > spec.tile_size: + return None + try: + return replace(spec, use_k_single_buffer=True) + except Exception: + return None + + +def _ldsfix_tile64(spec): + if spec.tile_size <= 64: + return None + try: + return replace(spec, tile_size=64) + except Exception: + return None + + +def _resolve_lds_budget(spec): + """Deterministically shrink an over-budget register-PV 2D spec until it fits + the arch LDS cap, using only compile-validated reductions. Returns the same + spec unchanged when it already fits or the path is not register-PV (so all + currently-compiling configs stay byte-identical).""" + # Validated on gfx950 (CDNA4): register-PV stages K/V tiles + the epilogue Acc_lds. The + # mechanism is arch-general (reads the cap dynamically); other arches' 2D + # footprint models are not yet validated, so engage only where proven. + if _resolve_attention_arch() != "gfx950" or not getattr( + spec, "use_register_pv", False + ): + return spec + cap = _lds_capacity_bytes() + if _lds_bytes_regpv(spec) <= cap: + return spec + single_k = _ldsfix_single_k(spec) + candidates = [single_k, _ldsfix_tile64(spec)] + if single_k is not None: + candidates.append(_ldsfix_tile64(single_k)) + for cand in candidates: + if cand is not None and _lds_bytes_regpv(cand) <= cap: + return cand + raise RuntimeError( + f"LDS budget: 2D register-PV D{spec.head_size} block_size={spec.block_size} " + f"T={spec.tile_size} needs {_lds_bytes_regpv(spec)} B > cap {cap} B on " + f"{_resolve_attention_arch()}; no validated reduction (single-K, T=64) fits." + ) + + def _tiled_spec_from_problem( problem: UnifiedAttentionProblem, ): @@ -1953,7 +2055,7 @@ def _tiled_spec_from_problem( # no-fp8 preconditions so it can never fire on an incompatible spec. if "use_k_single_buffer" in _spec_field_names and _enable_k_single_buffer(problem): _gfx950_schedule_fields["use_k_single_buffer"] = True - return UnifiedAttention2DTiledSpec( + _spec = UnifiedAttention2DTiledSpec( head_size=problem.head_size, block_size=problem.block_size, num_query_heads=problem.num_query_heads, @@ -2016,6 +2118,7 @@ def _tiled_spec_from_problem( # below -- gfx942's spec class does not declare it.) **_gfx950_schedule_fields, ) + return _resolve_lds_budget(_spec) def _select_2d_block_m_per_warp(problem: UnifiedAttentionProblem) -> int: diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py new file mode 100644 index 000000000000..26650c8d7d34 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py @@ -0,0 +1,87 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Unit tests for the 2D-tiled LDS-budget resolver. + +Pure codegen (no GPU, no subprocess). The resolver deterministically shrinks an +over-budget register-PV 2D spec until it fits the arch LDS cap, and is a strict +no-op for every already-fitting / non-register-PV / non-gfx950 config. + +The load-bearing case: bf16 head_dim=256 long prefill on gfx950 overflows LDS at +the default (K double-buffered) geometry -- 204800 B > the 163840 B cap -- and the +resolver must single-buffer K so it compiles. The pre-fix codegen matrix only +covered D256 in fp16, so this path was previously untested. +""" + +from __future__ import annotations + +import unittest +from dataclasses import replace +from unittest import mock + +import kernels.common.attention_unified as au +from kernels import UnifiedAttentionProblem + + +def _d256_bf16_long_prefill(block_size: int = 64, max_seqlen: int = 4096): + # No sinks / sliding-window / softcap / alibi / qq_bias, so use_register_pv + # is eligible and the spec routes onto the register-PV 2D path. + return UnifiedAttentionProblem( + total_q=max_seqlen, + num_seqs=1, + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=block_size, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + dtype="bf16", + ) + + +class TestLdsBudgetResolver(unittest.TestCase): + def test_d256_gfx950_bf16_prefill_shrinks_to_fit(self): + """D256 bf16 register-PV overflows at the default geometry; the resolver + must single-buffer K so the resolved spec fits the gfx950 160 KB cap.""" + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + self.assertTrue(spec.use_register_pv) + # The cheapest lever (single-buffer K) was applied ... + self.assertTrue(spec.use_k_single_buffer) + # ... and the resolved geometry fits the arch cap. + self.assertLessEqual(au._lds_bytes_regpv(spec), au._lds_capacity_bytes()) + + def test_pre_resolve_geometry_actually_overflows(self): + """Proves the resolver was necessary: the K-double-buffered geometry it + started from does exceed the cap (so the no-op guard didn't fire).""" + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + k_double = replace(spec, use_k_single_buffer=False) + self.assertGreater(au._lds_bytes_regpv(k_double), au._lds_capacity_bytes()) + + def test_resolver_is_noop_when_already_fitting(self): + """A spec that already fits is returned byte-identical (same object).""" + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + fitted = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + self.assertLessEqual(au._lds_bytes_regpv(fitted), au._lds_capacity_bytes()) + self.assertIs(au._resolve_lds_budget(fitted), fitted) + + def test_resolver_is_noop_off_gfx950(self): + """On a non-gfx950 arch the resolver never engages -- it returns the spec + unchanged (the footprint model is only validated for gfx950).""" + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx942"): + self.assertIs(au._resolve_lds_budget(spec), spec) + + def test_acc_lds_scales_with_geometry(self): + """The auxiliary term is the epilogue Acc_lds, computed from geometry + (BLOCK_M x OUT_STRIPE_COLS x 2), not a hard-coded constant.""" + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + block_m = spec.num_warps * spec.block_m_per_warp + out_stripe = 32 if spec.head_size <= 64 else spec.head_size + self.assertEqual(au._acc_lds_bytes(spec), block_m * out_stripe * 2) + + +if __name__ == "__main__": + unittest.main() From 3738343f4c310aa9084d66e11312266dab072424 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Tue, 7 Jul 2026 05:01:30 -0400 Subject: [PATCH 02/23] feat(rocke/attention): mirror LDS-budget resolver in the C++ engine + dedup footprint model The Python D256 gfx950 register-PV resolver had no C++ twin: once attention is wired through the shipped provider, its spec-assembly path would select the overflowing K-double geometry and fail comgr CODEGEN. Mirror the selection-layer resolver in the C++ selectors so the provider shrinks the geometry the same way (single-buffer K, then T=64) before build. - rocke_unified_attn_tiled_2d_lds_bytes: single source of truth for the tiled-2D LDS footprint, shared by the gfx942 admission gate and the register-PV budget resolver in both engines -- removes the duplicated per-buffer arithmetic that could drift. - rocke_unified_attn_resolve_lds_budget: byte-faithful C twin of _resolve_lds_budget; strict no-op for fitting / non-register-PV / non-gfx950. - Python: _tiled_2d_lds_bytes now backs both _lds_bytes_regpv and the gfx942 gate (byte-identical numbers); the shrink levers narrow their except to ValueError and thread the reject reason into the resolver diagnostic. - Tests: footprint single-source, lever reason-threading, infeasible-raises. --- .../kernels/common/attention_unified.py | 122 ++++++++++++--- .../kernels/gfx942/attention_tiled_2d.py | 21 ++- .../library/tests/test_lds_budget_resolver.py | 62 +++++++- ...ances.common.attention_unified_selectors.h | 33 ++++ .../common/attention_unified_selectors.cpp | 147 ++++++++++++++++++ .../attention_tiled_2d_public_entry_glue.cpp | 20 ++- 6 files changed, 361 insertions(+), 44 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 717f5e7b6f57..c2691860c028 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -1832,31 +1832,76 @@ def _enable_i64_kv_addr(problem: UnifiedAttentionProblem) -> bool: def _kv_lds_elem_bytes(spec) -> int: + # Byte width of one K/V element as stored in LDS: 1 for the fp8 cache, else 2 (bf16/fp16). if getattr(spec, "kv_storage_dtype", None) in ("fp8", "bf8", "e4m3", "e5m2"): return 1 return 2 # bf16 / fp16 K/V in LDS +def _out_stripe_cols(head_size: int) -> int: + # Epilogue Acc_lds column count: the narrow 32-col stripe for d<=64, else HD. + return 32 if head_size <= 64 else head_size + + +def _tiled_2d_lds_bytes( + *, + tile_size: int, + head_size: int, + block_m: int, + kv_elem_bytes: int, + k_slots: int, + v_slots: int = 1, + include_q_lds: bool = False, + include_p_lds: bool = False, + v_pad: int = 0, +) -> int: + """Static LDS footprint (bytes) of a tiled-2D attention geometry -- the single + source of truth for the per-buffer tile arithmetic shared by the gfx942 + admission gate (conservative: ``Q_lds``/``P_lds`` staged) and the gfx950 + register-PV budget resolver (exact: Q/P^T register-resident). Callers pick + which buffers are LDS-resident; the tile math (K/V slots, the ``block_m<=2T`` + Q alias, the ``OUT_STRIPE`` rule, the 16-bit Acc staging width) lives here so + the models cannot silently drift.""" + k_bytes = k_slots * tile_size * head_size * kv_elem_bytes + v_bytes = v_slots * (tile_size + v_pad) * head_size * kv_elem_bytes + p_bytes = block_m * (tile_size + 8) * kv_elem_bytes if include_p_lds else 0 + # Q aliases the K slab when it fits under the double-buffer window. + q_bytes = ( + block_m * head_size * kv_elem_bytes + if include_q_lds and block_m > 2 * tile_size + else 0 + ) + acc_bytes = block_m * _out_stripe_cols(head_size) * _ACC_LDS_ELEM_BYTES + return k_bytes + v_bytes + p_bytes + q_bytes + acc_bytes + + def _acc_lds_bytes(spec) -> int: """Epilogue Acc_lds staging buffer -- ``smem_alloc([BLOCK_M, OUT_STRIPE_COLS])`` of the compute dtype (see ``gfx950/attention_tiled_2d.py`` ~L1473/1248). In the register-PV path this is the *only* LDS beyond the K/V tiles (Q, P^T and the running softmax stats are register-resident), so it is the full auxiliary term. Scales with ``num_warps``, ``block_m_per_warp`` and ``head_size``.""" - out_stripe_cols = 32 if spec.head_size <= 64 else spec.head_size block_m = spec.num_warps * spec.block_m_per_warp - return block_m * out_stripe_cols * _ACC_LDS_ELEM_BYTES + return block_m * _out_stripe_cols(spec.head_size) * _ACC_LDS_ELEM_BYTES def _lds_bytes_regpv(spec) -> int: """Exact LDS footprint of the register-PV 2D path: the K/V tiles plus the - epilogue Acc_lds. Verified against comgr for the D256 T=128 K-double geometry: - ``(2+1)*128*256*2 (K/V) + 16*256*2 (Acc) = 196608 + 8192 = 204800``.""" - wb = _kv_lds_elem_bytes(spec) - kbuf = 1 if spec.use_k_single_buffer else 2 - vbuf = 2 if spec.use_v_double_buffer else 1 - kv_bytes = (kbuf + vbuf) * spec.tile_size * spec.head_size * wb - return kv_bytes + _acc_lds_bytes(spec) + epilogue Acc_lds (Q/P^T/O register-resident). Verified against comgr for the + D256 T=128 K-double geometry: ``(2+1)*128*256*2 (K/V) + 16*256*2 (Acc) = + 196608 + 8192 = 204800``.""" + return _tiled_2d_lds_bytes( + tile_size=spec.tile_size, + head_size=spec.head_size, + block_m=spec.num_warps * spec.block_m_per_warp, + kv_elem_bytes=_kv_lds_elem_bytes(spec), + # getattr: the gfx942 spec class declares fewer schedule fields than gfx950, + # so default the absent ones (K double-buffered, V single-buffered). + k_slots=1 if getattr(spec, "use_k_single_buffer", False) else 2, + v_slots=2 if getattr(spec, "use_v_double_buffer", False) else 1, + include_q_lds=False, # register-PV keeps Q in registers + include_p_lds=False, # register-PV keeps P^T in registers + ) def _lds_capacity_bytes() -> int: @@ -1867,24 +1912,29 @@ def _lds_capacity_bytes() -> int: def _ldsfix_single_k(spec): + """Candidate: drop K's second (prefetch) buffer. Returns ``(spec, None)`` on + success, or ``(None, reason)`` when the lever cannot apply -- the reason is + surfaced in the resolver's diagnostic rather than silently swallowed.""" if spec.use_k_single_buffer: - return None + return None, "K already single-buffered" # K-single needs Q to fit the lone K slot: block_m <= tile_size. if spec.num_warps * spec.block_m_per_warp > spec.tile_size: - return None + return None, "block_m > tile_size (Q would not fit the lone K slot)" try: - return replace(spec, use_k_single_buffer=True) - except Exception: - return None + return replace(spec, use_k_single_buffer=True), None + except ValueError as e: # __post_init__ rejected the new combo + return None, f"rejected by spec validation: {e}" def _ldsfix_tile64(spec): + """Candidate: shrink the KV tile to T=64. Returns ``(spec, None)`` or + ``(None, reason)`` (see :func:`_ldsfix_single_k`).""" if spec.tile_size <= 64: - return None + return None, "tile_size already <= 64" try: - return replace(spec, tile_size=64) - except Exception: - return None + return replace(spec, tile_size=64), None + except ValueError as e: # __post_init__ rejected the new combo + return None, f"rejected by spec validation: {e}" def _resolve_lds_budget(spec): @@ -1902,17 +1952,39 @@ def _resolve_lds_budget(spec): cap = _lds_capacity_bytes() if _lds_bytes_regpv(spec) <= cap: return spec - single_k = _ldsfix_single_k(spec) - candidates = [single_k, _ldsfix_tile64(spec)] - if single_k is not None: - candidates.append(_ldsfix_tile64(single_k)) - for cand in candidates: - if cand is not None and _lds_bytes_regpv(cand) <= cap: + + tried = [] # per-lever diagnostics, surfaced if nothing fits + + def _consider(label, cand, why): + if cand is None: + tried.append(f"{label} n/a ({why})") + return None + used = _lds_bytes_regpv(cand) + if used <= cap: return cand + tried.append(f"{label} still {used} B > {cap} B") + return None + + # Cheapest-first ladder: single-buffer K, then T=64, then both. + single_k, why = _ldsfix_single_k(spec) + hit = _consider("single-K", single_k, why) + if hit is not None: + return hit + t64, why = _ldsfix_tile64(spec) + hit = _consider("T=64", t64, why) + if hit is not None: + return hit + if single_k is not None: + both, why = _ldsfix_tile64(single_k) + hit = _consider("single-K+T=64", both, why) + if hit is not None: + return hit + raise RuntimeError( f"LDS budget: 2D register-PV D{spec.head_size} block_size={spec.block_size} " f"T={spec.tile_size} needs {_lds_bytes_regpv(spec)} B > cap {cap} B on " - f"{_resolve_attention_arch()}; no validated reduction (single-K, T=64) fits." + f"{_resolve_attention_arch()}; no validated reduction fits " + f"[{'; '.join(tried)}]." ) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py index e4df037aa2ff..cac939b9edbc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_tiled_2d.py @@ -1141,13 +1141,20 @@ def supports_tiled_2d( f"the {arch} {_LDS_CAPACITY_BYTES} B LDS budget", ) return True, "supported" - _out_stripe = 32 if head_size <= 64 else head_size - _lds = ( - 2 * _t_eff * head_size * _BPE # K_lds (double-buffered) - + _t_eff * head_size * _BPE # V_lds - + _block_m * (_t_eff + 8) * _BPE # P_lds - + (0 if _block_m <= 2 * _t_eff else _block_m * head_size * _BPE) # Q_lds - + _block_m * _out_stripe * _BPE # Acc_lds + # Default-path footprint via the shared model (single source of truth, + # also used by the gfx950 register-PV resolver). Conservative here: + # Q_lds/P_lds are staged (the register-PV path drops them). + from ..common.attention_unified import _tiled_2d_lds_bytes + + _lds = _tiled_2d_lds_bytes( + tile_size=_t_eff, + head_size=head_size, + block_m=_block_m, + kv_elem_bytes=_BPE, + k_slots=2, + v_slots=1, + include_q_lds=True, + include_p_lds=True, ) if _lds > _LDS_CAPACITY_BYTES: return ( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py index 26650c8d7d34..25fcb29d0d1e 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_lds_budget_resolver.py @@ -73,14 +73,66 @@ def test_resolver_is_noop_off_gfx950(self): with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx942"): self.assertIs(au._resolve_lds_budget(spec), spec) - def test_acc_lds_scales_with_geometry(self): - """The auxiliary term is the epilogue Acc_lds, computed from geometry - (BLOCK_M x OUT_STRIPE_COLS x 2), not a hard-coded constant.""" + def test_regpv_footprint_uses_shared_helper(self): + """W: the register-PV footprint is the shared ``_tiled_2d_lds_bytes`` model + parameterised for the register-PV layout (Q/P^T dropped) -- a single source + of truth, so it cannot drift from the gfx942 admission gate.""" with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) block_m = spec.num_warps * spec.block_m_per_warp - out_stripe = 32 if spec.head_size <= 64 else spec.head_size - self.assertEqual(au._acc_lds_bytes(spec), block_m * out_stripe * 2) + expected = au._tiled_2d_lds_bytes( + tile_size=spec.tile_size, + head_size=spec.head_size, + block_m=block_m, + kv_elem_bytes=au._kv_lds_elem_bytes(spec), + k_slots=1 if spec.use_k_single_buffer else 2, + v_slots=2 if spec.use_v_double_buffer else 1, + ) + self.assertEqual(au._lds_bytes_regpv(spec), expected) + + def test_shared_helper_models_both_layouts(self): + """W: one helper serves both consumers -- the conservative gfx942 gate + layout (Q_lds/P_lds staged) and the exact register-PV layout (dropped).""" + common = dict(tile_size=128, head_size=256, block_m=16, kv_elem_bytes=2) + conservative = au._tiled_2d_lds_bytes( + k_slots=2, v_slots=1, include_q_lds=True, include_p_lds=True, **common + ) + regpv = au._tiled_2d_lds_bytes(k_slots=2, v_slots=1, **common) + # Dropping the P_lds term makes the register-PV footprint the smaller of + # the two; the conservative gate estimate never under-counts it. + self.assertGreater(conservative, regpv) + self.assertEqual(regpv, 204800) # matches the comgr-verified D256 number + + def test_ldsfix_levers_report_reasons(self): + """S: the shrink levers return ``(None, reason)`` when they cannot apply, + so the resolver surfaces *why* rather than swallowing it.""" + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + # single-K applies (returns a spec, no reason) ... + cand, why = au._ldsfix_single_k(replace(spec, use_k_single_buffer=False)) + self.assertIsNotNone(cand) + self.assertIsNone(why) + # ... but is a no-op once K is already single-buffered (reason given). + cand, why = au._ldsfix_single_k(replace(spec, use_k_single_buffer=True)) + self.assertIsNone(cand) + self.assertIn("single-buffered", why) + # tile64 is a no-op at T<=64. + cand, why = au._ldsfix_tile64(replace(spec, tile_size=64)) + self.assertIsNone(cand) + self.assertIn("64", why) + + def test_infeasible_budget_raises_with_lever_diagnostics(self): + """S: when no lever fits, the resolver raises with each attempted lever's + result -- a diagnostic, not a cryptic downstream comgr abort.""" + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + # Force an impossibly small cap so even single-K + T=64 overflows. + with mock.patch.object(au, "_lds_capacity_bytes", return_value=1024): + with self.assertRaises(RuntimeError) as ctx: + au._resolve_lds_budget(replace(spec, use_k_single_buffer=False)) + msg = str(ctx.exception) + self.assertIn("single-K", msg) + self.assertIn("T=64", msg) if __name__ == "__main__": diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h index d08a2ca651a6..96a9cb328c5f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/include/rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h @@ -37,6 +37,9 @@ #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/ir.h" +#include /* size_t */ +/* rocke_attention_tiled_2d_spec_t (the shared tiled-2D spec mirror). */ +#include "rocke/helper_helper_rocke.instances.gfx942.attention_tiled_2d.h" #ifdef __cplusplus extern "C" { @@ -119,6 +122,36 @@ const char* rocke_unified_attn_kv_storage_dtype(const rocke_unified_attn_problem */ bool rocke_unified_attn_select_2d_waves_per_eu(const rocke_unified_attn_problem_t* p, int* out_wpe); +/* ------------------------------------------------------- LDS-budget resolver * + * Selection-layer LDS accounting shared by the gfx942 admission gate and the + * gfx950 register-PV budget resolver (mirrors the Python helpers in + * rocke/instances/common/attention_unified.py). */ + +/* Python: _tiled_2d_lds_bytes -- static LDS footprint (bytes) of a tiled-2D + * geometry. Single source of truth for the per-buffer tile arithmetic; callers + * pick which buffers are LDS-resident (K/V slots, Q/P_lds, the block_m<=2T Q + * alias, the OUT_STRIPE rule, the 16-bit Acc staging width). */ +int rocke_unified_attn_tiled_2d_lds_bytes(int tile_size, + int head_size, + int block_m, + int kv_elem_bytes, + int k_slots, + int v_slots, + bool include_q_lds, + bool include_p_lds, + int v_pad); + +/* Python: _resolve_lds_budget -- deterministically shrink an over-budget gfx950 + * register-PV 2D spec (single-buffer K, then T=64) in place until it fits the + * arch LDS cap, using only compile-validated reductions. A strict no-op for + * already-fitting / non-register-PV / non-gfx950 specs (returns true, spec + * unchanged). Returns false when no validated reduction fits, filling ``reason`` + * (up to ``reason_cap`` bytes) with a diagnostic. The provider's spec-assembly + * path calls this right after building a spec from the selectors above. */ +bool rocke_unified_attn_resolve_lds_budget(rocke_attention_tiled_2d_spec_t* spec, + char* reason, + size_t reason_cap); + /* ------------------------------------------------ 2D feature-gate predicates * * Exposed mirrors of the Python _enable_* gates that _tiled_spec_from_problem * consults when it builds the per-shape UnifiedAttention2DTiledSpec. The diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp index 864c7133a0ac..408a85df388c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp @@ -19,6 +19,7 @@ #include #include +#include "rocke/arch_target.h" /* rocke_arch_target_from_gfx, lds_capacity_bytes */ #include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" #include "rocke/helper_rocke.helpers.transforms.h" #include "rocke/ir.h" @@ -753,6 +754,152 @@ bool rocke_unified_attn_enable_early_v_schedule(const rocke_unified_attn_problem return enable_early_v_schedule(p); } +/* --------------------------------------------------- LDS-budget resolver */ +/* Python: _out_stripe_cols / _tiled_2d_lds_bytes / _kv_lds_elem_bytes / + * _lds_bytes_regpv / _ldsfix_single_k / _ldsfix_tile64 / _resolve_lds_budget + * (rocke/instances/common/attention_unified.py). Mirrors the selection-layer + * LDS-budget shrink so a gfx950 register-PV D256 spec that would overflow the + * 160 KB cap is deterministically reduced (single-buffer K, then T=64) instead + * of failing comgr CODEGEN. Byte-faithful with the Python arithmetic; the + * provider spec-assembly path calls the resolver right after it builds a spec + * from the selectors above. */ + +/* Acc is always 16-bit: the fp8 lever changes only the K/V *cache* width, never + * the output accumulator (mirrors Python _ACC_LDS_ELEM_BYTES). */ +#define ROCKE_ATTN_ACC_LDS_ELEM_BYTES 2 + +static int attn_out_stripe_cols(int head_size) +{ + return (head_size <= 64) ? 32 : head_size; +} + +int rocke_unified_attn_tiled_2d_lds_bytes(int tile_size, + int head_size, + int block_m, + int kv_elem_bytes, + int k_slots, + int v_slots, + bool include_q_lds, + bool include_p_lds, + int v_pad) +{ + int k_bytes = k_slots * tile_size * head_size * kv_elem_bytes; + int v_bytes = v_slots * (tile_size + v_pad) * head_size * kv_elem_bytes; + int p_bytes = include_p_lds ? block_m * (tile_size + 8) * kv_elem_bytes : 0; + /* Q aliases the K slab when it fits under the double-buffer window. */ + int q_bytes + = (include_q_lds && block_m > 2 * tile_size) ? block_m * head_size * kv_elem_bytes : 0; + int acc_bytes = block_m * attn_out_stripe_cols(head_size) * ROCKE_ATTN_ACC_LDS_ELEM_BYTES; + return k_bytes + v_bytes + p_bytes + q_bytes + acc_bytes; +} + +/* Byte width of one K/V element as stored in LDS: 1 for the fp8 cache, else 2 (bf16/fp16). */ +static int attn_kv_lds_elem_bytes(const rocke_attention_tiled_2d_spec_t* s) +{ + const char* kv = s->kv_storage_dtype; + if(kv != NULL + && (strcmp(kv, "fp8") == 0 || strcmp(kv, "bf8") == 0 || strcmp(kv, "e4m3") == 0 + || strcmp(kv, "e5m2") == 0)) + return 1; + return 2; /* bf16 / fp16 K/V in LDS */ +} + +static int attn_spec_tile_size_eff(const rocke_attention_tiled_2d_spec_t* s) +{ + return s->has_tile_size ? s->tile_size : s->block_size; +} + +/* Python: _lds_bytes_regpv -- exact register-PV footprint (K/V tiles + Acc; + * Q/P^T/O register-resident). */ +static int attn_lds_bytes_regpv(const rocke_attention_tiled_2d_spec_t* s) +{ + int block_m = s->num_warps * s->block_m_per_warp; + return rocke_unified_attn_tiled_2d_lds_bytes(attn_spec_tile_size_eff(s), + s->head_size, + block_m, + attn_kv_lds_elem_bytes(s), + s->use_k_single_buffer ? 1 : 2, + s->use_v_double_buffer ? 2 : 1, + /*include_q_lds=*/false, + /*include_p_lds=*/false, + /*v_pad=*/0); +} + +bool rocke_unified_attn_resolve_lds_budget(rocke_attention_tiled_2d_spec_t* spec, + char* reason, + size_t reason_cap) +{ + const char* arch = rocke_unified_attn_resolve_arch(); + /* Validated on gfx950 (CDNA4) register-PV only; other arches' 2D footprint + * models are not yet validated, so engage only where proven. No-op = fits, + * which keeps every currently-compiling spec byte-identical. */ + if(strcmp(arch, "gfx950") != 0 || !spec->use_register_pv) + return true; + const rocke_arch_target_t* t = rocke_arch_target_from_gfx(arch); + int cap = (t != NULL) ? t->lds_capacity_bytes : 163840; + if(attn_lds_bytes_regpv(spec) <= cap) + return true; + + /* Cheapest-first ladder: single-buffer K, then T=64, then both. Each lever + * mirrors the Python precondition guards (single-K needs the lone K slot to + * hold Q: block_m <= tile_size; T=64 only shrinks a larger tile). The + * dataclass __post_init__ validate() is the builder's backstop at build + * time; these preconditions are what actually gate the levers here. */ + int block_m = spec->num_warps * spec->block_m_per_warp; + int t_eff = attn_spec_tile_size_eff(spec); + bool single_k_ok = !spec->use_k_single_buffer && block_m <= t_eff; + + /* Lever 1: single-buffer K. */ + if(single_k_ok) + { + rocke_attention_tiled_2d_spec_t c = *spec; + c.use_k_single_buffer = true; + if(attn_lds_bytes_regpv(&c) <= cap) + { + *spec = c; + return true; + } + } + /* Lever 2: shrink the KV tile to T=64. */ + if(t_eff > 64) + { + rocke_attention_tiled_2d_spec_t c = *spec; + c.tile_size = 64; + c.has_tile_size = true; + if(attn_lds_bytes_regpv(&c) <= cap) + { + *spec = c; + return true; + } + } + /* Lever 3: single-K + T=64 (only when single-K's precondition held). */ + if(single_k_ok && t_eff > 64) + { + rocke_attention_tiled_2d_spec_t c = *spec; + c.use_k_single_buffer = true; + c.tile_size = 64; + c.has_tile_size = true; + if(attn_lds_bytes_regpv(&c) <= cap) + { + *spec = c; + return true; + } + } + + if(reason != NULL && reason_cap > 0) + snprintf(reason, + reason_cap, + "LDS budget: 2D register-PV D%d block_size=%d T=%d needs %d B > cap %d B " + "on %s; no validated reduction (single-K, T=64) fits", + spec->head_size, + spec->block_size, + t_eff, + attn_lds_bytes_regpv(spec), + cap, + arch); + return false; +} + /* ----------------------------------------------------------- magic div */ rocke_value_t* diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp index 07f83b702125..55b39541ed9c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/gfx942/attention_tiled_2d_public_entry_glue.cpp @@ -38,6 +38,7 @@ #include "rocke/instance_gfx942_attention_tiled_2d.h" #include "rocke/instance_gfx942_attention_tiled_2d_internal.h" +#include "rocke/helper_helper_rocke.instances.common.attention_unified_selectors.h" /* rocke_unified_attn_tiled_2d_lds_bytes */ #include "rocke/helper_rocke.helpers.spec.h" /* rocke_kernel_name_join */ #include "rocke/error_boundary.hpp" /* ckc::guard_builder boundary shim */ @@ -373,13 +374,18 @@ bool rocke_gfx942_attention_tiled_2d_supports( } } { - int _out_stripe = (args->head_size <= 64) ? 32 : args->head_size; - int _lds - = 2 * _t_eff * args->head_size * _BPE /* K_lds (double) */ - + _t_eff * args->head_size * _BPE /* V_lds */ - + _block_m * (_t_eff + 8) * _BPE /* P_lds */ - + ((_block_m <= 2 * _t_eff) ? 0 : _block_m * args->head_size * _BPE) /* Q_lds */ - + _block_m * _out_stripe * _BPE; /* Acc_lds */ + /* Default-path footprint via the shared model (single source of + * truth, also used by the gfx950 register-PV resolver). Conservative + * here: Q_lds/P_lds staged (the register-PV path drops them). */ + int _lds = rocke_unified_attn_tiled_2d_lds_bytes(_t_eff, + args->head_size, + _block_m, + _BPE, + /*k_slots=*/2, + /*v_slots=*/1, + /*include_q_lds=*/true, + /*include_p_lds=*/true, + /*v_pad=*/0); if(_lds > _LDS_CAPACITY_BYTES) { snprintf(buf, From 53824f8fd264e417d2d83157934748704dc72dde Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 8 Jul 2026 21:47:28 -0400 Subject: [PATCH 03/23] style(rocke/attention): black-format merged _tiled_spec_from_problem --- .../rocke/library/kernels/common/attention_unified.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index a288ee6dd145..946347ccf641 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -2036,6 +2036,7 @@ def _tiled_spec_from_problem( from builders.common.attention_spec_builder import ( _tiled_spec_from_problem as _impl, ) + _spec = _impl(problem) return _resolve_lds_budget(_spec) From a7c8fb2353d13124e947a0dcda34dab9485ce1ba Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 8 Jul 2026 22:25:27 -0400 Subject: [PATCH 04/23] fix(rocke/attention): recognize fp8e4m3/bf8e5m2 in shared LDS kv-elem-bytes The shared _tiled_2d_lds_bytes footprint model counted fp8 K/V cache tiles at 2 B/elem: _kv_storage_dtype() emits "fp8e4m3" but _kv_lds_elem_bytes() (Py) and attn_kv_lds_elem_bytes() (C++) only matched "fp8"/"bf8"/"e4m3"/"e5m2". Add the full dtype strings so fp8-cache specs are modeled at 1 B/elem in both engines. bf16/fp16 (incl. D256) unaffected; byte-identity green. --- .../rocke/library/kernels/common/attention_unified.py | 9 ++++++++- .../cpp/instances/common/attention_unified_selectors.cpp | 4 ++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 946347ccf641..2bd73cb0a1dc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -1875,7 +1875,14 @@ def _enable_i64_kv_addr(problem: UnifiedAttentionProblem) -> bool: def _kv_lds_elem_bytes(spec) -> int: # Byte width of one K/V element as stored in LDS: 1 for the fp8 cache, else 2 (bf16/fp16). - if getattr(spec, "kv_storage_dtype", None) in ("fp8", "bf8", "e4m3", "e5m2"): + if getattr(spec, "kv_storage_dtype", None) in ( + "fp8e4m3", + "bf8e5m2", + "fp8", + "bf8", + "e4m3", + "e5m2", + ): return 1 return 2 # bf16 / fp16 K/V in LDS diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp index 408a85df388c..26bfcab03a73 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp +++ b/dnn-providers/hip-kernel-provider/rocke/platform/cpp/instances/common/attention_unified_selectors.cpp @@ -798,8 +798,8 @@ static int attn_kv_lds_elem_bytes(const rocke_attention_tiled_2d_spec_t* s) { const char* kv = s->kv_storage_dtype; if(kv != NULL - && (strcmp(kv, "fp8") == 0 || strcmp(kv, "bf8") == 0 || strcmp(kv, "e4m3") == 0 - || strcmp(kv, "e5m2") == 0)) + && (strcmp(kv, "fp8e4m3") == 0 || strcmp(kv, "bf8e5m2") == 0 || strcmp(kv, "fp8") == 0 + || strcmp(kv, "bf8") == 0 || strcmp(kv, "e4m3") == 0 || strcmp(kv, "e5m2") == 0)) return 1; return 2; /* bf16 / fp16 K/V in LDS */ } From bf086c17232d7e97a991c15dbd49e8e380dd8e6c Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Thu, 9 Jul 2026 13:23:01 -0400 Subject: [PATCH 05/23] feat(rocke/attention): D256 gfx950 bf16 prefill 32x32 + softmax-MFMA interleave fast path Route the gfx950 D256 bf16 prefill cohort to the 32x32 transposed fast path plus the FA3-style softmax<->MFMA interleave (mode2, groups=4). Gates the three launch-critical geometry selectors (num_warps / tile_size / block_m_per_warp) so the builder spec, _tiled_cache_key and _get_2d_launch_meta all agree; the codegen flags (mfma_32x32, transposed stack, register_pv off, K single-buffer, interleave) are applied by the spec override in _tiled_spec_from_problem. Narrow + arch/feature-guarded: every other shape is byte-identical; the kernel codegen the flags trigger already ships on the base. Validated on MI355X (direct-launch sweep, fp32-ref max_abs=1.5625e-02): +9.1% @ Sq4096 (1.29x torch-flash) and +10.1% @ Sq8192 over the 32x32 base. --- .../kernels/common/attention_unified.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 2bd73cb0a1dc..a15eb4e84e58 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -585,6 +585,34 @@ def _enable_k_single_buffer(problem: UnifiedAttentionProblem) -> bool: return _enable_d128_small_tile(problem) and problem.block_size >= 32 +def _d256_gfx950_fast(problem: "UnifiedAttentionProblem") -> bool: + """Route the D256 gfx950 bf16 prefill cohort to the 32x32 + transposed fast path + FA3-style softmax<->MFMA interleave. + + Gates the three launch-critical geometry selectors (num_warps / tile_size / + block_m_per_warp) so the builder spec, ``_tiled_cache_key`` and + ``_get_2d_launch_meta`` all agree; the remaining codegen flags (32x32 MFMA, + transposed stack, register_pv off, K single-buffer, interleave mode2/g4) are + applied by the spec override in ``_tiled_spec_from_problem``. Narrow + + arch/feature-guarded: every other shape is byte-identical. + + Perf (MI355X, direct-launch sweep, fp32-ref max_abs=1.5625e-02): +interleave + mode2 g4 gives +9.1% @ Sq4096 and +10.1% @ Sq8192 over the 32x32 base. + """ + return ( + _resolve_attention_arch() == "gfx950" + and problem.head_size == 256 + and problem.dtype == "bf16" + and not problem.use_fp8 + and problem.sliding_window == 0 + and problem.softcap == 0 + and not problem.use_sinks + and not problem.use_alibi + and not problem.use_qq_bias + and problem.max_seqlen_q > 1 + ) + + def _select_2d_tile_size(problem: UnifiedAttentionProblem) -> int: """Choose ``tile_size`` (T) for the tiled 2D kernel. @@ -612,6 +640,8 @@ def _select_2d_tile_size(problem: UnifiedAttentionProblem) -> int: ``T = block_size`` (single block per iter, 32 tokens) — measured 1.15-1.30× win on every FP8 SW long-prefill shape. """ + if _d256_gfx950_fast(problem): + return 64 if _resolve_attention_arch() == "gfx1250": # gfx1250 v1 consumes exactly one 32-token paged-KV block per WMMA # iteration; wider T needs separate multi-block block-table handling. @@ -754,6 +784,8 @@ def _select_2d_num_warps(problem: UnifiedAttentionProblem) -> int: - an LDS-budget check (``<= 96 KiB`` so we keep >= 1 CTA/CU on MI355X comfortably). """ + if _d256_gfx950_fast(problem): + return 2 if _resolve_attention_arch() == "gfx1250": # A gfx1250 workgroup is one wave32 in the v1 WMMA tiled path. return 1 @@ -2045,6 +2077,39 @@ def _tiled_spec_from_problem( ) _spec = _impl(problem) + # D256 gfx950 bf16 prefill fast route. The default builder routes + # head_size==256 to the 16x16 register_pv body (~0.24x torch.compile). The + # 32x32 transposed constellation is LEGAL for D256 (use_mfma_32x32 needs only + # head_size%16==0) and ~4.5-5.9x faster; the FA3-style softmax<->MFMA + # interleave (mode2 g4) adds a further +9-10% by hiding the exposed v_exp + # behind the QK/PV MFMA at the kernel's low (4 waves/CU) occupancy. Geometry + # (num_warps / tile_size / block_m_per_warp) is set by the gated selectors + # above so the cache key + launch meta agree; here we set the codegen flags. + # The cohort is uniquely keyed by head_size==256, so the (unkeyed) interleave + # + codegen overrides cannot collide with any other shape. _resolve_lds_budget + # no-ops here (register_pv=False, fits the 160 KB cap at 64 KB). + if _d256_gfx950_fast(problem): + _spec = replace( + _spec, + use_mfma_32x32=True, + use_transposed_qk_32x32=True, + use_q_direct_reg=True, + use_transposed_half_local_pv=True, + use_transposed_scalar_state=True, + use_transposed_mask_once=True, + use_transposed_mask_limit=True, + use_mask_phase_split=True, + use_register_pv=False, + use_k_single_buffer=True, + use_v_double_buffer=False, + use_early_v_schedule=False, + use_sched_barrier=False, + use_softmax_mfma_interleave=True, + softmax_interleave_mode=2, + softmax_interleave_groups=4, + use_fast_paged_kv_desc=False, + use_mfma32_skip_legacy_qreg=False, + ) return _resolve_lds_budget(_spec) @@ -2089,6 +2154,8 @@ def _select_2d_block_m_per_warp(problem: UnifiedAttentionProblem) -> int: Gate: ``max_seqlen_q > 256 and num_seqs >= 2``. Below this, mw=16 is consistently within noise of mw=32 in the per-shape sweep. """ + if _d256_gfx950_fast(problem): + return 32 if _resolve_attention_arch() == "gfx1250": return 16 # mw=32 (BLOCK_M = 32 * num_warps) only pays off when a path actually From 2dff7af75503d89e363b8241f9adae88dc10a5af Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Thu, 9 Jul 2026 13:19:37 -0400 Subject: [PATCH 06/23] =?UTF-8?q?docs(runbook):=20Step=200=20caveat=20?= =?UTF-8?q?=E2=80=94=20the=20preset=20bench=20is=20not=20the=20exhaustive?= =?UTF-8?q?=20lever=20sweep?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The live preset bench (_variant_flags menu) can't express default-off levers (use_softmax_mfma_interleave / use_k_single_buffer / use_register_pv), so running it does not satisfy Step 0. Point to the raw-flag sweep (probe_config_sweep.py / _select_2d_* monkey-patch) and the autotuner, with the D256 interleave miss as the concrete failure. --- .../optimization/optimization_runbook.md | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md index 642dc8d8f251..df56c5a1733c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md +++ b/dnn-providers/hip-kernel-provider/rocke/platform/dsl_docs/optimization/optimization_runbook.md @@ -46,6 +46,36 @@ For the specific shape you are optimizing: has told you which config to redesign *from* and against which resource budget. +> **Caveat — the shipped preset bench is NOT this sweep.** The live benchmark +> most people reach for +> (`library/benchmarks/gfx{942,950}/attention/prefill/benchmark_prefill2d_live.py`) +> sweeps a *curated menu of named presets* (`_variant_flags`: +> `prod`/`combo`/`fallback`/`r4_t32`/…), **not** the full lever space above. That +> grammar exposes geometry (`num_warps`, `tile_size`, `block_m_per_warp`, MFMA +> atom, transpose flags) but **cannot express several default-off flags** — +> notably `use_softmax_mfma_interleave` (+ `softmax_interleave_mode` / +> `softmax_interleave_groups`), `use_k_single_buffer`, and `use_register_pv`. So +> running the preset bench does **not** satisfy Step 0: it systematically skips +> exactly the default-off levers this step flags as most likely mis-picked. The +> space it measures is `union(preset menu, production-dispatcher output)`; +> anything outside that — including winning configs — stays invisible until a +> human adds it as a candidate. +> +> To actually exhaust the space, sweep the **raw flags**, bypassing both the +> preset grammar and the dispatcher: build the spec from arbitrary flags and +> launch it directly +> ([`utilities/tools/dsl_probes/probe_config_sweep.py`](utilities/tools/dsl_probes/probe_config_sweep.py), +> or the `_select_2d_*` monkey-patch pattern in [§4.3 Attention](#43-attention)), +> or drive the autotuner ([§12 Autotune](#12-autotuning-strategy)) over the full +> cartesian product. +> +> *Real failure:* the D256 gfx950 softmax↔MFMA interleave win (+9–10% at +> Sq 4096/8192, correctness-clean) was invisible to the preset-bench sweep for a +> full optimization loop, because `use_softmax_mfma_interleave` is not in the +> `_variant_flags` grammar. A raw-lever direct-launch sweep found it in one pass +> once the profiler pointed at the exposed softmax on a low-occupancy +> (4 waves/CU) kernel. + This ordering is not optional. Real failures from skipping it: a selector routed a shape to a ~2× slower configuration that an exhaustive sweep beat immediately; a gap assumed "structural" turned out to be a single mis-set From 642e5dc5d1dc1cc4f5cf253ec2b813c195cb56c4 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Thu, 9 Jul 2026 20:35:08 -0400 Subject: [PATCH 07/23] fix(rocke/attention): key the D256 fast-path override in the tiled cache key The D256 gfx950 fast route applies a codegen override in _tiled_spec_from_problem (32x32 transposed stack + softmax<->MFMA interleave), but _tiled_cache_key recomputed those flags from the ungated _enable_* helpers, so the key recorded e.g. use_transposed_qk_32x32=False / use_register_pv=True and omitted the interleave flags entirely -- the key no longer described the built kernel. Add the _d256_gfx950_fast(problem) cohort bit to the key. It flips exactly when the override applies, so key === built kernel again, without duplicating the override's ~15 flags. Keying on the predicate (not the flags) deliberately keeps the override a Layer-2 hand-pinned exception ABOVE the autotuner-regenerable _enable_* helpers, rather than folding it into tables a future sweep->gate is meant to overwrite. --- .../rocke/library/kernels/common/attention_unified.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index a15eb4e84e58..e4d227b1604f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -1027,6 +1027,13 @@ def _tiled_cache_key(problem: UnifiedAttentionProblem) -> Tuple: _select_2d_tile_size(problem), _select_2d_waves_per_eu(problem), _select_2d_block_m_per_warp(problem), + # Layer-2 override marker: flips the key exactly when the D256 gfx950 + # codegen override applies in `_tiled_spec_from_problem`, so the key stays + # faithful to the built kernel (key === built kernel) without duplicating + # the override's flags here. The override is a hand-pinned exception that + # deliberately lives ABOVE the autotuner-regenerable `_enable_*` helpers + # below; keying on the predicate (not the flags) keeps that layering. + _d256_gfx950_fast(problem), _enable_mfma_32x32(problem), _enable_transposed_qk_32x32(problem), _enable_transposed_half_local_pv(problem), From 2f4dc9390ecfff938aee8c7656a41fe91f9d48cb Mon Sep 17 00:00:00 2001 From: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com> Date: Thu, 9 Jul 2026 17:44:28 -0700 Subject: [PATCH 08/23] feat(rocke): gfx950 slab-granularity K_lds pad bank conflict fix (#9260) Break the row-aliased LDS bank conflict on the D256 gfx950 bf16 prefill QK K read by padding K_lds between the 2-row async-DMA slabs. The async DMA writes one WAVE_BYTES (2-row) contiguous slab per wave/call, so padding BETWEEN slabs keeps the write contiguous (DMA-safe) while shifting each slab's banks; the QK read remaps (row,col) -> (row/2, (row%2)*HD + col) so write and read stay consistent (numerics preserved). Gated behind use_kq_lds_pad / kq_lds_pad_halves (default off -> byte identical; byte-identity gate GREEN). Wired pad=16 into the D256 route in _tiled_spec_from_problem. Measured on MI355X (direct-launch, fp32-ref max_abs=1.5625e-2): SQ_LDS_BANK_CONFLICT rate 897 -> 660 per 1000 (-4.5x conflict cycles) Sq8192 ~1592 -> ~1195 us (~25% faster); vs torch-flash 0.86x -> 1.14x Sq4096 vs torch-flash 1.27x -> ~1.6-2.0x Occupancy-neutral: LDS 65 KB <= 80 KB budget keeps 2 WG/CU; VGPR unchanged. Python-engine only for now (the C++ twin NOTIMPL-rejects the host use_q_direct_reg / use_softmax_mfma_interleave path). Also adds the parity harness lanes used to develop this: direct-launch D256 lane, occupancy probe, torch-flash comparison, interleave toggle, and pad/xor-swizzle probes. ISSUE ID : AICK-1495 ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../prefill/parity_unified_attention.py | 381 ++++++++++++++++++ .../kernels/common/attention_unified.py | 7 + .../kernels/gfx950/attention_tiled_2d.py | 97 ++++- 3 files changed, 481 insertions(+), 4 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index aa144f0d6814..81cb7306cb5f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -852,6 +852,104 @@ def creative_scenarios() -> List[Scenario]: block_size=16, dtype=torch.bfloat16, ), + # --- PR #9233: D256 gfx950 bf16 prefill fast-path cohort (Sq bs64) --- + # These are the shapes the PR perf table reports (+9.1% @ Sq4096, + # +10.1% @ Sq8192 over the 32x32 base). bf16, head_size=256, causal + # prefill (query_len == kv_len), batch 64. num_blocks covers 64 seqs * + # (Sq / block_size) unique blocks so the paged block-table is dense. + Scenario( + name="pr9233_d256_bf16_prefill_sq4096_bs64", + seq_lens=[(4096, 4096)] * 64, + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=64 * (4096 // 16), + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_bs64", + seq_lens=[(8192, 8192)] * 64, + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=64 * (8192 // 16), + ), + # bs64 == block_size=64 interpretation, single long sequence (batch 1): + # a low-occupancy compute-bound prefill where the exposed softmax v_exp + # is the bottleneck the interleave targets. + Scenario( + name="pr9233_d256_bf16_prefill_sq4096_b1_blk64", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=4096 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_blk64", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_h8", + seq_lens=[(8192, 8192)], + num_query_heads=8, + num_kv_heads=8, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_gqa16x2", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_h1", + seq_lens=[(8192, 8192)], + num_query_heads=1, + num_kv_heads=1, + head_size=256, + block_size=64, + dtype=torch.bfloat16, + num_blocks=8192 // 64, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq4096_b1_blk16", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=4096 // 16, + ), + Scenario( + name="pr9233_d256_bf16_prefill_sq8192_b1_blk16", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=16, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + num_blocks=8192 // 16, + ), Scenario( name="creative_d256_prefill", seq_lens=[(128, 1024), (256, 2048)], @@ -1186,9 +1284,88 @@ def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): _attn_signature, _attn_values, _select_2d_compile_backend, + _d256_gfx950_fast, + _tiled_spec_from_problem, ) from rocke.runtime import KernelLauncher, LaunchConfig + # PR #9233 direct-launch lane: for the D256 gfx950 bf16 prefill cohort, + # build and launch the PRODUCTION fast spec directly (GPU-time only, no + # run_unified_attention_torch dispatcher overhead). Honors the + # PARITY_NO_INTERLEAVE monkeypatch so the same lane measures both the + # interleave fast path and the 32x32 base. + if _d256_gfx950_fast(problem): + spec = _tiled_spec_from_problem(problem) + if os.environ.get("PARITY_KQ_SWIZZLE") == "1": + from dataclasses import replace as _replace + + spec = _replace(spec, use_kq_xor_swizzle=True) + print( + " [2d-direct] KQ XOR swizzle ON (read-only probe; " + "numerics intentionally wrong)" + ) + if os.environ.get("PARITY_KQ_PAD") == "1": + from dataclasses import replace as _replace + + _w = int(os.environ.get("PARITY_KQ_PAD_W", "8")) + spec = _replace(spec, use_kq_lds_pad=True, kq_lds_pad_halves=_w) + print( + f" [2d-direct] KQ slab-gran pad ON (pad={_w} halves, " + "correct: numerics preserved)" + ) + kernel = build_unified_attention_2d_tiled(spec) + if _select_2d_compile_backend(problem) == "hipcc": + from rocke.helpers.compile import compile_kernel_via_hipcc + + artifact = compile_kernel_via_hipcc(kernel) + else: + artifact = compile_kernel(kernel, capture_ir_text=False) + launcher = KernelLauncher( + hsaco=artifact.hsaco, + kernel_name=artifact.kernel_name, + signature=_attn_signature( + dtype_str, include_bt_stride=True, include_qq_bias_stride=True + ), + cache_key=("d256_fast_direct", spec.kernel_name()), + ) + if os.environ.get("PARITY_PROBE_OCC") == "1": + _probe_print_occupancy(artifact.hsaco, spec.num_warps) + vals = _attn_values( + problem=problem, + q=q, + k=data["key_cache"], + v=data["value_cache"], + out=output, + cu_seqlens_q=data["cu_q"], + seqused_k=data["kv_lens"], + softmax_scale=data["scale"], + block_table=data["block_tables"], + softcap=float(s.softcap), + sinks=data["sinks"], + bt_stride=int(data["block_tables"].stride(0)), + include_bt_stride=True, + alibi_slopes=data["alibi_slopes"], + qq_bias=qq_bias, + qq_bias_stride_0=qq_bias_stride_0, + include_qq_bias_stride=True, + ) + block_q = spec.block_q + total_num_q_blocks = q.shape[0] // block_q + len(s.seq_lens) + cfg = LaunchConfig( + grid=(int(s.num_kv_heads), int(total_num_q_blocks), 1), + block=(64 * spec.num_warps, 1, 1), + stream=hip_stream, + ) + print(f" [2d-direct] {spec.kernel_name()}") + + def call_once(): + launcher(vals, config=cfg) + + ms = _time_lane_ms( + call_once, warmup=warmup, attempts=attempts, stream=hip_stream + ) + return output, ms + ok, reason = supports_tiled_2d( head_size=s.head_size, block_size=s.block_size, @@ -1346,6 +1523,142 @@ def call_once(): return output, ms +def _probe_print_occupancy(hsaco: bytes, num_warps: int) -> None: + """Print VGPR/AGPR/SGPR/LDS + coarse waves-per-CU for a compiled kernel. + + Reuses the static occupancy probe (llvm-readelf on the HSACO notes); no + kernel launch, so it is safe to call during the build step. + """ + try: + import sys as _sys + + from rocke.assets import dsl_docs_dir + + probe_dir = ( + dsl_docs_dir() / "optimization" / "utilities" / "tools" / "dsl_probes" + ) + if str(probe_dir) not in _sys.path: + _sys.path.insert(0, str(probe_dir)) + from probe_occupancy import ( # type: ignore + ARCH_GFX950, + estimate_occupancy, + parse_hsaco_notes, + ) + + notes = parse_hsaco_notes(hsaco) + occ = estimate_occupancy(notes=notes, waves_per_wg=num_warps, arch=ARCH_GFX950) + print( + f" [occupancy] vgpr={notes.get('vgpr_count')} " + f"agpr={notes.get('agpr_count')} sgpr={notes.get('sgpr_count')} " + f"spill={notes.get('vgpr_spill_count')} " + f"lds={notes.get('lds_size')}B " + f"waves/CU={occ['waves_per_cu']} wg/CU={occ['wgs_per_cu']} " + f"limit={occ['limited_by']}" + ) + except Exception as e: # noqa: BLE001 + print(f" [occupancy] unavailable: {type(e).__name__}: {e}") + + +def _run_torch_flash(s: Scenario, data, *, warmup: int, attempts: int): + """Time torch's flash-attention (F.scaled_dot_product_attention) on the + same shape, as the "vs torch-flash" reference in PR #9233. + + Reconstructs the dense per-sequence causal Q/K/V from the paged KV cache + (so the numbers match the reference paged attention), then times F.sdpa + with the flash backend. Returns (out[total_q, H, D], ms). + """ + import torch + import torch.nn.functional as F + from torch.nn.attention import sdpa_kernel, SDPBackend + + q = data["query"] + key_cache = data["key_cache"] + value_cache = data["value_cache"] + block_tables = data["block_tables"] + query_lens = data["query_lens"] + kv_lens_list = data["kv_lens_list"] + H = s.num_query_heads + Hkv = s.num_kv_heads + D = s.head_size + scale = data["scale"] + dev = q.device + nqpkv = H // Hkv + + # Gather dense per-seq tensors once (outside the timed region). + seqs = [] + start = 0 + for i, (ql, kvl) in enumerate(zip(query_lens, kv_lens_list)): + if ql <= 0: + continue + q_i = q[start : start + ql] # [ql, H, D] + start += ql + nblk = (kvl + s.block_size - 1) // s.block_size + blk = block_tables[i, :nblk] + k_i = key_cache[blk].reshape(-1, Hkv, D)[:kvl] # [kvl, Hkv, D] + v_i = value_cache[blk].reshape(-1, Hkv, D)[:kvl] + # [1, H, L, D] + qh = q_i.permute(1, 0, 2).unsqueeze(0) + kh = k_i.permute(1, 0, 2).unsqueeze(0) + vh = v_i.permute(1, 0, 2).unsqueeze(0) + if nqpkv > 1: + kh = kh.repeat_interleave(nqpkv, dim=1) + vh = vh.repeat_interleave(nqpkv, dim=1) + seqs.append((qh, kh, vh, ql, kvl)) + + # Pick a working flash-first backend once. + backend = None + for be, name in ( + (SDPBackend.FLASH_ATTENTION, "flash"), + (SDPBackend.EFFICIENT_ATTENTION, "efficient"), + (SDPBackend.MATH, "math"), + ): + try: + with sdpa_kernel(be): + _ = F.scaled_dot_product_attention( + seqs[0][0][:, :, :64], + seqs[0][1][:, :, :64], + seqs[0][2][:, :, :64], + is_causal=True, + scale=scale, + ) + backend = (be, name) + break + except Exception: # noqa: BLE001 + continue + if backend is None: + raise NotImplementedError("no working torch sdpa backend") + _LAST_TORCH_FLASH["backend"] = backend[1] + + outs = [None] * len(seqs) + + def call_once(): + with sdpa_kernel(backend[0]): + for j, (qh, kh, vh, ql, kvl) in enumerate(seqs): + if ql == kvl: + o = F.scaled_dot_product_attention( + qh, kh, vh, is_causal=True, scale=scale + ) + else: + # query is a suffix of the kv stream (aiter convention) + mask = torch.ones(ql, kvl, device=dev, dtype=torch.bool) + mask = torch.triu(mask, diagonal=kvl - ql + 1).logical_not() + o = F.scaled_dot_product_attention( + qh, kh, vh, attn_mask=mask, scale=scale + ) + outs[j] = o + + hip_stream = _bench_stream_handle() + ms = _time_lane_ms(call_once, warmup=warmup, attempts=attempts, stream=hip_stream) + + # Assemble [total_q, H, D] in scenario order for the correctness compare. + flat = [o.squeeze(0).permute(1, 0, 2) for o in outs] # each [L, H, D] + out = torch.cat(flat, dim=0).to(q.dtype) + return out, ms + + +_LAST_TORCH_FLASH = {"backend": None} + + def run_unified(backend: str, s: Scenario, data, warmup: int = 3, attempts: int = 10): """Backwards-compatible wrapper: backend in {triton, rocke}; path auto.""" if backend == "triton": @@ -1460,6 +1773,12 @@ def main() -> int: parser.add_argument("--warmup", type=int, default=3) parser.add_argument("--report", type=Path, default=None) parser.add_argument("--skip-ck", action="store_true") + parser.add_argument( + "--torch-flash", + action="store_true", + help="also time torch F.sdpa (flash) and report the CK-auto vs " + "torch-flash ratio (PR #9233 'vs torch-flash' column)", + ) parser.add_argument( "--skip-triton", action="store_true", @@ -1494,6 +1813,36 @@ def main() -> int: return 1 print("device:", torch.cuda.get_device_name(0)) + # PR #9233 reproduction lever: force the softmax<->MFMA interleave OFF on the + # production ("auto") dispatch so the D256 gfx950 fast path degrades to the + # "32x32 base" it is measured against. Run the harness twice (with and + # without this env) and take the ck-auto latency ratio to recover the + # "+N% over 32x32 base" number. Only the interleave codegen flags are + # dropped; the gated geometry (num_warps/tile_size/block_m_per_warp) is + # unchanged, so grid/launch-meta still agree. + if os.environ.get("PARITY_NO_INTERLEAVE") == "1": + from dataclasses import replace as _replace + import kernels.common.attention_unified as _au + + _orig_tiled_spec = _au._tiled_spec_from_problem + + def _tiled_spec_no_interleave(problem): + spec = _orig_tiled_spec(problem) + if getattr(spec, "use_softmax_mfma_interleave", False): + spec = _replace( + spec, + use_softmax_mfma_interleave=False, + softmax_interleave_mode=0, + softmax_interleave_groups=1, + ) + return spec + + _au._tiled_spec_from_problem = _tiled_spec_no_interleave + print( + "[parity] PARITY_NO_INTERLEAVE=1: softmax<->MFMA interleave DISABLED " + "(auto lane -> 32x32 base)" + ) + if args.set == "default": scenarios = default_scenarios() elif args.set == "creative": @@ -1604,6 +1953,38 @@ def main() -> int: row["ck_auto_status"] = err_ck _row_print("ck-auto", ck_auto, ref_out, None) _isolate_benchmark_lane() + + if args.torch_flash: + tf, err_tf = _safe_run( + lambda: _run_torch_flash( + s, + data, + warmup=args.warmup, + attempts=args.attempts, + ) + ) + if tf: + tf_out, tf_ms = tf + row["torch_flash_ms"] = tf_ms + row["torch_flash_backend"] = _LAST_TORCH_FLASH["backend"] + tf_diff = compare(ref_out, tf_out) + row["torch_flash_vs_ref"] = tf_diff["max_abs"] + print( + f" {'torch-flash':14s}: {tf_ms * 1000:9.2f} us " + f"max_abs={tf_diff['max_abs']:.5f} " + f"backend={_LAST_TORCH_FLASH['backend']}" + ) + if row.get("ck_auto_ms"): + ratio = tf_ms / row["ck_auto_ms"] + row["ck_auto_vs_torch_flash"] = ratio + print( + f" {'ck-auto vs torch-flash':22s}: {ratio:.3f}x " + f"(ck {'faster' if ratio > 1 else 'slower'})" + ) + elif err_tf: + row["torch_flash_status"] = err_tf + print(f" torch-flash FAILED: {err_tf}") + _isolate_benchmark_lane() continue # Force-path: Triton on `path`, CK DSL on `path`. diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index e4d227b1604f..706550a7621b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -2116,6 +2116,13 @@ def _tiled_spec_from_problem( softmax_interleave_groups=4, use_fast_paged_kv_desc=False, use_mfma32_skip_legacy_qreg=False, + # Slab-granularity K_lds pad (16 halves): breaks the row-aliased + # bank conflict on the QK K read (SQ_LDS_BANK_CONFLICT rate 897->660) + # for a measured ~25% latency win at Sq8192. Occupancy-neutral (LDS + # 65 KB <= 80 KB budget keeps 2 WG/CU; VGPR unchanged). DMA-safe: + # padding is between the 2-row async-DMA slabs, write+read consistent. + use_kq_lds_pad=True, + kq_lds_pad_halves=16, ) return _resolve_lds_budget(_spec) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index cb724841ce94..d6fe1b11243c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -342,6 +342,26 @@ class UnifiedAttention2DTiledSpec: # K-axis mostly lives in registers. This flag tracks that orientation # independently while it is brought up. use_transposed_qk_32x32: bool = False + # EXPERIMENTAL / measurement-only: apply a vanilla XOR swizzle to the QK + # K_lds read column offset (``col ^= (row & 7) << 3``) to spread the + # row-aliased banks. This is a *read-only* probe: the async-DMA K write is + # contiguous and is NOT swizzled to match, so numeric results are INTENTIONALLY + # WRONG when this is on -- its sole purpose is to measure whether + # ``SQ_LDS_BANK_CONFLICT`` drops for the swizzled access pattern. Off by + # default; never enable in production. + use_kq_xor_swizzle: bool = False + # EXPERIMENTAL / measurement-only: pad the K_lds row stride by 8 halves + # (16 B) to break the row-aliased bank conflict on the QK K read. The read + # picks up the padded stride via smem_load_vN; the raw-address async-DMA K + # write is NOT pad-aware, so numeric results are INTENTIONALLY WRONG when on + # -- purpose is to measure the SQ_LDS_BANK_CONFLICT delta AND confirm the + # padded LDS still fits the 2-WG/CU budget (<= 80 KB/WG). The DMA-safe + # production form is a slab-granularity pad (pad every 2-row DMA slab). + use_kq_lds_pad: bool = False + # Pad width in halves for the slab-granularity K_lds pad (must be a multiple + # of 8 to keep ds_read/ds_write b128 16-byte alignment). Only used when + # ``use_kq_lds_pad`` is set. + kq_lds_pad_halves: int = 8 # Transposed 32x32 softmax has one online-softmax state per query lane, # not one per output-dimension accumulator register. Keep a single m/l # loop-carried state and broadcast alpha across the 16 output regs. @@ -1098,6 +1118,23 @@ def build_unified_attention_2d_tiled( FP8_NATIVE_QK = False REGISTER_PV = spec.use_register_pv TRANSPOSED_QK_32X32 = spec.use_transposed_qk_32x32 + KQ_XOR_SWIZZLE = spec.use_kq_xor_swizzle + # Slab-granularity K_lds pad. The async DMA writes one WAVE_BYTES (=1024 B = + # SLAB_ROWS*HD*2) contiguous slab per wave per call, so padding BETWEEN slabs + # keeps the write contiguous while breaking the row-aliased bank conflict on + # the QK read. Correct (write+read consistent) -> numerics preserved. + KQ_LDS_PAD = int(spec.kq_lds_pad_halves) if spec.use_kq_lds_pad else 0 + # SLAB_ROWS = rows filled by one wave's 1024-B (WAVE=64 lanes x 16 B) slab. + KQ_SLAB_ROWS = (512 // HD) if KQ_LDS_PAD else 0 + if KQ_LDS_PAD and ( + 512 % HD != 0 + or KQ_SLAB_ROWS == 0 + or T % KQ_SLAB_ROWS != 0 + or KQ_LDS_PAD % 8 != 0 + ): + # Unsupported geometry for the slab pad -> disable (byte-identical). + KQ_LDS_PAD = 0 + KQ_SLAB_ROWS = 0 KV_BYTES = 1 if KV_FP8 else 2 kv_io_dtype = FP8E4M3 if KV_FP8 else dtype @@ -1358,7 +1395,21 @@ def build_unified_attention_2d_tiled( # with the ring (>2) and V-double-buffer paths in __post_init__. K_SINGLE_BUFFER = bool(spec.use_k_single_buffer) K_BUFS = 1 if K_SINGLE_BUFFER else (KV_RING_DEPTH if KV_RING_DEPTH > 2 else 2) - K_lds = b.smem_alloc(K_LDS_DTYPE, [K_BUFS, T, HD], name_hint="Klds") + if KQ_LDS_PAD and K_BUFS != 1: + # Slab pad only wired for the single-buffer K path; disable otherwise. + KQ_LDS_PAD = 0 + KQ_SLAB_ROWS = 0 + if KQ_LDS_PAD: + # Slab-granularity layout: [K_BUFS, T/SLAB_ROWS, SLAB_ROWS*HD + PAD]. + # Each DMA wave-slab (SLAB_ROWS contiguous rows) lands in one padded + # slab; the read remaps (row,col) -> (slab, row%SLAB_ROWS*HD + col). + _KQ_NSLAB = T // KQ_SLAB_ROWS + _KQ_SLAB_COLS = KQ_SLAB_ROWS * HD + KQ_LDS_PAD + K_lds = b.smem_alloc( + K_LDS_DTYPE, [K_BUFS, _KQ_NSLAB, _KQ_SLAB_COLS], name_hint="Klds" + ) + else: + K_lds = b.smem_alloc(K_LDS_DTYPE, [K_BUFS, T, HD], name_hint="Klds") # V double-buffer: 2 slots let V[i+1] prefetch into the alternate slot # during iter i (reusing ``cur_buf``), giving V the same 1-deep prefetch # K has. Costs only LDS (async DMA bypasses VGPR), occupancy-neutral. @@ -1863,6 +1914,26 @@ def _acc_idx(n: int, atom: int) -> int: wave_lds_offset_i32 = b.to_sgpr_u32(b.mul(wave_id, b.const_i32(WAVE_BYTES))) wave_lds_offset_i64 = b.zext(wave_lds_offset_i32, I64) + # K-specific slab-padded strides. One wave-slab == WAVE_BYTES contiguous B; + # inserting KQ_LDS_PAD*2 bytes between slabs shifts each slab's banks while + # leaving the intra-slab write contiguous (DMA-safe). V is untouched. + if KQ_LDS_PAD: + _KQ_PAD_BYTES = KQ_LDS_PAD * 2 + WAVE_BYTES_K = WAVE_BYTES + _KQ_PAD_BYTES + bytes_per_call_K = NUM_WARPS * WAVE_BYTES_K + bytes_per_buf_K = (T // KQ_SLAB_ROWS) * WAVE_BYTES_K + if NUM_WARPS == 1: + wave_lds_offset_K_i64 = b.const_i64(0) + else: + wave_lds_offset_K_i64 = b.zext( + b.to_sgpr_u32(b.mul(wave_id, b.const_i32(WAVE_BYTES_K))), I64 + ) + else: + WAVE_BYTES_K = WAVE_BYTES + bytes_per_call_K = bytes_per_call + bytes_per_buf_K = bytes_per_buf + wave_lds_offset_K_i64 = wave_lds_offset_i64 + # ---- Paged KV byte descriptor (full transform DAG) ---- # The paged-KV cache is laid out ``[num_blocks, BS, NUM_KV, HD]`` with # *byte* strides. The kernel addresses it via a chain of coordinate @@ -2023,10 +2094,10 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: offset (lanes 64..127 have `tid*8 / HD` advanced by T/NUM_WARPS), the cooperative load fills the full `[T, HD]` slab correctly. """ - buf_off_i32 = b.mul(buf_idx, b.const_i32(bytes_per_buf)) + buf_off_i32 = b.mul(buf_idx, b.const_i32(bytes_per_buf_K)) buf_off_i64 = b.zext(buf_off_i32, I64) K_buf_base = b.smem_ptr_add(K_lds_addr, buf_off_i64) - K_wave_base = b.smem_ptr_add(K_buf_base, wave_lds_offset_i64) + K_wave_base = b.smem_ptr_add(K_buf_base, wave_lds_offset_K_i64) if FAST_PAGED_KV_DESC: fast_block0, fast_block1 = _fast_paged_kv_blocks(kv_tile_idx) for call in range(kv_calls_per_tile): @@ -2059,7 +2130,7 @@ def _issue_k_load_runtime(kv_tile_idx: Value, buf_idx: Value) -> None: linear_half=linear_half, kv_head=kv_head_idx, ) - k_dst = b.smem_ptr_add(K_wave_base, b.const_i64(call * bytes_per_call)) + k_dst = b.smem_ptr_add(K_wave_base, b.const_i64(call * bytes_per_call_K)) # CACHE_STREAM (SLC): one-shot streaming load, never re-read # within this kernel. Documented in # ``dsl_docs/primitives/intrinsics_and_primitives.md`` as the @@ -2625,6 +2696,24 @@ def _read_k8_mfma_operand(buf_idx: Value, k_row: Value, k_off: Value) -> Value: Otherwise read bf16 directly. Shared by the 16x16 and 32x32 QK paths so both get the fp8 LDS-footprint win. """ + if KQ_XOR_SWIZZLE: + # Measurement-only vanilla XOR swizzle: permute the read column by + # the low 3 bits of the row so 8 consecutive rows land on distinct + # banks (col ^= (row & 7) << 3, i.e. (row%8)*8 halves; stays 8-half + # aligned + in [0, HD)). Read-only -> numerically wrong; used solely + # to measure the SQ_LDS_BANK_CONFLICT delta. + k_off = b.xor(k_off, b.mul(b.mod(k_row, b.const_i32(8)), b.const_i32(8))) + if KQ_LDS_PAD and not K_FP8_MFMA: + # Slab-granularity padded layout: K_lds is + # [K_BUFS, T/SLAB_ROWS, SLAB_ROWS*HD + PAD]. Map the logical + # (row, col) to (slab, row_in_slab*HD + col) so the read hits the + # same padded slot the DMA wrote (write+read consistent). + slab = b.div(k_row, b.const_i32(KQ_SLAB_ROWS)) + col_in_slab = b.add( + b.mul(b.mod(k_row, b.const_i32(KQ_SLAB_ROWS)), b.const_i32(HD)), + k_off, + ) + return b.smem_load_vN(K_lds, buf_idx, slab, col_in_slab, dtype=dtype, n=8) if not K_FP8_MFMA: return b.smem_load_vN(K_lds, buf_idx, k_row, k_off, dtype=dtype, n=8) if FP8_NATIVE_QK: From 35dd1006331239814850723d647db6783bd073fd Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Mon, 13 Jul 2026 11:25:34 -0400 Subject: [PATCH 09/23] refactor(rocke): author D256 gfx950 winning spec in the builder (#9336) Move the D256 gfx950 bf16-prefill codegen constellation from a post-build replace() override in kernels.common _tiled_spec_from_problem into the builder (builders.common.attention_spec_builder), per review: the winning spec is created in the builder, not baked into the dispatch layer. The _d256_gfx950_fast cache-key discriminator is retained so the key stays faithful to the built kernel. The predicate is called via the module (_kau._d256_gfx950_fast) so tests patching attention_unified._d256_gfx950_fast still steer the builder branch. Verified byte-identical: built spec + _tiled_cache_key unchanged for the D256 fast path and every fallback/non-D256 shape; fast vs feature-gated fallback keys stay distinct; byte-identity gate GREEN (10/10 attention families, Python==C++); attention routing + lds-budget unit tests pass. ## Motivation ## Technical Details ## Test Plan ## Test Result ## Submission Checklist - [ ] Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests. --- .../builders/common/attention_spec_builder.py | 43 ++++++++++++++++++- .../kernels/common/attention_unified.py | 40 ----------------- 2 files changed, 41 insertions(+), 42 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py index f2347cf165a6..440652c0ea8b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py @@ -14,7 +14,7 @@ """ from __future__ import annotations -from dataclasses import fields +from dataclasses import fields, replace from kernels.common.attention_unified import ( UnifiedAttentionProblem, @@ -59,6 +59,10 @@ _tiled_2d_impl, _tiled_3d_impl, ) +# Imported as a module (not a bound symbol) so tests that +# ``mock.patch.object(attention_unified, "_d256_gfx950_fast", ...)`` still steer +# the builder's fast-route branch below (a bound import would freeze the ref). +from kernels.common import attention_unified as _kau def _tiled_spec_from_problem( @@ -213,7 +217,7 @@ def _tiled_spec_from_problem( # no-fp8 preconditions so it can never fire on an incompatible spec. if "use_k_single_buffer" in _spec_field_names and _enable_k_single_buffer(problem): _gfx950_schedule_fields["use_k_single_buffer"] = True - return UnifiedAttention2DTiledSpec( + _spec = UnifiedAttention2DTiledSpec( head_size=problem.head_size, block_size=problem.block_size, num_query_heads=problem.num_query_heads, @@ -276,6 +280,41 @@ def _tiled_spec_from_problem( # below -- gfx942's spec class does not declare it.) **_gfx950_schedule_fields, ) + if _kau._d256_gfx950_fast(problem): + # D256 gfx950 bf16 prefill fast route. Authored here in the builder + # (was a post-build override in kernels.common ``_tiled_spec_from_problem``, + # so the winning spec is created in the builder, not + # baked into the dispatch layer). Geometry (num_warps / tile_size / + # block_m_per_warp) already comes from the gated selectors above; this + # pins the 32x32 transposed + FA3 softmax<->MFMA-interleave codegen + # constellation. The cohort is discriminated in ``_tiled_cache_key`` by + # ``_d256_gfx950_fast`` so the key stays faithful to the built kernel. + _spec = replace( + _spec, + use_mfma_32x32=True, + use_transposed_qk_32x32=True, + use_q_direct_reg=True, + use_transposed_half_local_pv=True, + use_transposed_scalar_state=True, + use_transposed_mask_once=True, + use_transposed_mask_limit=True, + use_mask_phase_split=True, + use_register_pv=False, + use_k_single_buffer=True, + use_v_double_buffer=False, + use_early_v_schedule=False, + use_sched_barrier=False, + use_softmax_mfma_interleave=True, + softmax_interleave_mode=2, + softmax_interleave_groups=4, + use_fast_paged_kv_desc=False, + use_mfma32_skip_legacy_qreg=False, + # Slab-granularity K_lds pad (16 halves): breaks the row-aliased + # bank conflict on the QK K read for a ~25% Sq8192 latency win. + use_kq_lds_pad=True, + kq_lds_pad_halves=16, + ) + return _spec def _tiled_3d_spec_from_problem( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index af8d3f86056d..6c5dd7476f66 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -2090,46 +2090,6 @@ def _tiled_spec_from_problem( ) _spec = _impl(problem) - # D256 gfx950 bf16 prefill fast route. The default builder routes - # head_size==256 to the 16x16 register_pv body (~0.24x torch.compile). The - # 32x32 transposed constellation is LEGAL for D256 (use_mfma_32x32 needs only - # head_size%16==0) and ~4.5-5.9x faster; the FA3-style softmax<->MFMA - # interleave (mode2 g4) adds a further +9-10% by hiding the exposed v_exp - # behind the QK/PV MFMA at the kernel's low (4 waves/CU) occupancy. Geometry - # (num_warps / tile_size / block_m_per_warp) is set by the gated selectors - # above so the cache key + launch meta agree; here we set the codegen flags. - # The cohort is uniquely keyed by head_size==256, so the (unkeyed) interleave - # + codegen overrides cannot collide with any other shape. _resolve_lds_budget - # no-ops here (register_pv=False, fits the 160 KB cap at 64 KB). - if _d256_gfx950_fast(problem): - _spec = replace( - _spec, - use_mfma_32x32=True, - use_transposed_qk_32x32=True, - use_q_direct_reg=True, - use_transposed_half_local_pv=True, - use_transposed_scalar_state=True, - use_transposed_mask_once=True, - use_transposed_mask_limit=True, - use_mask_phase_split=True, - use_register_pv=False, - use_k_single_buffer=True, - use_v_double_buffer=False, - use_early_v_schedule=False, - use_sched_barrier=False, - use_softmax_mfma_interleave=True, - softmax_interleave_mode=2, - softmax_interleave_groups=4, - use_fast_paged_kv_desc=False, - use_mfma32_skip_legacy_qreg=False, - # Slab-granularity K_lds pad (16 halves): breaks the row-aliased - # bank conflict on the QK K read (SQ_LDS_BANK_CONFLICT rate 897->660) - # for a measured ~25% latency win at Sq8192. Occupancy-neutral (LDS - # 65 KB <= 80 KB budget keeps 2 WG/CU; VGPR unchanged). DMA-safe: - # padding is between the 2-row async-DMA slabs, write+read consistent. - use_kq_lds_pad=True, - kq_lds_pad_halves=16, - ) return _resolve_lds_budget(_spec) From 6d76dc860d4a1294be4602940df35f1baf03eda4 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Mon, 13 Jul 2026 13:53:45 -0400 Subject: [PATCH 10/23] test(rocke): wire D256 gfx950 prefill into parity + benchmark harnesses (#9343) JIRA ID : AICK-1495 ## What Hooks the D256 gfx950 bf16 prefill fast path (`_d256_gfx950_fast`: 32x32 transposed + FA3-style softmax<->MFMA interleave) into the two repeatable harnesses, per review request to wire new kernels into the dispatcher **and** the benchmark harnesses. The dispatcher hookup already shipped in #9233; this closes the harness side. - **Parity harness** (`builders/gfx950/attention/prefill/parity_unified_attention.py`): adds `prefill_d256_gqa8_b16_sq4096` / `prefill_d256_gqa8_b16_sq8192` to `default_scenarios()` (Qwen3-Next-80B-A3B cohort: GQA 16/2, head_size 256, bf16 - the two measured points). The pre-existing only head_size==256 scenario was decode-class and does not trip the prefill fast-path gate. - **Benchmark shapes** (`benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json`): committed aiter-trace-schema shape set (Sq 2048/4096/8192) for the shape-driven `benchmark_prefill2d_live.py` sweep. No committed head_size==256 prefill shapes existed before. ## Verification (gfx950 / MI350X, kreb rocm7.13, real GPU run) - **Parity: PASS.** Both scenarios route to the D256 fast path (`dispatch_path=2d`, kernel `rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_..._mphsplit_smxil2g4_qdreg`) and match the fp32 `ref_paged_attn` reference at `max_abs = 1.5625e-2` (= the documented parity value; within bf16 `atol=1e-1`). Identical `max_abs` across the auto / 2d / 3d paths. - **vs AOTriton (the ticket baseline): CK wins.** `--torch-flash` (torch SDPA FLASH = AOTriton on ROCm): Sq4096 = **1.59x**, Sq8192 = **1.03x** (ck-auto 431 / 1429 us vs torch-flash 685 / 1470 us). Consistent with "beats AOTriton at every ticket shape". - **Benchmark sweep: runs all 3 shapes, `correct=3 incorrect=0`,** and the `prod` variant is confirmed to dispatch the D256 fast kernel (`dispatch_path=2d`, `_d256_gfx950_fast=True`, same kernel name as above). - Note on the sweep's speedup column: `benchmark_prefill2d_live.py` on this branch compares against **AITER-Triton** (`aiter.ops.triton.unified_attention`), not AOTriton, so its "speedup vs Triton" is a different comparison than the AOTriton number above. It also times CK via the provider path (`run_unified_attention_torch`) while timing the Triton lane as a raw kernel call; the same CK kernel measures ~431 us direct (parity) vs ~1221 us via the provider path in the sweep - a measurement-path difference (mechanism not yet isolated). The AOTriton comparison above (direct GPU-time) is the representative metric. (An AOTriton baseline lane for this harness lives on `attn-bench-framework-toggles` @ f8f05cc, not in this line.) Stacked on the #9233 branch (`users/avirgoel/ck/d256-gfx950-opt`). --- .../prefill/benchmark_prefill2d_live.py | 16 +++++++++++--- .../prefill/d256_gfx950_prefill_shapes.json | 3 +++ .../prefill/parity_unified_attention.py | 21 +++++++++++++++++++ 3 files changed, 37 insertions(+), 3 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py index 70c14294dea1..3889cf9bd4ed 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py @@ -33,7 +33,7 @@ import sys import traceback from pathlib import Path -from typing import Any +from typing import Any, Optional from rocke.assets import shape_utils_dir @@ -201,7 +201,7 @@ def _variant_flags(name: str, *, sliding_window: int, dtype: str, is_fp8: bool) class CkVariantBench: - def __init__(self, *, compile_backend: str = "llvm", num_sms: int = 256): + def __init__(self, *, compile_backend: Optional[str] = None, num_sms: int = 256): self.compile_backend = compile_backend self.num_sms = num_sms self._launchers: dict[tuple, Any] = {} @@ -433,6 +433,15 @@ def main() -> int: # cap=8192 vs 1.11x at cap=65536. ap.add_argument("--cap-blocks", type=int, default=65536) ap.add_argument("--num-sms", type=int, default=256) + ap.add_argument( + "--compile-backend", + choices=("auto", "llvm", "hipcc"), + default="auto", + help="Compile backend for the prod lane. 'auto' (default) matches " + "production dispatch (_select_2d_compile_backend); 'llvm'/'hipcc' force " + "a backend for A/B. Forcing 'llvm' mismeasures cohorts whose production " + "backend is hipcc (e.g. D256 large prefill) by ~3x.", + ) ap.add_argument("--tol", type=float, default=5e-2) ap.add_argument("--shape-utils-path", type=Path, default=DEFAULT_SHAPE_UTILS) ap.add_argument( @@ -463,7 +472,8 @@ def main() -> int: print(f"device: {torch.cuda.get_device_name(0)}") print(f"shapes: {len(shapes)} variants: {args.variants}") - bench = CkVariantBench(num_sms=args.num_sms) + _cb = None if args.compile_backend == "auto" else args.compile_backend + bench = CkVariantBench(num_sms=args.num_sms, compile_backend=_cb) results = [] for i, shape in enumerate(shapes, 1): sw = shape.window_size[0] + 1 if shape.window_size[0] >= 0 else 0 diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json new file mode 100644 index 000000000000..5a30aaabe1a3 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/d256_gfx950_prefill_shapes.json @@ -0,0 +1,3 @@ +{"call_idx": 0, "kind": "2d", "ALL_DECODE": false, "num_seqs": 1, "num_query_heads": 16, "num_kv_heads": 2, "num_queries_per_kv": 8, "head_size": 256, "block_size": 16, "max_seqlen_q": 2048, "max_seqlen_k": 2048, "softmax_scale": 0.0625, "softcap": 0.0, "has_sinks": false, "has_alibi": false, "has_output_scale": false, "window_size": [-1, -1], "q_dtype": "torch.bfloat16", "k_dtype": "torch.bfloat16", "v_dtype": "torch.bfloat16", "out_dtype": "torch.bfloat16", "q_shape": [2048, 16, 256], "k_shape": [128, 16, 2, 256], "v_shape": [128, 16, 2, 256], "out_shape": [2048, 16, 256], "block_table_shape": [1, 128], "cu_seqlens_q_shape": [2], "seqused_k_shape": [1]} +{"call_idx": 1, "kind": "2d", "ALL_DECODE": false, "num_seqs": 1, "num_query_heads": 16, "num_kv_heads": 2, "num_queries_per_kv": 8, "head_size": 256, "block_size": 16, "max_seqlen_q": 4096, "max_seqlen_k": 4096, "softmax_scale": 0.0625, "softcap": 0.0, "has_sinks": false, "has_alibi": false, "has_output_scale": false, "window_size": [-1, -1], "q_dtype": "torch.bfloat16", "k_dtype": "torch.bfloat16", "v_dtype": "torch.bfloat16", "out_dtype": "torch.bfloat16", "q_shape": [4096, 16, 256], "k_shape": [256, 16, 2, 256], "v_shape": [256, 16, 2, 256], "out_shape": [4096, 16, 256], "block_table_shape": [1, 256], "cu_seqlens_q_shape": [2], "seqused_k_shape": [1]} +{"call_idx": 2, "kind": "2d", "ALL_DECODE": false, "num_seqs": 1, "num_query_heads": 16, "num_kv_heads": 2, "num_queries_per_kv": 8, "head_size": 256, "block_size": 16, "max_seqlen_q": 8192, "max_seqlen_k": 8192, "softmax_scale": 0.0625, "softcap": 0.0, "has_sinks": false, "has_alibi": false, "has_output_scale": false, "window_size": [-1, -1], "q_dtype": "torch.bfloat16", "k_dtype": "torch.bfloat16", "v_dtype": "torch.bfloat16", "out_dtype": "torch.bfloat16", "q_shape": [8192, 16, 256], "k_shape": [512, 16, 2, 256], "v_shape": [512, 16, 2, 256], "out_shape": [8192, 16, 256], "block_table_shape": [1, 512], "cu_seqlens_q_shape": [2], "seqused_k_shape": [1]} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index 81cb7306cb5f..b89f30825d4b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -264,6 +264,27 @@ def default_scenarios() -> List[Scenario]: block_size=16, dtype=torch.float16, ), + # D256 gfx950 bf16 prefill fast path (Qwen3-Next-80B-A3B cohort, + # GQA 16/2). Exercises the _d256_gfx950_fast route (32x32 transposed + # + FA3-style softmax<->MFMA interleave) at the two measured points. + Scenario( + name="prefill_d256_gqa8_b16_sq4096", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + ), + Scenario( + name="prefill_d256_gqa8_b16_sq8192", + seq_lens=[(8192, 8192)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + dtype=torch.bfloat16, + ), Scenario( name="prefill_d128_b16", seq_lens=[(64, 64), (128, 256), (32, 256)], From b306afa63bc746f324ace872df5a8095dcc8b099 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Mon, 13 Jul 2026 14:01:38 -0400 Subject: [PATCH 11/23] style(rocke/attention): black-format attention_spec_builder.py Blank line after the import block flagged by black --check in CI pre-commit. --- .../rocke/library/builders/common/attention_spec_builder.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py index 440652c0ea8b..d01a279bf92f 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py @@ -59,6 +59,7 @@ _tiled_2d_impl, _tiled_3d_impl, ) + # Imported as a module (not a bound symbol) so tests that # ``mock.patch.object(attention_unified, "_d256_gfx950_fast", ...)`` still steer # the builder's fast-route branch below (a bound import would freeze the ref). From beaaecbc3a98ba91c1e30b474c38ce81e07f68b3 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Mon, 13 Jul 2026 14:52:06 -0400 Subject: [PATCH 12/23] fix(rocke/attention): address review comments (D256 gfx950 prefill) - attention_tiled_2d.py: rewrite the `use_kq_lds_pad` field comment. It described the experimental measurement-only pad (intentionally-wrong numerics); production wires the DMA-safe slab-granularity pad (numerics preserved, enabled by the D256 fast route). Comment now matches behavior. - parity_unified_attention.py: the `pr9233_d256_*_bs64` scenarios set num_kv_heads=16 (MHA) but the PR cohort is GQA 16/2. Set num_kv_heads=2 so they match the reported cohort (num_queries_per_kv=8). - benchmark_prefill2d_live.py: `CkVariantBench.build()` compiled every variant via LLVM-direct, so --compile-backend only affected the prod lane. Route build() through the backend (self.compile_backend or _select_2d_compile_backend); the flag now applies to all lanes and the variant lanes match production compile by default. --- .../prefill/benchmark_prefill2d_live.py | 22 ++++++++++++++----- .../prefill/parity_unified_attention.py | 4 ++-- .../kernels/gfx950/attention_tiled_2d.py | 14 ++++++------ 3 files changed, 25 insertions(+), 15 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py index 3889cf9bd4ed..e9658447fbdc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/benchmarks/gfx950/attention/prefill/benchmark_prefill2d_live.py @@ -236,7 +236,10 @@ def build(self, shape, variant: str, sliding_window: int, is_fp8: bool): build_unified_attention_2d_tiled, supports_tiled_2d, ) - from kernels.common.attention_unified import _attn_signature + from kernels.common.attention_unified import ( + _attn_signature, + _select_2d_compile_backend, + ) from rocke.runtime import KernelLauncher dtype = "bf16" if shape.q_dtype == "torch.bfloat16" else "fp16" @@ -297,7 +300,13 @@ def build(self, shape, variant: str, sliding_window: int, is_fp8: bool): key = (shape.signature, variant, spec.kernel_name(), self.compile_backend) if key not in self._launchers: kernel = build_unified_attention_2d_tiled(spec) - artifact = compile_kernel(kernel, capture_ir_text=False) + backend = self.compile_backend or _select_2d_compile_backend(problem) + if backend == "hipcc": + from rocke.helpers.compile import compile_kernel_via_hipcc + + artifact = compile_kernel_via_hipcc(kernel) + else: + artifact = compile_kernel(kernel, capture_ir_text=False) self._launchers[key] = ( KernelLauncher( hsaco=artifact.hsaco, @@ -437,10 +446,11 @@ def main() -> int: "--compile-backend", choices=("auto", "llvm", "hipcc"), default="auto", - help="Compile backend for the prod lane. 'auto' (default) matches " - "production dispatch (_select_2d_compile_backend); 'llvm'/'hipcc' force " - "a backend for A/B. Forcing 'llvm' mismeasures cohorts whose production " - "backend is hipcc (e.g. D256 large prefill) by ~3x.", + help="Compile backend for all CK DSL lanes (prod + build() variants). " + "'auto' (default) matches production dispatch " + "(_select_2d_compile_backend); 'llvm'/'hipcc' force a backend for A/B. " + "Forcing 'llvm' mismeasures cohorts whose production backend is hipcc " + "(e.g. D256 large prefill) by ~3x.", ) ap.add_argument("--tol", type=float, default=5e-2) ap.add_argument("--shape-utils-path", type=Path, default=DEFAULT_SHAPE_UTILS) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index b89f30825d4b..81187e604b25 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -882,7 +882,7 @@ def creative_scenarios() -> List[Scenario]: name="pr9233_d256_bf16_prefill_sq4096_bs64", seq_lens=[(4096, 4096)] * 64, num_query_heads=16, - num_kv_heads=16, + num_kv_heads=2, head_size=256, block_size=16, dtype=torch.bfloat16, @@ -892,7 +892,7 @@ def creative_scenarios() -> List[Scenario]: name="pr9233_d256_bf16_prefill_sq8192_bs64", seq_lens=[(8192, 8192)] * 64, num_query_heads=16, - num_kv_heads=16, + num_kv_heads=2, head_size=256, block_size=16, dtype=torch.bfloat16, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index d6fe1b11243c..c1a1e50998fc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -350,13 +350,13 @@ class UnifiedAttention2DTiledSpec: # ``SQ_LDS_BANK_CONFLICT`` drops for the swizzled access pattern. Off by # default; never enable in production. use_kq_xor_swizzle: bool = False - # EXPERIMENTAL / measurement-only: pad the K_lds row stride by 8 halves - # (16 B) to break the row-aliased bank conflict on the QK K read. The read - # picks up the padded stride via smem_load_vN; the raw-address async-DMA K - # write is NOT pad-aware, so numeric results are INTENTIONALLY WRONG when on - # -- purpose is to measure the SQ_LDS_BANK_CONFLICT delta AND confirm the - # padded LDS still fits the 2-WG/CU budget (<= 80 KB/WG). The DMA-safe - # production form is a slab-granularity pad (pad every 2-row DMA slab). + # Slab-granularity K_lds pad: pad the K_lds allocation once per 2-row + # async-DMA slab (pad width ``kq_lds_pad_halves``) to break the row-aliased + # bank conflict on the QK K read (drops ``SQ_LDS_BANK_CONFLICT``). DMA-safe + # and numerics-preserving -- the async-DMA K write and the ``smem_load_vN`` + # read use the same padded addressing. The padded LDS stays within the + # 2-WG/CU budget (<= 80 KB/WG). Enabled by the D256 gfx950 fast route + # (``kq_lds_pad_halves=16``) for a ~25% Sq8192 latency win. use_kq_lds_pad: bool = False # Pad width in halves for the slab-granularity K_lds pad (must be a multiple # of 8 to keep ds_read/ds_write b128 16-byte alignment). Only used when From 01db86a61325edb683514197bce5011e43a6c863 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Tue, 14 Jul 2026 17:39:31 -0400 Subject: [PATCH 13/23] feat(rocke): register gfx950 D256 bf16 prefill in the attention dispatcher (#9405) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit JIRA ID : AICK-1495 This PR adds an optional **`AttentionSpec.tiled_overrides`** field to the attention dispatcher (`library/dispatch/attention.py`). That is an **intentional extension of the dispatcher's deliberately-coarse parity identity** — the module docstring states the CTA geometry/knobs are *"intentionally OUT of the parity identity"* (which is `(path, head_size, block_size)`). ## Commit 2 — knob handoff (MVP) - `_d256_gfx950_spec_overrides()` = single source of truth for the 32×32-transposed + softmax↔MFMA-interleave (mode2/g4) + slab-pad constellation. Both the builder override (`attention_spec_builder`) and the dispatch candidate consume it → dispatched spec and built kernel cannot drift. **Byte-identical** to the previous inline `.replace(...)`. - `_tiled_spec_from_problem(problem, *, overrides=None)` applies dispatcher-supplied overrides on top of the heuristic base and **still runs `_resolve_lds_budget` in all cases**. The no-arg (production) path is unchanged. - `AttentionSpec.tiled_overrides` carries the pinned knobs (see the reviewer-attention note above). ## Honest ROI / scope **Behavior-neutral today.** Production `run_unified_attention_torch` still **bypasses the dispatcher** (it uses `select_path` + `_tiled_spec_from_problem`), and for the D256 cohort the dispatcher's overrides are **idempotent** with the builder's existing `_impl` override. So nothing routes or runs differently today. The value is **de-risking the eventual production reroute**: it proves the "dispatcher declares knobs → orchestrator accepts" path so the later flip (route production through dispatch + delete the `_impl` override) is a small, safe change. ## Verification - 13/13 gfx950 dispatch tests (8 routing + 5 knob-handoff: SSOT constellation, dispatched spec carries the overrides, generic candidate carries none, orchestrator accepts overrides [SSOT reproduces the spec], invalid override rejected by the spec `__post_init__`); existing dispatch suite 12/12 (no regression). rocKE_dsl --- .../builders/common/attention_spec_builder.py | 26 +-- .../rocke/library/dispatch/attention.py | 86 ++++++++++ .../kernels/common/attention_unified.py | 70 ++++++-- .../attention/test_d256_gfx950_routing.py | 158 ++++++++++++++++++ 4 files changed, 304 insertions(+), 36 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py index d01a279bf92f..902f1edcb647 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py @@ -290,31 +290,7 @@ def _tiled_spec_from_problem( # pins the 32x32 transposed + FA3 softmax<->MFMA-interleave codegen # constellation. The cohort is discriminated in ``_tiled_cache_key`` by # ``_d256_gfx950_fast`` so the key stays faithful to the built kernel. - _spec = replace( - _spec, - use_mfma_32x32=True, - use_transposed_qk_32x32=True, - use_q_direct_reg=True, - use_transposed_half_local_pv=True, - use_transposed_scalar_state=True, - use_transposed_mask_once=True, - use_transposed_mask_limit=True, - use_mask_phase_split=True, - use_register_pv=False, - use_k_single_buffer=True, - use_v_double_buffer=False, - use_early_v_schedule=False, - use_sched_barrier=False, - use_softmax_mfma_interleave=True, - softmax_interleave_mode=2, - softmax_interleave_groups=4, - use_fast_paged_kv_desc=False, - use_mfma32_skip_legacy_qreg=False, - # Slab-granularity K_lds pad (16 halves): breaks the row-aliased - # bank conflict on the QK K read for a ~25% Sq8192 latency win. - use_kq_lds_pad=True, - kq_lds_pad_halves=16, - ) + _spec = replace(_spec, **_kau._d256_gfx950_spec_overrides()) return _spec diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py index a0bd12d5a44f..f40b9780c6b4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py @@ -150,6 +150,11 @@ class AttentionSpec: num_query_heads: int num_kv_heads: int name: str = "rocke_attention_unified" + # Optional pinned codegen knobs a specialized candidate wants the builder to + # apply (sorted (key, value) pairs; empty for generic candidates). See + # ``attention_unified._d256_gfx950_spec_overrides``; the builder consumes them + # via ``_tiled_spec_from_problem(problem, overrides=...)``. + tiled_overrides: Tuple[Tuple[str, object], ...] = () def kernel_name(self) -> str: from rocke.helpers.spec import kernel_name_join @@ -301,6 +306,87 @@ def select(req: OperatorRequest) -> AttentionSpec: ATTENTION_REGISTRY.register(_make_gfx942_dense_pipe_candidate()) +def _make_gfx950_d256_candidate() -> KernelCandidate: + """Fast gfx950 bf16 head_size-256 prefill kernel — 32x32 transposed stack + with FA3-style softmax<->MFMA interleave (mode2/g4) + slab-padded K_lds. + + Registered at priority 5 so it outranks the generic unified_2d candidate + (priority 10) for the gfx950 bf16 D256 prefill cohort. The registry sorts + ascending (lower = higher precedence); gfx950-only, so it never competes + with the gfx942 dense_pipe candidate. Callers can also force this path + explicitly via algorithm="d256_gfx950". + + The cohort is the single source of truth + ``kernels.common.attention_unified._d256_gfx950_cohort`` — the same predicate + the orchestrator's ``_d256_gfx950_fast`` override uses — so dispatch selection + and the built spec cannot drift. Only the arch gate differs (request arch + here vs resolved device arch there). + """ + spec_id = "gfx950_d256" + name = "attention_gfx950_d256" + + def support(req: OperatorRequest) -> Tuple[bool, str]: + errors = _request_errors(req) + if errors: + return False, "; ".join(errors) + assert isinstance(req, AttentionRequest) + if req.arch != "gfx950": + return False, f"d256_gfx950 requires gfx950 (got {req.arch!r})" + if req.dtype != "bf16": + return False, f"d256_gfx950 is bf16-only (got {req.dtype!r})" + ok, why = _selector_matches(req, candidate) + if not ok: + return False, why + problem = _problem(req) + ok, why = supports_native_unified_attention(problem) + if not ok: + return False, why + if problem.select_path() != "2d": + return False, "problem routes to 3D, not 2D" + from kernels.common.attention_unified import _d256_gfx950_cohort + + if not _d256_gfx950_cohort(problem): + return False, "not the gfx950 bf16 D256 prefill fast-path cohort" + return True, "ok" + + def select(req: OperatorRequest) -> AttentionSpec: + ok, why = support(req) + if not ok: + raise ValueError(f"{name} does not support request: {why}") + assert isinstance(req, AttentionRequest) + problem = _problem(req) + from kernels.common.attention_unified import _d256_gfx950_spec_overrides + return AttentionSpec( + path="2d", + head_size=problem.head_size, + block_size=problem.block_size, + dtype=problem.dtype, + num_query_heads=problem.num_query_heads, + num_kv_heads=problem.num_kv_heads, + name="rocke_attention_gfx950_d256", + tiled_overrides=tuple(sorted(_d256_gfx950_spec_overrides().items())), + ) + + candidate = KernelCandidate( + name=name, + family=_FAMILY, + algorithm="d256_gfx950", + spec_id=spec_id, + abi_version=ATTENTION_ABI_VERSION, + priority=5, + supports=support, + select_spec=select, + signature=lambda _spec: (), + grid=lambda spec, req: (0, 0, 0), + block=lambda spec: (0, 0, 0), + sweep_space=lambda req: (select(req),) if support(req)[0] else (), + ) + return candidate + + +ATTENTION_REGISTRY.register(_make_gfx950_d256_candidate()) + + def attention_candidates() -> Tuple[KernelCandidate, ...]: return ATTENTION_REGISTRY.candidates() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index 6c5dd7476f66..5e54f84858b5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -585,6 +585,26 @@ def _enable_k_single_buffer(problem: UnifiedAttentionProblem) -> bool: return _enable_d128_small_tile(problem) and problem.block_size >= 32 +def _d256_gfx950_cohort(problem: "UnifiedAttentionProblem") -> bool: + """Arch-agnostic cohort for the gfx950 bf16 D256 prefill fast path. + + The arch gate is applied by callers (``_d256_gfx950_fast`` uses the resolved + device arch; the ``dispatch.attention`` candidate uses the request arch), so + this stays CPU-pure and is the single source of truth for the cohort shape. + """ + return ( + problem.head_size == 256 + and problem.dtype == "bf16" + and not problem.use_fp8 + and problem.sliding_window == 0 + and problem.softcap == 0 + and not problem.use_sinks + and not problem.use_alibi + and not problem.use_qq_bias + and problem.max_seqlen_q > 1 + ) + + def _d256_gfx950_fast(problem: "UnifiedAttentionProblem") -> bool: """Route the D256 gfx950 bf16 prefill cohort to the 32x32 transposed fast path + FA3-style softmax<->MFMA interleave. @@ -599,17 +619,41 @@ def _d256_gfx950_fast(problem: "UnifiedAttentionProblem") -> bool: Perf (MI355X, direct-launch sweep, fp32-ref max_abs=1.5625e-02): +interleave mode2 g4 gives +9.1% @ Sq4096 and +10.1% @ Sq8192 over the 32x32 base. """ - return ( - _resolve_attention_arch() == "gfx950" - and problem.head_size == 256 - and problem.dtype == "bf16" - and not problem.use_fp8 - and problem.sliding_window == 0 - and problem.softcap == 0 - and not problem.use_sinks - and not problem.use_alibi - and not problem.use_qq_bias - and problem.max_seqlen_q > 1 + return _resolve_attention_arch() == "gfx950" and _d256_gfx950_cohort(problem) + + +def _d256_gfx950_spec_overrides() -> dict: + """Single source of truth for the D256 gfx950 bf16 prefill knob constellation. + + The 32x32 transposed + FA3-style softmax<->MFMA-interleave (mode2/g4) + + slab-padded K_lds codegen pins. Consumed by BOTH the builder override in + ``builders.common.attention_spec_builder._tiled_spec_from_problem`` (the + production path) and the ``dispatch.attention`` gfx950 D256 candidate, so the + dispatched spec and the built kernel cannot drift. + """ + return dict( + use_mfma_32x32=True, + use_transposed_qk_32x32=True, + use_q_direct_reg=True, + use_transposed_half_local_pv=True, + use_transposed_scalar_state=True, + use_transposed_mask_once=True, + use_transposed_mask_limit=True, + use_mask_phase_split=True, + use_register_pv=False, + use_k_single_buffer=True, + use_v_double_buffer=False, + use_early_v_schedule=False, + use_sched_barrier=False, + use_softmax_mfma_interleave=True, + softmax_interleave_mode=2, + softmax_interleave_groups=4, + use_fast_paged_kv_desc=False, + use_mfma32_skip_legacy_qreg=False, + # Slab-granularity K_lds pad (16 halves): breaks the row-aliased bank + # conflict on the QK K read for a ~25% Sq8192 latency win. + use_kq_lds_pad=True, + kq_lds_pad_halves=16, ) @@ -2084,12 +2128,16 @@ def _consider(label, cand, why): def _tiled_spec_from_problem( problem: UnifiedAttentionProblem, + *, + overrides=None, ): from builders.common.attention_spec_builder import ( _tiled_spec_from_problem as _impl, ) _spec = _impl(problem) + if overrides is not None: + _spec = replace(_spec, **overrides) return _resolve_lds_budget(_spec) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py new file mode 100644 index 000000000000..13ca34b03fe6 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py @@ -0,0 +1,158 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Routing tests for the gfx950 bf16 D256 prefill dispatch candidate. + +CPU-only: exercises the ``attention_gfx950_d256`` candidate's selection through +``dispatch_attention`` (filter -> priority -> pick). No GPU required. +""" + +from __future__ import annotations + +import unittest + +from dispatch.attention import AttentionRequest, _problem, dispatch_attention +from kernels.common import attention_unified as au + + +def _d256(arch="gfx950", **kw): + base = dict( + batch=2, + nhead_q=32, + nhead_k=8, + seqlen_q=4096, + seqlen_k=4096, + hdim_q=256, + hdim_v=256, + arch=arch, + dtype="bf16", + mask_type=1, # causal prefill + ) + base.update(kw) + return AttentionRequest(**base) + + +def _routed_spec_id(req): + """spec_id of the winning candidate, or None if nothing supports the req.""" + try: + return dispatch_attention(req).candidate.spec_id + except ValueError: + return None + + +class TestD256Gfx950Routing(unittest.TestCase): + def test_selects_gfx950_d256_for_cohort(self): + r = dispatch_attention(_d256()) + self.assertEqual(r.candidate.spec_id, "gfx950_d256") + self.assertEqual(r.spec.path, "2d") + + def test_outranks_generic_2d(self): + # priority 5 must beat the generic unified_2d (priority 10). + self.assertEqual(dispatch_attention(_d256()).candidate.priority, 5) + + def test_force_by_algorithm(self): + r = dispatch_attention(_d256(algorithm="d256_gfx950")) + self.assertEqual(r.candidate.spec_id, "gfx950_d256") + + def test_rejects_non_gfx950(self): + self.assertNotEqual(_routed_spec_id(_d256(arch="gfx942")), "gfx950_d256") + + def test_rejects_non_bf16(self): + self.assertNotEqual(_routed_spec_id(_d256(dtype="fp16")), "gfx950_d256") + + def test_rejects_non_d256(self): + self.assertNotEqual( + _routed_spec_id(_d256(hdim_q=128, hdim_v=128)), "gfx950_d256" + ) + + def test_rejects_decode(self): + # q=1 decode over long kv routes to the 3D split-KV path, not our cohort. + self.assertNotEqual( + _routed_spec_id( + _d256(batch=1, nhead_q=16, nhead_k=16, seqlen_q=1, seqlen_k=8192) + ), + "gfx950_d256", + ) + + def test_rejects_sliding_window(self): + self.assertNotEqual( + _routed_spec_id(_d256(sliding_window=256)), "gfx950_d256" + ) + + +class TestD256Gfx950SpecOverrides(unittest.TestCase): + """MVP: the dispatcher declares the knobs; the orchestrator accepts them.""" + + EXPECTED_KEYS = { + "use_mfma_32x32", + "use_transposed_qk_32x32", + "use_q_direct_reg", + "use_transposed_half_local_pv", + "use_transposed_scalar_state", + "use_transposed_mask_once", + "use_transposed_mask_limit", + "use_mask_phase_split", + "use_register_pv", + "use_k_single_buffer", + "use_v_double_buffer", + "use_early_v_schedule", + "use_sched_barrier", + "use_softmax_mfma_interleave", + "softmax_interleave_mode", + "softmax_interleave_groups", + "use_fast_paged_kv_desc", + "use_mfma32_skip_legacy_qreg", + "use_kq_lds_pad", + "kq_lds_pad_halves", + } + + def setUp(self): + self._saved = getattr(au, "_RESOLVED_ATTENTION_ARCH", None) + au._RESOLVED_ATTENTION_ARCH = "gfx950" + + def tearDown(self): + au._RESOLVED_ATTENTION_ARCH = self._saved + + def test_ssot_constellation(self): + ov = au._d256_gfx950_spec_overrides() + self.assertEqual(set(ov), self.EXPECTED_KEYS) + self.assertIs(ov["use_softmax_mfma_interleave"], True) + self.assertEqual(ov["softmax_interleave_mode"], 2) + self.assertIs(ov["use_register_pv"], False) + + def test_dispatched_spec_carries_overrides(self): + # the candidate declares the knobs on the dispatched spec (via the SSOT). + spec = dispatch_attention(_d256()).spec + self.assertEqual( + spec.tiled_overrides, + tuple(sorted(au._d256_gfx950_spec_overrides().items())), + ) + + def test_generic_candidate_has_no_overrides(self): + # a non-cohort shape (hd128) routes to the generic candidate -> empty. + spec = dispatch_attention(_d256(hdim_q=128, hdim_v=128)).spec + self.assertEqual(spec.tiled_overrides, ()) + + def test_orchestrator_accepts_overrides(self): + problem = _problem(_d256()) + # no-arg path: _impl already applies the D256 override, and + # _resolve_lds_budget still runs (a valid spec is returned). + base = au._tiled_spec_from_problem(problem) + self.assertIs(base.use_softmax_mfma_interleave, True) + self.assertIs(base.use_register_pv, False) + # passing the SSOT overrides explicitly reproduces the same valid spec + # (proves overrides= is plumbed + applied, and the resolver still runs). + ov = au._tiled_spec_from_problem( + problem, overrides=au._d256_gfx950_spec_overrides() + ) + self.assertEqual(ov, base) + + def test_overrides_reach_the_spec(self): + # an override that violates a spec invariant is raised by the spec's + # __post_init__ -> proves overrides= is really applied, not ignored. + problem = _problem(_d256()) + with self.assertRaises(ValueError): + au._tiled_spec_from_problem(problem, overrides={"use_register_pv": True}) + + +if __name__ == "__main__": + unittest.main() From 70c4a7bacc995d726618481c35159907b3d25bf0 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 15:09:07 -0400 Subject: [PATCH 14/23] refactor(rocke/attention): hoist torch-flash + occupancy probe to builders/common; PARITY_* env->CLI Addresses two #9233 review comments on the D256 gfx950 parity harness. Comment (hoist): _run_torch_flash and _probe_print_occupancy are not D256/gfx950-specific. Move the reusable cores to builders/common so the gfx942/decode harnesses can consume one implementation instead of drifting into near-identical copies: - builders/common/torch_flash_reference.py: run_torch_flash(...), decoupled from the harness-local Scenario/data (plain tensors/params; timing injected). - builders/common/occupancy_probe.py: print_occupancy(hsaco, num_warps, arch=) with an arch arg (was hardcoded gfx950). gfx950 keeps a thin _run_torch_flash adapter and calls print_occupancy directly. gfx942 is left untouched -- it imports these if/when it needs them. Comment (PARITY_* levers): the experimental env flags are all read in one main(), so they become real argparse flags (discoverable in --help, typed): PARITY_NO_INTERLEAVE -> --no-interleave PARITY_KQ_SWIZZLE -> --kq-swizzle PARITY_KQ_PAD/_W -> --kq-pad HALVES PARITY_PROBE_OCC -> --probe-occupancy The --no-interleave lever still patches the production _tiled_spec_from_problem (the auto/production-dispatch lane builds the spec internally, so a local _replace can't reach it), but the global mutation is now restored in a main() try/finally wrapper -- no leak to later in-process callers. GPU-verified on gfx950 (MI355X), prefill_d256_gqa8_b16_sq4096: - baseline unchanged (ck-2d 394 us, max_abs 1.56e-2); - --probe-occupancy prints the hoisted occupancy line; - --no-interleave drops the interleave (kernel name loses smxil2g4; 443 us 32x32 base); - --torch-flash uses the hoisted run_torch_flash (torch-flash 626 us, backend=flash; ck-auto vs torch-flash 1.58x). --- .../builders/common/occupancy_probe.py | 46 ++++ .../builders/common/torch_flash_reference.py | 117 +++++++++ .../prefill/parity_unified_attention.py | 245 +++++++----------- 3 files changed, 261 insertions(+), 147 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/common/occupancy_probe.py create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/occupancy_probe.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/occupancy_probe.py new file mode 100644 index 000000000000..b48e8ce48617 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/occupancy_probe.py @@ -0,0 +1,46 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Shared static occupancy probe for the attention parity/bench harnesses. + +Wraps the DSL static occupancy probe (``dsl_probes/probe_occupancy.py``, +``llvm-readelf`` on the HSACO notes -- no kernel launch, safe during the build +step) so the per-arch harnesses consume one implementation instead of drifting +into near-identical copies. +""" +from __future__ import annotations + + +def print_occupancy(hsaco: bytes, num_warps: int, *, arch: str = "gfx950") -> None: + """Print VGPR/AGPR/SGPR/LDS + coarse waves-per-CU for a compiled kernel. + + Static (reads the HSACO notes); no launch, so it is safe to call during the + build step. ``arch`` selects the occupancy model (``gfx950`` / ``gfx942`` / + ``gfx90a``); unknown values fall back to gfx950. + """ + try: + import sys as _sys + + from rocke.assets import dsl_docs_dir + + probe_dir = ( + dsl_docs_dir() / "optimization" / "utilities" / "tools" / "dsl_probes" + ) + if str(probe_dir) not in _sys.path: + _sys.path.insert(0, str(probe_dir)) + import probe_occupancy as _po # type: ignore + + arch_caps = getattr(_po, f"ARCH_{arch.upper()}", _po.ARCH_GFX950) + notes = _po.parse_hsaco_notes(hsaco) + occ = _po.estimate_occupancy( + notes=notes, waves_per_wg=num_warps, arch=arch_caps + ) + print( + f" [occupancy] vgpr={notes.get('vgpr_count')} " + f"agpr={notes.get('agpr_count')} sgpr={notes.get('sgpr_count')} " + f"spill={notes.get('vgpr_spill_count')} " + f"lds={notes.get('lds_size')}B " + f"waves/CU={occ['waves_per_cu']} wg/CU={occ['wgs_per_cu']} " + f"limit={occ['limited_by']}" + ) + except Exception as e: # noqa: BLE001 + print(f" [occupancy] unavailable: {type(e).__name__}: {e}") diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py new file mode 100644 index 000000000000..36b6215e6a78 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py @@ -0,0 +1,117 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Shared torch flash-attention reference timer for the parity/bench harnesses. + +Reconstructs the dense per-sequence causal Q/K/V from a paged KV cache (so the +numbers match the reference paged attention), then times torch's +``F.scaled_dot_product_attention`` (flash backend) on it. Kept arch-agnostic and +decoupled from any harness-local ``Scenario``/``data`` shape so the gfx942/gfx950 +harnesses can consume one implementation instead of near-identical copies. +""" +from __future__ import annotations + +from typing import Callable + + +def run_torch_flash( + *, + query, + key_cache, + value_cache, + block_tables, + query_lens, + kv_lens_list, + num_query_heads: int, + num_kv_heads: int, + head_size: int, + block_size: int, + scale: float, + warmup: int, + attempts: int, + bench_stream: int, + time_lane_ms: Callable, +): + """Time torch F.sdpa (flash) on the same shape as the paged reference. + + ``time_lane_ms(call_once, *, warmup, attempts, stream)`` and ``bench_stream`` + are injected by the caller so this stays free of harness-local timing infra. + Returns ``(out[total_q, H, D], ms, backend_name)``. + """ + import torch + import torch.nn.functional as F + from torch.nn.attention import sdpa_kernel, SDPBackend + + H = num_query_heads + Hkv = num_kv_heads + D = head_size + dev = query.device + nqpkv = H // Hkv + + # Gather dense per-seq tensors once (outside the timed region). + seqs = [] + start = 0 + for i, (ql, kvl) in enumerate(zip(query_lens, kv_lens_list)): + if ql <= 0: + continue + q_i = query[start : start + ql] # [ql, H, D] + start += ql + nblk = (kvl + block_size - 1) // block_size + blk = block_tables[i, :nblk] + k_i = key_cache[blk].reshape(-1, Hkv, D)[:kvl] # [kvl, Hkv, D] + v_i = value_cache[blk].reshape(-1, Hkv, D)[:kvl] + # [1, H, L, D] + qh = q_i.permute(1, 0, 2).unsqueeze(0) + kh = k_i.permute(1, 0, 2).unsqueeze(0) + vh = v_i.permute(1, 0, 2).unsqueeze(0) + if nqpkv > 1: + kh = kh.repeat_interleave(nqpkv, dim=1) + vh = vh.repeat_interleave(nqpkv, dim=1) + seqs.append((qh, kh, vh, ql, kvl)) + + # Pick a working flash-first backend once. + backend = None + for be, name in ( + (SDPBackend.FLASH_ATTENTION, "flash"), + (SDPBackend.EFFICIENT_ATTENTION, "efficient"), + (SDPBackend.MATH, "math"), + ): + try: + with sdpa_kernel(be): + _ = F.scaled_dot_product_attention( + seqs[0][0][:, :, :64], + seqs[0][1][:, :, :64], + seqs[0][2][:, :, :64], + is_causal=True, + scale=scale, + ) + backend = (be, name) + break + except Exception: # noqa: BLE001 + continue + if backend is None: + raise NotImplementedError("no working torch sdpa backend") + + outs = [None] * len(seqs) + + def call_once(): + with sdpa_kernel(backend[0]): + for j, (qh, kh, vh, ql, kvl) in enumerate(seqs): + if ql == kvl: + o = F.scaled_dot_product_attention( + qh, kh, vh, is_causal=True, scale=scale + ) + else: + # query is a suffix of the kv stream (aiter convention) + mask = torch.ones(ql, kvl, device=dev, dtype=torch.bool) + mask = torch.triu(mask, diagonal=kvl - ql + 1).logical_not() + o = F.scaled_dot_product_attention( + qh, kh, vh, attn_mask=mask, scale=scale + ) + outs[j] = o + + ms = time_lane_ms(call_once, warmup=warmup, attempts=attempts, stream=bench_stream) + + # Assemble [total_q, H, D] in scenario order for the correctness compare. + flat = [o.squeeze(0).permute(1, 0, 2) for o in outs] # each [L, H, D] + out = torch.cat(flat, dim=0).to(query.dtype) + return out, ms, backend[1] diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index 81187e604b25..68492b410239 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -1239,7 +1239,17 @@ def call_once(): _force_triton_path("auto") -def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): +def _run_rocke( + s: Scenario, + data, + *, + path: str, + warmup: int, + attempts: int, + kq_swizzle: bool = False, + kq_pad: int = 0, + probe_occupancy: bool = False, +): """Run CK DSL `run_unified_attention_torch` with the requested path forced. Both backends share the default bench stream (torch's current @@ -1312,12 +1322,12 @@ def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): # PR #9233 direct-launch lane: for the D256 gfx950 bf16 prefill cohort, # build and launch the PRODUCTION fast spec directly (GPU-time only, no - # run_unified_attention_torch dispatcher overhead). Honors the - # PARITY_NO_INTERLEAVE monkeypatch so the same lane measures both the - # interleave fast path and the 32x32 base. + # run_unified_attention_torch dispatcher overhead). --no-interleave + # (applied in main via a scoped patch of _tiled_spec_from_problem) makes + # this lane measure both the interleave fast path and the "32x32 base". if _d256_gfx950_fast(problem): spec = _tiled_spec_from_problem(problem) - if os.environ.get("PARITY_KQ_SWIZZLE") == "1": + if kq_swizzle: from dataclasses import replace as _replace spec = _replace(spec, use_kq_xor_swizzle=True) @@ -1325,13 +1335,12 @@ def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): " [2d-direct] KQ XOR swizzle ON (read-only probe; " "numerics intentionally wrong)" ) - if os.environ.get("PARITY_KQ_PAD") == "1": + if kq_pad > 0: from dataclasses import replace as _replace - _w = int(os.environ.get("PARITY_KQ_PAD_W", "8")) - spec = _replace(spec, use_kq_lds_pad=True, kq_lds_pad_halves=_w) + spec = _replace(spec, use_kq_lds_pad=True, kq_lds_pad_halves=kq_pad) print( - f" [2d-direct] KQ slab-gran pad ON (pad={_w} halves, " + f" [2d-direct] KQ slab-gran pad ON (pad={kq_pad} halves, " "correct: numerics preserved)" ) kernel = build_unified_attention_2d_tiled(spec) @@ -1349,8 +1358,10 @@ def _run_rocke(s: Scenario, data, *, path: str, warmup: int, attempts: int): ), cache_key=("d256_fast_direct", spec.kernel_name()), ) - if os.environ.get("PARITY_PROBE_OCC") == "1": - _probe_print_occupancy(artifact.hsaco, spec.num_warps) + if probe_occupancy: + from builders.common.occupancy_probe import print_occupancy + + print_occupancy(artifact.hsaco, spec.num_warps, arch="gfx950") vals = _attn_values( problem=problem, q=q, @@ -1544,136 +1555,33 @@ def call_once(): return output, ms -def _probe_print_occupancy(hsaco: bytes, num_warps: int) -> None: - """Print VGPR/AGPR/SGPR/LDS + coarse waves-per-CU for a compiled kernel. - - Reuses the static occupancy probe (llvm-readelf on the HSACO notes); no - kernel launch, so it is safe to call during the build step. - """ - try: - import sys as _sys - - from rocke.assets import dsl_docs_dir - - probe_dir = ( - dsl_docs_dir() / "optimization" / "utilities" / "tools" / "dsl_probes" - ) - if str(probe_dir) not in _sys.path: - _sys.path.insert(0, str(probe_dir)) - from probe_occupancy import ( # type: ignore - ARCH_GFX950, - estimate_occupancy, - parse_hsaco_notes, - ) - - notes = parse_hsaco_notes(hsaco) - occ = estimate_occupancy(notes=notes, waves_per_wg=num_warps, arch=ARCH_GFX950) - print( - f" [occupancy] vgpr={notes.get('vgpr_count')} " - f"agpr={notes.get('agpr_count')} sgpr={notes.get('sgpr_count')} " - f"spill={notes.get('vgpr_spill_count')} " - f"lds={notes.get('lds_size')}B " - f"waves/CU={occ['waves_per_cu']} wg/CU={occ['wgs_per_cu']} " - f"limit={occ['limited_by']}" - ) - except Exception as e: # noqa: BLE001 - print(f" [occupancy] unavailable: {type(e).__name__}: {e}") - - def _run_torch_flash(s: Scenario, data, *, warmup: int, attempts: int): - """Time torch's flash-attention (F.scaled_dot_product_attention) on the - same shape, as the "vs torch-flash" reference in PR #9233. + """Time torch flash SDPA on the same shape (the "vs torch-flash" reference). - Reconstructs the dense per-sequence causal Q/K/V from the paged KV cache - (so the numbers match the reference paged attention), then times F.sdpa - with the flash backend. Returns (out[total_q, H, D], ms). + Thin arch-local adapter over the shared + ``builders.common.torch_flash_reference.run_torch_flash`` (kept generic so + the gfx942/decode harnesses can reuse it). Returns (out[total_q, H, D], ms). """ - import torch - import torch.nn.functional as F - from torch.nn.attention import sdpa_kernel, SDPBackend - - q = data["query"] - key_cache = data["key_cache"] - value_cache = data["value_cache"] - block_tables = data["block_tables"] - query_lens = data["query_lens"] - kv_lens_list = data["kv_lens_list"] - H = s.num_query_heads - Hkv = s.num_kv_heads - D = s.head_size - scale = data["scale"] - dev = q.device - nqpkv = H // Hkv - - # Gather dense per-seq tensors once (outside the timed region). - seqs = [] - start = 0 - for i, (ql, kvl) in enumerate(zip(query_lens, kv_lens_list)): - if ql <= 0: - continue - q_i = q[start : start + ql] # [ql, H, D] - start += ql - nblk = (kvl + s.block_size - 1) // s.block_size - blk = block_tables[i, :nblk] - k_i = key_cache[blk].reshape(-1, Hkv, D)[:kvl] # [kvl, Hkv, D] - v_i = value_cache[blk].reshape(-1, Hkv, D)[:kvl] - # [1, H, L, D] - qh = q_i.permute(1, 0, 2).unsqueeze(0) - kh = k_i.permute(1, 0, 2).unsqueeze(0) - vh = v_i.permute(1, 0, 2).unsqueeze(0) - if nqpkv > 1: - kh = kh.repeat_interleave(nqpkv, dim=1) - vh = vh.repeat_interleave(nqpkv, dim=1) - seqs.append((qh, kh, vh, ql, kvl)) - - # Pick a working flash-first backend once. - backend = None - for be, name in ( - (SDPBackend.FLASH_ATTENTION, "flash"), - (SDPBackend.EFFICIENT_ATTENTION, "efficient"), - (SDPBackend.MATH, "math"), - ): - try: - with sdpa_kernel(be): - _ = F.scaled_dot_product_attention( - seqs[0][0][:, :, :64], - seqs[0][1][:, :, :64], - seqs[0][2][:, :, :64], - is_causal=True, - scale=scale, - ) - backend = (be, name) - break - except Exception: # noqa: BLE001 - continue - if backend is None: - raise NotImplementedError("no working torch sdpa backend") - _LAST_TORCH_FLASH["backend"] = backend[1] - - outs = [None] * len(seqs) - - def call_once(): - with sdpa_kernel(backend[0]): - for j, (qh, kh, vh, ql, kvl) in enumerate(seqs): - if ql == kvl: - o = F.scaled_dot_product_attention( - qh, kh, vh, is_causal=True, scale=scale - ) - else: - # query is a suffix of the kv stream (aiter convention) - mask = torch.ones(ql, kvl, device=dev, dtype=torch.bool) - mask = torch.triu(mask, diagonal=kvl - ql + 1).logical_not() - o = F.scaled_dot_product_attention( - qh, kh, vh, attn_mask=mask, scale=scale - ) - outs[j] = o - - hip_stream = _bench_stream_handle() - ms = _time_lane_ms(call_once, warmup=warmup, attempts=attempts, stream=hip_stream) + from builders.common.torch_flash_reference import run_torch_flash - # Assemble [total_q, H, D] in scenario order for the correctness compare. - flat = [o.squeeze(0).permute(1, 0, 2) for o in outs] # each [L, H, D] - out = torch.cat(flat, dim=0).to(q.dtype) + out, ms, backend = run_torch_flash( + query=data["query"], + key_cache=data["key_cache"], + value_cache=data["value_cache"], + block_tables=data["block_tables"], + query_lens=data["query_lens"], + kv_lens_list=data["kv_lens_list"], + num_query_heads=s.num_query_heads, + num_kv_heads=s.num_kv_heads, + head_size=s.head_size, + block_size=s.block_size, + scale=data["scale"], + warmup=warmup, + attempts=attempts, + bench_stream=_bench_stream_handle(), + time_lane_ms=_time_lane_ms, + ) + _LAST_TORCH_FLASH["backend"] = backend return out, ms @@ -1788,12 +1696,51 @@ def _row_print(label: str, out_ms, ref_out, t_out): def main() -> int: + # Guarantee any --no-interleave patch of the production + # _tiled_spec_from_problem is restored, even on exceptions (no global leak). + import kernels.common.attention_unified as _au + + _orig_tiled_spec = _au._tiled_spec_from_problem + try: + return _main_impl() + finally: + _au._tiled_spec_from_problem = _orig_tiled_spec + + +def _main_impl() -> int: parser = argparse.ArgumentParser() parser.add_argument("--scenario", default=None, action="append") parser.add_argument("--attempts", type=int, default=10) parser.add_argument("--warmup", type=int, default=3) parser.add_argument("--report", type=Path, default=None) parser.add_argument("--skip-ck", action="store_true") + parser.add_argument( + "--no-interleave", + action="store_true", + help="build the D256 fast spec with softmax<->MFMA interleave OFF (the " + "'32x32 base'); run with/without to recover the +N%% interleave delta " + "(was PARITY_NO_INTERLEAVE)", + ) + parser.add_argument( + "--kq-swizzle", + action="store_true", + help="direct-launch 2D lane: KQ XOR swizzle ON (read-only probe; " + "numerics intentionally wrong; was PARITY_KQ_SWIZZLE)", + ) + parser.add_argument( + "--kq-pad", + type=int, + default=0, + metavar="HALVES", + help="direct-launch 2D lane: KQ slab-granularity LDS pad width in " + "halves (0=off, correct; was PARITY_KQ_PAD / PARITY_KQ_PAD_W)", + ) + parser.add_argument( + "--probe-occupancy", + action="store_true", + help="direct-launch 2D lane: print static VGPR/LDS/waves occupancy for " + "the built kernel (was PARITY_PROBE_OCC)", + ) parser.add_argument( "--torch-flash", action="store_true", @@ -1834,14 +1781,15 @@ def main() -> int: return 1 print("device:", torch.cuda.get_device_name(0)) - # PR #9233 reproduction lever: force the softmax<->MFMA interleave OFF on the - # production ("auto") dispatch so the D256 gfx950 fast path degrades to the - # "32x32 base" it is measured against. Run the harness twice (with and - # without this env) and take the ck-auto latency ratio to recover the - # "+N% over 32x32 base" number. Only the interleave codegen flags are - # dropped; the gated geometry (num_warps/tile_size/block_m_per_warp) is - # unchanged, so grid/launch-meta still agree. - if os.environ.get("PARITY_NO_INTERLEAVE") == "1": + # PR #9233 reproduction lever (--no-interleave): force the softmax<->MFMA + # interleave OFF on the production ("auto") dispatch AND the direct-launch + # lane so the D256 gfx950 fast path degrades to the "32x32 base" it is + # measured against. Run the harness twice (with and without) and take the + # latency ratio to recover the "+N% over 32x32 base" number. Only the + # interleave codegen flags are dropped; the gated geometry + # (num_warps/tile_size/block_m_per_warp) is unchanged. The patch is restored + # in main()'s finally (no global leak). + if args.no_interleave: from dataclasses import replace as _replace import kernels.common.attention_unified as _au @@ -1860,8 +1808,8 @@ def _tiled_spec_no_interleave(problem): _au._tiled_spec_from_problem = _tiled_spec_no_interleave print( - "[parity] PARITY_NO_INTERLEAVE=1: softmax<->MFMA interleave DISABLED " - "(auto lane -> 32x32 base)" + "[parity] --no-interleave: softmax<->MFMA interleave DISABLED " + "(auto + direct-launch -> 32x32 base)" ) if args.set == "default": @@ -2036,6 +1984,9 @@ def _tiled_spec_no_interleave(problem): path=p, warmup=args.warmup, attempts=args.attempts, + kq_swizzle=args.kq_swizzle, + kq_pad=args.kq_pad, + probe_occupancy=args.probe_occupancy, ) ) _row_print(f"ck-{path}", ck_p, ref_out, None) From 56431d877e6129dd61bc0c93a4be3f00ac10733b Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 16:10:53 -0400 Subject: [PATCH 15/23] refactor(rocke/attention): move measurement-only KQ XOR swizzle to a harness-local ProbeSpec Addresses the #9233 review: use_kq_xor_swizzle produces INTENTIONALLY WRONG numerics (a read-only bank-conflict probe) but lived on the production UnifiedAttention2DTiledSpec guarded only by a comment -- one typo from silently corrupt output. Move it off the production spec so it can never be set there by accident: - drop the field from UnifiedAttention2DTiledSpec; the kernel emit now reads it via getattr(spec, "use_kq_xor_swizzle", False) -- False for every production spec, so production/normal paths are byte-identical. - the parity harness owns a frozen _ProbeTiledSpec(UnifiedAttention2DTiledSpec) subclass; --kq-swizzle builds it via _with_kq_xor_swizzle(spec). Verified (CPU): constructing a production spec with use_kq_xor_swizzle=True now raises TypeError; the ProbeSpec carries the flag, survives dataclasses.replace (kq-pad-after-swizzle path), and kernel_name() works. --- .../prefill/parity_unified_attention.py | 27 ++++++++++++++++--- .../kernels/gfx950/attention_tiled_2d.py | 10 +------ 2 files changed, 25 insertions(+), 12 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index 68492b410239..df112fc50f46 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -1239,6 +1239,29 @@ def call_once(): _force_triton_path("auto") +def _with_kq_xor_swizzle(spec): + """Return a harness-only ProbeSpec with the measurement-only KQ XOR swizzle ON. + + The XOR swizzle produces INTENTIONALLY WRONG numerics (a read-only + bank-conflict probe), so the flag lives ONLY on this harness-local subclass -- + never on the production ``UnifiedAttention2DTiledSpec`` -- and the kernel emit + reads it via ``getattr(spec, "use_kq_xor_swizzle", False)``. + """ + import dataclasses + + from kernels import UnifiedAttention2DTiledSpec + + @dataclasses.dataclass(frozen=True) + class _ProbeTiledSpec(UnifiedAttention2DTiledSpec): + use_kq_xor_swizzle: bool = False + + fields = { + f.name: getattr(spec, f.name) + for f in dataclasses.fields(UnifiedAttention2DTiledSpec) + } + return _ProbeTiledSpec(**fields, use_kq_xor_swizzle=True) + + def _run_rocke( s: Scenario, data, @@ -1328,9 +1351,7 @@ def _run_rocke( if _d256_gfx950_fast(problem): spec = _tiled_spec_from_problem(problem) if kq_swizzle: - from dataclasses import replace as _replace - - spec = _replace(spec, use_kq_xor_swizzle=True) + spec = _with_kq_xor_swizzle(spec) print( " [2d-direct] KQ XOR swizzle ON (read-only probe; " "numerics intentionally wrong)" diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index c1a1e50998fc..449e804e09cb 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -342,14 +342,6 @@ class UnifiedAttention2DTiledSpec: # K-axis mostly lives in registers. This flag tracks that orientation # independently while it is brought up. use_transposed_qk_32x32: bool = False - # EXPERIMENTAL / measurement-only: apply a vanilla XOR swizzle to the QK - # K_lds read column offset (``col ^= (row & 7) << 3``) to spread the - # row-aliased banks. This is a *read-only* probe: the async-DMA K write is - # contiguous and is NOT swizzled to match, so numeric results are INTENTIONALLY - # WRONG when this is on -- its sole purpose is to measure whether - # ``SQ_LDS_BANK_CONFLICT`` drops for the swizzled access pattern. Off by - # default; never enable in production. - use_kq_xor_swizzle: bool = False # Slab-granularity K_lds pad: pad the K_lds allocation once per 2-row # async-DMA slab (pad width ``kq_lds_pad_halves``) to break the row-aliased # bank conflict on the QK K read (drops ``SQ_LDS_BANK_CONFLICT``). DMA-safe @@ -1118,7 +1110,7 @@ def build_unified_attention_2d_tiled( FP8_NATIVE_QK = False REGISTER_PV = spec.use_register_pv TRANSPOSED_QK_32X32 = spec.use_transposed_qk_32x32 - KQ_XOR_SWIZZLE = spec.use_kq_xor_swizzle + KQ_XOR_SWIZZLE = getattr(spec, "use_kq_xor_swizzle", False) # Slab-granularity K_lds pad. The async DMA writes one WAVE_BYTES (=1024 B = # SLAB_ROWS*HD*2) contiguous slab per wave per call, so padding BETWEEN slabs # keeps the write contiguous while breaking the row-aliased bank conflict on From 23d7e02362a2235116471fed8c4c9da54543497b Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 16:23:39 -0400 Subject: [PATCH 16/23] fix(rocke/attention): disable K_lds slab pad under fp8 (write/read asymmetry) The K_lds slab pad's async-DMA write always pads when KQ_LDS_PAD is set, but the read only remaps for the non-fp8 K path -- an fp8+pad combo would pad the write and NOT remap the read, addressing different slots -> silent numeric corruption. The pad is a bf16 bank-conflict optimization and fp8+pad is not a supported combo (the D256 gfx950 cohort that uses the pad is bf16-only), so disable the pad under fp8 in the geometry guard, keeping write and read consistent. Byte-identical for the supported bf16 D256 path. Addresses the #9233 review (K_lds pad fp8 write/read asymmetry). --- .../rocke/library/kernels/gfx950/attention_tiled_2d.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py index 449e804e09cb..e441a0b6dca4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx950/attention_tiled_2d.py @@ -1118,13 +1118,19 @@ def build_unified_attention_2d_tiled( KQ_LDS_PAD = int(spec.kq_lds_pad_halves) if spec.use_kq_lds_pad else 0 # SLAB_ROWS = rows filled by one wave's 1024-B (WAVE=64 lanes x 16 B) slab. KQ_SLAB_ROWS = (512 // HD) if KQ_LDS_PAD else 0 + # The slab pad's remap is only implemented for the non-fp8 K read (see the K + # read path below). The async-DMA write pads whenever KQ_LDS_PAD is set, but + # the fp8 K read does NOT remap -- an fp8+pad combo would address different + # slots -> silent numeric corruption. So disable the pad under fp8 (and under + # unsupported geometry) so write and read stay consistent. Byte-identical for + # the supported bf16 D256 path (FP8_MFMA_QK is False there). if KQ_LDS_PAD and ( - 512 % HD != 0 + FP8_MFMA_QK + or 512 % HD != 0 or KQ_SLAB_ROWS == 0 or T % KQ_SLAB_ROWS != 0 or KQ_LDS_PAD % 8 != 0 ): - # Unsupported geometry for the slab pad -> disable (byte-identical). KQ_LDS_PAD = 0 KQ_SLAB_ROWS = 0 KV_BYTES = 1 if KV_FP8 else 2 From 5fca59381f53cffa45c82d9fe0173fd06caeecc2 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 16:47:34 -0400 Subject: [PATCH 17/23] test(rocke/attention): add bs=32 D256 gfx950 parity scenario (2 blocks/tile) Closes the block_size coverage gap flagged in #9233 review: the D256 gfx950 cohort predicate has no block_size constraint and the fast route forces T=64, so bs=32 (2 blocks/tile) took the fast path while only bs=16 (4/tile) was validated. GPU-verified on gfx950/MI355X: bs=32 builds the fast spec (b32_t64...) and is correct (max_abs=0.01562, bf16 band) -> bs=32 stays enabled, gap closed by coverage rather than by narrowing the predicate. --- .../attention/prefill/parity_unified_attention.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index df112fc50f46..dec1fcf23497 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -285,6 +285,17 @@ def default_scenarios() -> List[Scenario]: block_size=16, dtype=torch.bfloat16, ), + # bs=32 D256: takes the same fast route (T=64 forced -> 2 blocks/tile), + # a distinct multi-block ratio from the bs=16 (4/tile) validated points. + Scenario( + name="prefill_d256_gqa8_b32_sq4096", + seq_lens=[(4096, 4096)], + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=32, + dtype=torch.bfloat16, + ), Scenario( name="prefill_d128_b16", seq_lens=[(64, 64), (128, 256), (32, 256)], From d61156a3a851581594e938abd86e25ab9863c0b4 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 17:08:43 -0400 Subject: [PATCH 18/23] test(rocke/attention): pin D256 gfx950 fast-spec IR in the golden Closes the #9233 review gap: no golden covered the D256 fast spec (use_kq_lds_pad + use_softmax_mfma_interleave). The slab pad's async-DMA write and padded read remap must agree byte-for-byte or numerics silently corrupt, but only the spec dataclass was asserted -- an addressing drift would pass every CPU test and fail only on GPU. Adds an "attention" IR-parity case (attention/gfx950/d256_pad_interleave) that builds the production fast spec and pins its lowered LLVM IR (sha256, both llvm20/llvm22 flavors) in rocke_representative_ir_sha256. _d256_gfx950_fast gates on the memoized runtime device arch via _resolve_attention_arch (not the per-case arch, which only feeds the lowerer), and the spec builder holds a bound import of it; build_attention patches both bindings wholesale to gfx950 so the fast spec is selected deterministically on any bless/CI host, and raises if the fast spec / pad+interleave is not selected (never silently pins the fallback). Bless is add-only: no existing case sha changed. --- .../rocke_representative_ir_sha256.json | 26 ++++++- .../instances/rocke_ir_parity_harness.py | 75 +++++++++++++++++++ 2 files changed, 97 insertions(+), 4 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 16748c2188ae..25350bc49af4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -2,6 +2,14 @@ "flavors": { "llvm20": { "cases": { + "attention/gfx950/d256_pad_interleave": { + "arch": "gfx950", + "bytes": 609602, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "4650f01f089d4cd21714515984b697f6dd307bf506c7993ede6a04286f50a14f", + "status": "ok" + }, "conv/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 42207, @@ -308,20 +316,29 @@ "gfx1250": 2, "gfx90a": 3, "gfx942": 4, - "gfx950": 14 + "gfx950": 15 }, "expected_failures": 0, "families": { + "attention": 1, "conv": 7, "deep_fused_conv": 9, "gemm": 12, "moe": 9 }, - "ok_cases": 37 + "ok_cases": 38 } }, "llvm22": { "cases": { + "attention/gfx950/d256_pad_interleave": { + "arch": "gfx950", + "bytes": 609706, + "family": "attention", + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "50a532b844ea16f58bdc3bbc5b8bc53434642fa4bfeb4355f47b8c4f58e093c4", + "status": "ok" + }, "conv/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 42350, @@ -628,16 +645,17 @@ "gfx1250": 2, "gfx90a": 3, "gfx942": 4, - "gfx950": 14 + "gfx950": 15 }, "expected_failures": 0, "families": { + "attention": 1, "conv": 7, "deep_fused_conv": 9, "gemm": 12, "moe": 9 }, - "ok_cases": 37 + "ok_cases": 38 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 93a4be210743..475ef9d7ece6 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -235,6 +235,71 @@ def _build(): return _build +def build_attention(arch): + """D256 gfx950 bf16 prefill *fast* spec: 32x32-transposed + + ``use_kq_lds_pad`` (slab-granularity K_lds pad) + ``use_softmax_mfma_interleave``. + + Pins the *emitted addressing* for the slab pad: its async-DMA write and the + padded read remap must agree byte-for-byte or numerics silently corrupt, and + only the spec dataclass is otherwise asserted -- a codegen drift in the + addressing would pass every CPU test and fail only on GPU (#9233 review). + + ``_d256_gfx950_fast`` gates on the *memoized runtime device arch* via + ``_resolve_attention_arch`` -- NOT the per-case ``arch`` (which only feeds the + lowerer). The spec builder also holds a *bound* import of it. Both bindings + are patched wholesale (per its docstring) to ``arch`` so the fast spec is + selected deterministically on any bless/CI host, GPU or not; otherwise a + non-gfx950 host would silently pin the FALLBACK spec (no pad/interleave). + """ + + def _build(): + import builders.common.attention_spec_builder as asb + import kernels.common.attention_unified as au + from kernels import build_unified_attention_2d_tiled + from kernels.common.attention_unified import ( + UnifiedAttentionProblem, + _d256_gfx950_fast, + _tiled_spec_from_problem, + ) + + # Validated D256 cohort point (GQA 16/2, hd256, bs16, sq4096 bf16). + problem = UnifiedAttentionProblem( + total_q=4096, + num_seqs=1, + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + max_seqlen_q=4096, + max_seqlen_k=4096, + dtype="bf16", + num_sms=120, + ) + pin = lambda: arch # noqa: E731 + o_au, o_asb = au._resolve_attention_arch, asb._resolve_attention_arch + au._resolve_attention_arch = pin + asb._resolve_attention_arch = pin + try: + if not _d256_gfx950_fast(problem): + raise RuntimeError( + f"D256 fast spec not selected under pinned arch {arch!r}; " + "golden would pin the fallback (no pad/interleave)" + ) + spec = _tiled_spec_from_problem(problem) + if not (spec.use_kq_lds_pad and spec.use_softmax_mfma_interleave): + raise RuntimeError( + "expected pad+interleave in the D256 fast spec; got " + f"pad={spec.use_kq_lds_pad} " + f"interleave={spec.use_softmax_mfma_interleave}" + ) + return build_unified_attention_2d_tiled(spec, arch=arch) + finally: + au._resolve_attention_arch = o_au + asb._resolve_attention_arch = o_asb + + return _build + + def cases(): out = [] @@ -812,6 +877,16 @@ def add(family, case_id, arch, build): native_int=True, ), ) + # Attention: D256 gfx950 bf16 prefill fast spec. Pins the slab-pad DMA + # write + padded read remap addressing and the softmax<->MFMA interleave -- + # coupled optimizations that only agree byte-for-byte in the emitted IR, so + # a CPU spec assert can't catch a drift (it would fail only on GPU). #9233. + add( + "attention", + "attention/gfx950/d256_pad_interleave", + "gfx950", + build_attention("gfx950"), + ) return out From 684949124d27fe6d7c5dc32c722c0824ebef35c6 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 17:11:39 -0400 Subject: [PATCH 19/23] fix(rocke/attention): guard run_torch_flash against all-empty query_lens Copilot review (#9471): if every query_len is <= 0, seqs stays empty and the seqs[0] backend probe IndexErrors into the generic "no working torch sdpa backend" NotImplementedError -- a misleading failure mode. Raise an explicit ValueError naming the query_lens so the failure is deterministic and actionable. --- .../library/builders/common/torch_flash_reference.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py index 36b6215e6a78..e7362740cac8 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/torch_flash_reference.py @@ -68,6 +68,15 @@ def run_torch_flash( vh = vh.repeat_interleave(nqpkv, dim=1) seqs.append((qh, kh, vh, ql, kvl)) + # Every sequence was empty (all query_lens <= 0) -> nothing to time. Fail + # explicitly here rather than letting the seqs[0] backend probe below + # IndexError into a misleading "no working torch sdpa backend". + if not seqs: + raise ValueError( + "run_torch_flash: no sequence with query_len > 0 " + f"(query_lens={list(query_lens)}); nothing to time" + ) + # Pick a working flash-first backend once. backend = None for be, name in ( From 75796c00092db7c5921262a03aabbd1ed846540b Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 17:25:05 -0400 Subject: [PATCH 20/23] test(rocke): move D256 fast-spec IR golden to the library test layer Self-review (rocke-review.md sec.4 / one-way dependency): the D256 IR golden added earlier lived in the platform IR-parity harness (platform/tests/instances/rocke_ir_parity_harness.py) and imported the library packages kernels/builders -- a platform-test-imports-library violation. Precedent: library/tests/test_ci_static_attention.py was moved out of platform/tests for exactly this reason ("references the library builders package and therefore belongs in the library test layer"). Relocate it: - Revert the platform harness + rocke_representative_ir_sha256.json to platform-pure (byte-identical to before; the harness only builds rocke.instances kernels again). - Add library/tests/test_d256_attention_ir_golden.py + its golden. Same pinning (both _resolve_attention_arch bindings -> gfx950, raises if the fast spec / pad+interleave is not selected) and same lowered IR sha (llvm20/llvm22), but importing only rocke.core.lower_llvm from platform (library -> platform is the allowed direction). CPU-only, no GPU/comgr. --- .../golden/d256_attention_ir_sha256.json | 25 +++ .../tests/test_d256_attention_ir_golden.py | 182 ++++++++++++++++++ .../rocke_representative_ir_sha256.json | 26 +-- .../instances/rocke_ir_parity_harness.py | 75 -------- 4 files changed, 211 insertions(+), 97 deletions(-) create mode 100644 dnn-providers/hip-kernel-provider/rocke/library/tests/golden/d256_attention_ir_sha256.json create mode 100755 dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/d256_attention_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/d256_attention_ir_sha256.json new file mode 100644 index 000000000000..c3beaa509869 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/golden/d256_attention_ir_sha256.json @@ -0,0 +1,25 @@ +{ + "flavors": { + "llvm20": { + "cases": { + "d256_pad_interleave/gfx950": { + "arch": "gfx950", + "bytes": 609602, + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "4650f01f089d4cd21714515984b697f6dd307bf506c7993ede6a04286f50a14f" + } + } + }, + "llvm22": { + "cases": { + "d256_pad_interleave/gfx950": { + "arch": "gfx950", + "bytes": 609706, + "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", + "sha256": "50a532b844ea16f58bdc3bbc5b8bc53434642fa4bfeb4355f47b8c4f58e093c4" + } + } + } + }, + "schema": "rocke.d256_attention_ir_sha256/v1" +} diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py new file mode 100755 index 000000000000..614775a06980 --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Library-side IR golden for the gfx950 D256 bf16 prefill *fast* spec. + +Pins the emitted LLVM IR (sha256, both ``llvm20``/``llvm22`` flavors) of the +D256 fast spec: 32x32-transposed + ``use_kq_lds_pad`` (slab-granularity K_lds +pad) + ``use_softmax_mfma_interleave``. The slab pad's async-DMA write and its +padded read remap must agree byte-for-byte or numerics silently corrupt, but +only the spec dataclass is otherwise asserted -- an addressing drift would pass +every CPU test and fail only on GPU (#9233 review). + +Lives in the LIBRARY test layer (not ``platform/tests``) because it builds a +library kernel (``kernels`` / ``builders``); platform tests must not import +library (one-way ``platform <- library`` dependency). It imports only the +pure-Python lowerer from ``rocke.core.lower_llvm`` (``library -> platform`` is +the allowed direction) with an EXPLICIT llvm flavor, so the recorded sha is a +function of the committed IR alone -- reproducible in CI on any host, GPU or +not (no ``rocke_engine`` / comgr / real device needed). + +Bless (only from a verified-good tree, alongside a reviewed IR change):: + + python3 library/tests/test_d256_attention_ir_golden.py --write +""" +from __future__ import annotations + +import hashlib +import json +import unittest +from pathlib import Path + +GOLDEN = Path(__file__).resolve().parent / "golden" / "d256_attention_ir_sha256.json" +GOLDEN_FLAVORS = ("llvm20", "llvm22") +GOLDEN_SCHEMA = "rocke.d256_attention_ir_sha256/v1" + + +def _sha(text: str) -> str: + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def _current_flavor() -> str: + """The llvm flavor this host autodetects (the golden stores both; the gate + compares only this one, so one committed golden is valid on every host).""" + from rocke.core.lower_llvm import _resolve_llvm_flavor + + return _resolve_llvm_flavor() + + +def _build_d256_fast(arch: str): + """Build the *production* D256 fast-spec kernel under a pinned arch. + + ``_d256_gfx950_fast`` gates on the *memoized runtime device arch* via + ``_resolve_attention_arch`` -- NOT any per-case argument -- and the spec + builder holds a *bound* import of it. Both bindings are patched wholesale + (per its docstring: "Tests that monkeypatch this function replace it + wholesale") to ``arch`` so the fast spec is selected deterministically on + any host; otherwise a non-gfx950 host silently pins the FALLBACK spec (no + pad/interleave) and the golden becomes host-dependent. Raises if the fast + spec / pad+interleave is not selected, so the golden can never pin the + fallback by accident. + """ + import builders.common.attention_spec_builder as asb + import kernels.common.attention_unified as au + from kernels import build_unified_attention_2d_tiled + from kernels.common.attention_unified import ( + UnifiedAttentionProblem, + _d256_gfx950_fast, + _tiled_spec_from_problem, + ) + + # Validated D256 cohort point (GQA 16/2, hd256, bs16, sq4096 bf16). + problem = UnifiedAttentionProblem( + total_q=4096, + num_seqs=1, + num_query_heads=16, + num_kv_heads=2, + head_size=256, + block_size=16, + max_seqlen_q=4096, + max_seqlen_k=4096, + dtype="bf16", + num_sms=120, + ) + pin = lambda: arch # noqa: E731 + o_au, o_asb = au._resolve_attention_arch, asb._resolve_attention_arch + au._resolve_attention_arch = pin + asb._resolve_attention_arch = pin + try: + if not _d256_gfx950_fast(problem): + raise RuntimeError( + f"D256 fast spec not selected under pinned arch {arch!r}; " + "golden would pin the fallback (no pad/interleave)" + ) + spec = _tiled_spec_from_problem(problem) + if not (spec.use_kq_lds_pad and spec.use_softmax_mfma_interleave): + raise RuntimeError( + "expected pad+interleave in the D256 fast spec; got " + f"pad={spec.use_kq_lds_pad} " + f"interleave={spec.use_softmax_mfma_interleave}" + ) + return build_unified_attention_2d_tiled(spec, arch=arch) + finally: + au._resolve_attention_arch = o_au + asb._resolve_attention_arch = o_asb + + +def _cases(): + """case_id -> (arch, build thunk). Add future D256 IR pins here.""" + return { + "d256_pad_interleave/gfx950": ("gfx950", lambda: _build_d256_fast("gfx950")), + } + + +def _lower(arch, build, flavor): + from rocke.core.lower_llvm import _lower_kernel_to_llvm_python + + kernel = build() + llvm = _lower_kernel_to_llvm_python(kernel, arch=arch, llvm_flavor=flavor) + return { + "arch": arch, + "kernel_name": getattr(kernel, "name", ""), + "sha256": _sha(llvm), + "bytes": len(llvm.encode("utf-8")), + } + + +def _run(flavor): + return { + cid: _lower(arch, build, flavor) for cid, (arch, build) in _cases().items() + } + + +def build_golden() -> dict: + return { + "schema": GOLDEN_SCHEMA, + "flavors": {fl: {"cases": _run(fl)} for fl in GOLDEN_FLAVORS}, + } + + +class TestD256AttentionIrGolden(unittest.TestCase): + def test_ir_matches_golden(self): + """CPU-only: recompute the D256 fast-spec IR and compare its sha256 to + the blessed golden for this host's llvm flavor. Any drift in the + emitted addressing (slab-pad write/read, interleave) fails here.""" + self.assertTrue( + GOLDEN.is_file(), + f"missing golden {GOLDEN}; bless with " + "`python3 library/tests/test_d256_attention_ir_golden.py --write`", + ) + doc = json.loads(GOLDEN.read_text()) + flavor = _current_flavor() + base = doc.get("flavors", {}).get(flavor) + self.assertIsNotNone( + base, f"golden has no entry for flavor {flavor!r} (have {sorted(doc.get('flavors', {}))})" + ) + cur = _run(flavor) + self.assertEqual( + set(base["cases"]), + set(cur), + "D256 IR golden case set changed; re-bless with --write", + ) + for cid, brec in base["cases"].items(): + crec = cur[cid] + self.assertEqual( + brec["sha256"], + crec["sha256"], + f"IR DRIFT [{flavor}] {cid}: {brec['sha256']} -> {crec['sha256']} " + f"(kernel {crec['kernel_name']}); re-bless only with a reviewed, " + "expected IR change", + ) + + +if __name__ == "__main__": + import sys + + if "--write" in sys.argv: + GOLDEN.parent.mkdir(parents=True, exist_ok=True) + GOLDEN.write_text(json.dumps(build_golden(), indent=2, sort_keys=True) + "\n") + summary = ", ".join(f"{fl}={len(_run(fl))} case(s)" for fl in GOLDEN_FLAVORS) + print(f"wrote {GOLDEN}: {summary}") + else: + unittest.main(verbosity=2) diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json index 25350bc49af4..16748c2188ae 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/golden/rocke_representative_ir_sha256.json @@ -2,14 +2,6 @@ "flavors": { "llvm20": { "cases": { - "attention/gfx950/d256_pad_interleave": { - "arch": "gfx950", - "bytes": 609602, - "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", - "sha256": "4650f01f089d4cd21714515984b697f6dd307bf506c7993ede6a04286f50a14f", - "status": "ok" - }, "conv/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 42207, @@ -316,29 +308,20 @@ "gfx1250": 2, "gfx90a": 3, "gfx942": 4, - "gfx950": 15 + "gfx950": 14 }, "expected_failures": 0, "families": { - "attention": 1, "conv": 7, "deep_fused_conv": 9, "gemm": 12, "moe": 9 }, - "ok_cases": 38 + "ok_cases": 37 } }, "llvm22": { "cases": { - "attention/gfx950/d256_pad_interleave": { - "arch": "gfx950", - "bytes": 609706, - "family": "attention", - "kernel_name": "rocke_uattn2d_tiled_d256_b16_t64_h16kv2_bf16_w2_mw32_mfma32_stqk_s1_mask1_hlpv_mlim_ksb_mphsplit_smxil2g4_qdreg", - "sha256": "50a532b844ea16f58bdc3bbc5b8bc53434642fa4bfeb4355f47b8c4f58e093c4", - "status": "ok" - }, "conv/gfx1151/n1h8c16k32r3": { "arch": "gfx1151", "bytes": 42350, @@ -645,17 +628,16 @@ "gfx1250": 2, "gfx90a": 3, "gfx942": 4, - "gfx950": 15 + "gfx950": 14 }, "expected_failures": 0, "families": { - "attention": 1, "conv": 7, "deep_fused_conv": 9, "gemm": 12, "moe": 9 }, - "ok_cases": 38 + "ok_cases": 37 } } }, diff --git a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py index 475ef9d7ece6..93a4be210743 100644 --- a/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py +++ b/dnn-providers/hip-kernel-provider/rocke/platform/tests/instances/rocke_ir_parity_harness.py @@ -235,71 +235,6 @@ def _build(): return _build -def build_attention(arch): - """D256 gfx950 bf16 prefill *fast* spec: 32x32-transposed + - ``use_kq_lds_pad`` (slab-granularity K_lds pad) + ``use_softmax_mfma_interleave``. - - Pins the *emitted addressing* for the slab pad: its async-DMA write and the - padded read remap must agree byte-for-byte or numerics silently corrupt, and - only the spec dataclass is otherwise asserted -- a codegen drift in the - addressing would pass every CPU test and fail only on GPU (#9233 review). - - ``_d256_gfx950_fast`` gates on the *memoized runtime device arch* via - ``_resolve_attention_arch`` -- NOT the per-case ``arch`` (which only feeds the - lowerer). The spec builder also holds a *bound* import of it. Both bindings - are patched wholesale (per its docstring) to ``arch`` so the fast spec is - selected deterministically on any bless/CI host, GPU or not; otherwise a - non-gfx950 host would silently pin the FALLBACK spec (no pad/interleave). - """ - - def _build(): - import builders.common.attention_spec_builder as asb - import kernels.common.attention_unified as au - from kernels import build_unified_attention_2d_tiled - from kernels.common.attention_unified import ( - UnifiedAttentionProblem, - _d256_gfx950_fast, - _tiled_spec_from_problem, - ) - - # Validated D256 cohort point (GQA 16/2, hd256, bs16, sq4096 bf16). - problem = UnifiedAttentionProblem( - total_q=4096, - num_seqs=1, - num_query_heads=16, - num_kv_heads=2, - head_size=256, - block_size=16, - max_seqlen_q=4096, - max_seqlen_k=4096, - dtype="bf16", - num_sms=120, - ) - pin = lambda: arch # noqa: E731 - o_au, o_asb = au._resolve_attention_arch, asb._resolve_attention_arch - au._resolve_attention_arch = pin - asb._resolve_attention_arch = pin - try: - if not _d256_gfx950_fast(problem): - raise RuntimeError( - f"D256 fast spec not selected under pinned arch {arch!r}; " - "golden would pin the fallback (no pad/interleave)" - ) - spec = _tiled_spec_from_problem(problem) - if not (spec.use_kq_lds_pad and spec.use_softmax_mfma_interleave): - raise RuntimeError( - "expected pad+interleave in the D256 fast spec; got " - f"pad={spec.use_kq_lds_pad} " - f"interleave={spec.use_softmax_mfma_interleave}" - ) - return build_unified_attention_2d_tiled(spec, arch=arch) - finally: - au._resolve_attention_arch = o_au - asb._resolve_attention_arch = o_asb - - return _build - - def cases(): out = [] @@ -877,16 +812,6 @@ def add(family, case_id, arch, build): native_int=True, ), ) - # Attention: D256 gfx950 bf16 prefill fast spec. Pins the slab-pad DMA - # write + padded read remap addressing and the softmax<->MFMA interleave -- - # coupled optimizations that only agree byte-for-byte in the emitted IR, so - # a CPU spec assert can't catch a drift (it would fail only on GPU). #9233. - add( - "attention", - "attention/gfx950/d256_pad_interleave", - "gfx950", - build_attention("gfx950"), - ) return out From 3839cd1a4619441ab92f41368a05c2f7dec5b706 Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Wed, 15 Jul 2026 18:20:48 -0400 Subject: [PATCH 21/23] style(rocke): apply black to changed files --- .../rocke/library/dispatch/attention.py | 1 + .../tests/dispatch/attention/test_d256_gfx950_routing.py | 4 +--- .../rocke/library/tests/test_d256_attention_ir_golden.py | 7 +++---- 3 files changed, 5 insertions(+), 7 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py index f40b9780c6b4..6e72163a7fb3 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/dispatch/attention.py @@ -356,6 +356,7 @@ def select(req: OperatorRequest) -> AttentionSpec: assert isinstance(req, AttentionRequest) problem = _problem(req) from kernels.common.attention_unified import _d256_gfx950_spec_overrides + return AttentionSpec( path="2d", head_size=problem.head_size, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py index 13ca34b03fe6..1997d4f4a5cc 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_d256_gfx950_routing.py @@ -74,9 +74,7 @@ def test_rejects_decode(self): ) def test_rejects_sliding_window(self): - self.assertNotEqual( - _routed_spec_id(_d256(sliding_window=256)), "gfx950_d256" - ) + self.assertNotEqual(_routed_spec_id(_d256(sliding_window=256)), "gfx950_d256") class TestD256Gfx950SpecOverrides(unittest.TestCase): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py index 614775a06980..141d9ac9739e 100755 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_d256_attention_ir_golden.py @@ -125,9 +125,7 @@ def _lower(arch, build, flavor): def _run(flavor): - return { - cid: _lower(arch, build, flavor) for cid, (arch, build) in _cases().items() - } + return {cid: _lower(arch, build, flavor) for cid, (arch, build) in _cases().items()} def build_golden() -> dict: @@ -151,7 +149,8 @@ def test_ir_matches_golden(self): flavor = _current_flavor() base = doc.get("flavors", {}).get(flavor) self.assertIsNotNone( - base, f"golden has no entry for flavor {flavor!r} (have {sorted(doc.get('flavors', {}))})" + base, + f"golden has no entry for flavor {flavor!r} (have {sorted(doc.get('flavors', {}))})", ) cur = _run(flavor) self.assertEqual( From 922a8b334c77937da16bc07140cc88e9bdb4a4fe Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Tue, 21 Jul 2026 16:13:49 -0400 Subject: [PATCH 22/23] fix(rocke/attention): gfx950 d64 combo forces T=64 so fast_paged_kv_desc is valid; drop redundant HLPV parity variant The d64/b32/h64kv8 bf16 single-seq long-prefill combo enables use_fast_paged_kv_desc, but _select_2d_tile_size picked T=128 for it -> the spec validator (fast_paged_kv_desc requires T=64) raised, so production crashed on these shapes on gfx950. num_warps was already 4; tile_size was the incompatibility. - _select_2d_tile_size: force T=64 for the bf16 h64kv8 no-SW single-batch d64 family so its fast_paged_kv_desc enablement is valid. Measured ~1.3-1.4x over the T=128 no-fast_paged path on MI355X; other d64 single-batch shapes keep the wider T=128. - attention_spec_builder: gate use_fast_paged_kv_desc on _select_2d_tile_size==64 (self-consistency -- never set the flag with an incompatible tile). - parity_unified_attention: the 2d lane now builds the production spec directly; removed the hand-tuned HLPV variant (use_hlpv_variant + the r4_hlpv_parity spec), now redundant because production emits the same T=64 + fast_paged kernel. Verified on MI355X (gfx950): the 4 family shapes build + run correct (bf16) with the T=64 + fast_paged kernel and match the former HLPV ceiling. --- .../builders/common/attention_spec_builder.py | 5 ++ .../prefill/parity_unified_attention.py | 87 +++---------------- .../kernels/common/attention_unified.py | 7 ++ 3 files changed, 22 insertions(+), 77 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py index 5be4d8f1a0ec..95620341a2da 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py @@ -279,6 +279,11 @@ def _tiled_spec_from_problem( and not problem.use_fp8 and problem.num_query_heads == 64 and problem.num_kv_heads == 8 + # self-consistency: fast_paged_kv_desc requires T==64. Only enable it + # when the tile selector actually picks 64 for this shape, so the flag + # can never be set with an incompatible tile (which trips the spec + # validator). _select_2d_tile_size forces T=64 for this family. + and _select_2d_tile_size(problem) == 64 ), use_register_pv=_enable_register_pv(problem), use_fp8_mfma_qk=_enable_fp8_mfma_qk(problem), diff --git a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py index f70010cfc364..169c88ccd4d5 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/builders/gfx950/attention/prefill/parity_unified_attention.py @@ -1374,16 +1374,15 @@ def _run_rocke( hip_stream = _bench_stream_handle() - # NOTE: the ``"2d"`` lane force-builds the hand-tuned MFMA-32x32 / - # half-local-PV 2D kernel directly so we can measure the *best* 2D - # variant for a shape (this is the aspirational ceiling for the 2D - # path, including the d64/b32/h64kv8 trace family). The ``"auto"`` - # lane must instead exercise the *production* dispatcher + # NOTE: the ``"2d"`` lane builds and launches the PRODUCTION 2D spec + # (``_tiled_spec_from_problem``) directly -- GPU-time only, no + # ``run_unified_attention_torch`` dispatcher overhead -- so it measures the + # exact kernel production selects for a shape. The ``"auto"`` lane instead + # exercises the *production* dispatcher # (``run_unified_attention_torch(backend="auto")`` -> ``select_path``), - # otherwise we'd be reporting forced-2D timings as if they were what - # production launches -- which mis-measures decode shapes by ~6x - # (production correctly routes them to the 3D split-KV path). Keep the - # two lanes strictly separate. + # otherwise we'd report forced-2D timings as if they were what production + # launches -- which mis-measures decode shapes by ~6x (production correctly + # routes them to the 3D split-KV path). Keep the two lanes strictly separate. if path == "2d": from rocke import compile_kernel from kernels import ( @@ -1492,66 +1491,7 @@ def call_once(): if not ok: raise NotImplementedError(reason) - def use_hlpv_variant() -> bool: - if dtype_str != "bf16": - return False - if s.head_size != 64 or s.block_size != 32: - return False - if s.num_query_heads != 64 or s.num_kv_heads != 8: - return False - if ( - s.softcap > 0 - or problem.use_fp8 - or problem.use_alibi - or problem.use_qq_bias - ): - return False - if data["max_query_len"] <= 256: - return False - # The half-local PV path regresses on the high-num-seq SW tail. - if (s.sliding_window or 0) > 0 and len(s.seq_lens) >= 450: - return False - return True - - use_hlpv = use_hlpv_variant() - # Measured best local policy for the d64/b32/h64kv8 bf16+sinks family: - # * no-SW: R4_s1mask_hlpv + mask-limit + fast paged-KV + skip legacy Q - # * SW: R4_s1mask_hlpv + fast paged-KV + skip legacy Q - # * SW high-num-seq tail falls back to plain R4 in use_hlpv_variant(). - use_transposed_mask_limit = use_hlpv and (s.sliding_window or 0) == 0 - use_mfma32_skip_legacy_qreg = use_hlpv - use_fast_paged_kv_desc = use_hlpv - # AGPR0 is still experimental and did not improve this path broadly; keep - # it as an explicit environment opt-in for microbench work only. - use_agpr_alloc_zero = ( - use_hlpv and os.environ.get("ROCKE_ATTENTION_AGPR_ALLOC_ZERO") == "1" - ) - spec = UnifiedAttention2DTiledSpec( - head_size=s.head_size, - block_size=s.block_size, - num_query_heads=s.num_query_heads, - num_kv_heads=s.num_kv_heads, - dtype=dtype_str, - use_sinks=data["sinks"] is not None, - sliding_window=s.sliding_window or 0, - has_softcap=s.softcap > 0, - use_alibi=data["alibi_slopes"] is not None, - use_qq_bias=qq_bias is not None, - num_seqs=len(s.seq_lens), - num_warps=4, - waves_per_eu=2, - tile_size=2 * s.block_size, - block_m_per_warp=32, - use_mfma_32x32=True, - use_transposed_qk_32x32=True, - use_transposed_scalar_state=use_hlpv, - use_transposed_mask_once=use_hlpv, - use_transposed_half_local_pv=use_hlpv, - use_mfma32_skip_legacy_qreg=use_mfma32_skip_legacy_qreg, - use_transposed_mask_limit=use_transposed_mask_limit, - use_fast_paged_kv_desc=use_fast_paged_kv_desc, - use_agpr_alloc_zero=use_agpr_alloc_zero, - ) + spec = _tiled_spec_from_problem(problem) kernel = build_unified_attention_2d_tiled(spec) if _select_2d_compile_backend(problem) == "hipcc": from rocke.helpers.compile import compile_kernel_via_hipcc @@ -1565,14 +1505,7 @@ def use_hlpv_variant() -> bool: signature=_attn_signature( dtype_str, include_bt_stride=True, include_qq_bias_stride=True ), - cache_key=( - "r4_hlpv_parity", - spec.kernel_name(), - use_agpr_alloc_zero, - use_mfma32_skip_legacy_qreg, - use_transposed_mask_limit, - use_fast_paged_kv_desc, - ), + cache_key=("prod_2d_direct", spec.kernel_name()), ) vals = _attn_values( problem=problem, diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index c6cc8b95e6f6..e86e60baeadf 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -812,6 +812,13 @@ def _select_2d_tile_size(problem: UnifiedAttentionProblem) -> int: # GQA-8 d64 S2048 (1.19-2.68x over prod). if _enable_single_batch_combo(problem): if problem.head_size == 64: + # fast_paged_kv_desc (needs T=64) is a ~1.3-1.4x win for the bf16 + # h64kv8 no-SW family -> beats the T=128 early-V path (measured on + # MI355X). Force T=64 there so the combo's fast_paged enablement is + # valid; other d64 single-batch shapes keep the wider T=128 tile. + if (problem.dtype == "bf16" and problem.num_query_heads == 64 + and problem.num_kv_heads == 8 and problem.sliding_window == 0): + return 64 return 128 # d128 occupancy lever (supersedes the small-tile pick for the # LONG-context holdout): at num_warps=2 the d128 combo is LDS-bound From c24e8ec9b9bbf43e7b76b54f6944ed2ad6ea03ec Mon Sep 17 00:00:00 2001 From: Aviral Goel Date: Tue, 21 Jul 2026 16:20:02 -0400 Subject: [PATCH 23/23] style(rocke/attention): black-format the d64 tile-fix conditional --- .../rocke/library/kernels/common/attention_unified.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py index e86e60baeadf..b2d1896eb4a4 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py @@ -816,8 +816,12 @@ def _select_2d_tile_size(problem: UnifiedAttentionProblem) -> int: # h64kv8 no-SW family -> beats the T=128 early-V path (measured on # MI355X). Force T=64 there so the combo's fast_paged enablement is # valid; other d64 single-batch shapes keep the wider T=128 tile. - if (problem.dtype == "bf16" and problem.num_query_heads == 64 - and problem.num_kv_heads == 8 and problem.sliding_window == 0): + if ( + problem.dtype == "bf16" + and problem.num_query_heads == 64 + and problem.num_kv_heads == 8 + and problem.sliding_window == 0 + ): return 64 return 128 # d128 occupancy lever (supersedes the small-tile pick for the