- F.linear(x, W) computes x @ W.T which caused shape mismatch when
W_gate was pre-transposed to [E, H]
- Use torch.matmul(x, W_gate) instead — computes x @ W directly, no
transpose needed, no FP32 conversion, fully graph-capturable
- W_gate stays as [H, E] (original checkpoint shape)
- Run GEMM in BF16 (not FP32) during graph capture — Blackwell tensor cores
handle BF16 natively; FP32 GEMM triggers cudaErrorStreamCaptureUnsupported
- Pre-transpose W_gate to [E, H] at load time — avoids .T view during capture
- Convert only logits output to FP32 for sqrt(softplus) numerical stability
- This fixes the graph capture failure at layer 0 Graph B
hidden_states.float() and gate_bf16.T.float() create new FP32 tensors
during CUDA graph capture, which is not graph-capturable.
Fix: run the linear in BF16 (Blackwell tensor cores handle BF16 natively),
then convert only the output logits to FP32 for numerical stability
in sqrt(softplus). The single logits.float() is graph-capturable
because it's a unary op with a pre-existing output buffer.
Python view operations (reshape, transpose, permute) are not
graph-capturable — they cause cudaErrorStreamCaptureUnsupported.
Added:
- dsv4/kernels/cuda/blackwell_swizzle.cu: custom CUDA kernel for 32_4_4 swizzle
- to_blocked(): detects graph capture, uses CUDA kernel instead of Python views
- MoE _assemble_scales_cudagraph_safe: same treatment
- Shared expert _assemble_scales_single_group: same treatment
- Linear _assemble_scales_single_group: same treatment
- Pre-allocated swizzled output buffers for all layers (avoids torch.empty_like)
The CUDA kernel writes to a pre-allocated buffer — no per-step allocations.
Eager path unchanged (still uses fast Python view operations).
Every run_nvfp4_grouped_gemm call must pass out= with a pre-allocated
buffer. During CUDA graph capture, torch.zeros() allocations are
forbidden — they cause 'cudaErrorStreamCaptureUnsupported' errors.
Added:
- shared_expert: _l2_out_buf for L2 GEMM
- shared_expert: pass out= for both L1 and L2 GEMM calls
- moe: _l2_out_buf for L2 GEMM
- moe: pass out= for unfused L1 GEMM (fused L1 already had it)
- moe: pass out= for L2 GEMM
- linear: _gemm_out_buf for all GEMM calls
- linear: pass out= for both run() and run_from_quantized() paths
grouped_linear already had _output_buf_padded — no changes needed.
The L1 GEMM produces gate+up combined output with 2*intermediate_size
BF16 columns, but _l1_out_buf was only allocated with intermediate_size
columns. The GEMM wrote past the buffer boundary, corrupting GPU memory
and causing cudaErrorInvalidValue on subsequent operations.
This was the root cause of ALL the cudaErrorInvalidValue errors in the
shared expert and MoE L2 paths — the corrupted memory from the L1 buffer
overflow propagated downstream.
Fix: _l1_out_buf shape (max_rows, 2*intermediate_size) instead of
(max_rows, intermediate_size). Applied to both shared_expert.py and moe.py.
Also removed all DEBUG sync/print statements from quantize.py and
shared_expert.py — the bug was not in the quantize kernels, it was
the buffer overflow.
The pattern causes
cudaErrorInvalidValue when gsa_gpu is a non-contiguous expanded view
(e.g., shape (9,) from quantize_nvfp4_gpu_fused during prefill with M>1).
Root cause: copy_() from an expanded/reshaped view can fail when the
source tensor has non-standard strides. The expand() operation creates
a view with stride-0 dimensions that copy_() may not handle correctly
on all CUDA versions.
Fix: Replace all gsa copy_ patterns with scalar assignment:
self._gsa_buf[0] = gsa_gpu[0] # scalar GPU→GPU, graph-capturable
This is simpler, avoids view issues, and is CUDA-graph-compatible.
Applied to: shared_expert.py, moe.py, linear.py, grouped_linear.py
- _output_buf_padded: (max_tokens * n_groups, o_lora_rank) — matches GEMM output
- Extraction: groups are stacked vertically, not horizontally
- Each group's output is (padded_rows, o_lora_rank) with o_lora_rank columns
- gemm_runner.py: Add out= parameter to run_nvfp4_grouped_gemm and
run_fused_swiglu_grouped_gemm to accept pre-allocated output buffers
- quantize.py: Replace torch.zeros_like/torch.zeros with scalar 0.0 in
torch.where() calls (graph-capturable, no memory allocation)
- Both fixes prevent 'Disallowed operation during CUDA stream capture'
errors during graph capture
Patch torch.cuda.current_device to return the tensor's device index
during from_dlpack calls inside CUDA graph capture. This bypasses the
device check in __dlpack__ without changing the CUDA stream (which
caused 'Capture must end on the same stream' in v1) and without
triggering a cross-device copy (which caused 'Cannot copy between
CPU and CUDA tensors' in v2).
Previous fix (set_device) caused 'Capture must end on the same stream'.
New fix: wrap tensor in _DLPatchTensor during graph capture, which forces
dl_device in __dlpack__ to bypass the device check without changing the stream.
This enables CUDA graph capture on all 8 GPUs, not just cuda:0.
When capturing CUDA graphs on non-default GPUs, torch.cuda.current_device()
may not match the tensor's device. from_dlpack() checks this and fails.
Fix: set the current device to match the tensor's device before from_dlpack.
This enables graph capture on all 8 GPUs, not just cuda:0.
The stride-0 expand view for gsa_gpu caused illegal memory access
in quantize_nvfp4_from_buffer kernel. The CUDA kernel may not handle
stride-0 tensors correctly.
Fix:
- M=1 decode (graph-captured): just reshape scalar to (1,) — no alloc
- M>1 prefill (not graph-captured): expand + contiguous — allocation OK
1. scatter_add_ requires int64 indices — ensure sorted_ids is .long()
2. Fixed the SECOND torch.bincount call (line 590) — same scatter_add_ pattern
3. Both code paths now use pre-allocated _tokens_per_expert_buf
1. grouped_linear.py: Pre-allocate _scale_a_buf for swizzle
- Same fix as linear.py — avoids torch.zeros per call
- Uses correctly-sized view for pad_and_swizzle_single
2. quantize.py: Replace torch.zeros_like with scalar 0.0
- torch.zeros_like allocates a full tensor every call
- torch.where(cond, 0.0, x) broadcasts scalar — no allocation
The pre-allocated buffer is max-sized, but pad_and_swizzle_single
operates on the full buffer dimensions. Fix: pass a correctly-sized
view (buf[:padded_rows, :padded_cols]) so the swizzle produces the
right output size.
Same fix applied to both linear.py and shared_expert.py.
Fixes from running Section A detector on B200:
1. single_shot_inference.py: Use pinned CPU buffers for token/position transfer
- dec_tid_buf[0] = python_int causes CPU→GPU sync
- Fixed: write to pinned CPU buffer, then copy_ (async, graph-capturable)
2. grouped_linear.py: Fix expert_offsets Python loop
- expert_offsets[g] = python_int * padded_rows → CPU→GPU sync per iteration
- Fixed: element-wise multiply with pre-allocated range tensor (GPU-only)
3. grouped_linear.py: Vectorized output extraction for T=1 decode
- Python loop z[:, g, :] = out[...] → CPU sync for each slice
- Fixed: GPU gather with pre-computed indices for T=1
4. grouped_linear.py: Pre-allocate output buffer
- torch.empty() per call → allocation inside graph
- Fixed: use self._output_buf (pre-allocated at max size)
5. grouped_linear.py: Pre-allocate expert_offsets_range_buf
- torch.arange() per call → allocation inside graph
- Fixed: compute once at init, reuse via element-wise multiply
1. mhc.py: Remove .item() from post_block (122 syncs/step eliminated)
- The X_next.abs().max().item() was syncing EVERY layer's post_block
- Diagnostics moved to caller (outside graph region)
2. linear.py: Pre-allocate _scale_a_buf in _ensure_buffer_size
- _assemble_scales_single_group now uses pre-allocated buffer
- Eliminates per-call torch.zeros() allocation (graph capture killer)
3. shared_expert.py: Same fix — use pre-allocated padded_x_sf_buf
- _assemble_scales_single_group no longer allocates
4. quantize.py: Remove .contiguous() from gsa expand
- expand() creates stride-0 view, CUDA kernel reads correctly
- No allocation on the hot path
5. Add CUDA_GRAPH_SYNC_INVENTORY.md with full violation catalog
The positional bias (ape/B) should only modulate the compression
softmax logits (Z + B), NOT be added to the KV content itself.
Paper equation: compressed = softmax(Z + B) · C
Bug was doing: compressed = softmax(Z + B) · (C + B) — poisons every
compressed KV entry with learned positional-bias content.
Fixed in both CSA (compress_csa_reduce_kernel) and HCA
(hca_compress_reduce_kernel) paths in compressor_reduce.cu.
The intermediate tensor from fused SwiGLU deinterleave is a column slice
(non-contiguous). When T>1, quantize_nvfp4_gpu_fused receives this and
the CUDA kernel crashes with 'input must be contiguous'.
Fix: add is_contiguous() check + .contiguous() in quantize_nvfp4_gpu_fused
and in SharedExpert._run_l2. This is the root cause, not a workaround —
CUDA kernels legitimately require contiguous memory.
The kernel was using head strides for the T (query row) dimension,
which happened to work for T=1 (qr=0 always) but was wrong for T>1.
For (B,H,T,NOPE) layout:
- Head stride = T*NOPE, but T stride = NOPE
- Scale head stride = T, but T stride = 1
- RoPE head stride = T*ROPE, but T stride = ROPE
Added q_nope_t_stride, q_scale_t_stride, q_rope_t_stride to params
struct, C API, and Python wrapper.
Replace complex n_sub-iterating read with the same HD/8 iteration
pattern as the proven decode kernel. Extract from lane qr%32 instead
of always lane 0. For qr>=32, use warp 1; for qr>=64, add TMEM offset.
This should fix the row 1 accuracy issue (was cos=0.94 vs decode).
prefill_read_qk_rows was reading from address 0 (sg_off + n * 8)
instead of tb + sg_off + n * 8. This caused garbage QK values,
explaining the 0.928 cosine for T=1 and NaN for T>1.
Two critical fixes:
1. prefill_read_pv_all_subs: was missing 'tb' base in TMEM read address
2. PV MMA ACCUMULATE: use pv_kt == 0 (not kv_tile==0 && pv_kt==0 && n_sub==0)
so each query row's PV starts fresh instead of accumulating into previous row's result
Critical bug: prefill_read_pv_row only read n_sub=0 (16 out of 512 HD dims).
Replaced with prefill_read_pv_all_subs that iterates over all 32 n_sub groups.
Also fixed TMEM row-group/warp mapping for rows 32-127.
