[Bugfix] Drop arch env vars from the JIT cache key (fix AOT cache misses)

[coderfeli]

Motivation

After #597 hardened the JIT cache key to fold whitelisted code-gen env vars into the key, AOT-precompiled-then-shipped caches stopped being hit at runtime. This surfaces downstream in aiter as RuntimeError: AOT cache miss for FlyDSL MoE kernels under its strict check_aot_cache=True tests (the cache dir is present with the right manager_key, but the inner key never matches).

Root cause

_cache_invalidating_env_values() folds these arch selectors into the key: ARCH, FLYDSL_GPU_ARCH, HSA_OVERRIDE_GFX_VERSION. Their only effect is to pick the target arch — which is already in the key via ("_target_", GPUTarget).

The arch can be supplied two ways, and detect_target() normalizes both to the same GPUTarget:

• AOT pre-compilation runs on a GPU-less build host and injects ARCH=gfx950 / FLYDSL_GPU_ARCH=gfx950 via env.
• Runtime has those env vars unset; arch comes from device detection.

So the resolved ("_target_", GPUTarget(arch="gfx950", ...)) is identical in both, but the raw ("_env_", (... "ARCH"="gfx950"/"" ...)) segment diverges ("gfx950" at build vs "" at runtime) → cache key mismatch → guaranteed miss for any shipped AOT cache.

Fix

Drop the three arch selectors from _CACHE_INVALIDATING_ENV_VARS. No key precision is lost — any arch change still changes ("_target_", ...). Genuine code-gen knobs (FLYDSL_COMPILE_OPT_LEVEL, FLYDSL_COMPILE_BACKEND, FLYDSL_COMPILE_LLVM_DIR, FLYDSL_DEBUG_ENABLE_DEBUG_INFO, FLYDSL_EXTRA_SOURCE_DIRS) stay in the key.

Tests

Existing cache-key tests are unaffected:

• test_compile_hints.py::test_different_arch_gives_different_key asserts the change flows through ("_target_", ...) (it resets _sig and checks t1.arch != t2.arch) — still holds, since ARCH still feeds detect_target().
• test_jit_cache_key_completeness.py env-drift tests use FLYDSL_COMPILE_OPT_LEVEL (retained).

Follow-up note

FLYDSL_COMPILE_LLVM_DIR / FLYDSL_EXTRA_SOURCE_DIRS are path-like and could similarly diverge build-vs-runtime for shipped AOT caches; they're kept here because their contents can affect codegen. A more robust long-term approach would key on their effect (e.g. resolved LLVM version / source content hash) rather than the raw path string. Out of scope for this fix.

🤖 Generated with Claude Code