diff --git a/WALKING_BACK_SOME_QUANTS.md b/WALKING_BACK_SOME_QUANTS.md new file mode 100644 index 00000000..b8d1e2de --- /dev/null +++ b/WALKING_BACK_SOME_QUANTS.md @@ -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 ... # 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. \ No newline at end of file diff --git a/single_shot_inference.py b/single_shot_inference.py index 9d10121e..725d9c76 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -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