[CuTe DSL] Raise a clear error when jit source is stale instead of silently misbehaving, issue #3395

[yunweili3]

Problem of #3395

@cute.jit AST preprocessing is lazy: it runs at the first call/compile of each function, and at that point it re-reads the function source from disk via inspect.getsourcelines, slicing the current file at the in-memory code object's co_firstlineno (base_dsl/ast_preprocessor.py::transform_function).

If the source file was modified after the module was imported — typically a package upgraded in place while a long-lived process (serving worker, orchestrator, notebook kernel) still holds the old modules — the extracted slice is misaligned, and there is no consistency check. What happens next depends on what text sits at the stale line number:

1. Slice has no DSL decorator: check_decorator returns False, transform_function silently returns [], run_preprocessor returns None, and _preprocess_and_replace_code silently skips the code replacement. The kernel then runs as plain Python with no control-flow staging, and its first dynamic if raises the confusing error[PHASE_DYNAMIC_TO_STATIC_BOOL]: Cannot use a Runtime value (Staged value) value where a Python value (Meta value) boolean is required ("Unable to convert dynamic Boolean value to bool" in older wording). Class methods hit this path.
2. Top-level functions: the (empty) transformed module is exec'd, the name resolves back to the original jit wrapper from the module globals, and code replacement dies with ValueError: f() requires a code object with 0 free vars, not 3.
3. Worst case — silent wrong-code swap: if the stale offset lands on a different same-named decorated function (e.g. __call__ methods of sibling classes in one file), the preprocessor transforms that function, exec binds it under the bare name, and the original function's __code__ is replaced with the wrong function's body — no error at all.

None of these point at the actual cause. Found while debugging Dao-AILab/flash-attention#2716, where attribution took ~20 experiment rounds across two engineers. Repros are GPU-independent and hit nvidia-cutlass-dsl 4.6.1 and 4.6.0.dev0 alike, on CPython 3.10 and 3.12.

Fix

Guard transform_function right after parsing the extracted slice:

1. If the slice is not a FunctionDef named func_name, raise a DSLRuntimeError naming the mismatch and the likely cause, with a restart/reload suggestion:

   DSLRuntimeError: Source extracted for function `load_KV`
   (flash_fwd_sm100.py:3028) defines `load_Q_non_tma` instead -- the source file
   appears to have been modified after the module was imported.
   suggestion: Restart the process (or re-import/reload the module) so the
   in-memory code and the on-disk source agree.
   

   <lambda> is exempt (no retrievable def block; already skipped by the decorator check).

2. If the name matches, cross-check the declared parameter list against the code object (co_varnames / co_argcount / co_kwonlyargcount / vararg / kwarg flags) — a stale offset can land on a same-named method of another class with a different signature.

3. Extend the parse-failure suggestion (previously REPL-only) to also mention in-place source modification.

No behavior change for consistent sources; single-file change, +58/−1.

Validation

• Minimal repros (no GPU; CPython 3.10 + 3.12; dsl 4.6.1 + 4.6.0.dev0): top-level function, class-method, and same-named-function variants — every previously-cryptic failure above now raises the clear stale-source error; untouched files compile exactly as before.
• Real-workload regression: flash-attention SM100 forward kernels on B200 — bf16 / fp8-e4m3 / fp8-e5m2 compiles plus a repeated warm call, results bit-identical with and without this patch.
• Stale-file scenarios on the same kernels (import → replace flash_fwd_sm100.py in place → first trace): previously Unable to convert dynamic Boolean value to bool at compile time, now the clear error quoted above.

Known residual (for discussion)

A same-name and same-signature swap (mode 3 with twin methods) is not detectable from the slice alone. We prototyped recompiling the slice and comparing code objects, but CPython emits different bytecode for identical source depending on compile context (class body vs wrapped block — the method-call specialization bits in LOAD_GLOBAL/LOAD_ATTR opargs differ), so bytecode equality is not a sound invariant; it false-positived on real kernels and was dropped. The complete fix would capture the source (or a hash of it) eagerly at decoration time — happy to explore that in a follow-up if maintainers prefer. f = cute.jit(g) (non-@ usage) silently skips preprocessing today and is unchanged by this PR.