[BUG][CuTe DSL] Stale source after import breaks lazy jit preprocessing silently

Describe the bug

@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 is modified after the module was imported — the typical case being 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. There is no consistency check, and the failure mode depends on what text happens to sit at the stale line number:

1. Slice has no DSL decorator (it lands on imports, or on a def without @jit): check_decorator returns False and transform_function silently returns []; run_preprocessor returns None; _preprocess_and_replace_code silently skips the code replacement (if fcn_ptr:). The kernel then executes as plain Python with no control-flow staging, and its first dynamic if raises the deeply confusing error[PHASE_DYNAMIC_TO_STATIC_BOOL]: Cannot use a Runtime value (Staged value) value where a Python value (Meta value) boolean is required (4.6.0.dev0 wording: DSLRuntimeError: Unable to convert dynamic 'Boolean' value to bool at compile time). Class methods hit this path.
2. Top-level functions: the (empty) transformed module is exec'd and the function name resolves back to the original jit wrapper from the module globals, so code replacement dies with ValueError: pick() requires a code object with 0 free vars, not 3.
3. Worst case — silent wrong-code swap: if the stale line number lands on a different same-named decorated function (think __call__ / forward / load 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, with no error raised at all.

All three modes reproduce on the latest released nvidia-cutlass-dsl==4.6.1 (and on 4.6.0.dev0), CPython 3.10 and 3.12. The repros below are GPU-independent.

Steps/Code to reproduce bug

# tiny_mod.py
import cutlass
import cutlass.cute as cute


@cute.jit
def pick(x: cutlass.Int32) -> cutlass.Int32:
    y = cutlass.Int32(0)
    if x > 0:
        y = x + 1
    else:
        y = x - 1
    return y

# driver.py — run both files from one directory: `python driver.py`
import pathlib, sys

here = pathlib.Path(__file__).parent
sys.path.insert(0, str(here))
import tiny_mod                    # decorators run here; preprocessing is LAZY

# Simulate an in-place package upgrade after import (any line shift will do):
p = here / "tiny_mod.py"
p.write_text("# shifted\n" * 5 + p.read_text())

import cutlass
import cutlass.cute as cute

cute.compile(tiny_mod.pick, cutlass.Int32(3))
# -> ValueError: pick() requires a code object with 0 free vars, not 3   (mode 2)

Comment out the write_text line and the same driver compiles fine. Put the identical function inside a class and compile the bound method instead (cute.compile(Picker().pick, cutlass.Int32(3))) and you get mode 1:

error[PHASE_DYNAMIC_TO_STATIC_BOOL]: Cannot use a Runtime value (Staged value) value
where a Python value (Meta value) boolean is required

Mode 3 (silent swap), self-checking:

# tiny_swap.py — two classes with a same-named, same-signature jit method
import cutlass
import cutlass.cute as cute


class B:
    @cute.jit
    def pick(self, x: cutlass.Int32) -> cutlass.Int32:
        y = cutlass.Int32(0)
        if x > 0:
            y = x - 100
        else:
            y = x + 100
        return y


class A:
    @cute.jit
    def pick(self, x: cutlass.Int32) -> cutlass.Int32:
        y = cutlass.Int32(0)
        if x > 0:
            y = x + 1
        else:
            y = x - 1
        return y

# driver_swap.py
import pathlib, sys

here = pathlib.Path(__file__).parent
sys.path.insert(0, str(here))
import tiny_swap

fa = tiny_swap.A.pick.__wrapped__
fb = tiny_swap.B.pick.__wrapped__
n = fa.__code__.co_firstlineno - fb.__code__.co_firstlineno
p = here / "tiny_swap.py"
p.write_text("# shifted\n" * n + p.read_text())   # A's stale lineno now = B's def

import cutlass
import cutlass.cute as cute

cute.compile(tiny_swap.A().pick, cutlass.Int32(3))   # no error raised!

def consts(code, acc):
    for c in code.co_consts:
        if isinstance(c, (int, float)):
            acc.add(c)
        elif hasattr(c, "co_consts"):
            consts(c, acc)
    return acc

assert 100 in consts(fa.__code__, set())   # A.pick now contains B.pick's body
print("silent code swap confirmed: A.pick was replaced by B.pick's code, no error")

The title is (the form pre-fills [BUG] , so you continue after it):

[BUG][CuTe DSL] Lazy AST preprocessing silently misbehaves when the source file changes after import — cryptic errors, or a silent wrong-function code swap

If that reads too long for your taste, a shorter equivalent:

[BUG][CuTe DSL] Stale source after in-place file update makes lazy jit preprocessing fail cryptically or swap in the wrong function

Expected behavior

Either preprocessing works from source captured at decoration time, or the DSL detects the inconsistency and says so. A clear error such as

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.

turns a multi-day debugging session into a one-line fix.

Environment details (please complete the following information):

• Environment location: Bare-metal (Linux x86_64) and containers; GPU-independent repro (also verified end-to-end on B200 with flash-attention SM100 kernels)
• nvidia-cutlass-dsl: 4.6.1 (pip) and 4.6.0.dev0
• CPython: 3.10.12 and 3.12.3

Additional context

How we hit this in the real world: while debugging Dao-AILab/flash-attention#2716 / #2717, a user's long-lived serverless orchestrator had imported flash-attention, the package was then reinstalled in place with a patched kernel file (a few added lines shift everything below), and the first attention call in that process failed with the dynamic-to-bool error on every dtype — while a fresh subprocess worked and the previous file worked. Nothing in the error points at stale source, so attribution took ~20 experiment rounds across two engineers. Any Python-source kernel library (flash-attention, quack, user kernels) upgraded under a live process can hit this.

We have a small validated patch for transform_function that raises the clear error above when the extracted slice does not define the expected function (name check, plus a declared-parameter/co_varnames cross-check for same-named functions; lambdas exempt), with no behavior change for consistent sources — verified against flash-attention's full SM100 forward kernels (bf16/fp8, compiles and numerics identical). Happy to open a PR.

Two honest limitations for discussion:

• A same-name and same-signature swap (mode 3 above) 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 and it false-positived on real kernels. The complete fix would be capturing the source (or a hash of it) eagerly at decoration time.
• f = cute.jit(g) (non-@ usage) silently skips preprocessing today and is unchanged by the proposed guard.

Start in base_dsl/ast_preprocessor.py at transform_function, then trace run_preprocessor and _preprocess_and_replace_code. Reproduce the stale-source cases with the tiny_mod.py/driver.py and tiny_swap.py/driver_swap.py examples; done means inconsistent source is reported clearly instead of being silently skipped, failing cryptically, or replacing code with another function.