- _can_run_minimax_h3_fusion now also requires Triton availability, so an
enabled fusion on a CUDA build without a working Triton falls back to
eager instead of hitting the strict wrappers' RuntimeError mid-forward.
- One-time logger.info of the resolved fusion set at model init (and a
warning when the set is requested without Triton), plus a
prepare_for_compile hook warning that torch.compile capture makes the
fusions inert inside compiled block forwards.
- Positive-engagement routing test: counts 1/1/2/1 fused-kernel calls in
one CUDA inference block forward and asserts a grad-enabled forward
leaves the counters unchanged, so a guard regression to always-eager
can no longer pass silently.
- Document the index-bounds contract ([0, table_rows)) on the modulation
wrappers; a device-side check would synchronize, and in-model callers
are safe by construction.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
_qknorm_partial_rope_kernel left tl.program_id(0) in int32, so
row * head_dim wrapped once the flattened input reached 2**31 elements
and the kernel read/wrote out of bounds (CUDA illegal memory access).
The PR's other two kernels (modulation.py, swiglu.py) already cast
tl.program_id(0).to(tl.int64); this one now matches, and seq_index /
table_offset inherit int64 from row.
Confirmed on GB200: (1, 8_500_000, 2, 128) bf16 (2.176e9 elements)
crashed before the cast and matches eager after it; the just-under-2**31
control shape matched all along. For H3 (56 heads x 128 head_dim) the
boundary is batch*seq >= 299_593 tokens per rank, reachable at SP=1.
Adds a GPU regression test at the over-2**31 shape that compares the
head and tail rows against eager.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
