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 b5f79fe0387b..62b81610965c 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 @@ -51,7 +51,6 @@ _gfx942_flash_wide_setting, _kv_storage_dtype, _num_segments, - _resolve_attention_arch, _resolve_gfx1250_tiled3d, _select_2d_block_m_per_warp, _select_2d_num_warps, @@ -66,13 +65,24 @@ # 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). +# +# ``_resolve_attention_arch`` MUST be reached through this module handle for the +# same reason, and is deliberately absent from the ``from ... import`` list +# above. It used to be bound, which silently defeated +# ``mock.patch.object(au, "_resolve_attention_arch", ...)``: the patch rebinds +# the attribute on the module, but a bound reference captured at import time +# still points at the original. The builder then resolved the REAL device arch, +# ``_tiled_2d_impl`` handed back that arch's spec class, and the gfx950-only +# override fields (e.g. ``use_q_direct_reg``) raised TypeError -- but only when +# this module was already imported before the patch was applied, so the +# breakage looked like flaky test ordering. from kernels.common import attention_unified as _kau def _tiled_spec_from_problem( problem: UnifiedAttentionProblem, ): - arch = _resolve_attention_arch() + arch = _kau._resolve_attention_arch() UnifiedAttention2DTiledSpec, _, _ = _tiled_2d_impl(arch) if arch == "gfx1250": return UnifiedAttention2DTiledSpec( @@ -327,7 +337,7 @@ def _tiled_spec_from_problem( def _tiled_3d_spec_from_problem( problem: UnifiedAttentionProblem, ): - arch = _resolve_attention_arch() + arch = _kau._resolve_attention_arch() UnifiedAttention3DTiledSpec, *_ = _tiled_3d_impl(arch) tile_size_override = _gfx942_3d_tile_size_override(problem) if arch == "gfx1250": diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/conftest.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/conftest.py index 2b679b204ed8..f8aa56856979 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/conftest.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/conftest.py @@ -15,6 +15,8 @@ import sys from pathlib import Path +import pytest + _LIBROOT = Path(__file__).resolve().parents[1] # tests -> rocke/library if str(_LIBROOT) not in sys.path: sys.path.insert(0, str(_LIBROOT)) @@ -22,3 +24,42 @@ _PYROOT = Path(__file__).resolve().parents[2] / "platform" / "python" if str(_PYROOT) not in sys.path: sys.path.insert(0, str(_PYROOT)) + + +@pytest.fixture(autouse=True) +def _restore_attention_arch_state(): + """Undo any write to the process-wide arch memo and its derived cache. + + ``attention_unified._RESOLVED_ATTENTION_ARCH`` memoizes the real device arch + on first resolve, and a dozen tests pin it directly to exercise a specific + arch. One that leaves it set -- or cleared -- changes what every later test + resolves, which surfaces as order-dependent failures on a GPU runner and not + at all on a CPU one. Restoring it here makes the leak class impossible + rather than relying on each test's own teardown. + + ``_2D_LAUNCH_META`` is *cleared*, not restored: it is a pure memo whose + entries are derived from whatever arch was resolved when they were computed, + so a pinned-gfx942 entry must not be served back to a gfx950 test. + ``test_attn_k_slice_hd`` and ``test_attn_bf16_d128_ring`` clear it by hand + today; this makes that unnecessary. + + Deliberately NOT covered here: ``HIPDNN_GFX942_*`` env overrides. An env leak + belongs to the test that pops them (``monkeypatch.delenv``), not to a blanket + conftest snapshot -- a snapshot would mask the leak rather than fix it. + """ + # Looked up, not imported: this fixture runs for every test in the tree, + # including ones with nothing to do with attention. Forcing the import here + # would turn any import regression in ``attention_unified`` into a tree-wide + # setup error. A test that never imported the module cannot have written the + # memo, and ``None`` is the module's pristine value, so the not-imported + # branch restores correctly either way. + _MOD = "kernels.common.attention_unified" + au = sys.modules.get(_MOD) + prev = au._RESOLVED_ATTENTION_ARCH if au is not None else None + try: + yield + finally: + au = sys.modules.get(_MOD) + if au is not None: + au._RESOLVED_ATTENTION_ARCH = prev + au._2D_LAUNCH_META.clear() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_arch_binding_guard.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_arch_binding_guard.py new file mode 100644 index 000000000000..1fd8111d4e3a --- /dev/null +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_arch_binding_guard.py @@ -0,0 +1,103 @@ +# Copyright (c) Advanced Micro Devices, Inc., or its affiliates. +# SPDX-License-Identifier: MIT +"""Guard: the spec builder must reach ``_resolve_attention_arch`` lazily. + +``builders.common.attention_spec_builder`` used to bind the resolver with +``from kernels.common.attention_unified import _resolve_attention_arch``. A bound +import freezes the reference at import time, so ``mock.patch.object(au, +"_resolve_attention_arch", ...)`` -- which rebinds the attribute on the *module* +-- never reached the builder. The builder then resolved the real device arch, +and the gfx950-only spec overrides raised ``TypeError`` against a gfx942 spec +class. The binding dates to #9057 and turned load-bearing in #9233. + +The rule is now: reach the resolver through the ``_kau`` module handle. Nothing +in the language enforces that, so these tests do -- one statically, one through +observable behaviour. + +Both tests are host-independent. The behavioural one pins *two* archs, because +the resolver falls back to ``"gfx950"`` with no device and ``_tiled_2d_impl`` +maps every arch except gfx942/gfx1250 onto the gfx950 impl -- so a single pin +can pass by coincidence on most hosts. Both builder entry points are covered: +``_tiled_spec_from_problem`` and ``_tiled_3d_spec_from_problem`` each hold their +own ``_kau._resolve_attention_arch()`` call site, and both were served by the +same bound import before the fix. + +gfx1250 is deliberately NOT pinned: it needs its own problem shape (2D rejects +``block_size=16``, 3D rejects ``head_size=128``), so it would couple this guard +to the evolving gfx1250 support matrix while adding no discriminating power -- +the gfx942 pin alone already fails on the unfixed builder. +""" + +from __future__ import annotations + +import unittest +from unittest import mock + +import builders.common.attention_spec_builder as asb +import kernels.common.attention_unified as au +from kernels import UnifiedAttentionProblem + +# Plain bf16 D128 prefill: no sinks / sliding-window / softcap / alibi / bias, so +# it routes onto the ordinary 2D tiled path on every arch under test and is not +# diverted by any arch-specific fast route. +_PROBLEM = UnifiedAttentionProblem( + total_q=1024, + num_seqs=2, + num_query_heads=16, + num_kv_heads=2, + head_size=128, + block_size=16, + max_seqlen_q=512, + max_seqlen_k=4096, + dtype="bf16", +) + +# Decode-shaped variant of the same plain bf16 problem, for the 3D entry point. +_PROBLEM_3D = UnifiedAttentionProblem( + total_q=2, + num_seqs=2, + num_query_heads=16, + num_kv_heads=2, + head_size=128, + block_size=16, + max_seqlen_q=1, + max_seqlen_k=4096, + dtype="bf16", +) + +_BOUND_IMPORT_TRAP = ( + "attention_spec_builder must reach _resolve_attention_arch through its " + "attention_unified module handle -- a bound import freezes the reference at " + "import time and silently defeats mock.patch.object on the module" +) + + +class TestArchResolverBinding(unittest.TestCase): + def test_builder_does_not_bind_the_resolver(self): + """Static form: the name must not be an attribute of the builder.""" + # assertFalse(hasattr(...)), not assertNotIn(..., vars(asb)) -- the + # latter prints the builder's entire module dict on failure. + self.assertFalse(hasattr(asb, "_resolve_attention_arch"), _BOUND_IMPORT_TRAP) + + def test_patching_the_module_steers_the_builder(self): + """Behavioural form: the pin must reach the spec class the builder picks.""" + entries = ( + ("2d", asb._tiled_spec_from_problem, _PROBLEM), + ("3d", asb._tiled_3d_spec_from_problem, _PROBLEM_3D), + ) + for arch in ("gfx942", "gfx950"): + for label, build, problem in entries: + with self.subTest(arch=arch, entry=label): + with mock.patch.object( + au, "_resolve_attention_arch", return_value=arch + ): + spec = build(problem) + self.assertTrue( + type(spec).__module__.startswith(f"kernels.{arch}."), + f"{_BOUND_IMPORT_TRAP}; {label} pinned {arch} but got " + f"{type(spec).__module__}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py index 52df9a59c0fd..239fd0ccbcdd 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attention_builds.py @@ -43,47 +43,21 @@ def _patch_resolved_arch(arch: str): - """Pin the resolved attention arch for a test, on every module that reads it. - - ``_resolve_attention_arch`` is defined in ``kernels.common.attention_unified`` - but imported *by name* into other modules (e.g. - ``builders.common.attention_spec_builder``), so each holds its own binding. - Patching only the defining module leaves the spec builder resolving the real - device arch -- which silently ignores the test's requested arch on any host - whose GPU differs (e.g. an ``arch='gfx950'`` case on a gfx942 box). Patch - every by-name importer so the pin actually reaches the builder. + """Pin the resolved attention arch for a test, on the module that defines it. + + ``_resolve_attention_arch`` lives in ``kernels.common.attention_unified`` and + ``builders.common.attention_spec_builder`` reaches it through that module + handle, so this one patch steers the builder too. A bound import in the + builder would freeze the reference at import time, leaving it on the real + device arch and silently ignoring the test's requested arch on any host whose + GPU differs (e.g. an ``arch='gfx950'`` case on a gfx942 box). That invariant + is pinned by ``test_arch_binding_guard.py``; it is not re-asserted here. """ from unittest import mock - import builders.common.attention_spec_builder as _asb import kernels.common.attention_unified as _au - targets = [_au] - if getattr(_asb, "_resolve_attention_arch", None) is not None: - targets.append(_asb) - return _MultiPatch( - [ - mock.patch.object(m, "_resolve_attention_arch", return_value=arch) - for m in targets - ] - ) - - -class _MultiPatch: - """Enter/exit a list of ``mock.patch`` context managers as one.""" - - def __init__(self, patches): - self._patches = patches - - def __enter__(self): - for p in self._patches: - p.start() - return self - - def __exit__(self, *exc): - for p in reversed(self._patches): - p.stop() - return False + return mock.patch.object(_au, "_resolve_attention_arch", return_value=arch) # --------------------------------------------------------------------- diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py index 7cc36cb8fb9a..239146d9be1b 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_attn_bf16_d128_ring.py @@ -32,7 +32,6 @@ from __future__ import annotations import importlib.util -import os from pathlib import Path import pytest @@ -190,7 +189,7 @@ def _gpu_ready(): @pytest.mark.parametrize("dtype,tol", [("bf16", 5e-2), ("fp16", 5e-2)]) @pytest.mark.parametrize("hq,hk", [(32, 8), (16, 16)]) @pytest.mark.parametrize("sq", [512, 1024, 4096]) -def test_d128_numeric_vs_fp32_oracle_at_magnitude(dtype, tol, hq, hk, sq): +def test_d128_numeric_vs_fp32_oracle_at_magnitude(monkeypatch, dtype, tol, hq, hk, sq): """Launch the production D128 kernel at realistic (unit-variance) magnitude and assert max_abs vs the fp32 paged-attn oracle. @@ -207,12 +206,16 @@ def test_d128_numeric_vs_fp32_oracle_at_magnitude(dtype, tol, hq, hk, sq): # Force the arch resolver so the production selector picks the gfx942 path. old_arch = au._RESOLVED_ATTENTION_ARCH au._RESOLVED_ATTENTION_ARCH = "gfx942" + # monkeypatch.delenv, not os.environ.pop: a bare pop leaks the cleared + # override into every later gfx942 routing test in the same process. This + # test is GPU-gated, so the leak is invisible on a CPU host and only shows up + # in the ordering runs the regression guard is meant to cover. for var in ( "HIPDNN_GFX942_K_SLICED_RING", "HIPDNN_GFX942_FLASH_WIDE", "HIPDNN_GFX942_FLASH_MLIM", ): - os.environ.pop(var, None) + monkeypatch.delenv(var, raising=False) try: from builders.common.attention_spec_builder import _tiled_spec_from_problem 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 b38a429a6a05..a7cb813c00d4 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 @@ -151,6 +151,25 @@ def test_infeasible_budget_raises_with_lever_diagnostics(self): self.assertIn("single-K", msg) self.assertIn("T=64", msg) + def test_pin_wins_over_a_leaked_memo(self): + """Regression guard, stated deterministically instead of by ordering. + + The original failure needed two conditions an earlier test supplies: + ``attention_spec_builder`` already imported, and + ``_RESOLVED_ATTENTION_ARCH`` holding a non-gfx950 arch. Under a bound + import the builder read the memo instead of the patch and handed back a + gfx942 spec, so every assertion in this class measured the host. Both + conditions are reproduced here so the guard cannot lapse when collection + order changes. + """ + import builders.common.attention_spec_builder # noqa: F401 (pre-import) + + au._RESOLVED_ATTENTION_ARCH = "gfx942" # restored by the autouse fixture + with mock.patch.object(au, "_resolve_attention_arch", return_value="gfx950"): + spec = au._tiled_spec_from_problem(_d256_bf16_long_prefill()) + self.assertEqual(type(spec).__module__, "kernels.gfx950.attention_tiled_2d") + self.assertTrue(spec.use_k_single_buffer) + class TestD256ProductionRouting(unittest.TestCase): """Production routing (``_d256_gfx950_fast`` live, not mocked): the diff --git a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py index 5689d0c53300..d5f05964ce3c 100644 --- a/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py +++ b/dnn-providers/hip-kernel-provider/rocke/library/tests/test_selector_parity.py @@ -39,24 +39,16 @@ def _patch_arch(arch: str): - """Pin _resolve_attention_arch in every module that imports it by name.""" - patches = [mock.patch.object(_au, "_resolve_attention_arch", return_value=arch)] - if getattr(_asb, "_resolve_attention_arch", None) is not None: - patches.append( - mock.patch.object(_asb, "_resolve_attention_arch", return_value=arch) - ) - - class _Multi: - def __enter__(self): - for p in patches: - p.start() - return self - - def __exit__(self, *exc): - for p in reversed(patches): - p.stop() - - return _Multi() + """Pin ``_resolve_attention_arch`` on the module that defines it. + + That is sufficient because the spec builder reaches the resolver through its + ``attention_unified`` module handle -- a bound import there would freeze the + reference and silently leave the builder on the real device arch, so these + cases would test the host, not ``arch``. That invariant is pinned, statically + and behaviourally, by ``test_arch_binding_guard.py``; it is not re-asserted + here. + """ + return mock.patch.object(_au, "_resolve_attention_arch", return_value=arch) def _spec(problem: UnifiedAttentionProblem, arch: str = "gfx950") -> dict: 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 12bc8872a792..d52142009dd2 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 @@ -406,23 +406,32 @@ def _pinned_attention_arch(arch): """Pin the *memoized runtime device arch* the D256 fast routes gate on. Those routes consult ``_resolve_attention_arch`` -- NOT any per-case argument - -- and the spec builder holds a *bound* import of it, so both bindings are - patched wholesale (per its docstring: "Tests that monkeypatch this function - replace it wholesale"). Without this a non-gfx950 host silently selects the - FALLBACK spec and the recorded sha becomes host-dependent. + -- so a non-gfx950 host would otherwise select the FALLBACK spec and the + recorded sha would become host-dependent. Two things are pinned, and the + second is the load-bearing one: + + * the function is rebound wholesale on its defining module (per its + docstring: "Tests that monkeypatch this function replace it wholesale"); + * ``_RESOLVED_ATTENTION_ARCH``, the process-wide memo the function returns + before doing anything else, is set too. The memo is what makes the pin + independent of *how* a consumer reached the resolver: the spec builder + goes through its ``attention_unified`` module handle today, but a bound + import anywhere would freeze the function reference and see only the memo. + This harness is installed standalone (``platform/CMakeLists.txt``) and run + by ``rocke_installed_golden_test.py`` without the library test tree, so it + cannot delegate that invariant to a guard test. """ - import builders.common.attention_spec_builder as asb import kernels.common.attention_unified as au pin = lambda: arch # noqa: E731 - o_au, o_asb = au._resolve_attention_arch, asb._resolve_attention_arch + o_fn, o_memo = au._resolve_attention_arch, au._RESOLVED_ATTENTION_ARCH au._resolve_attention_arch = pin - asb._resolve_attention_arch = pin + au._RESOLVED_ATTENTION_ARCH = arch try: yield finally: - au._resolve_attention_arch = o_au - asb._resolve_attention_arch = o_asb + au._resolve_attention_arch = o_fn + au._RESOLVED_ATTENTION_ARCH = o_memo def build_attention_d256_gfx950(arch):