diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py index 65fdf854100b..b5d771ea6b17 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/common/attention_dense_spec.py @@ -48,6 +48,30 @@ INT32_LIMIT = 2**31 +def partial_tile_reason(spec: "AttentionDenseSpec") -> str: + """Say which sequence length leaves a partial tile on a ``ragged`` spec. + + The dense ``_base_spec`` factories set ``ragged`` for an ordinary graph whose + lengths are not tile multiples (Sq == Skv == 384 at block_m 256, say), so a + refusal that only says "ragged" reads as a varlen problem. This names the + length and the tile instead. + """ + parts = [] + if spec.seqlen_q % spec.block_m: + parts.append( + f"seqlen_q {spec.seqlen_q} is not a multiple of block_m {spec.block_m}" + ) + if spec.seqlen_kv % spec.block_n: + parts.append( + f"seqlen_kv {spec.seqlen_kv} is not a multiple of block_n {spec.block_n}" + ) + return "; ".join(parts) or ( + f"ragged=True at seqlen_q {spec.seqlen_q} and seqlen_kv {spec.seqlen_kv}, " + f"which are already multiples of block_m {spec.block_m} and block_n " + f"{spec.block_n}" + ) + + @dataclass(frozen=True) class AttentionDenseSpec: """Shared compile-time problem and geometry for dense attention.""" @@ -164,7 +188,11 @@ def __post_init__(self) -> None: if self.varlen: raise ValueError("ragged is not supported with varlen") if self.sliding_window > 0: - raise ValueError("ragged is not supported with sliding_window") + raise ValueError( + "sliding_window needs whole tiles: " + f"{partial_tile_reason(self)} (the partial-tile path has no " + "window mask)" + ) else: if self.seqlen_q % self.block_m != 0: raise ValueError( @@ -629,4 +657,5 @@ def check_dense_spec_preflight(spec: AttentionDenseSpec) -> tuple[bool, str]: "INT32_LIMIT", "attention_dense_cache_key", "check_dense_spec_preflight", + "partial_tile_reason", ] diff --git a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py index 45da8083eae6..c2ac71355704 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/kernels/gfx942/attention_dense.py @@ -212,6 +212,7 @@ DENSE_TILE_GEOMETRIES, attention_dense_cache_key, check_dense_spec_preflight, + partial_tile_reason, ) # C-output lane maps: IDENTICAL between the 32x32x8 (gfx942) and 32x32x16 (gfx950) @@ -1184,7 +1185,12 @@ def supports_attention_dense( if spec.varlen: return False, "gfx942 attention_dense: varlen not yet supported" if spec.ragged: - return False, "gfx942 attention_dense: ragged not yet supported" + # The factory sets ragged for any graph whose lengths leave a partial tile; + # name the length and the tile rather than calling an ordinary graph ragged. + return False, ( + f"gfx942 attention_dense needs whole tiles: {partial_tile_reason(spec)} " + f"(the partial-tile path is not yet supported)" + ) # sliding_window is supported (KV-loop prune + window mask); the shared spec # __post_init__ re-run above enforces its constraints (W % block_n, causal). # use_sinks is supported here (m/l seed per work item). The shared __post_init__ diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py index cd4798ab1324..66c5086a2509 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx942_dense_wiring.py @@ -35,6 +35,7 @@ iter_registered_attention_combos, tuning_spec_with_knobs, ) +from dispatch.attention.gfx942_dense import _base_spec from kernels.common.attention_dense_spec import ( AttentionDenseSpec, attention_dense_cache_key, @@ -205,16 +206,39 @@ def test_admits_sliding_window_with_sinks(self): ok, why = _candidate().admits(_req(sliding_window=128, use_sinks=True)) self.assertTrue(ok, why) - def test_rejects_ragged_sequence_length(self): - """The default spec sets ragged=True for any non-256-multiple - self-attention length. The kernel must decline, not select-then-fail. + def test_rejects_a_partial_tile_and_names_the_length(self): + """The default spec sets ragged=True for any self-attention length that is + not a tile multiple. The kernel must decline, not select-then-fail. Capability cannot see this one: it is a property of the BUILT spec, so - it stays in the predicate.""" - with _Gfx942Arch(): - ok, why = _candidate().admits(_req(seqlen_q=1000, seqlen_k=1000)) - self.assertFalse(ok) - self.assertNotIn("capability", why) - self.assertIn("ragged", why) + it stays in the predicate. The graph is not ragged, so the reason names + the length and the tile instead of saying "ragged".""" + # The default tile depends on the shape (128 rows when the 256-row grid + # would leave CUs idle), so the expected tile is read from the default spec + # the factory builds (dispatch refuses it, so it cannot come from dispatch). + # 1152 is a block_n multiple, so only the query length is named there. + with _Gfx942Arch(): + for seqlen in (1152, 1000): + with self.subTest(seqlen=seqlen): + req = _req(seqlen_q=seqlen, seqlen_k=seqlen) + spec = _base_spec(req) + self.assertNotEqual(seqlen % spec.block_m, 0) + ok, why = _candidate().admits(req) + self.assertFalse(ok) + self.assertNotIn("capability", why) + self.assertNotIn("ragged", why) + self.assertIn( + f"seqlen_q {seqlen} is not a multiple of block_m " + f"{spec.block_m}", + why, + ) + kv_part = ( + f"seqlen_kv {seqlen} is not a multiple of block_n " + f"{spec.block_n}" + ) + if seqlen % spec.block_n: + self.assertIn(kv_part, why) + else: + self.assertNotIn("seqlen_kv", why) class TestGfx942BottomRightSafety(unittest.TestCase): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py index d5d757a67389..c80676b0b0ce 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_gfx950_dense_wiring.py @@ -18,6 +18,7 @@ from __future__ import annotations +import re import unittest from dataclasses import replace from itertools import islice @@ -619,15 +620,32 @@ def test_sliding_window_with_persistent_mode(self): self.assertIn("swa128", kname) self.assertIn("persist", kname) - def test_sliding_window_with_ragged_shape_rejected_at_dispatch(self): - """Ragged + sliding_window is rejected by the spec validator, so no - dense variant admits the request.""" + def test_sliding_window_with_partial_tile_rejected_at_dispatch(self): + """A partial-tile length plus sliding_window is rejected by the spec + validator, so no dense variant admits the request. The graph is not + ragged, so the reason names the length, the tile and the window.""" req = _gfx950_dense_req(seqlen_q=500, seqlen_k=500, sliding_window=128) with self.assertRaises(ValueError) as cm: _spec(req) - err_msg = str(cm.exception).lower() - self.assertIn("ragged", err_msg) + err_msg = str(cm.exception) + # The spec never gets built, so read the tile out of the refusal. The + # sliding-window tile can be 128 rows, so assert it is a registered tile + # that really leaves a partial block rather than naming one value. + tile = re.search( + r"seqlen_q 500 is not a multiple of block_m (\d+); " + r"seqlen_kv 500 is not a multiple of block_n (\d+)", + err_msg, + ) + self.assertIsNotNone(tile, err_msg) + block_m, block_n = int(tile.group(1)), int(tile.group(2)) + self.assertIn( + (block_m, block_n), + {(g["block_m"], g["block_n"]) for g in DENSE_TILE_GEOMETRIES.values()}, + ) + self.assertNotEqual(500 % block_m, 0) + self.assertNotEqual(500 % block_n, 0) self.assertIn("sliding_window", err_msg) + self.assertNotIn("ragged", err_msg) def test_sliding_window_without_ragged_accepted(self): spec = _spec( diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_mask_wiring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_mask_wiring.py index 10313039b5d0..5ab6a4457308 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_mask_wiring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/dispatch/attention/test_unified_mask_wiring.py @@ -13,6 +13,7 @@ from __future__ import annotations +import re import unittest from dataclasses import replace from unittest import mock @@ -30,6 +31,7 @@ _attention_band, normalize_band, ) +from kernels.common.attention_dense_spec import DENSE_TILE_GEOMETRIES NO_MASK = AttentionMaskType.NO_MASK TOP_LEFT = AttentionMaskType.TOP_LEFT_CAUSAL @@ -354,9 +356,10 @@ def test_kernel_name_carries_the_window_for_every_mask(self): self.assertEqual(len(set(kernel.values())), 3, kernel) self.assertIn("sw256", kernel["w256"].split("_")) - def test_ragged_window_is_declined_not_served(self): + def test_partial_tile_window_is_declined_not_served(self): # The dense window loops do not mask a partial last KV tile; the spec rejects - # ragged + window so they never see one. + # a partial-tile length plus a window so they never see one. The graph is not + # ragged, so the reason names the lengths, the tiles and the window. from dispatch.attention import attention_candidates from rocke.dispatch.core import opt_in_probe @@ -373,7 +376,27 @@ def test_ragged_window_is_declined_not_served(self): with self.subTest(kw=kw, candidate=c.name): ok, why = c.admits(opt_in_probe(req, c)) self.assertFalse(ok) - self.assertIn("ragged", why) + # The tile depends on the shape and the window, and the + # spec is never built, so read it out of the refusal and + # check it is a registered tile that leaves a partial block. + tile = re.search( + r"seqlen_q 1000 is not a multiple of block_m (\d+); " + r"seqlen_kv 1050 is not a multiple of block_n (\d+)", + why, + ) + self.assertIsNotNone(tile, why) + block_m, block_n = int(tile.group(1)), int(tile.group(2)) + self.assertIn( + (block_m, block_n), + { + (g["block_m"], g["block_n"]) + for g in DENSE_TILE_GEOMETRIES.values() + }, + ) + self.assertNotEqual(1000 % block_m, 0) + self.assertNotEqual(1050 % block_n, 0) + self.assertIn("sliding_window", why) + self.assertNotIn("ragged", why) class TestBandNormalization(unittest.TestCase): diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py index 1e0814e01431..6a778a6c4808 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_dense_gfx942.py @@ -848,7 +848,10 @@ def test_supports_rejects_non_gfx942(): # the persistent build/decode tests and the SWA coverage below. use_sinks # is supported too; see the sink tests below. (dict(varlen=True), "varlen"), - (dict(seqlen_q=1000, seqlen_kv=1000, ragged=True), "ragged"), + ( + dict(seqlen_q=1000, seqlen_kv=1000, ragged=True), + "seqlen_q 1000 is not a multiple of block_m 256", + ), ], ) def test_supports_rejects_modes_deferred_to_later_phases(kw, marker):