nvfp4-megamoe-kernel/NEXT_SESSION.md at eb5ef93bf1390d4129d5eb91b19ccde29d49ec53

Files

biondizzle eb5ef93bf1 Add A/B comparison mode for P4 fused vs unfused RMSNorm+quantize

- Added --ab-compare flag to run both fused and unfused paths for first 3 layers
- Compares x_normed, gsa values, FP4 data, and GEMM outputs (q_a, kv)
- Added --no-fused-rmsnorm to disable P4 and use unfused path
- This will help diagnose the correctness regression introduced by P4

2026-06-02 18:49:30 +00:00

2.9 KiB

Raw Blame History

Here's the summary for next session:

Summary: What Was Done & What to Investigate

What was done today

P4 — Fused RMSNorm + NVFP4 quantize: New CUDA kernel (2 launches vs 6+). Integrated into single_shot_inference.py for the attention path. Unit test passes (cos=0.996-0.999 vs unfused). Decode speed improved 0.53→0.43s/token.
P5 — Fused mHC + RMSNorm + NVFP4 quantize: New CUDA kernel (2 launches vs 7+). Unit test passes on B200 (cos=0.996-0.999). NOT yet integrated into single_shot — needs moe_forward refactoring.
CRITICAL FIX — Indexer deadlock: __syncthreads() inside a strided loop caused deadlock at production context lengths (always deadlocked when num_valid % 128 != 0). Fixed with per-thread local top-k + block-level merge. Both copies updated.

The regression

After P4 integration, model output went from coherent English to garbled nonsense. This is a NEW regression — output was correct last night after P0-P3 + KV work.

What to investigate (ordered by likelihood)

run_from_quantized gsa shape bug (MOST LIKELY): The fused kernel produces per-row gsa (shape (M,)). run_from_quantized passes this to the CuTeDSL NVFP4 GEMM as global_scale_a. But the GEMM expects a single scalar — it's one scale for the entire A matrix. For M=1 decode, shape (1,) may work as a scalar, BUT the gsa VALUE differs from the unfused path's scalar gsa because:
- Unfused: quantize computes gsa from the full (M, N) tensor — single scalar
- Fused: computes gsa per-row, then [:1].reshape(1) takes only the first row's gsa
- These differ when rows have different magnitudes
Dequant→requant noise for compressor: The compressor gets x_normed by dequantizing the fused kernel's FP4 output. This introduces ~0.5% quantization error (cos=0.994) that wasn't present before. This noise propagates into compression scores and indexer queries.
A/B test first: Run with _use_fused_rmsnorm_quantize=False to confirm P4 is the cause. If output is correct with unfused, the bug is in the P4 path.

Key code paths

dsv4/layers/linear.py: run_from_quantized() — the gsa handling is suspicious
dsv4/ops/gemm_runner.py: run_nvfp4_grouped_gemm() — how global_scale_a is consumed
dsv4/kernels/gemm/grouped.py: CuTeDSL GEMM — global_scale_a is scalar
single_shot_inference.py:855-870: P4 integration point
dsv4/kernels/cuda/fused_rmsnorm_quantize.cu: the fused kernel

Quick fix ideas

Use scalar gsa (reduce per-row gsa to a single max) for the GEMM — keeps the fusion but makes gsa compatible
Or: don't use run_from_quantized — just use the fused kernel for the BF16 output (dequant) and let each linear re-quantize with its own scalar gsa (saves rmsnorm launches but not quantize launches)
Or: fix the CuTeDSL GEMM to support per-row global_scale_a (THIS IS COMPLEX, BUT IF IT IS THE CORRECT WAY OF DOING THINGS. I WANT IT DONE THAT WAY!!!!!)

2.9 KiB Raw Blame History