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 56c2ed7f953f..0ccecdc7d890 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 @@ -1856,6 +1856,193 @@ 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, ): @@ -1863,7 +2050,8 @@ def _tiled_spec_from_problem( _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: 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 new file mode 100644 index 000000000000..aea87ee96fee --- /dev/null +++ b/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() 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..674f8ce4e749 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,156 @@ 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, "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 */ +} + +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(); + /* Arch-agnostic: key off the register-PV path + the target arch's LDS cap + * read dynamically from the arch-target API (no hard-coded arch or capacity). + * Validated on gfx950 (CDNA4); a strict no-op on other arches because their + * over-budget 2D specs are filtered upstream by the (more conservative) + * supports_tiled_2d gate, so any spec reaching the resolver already fits. */ + if(!spec->use_register_pv) + return true; + const rocke_arch_target_t* t = rocke_arch_target_from_gfx(arch); + if(t == NULL) + return true; /* arch has no declared LDS cap -> nothing to resolve against */ + int cap = t->lds_capacity_bytes; + 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,