fix(rocke): reach _resolve_attention_arch through the module handle

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":

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

@@ -15,10 +15,51 @@
 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))
 
 _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()

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()

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)
 
 
 # ---------------------------------------------------------------------

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
 

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

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:

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):