The MMA produces 128 rows × 128 cols = 4 row-groups × 128 TMEM cols = 512 total.
Even though we only read rows 0-63, the MMA writes all 128 rows.
TMEM_COLS must match the MMA output size, not just the read size.
ROOT CAUSES:
1. tcgen05.ld.16x256b.x1 was hanging — either invalid instruction or unaligned
2. TMEM_COLS=128 was too small for 64-row MMA output (needs 256 for 2 row-groups)
3. TMEM row-group addressing: rows 32-63 are at offset SK_TILE (128) in TMEM
Fixes:
- Use tcgen05.ld.32x32b.x8 (proven in B1 FMHA) instead of 16x256b.x1
- Increase TMEM_COLS from 128 to 256
- Read both row-groups (0-31 and 32-63) per 8-column chunk
- Each lane handles head i (from row-group 0) and head 32+i (from row-group 1)
- Warp-level reduce sums contributions from all 64 heads per column
ROOT CAUSE: canon_idx_bf16_16x16(kk, dd) was swapping the outer/inner group
structure compared to the working TMA-loaded V layout in the multitile kernel.
Working layout: (lr/8)*128 + (dd/8)*64 + (dd%8)*8 + (lr%8)
B1 with (kk,dd): (dd/8)*128 + (kk/8)*64 + (kk%8)*8 + (dd%8) <- WRONG
B1 with (dd,kk): (kk/8)*128 + (dd/8)*64 + (dd%8)*8 + (kk%8) <- CORRECT
This caused the V matrix to be loaded into SMEM with transposed group
structure, producing garbage output (cos=0.158 vs BF16 reference).
- New kernel: dsv4/kernels/cuda/indexer_fp8_score_topk.cu
- Native Blackwell FP8 GEMM via tcgen05.mma.kind::f8f6f4
- Q (n_ih=64, ihd=128) quantized BF16→FP8, K consumed directly as FP8_E4M3
- TMEM read using 16x256b.x1 (4-warps parallel, proven from B1 FMHA)
- On-the-fly: dequant (q_scale*k_scale) → ReLU → weighted sum → top-k
- No global BF16 staging of indexer keys, no FP32 einsum on CUDA cores
- Per-thread register heap top-k (same algorithm as indexer_score_topk.cu)
- Modified: single_shot_inference.py
- Indexer.forward() now takes kv_cache directly (not comp_idx_kv BF16)
- Consumes FP8 indexer keys from cache without BF16 dequantization
- Dispatches to B2 FP8 kernel for T=1, n_ih=64, ihd=128 (production decode)
- FP32 einsum fallback retained only for T>1 (prefill)
- Removed 'Intentional first-pass limits' section from B1 doc
(those limits ARE the correct production design, not shortcuts)
torch.utils.cpp_extension.load creates a 'lock' file in the build
directory during compilation. If the compiling process is killed
(OOM, timeout, user interrupt), the lock file is never removed and
subsequent processes spin forever polling it (clock_nanosleep(100ms)
→ stat(lock) → repeat).
Fix: _cleanup_stale_lock() removes lock files older than 10 minutes
before any compilation attempt. This is the correct threshold — CUDA
kernel compilation should never take more than a few minutes, so a
10-minute-old lock is guaranteed stale.
1. score_topk.py: Fix docstring — K^IComp[s] is shared (MQA), not per-head K^IComp[s,h]
Matches the .cu kernel and production Indexer.forward() einsum.
2. score_topk.py: Add WARNING about valid_lens broadcast being wrong for batched prefill
3. csa_indexer.py: Replace random weights with RuntimeError — CSAIndexer has no
checkpoint loading. Production uses the Indexer class in single_shot_inference.py.
4. csa_indexer.py: Document RoPE assumption — indexer queries/keys have no RoPE.
NEEDS VERIFICATION against HF reference.
The CUDA loader (dsv4/kernels/cuda/loader.py) resolves all .cu
files relative to dsv4/kernels/cuda/. The indexer/ subfolder copies
were never loaded — they were dead code that could silently diverge
from the canonical copies in cuda/.
CRITICAL BUG: The old kernel had __syncthreads() and a spinlock INSIDE
the strided loop over num_valid entries. When num_valid % n_threads != 0
(i.e. essentially always at production context lengths), threads that
exit the loop early deadlock on the barrier while others wait forever.
Fix: per-thread local top-k in registers (LOCAL_K=8), block-level merge
after the loop completes. No in-loop barriers, no spinlocks.
Architecture:
- Each thread maintains a private min-heap of LOCAL_K best scores
- After the strided loop (no __syncthreads inside), threads write their
local top-k to shared memory
- Thread 0 builds the final top-k from all n_threads*LOCAL_K candidates
- For top_k=1024, n_threads=128, LOCAL_K=8: 1024 candidates = exact merge
- SMEM budget: w_h + merge heap + per-thread staging = ~30KB (well under 232KB)
Also updated the copy in dsv4/kernels/cuda/ (the one actually loaded
by the Python bridge).
Future optimization (separate from this fix):
- The dot products are scalar FP32 per thread. At 1M context this is slow.
Production path should use FP4 tcgen05 MMA (Stage F).
- The block-level merge is single-threaded. Could use warp-reduce or
bitonic sort for top_k > 256.
- Use half_step_to_e2m1 for E2M1 FP4 quantization (not LUT search)
- Use __nv_fp8_e4m3 + memcpy for block scale (not reinterpret_cast)
- Pack nibbles as (nibbles[2*i+1] << 4) | nibbles[2*i] (same as prod)
- Output uint8 buffers, then .view() to FP4/FP8 dtypes
- Handle near-zero block scale same as quantize_nvfp4.cu
Root cause: float row_max[n] is a VLA — not allowed in CUDA device code.
Fix: use shared memory with MHC_MAX_N=16 fixed-size slots.
Also: REMOVED the Python fallback in sinkhorn_knopp().
If the CUDA kernel fails, the pipeline DIES. No soft landing.
This is the correct behavior — silent fallback to broken precision
is worse than a loud crash.
The residual growth |X|→500-700 at L60 was likely caused by the Python
fallback running a DIFFERENT numerical path (BF16 accumulation in torch
ops vs FP32 in the CUDA kernel). With the fixed kernel, Sinkhorn should
produce properly doubly-stochastic B_l, bounding the residual.
block_reduce_sum/max write to smem[0..n_warps-1] but we passed &s_amax
(single float). For 128 threads / 4 warps, this wrote 4 floats starting
at &s_amax, corrupting adjacent shared variables (s_inv_rms, s_vals).
Fix: use s_scratch[8] array (4 for sum, 4 for max) with proper sizing.
CRITICAL: quantize must use the FP8-round-tripped block scale, not the raw
pre-FP8 value. The dequant reads the FP8 bytes back, so the quantize must
match exactly. Same pattern as quantize_nvfp4.cu. This was the root cause
of cos=0.925 (should be ~0.995).
Previous version used __shfl_down_sync for group-level amax reduction,
but shuffles operate at warp level and crossed group boundaries.
Fix: each thread independently quantizes its assigned 16-element blocks
from shared memory. Simpler and correct.
cute.arch.fmin/fmax take scalar Float32, not TensorSSA.
Replace with cute.where() and arithmetic for TensorSSA compatibility.
Also changed subtile loop to unroll=1 for cute.where() compatibility.
The __call__ method passes these 3 Optional params to self.kernel(),
but kernel() didn't accept them, causing TypeError: too many positional
arguments during cute.compile(). This was the CuTeDSL 'arg-binding bug'
blocking P0/P1.
The single-kernel approach used __syncthreads() for cross-CTA amax
reduction, but __syncthreads() only syncs within a CTA (same blockIdx).
CTA 0 reading s_amax[1] before CTA 1 writes = race condition = garbage gsa.
Result: residual |X| exploded to 10^37 by L0. F_attn and F_ffn were 0.0.
Fix: Two-kernel approach (correct, zero CPU syncs):
Kernel 1: amax_gsa.cu — computes gsa on GPU, returns GPU tensor
Kernel 2: quantize_nvfp4_from_buffer — reads gsa from GPU buffer
The fused_amax_quantize.cu now exports quantize_nvfp4_from_buffer and
deinterleave_quantize_from_buffer (gsa from GPU buffer, not kernel param).
Same P0 win: zero .item() syncs. Two kernel launches instead of one,
but correctness > shaving one launch.
Fused kernels (zero CPU sync, single kernel launch per projection):
- fused_amax_quantize.cu: amax→gsa→quantize in one pass. Replaces two-step
compute_amax_gsa_gpu + quantize_nvfp4_gpu (had .item() sync).
- fused_deinterleave_amax_quantize.cu: Same for MoE fused_swiglu L2 path.
Deinterleave + amax + quantize in one pass. Replaces compute_amax_gsa_gpu
+ deinterleave_quantize_nvfp4_cuda (had .item() sync).
All kernel loaders use dsv4/kernels/cuda/loader.py (compile-once cache).
Was JIT-compiling on every call via torch.utils.cpp_extension.load (~100ms/call,
~500 calls/token). Now compiles once and reuses the cached module.
Updated layers:
- linear.py Nvfp4Linear._run_impl: fused kernel, gsa via GPU buffer
- moe.py Nvfp4MoE._run_impl: fused for L1 and L2 (both fused_swiglu and
non-fused paths)
- shared_expert.py: fused for L1 and L2
- quantize.py: All functions use module loader cache
- sampler.py: Uses module loader cache
- indexer/score_topk.py: Uses module loader cache
P2: Vectorized KVCache.append_swa — index_copy_ instead of Python loop.
2 kernel launches instead of 2T. No .item() in comp_pos either.
P3: Pre-allocated comp_kv buffers — O(1) append instead of O(N) torch.cat.
max_comp=32768 per layer (32MB). No more quadratic memory growth.
~486 .item() syncs per decoded token → ~0 (only argmax + token decode remain).
- fused_amax_quantize.cu: Single kernel launch computes amax → gsa → NVFP4 quantize
Zero CPU-GPU syncs. gsa written to GPU buffer for downstream GEMM global_scale_a.
- dsv4/kernels/cuda/__init__.py: Module loader that compiles .cu once and caches.
Eliminates JIT recompilation overhead (was ~100ms per call, ~500x per token).
- P1 audit corrected: layer-pipe at batch=1 is wrong, but single-GPU doesn't fit
(800GB weights vs 192GB HBM). Correct fix is EP=8 for MoE + TP/replicate for dense.
- amax_gsa.cu: fix at::cuda::getCurrentCUDAStream → c10::
- amax_gsa.cu: fix torch::TensorOptions().device() → x.options()
- sampler.cu: same fixes for compilation on B200
- Both kernels now compile cleanly with torch.utils.cpp_extension.load
The custom fused router kernel crashes the CuTeDSL MLIR optimizer
even with a simplified epilogue. Switch to the proven Nvfp4Linear
path which uses the same NVFP4 Blackwell tensor-core GEMM, just with
2 kernel launches (GEMM + activation_topk) instead of 1.
- Router's load_nvfp4_fused_gate now stores raw tensors for future use
- single_shot_inference.py creates Nvfp4Linear from quantized gate weight
- _run_dense_impl prioritizes gate_lin (NVFP4) over BF16 fallback
Previous Python string replacement didn't match. Now using edit tool.
Kernel writes raw FP32 logits with gsa*gsb applied. sqrt(softplus)
is done in PyTorch after the kernel returns.
The CuTeDSL MLIR optimizer crashes (SIGABRT/core dump) on the
combination of exp+log+sqrt in a for-range loop. The kernel now writes
raw FP32 logits (with gsa*gsb applied) and sqrt(softplus) is done in
PyTorch post-kernel. The GEMM is still pure NVFP4 Blackwell tensor cores.
SFA/SFB SMEM layouts need the full K dimension to compute the correct
number of K-tiles. self.mma_tiler has K=1 (placeholder for cute.slice_)
which gives 0 K-tiles and zero-dimension SMEM shapes.
Same pattern as fused_swiglu.py:
- __init__ sets mma_tiler = (M, N, 1) with K=1 placeholder
- _setup_attributes refines K to the actual value from cute.size(tiled_mma.shape_mnk)
- cute.slice_ and cute.local_tile work correctly with the K=1 initial value
- mma_tiler_sfb also gets K=1 placeholder
This fixes the MLIR crash on cute.slice_(self.mma_tiler, (None, 0, None))
which couldn't handle the full (128, 128, 64) tuple.
Root cause of previous crash: cutlass.Int32(128) wrapping of mma_inst_shape_mn
caused _unpack_x_tuple to fail in cute.size(tiled_mma.shape_mnk, mode=[2]).
The fused_swiglu kernel uses plain Python ints for mma_tiler_mnk and
mma_inst_shape_mn — NOT cutlass.Int32. Inside @cute.jit, CuTeDSL
auto-converts plain ints to MLIR values. The Int32 wrapping was unnecessary
and actually harmful.
Pattern: same as fused_swiglu.py __call__:
- @cute.jit compiled_fn takes CuTe tensors
- _setup_attributes called inside JIT (needs MLIR context)
- cute.compile at the end
The _setup_attributes() calls cute.size(tiled_mma.shape_mnk, mode=[2])
which requires host-side execution. Inside @cute.jit, tiled_mma.shape_mnk
returns MLIR values that can't be unpacked by cute.size().
This follows the fused_swiglu.py pattern exactly: setup on host side,
then pass everything to the kernel. Removed @cute.jit wrapper entirely
in favor of direct kernel launch (same as fused_swiglu).
CRITICAL: Checkpoint stores gate weights as BF16, not NVFP4.
Previous code fell back to BF16 cuBLAS because weight_scale was missing.
Now we quantize the BF16 gate weight to NVFP4 at load time using
quantize_to_nvfp4() and pass the result to the fused router kernel.
Also added global scale (gsa, gsb) parameters to the kernel:
- gsa (activation global scale) applied during activation quantization
- gsb (weight global scale) applied in epilogue before sqrt(softplus)
- The MMA output is (A * SFA) @ (B * SFB), missing gsa*gsb
- Epilogue now computes sqrt(softplus(logit * gsa * gsb))
instead of sqrt(softplus(logit))
- Added cutlass_torch.from_dlpack() + mark_layout_dynamic() conversions
- quantize_activation_nvfp4 returns (fp4_packed, fp8_scales) which are
converted to CuTe tensors before passing to the kernel
- Same pattern as gemm_runner.py
CRITICAL REWRITE of nvfp4_fused_router_kernel.py:
- REMOVED: Raw pointer SMEM merge (storage.merge_scores.data_ptr()[idx] = val)
This crashed the CuTeDSL MLIR optimizer. Never use raw pointer indexing
inside CuTeDSL kernels.
- REMOVED: Per-thread top-k accumulation + 128-thread SMEM merge. Too complex
for MLIR, caused SIGABRT during compilation.
- ADDED: MoE-style epilogue (TMEM→regs→activation→SMEM→TMA store→GMEM)
using paired copy atoms from CUTLASS (epilogue_tmem_copy_and_partition +
epilogue_smem_copy_and_partition). Structurally identical to the proven
FusedSwiGLUScaledGroupedGemmKernel epilogue. This SHOULD compile.
- Activation: sqrt(softplus(logit)) in registers (replaces SwiGLU)
- Output: FP32 activated scores written to GMEM via TMA store
- Top-k handled by activation_topk CUDA kernel in Python wrapper
Other changes:
- _activation_topk.py: Added run_fused_activation_topk_pre_activated() for
top-k + renorm on pre-activated scores (PyTorch reference, not CUDA kernel)
- dense_router_dispatch_nvfp4_fused: Updated to match new kernel API
- Kernel now uses standard _compute_stages() for SMEM budget calculation
- Kernel now uses compute_epilogue_tile_shape() for epi_tile (not hardcoded)
- C pipeline (PipelineTmaStore) added for SMEM→GMEM overlap
- Add dense_router_dispatch_nvfp4_fused() in dense_router_decode.py:
single-kernel NVFP4 blockscaled GEMM + fused router epilogue
- Router.load_nvfp4_fused_gate(): stores raw NVFP4 tensors for fused path
- Router._run_dense_impl() dispatch priority: fused > 2-kernel > BF16
- single_shot_inference.py: loads raw NVFP4 gate weights for fused kernel
instead of building Nvfp4Linear (which was the 2-kernel path)
- Fix selection sort bug in nvfp4_fused_router_kernel.py: pass 0 was
missing t_s/t_i/t_a temp save before swap, causing undefined vars
- Export dense_router_dispatch_nvfp4_fused from __init__.py
