Skip to content

fix(rocke): reach _resolve_attention_arch through the module handle #10508

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
AnaghaRaoAMD merged 6 commits into from
Aug 17, 2026
Merged

fix(rocke): reach _resolve_attention_arch through the module handle #10508

Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
756dc0e
fix(rocke): reach _resolve_attention_arch through the module handle
AnaghaRaoAMD Aug 7, 2026
818c75f
fix(rocke): restore attention_d256 in the golden IR parity gate
aosewski Aug 11, 2026
adfc04a
test(rocke): contain attention arch state leaks between tests
aosewski Aug 11, 2026
7c47e74
test(rocke): pin the lazy arch-resolver binding with a guard test
aosewski Aug 11, 2026
55c3ecb
test(rocke): restate the arch-pin failure deterministically
aosewski Aug 11, 2026
b22397f
Merge branch 'develop' into users/anarao/fix-attention-arch-patch-seam
aosewski Aug 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions dnn-providers/hip-kernel-provider/rocke/library/builders/common/attention_spec_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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.
Comment thread
AnaghaRaoAMD marked this conversation as resolved.
Resolved
from kernels.common import attention_unified as _kau
Comment thread
AnaghaRaoAMD marked this conversation as resolved.
Resolved


def _tiled_spec_from_problem(
problem: UnifiedAttentionProblem,
):
arch = _resolve_attention_arch()
arch = _kau._resolve_attention_arch()
Comment thread
AnaghaRaoAMD marked this conversation as resolved.
Resolved
UnifiedAttention2DTiledSpec, _, _ = _tiled_2d_impl(arch)
if arch == "gfx1250":
return UnifiedAttention2DTiledSpec(
Expand Down Expand Up @@ -320,7 +330,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":
Expand Down
Loading