biondizzle
a51edd238e
The key insight: process_weights_after_loading runs AFTER load_weights and sets up FlashInferCutlassNvFp4LinearKernel with broken input_global_scale_inv. Any fix inside load_weights gets overwritten. Solution: register a one-shot forward pre-hook that runs on the first forward call (guaranteed after all init). It dequantizes attention NVFP4 weights to BF16 and replaces quant_method with UnquantizedLinearMethod. Since process_weights_after_loading already ran, our changes won't be overwritten. Standalone test confirmed: all attention weights produce valid non-NaN output when dequantized to BF16.
105 lines
4.4 KiB
Markdown
105 lines
4.4 KiB
Markdown
# Current Bug: vLLM produces empty/garbage output
|
|
|
|
**Status:** Weights confirmed good — bug is in vLLM's quant pipeline for attention
|
|
**Date:** 2026-05-18
|
|
|
|
## Symptom
|
|
- vLLM server starts, loads model, processes requests (200 OK)
|
|
- Chat completions return `content: ""` with `finish_reason: "length"`
|
|
- 20 completion tokens generated but all produce empty/NaN logits
|
|
- With enforce-eager + diagnostics: **NaN from layer 0 onward** on real requests
|
|
|
|
## ✅ Confirmed: Weights produce valid output
|
|
|
|
Standalone test (`test_attn_moe_chain.py`) running directly on B200:
|
|
|
|
| Step | Operation | amax | NaN? |
|
|
|------|-----------|------|------|
|
|
| 1 | Embed tokens | 1.27 | No |
|
|
| 2 | hc_mult expansion | 1.27 | No |
|
|
| 3 | RMSNorm | 0.20 | No |
|
|
| 4 | q_a_proj (NVFP4→BF16 dequant + matmul) | 0.50 | No |
|
|
| 5 | kv_proj (NVFP4→BF16 dequant + matmul) | 1.30 | No |
|
|
| 6 | q_norm + kv_norm | 0.11 / 1.87 | No |
|
|
| 7 | q_b_proj (NVFP4→BF16 dequant + matmul) | 1.10 | No |
|
|
| 8 | MoE CuTeDSL runner (with warmup gs) | cosine 0.988 | No |
|
|
|
|
**Every step produces valid, non-NaN, non-zero output.** The problem is NOT the weights.
|
|
|
|
## ❌ Root Cause: vLLM's `process_weights_after_loading` breaks attention
|
|
|
|
### The timeline
|
|
|
|
1. `load_weights()` → our `_convert_nvfp4_post_load()` runs
|
|
2. `process_weights_after_loading()` → vLLM's quant method runs AFTER, **overwriting our fixes**
|
|
3. `FlashInferCutlassNvFp4LinearKernel` gets set up with broken `input_global_scale_inv`
|
|
|
|
### What the quant method does
|
|
|
|
`CompressedTensorsW4A4Fp4.process_weights_after_loading()`:
|
|
```python
|
|
input_global_scale_inv = layer.input_scale.max() # = 0.00025141 (WRONG)
|
|
layer.input_global_scale = 1.0 / input_global_scale_inv # = 3977.6
|
|
layer.input_global_scale_inv = input_global_scale_inv # = 0.00025141
|
|
layer.alpha = input_global_scale * weight_global_scale
|
|
```
|
|
|
|
At runtime: `scaled_fp4_quant(x, input_global_scale_inv=0.00025141)` divides by 0.00025141 → multiplies by 3977.6 → massive overflow → NaN.
|
|
|
|
### Why our fixes didn't work
|
|
|
|
| Attempt | Why it failed |
|
|
|---------|---------------|
|
|
| BF16 dequant + `UnquantizedLinearMethod` | `process_weights_after_loading` overwrites `quant_method` back to `FlashInferCutlassNvFp4LinearKernel` |
|
|
| Fix `input_scale` before quant method | Runs too early — quant method reads `input_scale` and overwrites our value |
|
|
| Fix `input_global_scale_inv` directly | Attribute doesn't exist yet when our code runs — it's set BY the quant method |
|
|
|
|
### The key insight
|
|
|
|
Our code runs **inside** `load_weights()`. The quant method's `process_weights_after_loading()` runs **after** `load_weights()` returns. Any changes we make get overwritten.
|
|
|
|
## Config values (corrected)
|
|
|
|
| Parameter | Value |
|
|
|-----------|-------|
|
|
| head_dim | 512 (NOT 56) |
|
|
| num_attention_heads | 128 |
|
|
| num_key_value_heads | 1 |
|
|
| q_lora_rank | 1536 |
|
|
| qk_rope_head_dim | 64 |
|
|
| o_lora_rank | 1024 |
|
|
| hc_mult | 4 |
|
|
| n_routed_experts | 384 (48 per EP rank) |
|
|
|
|
## Next step: Post-init hook
|
|
|
|
The fix must run AFTER `process_weights_after_loading` and BEFORE the first inference. Options:
|
|
|
|
**Option A: Override `input_global_scale_inv` post-init**
|
|
- Add a `_fix_nvfp4_activation_scales()` method
|
|
- Call it from the right hook point (after quant method setup, before inference)
|
|
- Compute correct `input_global_scale_inv` from BF16 warmup
|
|
- Override the Parameter on each attention module
|
|
|
|
**Option B: Replace quant_method with UnquantizedLinearMethod post-init**
|
|
- After `process_weights_after_loading`, dequant weights to BF16
|
|
- Swap `quant_method` on attention modules to `UnquantizedLinearMethod`
|
|
- This time the quant method won't overwrite us (it already ran)
|
|
|
|
**Option C: Override the quant config to skip attention modules**
|
|
- Tell `CompressedTensorsW4A4Fp4` to skip attention projections
|
|
- Then dequantize to BF16 ourselves
|
|
- Cleanest but requires modifying the quant config
|
|
|
|
Option B is most straightforward. The quant method already ran and set up its attributes. We can then come in and replace everything with BF16.
|
|
|
|
## Architecture notes
|
|
|
|
- Attention uses MLA (Multi-head Latent Attention) with 2-step Q projection (q_a → q_b)
|
|
- `fused_wqa_wkv` = MergedColumnParallelLinear(q_a + kv fused)
|
|
- `wo_a` = FP8 via fp8_einsum (no input_scale, weight-only)
|
|
- `wo_b` = standard ColumnParallelLinear
|
|
- `hc_pre` / `hc_post` = Head-Conditioned mixing (tilelang custom ops)
|
|
- Dummy run zeros attention output by design (`out.zero_(); return`)
|
|
- FlashMLA handles the actual MLA attention kernel
|