diff --git a/DEBUG_LOG.md b/DEBUG_LOG.md deleted file mode 100644 index 8b3da246..00000000 --- a/DEBUG_LOG.md +++ /dev/null @@ -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? diff --git a/README.md b/README.md index d5843cf4..da2e9922 100644 --- a/README.md +++ b/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<16, 4>> -Atom Stride: Stride, 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 diff --git a/REWRITE_PLAN.md b/REWRITE_PLAN.md deleted file mode 100644 index dc8ab005..00000000 --- a/REWRITE_PLAN.md +++ /dev/null @@ -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 | diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py deleted file mode 100644 index fcfe2a1b..00000000 --- a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py +++ /dev/null @@ -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() diff --git a/vllm/nvfp4_cutedsl.py b/vllm/nvfp4_cutedsl.py new file mode 100644 index 00000000..fa5456bc --- /dev/null +++ b/vllm/nvfp4_cutedsl.py @@ -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