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
- Checkpoint path: load NVFP4 gate weight directly into Nvfp4Linear
- BF16 path: quantize and load into Nvfp4Linear
- Both paths use proven production GEMM (no custom kernel)
- load_nvfp4_fused_gate now creates Nvfp4Linear from BF16 weight
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
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))
- 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
The attention output projection first half (wo_a) was using BF16
grouped BMM (torch.bmm). Now uses production Nvfp4GroupedLinear
which performs the same grouped GEMM with NVFP4 tensor-core
acceleration on Blackwell.
The weight is loaded from NVFP4 checkpoint if available, otherwise
quantized from BF16 via set_bf16_weight().
Also includes:
- NVFP4 gate projection for router (from previous commit)
- Compressor position_bias in CUDA kernel (from earlier fix)
The dense router now uses NVFP4 GEMM via Nvfp4Linear for the gate
projection when NVFP4 scales are available in the checkpoint. This
replaces the BF16 cuBLAS GEMM with Blackwell SM100 tensor-core
NVFP4 acceleration.
Changes:
- dsv4/layers/router.py: add gate_lin (Nvfp4Linear) alongside W_gate
fallback. New load_nvfp4_gate() method.
- dsv4/kernels/router/dense_router_decode.py: add
dense_router_dispatch_nvfp4() using Nvfp4Linear + activation_topk
- dsv4/kernels/router/__init__.py: export new function
- single_shot_inference.py: load NVFP4 gate weights when available,
fall back to BF16 when not
The compressor_reduce.cu kernel now adds position_bias to BOTH kv and
gate values, matching the PyTorch reference. Previously the kernel only
added it to gate, and a Python workaround loop was adding it to both
before the kernel call (then passing None to the kernel).
Changes:
- compressor_reduce.cu: add position_bias to kv_val in pass 2 (CSA + HCA)
- single_shot_inference.py: remove Python position_bias loop, pass
self.ape directly to csa/hca_compress_production
- production_compress.py: already supports position_bias passthrough
- New compressor_reduce.cu: CSA/HCA token-level softmax + weighted sum + kv_norm
One block per compressed entry, 128 threads, FP32 accumulation
CSA: overlapping Ca/Cb streams (2m tokens per block)
HCA: single stream (m tokens per block)
Includes apply_kv_norm kernel (unweighted RMSNorm + weight)
- New production_compress.py: Python wrapper for CUDA kernels
- single_shot_inference.py: Compressor/Indexer now use production Nvfp4Linear
for kv_proj, gate_proj, q_b_proj, weights_proj projections
Then CUDA reduce kernel for softmax + weighted sum
No more PyTorch reference nvfp4_linear_ref in compressor/indexer path
Critical bug: checkpoint weights are (N_packed, K_packed) N-major format,
but make_b_k_major expects (E, K_packed, N_packed) input. Without the
permute, the K and N dimensions are swapped, producing garbage output
with wrong dimensions (e.g., q_a output was 3584 instead of 1536).
Also fix scale assembly: checkpoint scales are (N, K_sf) which should
use assemble_raw_scales_2d3d_3d_side (no transpose), not
assemble_scales_3d_side (which incorrectly transposes K_sf↔N).
Critical bug fix: weight_scale_2 (the second-level NVFP4 scale) was
being dropped entirely in the production pipeline. The dequant formula
is lut[w] * weight_scale * weight_scale_2, so weight_scale_2 must be
folded into the GEMM's global_scale_b parameter.
Fixes in:
- Nvfp4Linear: ws2 field, folded in finalize_weights()
- Nvfp4MoE: l1_ws2/l2_ws2 lists, folded in _ensure_stacked()
- Nvfp4SharedExpert: l1_ws2/l2_ws2 lists, folded in finalize_weights()
- single_shot_inference.py: pass weight_scale_2 through all loading paths
- Also fix missing o_a_prod key fallback in attention output
