[BUG] CUTLASS DSL
bug? - Needs TriageCuTe DSL
### Which component has the problem?
CuTe DSL
### Bug Report
**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**
```python
# 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
```
```python
# 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:
```python
# 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
```
```python
# 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.
0 条评论