- grouped_linear.py: Replace .item() gsa + Python quantize with
quantize_nvfp4_gpu_fused (zero CPU syncs). Flatten all groups
into (G*T, D), single fused kernel launch, GPU-only gsa copy.
- single_shot_inference.py: Reduce torch.cuda.synchronize() to
every 20 steps instead of every step. Gate per-layer diagnostics
to li<3 or li>=58 (avoid 61 .item() calls per decode step).
NVFP4 quantize_from_buffer produces CUDA error on large-magnitude
inputs (|X|>500 at L60 output). BF16 lm_head is correct and only
runs once per decode step — not a bottleneck.
TODO: debug the NVFP4 path for large activations and re-enable.
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.
Eliminates 183 kernel launches per decoded token from pointless memcpy.
Operates on rope dims in-place via views instead of cloning the full tensor
and allocating an empty_like buffer.
The single_shot is a reference for vLLM/SGLang integration. The layer-pipeline
sharding (gpu = li % NUM_GPUS) is the right pattern for this reference.
EP/TP sharding belongs in the actual vLLM integration, not here.
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
1. o_a_proj (Nvfp4GroupedLinear): Added load_nvfp4_weight() method
that loads checkpoint NVFP4 weights directly — no more dequant→BF16→requant.
Each group's weight is transposed from (N, K_packed) checkpoint layout
to (K_packed, N) layout expected by the grouped GEMM.
2. lm_head: Quantize BF16 weight to NVFP4 at load time, use production
Nvfp4Linear GEMM instead of F.linear. Runtime gsa for activation.
Frees the 1.8GB BF16 weight after quantization.
3. Hash router (L0-2): Already optimal — tid2eid is an int32 lookup,
no GEMM to accelerate.
The checkpoint-path gate was using the checkpoint's input_scale as gsa
— the same E4M3 overflow bug we fixed in Nvfp4Linear/Nvfp4MoE/etc.
The runtime-quantized BF16 path was using 1/(6*448) as a fixed gsa.
Both now compute gsa from actual activation magnitude at runtime.
The checkpoint's input_scale was designed for training-time FP8 quantization,
not NVFP4 activation quantization. Using it as gsa causes x/gsa to exceed
the E4M3 block scale maximum (448), leading to systematic magnitude loss
in every projection. This accumulates over 61 layers, compressing the
logit range and producing garbage tokens.
Fix: compute gsa at runtime from actual activation magnitude:
gsa = max(|x|) / (6.0 * 448.0)
This ensures x/gsa ≤ 2688 (the maximum representable in E4M3 block scales).
Applied to: Nvfp4Linear, Nvfp4GroupedLinear, Nvfp4MoE, Nvfp4SharedExpert, Router gate
