[jit] Rebuild a module when its prebuilt .so lacks the running GPU arch (fixes gfx1201 crash)

[carlushuang]

Summary

Make the JIT loader rebuild a module for the running GPU when its prebuilt .so doesn't cover that arch, instead of importing a wrong-arch prebuilt that then faults at kernel launch.

Problem

A prebuilt module .so is a valid host Python extension on any GPU. So get_module_custom_op's importlib.import_module succeeds even when the .so was built for a different arch (e.g. a gfx942;gfx950 wheel imported on gfx1201). The missing device code object then only surfaces later, at kernel launch, as an uncatchable HIP fault (SIGSEGV) — never as the ModuleNotFoundError that would make compile_ops JIT-rebuild. Net: a CDNA-only prebuilt silently shadows the JIT and crashes instead of rebuilding for the running arch.

This is the root cause behind the gfx1201 (RDNA4) KV-cache crash, and it generalizes to any wheel whose prebuilt set omits the user's GPU.

Fix

Add an arch-coverage check before importing (a stale extension can't be reloaded once loaded into the process):

run_arch = get_gfx_runtime()
archs = _so_offload_archs(os.path.join(get_user_jit_dir(), f"{md_name}.so"))  # parse clang offload targets
if archs and run_arch not in archs:
    raise ModuleNotFoundError(md_name)   # -> compile_ops rebuilds for the native arch

• _so_offload_archs parses the .so's clang offload-bundle entry ids (the gfxNNNN in …amdhsa--gfxNNNN).
• Host-only modules (no offload targets) and arch-undetectable environments fall through to a normal import — no behavior change.
• When triggered, build_module overwrites the .so with a native-arch build and the subsequent import loads it.

Any reachable module now auto-JITs for the running GPU instead of faulting.

Validation (RX 9070 XT, gfx1201)

Unit (op_tests/test_jit_arch_guard.py): offload-target parser — multi-arch, host-only, missing-file.

End-to-end, the real shipped-prebuilt scenario: restored a CDNA-only module_cache.so (gfx942;gfx950, no gfx1201), then ran ATOM + aiter on gfx1201. The loader logged

[module_cache] prebuilt .so targets ['gfx942', 'gfx950'] but not the running arch gfx1201; rebuilding for gfx1201

rebuilt in place, and the server booted clean (no SIGSEGV):

Model	gsm8k strict/flex (n=200)	TPOT (549/256)	tok/s
Qwen3-8B-FP8	0.885 / 0.895	18.64 ms	52.8
Ministral-3-8B	0.775 / 0.775	21.74 ms	45.5

Matches prior runs exactly — correctness and perf unchanged; the only difference is it no longer crashes.

Relationship to other PRs

• Supersedes the need for [gfx1201] Dedicated CK-free reshape_and_cache (paged KV write) for RDNA4 #3837 (dedicated CK-free reshape module) and for [triton] kv_cache.reshape_and_cache (HND default, NHD option) + HIP→triton fallback in cache.reshape_and_cache #3319's _hip_cache_supported() allowlist — with this guard, gfx1201 just JITs module_cache. I'll close [gfx1201] Dedicated CK-free reshape_and_cache (paged KV write) for RDNA4 #3837 if this lands.
• [triton] kv_cache.reshape_and_cache (HND default, NHD option) + HIP→triton fallback in cache.reshape_and_cache #3319's triton fallback is still useful for installs with no hipcc/CK toolchain at runtime (can't JIT at all) — there triton is the zero-build escape hatch. This guard is for from-source / toolchain-available environments.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 3839 --add-label <label>

[carlushuang]

Closing — not needed for ATOM ROCm/ATOM#811.

The simplest resolution for gfx1201 is to build aiter with gfx1201 in GPU_ARCHS so the prebuilt module_cache ships a gfx1201 code object. A native from-source build on gfx1201 (GPU_ARCHS=native) does this automatically; for a multi-arch wheel, add gfx1201 to the prebuild GPU_ARCHS. The KV-cache kernels are arch-clean and run correctly once built for the arch (validated bitwise + end-to-end gsm8k/perf on RX 9070 XT).

This loader arch-coverage guard is still a reasonable general robustness improvement — auto-rebuild (or a clear error) when an installed prebuilt wheel lacks the running GPU's arch, instead of an uncatchable SIGSEGV at launch. The underlying behavior is tracked in #3846, which stays open. If maintainers want that auto-rebuild behavior, this can be revived; it's just not a dependency for #811.