Repository navigation
fix(rocke): reach _resolve_attention_arch through the module handle - #10508
Conversation
`builders/common/attention_spec_builder.py` bound-imported `_resolve_attention_arch` from `kernels.common.attention_unified`, which froze the reference at import time and silently defeated `mock.patch.object(au, "_resolve_attention_arch", ...)`. With the patch bypassed, `_tiled_spec_from_problem` 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. The failure only triggered when this module was already in `sys.modules` before the patch was applied, so it presented as flaky test ordering. Route both call sites through the existing `_kau` module handle -- the same handle, and the same reason, already documented one line above for `_d256_gfx950_fast`. Import plumbing only; no emission change. Fixes 5 failures in `library/tests/test_lds_budget_resolver.py`: before: 5 failed, 499 passed, 34 skipped, 208 subtests passed after: 504 passed, 34 skipped, 208 subtests passed Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
project: hip-kernel-provider
rocke-lib: attention
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
|
🎉 All checks passed! This PR is ready for review. |
AnaghaRaoAMD
changed the title
AnaghaRaoAMD
changed the title
aosewski
left a comment
aosewski
left a comment
There was a problem hiding this comment.
The mechanism here is right, and it reproduces cleanly: pre-import the builder, set au._RESOLVED_ATTENTION_ARCH = "gfx942", then run test_lds_budget_resolver.py — 3 failures + 2 errors before this change, 9/9 green after. Your point about the memo also holds up under a two-factor check: with a cold memo on a real gfx942 the 5 failures still appear, and with a polluted memo on a non-gfx942 host they don't. So the bound import is the whole mechanism and AICK-1898's title is indeed misnamed.
A few things worth a second look before this goes in.
The verification didn't cover the lane that this change affects. library/kernels/ having no C++ mirror isn't quite the relevant axis — library/tests/test_attention_ir_cpp_parity.py is the dual-engine gate for these kernels (_FAMILIES = ("attention_dense", "attention_d256")), and the golden-SHA gate lives in platform/tests, which pytest library/tests never collects. On a box without rocke_engine built, the parity test skips, so that run produced no evidence either way. Worth re-running pytest platform/tests/test_rocke_ci_static.py and the parity gate with the engine built. Good news: the change is genuinely IR-neutral — I re-lowered both attention_d256 cases against the committed golden and got 6/6 SHA matches across llvm20/22/23, so no re-bless is needed.
test_attention_ir_cpp_parity.py:68 masks failures. The pytest.skip sits inside the per-case loop, so the first case with an unavailable C++ lowering aborts the whole gate. The 18 attention_dense cases are ordered ahead of the 2 attention_d256 ones, so anything wrong with d256 is unreachable in practice. Collecting unavailability per case and skipping only if all were unavailable would make it honest.
The test_attention_builds.py failures are still there. Two of them (TestAttentionHarnessTimers) fail identically on the base commit and here: builders.gfx950.attention.parity_unified_attention is missing the .prefill segment, so find_spec returns None. Same lead the ticket flagged. They're not this PR's doing, but they're AICK-1898 deliverable 3 and they keep the lane red — worth either the one-word path fix or a linked ticket.
Minor: the description carries internal Jira URLs and the ticket ID. platform/AGENTS.md asks for internal tracker links to stay out of public artifacts; the commit message itself is clean, so it's just the PR body. The GitHub refs are fine to keep.
The harness read `asb._resolve_attention_arch` to pin the arch the D256 fast routes gate on. #10508 removed that bound import from the spec builder, so `_pinned_attention_arch` now raises AttributeError and both `attention_d256` cases drop out of the golden set -- `rocke_golden_static` goes red with 12 drift entries. Drop the `asb` half and additionally pin `_RESOLVED_ATTENTION_ARCH`, the process-wide memo `_resolve_attention_arch` returns before doing anything else. Pinning the memo makes the pin independent of *how* a consumer reaches the resolver, which matters here specifically: the 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. No re-bless: IR parity OK [llvm20, llvm22, llvm23]. Follow-up to #10508. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
`attention_unified._RESOLVED_ATTENTION_ARCH` memoizes the real device arch, and a dozen tests pin it directly. One that leaves it set changes what every later test resolves, which shows up as order-dependent failures on a GPU runner and not at all on a CPU one. An autouse fixture restores it and clears the derived `_2D_LAUNCH_META` memo, making the leak class impossible instead of relying on each test's own teardown. The module is looked up in sys.modules rather than imported, so a test with nothing to do with attention does not pay for the import -- and an import regression in `attention_unified` stays an attention failure instead of erroring every test in the tree during fixture setup. Env overrides are deliberately left out of the fixture: a blanket snapshot would mask an env leak rather than fix it. The one real leak, in the GPU-gated D128 ring test, is fixed at its source by switching `os.environ.pop` to `monkeypatch.delenv`. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Nothing in the language stops the spec builder from re-acquiring a bound import of `_resolve_attention_arch`, and the regression is silent: the builder falls back to the real device arch and the test measures the host. Assert it once, statically (`hasattr`) and behaviourally (two archs across both the 2D and 3D builder entry points, each of which holds its own call site). Two archs, not one, because the resolver falls back to "gfx950" with no device and `_tiled_2d_impl` maps every arch except gfx942/gfx1250 onto the gfx950 impl. Which pin catches the regression therefore depends on the host: on a gfx90a box the gfx942 pin fails, on a gfx942 box the gfx950 pin fails. Neither alone is a guard; the pair is. With the invariant centralized, the per-test `getattr`-guarded dual-module patch scaffolding in `test_selector_parity` and `test_attention_builds` is dead weight; drop it and point the docstrings at the guard. Verified to fail on the pre-#10508 builder. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The original breakage needed two conditions an earlier test happened to supply: `attention_spec_builder` already in sys.modules, and `_RESOLVED_ATTENTION_ARCH` holding a non-gfx950 arch. Reproducing both in-test means the guard cannot lapse when collection order changes. Verified to fail on the pre-#10508 builder. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
I've pushed four commits to this branch (
The first is the blocking one — it clears the red Your diagnosis reproduces exactly on real hardware. On a gfx942 MI300X node at the base commit: the resolver test alone gives 9 passed, and after One thing worth knowing about the guard test: it pins two archs deliberately. The resolver falls back to Two further things I found are not yours and are out in #10625 instead of here, to keep this PR reviewable: the parity gate's skip-on-first-case behaviour (which is why the harness breakage went unnoticed in the first place), and a stale Verified before pushing: |
AnaghaRaoAMD
changed the title
AnaghaRaoAMD
changed the title
…ess module path (#10625) Two independent pre-existing fixes in the rocKE attention test lane. Neither depends on the other; they are together only because both surfaced while reviewing #10508, and both are separable from it. ## 1. The dual-engine parity gate could not see most of its cases `library/tests/test_attention_ir_cpp_parity.py` is the C++/Python byte-identity gate for the `attention_dense` and `attention_d256` kernels. Its per-case handler called `pytest.skip` on **any** exception from the C++ lowering. Since `pytest.skip` aborts the whole test, the first case with an unavailable engine ended the run — and the 18 `attention_dense` cases sort ahead of the 2 `attention_d256` ones, so d256 was effectively unreachable on any box without `rocke_engine` built. That is not hypothetical. It is exactly why a breakage in the `attention_d256` harness builds reached CI unflagged: the gate that should have caught it skipped first. Failures are now classified by exception type instead of lumped together: | Outcome | Verdict | |---|---| | `BackendCoverageGap` — an arch in `backend.CPP_UNPORTED_ARCHES`, the one sanctioned gap taxonomy | collect and continue, report as skip | | `BackendError` — the engine is not importable at all | whole-run skip; an absent engine cannot appear mid-loop | | anything else | propagates — the engine is present and will not lower a case, which is a defect | `assert cases` is also tightened into a per-family presence check. A renamed or dropped family would still leave a non-empty list, and harness drift must not be reported as an engine gap. ## 2. Stale `parity_unified_attention` module path The harness moved under `attention/prefill/`, but three call sites still resolve `builders.gfx950.attention.parity_unified_attention`. Two are hard failures on `develop` today: - `test_attention_builds.py::TestAttentionHarnessTimers` (2 tests) — `find_spec` returns `None`, so `.origin` raises `AttributeError` - `test_gfx950_smoke_attention.py` — `python -m` on a module that does not exist The other two are stale copy-paste targets in docstrings. ## Verification On `develop` tip (`1781f84e4b`), CPU lane: | | before | after | |---|---|---| | `pytest library/tests` | 2 failed, 502 passed | **504 passed**, 0 failed | | `pytest platform/tests` | 685 passed | 685 passed | `black` 25.12.0 clean. No new `ruff` findings. No golden re-bless — this touches no emission path. ## Scope Test-lane only. No kernel, spec, lever, or emission change; no C++ mirror needed. A related follow-up covers ~15 remaining references to a dead `examples/gfx950/attention/parity_unified_attention.py` documentation path across 8 doc files — kept out of here to stay reviewable. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
ISSUE ID : AICK-1898
Problem
attention_spec_builder.pybound-imported_resolve_attention_arch. Afrom ... import symfreezes the reference at import time, somock.patch.object(au, "_resolve_attention_arch", ...)never reached the builder. It then resolved the real device arch,_tiled_2d_implreturned that arch's spec class, and the gfx950-only override fields hit:Only fires when the builder is already in
sys.modulesbefore the patch, so it looked like flaky ordering. Pre-existing on develop; introduced by #9233, which added the_kaumodule handle and a comment about this exact hazard but applied it to_d256_gfx950_fastonly.Fix
Route both call sites through the existing
_kaumodule handle. Import plumbing only — no emission or spec change.Four follow-up commits from @aosewski (
756dc0e3..55c3ecb4):818c75frepairsrocke_ir_parity_harness.py, which read the removed attribute and tookrocke_golden_staticred with 12 drift entries;adfc04acontains_RESOLVED_ATTENTION_ARCHleaks via a conftest fixture;7c47e74addstest_arch_binding_guard.py(verified to fail on the unfixed builder);55c3ecbmakes the resolver test reproduce both preconditions in-test.Verification
pytest library/tests -m "not gpu"on gfx942:223bff09f9055c3ecbAlso at head:
platform/tests668 passed, goldentest_rocke_ci_static.py5 passed (IR parity OK [llvm20, llvm22, llvm23], no re-bless), and 520 passed / 0 failed on a gfx942 MI300X node.test_attention_ir_cpp_parity.pyis the dual-engine gate for these families; it skips whererocke_engineisn't built, so it did not execute in the runs above.Side effect
prove_parity_binding.pypatches the resolver at :1465 but imports_tiled_spec_from_problemat :1431, sorun_selector_parity(arch=...)previously compared a Python spec for the host arch against a C++ spec for the requested arch. Fixed as a side effect, but earlier selector-parity numbers aren't comparable and masked mismatches may surface — worth a re-run across arches.Two adjacent defects are out in #10625, not here: the parity gate's skip-on-first-case behaviour, and a stale
parity_unified_attentionpath missing its.prefillsegment.