feat(rocke): enable head-dim-256 bf16 flash prefill on gfx942 + gfx950

dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_unified.py

@@ -1856,14 +1856,202 @@ 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:
+    # 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 (
+        "fp8e4m3",
+        "bf8e5m2",
+        "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``."""
+    block_m = spec.num_warps * spec.block_m_per_warp
+    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 (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:
+    # 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):
+    """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, "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, "block_m > tile_size (Q would not fit the lone K slot)"
+    try:
+        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, "tile_size already <= 64"
+    try:
+        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):
+    """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)."""
+    # Arch-agnostic: the resolver keys off the register-PV path (whose footprint
+    # model it owns) and the target arch's LDS cap read *dynamically* from the
+    # arch-target API -- no hard-coded arch name or capacity. It engages for any
+    # register-PV 2D spec over that arch's cap. Validated on gfx950 (CDNA4); a
+    # strict no-op on other arches because their over-budget 2D specs are already
+    # filtered upstream by the (more conservative, Q/P-staged) ``supports_tiled_2d``
+    # gate, so any spec that reaches the resolver there already fits -> every
+    # currently-compiling config stays byte-identical.
+    if not getattr(spec, "use_register_pv", False):
+        return spec
+    try:
+        cap = _lds_capacity_bytes()
+    except Exception:
+        return spec  # arch has no declared LDS cap -> nothing to resolve against
+    if _lds_bytes_regpv(spec) <= cap:
+        return spec
+
+    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 fits "
+        f"[{'; '.join(tried)}]."
+    )
+
+
 def _tiled_spec_from_problem(
     problem: UnifiedAttentionProblem,
 ):
     from builders.common.attention_spec_builder import (
         _tiled_spec_from_problem as _impl,
     )
 
-    return _impl(problem)
+    _spec = _impl(problem)
+    return _resolve_lds_budget(_spec)
 
 
 def _select_2d_block_m_per_warp(problem: UnifiedAttentionProblem) -> int:

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 (

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

@@ -0,0 +1,145 @@
+# 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 or non-register-PV config, on any arch.
+
+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_arch_agnostic_targets_dynamic_cap(self):
+        """Arch-agnostic: the resolver targets whatever LDS cap the arch-target
+        API reports -- not a hard-coded gfx950 / 163840 constant. Same D256
+        K-double geometry (204800 B): with a larger reported cap it already fits
+        and is returned byte-identical; with a tiny cap it engages and raises."""
+        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)  # 204800 B
+            with mock.patch.object(au, "_lds_capacity_bytes", return_value=262144):
+                self.assertIs(au._resolve_lds_budget(k_double), k_double)
+            with mock.patch.object(au, "_lds_capacity_bytes", return_value=65536):
+                with self.assertRaises(RuntimeError):
+                    au._resolve_lds_budget(replace(spec, use_k_single_buffer=False))
+
+    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
+        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__":
+    unittest.main()