biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix: Use PyTorch dequant_nvfp4 for weight dequantization (compressor/indexer/router gate)

Browse Source 
The CUDA dequantize_nvfp4 (dsv4/ops/quantize.py) was designed for
activations/KV and assumes row-major (M, N/16) scale layout. Using it
for weight dequantization caused async illegal memory access because
weight scales don't match the kernel's expected layout. The kernel only
validates row count, not width or contiguity.

All 4 call sites now use the PyTorch dequant_nvfp4 (defined in
single_shot_inference.py) which handles weight_scale_2 and input_scale
correctly and cannot cause OOB access:
- Compressor.load: kv_proj, gate_proj
- Indexer.load: weights_proj
- Router gate dequantization in main()
This commit is contained in:
biondizzle
2026-06-03 14:57:40 +00:00
parent 1121cd7b47
commit f577ed97f4
2 changed files with 82 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

69
WALKING_BACK_SOME_QUANTS.md Normal file
View File

@@ -0,0 +1,69 @@
# DSV4 Precision Floor — PyTorch Validation (PART 1) + Native Port (PART 2)
**What we learned:** the NVFP4 precision floor for this model is — keep **LM head** BF16, **router gate** BF16, and the **compressor/indexer helper projections** BF16, with the **one exception** that the **CSA indexer QK path stays FP4** (it was explicitly FP4-QATed; the other compressor projections were not, so PTQ-ing them to FP4 breaks). We validated each individually. Now do all of them together, simple-PyTorch first, then native.
---
## ⚠️ First: the CUDA illegal-memory-access (you're calling the wrong dequant)
There are **two** functions with nearly the same name:
- `single_shot_inference.py:238` — `dequant_nvfp4(weight, weight_scale, weight_scale_2, input_scale)` — **pure PyTorch** (does `weight_scale.repeat_interleave(16,1) * scales`). This is what `nvfp4_linear_ref` uses — your **validated reference**. It cannot cause an illegal access.
- `dsv4/ops/quantize.py:377` — `dequantize_nvfp4(x_fp4, x_sf, gsa)` — calls the **CUDA kernel** `dequant_nvfp4.cu`. **This is the one crashing.**
The precision-floor code (lines 328 / 333 / 426: kv_proj, gate_proj, wp) imports the **CUDA** one and feeds it **weights**. But that kernel was written for the **activation / KV-gather** path — read its own docstring: *"compressed KV is stored as NVFP4, dequantized on-the-fly."* It assumes row-major `(M, N/16)` block scales, per-row `gsa`, `N=512`.
The host wrapper only does `TORCH_CHECK(sf_data.size(0) == M)` — it validates the scale's **row count and nothing else** (not width, not total size, not contiguity). The kernel then indexes `sf_data[m*(N/16) + n_block]` flat. For a weight whose scale isn't *exactly* contiguous row-major `(M, N/16)` — different width, padding, non-contiguous `.to(dev)` view, or the GEMM swizzle — that index walks off the allocation → **async illegal access, surfacing at the next sync (the compressor load).** The activation/KV path never tripped it because those scales already match the assumed layout.
**Confirm it in 2 minutes** (the error is async, so do this to localize it):
```bash
compute-sanitizer --tool memcheck <your harness> ... # will name dequant_nvfp4_kernel + the sf_data read
# or: CUDA_LAUNCH_BLOCKING=1 to move the report to the offending launch
```
And add these guards to `dequant_nvfp4_cuda` in `dequant_nvfp4.cu` — they turn the async crash into an immediate, located error and print the size mismatch:
```cpp
TORCH_CHECK(fp4_data.is_contiguous() && sf_data.is_contiguous(), "dequant inputs must be contiguous");
TORCH_CHECK(sf_data.numel() >= (int64_t)M * (N/16), "sf too small: have ", sf_data.numel(), " need ", (int64_t)M*(N/16));
TORCH_CHECK(fp4_data.numel() >= (int64_t)M * (N/2), "fp4 too small: have ", fp4_data.numel(), " need ", (int64_t)M*(N/2));
```
You don't need the CUDA kernel here at all (see PART 1) — these weights are dequanted **once at load**, so there's zero performance reason to use a custom kernel for them.
---
## PART 1 — PyTorch quick version (all floor fixes together, simple, no crash)
Goal: one combined config, pure PyTorch, prove correctness end-to-end. This also sidesteps the OOB by not using the CUDA dequant for weights.
1. **Swap the three weight-dequant call sites (328/333/426) to the PyTorch reference.** The CUDA `dequantize_nvfp4(kv_w, kv_ws, gsa)` becomes the PyTorch `dequant_nvfp4(kv_w, kv_ws, kv_ws2, kv_isc)` — and you can delete the manual `gsa = torch.tensor([ws2_v]*shape[0])` lines, because the PyTorch version handles `weight_scale_2` / `input_scale` internally. Be explicit about *which* function you import (they're nearly identically named — that's how this got crossed). Example:
```python
from single_shot_inference import dequant_nvfp4 as dequant_nvfp4_torch # the pure-PyTorch one
# kv_proj:
self._kv_bf16 = dequant_nvfp4_torch(kv_w.to(dev), kv_ws.to(dev), kv_ws2, kv_isc).to(dev).contiguous()
# gate_proj, wp: same pattern
```
2. **LM head → BF16, router gate → BF16.** Dequant their FP4 weights to BF16 once at load via the same PyTorch path, then run them as plain `F.linear`. (The gate is tiny; the LM head is the only sizable one and it's ~1.4 GB — negligible against the KV/concurrency budget.)
3. **Keep the CSA indexer QK path in FP4 — do NOT dequant it.** Only the QK projection of the indexer was QATed. Its non-QATed siblings in the compressor go to BF16 with everything else.
4. **Run a clean generation** with the fixed chat template (the official `encoding/encoding_dsv4.py`, not the hand-rolled path). Confirm: coherent, **no repetition loop**, **clean stop**, Paris top-1 on the canonical probe, and run **≥ a few hundred tokens** so HCA actually engages (HCA's first compressed entry only forms at 128 tokens).
5. **A/B insurance:** this is the all-at-once config. If it regresses versus the individual fixes, flip one component FP4↔BF16 at a time to find the interaction — and record which ones were necessary (that table is the NVIDIA-writeup evidence).
---
## PART 2 — Native CuteDSL / CUDA version
Only after PART 1 validates the combined config (it becomes your reference for it).
1. **Fix the weight dequant path** (you have two options; pick one):
- *Simplest:* keep dequanting these few weights to BF16 **at load in PyTorch** (PART 1) even in the native build. It's a one-time load op — no hot-path cost — so there's no need to native-ize it at all.
- *If you insist on the CUDA kernel for load:* add the `numel`/contiguity guards above, then make the scale match what the kernel reads. The raw checkpoint `weight_scale` appears row-major **before** `finalize_weights` (the production GEMM swizzles at finalize — see the "K-major + swizzle" step ~line 1352 — so the *raw* scale is unswizzled). The guards will tell you if it's actually `(M, N/16)` contiguous; if not, make it contiguous before launch or teach the kernel the real stride. Also: the kernel was built around `N=512`; for weights `N=in` (≈7168) — make sure nothing downstream hardcodes 512.
2. **Hot-path natives are unchanged:** FP8 FMHA, FP4 MoE, and the **FP4 CSA indexer QK** all stay as they are. The floor change only touches load-time weight handling + two small GEMMs (gate, lm_head) that run as native **BF16** (cuBLAS/standard), not FP4.
3. **Re-validate per-layer cosine** of the native build against the PART 1 PyTorch combined-config reference before declaring done.
---
## Guardrails
- Don't reintroduce the **CUDA** `dequantize_nvfp4` for **weights** until the wrapper guards are in and the scale layout is confirmed — for now the PyTorch dequant is correct and crash-proof.
- The two functions `dequant_nvfp4` (PyTorch, weights) and `dequantize_nvfp4` (CUDA, activations/KV) are a foot-gun. Consider renaming the CUDA one to `dequantize_nvfp4_kvcache` so this can't recur.
- Only the **CSA indexer QK** path is FP4-QATed — do not let FP4 creep onto its non-QATed siblings.
- Validate end-to-end (coherent + non-looping + clean stop + HCA-depth) **before** calling it done.

38
single_shot_inference.py
View File

@@ -317,22 +317,17 @@ class Compressor:
def load(self, w, pfx, dev=None):
"""Load weights and build BF16 projections (dequantized from NVFP4)."""
if dev is None: dev = self.device
# Compressor projections are NOT explicitly FP4-QATed — use BF16 F.linear
from dsv4.ops.quantize import dequantize_nvfp4
# Compressor projections are NOT explicitly FP4-QATed — dequant to BF16, use F.linear
# CRITICAL: Use the PyTorch dequant_nvfp4 (defined in this file), NOT the CUDA
# dequantize_nvfp4 from dsv4/ops/quantize.py. The CUDA kernel assumes
# activation/KV scale layout (row-major (M, N/16)) and crashes on weight scales
# that don't match — async illegal memory access surfaces at next sync.
kv_w, kv_ws, kv_ws2, kv_isc = get_nvfp4_weight(w, pfx, 'kv_proj')
gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(w, pfx, 'gate_proj')
if kv_w is not None:
ws2_v = kv_ws2.float().item() if kv_ws2 is not None else 1.0
# For weight dequantization: gsa = ws2 (NOT input_scale * ws2)
# input_scale is the activation global scale, only used in GEMM's gsb computation
gsa = torch.tensor([ws2_v] * kv_w.shape[0], device=dev, dtype=torch.float32)
kv_bf16 = dequantize_nvfp4(kv_w.to(dev), kv_ws.to(dev), gsa) # (out, in)
self._kv_bf16 = kv_bf16.to(dev).contiguous()
self._kv_bf16 = dequant_nvfp4(kv_w.to(dev), kv_ws.to(dev), kv_ws2, kv_isc).to(dev).contiguous()
if gate_w is not None:
ws2_v = gate_ws2.float().item() if gate_ws2 is not None else 1.0
gsa = torch.tensor([ws2_v] * gate_w.shape[0], device=dev, dtype=torch.float32)
gate_bf16 = dequantize_nvfp4(gate_w.to(dev), gate_ws.to(dev), gsa)
self._gate_bf16 = gate_bf16.to(dev).contiguous()
self._gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev), gate_ws2, gate_isc).to(dev).contiguous()
self.ape = w.get(f"{pfx}.position_bias")
self.kv_norm_w = w.get(f"{pfx}.kv_norm.weight")
@@ -417,15 +412,10 @@ class Indexer:
qb_out = qb_w.shape[0]
qb_in = qb_w.shape[1] * 2
self.q_b_lin = make_nvfp4_linear(qb_in, qb_out, dev, w, pfx, 'q_b_proj')
# weights_proj is NOT FP4-QATed — use BF16 F.linear
# weights_proj is NOT FP4-QATed — dequant to BF16 via PyTorch reference
# CRITICAL: Use PyTorch dequant_nvfp4, NOT CUDA dequantize_nvfp4 (see Compressor.load)
if wp_w is not None:
from dsv4.ops.quantize import dequantize_nvfp4
ws2_v = wp_ws2.float().item() if wp_ws2 is not None else 1.0
isc_v = wp_isc.float().item() if wp_isc is not None else 1.0/(6.0*448.0)
gsb = isc_v * ws2_v # global_scale_b = input_scale * weight_scale_2
gsa = torch.tensor([gsb] * wp_w.shape[0], device=dev, dtype=torch.float32)
wp_bf16 = dequantize_nvfp4(wp_w.to(dev), wp_ws.to(dev), gsa)
self._wp_bf16 = wp_bf16.to(dev).contiguous()
self._wp_bf16 = dequant_nvfp4(wp_w.to(dev), wp_ws.to(dev), wp_ws2, wp_isc).to(dev).contiguous()
# Indexer compressor weights are directly under the indexer prefix
# (e.g. *.indexer.kv_proj.weight), NOT nested under *.indexer.compressor.
if f"{pfx}.kv_proj.weight" in w:
@@ -1323,11 +1313,9 @@ def main():
gate_w, gate_ws, gate_ws2, gate_isc = get_nvfp4_weight(all_w, pfx, 'gate')
if gate_w is not None and gate_ws is not None:
# Checkpoint has NVFP4 gate weight — dequantize to BF16
from dsv4.ops.quantize import dequantize_nvfp4
ws2_v = gate_ws2.float().item() if gate_ws2 is not None else 1.0
# For weight dequantization: gsa = ws2 (NOT input_scale * ws2)
gsa = torch.tensor([ws2_v] * gate_w.shape[0], device=dev, dtype=torch.float32)
gate_bf16 = dequantize_nvfp4(gate_w.to(dev), gate_ws.to(dev), gsa) # (E_packed*2, H)
# CRITICAL: Use PyTorch dequant_nvfp4, NOT CUDA dequantize_nvfp4
# (same fix as Compressor.load — CUDA kernel crashes on weight scale layouts)
gate_bf16 = dequant_nvfp4(gate_w.to(dev), gate_ws.to(dev), gate_ws2, gate_isc)
router.W_gate = gate_bf16.T.contiguous().to(dev) # (H, E) for F.linear(x, W_gate.T)
else:
# BF16 gate weight from checkpoint