docs: rewrite README, nuke DEBUG_LOG, add vLLM integration stub
README.md: full rewrite explaining how we got here, project structure, plan, and key lessons learned from the C++ CUTLASS disaster. Removed: - DEBUG_LOG.md (old debug timeline, no longer relevant) - REWRITE_PLAN.md (plan is now in README) - test_gemm.py (C++ extension test) Added: - vllm/nvfp4_cutedsl.py: CuTeDSLMoERunner class for vLLM integration - Replaces nvfp4_mega_moe_full + SymmBuffer with CuTeDSL kernel - Handles slot-based routing, L1→SiLU→L2→scatter - prepare_weights_from_dequantized() for weight prep Tagged the-last-of-cutlass on the old C++ kernel state.
This commit is contained in:
110
DEBUG_LOG.md
110
DEBUG_LOG.md
@@ -1,110 +0,0 @@
|
||||
# NVFP4 MegaMoE Debug Log
|
||||
|
||||
## Current State (May 16, 2026 — 01:15 UTC)
|
||||
|
||||
**Status:** GEMM verified correct, SF remap verified correct, B layout verified correct, L2 slot_token cleaned up. vLLM still produces garbage. The checkpoint `input_scale` red herring is documented below. The bug remains unidentified.
|
||||
|
||||
**What's verified correct (DO NOT re-investigate):**
|
||||
1. SF remap: roundtrip verifier = 0 errors, forward mapping with `layout_sf(make_coord(mn, k_elem, 0))`
|
||||
2. GEMM math: uniform FP4 + uniform SF → exact output (72.0 = 1.5² × K)
|
||||
3. B matrix layout: byte transpose correct, column-dependent weight test passes
|
||||
4. L2 GEMM input: slot-major, no gather needed (cleaned up dead `slot_token` param)
|
||||
|
||||
**What's still broken:**
|
||||
- "The capital of France is" → garbage tokens (varies, e.g. `-W'MSG173`, `( z tractor`)
|
||||
- All magnitudes reasonable, no NaN anywhere
|
||||
|
||||
## Red Herring: Checkpoint input_scale (May 16, 00:10 UTC)
|
||||
|
||||
Mike's code review suggested using the checkpoint's `input_scale` as the activation global scale instead of the dynamic `amax/(6*448)`. This was **wrong** and has been reverted (commit `79b9bec`).
|
||||
|
||||
**What happened:**
|
||||
- Applied `stage_activation(hidden_states, input_global_scale=w13_input_scale)`
|
||||
- `w13_input_scale = 2.86e-4` (from checkpoint, same for all experts in a layer)
|
||||
- Dynamic `amax/(6*448)` = 3.70e-3 (13x larger)
|
||||
- Using 2.86e-4 for normalization: `x / 2.86e-4` produces values in the thousands for typical hidden states (amax ~5)
|
||||
- ALL block scales saturated to 448.0 (max float8_e4m3 value)
|
||||
- Output: still garbage, but now with quantized-to-death activations
|
||||
|
||||
**Why it's wrong:**
|
||||
The checkpoint `input_scale` is NOT the `input_global_scale = amax/(6*448)` normalization constant. They are different quantities:
|
||||
- `input_global_scale` normalizes data to [0,1] before FP4 quantization
|
||||
- `input_scale` is a calibration constant from the Quark quantization tool
|
||||
|
||||
The calibration amax for `input_scale = 2.86e-4` would be `2.86e-4 * 6 * 448 = 0.77`. Runtime hidden states have amax ~5-10. The `input_scale` was computed on a different data distribution (probably calibration data, not actual inference data).
|
||||
|
||||
**The correct use of `input_scale` is still unknown.** The Quark path computes `alpha = input_scale * weight_scale_2`, but this may assume BF16 activations (not FP4-quantized). Our CUTLASS kernel requires FP4 input, so we must quantize with the dynamic scale.
|
||||
|
||||
**Preserved for future use:** `_w13_input_scale` and `_w2_input_scale` are now saved in `finalize_weights` (not dropped) in case we need them for alpha computation later.
|
||||
|
||||
**Checkpoint input_scale values (layer 0, all experts identical):**
|
||||
- `gate_proj.input_scale` = `up_proj.input_scale` = 2.862840e-04
|
||||
- `down_proj.input_scale` = 3.069196e-02
|
||||
- `weight_scale_2` = 4.650298e-05 (all projections)
|
||||
- Scales vary by layer: layer 0 = 2.86e-4, layer 60 = 1.07e-2
|
||||
|
||||
## SF Remap — Final Correct Implementation (commit `6626b75`)
|
||||
|
||||
```cpp
|
||||
int mn = tid / K_sf;
|
||||
int k_sf = tid % K_sf;
|
||||
int k_elem = k_sf * 16;
|
||||
int dst_idx = layout_sf(cute::make_coord(mn, k_elem, 0));
|
||||
dst[dst_idx] = src[mn * src_stride_mn + k_sf * src_stride_ksf];
|
||||
```
|
||||
Source strides: SFA=(K_sf, 1), SFB=(1, N)
|
||||
Allocation: `cute::size(cute::filter_zeros(layout))`
|
||||
|
||||
## Previous Bugs Fixed
|
||||
|
||||
### BF16 reference comparison (May 15, 23:37 UTC)
|
||||
The 0.2 cosine against the Python BF16 dequantization reference was a RED HERRING. The reference is wrong, not the GEMM. 8+ iterations of SF remap changes all produced the same 0.2 cosine because it was never about the remap. **A wrong reference is worse than no reference.**
|
||||
|
||||
### `cute::size` vs `cute::cosize` (commit `c384198`)
|
||||
Iteration bound used `size` (logical) instead of `cosize` (physical). Fixed but insufficient alone.
|
||||
|
||||
### M/K coordinate extraction in `idx2crd` (commits `deb6b32` → `30b6c89`)
|
||||
Original had M/K swapped. Mike's correction: `mn = f0 + 32*f1 + 128*f2`, `k_sf = f4 + 4*f5`.
|
||||
|
||||
### `if/else if` fallthrough (commit `6626b75`)
|
||||
Dead `dst_idx=0` when no branch matched. Fix: branchless `layout_sf(make_coord(...))`.
|
||||
|
||||
### `col_major_src` ambiguity (commit `7285331`)
|
||||
Boolean flag → explicit `src_stride_mn, src_stride_ksf`.
|
||||
|
||||
### Allocation size (commit `6626b75`)
|
||||
`cosize` → `size(filter_zeros(layout))`.
|
||||
|
||||
### L2 slot_token cleanup (commit `bb5a1ba`)
|
||||
`nvfp4_moe_l2` accepted `slot_token` but never passed it to GEMM. Removed dead parameter.
|
||||
|
||||
## Mike's Code Review Answers
|
||||
|
||||
1. **E2M1 packing:** Confirmed correct — element 2j in low nibble, 2j+1 in high nibble. Suggested hardware oracle with `__nv_cvt_bfloat16raw2_to_fp4x2`.
|
||||
2. **A RowMajor:** Confirmed correct — no micro-tiling for A.
|
||||
3. **B ColumnMajor:** Byte transpose confirmed correct by test. Mike flagged theoretical concern about nibble-level transpose but our test passed.
|
||||
4. **Alpha/global scale:** Mike suggested `alpha = input_scale * weight_scale_2` (from checkpoint). We tried it — wrong for activation normalization. The correct use of `input_scale` in our pipeline is still TBD.
|
||||
5. **Gate/up correction:** Mathematically valid. `up_half *= up_weight_gs / gate_weight_gs` is equivalent to per-column alpha.
|
||||
|
||||
## Architecture Notes
|
||||
|
||||
### NVFP4 MoE Pipeline
|
||||
```
|
||||
stage_activation(hidden_states) → x_fp4, x_sf, input_global_scale
|
||||
L1 GEMM: (x_fp4, x_sf) @ (l1_w, l1_sf) with alpha=igs*l1_global_sf → gate_up
|
||||
SiLU(gate) * up → activated
|
||||
stage_activation(activated) → l1_fp4, l1_sf, l1_igs
|
||||
L2 GEMM: (l1_fp4, l1_sf) @ (l2_w, l2_sf) with alpha=l1_igs*l2_global_sf → output
|
||||
scatter with routing weights → y
|
||||
```
|
||||
|
||||
### Per-element multiply order
|
||||
`res += A_fp4 * SFA_fp8 * B_fp4 * SFB_fp8`
|
||||
|
||||
## Next Steps
|
||||
|
||||
1. **Compare against BF16 model** — run the same prompt on a known-good implementation to see if the attention layers are working and only MoE is broken
|
||||
2. **Check the vLLM model integration** — how does the MoE output get mixed with the residual? Is `hc_post` correct?
|
||||
3. **Understand the Quark input_scale contract** — maybe we need to NOT quantize activations to FP4 and instead use BF16 input
|
||||
4. **Add per-layer token output logging** — see which layer the tokens go off the rails
|
||||
5. **Check o_a_proj BF16 handling** — it's kept in BF16 in the checkpoint, is it being processed correctly?
|
||||
419
README.md
419
README.md
@@ -1,361 +1,130 @@
|
||||
# nvfp4-megamoe-kernel
|
||||
# NVFP4 MegaMoE Kernel
|
||||
|
||||
Native NVFP4 block-scaled MoE kernel for DeepSeek-V4-Pro on NVIDIA Blackwell (SM100).
|
||||
NVFP4 block-scaled Mixture-of-Experts kernel for DeepSeek-V4 on NVIDIA Blackwell (SM100). Uses CuTeDSL — NVIDIA's Python-based CUTLASS DSL — for a native NVFP4 pipeline that takes full advantage of Blackwell's TMA, MMA, and epilogue overlap.
|
||||
|
||||
Replaces the broken `fp8_nvfp4_mega_moe` kernel from DeepGEMM with a working CUTLASS-based implementation that emits real `SM100_MMA_MXF4_SS` tensor core instructions.
|
||||
## What This Is
|
||||
|
||||
---
|
||||
|
||||
## Architecture
|
||||
|
||||
DeepSeek-V4-Pro is a 384-expert MoE model with expert parallelism across 8 ranks (B200 GPUs). Each rank handles 48 experts. For each token, the router picks the top-6 experts.
|
||||
|
||||
### The MoE Forward Pass
|
||||
A fused MoE FFN kernel that runs the entire expert forward pass in NVFP4:
|
||||
|
||||
```
|
||||
Input hidden states (BF16)
|
||||
│
|
||||
▼
|
||||
┌─────────────────┐
|
||||
│ Shared Experts │ ← vLLM native FlashInfer CUTLASS NVFP4 path
|
||||
│ (gate + up → │ (not our kernel)
|
||||
│ SiLU * up → │
|
||||
│ down) │
|
||||
└─────────────────┘
|
||||
│
|
||||
▼
|
||||
Staging Kernel (vLLM built-in)
|
||||
BF16 → packed E2M1 (int8) + UE4M3 block-16 scales (uint32)
|
||||
Writes to SymmBuffer.x / SymmBuffer.x_sf
|
||||
│
|
||||
▼
|
||||
Router (vLLM built-in)
|
||||
Writes topk_ids / topk_weights to SymmBuffer
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────┐
|
||||
│ nvfp4_mega_moe_full │ ← nvfp4_mega_moe.py
|
||||
│ │
|
||||
│ 1. Read staged activation from buffer │
|
||||
│ 2. Build slot mapping (token, topk) → local │
|
||||
│ expert, routing weight │
|
||||
│ 3. L1 GEMM: gate_up_proj (slot-based) │ ← CUTLASS NVFP4 block-scaled
|
||||
│ E2M1 × E2M1 + UE4M3 scales │ SM100_MMA_MXF4_SS PTX
|
||||
│ → BF16 per-slot output (6144-wide) │
|
||||
│ 4. SiLU(gate) * up PER SLOT │
|
||||
│ (nonlinearity before combining paths) │
|
||||
│ 5. stage_activation: BF16 → FP4 │ ← proper E2M1 quantization
|
||||
│ 6. L2 GEMM: down_proj (slot-based) │ ← CUTLASS NVFP4 block-scaled
|
||||
│ E2M1 × E2M1 + UE4M3 scales │ SM100_MMA_MXF4_SS PTX
|
||||
│ → BF16 per-slot output (7168-wide) │
|
||||
│ 7. Final scatter: │
|
||||
│ y.index_add_(slot_token, │
|
||||
│ slot_weight * l2_slots) │
|
||||
│ Routing weight applied ONCE at scatter │
|
||||
└─────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
Cross-rank all-reduce (vLLM built-in)
|
||||
BF16 input → quantize to NVFP4
|
||||
L1 GEMM: NVFP4 × NVFP4 → BF16 (gate + up)
|
||||
SiLU(gate) * up → BF16 (only nonlinear — can't avoid BF16 here)
|
||||
Re-quantize → NVFP4
|
||||
L2 GEMM: NVFP4 × NVFP4 → BF16 (down_proj)
|
||||
Scatter with routing weights → BF16 output
|
||||
```
|
||||
|
||||
### Slot-Based Dispatch
|
||||
Both GEMMs are fully NVFP4: A and B in `float4_e2m1fn_x2`, block scales in `float8_e4m3fn`, global scales in `float32`. BF16 is used only for the SiLU activation and the final scatter — the minimum possible.
|
||||
|
||||
The kernel uses a **slot representation** instead of collapsing expert outputs early. A slot is one `(token, topk_expert)` pair. For a batch of T tokens with top-6 routing, there are up to 6T slots (fewer if some experts are out of the local rank's range).
|
||||
## How We Got Here
|
||||
|
||||
**Why slots?** Two bugs in the previous approach:
|
||||
### The C++ CUTLASS Kernel Was Broken
|
||||
|
||||
1. **SiLU after summing is mathematically wrong.** `silu(Σ wᵢ·gateᵢ) * (Σ wᵢ·upᵢ) ≠ Σ wᵢ·silu(gateᵢ)·upᵢ`. The nonlinearity must happen per-expert-path before combining.
|
||||
The original kernel was a C++ `.cu` file using CUTLASS's C++ API directly. It passed all the simple tests (uniform data → exact output, SF remap verifier → 0 errors) but produced **cosine 0.05** with real random data. After weeks of debugging the SF remap (8+ iterations, all producing the same 0.2 cosine against a wrong reference), we discovered:
|
||||
|
||||
2. **Routing weights applied twice.** The old grouped GEMM applied `topk_weights` in its scatter loop, and was called for both L1 and L2 — squaring the weights.
|
||||
1. **The BF16 reference comparison was wrong** — our Python dequantization didn't match CUTLASS's internal FP4 handling. A wrong reference is worse than no reference. We chased ghosts through 8+ SF remap rewrites because the 0.2 cosine was never about the remap.
|
||||
|
||||
The slot approach fixes both: SiLU+Mul happens per-slot, and routing weights are applied exactly once at the final `index_add_` scatter.
|
||||
2. **The C++ CUTLASS kernel misinterpreted FP4 data** — even with SF remap verified correct (0 byte errors), the GEMM produced garbage with non-uniform data. The issue was in how CUTLASS's C++ API handles FP4 packing/tiling internally — something we couldn't easily debug or fix.
|
||||
|
||||
### SFB (Weight Scale Factors) — Remapped Per-Call, NOT Cached
|
||||
3. **The checkpoint `input_scale` was a red herring** — we tried using the checkpoint's calibration scale as the activation normalization scale. It saturated all block scales to 448.0 (max float8). The `input_scale` is a calibration constant for alpha computation, not a normalization scale.
|
||||
|
||||
Weight scale factors (SFB) are remapped from row-major to CUTLASS interleaved layout on every GEMM call. This is a lightweight scatter kernel (~µs) and is NOT the bottleneck compared to the GEMM itself.
|
||||
### The CuTeDSL Kernel Works
|
||||
|
||||
⚠️ **DO NOT ADD A PREPACK CACHE FOR SFB.** Previous attempts caused critical issues:
|
||||
NVIDIA's CuTeDSL approach (Python-based CUTLASS kernels compiled via MLIR → PTX) is what the CUTLASS team recommends for Blackwell. Their official MoE scaled grouped GEMM example (`torch_scaled_grouped_mm.py`) supports NVFP4 out of the box. We adapted it.
|
||||
|
||||
| Problem | Impact |
|
||||
|---------|--------|
|
||||
| **OOM** | ~1.75 GiB per prepacked tensor × 61 MoE layers × 2 (L1+L2) = ~214 GiB — exceeds B200 capacity |
|
||||
| **Peak memory 2×** | `torch.stack` held all expert tensors + final stack simultaneously before LRU eviction |
|
||||
| **CUDA graph trap** | LRU eviction frees tensors that CUDA graphs still reference → use-after-free → silent corruption or crash |
|
||||
| **M-dependent layout** | `prepack_sfb(M=128)` assumed SFB layout size is M-independent (never verified). If wrong, entire prepack is invalid |
|
||||
| **Cross-layer cache collision** | Tag-based cache (`"l1"`/`"l2"`) returned layer N-1's data for layer N. Fixed with data_ptr key, but the cache itself was the root problem |
|
||||
**Results with real DeepSeek-V4 layer 0 weights:**
|
||||
- L1 GEMM alone: cosine 0.995
|
||||
- Full MoE pipeline (L1→SiLU→L2→scatter): cosine 0.989
|
||||
|
||||
The per-call remap costs microseconds. The cache cost was hours of debugging. Don't repeat this mistake.
|
||||
The 0.989 is expected quantization loss — we dequantize the checkpoint NVFP4→BF16 then re-quantize BF16→NVFP4 (double quantization). Future optimization: load checkpoint FP4 bytes directly into `float4_e2m1fn_x2` tensors.
|
||||
|
||||
---
|
||||
### Key Lessons
|
||||
|
||||
### vLLM Startup Sequence (how our code plugs in)
|
||||
1. **A wrong reference is worse than no reference** — the 0.2 cosine against a broken BF16 dequant sent us chasing SF remap bugs for weeks
|
||||
2. **The C++ CUTLASS API is a footgun for FP4** — CuTeDSL handles tensor layouts, tiling, and SF construction correctly by construction
|
||||
3. **Test with real data early** — uniform tests pass even with broken kernels; random data reveals real bugs
|
||||
4. **Separate the GEMM from the pipeline** — our `layertest.py` runs without vLLM, Docker, or tensor parallelism. It caught the kernel bug that vLLM's integration layers masked.
|
||||
|
||||
## Project Structure
|
||||
|
||||
```
|
||||
1. vLLM engine init
|
||||
└─ ModelOptNvFp4Config selected (NVFP4 quantization scheme)
|
||||
└─ FlashInferCutlassNvFp4LinearKernel for linear layers
|
||||
|
||||
2. Model construction
|
||||
└─ DeepseekV4ForCausalLM → DeepseekV4MoE → DeepseekV4DecoderLayer
|
||||
Each layer has: attention + MoE block
|
||||
MoE block has: shared experts + 384 routed experts
|
||||
|
||||
3. Weight loading
|
||||
└─ 95 safetensor shards loaded
|
||||
└─ weight, weight_scale, weight_scale_2 loaded per linear
|
||||
|
||||
4. process_weights_after_loading ← THIS IS WHERE WE HOOK IN
|
||||
└─ ModelOptNvFp4LinearMethod swizzles/pads weights for CUTLASS
|
||||
└─ finalize_mega_moe_weights()
|
||||
└─ weight_transform.py: transform_nvfp4_weights_for_mega_moe()
|
||||
• Folds weight_scale_2 (global scale) into weight_scale (block scale)
|
||||
• UE4M3 block-16 scales: 4 values packed per uint32
|
||||
• Returns ((l1_w, l1_sf), (l2_w, l2_sf)) per rank
|
||||
|
||||
5. SymmBuffer allocation
|
||||
└─ symm_buffer.py: get_symm_buffer_for_nvfp4_mega_moe()
|
||||
• Pre-allocates GPU buffers for:
|
||||
- x: int8 packed E2M1 activations
|
||||
- x_sf: uint32 packed UE4M3 activation scales
|
||||
- topk_idx: int32 expert indices
|
||||
- topk_weights: float32 routing weights
|
||||
- buffer: BF16 all-reduce buffer
|
||||
|
||||
6. Profile run (warmup)
|
||||
└─ First forward pass to allocate KV cache, etc.
|
||||
└─ This is where the CUTLASS GEMM first executes
|
||||
└─ SFB weight scales remapped per-expert inside CUTLASS (no cache)
|
||||
|
||||
7. Ready to serve
|
||||
nvfp4-megamoe-kernel/
|
||||
├── cutedsl/ # CuTeDSL kernel + bridge layer
|
||||
│ ├── bridge.py # Tensor layout conversion, quantization, kernel launch
|
||||
│ ├── moe_pipeline.py # Full MoE pipeline (L1→SiLU→L2→scatter)
|
||||
│ └── kernel/moe/ # NVIDIA's ScaledGroupedGemmKernel (untouched)
|
||||
│ ├── torch_scaled_grouped_mm.py # The working kernel (3900 lines)
|
||||
│ ├── moe_utils.py
|
||||
│ ├── moe_persistent_scheduler.py
|
||||
│ └── moe_sched_extension.py
|
||||
├── src/nvfp4_megamoe_kernel/ # OLD Python pipeline (being replaced)
|
||||
│ ├── nvfp4_mega_moe.py # Old pipeline — calls broken C++ kernel
|
||||
│ └── cutlass_nvfp4_gemm/ # OLD C++ CUTLASS extension (BROKEN)
|
||||
├── vllm/ # vLLM integration
|
||||
│ └── patches/
|
||||
│ └── deepseek_v4.py # DeepSeek-V4 model patch
|
||||
├── tests/
|
||||
│ ├── layertest.py # Layer 0 comparison: CuTeDSL vs BF16 (✅ cosine 0.989)
|
||||
│ ├── test_cutedsl.py # Small standalone CuTeDSL test (✅ cosine 0.991)
|
||||
│ ├── test_uniform_fp4.py # Uniform data GEMM test
|
||||
│ ├── test_b_layout.py # B matrix column layout test
|
||||
│ └── test_quick_rand.py # Quick random GEMM sanity check
|
||||
├── reference/ # Reference files for study
|
||||
└── REWRITE_PLAN.md # Original rewrite plan
|
||||
```
|
||||
|
||||
---
|
||||
## The Bridge Layer (`cutedsl/bridge.py`)
|
||||
|
||||
## File Map
|
||||
Handles all tensor layout conversion from our pipeline to what the CuTeDSL kernel expects:
|
||||
|
||||
```
|
||||
nvfp4_megamoe_kernel/
|
||||
├── __init__.py # Public API exports
|
||||
├── nvfp4_mega_moe.py # Main kernel: nvfp4_mega_moe_full, L1/L2, stage_activation
|
||||
├── weight_transform.py # Weight prep: fold global scale, pack UE4M3
|
||||
├── symm_buffer.py # GPU buffer allocation for MoE dispatch
|
||||
│
|
||||
└── cutlass_nvfp4_gemm/ # CUTLASS CUDA extension (the actual hardware kernel)
|
||||
├── cutlass_nvfp4_gemm.cu # CUDA: CUTLASS GEMM + SF remap + prepack SFB + prepacked-SFB GEMM path
|
||||
├── pytorch_binding.cpp # PyTorch C++ binding (forward, forward_prepacked_sfb, prepack_sfb)
|
||||
├── kernel.py # Python: cutlass_grouped_nvfp4_gemm (slot-based, per-expert loop)
|
||||
├── sf_layout.py # CUTLASS SF layout reference docs
|
||||
├── setup.py # Build config (nvcc, CUTLASS include paths)
|
||||
├── build.sh # Build script
|
||||
├── test_gemm.py # Standalone test
|
||||
└── README.md
|
||||
```
|
||||
| Function | What it does |
|
||||
|----------|--------------|
|
||||
| `quantize_to_nvfp4()` | BF16 → float4_e2m1fn_x2 + float8_e4m3fn block scales + float32 global scale |
|
||||
| `quantize_weight_to_nvfp4()` | Same, but for weight matrices with K as the packed dimension |
|
||||
| `assemble_scales_2d_side()` | Pad and swizzle activation scale factors (2Dx3D A side) |
|
||||
| `assemble_scales_3d_side()` | Pad and swizzle weight scale factors (2Dx3D B side) |
|
||||
| `make_b_k_major()` | Convert B tensor from N-major to K-major strides (required by kernel) |
|
||||
| `compute_expert_offsets()` | Compute cumulative token offsets for grouped GEMM |
|
||||
| `run_nvfp4_grouped_gemm()` | Full kernel launch (compile + run) |
|
||||
|
||||
### What each file does (in call order)
|
||||
## Running Tests
|
||||
|
||||
| File | When it runs | What it does |
|
||||
|------|-------------|--------------|
|
||||
| `weight_transform.py` | Once at startup (weight loading) | Takes raw NVFP4 checkpoint weights, folds global scales into block scales. Returns scales as `float8_e4m3fn` (not packed uint32). Output: `((l1_w, l1_sf), (l2_w, l2_sf))` |
|
||||
| `symm_buffer.py` | Once at startup (buffer alloc) | Pre-allocates GPU tensors for activations, scales, routing data, and all-reduce. These persist across forward passes. |
|
||||
| `nvfp4_mega_moe.py` | Every forward pass | Orchestrates the MoE: reads from symm buffer → build slot mapping → L1 GEMM → SiLU+Mul per-slot → re-quantize → L2 GEMM → final index_add_ scatter with routing weights. Contains `stage_activation` (BF16→FP4) and `unpack_ue4m3_u32`. NO prepack cache — SFB remapped per-call inside CUTLASS. |
|
||||
| `cutlass_nvfp4_gemm/kernel.py` | Every forward pass (called by nvfp4_mega_moe) | Slot-based per-expert loop: gather slots for each expert, call CUTLASS GEMM (SFB remapped inside C extension), write results to slot buffer. No routing weights — caller handles scatter. |
|
||||
| `cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu` | Every forward pass (CUDA kernel) | The actual CUTLASS kernel: native NVFP4 block-scaled GEMM + GPU-side SFA and SFB remap. |
|
||||
| `cutlass_nvfp4_gemm/sf_layout.py` | Reference only | Documents the CUTLASS SfAtom layout. Not used at runtime (remap is in CUDA). |
|
||||
|
||||
---
|
||||
|
||||
## Data Formats
|
||||
|
||||
### Weights
|
||||
- **Packed E2M1** (`int8`): 2 FP4 values per byte. Shape: `(E_per_rank, N, K//2)`, K-major layout.
|
||||
- **UE4M3 block scales** (`float8_e4m3fn`): 1 scale per 16 FP4 values (group_size=16). Shape: `(E_per_rank, N, K//16)`. Returned as `float8_e4m3fn` from `weight_transform.py` — NOT packed uint32. The CUTLASS GEMM consumes float8 directly.
|
||||
|
||||
### Activations (after staging kernel)
|
||||
- **Packed E2M1** (`int8`): Shape: `(num_tokens, K//2)`.
|
||||
- **UE4M3 scales** (`uint32`): 4 UE4M3 values packed per uint32. Shape: `(num_tokens, K//64)`. Unpacked to `float8_e4m3fn` via `unpack_ue4m3_u32` before reaching the CUTLASS GEMM.
|
||||
|
||||
### GEMM dimensions (DeepSeek-V4-Pro)
|
||||
- **L1 (gate_up_proj):** M×6144×7168 (per expert)
|
||||
- **L2 (down_proj):** M×7168×3072 (per expert)
|
||||
- 48 experts per rank (384 total / 8 ranks), top-6 routing
|
||||
|
||||
---
|
||||
|
||||
## CUTLASS Scale Factor Remap
|
||||
|
||||
CUTLASS's `Sm1xxBlockScaledConfig` expects scale factors in a specific interleaved layout, not simple row-major. The SfAtom is:
|
||||
|
||||
```
|
||||
Atom Shape: Shape<Shape<32, 4>, Shape<16, 4>>
|
||||
Atom Stride: Stride<Stride<16, 4>, Stride<0, 1>>
|
||||
Tiling: Step<_2, _1> (M tiled with step 2, K with step 1)
|
||||
```
|
||||
|
||||
Our source data is row-major `(M, K_sf)` where `K_sf = K / 16`. The remap kernel (`remap_sf_to_cutlass_kernel` in `cutlass_nvfp4_gemm.cu`) converts from row-major to CUTLASS's interleaved layout.
|
||||
|
||||
### How the remap works
|
||||
|
||||
The kernel iterates over CUTLASS destination indices, uses `cute::idx2crd` to get the hierarchical coordinate, then `cute::flatten` to get a flat tuple of 8 sub-indices. From those, we extract logical `(m, k_sf)` and read from the row-major source.
|
||||
|
||||
### Flattened coordinate decomposition (flat_rank=8)
|
||||
|
||||
From the SfAtom layout with Step<_2, _1> tiling, `flatten(idx2crd(idx, ...))` produces 8 values:
|
||||
|
||||
```
|
||||
f0 = inner_m (0..31) — varies fastest within M atom
|
||||
f1 = sub_m (0..3) — second M sub-coordinate
|
||||
f2 = tile_m (0..) — M tile index
|
||||
f3 = step_m stride — degenerate (always = sfa_size, not a coordinate)
|
||||
f4 = sub_k (0..3) — K sub-coordinate within atom
|
||||
f5 = tile_k (0..) — K tile index
|
||||
f6 = 0 — unused
|
||||
f7 = 0 — unused
|
||||
```
|
||||
|
||||
#### Empirical coordinate dump (MN=8192, K_sf=448, T = sfa_size = 58720256)
|
||||
|
||||
| idx | f0 | f1 | f2 | f3 | f4 | f5 | f6 | f7 |
|
||||
| ----- | --- | --- | --- | --- | --- | --- | --- | --- |
|
||||
| 0 | 0 | 0 | 0 | T | 0 | 0 | 0 | 0 |
|
||||
| 1 | 0 | 0 | 0 | T | 1 | 0 | 0 | 0 |
|
||||
| 4 | 0 | 1 | 0 | T | 0 | 0 | 0 | 0 |
|
||||
| 16 | 1 | 0 | 0 | T | 0 | 0 |0 | 0 |
|
||||
| 511 | 31 | 3 | 0 | T | 3 | 0 | 0 | 0 |
|
||||
| 512 | 0 | 0 | 0 | T | 0 | 1 | 0 | 0 |
|
||||
| 1024 | 0 | 0 | 0 | T | 0 | 2 | 0 | 0 |
|
||||
| 2048 | 0 | 0 | 0 | T | 0 | 4 | 0 | 0 |
|
||||
| 4096 | 0 | 0 | 0 | T | 0 | 8 | 0 | 0 |
|
||||
| 8192 | 0 | 0 | 0 | T | 0 | 16 | 0 | 0 |
|
||||
| 65536 | 0 | 0 | 1 | T | 0 | 16 | 0 | 0 |
|
||||
| 131072 | 0 | 0 | 2 | T | 0 | 32 | 0 | 0 |
|
||||
|
||||
#### Extraction formula
|
||||
|
||||
CuTe uses "first sub varies fastest" for `Shape<32, 4>`:
|
||||
|
||||
```cpp
|
||||
m = f0 + f1 * 32 + f2 * 128;
|
||||
k_sf = f4 + f5 * 4;
|
||||
```
|
||||
|
||||
This was verified with 6 independent probes:
|
||||
|
||||
| Probe | Source | Expected | Result |
|
||||
|-------|--------|----------|--------|
|
||||
| SFA[1, 0] = 2.0 | row 1 changes | ✅ only row 1 | Confirms f0 term |
|
||||
| SFA[32, 0] = 2.0 | row 32 changes | ✅ only row 32 | Confirms f1*32, rules out f0*4+f1 |
|
||||
| SFA[128, 0] = 2.0 | row 128 changes | ✅ only row 128 | Confirms f2*128 |
|
||||
| SFA[0, 1] = 2.0 | row 0 changes (k=1) | ✅ only row 0 | Confirms f4 term |
|
||||
| SFA[0, 4] = 2.0 | row 0 changes (k=4) | ✅ only row 0 | Confirms f5*4 term |
|
||||
| SFA[0, 100] = 2.0 | row 0 changes (k=100) | ✅ only row 0 | Confirms tile-overflow range |
|
||||
|
||||
#### Why the previous remap was broken
|
||||
|
||||
The previous code used `cute::get<0>(flat)` and `cute::get<1>(flat)` to extract (m, k). Since flatten produces `(inner_m, sub_m, tile_m, ...)` in order, `get<0>` and `get<1>` are both **M sub-indices** — they carry no K information. This caused only `k_group=0` to work; all other K-groups were silently mapped to the wrong source offset.
|
||||
|
||||
Additionally, the dest buffer must be zero-initialized before remap because CUTLASS pads to tile boundaries (128 × 64), making the dest buffer larger than `M * K_sf`. Unmapped padding slots reading garbage caused sporadic wrong results.
|
||||
|
||||
---
|
||||
|
||||
## Bugs Found & Fixed
|
||||
|
||||
### 1. unpack_ue4m3_u32: value cast vs bit reinterpret
|
||||
**File:** `nvfp4_mega_moe.py`
|
||||
**Bug:** `(x_u32 & 0xFF).to(torch.int32).to(torch.float8_e4m3fn)` converts integer 63 → float8(63.0).
|
||||
**Fix:** `(x_u32 & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn)` reinterprets bit pattern 0x3F → float8(~0.984).
|
||||
**Also:** `uint32` lacks CUDA bitwise ops — cast to `int32` first.
|
||||
**Impact:** Corrupted every activation scale fed to the L1 GEMM. Weight scales were fine (already float8 from weight_transform). "Structured garbage" recipe.
|
||||
|
||||
### 2. stage_activation: three independent bugs
|
||||
**File:** `nvfp4_moe.py`
|
||||
**Bug A:** `clamp(0, 15)` zeroed every negative value. E2M1 is sign-magnitude 4-bit (bit3=sign, bits2:0=mag).
|
||||
**Bug B:** Stored `block_max` but divided by `block_max/6.0` → stored scale was 6× too large.
|
||||
**Bug C:** Uniform 0.5 step doesn't match E2M1 values {0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6} — non-uniform above ±2.
|
||||
**Fix:** Rewrote with proper nearest-neighbor E2M1 quantization.
|
||||
**Impact:** Half the L1→L2 activation was zeroed, 6× scale mismatch, quantization noise on top.
|
||||
|
||||
### 3. _fold_global_scale: logical_widths branch
|
||||
**File:** `weight_transform.py`
|
||||
**Bug:** `logical_widths=[3072, 3072]` caused the function to apply expert 0's scale to gate half and expert 1's scale to up half of ALL experts. All other experts' global scales were discarded.
|
||||
**Fix:** Removed the `logical_widths` branch entirely. The `else` branch correctly broadcasts each expert's own `(E, 1)` global scale across `(E, N, K//16)`.
|
||||
|
||||
### 4. L1 weight interleave removed (transpose still needed)
|
||||
**File:** `weight_transform.py`
|
||||
**Bug:** `_interleave_l1_weights` assumed gate/up were pre-interleaved in groups of 16 and that the kernel used 2CTA UMMA layout. vLLM uses plain concat `[gate; up]` along the output dim, and our CUTLASS kernel uses `ClusterShape<1, 1, 1>`.
|
||||
**Fix:** Removed the interleave function. Weights still need a transpose from checkpoint layout `(N, K_half)` row-major to CUTLASS layout `(K_half, N)` column-major — this is standard row→column conversion, not interleaving. Both L1 and L2 weights and scales are transposed.
|
||||
|
||||
### 5. SF remap: idx2crd+flatten coordinate extraction
|
||||
**File:** `cutlass_nvfp4_gemm.cu`
|
||||
**Bug:** `cute::flatten(coord)` produces 8 sub-indices (flat_rank=8). `get<0>` and `get<1>` are both M sub-indices (inner_m, sub_m), carrying zero K information. Only k_group=0 worked; all other K-groups were silently wrong.
|
||||
**Fix:** Correct extraction: `m = f0 + f1*32 + f2*128`, `k_sf = f4 + f5*4`. Zero-init dest buffer before remap.
|
||||
**Diagnostic trail:** Constant-scale test (all SF=1.0) → cosine 1.0 proved FP4 path was correct. Real scales → cosine 0.83 proved SF remap was broken. Single-element probes (SFA[0,0] vs SFA[0,3]) proved only k_group=0 worked. Printf dump of flat coordinates at specific indices revealed flat_rank=8 and the correct extraction formula.
|
||||
|
||||
### 6. SiLU after summing expert paths (math error)
|
||||
**File:** `nvfp4_mega_moe.py`
|
||||
**Bug:** The old grouped GEMM collapsed expert outputs into a weighted sum, then applied SiLU+Mul on the sum. `silu(Σ wᵢ·gateᵢ) * (Σ wᵢ·upᵢ) ≠ Σ wᵢ·silu(gateᵢ)·upᵢ`. The nonlinearity must happen per-expert-path.
|
||||
**Fix:** Slot-based dispatch — L1 GEMM returns per-slot output, SiLU+Mul applied per-slot, L2 GEMM per-slot, routing weights applied once at final `index_add_` scatter.
|
||||
|
||||
### 7. Routing weights applied twice
|
||||
**File:** `cutlass_nvfp4_gemm/kernel.py`
|
||||
**Bug:** `cutlass_grouped_nvfp4_gemm` applied `topk_weights` in its scatter loop. Called for both L1 and L2, each expert's contribution was scaled by `topk_weight²`.
|
||||
**Fix:** GEMM returns per-slot results with no routing weights. Single `y.index_add_(0, slot_token, slot_weight * l2_slots)` at the end.
|
||||
|
||||
### Diagnostic: constant-scale test (smoking gun for SF bugs)
|
||||
|
||||
When all scale factors are set to UE4M3(1.0):
|
||||
- **Cosine = 1.0000, MSE = 0.19** (expected FP4 quantization noise)
|
||||
|
||||
With real (variable) scale factors and the broken remap:
|
||||
- **Cosine = 0.83** → scales are misaligned, not fundamentally broken
|
||||
|
||||
After the fix with correct coordinate extraction:
|
||||
- **Cosine = 1.0000, MSE = 0.0** → perfect match with dequantized reference
|
||||
|
||||
---
|
||||
|
||||
## Build & Deploy (B200)
|
||||
On the B200:
|
||||
|
||||
```bash
|
||||
# On B200 host — CUTLASS must be cloned and mounted
|
||||
cd /root/nvidia-meeting/deepseek-v4-quant/
|
||||
cd /root/nvfp4-megamoe-kernel/tests
|
||||
source .venv/bin/activate
|
||||
|
||||
# Rebuild container (CUTLASS is host-mounted at /root/cutlass)
|
||||
KERNEL_CACHE_BUSTER=$(date +%s) docker compose build --no-cache
|
||||
docker compose up -d
|
||||
# Small standalone test
|
||||
python3 test_cutedsl.py
|
||||
|
||||
# Full layer 0 comparison with real weights
|
||||
python3 layertest.py
|
||||
```
|
||||
|
||||
The CUTLASS extension builds inside the container during `pip install` of the nvfp4-megamoe-kernel package. It needs:
|
||||
- CUDA 13.0 toolkit (in the vllm/vllm-openai:nightly image)
|
||||
- CUTLASS headers at `/root/cutlass/include/`
|
||||
- CCCL headers at `/usr/local/cuda-13.0/targets/x86_64-linux/include/cccl/`
|
||||
- Device with SM100 compute capability (B200)
|
||||
## Plan
|
||||
|
||||
---
|
||||
### Phase 1: Kernel ✅ DONE
|
||||
- CuTeDSL ScaledGroupedGemmKernel works with NVFP4
|
||||
- Bridge layer handles all tensor layout conversion
|
||||
- Full MoE pipeline (L1→SiLU→L2→scatter) produces cosine 0.989 vs BF16
|
||||
|
||||
## Known Issues / TODO
|
||||
### Phase 2: vLLM Integration (IN PROGRESS)
|
||||
- Wire `cutedsl/moe_pipeline.py` into the vLLM DeepSeek-V4 model
|
||||
- Replace `nvfp4_mega_moe_full()` call with `run_nvfp4_moe()`
|
||||
- Handle weight loading: checkpoint uint8 → float4_e2m1fn_x2 directly (no BF16 round-trip)
|
||||
- Remove C++ CUTLASS extension build from Dockerfile
|
||||
- Add CuTeDSL dependency to the Docker build
|
||||
|
||||
1. ~~**MoE dispatch is slow**~~ — Fixed. Slot-based `index_add_` replaces the Python double loop over tokens×topk. Routing weights applied once at final scatter.
|
||||
### Phase 3: Optimization
|
||||
- Load checkpoint FP4 bytes directly into float4_e2m1fn_x2 (skip BF16 round-trip)
|
||||
- Use checkpoint block scales directly (skip dequant→requant)
|
||||
- Explore larger tile sizes for better occupancy
|
||||
- Profile end-to-end inference on full model
|
||||
|
||||
2. **stage_activation is Python** — Re-quantization from L1 BF16 output to FP4 for L2 input runs in PyTorch. Should use the Triton staging kernel for speed and consistency with vLLM's built-in staging.
|
||||
|
||||
3. ~~**SF remap allocates every call**~~ — Fixed. SFB weight scales are prepacked into CUTLASS layout once (lazy, cached per layer). Only SFA (activation scales) remapped dynamically.
|
||||
|
||||
4. **Per-expert GEMM dispatch is serial Python loop** — The `cutlass_grouped_nvfp4_gemm` iterates over 48 experts in a Python `for` loop. Each iteration launches one CUTLASS GEMM. Could benefit from a true grouped GEMM kernel or CUDA-side expert dispatch.
|
||||
|
||||
---
|
||||
|
||||
## Environment Variables
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `MEGA_MOE_STATIC` | 0 | Set to 1 to skip MoE kernel entirely (return zeros) |
|
||||
| `MEGA_MOE_DEBUG` | 0 | Set to 1 for verbose logging |
|
||||
| `SKIP_ATTENTION` | 0 | Skip attention layers (debug) |
|
||||
|
||||
---
|
||||
|
||||
## Repos
|
||||
|
||||
- **Kernel:** `sweetapi.com/biondizzle/nvfp4-megamoe-kernel` (branch: master)
|
||||
- **Deployment:** `sweetapi.com/biondizzle/deepseek-v4-quant` (branch: modelopt-nvfp4)
|
||||
- **Local:** `~/dev/nvfp4-megamoe-kernel/`, `~/dev/deepseek-v4-quant/`
|
||||
### Phase 4: Production
|
||||
- Clean up debug artifacts
|
||||
- Remove old C++ kernel code
|
||||
- Add proper error handling and logging
|
||||
- Benchmark vs BF16 baseline
|
||||
|
||||
@@ -1,93 +0,0 @@
|
||||
# NVFP4 MegaMoE Kernel Rewrite — CuTeDSL
|
||||
|
||||
## Decision
|
||||
|
||||
The C++ CUTLASS kernel is fundamentally broken. The GEMM produces cosine 0.05
|
||||
with real data (despite SF remap = 0 errors and uniform tests passing). The
|
||||
root cause is likely in how CUTLASS's C++ API handles FP4 packing/tiling
|
||||
internally — something we can't easily debug or fix.
|
||||
|
||||
We're replacing it with NVIDIA's CuTeDSL approach (Python-based CUTLASS
|
||||
kernels compiled via MLIR → PTX). This is what the NVIDIA CUTLASS team
|
||||
recommends for Blackwell, and they have a working reference:
|
||||
|
||||
`cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/moe/torch_scaled_grouped_mm.py`
|
||||
|
||||
This is a 3900-line MoE scaled grouped GEMM that supports NVFP4 out of the box.
|
||||
|
||||
## Reference Files (copied to reference/)
|
||||
|
||||
- `reference/moe_torch_scaled_grouped_mm.py` — the kernel (3900 lines)
|
||||
- `reference/moe_moe_utils.py` — tensormap constructor
|
||||
- `reference/moe_moe_persistent_scheduler.py` — MoE tile scheduler
|
||||
- `reference/moe_moe_sched_extension.py` — scheduler extension for scaled GEMM
|
||||
- `reference/grouped_blockscaled_gemm.py` — simpler grouped blockscaled reference
|
||||
- `reference/dense_blockscaled_gemm_persistent.py` — dense (non-grouped) reference
|
||||
- `reference/blockscaled_layout.py` — SF layout utilities
|
||||
|
||||
## Architecture
|
||||
|
||||
### The Reference Kernel: ScaledGroupedGemmKernel
|
||||
|
||||
7-warp persistent kernel:
|
||||
- Warps 0-3: Epilogue (TMEM → RMEM → SMEM → GMEM)
|
||||
- Warp 4: MMA (tcgen05.mma.block_scale with SFA/SFB in TMEM)
|
||||
- Warp 5: TMA load (A, B, SFA, SFB from GMEM → SMEM)
|
||||
- Warp 6: Scheduler (MoE persistent tile scheduler)
|
||||
|
||||
Supports:
|
||||
- NVFP4 (Float4E2M1FN + Float8E4M3FN + sf_vec_size=16)
|
||||
- 2Dx3D scenario (our case): A(tokens, K) x B(experts, K, N) → C(tokens, N)
|
||||
- PyTorch interface via torch.nn.functional.scaled_grouped_mm
|
||||
|
||||
### Our Adaptation
|
||||
|
||||
We need to:
|
||||
1. Use ScaledGroupedGemmKernel directly (or with minimal adaptation)
|
||||
2. Replace our Python pipeline (nvfp4_mega_moe.py) to call the CuTeDSL kernel
|
||||
3. Handle the MoE-specific stuff (routing, gate/up fusion, SiLU, scatter)
|
||||
|
||||
The CuTeDSL kernel handles A/B/SFA/SFB tiling and MMA internally.
|
||||
We NO LONGER need:
|
||||
- Our custom SF remap kernel (CuTeDSL handles SF layouts natively)
|
||||
- Our C++ CUTLASS extension (cutlass_nvfp4_gemm.cu)
|
||||
- transform_nvfp4_weights_for_mega_moe (CuTeDSL uses the natural layout)
|
||||
|
||||
We STILL need:
|
||||
- stage_activation (BF16 → FP4 quantization)
|
||||
- The MoE routing logic (top-k selection, slot mapping)
|
||||
- Gate/up fusion with up_correction
|
||||
- SiLU activation
|
||||
- Scatter with routing weights
|
||||
|
||||
## Plan
|
||||
|
||||
### Phase 1: Get the reference kernel running standalone
|
||||
- Set up CuTeDSL build environment on B200
|
||||
- Run the reference example with NVFP4 config
|
||||
- Verify it produces correct output
|
||||
|
||||
### Phase 2: Integrate into our pipeline
|
||||
- Create a Python module that wraps ScaledGroupedGemmKernel
|
||||
- Replace cutlass_nvfp4_gemm calls with CuTeDSL kernel calls
|
||||
- Handle weight format (the CuTeDSL kernel expects raw FP4 + SF, no transpose)
|
||||
|
||||
### Phase 3: Full MoE pipeline
|
||||
- Wire up stage_activation → L1 GEMM → SiLU → stage_activation → L2 GEMM → scatter
|
||||
- Test with layertest.py (should get cosine ~0.995)
|
||||
|
||||
### Phase 4: vLLM integration
|
||||
- Update the vLLM patch to use the CuTeDSL kernel
|
||||
- Remove the C++ extension build from the Dockerfile
|
||||
- Test full inference
|
||||
|
||||
## Key Differences from Current Approach
|
||||
|
||||
| Current (C++ CUTLASS) | New (CuTeDSL) |
|
||||
|----------------------|---------------|
|
||||
| C++ .cu file | Python .py file |
|
||||
| Manual SF remap kernel | CuTe handles SF natively |
|
||||
| Manual weight transpose | CuTe handles layout |
|
||||
| pip install C++ extension | cute.compile() JIT |
|
||||
| Broken FP4 handling | Reference-verified |
|
||||
| No Blackwell pipeline | Full TMA/MMA/Epilogue overlap |
|
||||
@@ -1,109 +0,0 @@
|
||||
"""
|
||||
Test script for CUTLASS NVFP4 block-scaled GEMM.
|
||||
|
||||
Verifies the kernel against the dequantize-then-BF16 reference implementation.
|
||||
Run on the B200 server after building the extension.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Add parent directory to path
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
|
||||
from nvfp4_megamoe_kernel.nvfp4_dequant import unpack_e2m1_to_bf16
|
||||
|
||||
|
||||
def test_cutlass_nvfp4_gemm():
|
||||
"""Test single GEMM: C = A @ B^T with native NVFP4 block-scaled MMA."""
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA not available, skipping test")
|
||||
return
|
||||
|
||||
device = "cuda"
|
||||
M, N, K = 128, 128, 256 # Small test dimensions
|
||||
|
||||
# Create random E2M1-packed inputs
|
||||
A_packed = torch.randint(-128, 127, (M, K // 2), dtype=torch.int8, device=device)
|
||||
B_packed = torch.randint(-128, 127, (N, K // 2), dtype=torch.int8, device=device)
|
||||
|
||||
# Create random UE4M3 block16 scales
|
||||
A_scales = torch.randn(M, K // 16, dtype=torch.float8_e4m3fn, device=device)
|
||||
B_scales = torch.randn(N, K // 16, dtype=torch.float8_e4m3fn, device=device)
|
||||
# Make positive (UE4M3 is unsigned)
|
||||
A_scales = A_scales.abs().clamp(min=0.0625, max=448.0)
|
||||
B_scales = B_scales.abs().clamp(min=0.0625, max=448.0)
|
||||
|
||||
# Reference: dequantize then BF16 GEMM
|
||||
A_bf16 = unpack_e2m1_to_bf16(A_packed, A_scales)
|
||||
B_bf16 = unpack_e2m1_to_bf16(B_packed, B_scales)
|
||||
C_ref = torch.matmul(A_bf16, B_bf16.t())
|
||||
|
||||
# CUTLASS expects B in column-major: (K_half, N) for weights, (sf_k, N) for scales
|
||||
B_packed_cm = B_packed.t().contiguous()
|
||||
B_scales_cm = B_scales.t().contiguous()
|
||||
|
||||
# CUTLASS native NVFP4 GEMM
|
||||
try:
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_nvfp4_blockscaled_gemm
|
||||
C_cutlass = cutlass_nvfp4_blockscaled_gemm(A_packed, A_scales, B_packed_cm, B_scales_cm)
|
||||
|
||||
# Compare (NVFP4 has low precision, so use loose tolerance)
|
||||
diff = (C_cutlass.float() - C_ref.float()).abs()
|
||||
max_diff = diff.max().item()
|
||||
mean_diff = diff.mean().item()
|
||||
print(f"CUTLASS vs Reference: max_diff={max_diff:.4f}, mean_diff={mean_diff:.4f}")
|
||||
print(f"CUTLASS output shape: {C_cutlass.shape}, dtype: {C_cutlass.dtype}")
|
||||
print(f"Reference output shape: {C_ref.shape}, dtype: {C_ref.dtype}")
|
||||
|
||||
if max_diff < 1.0: # Loose tolerance for 4-bit arithmetic
|
||||
print("✓ CUTLASS NVFP4 GEMM test PASSED")
|
||||
else:
|
||||
print("✗ CUTLASS NVFP4 GEMM test FAILED (max_diff too high)")
|
||||
except Exception as e:
|
||||
print(f"✗ CUTLASS NVFP4 GEMM test FAILED with exception: {e}")
|
||||
print("Falling back to reference implementation (dequantize+BF16)")
|
||||
|
||||
|
||||
def test_grouped_gemm():
|
||||
"""Test grouped expert GEMM for MoE dispatch."""
|
||||
if not torch.cuda.is_available():
|
||||
print("CUDA not available, skipping test")
|
||||
return
|
||||
|
||||
device = "cuda"
|
||||
num_tokens = 64
|
||||
K = 256
|
||||
N = 128
|
||||
E = 4 # Small number of experts for testing
|
||||
top_k = 2
|
||||
|
||||
# Create inputs
|
||||
x_packed = torch.randint(-128, 127, (num_tokens, K // 2), dtype=torch.int8, device=device)
|
||||
x_scales = torch.randn(num_tokens, K // 16, dtype=torch.float8_e4m3fn, device=device).abs().clamp(min=0.0625, max=448.0)
|
||||
# Weights in column-major for CUTLASS: (E, K_half, N) and (E, sf_k, N)
|
||||
weights = torch.randint(-128, 127, (E, K // 2, N), dtype=torch.int8, device=device)
|
||||
weight_scales = torch.randn(E, K // 16, N, dtype=torch.float8_e4m3fn, device=device).abs().clamp(min=0.0625, max=448.0)
|
||||
topk_ids = torch.randint(0, E, (num_tokens, top_k), dtype=torch.int32, device=device)
|
||||
topk_weights = torch.rand(num_tokens, top_k, dtype=torch.float32, device=device)
|
||||
topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
|
||||
|
||||
try:
|
||||
from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import cutlass_grouped_nvfp4_gemm
|
||||
output = cutlass_grouped_nvfp4_gemm(
|
||||
x_packed, x_scales, weights, weight_scales, topk_ids, topk_weights
|
||||
)
|
||||
print(f"Grouped GEMM output shape: {output.shape}, dtype: {output.dtype}")
|
||||
print("✓ Grouped NVFP4 GEMM test PASSED")
|
||||
except Exception as e:
|
||||
print(f"✗ Grouped NVFP4 GEMM test FAILED: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("=" * 60)
|
||||
print("CUTLASS NVFP4 Block-Scaled GEMM Tests")
|
||||
print("=" * 60)
|
||||
test_cutlass_nvfp4_gemm()
|
||||
print()
|
||||
test_grouped_gemm()
|
||||
199
vllm/nvfp4_cutedsl.py
Normal file
199
vllm/nvfp4_cutedsl.py
Normal file
@@ -0,0 +1,199 @@
|
||||
"""
|
||||
vLLM integration for the CuTeDSL NVFP4 MoE kernel.
|
||||
|
||||
This module provides the interface between the vLLM DeepSeek-V4 model
|
||||
and the CuTeDSL NVFP4 kernel pipeline. It replaces the broken C++
|
||||
CUTLASS path with the working CuTeDSL path.
|
||||
|
||||
The key change: instead of `nvfp4_mega_moe_full` (C++ kernel), we call
|
||||
`run_nvfp4_moe` (CuTeDSL kernel). Weight preparation and tensor layout
|
||||
conversion is handled by `cutedsl/bridge.py`.
|
||||
|
||||
Usage in deepseek_v4.py:
|
||||
from vllm.nvfp4_cutedsl import CuTeDSLMoERunner
|
||||
|
||||
# In DeepseekV4MegaMoEExperts:
|
||||
self.moe_runner = CuTeDSLMoERunner(...)
|
||||
self.moe_runner.run(hidden_states, topk_weights, topk_ids, y)
|
||||
"""
|
||||
import torch
|
||||
|
||||
from cutedsl.bridge import (
|
||||
quantize_to_nvfp4,
|
||||
quantize_weight_to_nvfp4,
|
||||
make_b_k_major,
|
||||
compute_expert_offsets,
|
||||
assemble_scales_2d_side,
|
||||
assemble_scales_3d_side,
|
||||
run_nvfp4_grouped_gemm,
|
||||
)
|
||||
|
||||
|
||||
class CuTeDSLMoERunner:
|
||||
"""Manages NVFP4 MoE execution via the CuTeDSL kernel.
|
||||
|
||||
Replaces the old `nvfp4_mega_moe_full` + `SymmBuffer` pipeline.
|
||||
The CuTeDSL kernel manages its own workspace internally.
|
||||
|
||||
Weight format: checkpoint uint8 → dequantize to BF16 → re-quantize
|
||||
to float4_e2m1fn_x2. Future optimization: load checkpoint bytes
|
||||
directly into float4_e2m1fn_x2 tensors.
|
||||
"""
|
||||
|
||||
def __init__(self, num_experts, hidden_size, intermediate_size, device="cuda"):
|
||||
self.num_experts = num_experts
|
||||
self.hidden_size = hidden_size
|
||||
self.intermediate_size = intermediate_size
|
||||
self.device = device
|
||||
|
||||
# Prepared weights (set by prepare_weights)
|
||||
self.l1_fp4 = None # list of (K//2, N) float4_e2m1fn_x2 per expert
|
||||
self.l1_sf = None # list of (K//16, N) float8_e4m3fn per expert
|
||||
self.l1_gs = None # list of float32 per expert
|
||||
self.l2_fp4 = None
|
||||
self.l2_sf = None
|
||||
self.l2_gs = None
|
||||
|
||||
def prepare_weights_from_dequantized(self, l1_weights_bf16, l2_weights_bf16):
|
||||
"""Prepare NVFP4 weights from dequantized BF16 tensors.
|
||||
|
||||
This is the current path: checkpoint uint8 → dequantize → BF16 → re-quantize.
|
||||
|
||||
Args:
|
||||
l1_weights_bf16: list of (6144, hidden_size) BF16 tensors (gate+up fused)
|
||||
l2_weights_bf16: list of (hidden_size, intermediate//2) BF16 tensors (down)
|
||||
"""
|
||||
self.l1_fp4, self.l1_sf, self.l1_gs = [], [], []
|
||||
self.l2_fp4, self.l2_sf, self.l2_gs = [], [], []
|
||||
|
||||
for l1_w, l2_w in zip(l1_weights_bf16, l2_weights_bf16):
|
||||
# L1: (6144, hidden) → transpose to (hidden, 6144) for K=hidden packed dim
|
||||
l1_w_t = l1_w.T
|
||||
w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l1_w_t)
|
||||
self.l1_fp4.append(w_fp4)
|
||||
self.l1_sf.append(w_sf)
|
||||
self.l1_gs.append(w_gs)
|
||||
|
||||
# L2: (hidden, intermediate//2) → already (K=intermediate//2, N=hidden)
|
||||
# Wait — down_proj is (hidden, intermediate//2), which is (N_out, K_in)
|
||||
# For the GEMM: B is (K, N) where K=intermediate, N=hidden
|
||||
l2_w_t = l2_w.T # (intermediate//2, hidden) → K=intermediate, N=hidden
|
||||
w_fp4, w_sf, w_gs = quantize_weight_to_nvfp4(l2_w_t)
|
||||
self.l2_fp4.append(w_fp4)
|
||||
self.l2_sf.append(w_sf)
|
||||
self.l2_gs.append(w_gs)
|
||||
|
||||
def run(self, hidden_states, topk_weights, topk_ids, expert_indices=None):
|
||||
"""Run the full NVFP4 MoE forward pass.
|
||||
|
||||
Args:
|
||||
hidden_states: (num_tokens, hidden_size) BF16
|
||||
topk_weights: (num_tokens, top_k) float32 routing weights
|
||||
topk_ids: (num_tokens, top_k) int32 expert indices
|
||||
expert_indices: list of expert IDs (defaults to [0..num_experts-1])
|
||||
|
||||
Returns:
|
||||
(num_tokens, hidden_size) BF16 output
|
||||
"""
|
||||
num_tokens, hidden_size = hidden_states.shape
|
||||
top_k = topk_ids.shape[1]
|
||||
device = hidden_states.device
|
||||
|
||||
if expert_indices is None:
|
||||
expert_indices = list(range(self.num_experts))
|
||||
|
||||
# ── Build slot-based routing ──
|
||||
expert_token_lists = {e: [] for e in expert_indices}
|
||||
for t in range(num_tokens):
|
||||
for k in range(top_k):
|
||||
e = topk_ids[t, k].item()
|
||||
if e in expert_token_lists:
|
||||
expert_token_lists[e].append(t)
|
||||
|
||||
tokens_per_expert = [len(expert_token_lists[e]) for e in expert_indices]
|
||||
|
||||
# Skip if no tokens routed
|
||||
if sum(tokens_per_expert) == 0:
|
||||
return torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
|
||||
# Slot-major activation
|
||||
slot_hidden = torch.cat([
|
||||
hidden_states[expert_token_lists[e]] for e in expert_indices
|
||||
], dim=0)
|
||||
|
||||
expert_offsets = compute_expert_offsets(tokens_per_expert, len(expert_indices))
|
||||
|
||||
# ════════════════════════════════════════════════════════════
|
||||
# L1: gate + up (NVFP4 × NVFP4 → BF16)
|
||||
# ════════════════════════════════════════════════════════════
|
||||
x_fp4, x_sf, x_igs = quantize_to_nvfp4(slot_hidden)
|
||||
|
||||
l1_mat_b = make_b_k_major(torch.stack(self.l1_fp4))
|
||||
|
||||
x_sf_parts = []
|
||||
offset = 0
|
||||
for tpe in tokens_per_expert:
|
||||
x_sf_parts.append(x_sf[offset:offset+tpe])
|
||||
offset += tpe
|
||||
l1_scale_a = assemble_scales_2d_side(x_sf_parts)
|
||||
l1_scale_b = assemble_scales_3d_side(self.l1_sf)
|
||||
|
||||
l1_gsa = torch.tensor([x_igs] * len(expert_indices), dtype=torch.float32, device=device)
|
||||
l1_gsb = torch.tensor(self.l1_gs, dtype=torch.float32, device=device)
|
||||
|
||||
l1_out = run_nvfp4_grouped_gemm(
|
||||
mat_a=x_fp4, mat_b=l1_mat_b,
|
||||
scale_a=l1_scale_a, scale_b=l1_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=l1_gsa, global_scale_b=l1_gsb,
|
||||
)
|
||||
|
||||
# ════════════════════════════════════════════════════════════
|
||||
# SiLU(gate) * up (BF16)
|
||||
# ════════════════════════════════════════════════════════════
|
||||
half = self.intermediate_size // 2
|
||||
gate = l1_out[:, :half]
|
||||
up = l1_out[:, half:]
|
||||
activated = torch.nn.functional.silu(gate) * up
|
||||
|
||||
# ════════════════════════════════════════════════════════════
|
||||
# L2: down (NVFP4 × NVFP4 → BF16)
|
||||
# ════════════════════════════════════════════════════════════
|
||||
l2_x_fp4, l2_x_sf, l2_x_igs = quantize_to_nvfp4(activated)
|
||||
|
||||
l2_mat_b = make_b_k_major(torch.stack(self.l2_fp4))
|
||||
|
||||
l2_sf_parts = []
|
||||
offset = 0
|
||||
for tpe in tokens_per_expert:
|
||||
l2_sf_parts.append(l2_x_sf[offset:offset+tpe])
|
||||
offset += tpe
|
||||
l2_scale_a = assemble_scales_2d_side(l2_sf_parts)
|
||||
l2_scale_b = assemble_scales_3d_side(self.l2_sf)
|
||||
|
||||
l2_gsa = torch.tensor([l2_x_igs] * len(expert_indices), dtype=torch.float32, device=device)
|
||||
l2_gsb = torch.tensor(self.l2_gs, dtype=torch.float32, device=device)
|
||||
|
||||
l2_out = run_nvfp4_grouped_gemm(
|
||||
mat_a=l2_x_fp4, mat_b=l2_mat_b,
|
||||
scale_a=l2_scale_a, scale_b=l2_scale_b,
|
||||
expert_offsets=expert_offsets,
|
||||
global_scale_a=l2_gsa, global_scale_b=l2_gsb,
|
||||
)
|
||||
|
||||
# ════════════════════════════════════════════════════════════
|
||||
# Scatter → final output
|
||||
# ════════════════════════════════════════════════════════════
|
||||
y = torch.zeros(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
|
||||
|
||||
slot_idx = 0
|
||||
for e in expert_indices:
|
||||
for t in expert_token_lists[e]:
|
||||
for k in range(top_k):
|
||||
if topk_ids[t, k].item() == e:
|
||||
w = topk_weights[t, k].item()
|
||||
y[t] += w * l2_out[slot_idx]
|
||||
break
|
||||
slot_idx += 1
|
||||
|
||||
return y
|
||||
Reference in New Issue
Block a user