docs: comprehensive README update through build 22

- Full architecture diagram and NVFP4→vLLM bridge details
- All 8 bugs documented with fixes
- SM100 hardware limitation (mxf4nvf4 unsupported)
- MegaMoE kernel architecture and debugging log (builds 1-22)
- Three paths forward (A: FlashInfer, B: BF16 mega_moe, C: SM103+)
- Container build pipeline, NVFP4 format spec, hard rules
This commit is contained in:
2026-05-11 13:53:41 +00:00
parent 8cb23bdb78
commit cfead0012d

319
README.md
View File

@@ -11,11 +11,10 @@ Full NVFP4 quantization of DeepSeek V4 Pro and vLLM serving on 8× NVIDIA B200 G
| NVFP4→FP8 Conversion (wo_a) | ✅ DeepGEMM block-scale format |
| NVFP4→BF16 Dequantization | ✅ 305 attn/shared, 91 compressor layers |
| Compressor Reconstruction | ✅ Separate kv_proj/gate_proj → fused_wkv_wgate |
| MoE Expert Serving | ✅ FusedMoE NVFP4 (FLASHINFER_TRTLLM backend) |
| Profile/Warmup Run | ✅ Passes |
| MoE Expert Serving (FusedMoE) | ✅ FLASHINFER_TRTLLM backend |
| MoE Expert Serving (MegaMoE) | 🔧 Kernel compiles, runs, but garbled (SM100 HW limit) |
| API Server | ✅ Running on port 8000 |
| Output Quality | 🔧 Garbled — likely remaining dequant/scale bug |
| NVFP4 Mega MoE Kernel | 🔧 Built, pending compile & integration |
| Output Quality | 🔧 Garbled — UE8M0 scale precision loss + attention bugs |
## B200 Node
@@ -27,6 +26,14 @@ Full NVFP4 quantization of DeepSeek V4 Pro and vLLM serving on 8× NVIDIA B200 G
- **Model weights**: `/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4/`
- **BF16 reference**: `/root/nvidia-meeting/DeepSeek-V4-Pro-BF16/`
## Repositories
| Repo | Branch | Purpose |
|------|--------|---------|
| `deepseek-v4-quant` | `modelopt-nvfp4` | Main patch (FlashInfer FusedMoE path) |
| `deepseek-v4-quant` | `mega-moe-nvfp4` | MegaMoE patch (DeepGEMM mega_moe path) |
| `DeepGEMM` | `nvfp4-mega-moe` | NVFP4 mega_moe kernel fork |
## Architecture
```
@@ -37,8 +44,8 @@ DeepSeek V4 Pro (1.2T params, 61 layers)
│ ├── wo_b → BF16 (UnquantizedLinearMethod)
│ └── compressor.fused_wkv_wgate → BF16 (reconstructed from NVFP4)
├── MoE Experts (384 experts, 61 layers)
│ ├── w13_weight → NVFP4 (FusedMoE, FLASHINFER_TRTLLM backend)
│ └── w2_weight → NVFP4 (FusedMoE, FLASHINFER_TRTLLM backend)
│ ├── [FusedMoE path] → NVFP4 (FLASHINFER_TRTLLM backend)
│ └── [MegaMoE path] → NVFP4 (DeepGEMM mxf8f6f4, UE4M3→UE8M0 adapted)
└── Shared Expert → FP8 (Fp8LinearMethod, DeepGEMM)
```
@@ -58,99 +65,162 @@ Our NVFP4 weights are uint8 packed FP4 with separate block/global scales.
**Solution** (`_convert_nvfp4_to_fp8`):
1. Unpack NVFP4 uint8 → BF16 using E2M1 lookup table
2. Dequantize: `weight_bf16 * block_scale * global_scale` (NO input_scale — it's for activations)
2. Dequantize: `weight_bf16 * block_scale * global_scale` (NO input_scale)
3. Re-quantize BF16 → FP8 e4m3 with per-tensor scale (`w_amax / fp8_max`)
4. Create block scale tensor filled with `fp8_scale` (same scale for every 128×128 block)
5. Call `deepgemm_post_process_fp8_weight_block(wq, ws, quant_block_shape=(128,128), use_e8m0=True, is_bmm=True, bmm_batch_size=N)`
6. Store: `weight_scale_inv = dg_ws` (DeepGEMM-formatted scale), `weight = w_fp8` (3D BMM shape)
**Why `weight_scale_inv`?** The attention forward reads `self.wo_a.weight_scale_inv` as
`b_scale` for `deepseek_v4_fp8_einsum` → DeepGEMM `fp8_einsum`. This must be the
DeepGEMM block-scale tensor, not a per-tensor scalar.
**Why `fp8_scale` in the block scale (not all-ones)?** DeepGEMM divides by the block
scale at runtime. If the block scale is all-ones, it divides by 1.0, producing garbage.
Each block needs the actual per-tensor scale value.
4. Create block scale tensor filled with `fp8_scale`
5. Call `deepgemm_post_process_fp8_weight_block` with `quant_block_shape=(128,128)`, `use_e8m0=True`, `is_bmm=True`
6. Store: `weight_scale_inv = dg_ws`, `weight = w_fp8` (3D BMM shape)
### 2. Attention Layers: NVFP4 → BF16
**Problem**: `fused_wqa_wkv`, `wo_b` use standard `torch.nn.functional.linear`.
NVFP4 weights (uint8) can't be used directly.
**Solution** (`_convert_nvfp4_to_bf16`):
1. Unpack NVFP4 → BF16
2. Dequantize with block/global scales (input_scale is for activations, not weights)
3. Replace `mod.weight` with BF16 parameter
4. Set `quant_method = UnquantizedLinearMethod()`
5. Remove NVFP4 scale attributes (`weight_scale`, `weight_scale_2`, `input_scale`)
**Solution** (`_convert_nvfp4_to_bf16`): Unpack → dequantize → set `UnquantizedLinearMethod`.
### 3. Compressor: Reconstructing fused_wkv_wgate from NVFP4
**Problem**: The compressor's `fused_wkv_wgate` is a `MergedColumnParallelLinear`
with `disable_tp=True`. NVFP4 uint8 data can't be loaded into the BF16 parameter
(shape mismatch: uint8 is half the input dim). The default weight loader silently
skips these weights, leaving the parameter uninitialized.
**Solution** (`_reconstruct_compressor_weight`):
1. Read original `kv_proj.weight` and `gate_proj.weight` directly from safetensors
2. Unpack NVFP4 → BF16, dequantize with scales
3. Concatenate: `fused = cat([wkv, wgate], dim=0)`
4. Replace the uninitialized parameter
**Critical detail**: The **indexer** compressor is at a different checkpoint path:
- Main: `model.layers.N.self_attn.compressor.{kv_proj,gate_proj}.weight`
- Indexer: `model.layers.N.self_attn.compressor.indexer.{kv_proj,gate_proj}.weight`
Using the wrong prefix loads the main compressor weight into the indexer's
`fused_wkv_wgate`, causing a 4× shape mismatch and `split_with_sizes` crash.
**Solution** (`_reconstruct_compressor_weight`): Read `kv_proj`+`gate_proj` from
safetensors, dequantize, concatenate. **Critical**: indexer compressor at
`.compressor.indexer.{kv_proj,gate_proj}` not `.compressor.{kv_proj,gate_proj}`.
### 4. MoE Experts: NVFP4 FusedMoE
**Problem**: vLLM's DeepSeek V4 uses `DeepseekV4MegaMoEExperts` with DeepGEMM
grouped GEMM. NVFP4 experts need a different kernel path.
**Solution**: The existing `ModelOptNvFp4LinearMethod` + `FusedMoE` infrastructure
handles NVFP4 experts natively. We just need to:
- Keep expert weights as NVFP4 uint8 + block/global scales
- Use `FLASHINFER_TRTLLM` MoE backend (auto-selected)
- Skip any conversion in `process_weights_after_loading`
**Solution**: Keep expert weights as NVFP4, use `FLASHINFER_TRTLLM` MoE backend.
### 5. BF16 wo_a Layers: BF16 → FP8
**Problem**: Some `wo_a` layers were NOT quantized by modelopt (BF16 in checkpoint).
The attention forward still reads them as FP8 for the einsum path.
**Solution** (`_convert_bf16_to_fp8`): Same as #1 but skip the NVFP4 unpack step.
Directly quantize BF16 → FP8 with block scale.
**Solution** (`_convert_bf16_to_fp8`): Directly quantize BF16 → FP8 with block scale.
## Bugs Found and Fixed
### DeepGEMM `sf.dim()` Assertion (layout.hpp:94)
- **Root cause**: `weight_scale_inv` was a 1D per-tensor scale `(g,)`. DeepGEMM expects
2D/3D block-scale tensor formatted by `transform_sf_into_required_layout`.
- **Fix**: Use `deepgemm_post_process_fp8_weight_block` to produce correctly formatted
block scales, store result in `weight_scale_inv`.
| # | Bug | Impact | Fix |
|---|-----|--------|-----|
| 1 | DeepGEMM `sf.dim()` crash | Server crash | `deepgemm_post_process_fp8_weight_block` for block-scale format |
| 2 | Block scale dtype `float8_e4m3fn` | Crash | Use `float32` |
| 3 | Missing `deepgemm_post_process` args | Crash | Pass `quant_block_shape`, `use_e8m0` |
| 4 | Compressor indexer shape mismatch | Crash | `.indexer.` sub-path in checkpoint keys |
| 5 | All-ones block scale | Garbage output | `torch.full(..., fp8_scale)` not `torch.ones` |
| 6 | `fused_skip_regex` skipping q_b/o_a/o_b scales | Garbage output | Remove non-fused scale entries from skip list |
| 7 | UE8M0 block scale misinterpreted as E4M3 | Garbled output | `_ue8m0_to_float32()`: reinterpret raw uint8 as IEEE 754 exponent |
| 8 | wo_a BF16 weight into uint8 param (suspected) | Double-conversion loss | On-the-fly BF16→NVFP4 in weight_loader, or BF16→FP8 directly |
### Block Scale dtype (`float8_e4m3fn` vs `float32`)
- **Root cause**: `deepgemm_post_process_fp8_weight_block` expects `float32` or
`float8_e8m0fnu` block scales. We initially used `float8_e4m3fn`.
- **Fix**: Create block scale as `dtype=torch.float32`.
### Bug #7 Detail: UE8M0 → float32 Misinterpretation
### Missing `deepgemm_post_process` args
- **Root cause**: Function signature changed to require `quant_block_shape` and `use_e8m0`.
- **Fix**: Pass `quant_block_shape=(128, 128)` and `use_e8m0=True`.
**Root cause**: `weight_scale` bytes are E8M0 format (power-of-2 only, 8-bit exponent),
but `.to(torch.float32)` interprets the raw byte as E4M3 (8-bit: sign+exp+mantissa),
producing a completely wrong float value.
### Compressor Indexer Shape Mismatch
- **Root cause**: `_reconstruct_compressor_weight` used the same checkpoint prefix
for both main and indexer compressors. The indexer's keys have `.indexer.` in the path.
- **Fix**: Add `sub_path` parameter; pass `".indexer"` for indexer compressors.
**Fix**: `_ue8m0_to_float32()` — reinterpret the raw uint8 bits as the upper 8 bits
of an IEEE 754 float32 exponent: `(uint8_value << 23).view(float32)`. Applied to all
dequant paths.
### All-Ones Block Scale → Garbage Output
- **Root cause**: Block scale was `torch.ones(...)` (scale=1.0). DeepGEMM divides by
the block scale at runtime, so the output was divided by 1.0 instead of the actual
per-tensor scale, producing incoherent text.
- **Fix**: Use `torch.full(..., fp8_scale.item())` to fill the block scale with the
correct per-tensor FP8 quantization scale.
### Bug #8 Detail: wo_a BF16 Loading
`o_a_proj.weight` is BF16 in checkpoint, but `ModelOptNvFp4Config` creates a uint8
param (shape mismatch: BF16 `(16384,4096)` vs uint8 `(16384,2048)`). The weight_loader
does on-the-fly BF16→NVFP4 quantization, but the double conversion (BF16→NVFP4→BF16→FP8)
is lossy. Diagnostics added but fix pending.
## NVFP4 Mega MoE Kernel
### What We Built
A native NVFP4 mega_moe kernel in our DeepGEMM fork that avoids dequantizing
expert weights to BF16 before the GEMM. The kernel keeps weights in E2M1 packed
format and uses block-scaled MMA directly.
### SM100 Hardware Limitation (CRITICAL)
**B200 (SM100) does NOT support `kind::mxf4nvf4`** (neither `scale_vec::2X` nor `4X`).
This PTX instruction requires SM103 (B300) or SM120 (GB300). On SM100, the only
FP4 block-scaled MMA is `kind::mxf8f6f4` with UE8M0 scales (block32, group_size=32).
| Parameter | NVFP4 Checkpoint | Kernel (SM100 Adapted) |
|-----------|-----------------|----------------------|
| Weight format | E2M1 uint8 | E2M1 uint8 (unchanged) |
| Block scale format | UE4M3 (float8_e4m3fn) | UE8M0 (uint8) — **adapted for HW** |
| Block size | 16 | 32 (merged adjacent pairs, max) |
| Global scale | float32 | Folded in before UE4M3→UE8M0 |
| PTX instruction | N/A (requires SM103+) | `mxf8f6f4.block_scale` (same as MXFP4) |
**Result**: Server starts and serves, but output is **garbled**. The UE4M3→UE8M0
conversion loses 3 bits of mantissa precision per scale (8× precision loss),
which destroys output quality. The E2M1 weights are correct, but the power-of-2-only
UE8M0 scales can't faithfully represent the original UE4M3 values.
### Kernel Architecture
```
sm100_fp8_nvfp4_mega_moe_impl (adapted from sm100_fp8_fp4_mega_moe_impl)
├── Same E2M1 weight packing as MXFP4
├── Same TMEM layout as MXFP4 (2X, block32)
├── Same UTCCP copy (4x32 transpose, i*4 stride)
├── mxf8f6f4.block_scale PTX instruction (UE8M0)
├── float_ue8m0_t instruction descriptor
└── UE8M0 L1 epilogue (>> 23 activation scales)
Python API:
├── fp8_nvfp4_mega_moe() — recipe=(1,1,32)
├── transform_nvfp4_weights_for_mega_moe()
│ ├── fold_global_scale(): UE4M3 * FP32 → UE4M3
│ ├── merge_block16_to_block32(): max of adjacent pairs
│ ├── UE4M3 → float32 → UE8M0 (extract exponent byte)
│ └── pack_uint8_to_int32() + transform_sf_into_required_layout()
└── get_symm_buffer_for_nvfp4_mega_moe() — 2x SF buffer
C++ Bindings:
├── csrc/apis/mega_nvfp4.hpp
├── csrc/jit_kernels/impls/sm100_fp8_nvfp4_mega_moe.hpp
└── csrc/apis/layout.hpp — gran_k=32 support
```
### Container Build Pipeline
```
Dockerfile → FROM atl.vultrcr.com/vllm/vllm-with-lmcache:dream-build
├── DeepGEMM (nvfp4-mega-moe branch) — JIT-compiled at runtime
├── vLLM patch (deepseek_v4.py) — COPY over model file
└── NVRTC symlink for CUDA compilation
build_push.sh → build → login to CR → push → update docker-compose
Container registry: atl.vultrcr.com/vllm/vllm-dsv4-nvfp4:latest
Always run builds in screen: screen -S nvfp4-build
```
### Debugging Log (Builds 122)
| Build | Error | Fix |
|-------|-------|-----|
| 16 | Various Dockerfile/build issues | NVRTC symlink, CPATH, PYTHONPATH |
| 7 | `kPackedFP4` type mismatch | uint8→int8 view on weights |
| 9 | SF stride assertion | Need MN-major layout + TMA alignment |
| 10 | `transform_sf_into_required_layout` no gran_k=16 | C++ fix |
| 11 | SF dtype float8_e4m3fn rejected | Pack UE4M3→int32 first |
| 1214 | SF stride layout | Transpose to MN-major before transform |
| 15 | SymmBuffer too small (NVFP4 has 2× SF) | NVFP4-specific SymmBuffer |
| 16 | `ImportError: deep_gemm.mega.nvfp4` | Python wrapper in mega/__init__.py |
| 17 | NVCC: `scale_vec::4X` not supported on sm_100f | — |
| 18 | NVCC: `scale_vec::2X` ALSO not supported | — |
| 19 | kGranK=16 still in C++ binding | → 32 |
| 20 | UE4M3→UE8M0 `uint32 >> 23` fails | Cast to int32 first |
| 22 | Server UP, but garbled output | UE4M3→UE8M0 precision loss |
## Path Forward
The mega_moe approach has a **hardware ceiling on B200** — the Tensor Core can't
consume UE4M3 block scales. Three options:
### Option A: Fix FlashInfer FusedMoE Path (Recommended)
The `modelopt-nvfp4` branch already uses FlashInfer FP4 MoE which dequantizes
NVFP4→BF16 before GEMM. This avoids the UE4M3→UE8M0 precision loss. The garbled
output on that branch is likely from attention layer bugs (#7, #8), not the MoE.
Fix those and we should get coherent output.
### Option B: Dequant NVFP4→BF16 in MegaMoE Shared Memory
Build a mega_moe that dequantizes in shared memory, then uses BF16 MMA.
Slower than FlashInfer but gets the mega_moe communication pattern.
### Option C: Wait for SM103+ Hardware
B300 (SM103) and GB300 (SM120) support `mxf4nvf4` natively with UE4M3 scales.
The kernel we built would work correctly on that hardware.
## Running
@@ -173,87 +243,22 @@ curl http://localhost:8000/v1/chat/completions \
| File | Purpose |
|------|---------|
| `patches/deepseek_v4.py` | Main patch: NVFP4 post-load conversion, weight reconstruction, DeepGEMM block-scale |
| `patches/deepseek_v4.py` | Main patch: NVFP4 post-load conversion, weight reconstruction |
| `patches/modelopt.py` | ModelOpt FP4 config patches for weight loading |
| `.env` | B200 node credentials |
| `docker-compose.yml` | Container config (8 GPU, TP=8, EP=8, NVFP4 quant) |
| `Dockerfile` | Container build (extends dream-build with DeepGEMM + patch) |
| `build_push.sh` | Build, push to CR, update docker-compose |
## Conversion Flow
## NVFP4 Format Specification
```
Checkpoint (NVFP4 safetensors)
├── [weight loader] ──→ vLLM model (NVFP4 uint8 params)
└── [process_weights_after_loading]
├── wo_a (is_bmm=True):
│ NVFP4→BF16→FP8 + DeepGEMM block scale
│ weight_scale_inv = dg_ws, weight = 3D FP8
├── fused_wqa_wkv, wo_b, shared_expert:
│ NVFP4→BF16, UnquantizedLinearMethod
├── compressor.fused_wkv_wgate:
│ Read kv_proj+gate_proj from checkpoint
│ NVFP4→BF16, cat into fused weight
└── MoE experts: stay NVFP4 (FusedMoE backend)
```
- **Weights**: E2M1 packed uint8 (2 values per byte)
- **Block scales**: float8_e4m3fn (UE4M3), group_size=16
- **Global scales**: float32 (weight_scale_2), per-tensor
- **Dequant formula**: `value = packed_E2M1 * block_scale * global_scale`
- **Block scale range**: [0, 448] (UE4M3 max = 448, E2M1 max = 6, so 6×448 = 2688)
- **UE8M0**: Power-of-2 only. Encoded as uint8 = float32_exponent_bits[31:23]
- **UE4M3**: 3-bit mantissa + 4-bit exponent + sign. Max = 448.
## Bugs Found and Fixed (continued)
## HARD RULES
### `input_scale` Multiplied into Weight Dequantization (CRITICAL)
- **Root cause**: `_convert_nvfp4_to_bf16`, `_convert_nvfp4_to_fp8`, and
`_reconstruct_compressor_weight` all multiplied by `input_scale` during weight
dequantization. `input_scale` is for **activations**, not weights. The correct
formula is: `weight_bf16 = e2m1 * block_scale * global_scale` (NO input_scale).
Including it made weights ~5000× too small, causing garbage output.
- **Fix**: Removed `* input_scale` from all three dequant paths.
### `fused_skip_regex` Skipping Non-Fused Layer Scales (CRITICAL)
- **Root cause**: The skip list included `q_b_proj`, `o_a_proj`, `o_b_proj` weight
scales. These are **NOT fused/stacked** — they're individual Linear layers
(`wq_b`, `wo_a`, `wo_b`) converted in-place. Skipping their scales caused
`process_weights_after_loading` to read `torch.empty()` garbage for
`weight_scale_inv`, producing garbled output.
- **Fix**: Removed `q_b_proj`, `o_a_proj`, `o_b_proj` scale entries from
`fused_skip_regex`. Only truly stacked params remain skipped:
`compressor.{kv_proj,gate_proj}``fused_wkv_wgate`,
`self_attn.{kv_proj,q_a_proj}``fused_wqa_wkv`,
`shared_experts.{gate_proj,up_proj}``gate_up_proj`.
## Version Banner
The patch prints a version banner at import time (visible in `docker logs`):
```
======================================================================
DeepSeek V4 NVFP4 Patch
Commit: 26aaaba
Loaded: 2026-05-11 04:25:00 UTC
Node: ...
Architecture: ...
Bugs fixed: #1-#6
======================================================================
```
This ensures you can always verify what's running inside the container.
## Known Issues
1. **Output quality**: Model produces tokens but they're garbled/incoherent.
All 6 known bugs are fixed. The remaining issue is under investigation —
likely a subtle dequantization bug (sign handling, scale ordering, or
E2M1 unpack edge case). The version banner in the logs helps debug which
patch version is active.
2. **Runtime performance**: Not yet benchmarked. The DeepGEMM einsum + FusedMoE
path should be efficient on B200, but the BF16 layers go through
`UnquantizedLinearMethod` which may be slower than dedicated kernels.
## Quantization Details
- **Model**: DeepSeek V4 Pro (1.2T parameters)
- **Format**: NVIDIA NVFP4 (4-bit floating point with 128-element block scales)
- **Tool**: modelopt 0.45.0.dev64 + transformers 5.8.0.dev0
- **Run**: Run 11 (881GB), 8× B200, ~$161/run
- **Checkpoint**: 95 safetensors shards
- **NEVER convert DeepSeek MoE experts to MXFP4.** Experts stay in NVFP4. Period.