59 lines
3.1 KiB
Markdown
59 lines
3.1 KiB
Markdown
# Current Bug: vLLM produces NaN from layer 0
|
|
|
|
**Status:** Active debugging — BF16 dequant fix in progress
|
|
**Date:** 2026-05-18
|
|
|
|
## Symptom
|
|
- vLLM server starts, loads model, but every inference produces NaN logits
|
|
- Diagnostic prints show **NaN from layer 0 onward** — no layer ever produces valid output
|
|
- Empty content in chat completions, NaN in logprobs
|
|
|
|
## Root Cause (in progress)
|
|
|
|
**The attention NVFP4 linear layers produce NaN immediately.**
|
|
|
|
The attention projections go through vLLM's `FlashInferCutlassNvFp4LinearKernel` which uses checkpoint `input_scale` as the activation global scale for `scaled_fp4_quant()`. The checkpoint `input_scale` values are wrong for this use case, causing overflow → NaN.
|
|
|
|
### What we've tried
|
|
|
|
1. ✅ **MoE kernel is NOT the problem** — `test_runner_vllm_style.py` with warmup gs gives cosine 0.988, no NaN
|
|
2. ❌ **Dequant ALL attn projections to BF16** — crashed: `wo_a.weight_scale_inv` missing (fp8_einsum needs it)
|
|
3. ❌ **Dequant all except wo_a (keep wo_a as FP8)** — still NaN from layer 0. `wq_a` and `wkv` don't exist as separate attrs — they're **fused as `fused_wqa_wkv`**
|
|
4. ❌ **Changed to dequant `fused_wqa_wkv`** — still NaN from layer 0. Debug prints added to check if the attrs are actually found.
|
|
|
|
### Current theory
|
|
|
|
The BF16 dequant code may not be finding `fused_wqa_wkv` on the attention module, so it silently skips the most important projection. Debug logging added in latest commit to verify.
|
|
|
|
### Attention architecture (DeepSeek V4 MLA)
|
|
|
|
- `fused_wqa_wkv` — MergedColumnParallelLinear (q_a + kv fused)
|
|
- `wq_b` — ColumnParallelLinear (second Q projection after RoPE)
|
|
- `wo_a` — ColumnParallelLinear (FP8 via fp8_einsum, weight-only, NO input_scale)
|
|
- `wo_b` — ColumnParallelLinear (final output projection)
|
|
- `compressor` — already handled (reconstructed to BF16 from checkpoint)
|
|
|
|
### Why `wo_a` is safe as FP8
|
|
|
|
`wo_a` uses `fp8_einsum` which does `output = fp8_act * fp8_weight * scale`. It's a **weight-only FP8** GEMM — no `input_scale` involved. The NaN comes from `scaled_fp4_quant(x, input_global_scale_inv)` in the other projections.
|
|
|
|
## Key evidence
|
|
|
|
- `q_a_proj.input_scale = 0.00025141` → `1/input_scale = 3977.6` → quantizing activations with amax ~2-8 by 3977.6x = massive overflow
|
|
- `q_b_proj.input_scale = 0.00006140` → `1/input_scale = 16287.1` → even worse
|
|
- Embedding values: amax=1.27, std=0.09 — very small values that get multiplied by thousands during quantization
|
|
|
|
## Next steps
|
|
|
|
1. Check debug logs to see which projections were actually dequantized
|
|
2. If `fused_wqa_wkv` wasn't found, fix the attribute path
|
|
3. If it was found and dequantized, the NaN source is elsewhere (wo_b? wq_b? something else?)
|
|
4. Consider: maybe the NaN is from the **KV cache FP8 quantization** or the **RoPE** implementation
|
|
|
|
## Docker/Build Notes
|
|
|
|
- Build: `screen -dmS build bash -c './build_and_run.sh 2>&1 | tee build.log'`
|
|
- Currently using `--enforce-eager` + `CLAWMINE_DEBUG=1` for diagnostics
|
|
- Don't hit the API with enforce-eager — JIT spikes crash the container
|
|
- For real testing: use compilation-config `{"cudagraph_mode": "NONE", "custom_ops": ["all"]}` instead of enforce-eager
|