In-process JIT caches can reuse stale compiled artifacts after helper changes

Summary

FlyDSL can keep using an already compiled in-process JIT artifact after a helper function used by the kernel has changed. This can make iterative kernel development confusing: a rerun in the same Python process may appear to use the new Python code, but actually executes the old compiled artifact until the relevant compile factory cache is cleared or the process restarts.

Reproducer

I added a local reproducer script while debugging PA decode:

FLYDSL_RUNTIME_ENABLE_CACHE=0 PYTHONPATH=./python:./ python scripts/repro_pa_decode_ps_stale_cache.py

The script does not rely on editing files during the run. Instead it simulates a helper change by monkeypatching _load_q_fragments after the first compile.

Flow:

1. Run one PA decode case to compile and cache the baseline launch.
2. Monkeypatch kernels.pa_decode_fp8._load_q_fragments to a sentinel helper that always raises.
3. Rerun the same shape without clearing caches.
4. Clear compile_pa_decode_ps.cache_clear() and rerun.

Observed output excerpt:

Step 1: compile and cache the baseline PA decode launch.
[baseline] passed with max_diff=0.005249, threshold=0.005123

Step 2: replace _load_q_fragments with a sentinel that must fail if retraced.
Step 3: rerun the same shape without clearing caches.
[cached_after_helper_change] passed with max_diff=0.005249, threshold=0.005123
[cached_after_helper_change] sentinel was not observed, so the launch reused the already compiled in-process artifact.

Step 4: clear compile_pa_decode_ps.cache_clear() and rerun.
[after_compile_cache_clear] observed sentinel after retrace. This confirms the previous pass used stale compiled code.

Why this matters

During iterative kernel development, changing helper code and rerunning a test in the same process can silently exercise old compiled code. This can look like an intermittent correctness issue or make a fix appear ineffective until caches are manually cleared.

This is especially easy to hit for large kernels that use @functools.lru_cache compile factories, for example:

@functools.lru_cache(maxsize=256)
def compile_pa_decode_ps(...):
    @flyc.kernel
    def pa_decode_ps_kernel(...):
        ... _load_q_fragments(...) ...

The outer compile factory returns the previously created @flyc.jit launch wrapper for the same static signature, so helper changes are not retraced.

Expected behavior

One of the following would make this safer:

• Include relevant helper source/code identity in the in-process compile factory cache key.
• Provide a standard dev-mode cache invalidation hook that clears both FlyDSL JIT caches and user-facing compile factory caches.
• Add documentation/warnings for kernel authors that functools.lru_cache compile factories must be manually cleared when helper code changes in the same process.
• In debug/development mode, disable or validate in-process reuse against dependency source hashes.

Notes

FLYDSL_RUNTIME_ENABLE_CACHE=0 disables disk/runtime cache, but it does not disable Python-level functools.lru_cache on compile factories or already constructed JIT wrapper caches. The reproducer demonstrates stale reuse even with runtime cache disabled.

Start with scripts/repro_pa_decode_ps_stale_cache.py and the compile_pa_decode_ps factory referenced in kernels.pa_decode_fp8. Reproduce the helper monkeypatch with runtime caching disabled, then trace the Python-level lru_cache and JIT wrapper reuse. Done should prevent stale compiled code after a helper change without requiring manual cache clearing, but the issue leaves the implementation strategy open.