biondizzle cc6b094450 fix: root-cause JIT memory corruption myth, add eager warmup, remove _needs_token_refill ...

Bug #1 fix: The _needs_token_refill workaround was a band-aid over a
misdiagnosis. cute.compile does NOT corrupt GPU memory (verified on B200).
The original corruption was from a different bug (likely OOB write or
weight loading issue).

Changes:
- bridge.py: Add warmup_compilation() for eager JIT before runtime buffers
  exist. Pre-allocate workspace per cache entry (no torch.full in hot path).
  Cache stores {compiled, workspace, workspace_size} instead of just compiled.
  CuTe tensor wrappers re-created per call (cheap metadata, avoids stale refs).
- runner.py: Remove _needs_token_refill hack. Add eager warmup call in
  _ensure_stacked() for both L1 and L2 GEMM shapes.
- nvfp4_linear.py: Add eager warmup in finalize_weights() for single GEMM.

The warmup approach ensures cute.compile runs exactly once per shape during
model init, before any forward pass. This is deterministic and eliminates
any possible interaction between JIT and runtime GPU memory.

2026-05-20 02:08:01 +00:00

..

kernel

refactor: copy CuTeDSL kernel into repo with local imports

2026-05-16 02:57:54 +00:00

__init__.py

refactor: copy CuTeDSL kernel into repo with local imports

2026-05-16 02:57:54 +00:00

blackwell_attention.py

Add blackwell_attention module and comprehensive test

2026-05-19 15:30:29 +00:00

bridge.py

fix: root-cause JIT memory corruption myth, add eager warmup, remove _needs_token_refill

2026-05-20 02:08:01 +00:00

csa_attention.py

Fix device reference in full_attention_reference

2026-05-19 08:01:31 +00:00

custom_ops.py

Replace autograd.Function with torch.library.custom_op for Dynamo compat

2026-05-19 01:54:48 +00:00

inverse_rope.py

Replace DeepGEMM fp8_einsum with CuTeDSL NVFP4 for wo_a (o_proj)

2026-05-19 02:36:30 +00:00

moe_pipeline.py

Add pipeline test with real model weights, add swiglu_limit to reference moe_pipeline

2026-05-17 18:07:44 +00:00

native_swa_decode.py

fix: handle 3D swa_indices and correct kv_bf16 expand dims

2026-05-20 01:36:27 +00:00

nvfp4_linear.py

fix: root-cause JIT memory corruption myth, add eager warmup, remove _needs_token_refill

2026-05-20 02:08:01 +00:00

runner.py

fix: root-cause JIT memory corruption myth, add eager warmup, remove _needs_token_refill

2026-05-20 02:08:01 +00:00

shared_expert_pipeline.py

Fix garbled shared_expert_pipeline.py: imports/class were merged

2026-05-19 07:18:10 +00:00

wo_a_grouped_linear.py

Fix wo_a: scatter each group's data at correct offset in padded buffer

2026-05-19 02:45:57 +00:00