diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000..a4df20e9 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,6 @@ +docker/ +scripts/ +*.egg-info/ +.git/ +.gitignore +README.md diff --git a/.gitignore b/.gitignore new file mode 100644 index 00000000..9f7983e8 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +__pycache__/ +*.pyc +*.egg-info/ diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 00000000..8155124b --- /dev/null +++ b/Dockerfile @@ -0,0 +1,50 @@ +# DeepSeek V4 NVFP4 vLLM + CUTLASS NVFP4 Mega MoE Kernel +FROM vllm/vllm-openai:nightly-x86_64 + +# Remove broken nixl_ep (built against CUDA 12, image is CUDA 13) +RUN pip uninstall -y nixl-ep; rm -rf /usr/local/lib/python3.12/dist-packages/nixl_ep + +RUN apt-get update && apt-get install -y git screen cmake libcusolver-dev-13-0 libcusparse-dev-13-0 libcublas-dev-13-0 libcurand-dev-13-0 libcufft-dev-13-0 libnvjitlink-dev-13-0 && rm -rf /var/lib/apt/lists/* + +# Remove the broken symlink if it exists +RUN rm -f /usr/local/cuda/lib64/libcudart.so.12 + +ENV CUDA_HOME=/usr/local/cuda +ENV TORCH_CUDA_ARCH_LIST="10.0" + +# Clone latest CUTLASS (has NVFP4 block-scaled MMA support) +RUN git clone --depth 1 https://github.com/NVIDIA/cutlass.git /root/cutlass + +ARG CACHE_BUSTER=${TIMESTAMP} + +# Copy and install the NVFP4 mega_moe kernel (from this repo) +COPY src/ /root/nvfp4-megamoe-kernel/src/ +COPY pyproject.toml /root/nvfp4-megamoe-kernel/pyproject.toml +RUN cd /root/nvfp4-megamoe-kernel && pip install -e . + +# Build the CUTLASS NVFP4 block-scaled GEMM extension +RUN cd /root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm && \ + mkdir -p cutlass_nvfp4_gemm && \ + CUTLASS_INCLUDE_DIR=/root/cutlass/include \ + TORCH_CUDA_ARCH_LIST=10.0 \ + python3 setup.py build_ext --inplace + +# Install TileLang (for potential future use) +RUN pip install tilelang + +ENV PYTHONPATH="/root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm:/root/nvfp4-megamoe-kernel:${PYTHONPATH}" + +# Patch vLLM — overwrite model files and register architecture +ARG VLLM_MODELS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models +ARG VLLM_LAYERS_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/layers + +COPY vllm/patches/deepseek_v4.py ${VLLM_MODELS_DIR}/deepseek_v4.py +COPY vllm/patches/deepseek_v4_attention.py ${VLLM_LAYERS_DIR}/deepseek_v4_attention.py + +RUN sed -i 's/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),\n "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),/' \ + ${VLLM_MODELS_DIR}/registry.py + +# Verify +RUN python3 -c "import torch; import cutlass_nvfp4_gemm._C; print('CUTLASS NVFP4 OK')" && \ + python3 -c "import vllm; print('vLLM OK')" && \ + python3 -c "import nvfp4_megamoe_kernel; print('NVFP4 kernel OK')" diff --git a/README.md b/README.md index 7b819617..d5843cf4 100644 --- a/README.md +++ b/README.md @@ -1,40 +1,361 @@ -# NVFP4 Mega MoE Kernel — Mojo Rewrite +# nvfp4-megamoe-kernel -Rewrite of the DeepGEMM `fp8_nvfp4_mega_moe` kernel in Mojo. +Native NVFP4 block-scaled MoE kernel for DeepSeek-V4-Pro on NVIDIA Blackwell (SM100). -## Why Mojo? -- Python-like syntax, C-level performance -- Direct GPU programming without PTX inline asm -- Safer than CUDA C++ (ownership, borrowing) -- Better ergonomics for complex kernel development +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. + +--- ## Architecture -The kernel performs NVFP4 (E2M1 + UE4M3 block16 scales) matrix multiply -for MoE (Mixture of Experts) with expert parallelism across NVLink. +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. -### Key operations: -1. **Staging** — quantize BF16 activation to FP4 (E2M1) with UE8M0 scales -2. **TMA load** — load packed FP4 weights and UE4M3 scales from global memory -3. **UMMA** — `mxf4nvf4` matrix multiply with block scaling -4. **Epilogue** — quantize L1 output (BF16 → FP4 + UE4M3 scales for L2) -5. **NVLink sync** — cross-rank barrier and buffer management +### The MoE Forward Pass -### NVFP4 specifics (vs MXFP4): -- group_size=16 (UE4M3 block scales), not group_size=32 (UE8M0) -- 2 SF K-columns per BLOCK_K (128/16/4=2), not 1 -- Weights are E2M1 packed int8 (2 values per byte) -- `mxf4nvf4` UMMA instruction with `scale_vec::4X` - -## Structure ``` -src/ - mega_moe.mojo — main kernel entry point - staging.mojo — activation quantization (BF16 → FP4) - tma.mojo — TMA descriptor creation and copy - umma.mojo — UMMA descriptor and MMA operations - epilogue.mojo — output quantization and TMA store - barrier.mojo — NVLink cluster sync and symm buffer - layout.mojo — weight transformation and SF layout - utils.mojo — math helpers, UE4M3 packing +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) ``` + +### Slot-Based Dispatch + +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). + +**Why slots?** Two bugs in the previous approach: + +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. + +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. + +The slot approach fixes both: SiLU+Mul happens per-slot, and routing weights are applied exactly once at the final `index_add_` scatter. + +### SFB (Weight Scale Factors) — Remapped Per-Call, NOT Cached + +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. + +⚠️ **DO NOT ADD A PREPACK CACHE FOR SFB.** Previous attempts caused critical issues: + +| 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 | + +The per-call remap costs microseconds. The cache cost was hours of debugging. Don't repeat this mistake. + +--- + +### vLLM Startup Sequence (how our code plugs in) + +``` +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 +``` + +--- + +## File Map + +``` +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 +``` + +### What each file does (in call order) + +| 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) + +```bash +# On B200 host — CUTLASS must be cloned and mounted +cd /root/nvidia-meeting/deepseek-v4-quant/ + +# Rebuild container (CUTLASS is host-mounted at /root/cutlass) +KERNEL_CACHE_BUSTER=$(date +%s) docker compose build --no-cache +docker compose up -d +``` + +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) + +--- + +## Known Issues / TODO + +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. + +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/` diff --git a/build_and_run.sh b/build_and_run.sh new file mode 100755 index 00000000..4431cf04 --- /dev/null +++ b/build_and_run.sh @@ -0,0 +1,22 @@ +#!/bin/bash +set -euo pipefail + +cd "$(dirname "$0")" + +# Bust the Docker build cache by injecting a timestamp +TIMESTAMP=$(date +%s) +sed -i -E "s/ARG CACHE_BUSTER=.*/ARG CACHE_BUSTER=${TIMESTAMP}/" Dockerfile + +echo "=== Stopping existing container ===" +docker compose down 2>/dev/null || true + +echo "=== Building (no cache) ===" +docker compose build --no-cache + +echo "=== Starting ===" +docker compose up -d + +# Restore Dockerfile so git diff stays clean +sed -i -E "s/ARG CACHE_BUSTER=.*/ARG CACHE_BUSTER=\${TIMESTAMP}/" Dockerfile + +echo "=== Done. Container: $(docker compose ps -q) ===" diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 00000000..a9c2c547 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,37 @@ +services: + vllm: + build: + context: . + dockerfile: Dockerfile + ports: + - "8000:8000" + environment: + - OMP_NUM_THREADS=128 + - CUDA_LAUNCH_BLOCKING=0 + - TORCH_SHOW_CPP_STACKTRACES=0 + - MEGA_MOE_DEBUG=0 + - MEGA_MOE_STATIC=0 + - NVFP4_DEBUG=0 + - NVFP4_DEBUG_SYNC=0 + - SKIP_ATTENTION=0 + - MEGA_MOE_USE_CUTLASS=1 + - DG_JIT_DEBUG=0 + - DEEP_GEMM_JIT_DEBUG=0 + command: + - /model + - --trust-remote-code + - --enable-expert-parallel + - --tensor-parallel-size=8 + - --enforce-eager + - --tokenizer-mode=deepseek_v4 + - --host=0.0.0.0 + - --port=8000 + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + volumes: + - /root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4:/model:ro diff --git a/src/layout.mojo b/src/layout.mojo deleted file mode 100644 index af284d20..00000000 --- a/src/layout.mojo +++ /dev/null @@ -1,32 +0,0 @@ -""" -NVFP4 weight transformation and SF layout utilities. - -Port of deep_gemm.mega.transform_nvfp4_weights_for_mega_moe -""" - -from math import ceil_div - -fn fold_global_scale_into_block_scales( - weight_scale: Tensor[float8_e4m3fn], # (N, K//16) UE4M3 block scales - weight_scale_2: Tensor[float32], # (num_logical,) or scalar global scale - logical_widths: List[int], # per-logical-weight row counts -) -> Tensor[float32]: - """Fold global scale into block scales: UE4M3 * FP32 -> FP32""" - # Convert UE4M3 to float32, multiply by global scale - # For MergedColumnParallelLinear, expand per-logical global scale - ... - -fn pack_ue4m3_to_int32(sf: Tensor[float8_e4m3fn]) -> Tensor[int32]: - """Pack 4 UE4M3 values (4 bytes) into one int32 for DeepGEMM TMA""" - # View as uint8, pack 4 consecutive bytes into int32 - ... - -fn transform_sf_into_required_layout( - sf_mn: Tensor[int32], # MN-major packed SF - N: int, K: int, - recipe: Tuple[int, int], # (gran_mn, gran_k) - num_groups: int, -) -> Tensor[int32]: - """Transform SF into TMA-aligned UTCCP layout for DeepGEMM""" - # Call into DeepGEMM's C++ layout transform - ... diff --git a/src/mega_moe.mojo b/src/mega_moe.mojo deleted file mode 100644 index d9056f5d..00000000 --- a/src/mega_moe.mojo +++ /dev/null @@ -1,24 +0,0 @@ -""" -NVFP4 Mega MoE Kernel — Mojo Rewrite - -This is a from-scratch rewrite of the DeepGEMM fp8_nvfp4_mega_moe kernel. -The CUDA C++ version crashes on B200 with CUDA_ERROR_LAUNCH_FAILED in the -SM100_MMA_MXF4NVF4 instruction. This Mojo rewrite aims to: -1. Produce a correct, working NVFP4 mega_moe kernel -2. Be more maintainable than 1200+ lines of CUDA template metaprogramming -3. Leverage Mojo's GPU programming model for cleaner TMA/UMMA integration - -NVFP4 format: -- Weights: E2M1 packed (2 x 4-bit values per byte), int8 container -- Block scales: UE4M3 (float8_e4m3fn), group_size=16 -- Global scale: float32 scalar -- Activation: FP8 e4m3fn with UE8M0 per-token scales - -Key differences from MXFP4 (group_size=32, UE8M0): -- 2 SF K-columns per BLOCK_K (128/16/4=2) instead of 1 -- mxf4nvf4 UMMA instruction with scale_vec::4X -- Different weight transformation (gran_k=16 vs gran_k=32) -""" - -# TODO: Implement as Mojo GPU kernel -# Waiting for Mojo GPU programming docs / SDK setup diff --git a/src/nvfp4_blockscaled_gemm.py b/src/nvfp4_blockscaled_gemm.py deleted file mode 100644 index 135c65bc..00000000 --- a/src/nvfp4_blockscaled_gemm.py +++ /dev/null @@ -1,256 +0,0 @@ -""" -NVFP4 Block-Scaled GEMM — TileLang -2CTA persistent kernel for Blackwell SM100 - -Adapted from tilelang/examples/blockscaled_gemm_sm100/gemm_mxfp8_blockscaled_1d1d.py - -Key NVFP4 differences from MXFP8: -- sf_granularity_k=16 (UE4M3 block16) instead of 128 (UE8M0) -- 2 SF uint32 words per BLOCK_K instead of 1 (sf_load_period=1, not 4) -- sf_a_id/sf_b_id cycle through 0,1,2,3 within each SF word -- Must call tcgen05_gemm_blockscaled twice per K-block with different SF words -- in_dtype="float4_e2m1fnx2" (packed FP4) instead of "float8_e4m3fn" (FP8) -""" - -import torch -import tilelang -import tilelang.language as T -from tilelang.carver.arch import driver -from tilelang.profiler import do_bench - -# NVFP4 constants -NVFP4_SF_GRAN_K = 16 # UE4M3 group_size -NVFP4_SF_PACK = 4 # 4 UE4M3 values per uint32 -NVFP4_SF_K_PER_WORD = NVFP4_SF_GRAN_K * NVFP4_SF_PACK # 64 K-elements per uint32 word -# For BLOCK_K=128: 128/64 = 2 SF words per K-block -# Each word covers 4 UMMA sub-atoms (sf_id 0-3) - - -@tilelang.jit -def nvfp4_blockscaled_gemm_2cta_persistent( - A, # (M, K) packed FP4 - B, # (N, K) packed FP4 (K-major for NN) - SFA, # (sf_k_groups * M) uint32 packed UE4M3 - SFB, # (sf_k_groups * N) uint32 packed UE4M3 - block_M=128, - block_N=256, - block_K=128, - in_dtype="float4_e2m1fnx2", - out_dtype="bfloat16", - accum_dtype="float32", - num_stages=2, - sf_granularity_k=NVFP4_SF_GRAN_K, - use_tma_store=True, - store_block_N=64, -): - M, N, K = T.const("M, N, K") - - assert block_M == 128 - assert block_N == 256 - assert block_K == 128 - - half_N = block_N // 2 - k_iters = T.ceildiv(K, block_K) - - # NVFP4 SF layout: - # sf_granularity_k=16, pack_factor=4 → 64 K-elements per uint32 - # BLOCK_K=128 → 2 SF words per K-block - # sf_load_period=1 (load every K-block, unlike MXFP8 which loads every 4) - sf_words_per_k_block = block_K // (sf_granularity_k * NVFP4_SF_PACK) # 2 - sf_k_groups = T.ceildiv(K, sf_granularity_k * NVFP4_SF_PACK) # total SF groups across K - - A: T.Tensor[[M, K], in_dtype] - B: T.Tensor[[N, K], in_dtype] - SFA: T.Tensor[[sf_k_groups * M], T.uint32] - SFB: T.Tensor[[sf_k_groups * N], T.uint32] - C = T.empty((M, N), out_dtype) - - sm_num = driver.get_num_sms() - num_clusters = sm_num // 2 - m_blocks = T.ceildiv(M, block_M) - m_clusters = m_blocks // 2 - n_blocks = T.ceildiv(N, block_N) - waves = T.ceildiv(m_blocks * n_blocks, sm_num) - group_size = 16 - assert n_blocks % (2 * group_size) == 0 - - with T.Kernel(sm_num, threads=256, cluster_dims=2) as (block_id): - cta_id = T.block_rank_in_cluster() - T.assume(cta_id < 2) - - # Shared memory — FP4 packed (K is halved because 2 values per byte) - A_shared = T.alloc_shared((num_stages, block_M, block_K // 2), "uint8") - B_shared = T.alloc_shared((num_stages, block_K // 2, half_N), "uint8") - - # Shared memory for SF — 2 uint32 words per K-block for NVFP4 - SFA_shared = T.alloc_shared((num_stages, block_M, sf_words_per_k_block), T.uint32) - SFB_shared = T.alloc_shared((num_stages, block_N, sf_words_per_k_block), T.uint32) - - # Tensor memory - C_tmem = T.alloc_tmem([block_M, block_N], accum_dtype) - SFA_tmem = T.alloc_tmem([block_M, block_M // 128 * 4], T.uint32) - SFB_tmem = T.alloc_tmem([block_M, block_N // 128 * 4], T.uint32) - - C_local = T.alloc_fragment((block_M, block_N), accum_dtype) - C_local_cast = T.alloc_fragment((block_M, block_N), out_dtype) - C_shared = T.alloc_shared((block_M, store_block_N), out_dtype) - - loaded = T.alloc_barrier([32] * num_stages) - with_sf_full = T.alloc_cluster_barrier([32 * 2] * num_stages) - consumed = T.alloc_cluster_barrier([1] * num_stages) - tmem_full = T.alloc_cluster_barrier([1]) - tmem_empty = T.alloc_cluster_barrier([128 * 2]) - - tx = T.get_thread_binding() - warp_idx = tx // 32 - - if warp_idx == 0: - # Warp 0: TMA load - for w in T.serial(waves): - cluster_id = block_id // 2 - tile_id = num_clusters * w + cluster_id - bx_cluster = (tile_id // group_size) % m_clusters - bx = bx_cluster * 2 + cta_id - by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size - - if bx * block_M < M and by * block_N < N: - for k in T.serial(k_iters): - phase = w * k_iters + k - stage = phase % num_stages - parity = (phase // num_stages) & 1 - - T.mbarrier_wait_parity(consumed[stage], parity ^ 1) - - # TMA load A (packed FP4) - T.tma_copy( - A[bx * block_M, k * block_K], - A_shared[stage, :, :], - barrier=loaded[stage], - ) - # TMA load B (packed FP4) - T.tma_copy( - B[k * block_K, by * block_N + cta_id * half_N], - B_shared[stage, :, :], - barrier=loaded[stage], - ) - # TMA load SFA — 2 words per K-block for NVFP4 - # Word 0 covers K[k*128 : k*128+64] - # Word 1 covers K[k*128+64 : k*128+128] - for sf_word in T.unroll(sf_words_per_k_block): - sf_group = k * sf_words_per_k_block + sf_word - T.tma_copy( - SFA[sf_group * M + bx * block_M : - sf_group * M + (bx + 1) * block_M], - SFA_shared[stage, :, sf_word], - barrier=loaded[stage], - ) - # TMA load SFB — 2 words per K-block - for sf_word in T.unroll(sf_words_per_k_block): - sf_group = k * sf_words_per_k_block + sf_word - T.tma_copy( - SFB[sf_group * N + by * block_N : - sf_group * N + (by + 1) * block_N], - SFB_shared[stage, :, sf_word], - barrier=loaded[stage], - ) - T.mbarrier_arrive(loaded[stage]) - - elif warp_idx == 1 and cta_id == 0: - # Warp 1: MMA issue + UTCCP - for w in T.serial(waves): - cluster_id = block_id // 2 - tile_id = num_clusters * w + cluster_id - bx_cluster = (tile_id // group_size) % m_clusters - bx = bx_cluster * 2 + cta_id - by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size - - if bx * block_M < M and by * block_N < N: - for k in T.serial(k_iters): - phase = w * k_iters + k - stage = phase % num_stages - parity = (phase // num_stages) & 1 - - T.mbarrier_wait_parity(with_sf_full[stage], parity) - - # NVFP4: 2 SF words per K-block → 2 sub-GEMM calls - # Each sub-GEMM covers 64 K-elements (BLOCK_K/2) - for sf_word in T.unroll(sf_words_per_k_block): - # UTCCP: copy SF word from shared to tensor memory - T.tcgen05_cp_warpx4( - SFA_shared[stage, :, sf_word], - SFA_tmem, - use_2cta=True, - ) - T.tcgen05_cp_warpx4( - SFB_shared[stage, :, sf_word], - SFB_tmem, - use_2cta=True, - ) - - # Block-scaled GEMM - # For NVFP4: sf_a_id selects which of 4 packed UE4M3 - # values in the uint32 word to use. - # With sf_granularity_k=16 and UMMA_K=64: - # 4 SF values per UMMA atom, sf_id=0 covers all 4 - # (the hardware auto-cycles through 0,1,2,3 internally) - T.tcgen05_gemm_blockscaled( - A_shared[stage, :, :], - B_shared[stage, :, :], - C_tmem, - SFA_tmem, - SFB_tmem, - transpose_B=True, - mbar=consumed[stage], - clear_accum=(k == 0 and sf_word == 0), - sf_a_id=0, - sf_b_id=0, - use_2cta=True, - ) - - T.tcgen05_mma_arrive(tmem_full, arrive_2cta=True) - - elif warp_idx == 2: - # Warp 2: SF transpose - for w in T.serial(waves): - cluster_id = block_id // 2 - tile_id = num_clusters * w + cluster_id - bx_cluster = (tile_id // group_size) % m_clusters - bx = bx_cluster * 2 + cta_id - by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size - - if bx * block_M < M and by * block_N < N: - for k in T.serial(k_iters): - phase = w * k_iters + k - stage = phase % num_stages - parity = (phase // num_stages) & 1 - - T.mbarrier_wait_parity(loaded[stage], parity) - # Transpose all SF words for this K-block - for sf_word in T.unroll(sf_words_per_k_block): - T.tcgen05_sf_warp_transpose(SFA_shared[stage, :, sf_word]) - T.tcgen05_sf_warp_transpose(SFB_shared[stage, :, sf_word]) - T.fence_proxy_async() - T.mbarrier_arrive(with_sf_full[stage], 0) - - # Epilogue: write C from tmem to global memory - for w in T.serial(waves): - cluster_id = block_id // 2 - tile_id = num_clusters * w + cluster_id - bx_cluster = (tile_id // group_size) % m_clusters - bx = bx_cluster * 2 + cta_id - by = (tile_id % group_size) + (tile_id // group_size) // m_clusters * group_size - - if bx * block_M < M and by * block_N < N: - T.mbarrier_wait_parity(tmem_full, w & 1) - T.copy(C_tmem, C_local) - T.copy(C_local, C_local_cast) - T.copy(C_local_cast, C_shared) - - if use_tma_store: - T.copy(C_shared, C[bx * block_M, by * block_N]) - else: - T.copy(C_shared, C[bx * block_M, by * block_N], disable_tma=True) - - T.mbarrier_arrive(tmem_empty, 0) - - return C diff --git a/src/nvfp4_mega_moe.py b/src/nvfp4_mega_moe.py deleted file mode 100644 index ee0e26f1..00000000 --- a/src/nvfp4_mega_moe.py +++ /dev/null @@ -1,163 +0,0 @@ -""" -NVFP4 Mega MoE Kernel — Full MoE with expert parallelism. - -This is the main kernel that replaces fp8_nvfp4_mega_moe from DeepGEMM. - -Architecture: -- L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with UE4M3 scales) -- SiLU+Mul activation -- L2 GEMM: down_proj (FP4 x FP4 → BF16 with UE4M3 scales) -- NVLink cross-rank sync via symm buffer -- Expert parallel: each rank handles NUM_EXPERTS/8 experts - -The kernel is written in TileLang, compiled to SM100 (Blackwell) CUBIN. -""" - -import torch -import tilelang -import tilelang.language as T -from tilelang.carver.arch import driver - -# DeepSeek-V4-Pro dimensions -HIDDEN = 7168 -INTERMEDIATE = 3072 -NUM_EXPERTS = 256 -NUM_RANKS = 8 -NUM_TOPK = 6 - -# Block sizes for the block-scaled GEMM -BLOCK_M = 192 # tokens per tile -BLOCK_K = 128 -BLOCK_N = 128 - -# NVFP4 scale parameters -SF_GRANULARITY_K = 16 # UE4M3 group_size -SF_PACK_FACTOR = 4 # 4 UE4M3 values per uint32 -SF_WORDS_PER_K_BLOCK = BLOCK_K // (SF_GRANULARITY_K * SF_PACK_FACTOR) # 2 - - -@tilelang.jit -def nvfp4_mega_moe_l1( - # Activation (packed FP4, from staging kernel) - X, # (num_tokens, K//2) int8 packed E2M1 - X_SF, # (num_tokens, sf_k_groups) uint32 packed UE4M3 - # L1 weights (pre-transformed for mega_moe) - L1_W, # (num_experts_per_rank, 2*INTERMEDIATE, K//2) int8 K-major - L1_SF, # (num_experts_per_rank, 2*INTERMEDIATE, sf_k_groups) uint32 TMA-aligned - # Routing - TopkIdx, # (num_tokens, NUM_TOPK) int32 - TopkW, # (num_tokens, NUM_TOPK) float32 - # Output - Y, # (num_tokens, 2*INTERMEDIATE) bfloat16 - num_experts_per_rank, -): - """ - L1 GEMM for mega_moe: gate_up_proj - X(FP4) @ L1_W(FP4)^T → Y(BF16) with UE4M3 block scaling - - This handles the gate_up_proj for all experts in the rank. - Each token is routed to NUM_TOPK experts, and the GEMM is computed - per expert group using persistent scheduling. - """ - num_tokens = T.const("num_tokens") - K = T.const("K") - INTER = T.const("INTERMEDIATE") - - k_iters = T.ceildiv(K, BLOCK_K) - sf_k_groups = T.ceildiv(K, SF_GRANULARITY_K * SF_PACK_FACTOR) - - with T.Kernel( - T.ceildiv(num_tokens, BLOCK_M), - T.ceildiv(2 * INTER, BLOCK_N), - threads=256, - cluster_dims=2, - ) as (bx, by): - cta_id = T.block_rank_in_cluster() - T.assume(cta_id < 2) - - # ... (persistent scheduling, expert routing, L1 GEMM with block scaling) - # This follows the same pattern as nvfp4_blockscaled_gemm_2cta_persistent - # but adds expert scheduling and topk weight scaling - - pass # Full implementation in progress - - -@tilelang.jit -def nvfp4_mega_moe_l2( - # L1 output (quantized to FP4 for L2 input) - X, # (num_tokens, INTER//2) int8 packed E2M1 - X_SF, # (num_tokens, sf_k_groups) uint32 packed UE4M3 - # L2 weights - L2_W, # (num_experts_per_rank, HIDDEN, INTER//2) int8 K-major - L2_SF, # (num_experts_per_rank, HIDDEN, sf_k_groups) uint32 TMA-aligned - # Routing - TopkIdx, # (num_tokens, NUM_TOPK) int32 - TopkW, # (num_tokens, NUM_TOPK) float32 - # Output - Y, # (num_tokens, HIDDEN) bfloat16 - num_experts_per_rank, -): - """ - L2 GEMM for mega_moe: down_proj - X(FP4) @ L2_W(FP4)^T → Y(BF16) with UE4M3 block scaling - - After SiLU+Mul on the L1 output, the result is quantized to FP4 - and fed into L2. - """ - pass # Symmetric to L1 - - -def nvfp4_mega_moe_full( - hidden_states, # (num_tokens, HIDDEN) bfloat16 - topk_weights, # (num_tokens, NUM_TOPK) float32 - topk_ids, # (num_tokens, NUM_TOPK) int32 - l1_weights, # L1 weights (transformed for mega_moe) - l1_scales, # L1 UE4M3 scales (transformed) - l2_weights, # L2 weights (transformed for mega_moe) - l2_scales, # L2 UE4M3 scales (transformed) - symm_buffer, # NVLink symm buffer for cross-rank sync -): - """ - Full mega_moe forward pass: - 1. Stage: quantize BF16 hidden_states → FP4 + UE4M3 scales - 2. L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with block scaling) - 3. SiLU + Mul (activation) - 4. Quantize L1 output → FP4 + UE4M3 scales - 5. L2 GEMM: down_proj (FP4 x FP4 → BF16 with block scaling) - 6. NVLink sync + reduce across ranks - """ - # Step 1: Stage activation (BF16 → FP4 quantization) - # This is the staging kernel from patches/staging_kernel.py - x_fp4, x_sf = stage_activation(hidden_states) - - # Step 2: L1 GEMM - l1_output = nvfp4_mega_moe_l1( - x_fp4, x_sf, l1_weights, l1_scales, topk_ids, topk_weights) - - # Step 3: SiLU + Mul - gate, up = l1_output.chunk(2, dim=-1) - activated = torch.nn.functional.silu(gate) * up - - # Step 4: Quantize L1 output → FP4 - l1_fp4, l1_sf = stage_activation(activated) - - # Step 5: L2 GEMM - l2_output = nvfp4_mega_moe_l2( - l1_fp4, l1_sf, l2_weights, l2_scales, topk_ids, topk_weights) - - # Step 6: NVLink reduce - output = nvlink_reduce(l2_output, topk_weights, symm_buffer) - - return output - - -def stage_activation(x_bf16): - """Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales. - - This replaces the Triton staging kernel from patches/staging_kernel.py. - """ - # E2M1 quantization with UE4M3 block scaling - # For now, use PyTorch reference implementation - # TODO: Write as TileLang kernel for full pipeline integration - from deep_gemm import per_token_cast_to_fp4 - return per_token_cast_to_fp4(x_bf16) diff --git a/src/nvfp4_megamoe_kernel.egg-info/PKG-INFO b/src/nvfp4_megamoe_kernel.egg-info/PKG-INFO new file mode 100644 index 00000000..59be8beb --- /dev/null +++ b/src/nvfp4_megamoe_kernel.egg-info/PKG-INFO @@ -0,0 +1,7 @@ +Metadata-Version: 2.4 +Name: nvfp4-megamoe-kernel +Version: 0.1.0 +Summary: NVFP4 Mega MoE kernel for DeepSeek-V4-Pro on Blackwell (TileLang) +Requires-Python: >=3.10 +Requires-Dist: torch>=2.5 +Requires-Dist: tilelang>=0.1 diff --git a/src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt b/src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt new file mode 100644 index 00000000..f3641fa0 --- /dev/null +++ b/src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt @@ -0,0 +1,11 @@ +README.md +pyproject.toml +src/nvfp4_megamoe_kernel/__init__.py +src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +src/nvfp4_megamoe_kernel/symm_buffer.py +src/nvfp4_megamoe_kernel/weight_transform.py +src/nvfp4_megamoe_kernel.egg-info/PKG-INFO +src/nvfp4_megamoe_kernel.egg-info/SOURCES.txt +src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt +src/nvfp4_megamoe_kernel.egg-info/requires.txt +src/nvfp4_megamoe_kernel.egg-info/top_level.txt \ No newline at end of file diff --git a/src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt b/src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt new file mode 100644 index 00000000..8b137891 --- /dev/null +++ b/src/nvfp4_megamoe_kernel.egg-info/dependency_links.txt @@ -0,0 +1 @@ + diff --git a/src/nvfp4_megamoe_kernel.egg-info/requires.txt b/src/nvfp4_megamoe_kernel.egg-info/requires.txt new file mode 100644 index 00000000..06d97fde --- /dev/null +++ b/src/nvfp4_megamoe_kernel.egg-info/requires.txt @@ -0,0 +1,2 @@ +torch>=2.5 +tilelang>=0.1 diff --git a/src/nvfp4_megamoe_kernel.egg-info/top_level.txt b/src/nvfp4_megamoe_kernel.egg-info/top_level.txt new file mode 100644 index 00000000..0c0c2376 --- /dev/null +++ b/src/nvfp4_megamoe_kernel.egg-info/top_level.txt @@ -0,0 +1 @@ +nvfp4_megamoe_kernel diff --git a/src/nvfp4_megamoe_kernel/__init__.py b/src/nvfp4_megamoe_kernel/__init__.py index 37f117d4..05cdc12d 100644 --- a/src/nvfp4_megamoe_kernel/__init__.py +++ b/src/nvfp4_megamoe_kernel/__init__.py @@ -1,4 +1,4 @@ -"""NVFP4 Mega MoE Kernel — TileLang implementation for DeepSeek-V4-Pro on Blackwell.""" +"""NVFP4 Mega MoE Kernel — CUTLASS implementation for DeepSeek-V4-Pro on Blackwell.""" from nvfp4_megamoe_kernel.nvfp4_mega_moe import ( nvfp4_mega_moe_full, @@ -7,10 +7,7 @@ from nvfp4_megamoe_kernel.nvfp4_mega_moe import ( stage_activation, ) from nvfp4_megamoe_kernel.weight_transform import ( - fold_global_scale, - pack_ue4m3_to_uint32, - interleave_l1_weights, - transform_nvfp4_weights_for_tilelang as transform_nvfp4_weights_for_mega_moe, + transform_nvfp4_weights_for_mega_moe, ) from nvfp4_megamoe_kernel.symm_buffer import ( SymmBuffer, @@ -20,12 +17,9 @@ from nvfp4_megamoe_kernel.symm_buffer import ( __all__ = [ "nvfp4_mega_moe_full", "nvfp4_mega_moe_l1", - "nvfp4_moe_l2", + "nvfp4_mega_moe_l2", "stage_activation", "transform_nvfp4_weights_for_mega_moe", - "fold_global_scale", - "pack_ue4m3_to_uint32", - "interleave_l1_weights", "SymmBuffer", "get_symm_buffer_for_nvfp4_mega_moe", ] diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md new file mode 100644 index 00000000..4cd75921 --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/README.md @@ -0,0 +1,99 @@ +# CUTLASS NVFP4 Block-Scaled GEMM Kernel + +Native Blackwell (SM100) NVFP4 block-scaled GEMM using CUTLASS 3.x. + +## Overview + +This kernel implements the DeepSeek-V4-Pro MoE GEMM operations using CUTLASS's +`MainloopSm100TmaUmmaWarpSpecializedBlockScaled` collective, which invokes the +native `mxf8f6f4.block_scale` tensor core instruction (`tcgen05.mma`) on NVIDIA +Blackwell GPUs. + +### Key Features + +- **Native NVFP4 MMA**: E2M1 × E2M1 with UE4M3 block-16 scaling entirely in hardware +- **No dequantization**: Avoids the costly dequantize-then-BF16-GEMM fallback path +- **TMA + UMMA**: Uses TMA for loading data into shared memory and UMMA for tensor core ops +- **TMEM scale loading**: UE4M3 scale factors loaded into tensor memory via `tcgen05.ld` +- **Grouped expert GEMM**: Per-expert dispatch for MoE with top-k routing + +### Architecture + +``` +E2M1 (int8, 2 vals/byte) + UE4M3 (float8_e4m3fn, group_size=16) + → TMA load to shared memory + → UMMA block-scaled MMA (mxf8f6f4.block_scale) + → float32 accumulator + → BF16 output +``` + +## Data Layout + +| Tensor | Shape | Type | Layout | +|--------|-------|------|--------| +| A (activation) | (M, K//2) | int8 | K-major (ColumnMajor) | +| SFA (activation scales) | (M, K//16) | float8_e4m3fn | K-major (Sm1xxBlockScaledConfig) | +| B (weight) | (N, K//2) | int8 | K-major (ColumnMajor) | +| SFB (weight scales) | (N, K//16) | float8_e4m3fn | K-major (Sm1xxBlockScaledConfig) | +| C (output) | (M, N) | bfloat16 | RowMajor | + +K//2 because E2M1 packs 2 values per byte. +K//16 because UE4M3 block scale has group_size=16. + +## Building on B200 + +```bash +# Inside the Docker container on the B200: +cd /root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm +bash build.sh +``` + +Or manually: + +```bash +export CUTLASS_INCLUDE_DIR=/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include +python3 setup.py build_ext --inplace +``` + +## Testing + +```bash +python3 test_gemm.py +``` + +## Usage in DeepSeek-V4-Pro + +The kernel is automatically used by `nvfp4_mega_moe.py` when: +1. `MEGA_MOE_USE_CUTLASS=1` (default) +2. The CUTLASS extension compiles successfully + +If CUTLASS is unavailable, it falls back to the TileLang or dequantize+BF16 path. + +## CUTLASS Internals + +### Dispatch Policy +`MainloopSm100TmaUmmaWarpSpecializedBlockScaled` + +### TiledMma +UMMA atom: `mxf8f6f4.block_scale` with SFVecSize=16 + +### Scale Factor Layout +Uses `Sm1xxBlockScaledConfig<16>` which defines: +- SfAtom layout for K-major scale factors +- `tile_atom_to_shape_SFA/SFB` for computing the global scale layout +- `deduce_smem_layoutSFA/SFB` for shared memory layout + +### Pipeline +1. TMA loads A, B, SFA, SFB into shared memory +2. UMMA warp-specialized MMA with block scaling +3. Scale factors loaded from shared memory to TMEM via UTCCP +4. Accumulator in float32, converted to BF16 in epilogue + +## Files + +- `cutlass_nvfp4_gemm.cu` — Standalone CUDA kernel (C API) +- `pytorch_binding.cpp` — PyTorch extension binding +- `kernel.py` — Python wrapper with compilation and fallback +- `setup.py` — Build configuration +- `build.sh` — Build script for B200 +- `test_gemm.py` — Test script diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py new file mode 100644 index 00000000..779f0deb --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/__init__.py @@ -0,0 +1,6 @@ +"""CUTLASS NVFP4 Block-Scaled GEMM for DeepSeek-V4-Pro on Blackwell (SM100).""" + +from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import ( + cutlass_nvfp4_blockscaled_gemm, + cutlass_grouped_nvfp4_gemm, +) diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh new file mode 100644 index 00000000..78c233ab --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/build.sh @@ -0,0 +1,44 @@ +#!/bin/bash +# Build script for CUTLASS NVFP4 block-scaled GEMM on B200 (Blackwell SM100). +# +# Run inside the Docker container: +# docker exec -it deepseek-v4-quant-vllm bash +# cd /path/to/cutlass_nvfp4_gemm && bash build.sh +# +# Or from outside: +# docker exec deepseek-v4-quant-vllm bash -c "cd /root/nvfp4-megamoe-kernel/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm && bash build.sh" + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +cd "$SCRIPT_DIR" + +# CUTLASS include path (inside the Docker container) +export CUTLASS_INCLUDE_DIR="${CUTLASS_INCLUDE_DIR:-/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include}" + +echo "=== CUTLASS NVFP4 GEMM Build ===" +echo "CUTLASS_INCLUDE_DIR: $CUTLASS_INCLUDE_DIR" + +# Verify CUTLASS headers +if [ ! -f "${CUTLASS_INCLUDE_DIR}/cutlass/cutlass.h" ]; then + echo "ERROR: CUTLASS headers not found at ${CUTLASS_INCLUDE_DIR}" + echo "Set CUTLASS_INCLUDE_DIR to point to the cutlass/include directory." + exit 1 +fi + +# Verify block-scaled MMA header +if [ ! -f "${CUTLASS_INCLUDE_DIR}/cutlass/gemm/collective/sm100_blockscaled_mma_warpspecialized.hpp" ]; then + echo "WARNING: Block-scaled MMA header not found. The CollectiveBuilder path will be used." +fi + +echo "Building PyTorch extension..." +python3 setup.py build_ext --inplace 2>&1 | tee build.log + +if [ $? -eq 0 ]; then + echo "=== Build SUCCESS ===" + echo "Extension built. Test with: python3 test_gemm.py" +else + echo "=== Build FAILED ===" + echo "Check build.log for errors." + exit 1 +fi diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu new file mode 100644 index 00000000..a8715762 --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/cutlass_nvfp4_gemm.cu @@ -0,0 +1,312 @@ +/*************************************************************************************************** + * CUTLASS NVFP4 Block-Scaled GEMM for DeepSeek-V4-Pro MoE + **************************************************************************************************/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define CUTLASS_CHECK(status) \ + do { \ + cutlass::Status _s = status; \ + if (_s != cutlass::Status::kSuccess) { \ + fprintf(stderr, "CUTLASS error at %s:%d: %s\n", \ + __FILE__, __LINE__, cutlassGetStatusString(_s)); \ + return -1; \ + } \ + } while (0) + +#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) + +using namespace cute; + +using ElementA = cutlass::nv_float4_t; +using LayoutATag = cutlass::layout::RowMajor; +constexpr int AlignmentA = 32; + +using ElementB = cutlass::nv_float4_t; +using LayoutBTag = cutlass::layout::ColumnMajor; +constexpr int AlignmentB = 32; + +using ElementD = cutlass::bfloat16_t; +using ElementC = float; +using LayoutCTag = cutlass::layout::RowMajor; +using LayoutDTag = cutlass::layout::RowMajor; +constexpr int AlignmentD = 128 / cutlass::sizeof_bits::value; +constexpr int AlignmentC = 128 / cutlass::sizeof_bits::value; + +using ElementAccumulator = float; +using ElementCompute = float; +using ArchTag = cutlass::arch::Sm100; +using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; + +using MmaTileShape = Shape<_128, _128, _256>; +using ClusterShape = Shape<_1, _1, _1>; + +constexpr int InputSFVectorSize = 16; + +using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, OperatorClass, + MmaTileShape, ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, ElementCompute, + ElementC, LayoutCTag, AlignmentC, + ElementD, LayoutDTag, AlignmentD, + cutlass::epilogue::collective::EpilogueScheduleAuto +>::CollectiveOp; + +using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, OperatorClass, + ElementA, LayoutATag, AlignmentA, + ElementB, LayoutBTag, AlignmentB, + ElementAccumulator, + MmaTileShape, ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout(sizeof(typename CollectiveEpilogue::SharedStorage))>, + cutlass::gemm::collective::KernelScheduleAuto +>::CollectiveOp; + +using GemmKernel = cutlass::gemm::kernel::GemmUniversal< + Shape, + CollectiveMainloop, + CollectiveEpilogue, + void>; + +using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + +using StrideA = typename Gemm::GemmKernel::StrideA; +using StrideB = typename Gemm::GemmKernel::StrideB; +using StrideC = typename Gemm::GemmKernel::StrideC; +using StrideD = typename Gemm::GemmKernel::StrideD; +using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; +using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; + +///////////////////////////////////////////////////////////////////////////////////////////////// +// Scale factor remap: source (row-major or col-major) -> CUTLASS interleaved layout +// +// Source-first remap: iterates over the logical (m, k_sf) source grid (which is +// a simple 2D row-major tensor), and for each valid element, computes the +// physical CUTLASS dest index using the layout's forward mapping. +// +// This is the inverse of the old idx2crd approach (which iterated dest → source +// and relied on fragile coordinate extraction from flattened hierarchical coords). +// Source-first is simpler, correct by construction, and avoids the idx2crd bugs. +// +// 2D kernel launch: blockDim = (32, 8), grid over (K_sf, MN). +// SFA: source is (MN, K_sf) row-major → src[m * K_sf + k_sf] +// SFB: source is (K_sf, MN) row-major → src[k_sf * MN + m] +///////////////////////////////////////////////////////////////////////////////////////////////// + +template +__global__ void remap_sf_to_cutlass_kernel( + const cutlass::float_ue4m3_t* __restrict__ src, + cutlass::float_ue4m3_t* __restrict__ dst, // CUTLASS interleaved layout (zero-initialized) + LayoutSF layout_sf, // CuTe layout for dst + int MN, int K_sf, // Source dimensions (in SF groups) + bool col_major_src = false // true if source is (K_sf, MN) row-major +) { + int k_sf = blockIdx.x * blockDim.x + threadIdx.x; + int m = blockIdx.y * blockDim.y + threadIdx.y; + if (m >= MN || k_sf >= K_sf) return; + + // Compute the CUTLASS physical index from logical (m, k_sf) via the layout + int dst_idx = layout_sf(cute::make_coord(m, k_sf)); + + // Read from source (row-major or col-major) and write to CUTLASS position + dst[dst_idx] = col_major_src ? src[k_sf * MN + m] : src[m * K_sf + k_sf]; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// +// C API +///////////////////////////////////////////////////////////////////////////////////////////////// + +extern "C" { + +int cutlass_nvfp4_gemm_run( + const void* A_ptr, const void* SFA_ptr, + const void* B_ptr, const void* SFB_ptr, + void* D_ptr, + int M, int N, int K, + float alpha, float beta, + cudaStream_t stream +) { + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1)); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1)); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1)); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1)); + + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1)); + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); + + using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA; + using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB; + using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + + int sfa_size = cute::cosize(layout_SFA); + int sfb_size = cute::cosize(layout_SFB); + int K_sf = K / InputSFVectorSize; + + cutlass::device_memory::allocation sfa_cutlass(sfa_size); + cutlass::device_memory::allocation sfb_cutlass(sfb_size); + cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); + cudaMemsetAsync(sfb_cutlass.get(), 0, sfb_size * sizeof(ElementSF), stream); + + dim3 block(32, 8); + dim3 grid_sfa((K_sf + block.x - 1) / block.x, (M + block.y - 1) / block.y); + dim3 grid_sfb((K_sf + block.x - 1) / block.x, (N + block.y - 1) / block.y); + + remap_sf_to_cutlass_kernel<<>>( + static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false); + remap_sf_to_cutlass_kernel<<>>( + static_cast(SFB_ptr), sfb_cutlass.get(), layout_SFB, N, K_sf, true); + + typename Gemm::Arguments arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + { + static_cast(A_ptr), stride_A, + static_cast(B_ptr), stride_B, + sfa_cutlass.get(), layout_SFA, + sfb_cutlass.get(), layout_SFB + }, + { + { alpha, beta }, + nullptr, stride_C, + static_cast(D_ptr), stride_D + } + }; + + Gemm gemm; + CUTLASS_CHECK(gemm.can_implement(arguments)); + + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream)); + CUTLASS_CHECK(gemm.run(stream)); + + return 0; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// +// SFB prepack: pre-remap weight scale factors once at load time +///////////////////////////////////////////////////////////////////////////////////////////////// + +extern "C" int cutlass_nvfp4_sfb_size( + int M, int N, int K, + int* out_size +) { + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); + *out_size = cute::cosize(layout_SFB); + return 0; +} + +extern "C" int cutlass_nvfp4_prepack_sfb_run( + const void* SFB_ptr, + void* SFB_cutlass_ptr, + int M, int N, int K, + cudaStream_t stream +) { + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); + using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + + int sfb_size = cute::cosize(layout_SFB); + int K_sf = K / InputSFVectorSize; + + cudaMemsetAsync(static_cast(SFB_cutlass_ptr), 0, sfb_size * sizeof(ElementSF), stream); + + dim3 block(32, 8); + dim3 grid((K_sf + block.x - 1) / block.x, (N + block.y - 1) / block.y); + remap_sf_to_cutlass_kernel<<>>( + static_cast(SFB_ptr), + static_cast(SFB_cutlass_ptr), + layout_SFB, + N, K_sf, true // SFB source is (K_sf, N) + ); + + return 0; +} + +///////////////////////////////////////////////////////////////////////////////////////////////// +// GEMM with prepacked SFB — skips SFB allocation, memset, and remap +///////////////////////////////////////////////////////////////////////////////////////////////// + +extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( + const void* A_ptr, const void* SFA_ptr, + const void* B_ptr, const void* SFB_cutlass_ptr, + void* D_ptr, + int M, int N, int K, + float alpha, float beta, + cudaStream_t stream +) { + StrideA stride_A = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(M, K, 1)); + StrideB stride_B = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(N, K, 1)); + StrideC stride_C = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, 1)); + StrideD stride_D = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, 1)); + + using Sm1xxBlkScaledConfig = typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; + LayoutSFA layout_SFA = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1)); + LayoutSFB layout_SFB = Sm1xxBlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1)); + + using ArrayElementA = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementA; + using ArrayElementB = typename Gemm::GemmKernel::CollectiveMainloop::ArrayElementB; + using ElementSF = typename Gemm::GemmKernel::CollectiveMainloop::ElementSF; + + int sfa_size = cute::cosize(layout_SFA); + int K_sf = K / InputSFVectorSize; + + // Only remap SFA (activation scales) — SFB is prepacked + cutlass::device_memory::allocation sfa_cutlass(sfa_size); + cudaMemsetAsync(sfa_cutlass.get(), 0, sfa_size * sizeof(ElementSF), stream); + + dim3 block(32, 8); + dim3 grid_sfa((K_sf + block.x - 1) / block.x, (M + block.y - 1) / block.y); + remap_sf_to_cutlass_kernel<<>>( + static_cast(SFA_ptr), sfa_cutlass.get(), layout_SFA, M, K_sf, false); + + typename Gemm::Arguments arguments { + cutlass::gemm::GemmUniversalMode::kGemm, + {M, N, K, 1}, + { + static_cast(A_ptr), stride_A, + static_cast(B_ptr), stride_B, + sfa_cutlass.get(), layout_SFA, + static_cast(SFB_cutlass_ptr), layout_SFB + }, + { + { alpha, beta }, + nullptr, stride_C, + static_cast(D_ptr), stride_D + } + }; + + Gemm gemm; + CUTLASS_CHECK(gemm.can_implement(arguments)); + + size_t workspace_size = Gemm::get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + + CUTLASS_CHECK(gemm.initialize(arguments, workspace.get(), stream)); + CUTLASS_CHECK(gemm.run(stream)); + + return 0; +} + +} // extern "C" + +#endif diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py new file mode 100644 index 00000000..124a8205 --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/kernel.py @@ -0,0 +1,146 @@ +""" +CUTLASS NVFP4 Block-Scaled GEMM — Native Blackwell SM100 kernel. + +Uses the pre-compiled PyTorch CUDA extension (cutlass_nvfp4_gemm._C) +which invokes native mxf8f6f4.block_scale tensor core instructions. +""" + +import os +import torch + +MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0")) + +try: + from cutlass_nvfp4_gemm import _C + _CUTLASS_AVAILABLE = True +except ImportError: + _CUTLASS_AVAILABLE = False + + +def cutlass_nvfp4_blockscaled_gemm( + A_packed, # (M, K_half) int8 packed E2M1 + SFA, # scale factors for A (float8_e4m3fn) + B_packed, # (K_half, N) int8 packed E2M1, column-major for CUTLASS + SFB, # scale factors for B — either (sf_k, N) float8_e4m3fn row-major, or prepacked CUTLASS layout + M, N, K, # Problem dimensions (K in FP4 elements) + alpha=1.0, # fp32 scalar applied in epilogue: D = alpha * A @ B + beta * C + sfb_prepacked=False, # True if SFB is already in CUTLASS layout +): + """Single NVFP4 block-scaled GEMM using CUTLASS. + + If sfb_prepacked=True, SFB is assumed to be in CUTLASS interleaved layout + (from prepack_sfb) and the SFB remap is skipped. + """ + if not _CUTLASS_AVAILABLE: + raise RuntimeError("CUTLASS NVFP4 GEMM extension not available") + if sfb_prepacked: + return _C.forward_prepacked_sfb(A_packed, SFA, B_packed, SFB, M, N, K, alpha) + else: + return _C.forward(A_packed, SFA, B_packed, SFB, M, N, K, alpha) + + +def prepack_sfb(SFB, M, N, K): + """Pre-remap SFB weight scales into CUTLASS interleaved layout. + + Call once after weight transform. Returns a tensor that can be passed + to cutlass_nvfp4_blockscaled_gemm with sfb_prepacked=True. + + M is used for layout sizing. Test with different M values to confirm + SFB layout is M-independent; if so, any valid M works (e.g. 128). + """ + if not _CUTLASS_AVAILABLE: + raise RuntimeError("CUTLASS NVFP4 GEMM extension not available") + return _C.prepack_sfb(SFB, M, N, K) + + +def cutlass_grouped_nvfp4_gemm( + x_fp4, # (num_slots_or_tokens, K_half) int8 packed E2M1 + x_sf, # (num_slots_or_tokens, sf_k) float8_e4m3fn block scales + weights, # (E_per_rank, K_half, N) int8 packed E2M1, column-major for CUTLASS + weight_sf, # (E_per_rank, sf_k, N) float8_e4m3fn, column-major + slot_expert_ids, # (num_slots,) int32 — per-slot local expert IDs + slot_token=None, # (num_slots,) int64 — per-slot token indices (default: arange) + alpha=1.0, # fp32 scalar: D = alpha * A @ B (from stage_activation global scale) +): + """Per-expert grouped GEMM for MoE dispatch using CUTLASS NVFP4. + + Takes 1D per-slot expert IDs and token indices (pre-built by caller). + SFB weight scales are remapped per-expert inside CUTLASS on each call. + NO prepack cache — see nvfp4_mega_moe.py for rationale. + + For L1: x_fp4 has num_tokens rows, slot_token maps slots→rows. + For L2: x_fp4 has num_slots rows, slot_token is just arange(num_slots). + + Returns: + slot_out: (num_slots, N) bfloat16 — per-slot GEMM results + slot_token: (num_slots,) int64 — token index for each slot + """ + num_slots = slot_expert_ids.shape[0] + K_half = x_fp4.shape[1] + K = K_half * 2 + N = weights.shape[2] + num_experts = weights.shape[0] + + if num_slots == 0: + slot_out = torch.empty(0, N, dtype=torch.bfloat16, device=x_fp4.device) + slot_token_out = torch.empty(0, dtype=torch.int64, device=x_fp4.device) + return slot_out, slot_token_out + + # Use provided slot_token or default to identity mapping + provided_slot_token = slot_token + + if provided_slot_token is None: + slot_token_out = torch.arange(num_slots, device=x_fp4.device) + slot_x = x_fp4 + slot_x_sf = x_sf + else: + slot_token_out = provided_slot_token + slot_x = x_fp4[provided_slot_token].contiguous() + slot_x_sf = x_sf[provided_slot_token].contiguous() + + if MEGA_MOE_DEBUG: + print(f"[cutlass_grouped_gemm] slots={num_slots} K={K} N={N} " + f"experts={num_experts}") + + slot_out = torch.empty(num_slots, N, dtype=torch.bfloat16, device=x_fp4.device) + + for e in range(num_experts): + expert_slots = (slot_expert_ids == e) + if not expert_slots.any(): + continue + + e_idx = expert_slots.nonzero(as_tuple=True)[0] + expert_x = slot_x[e_idx] + expert_x_sf = slot_x_sf[e_idx] + expert_w = weights[e] + expert_w_sf = weight_sf[e] + M_expert = e_idx.shape[0] + + if MEGA_MOE_DEBUG and e < 3 and M_expert > 0: + print(f"[GEMM-IN] expert={e} M={M_expert} N={N} K={K} " + f"w shape={expert_w.shape}") + + # Shape/dtype contract asserts — SFB bugs hide in silent shape mismatches + assert expert_x.shape == (M_expert, K // 2), f"expert_x shape {expert_x.shape} != ({M_expert}, {K // 2})" + assert expert_x_sf.shape == (M_expert, K // 16), f"SFA shape {expert_x_sf.shape} != ({M_expert}, {K // 16})" + assert expert_w.shape == (K // 2, N), f"expert_w shape {expert_w.shape} != ({K // 2}, {N})" + assert expert_w_sf.shape == (K // 16, N), f"SFB shape {expert_w_sf.shape} != ({K // 16}, {N})" + assert expert_x_sf.dtype == torch.float8_e4m3fn, f"SFA dtype {expert_x_sf.dtype}" + assert expert_w_sf.dtype == torch.float8_e4m3fn, f"SFB dtype {expert_w_sf.dtype}" + + expert_out = cutlass_nvfp4_blockscaled_gemm( + expert_x, expert_x_sf, + expert_w, expert_w_sf, + M_expert, N, K, + alpha=alpha, + ) + + if MEGA_MOE_DEBUG: + if torch.isnan(expert_out).any() or torch.isinf(expert_out).any(): + raise RuntimeError( + f"expert {e} of {num_experts}: GEMM emitted NaN/Inf. " + f"M={M_expert} N={N} K={K}") + + slot_out[e_idx] = expert_out + + return slot_out, slot_token_out diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp new file mode 100644 index 00000000..774502df --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/pytorch_binding.cpp @@ -0,0 +1,131 @@ +/** PyTorch binding for CUTLASS NVFP4 block-scaled GEMM */ + +#include +#include +#include + +extern "C" int cutlass_nvfp4_gemm_run( + const void* A_ptr, const void* SFA_ptr, + const void* B_ptr, const void* SFB_ptr, + void* D_ptr, + int M, int N, int K, + float alpha, float beta, + cudaStream_t stream +); + +extern "C" int cutlass_nvfp4_gemm_run_prepacked_sfb( + const void* A_ptr, const void* SFA_ptr, + const void* B_ptr, const void* SFB_cutlass_ptr, + void* D_ptr, + int M, int N, int K, + float alpha, float beta, + cudaStream_t stream +); + +extern "C" int cutlass_nvfp4_sfb_size( + int M, int N, int K, + int* out_size +); + +extern "C" int cutlass_nvfp4_prepack_sfb_run( + const void* SFB_ptr, + void* SFB_cutlass_ptr, + int M, int N, int K, + cudaStream_t stream +); + +torch::Tensor cutlass_nvfp4_gemm_forward( + torch::Tensor A_packed, + torch::Tensor SFA, + torch::Tensor B_packed, + torch::Tensor SFB, + int64_t M, int64_t N, int64_t K, + double alpha = 1.0 +) { + auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device())); + + auto stream = c10::cuda::getCurrentCUDAStream(); + cudaStream_t cuda_stream = stream.stream(); + + int rc = cutlass_nvfp4_gemm_run( + A_packed.data_ptr(), SFA.data_ptr(), + B_packed.data_ptr(), SFB.data_ptr(), + D.data_ptr(), + static_cast(M), static_cast(N), static_cast(K), + static_cast(alpha), 0.0f, + cuda_stream + ); + + TORCH_CHECK(rc == 0, "CUTLASS NVFP4 GEMM failed with error code ", rc); + + return D; +} + +torch::Tensor cutlass_nvfp4_gemm_forward_prepacked_sfb( + torch::Tensor A_packed, + torch::Tensor SFA, + torch::Tensor B_packed, + torch::Tensor SFB_cutlass, + int64_t M, int64_t N, int64_t K, + double alpha = 1.0 +) { + auto D = torch::empty({M, N}, torch::dtype(torch::kBFloat16).device(A_packed.device())); + + auto stream = c10::cuda::getCurrentCUDAStream(); + cudaStream_t cuda_stream = stream.stream(); + + int rc = cutlass_nvfp4_gemm_run_prepacked_sfb( + A_packed.data_ptr(), SFA.data_ptr(), + B_packed.data_ptr(), SFB_cutlass.data_ptr(), + D.data_ptr(), + static_cast(M), static_cast(N), static_cast(K), + static_cast(alpha), 0.0f, + cuda_stream + ); + + TORCH_CHECK(rc == 0, "CUTLASS NVFP4 GEMM (prepacked SFB) failed with error code ", rc); + + return D; +} + +torch::Tensor prepack_sfb( + torch::Tensor SFB, + int64_t M, + int64_t N, + int64_t K +) { + int size = 0; + int rc = cutlass_nvfp4_sfb_size( + static_cast(M), + static_cast(N), + static_cast(K), + &size + ); + TORCH_CHECK(rc == 0, "sfb_size failed"); + + auto out = torch::empty( + {size}, + torch::dtype(SFB.dtype()).device(SFB.device()) + ); + + auto stream = c10::cuda::getCurrentCUDAStream(); + + rc = cutlass_nvfp4_prepack_sfb_run( + SFB.data_ptr(), + out.data_ptr(), + static_cast(M), + static_cast(N), + static_cast(K), + stream.stream() + ); + + TORCH_CHECK(rc == 0, "prepack_sfb failed"); + + return out; +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("forward", &cutlass_nvfp4_gemm_forward, "CUTLASS NVFP4 block-scaled GEMM forward"); + m.def("forward_prepacked_sfb", &cutlass_nvfp4_gemm_forward_prepacked_sfb, "CUTLASS NVFP4 GEMM forward with prepacked SFB"); + m.def("prepack_sfb", &prepack_sfb, "Pre-remap SFB weight scales into CUTLASS layout"); +} diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py new file mode 100644 index 00000000..08de1df0 --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/setup.py @@ -0,0 +1,65 @@ +"""Setup script for CUTLASS NVFP4 block-scaled GEMM PyTorch extension.""" + +import os +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + +# CUTLASS include directory — prefer the latest from GitHub +CUTLASS_INCLUDE_DIR = os.environ.get( + "CUTLASS_INCLUDE_DIR", + "/root/cutlass/include" +) +if not os.path.exists(os.path.join(CUTLASS_INCLUDE_DIR, "cutlass", "cutlass.h")): + for alt in [ + "/root/cutlass/include", + "/usr/local/lib/python3.12/dist-packages/tilelang/3rdparty/cutlass/include", + "/usr/local/include/cutlass", + "/opt/cutlass/include", + ]: + if os.path.exists(os.path.join(alt, "cutlass", "cutlass.h")): + CUTLASS_INCLUDE_DIR = alt + break + +CUTLASS_UTIL_INCLUDE = os.path.join(os.path.dirname(CUTLASS_INCLUDE_DIR), "tools", "util", "include") + +include_dirs = [CUTLASS_INCLUDE_DIR] +if os.path.exists(CUTLASS_UTIL_INCLUDE): + include_dirs.append(CUTLASS_UTIL_INCLUDE) + +# CCCL / libcu++ headers (required by CUTLASS 3.x) +CCCL_INCLUDE = "/usr/local/cuda-13.0/targets/x86_64-linux/include/cccl" +if os.path.exists(CCCL_INCLUDE): + include_dirs.append(CCCL_INCLUDE) + +setup( + name="cutlass_nvfp4_gemm", + ext_modules=[ + CUDAExtension( + name="cutlass_nvfp4_gemm._C", + sources=[ + "pytorch_binding.cpp", + "cutlass_nvfp4_gemm.cu", + ], + include_dirs=include_dirs, + extra_compile_args={ + "cxx": [ + "-O3", + "-std=c++17", + "-DCUTLASS_ENABLE_GEMP_OPERATION=1", + "-DCUTLASS_ARCH_SM100_ENABLED=1", + ], + "nvcc": [ + "-gencode=arch=compute_100a,code=sm_100a", + "--expt-relaxed-constexpr", + "-DCUTLASS_ENABLE_GEMP_OPERATION=1", + "-DCUTLASS_ARCH_SM100_ENABLED=1", + "--ptxas-options=-v", + "--ptxas-options=-allow-expensive-optimizations=true", + ], + }, + ), + ], + cmdclass={ + "build_ext": BuildExtension, + }, +) diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py new file mode 100644 index 00000000..ec74b4cf --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/sf_layout.py @@ -0,0 +1,21 @@ +""" +CUTLASS NVFP4 scale factor layout — reference documentation. + +CUTLASS's Sm1xxBlockScaledConfig expects scale factors in a specific +interleaved layout (not simple row-major). The layout is defined by: + + SfAtom = Shape, Shape> + with Stride, Stride<_0, _1>> + (SFVecSize=16 for NVFP4 UE4M3 block-16) + + layout_SFA = tile_to_shape(SfAtom{}, make_shape(M, K), Step<_2, _1>) + layout_SFB = tile_to_shape(SfAtom{}, make_shape(N, K), Step<_2, _1>) + +The actual remap from row-major → CUTLASS interleaved layout happens +in the CUDA kernel (remap_sf_to_cutlass_kernel in cutlass_nvfp4_gemm.cu), +NOT in Python. This file exists for reference only. + +The CUDA remap uses cute::idx2crd() to invert the CUTLASS layout: +for each linear index in the CUTLASS layout, it computes the logical +(m, k) coordinate and reads from the corresponding row-major position. +""" diff --git a/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py new file mode 100644 index 00000000..fcfe2a1b --- /dev/null +++ b/src/nvfp4_megamoe_kernel/cutlass_nvfp4_gemm/test_gemm.py @@ -0,0 +1,109 @@ +""" +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/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py index cf32a293..72576c24 100644 --- a/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py +++ b/src/nvfp4_megamoe_kernel/nvfp4_mega_moe.py @@ -5,17 +5,70 @@ This is the main kernel that replaces fp8_nvfp4_mega_moe from DeepGEMM. Architecture: - L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with UE4M3 scales) -- SiLU+Mul activation +- SiLU+Mul activation (per-slot, BEFORE combining expert paths) - L2 GEMM: down_proj (FP4 x FP4 → BF16 with UE4M3 scales) -- NVLink cross-rank sync via symm buffer +- Routing weights applied ONCE at final scatter +- NVLink cross-rank sync handled by caller (not this kernel) - Expert parallel: each rank handles NUM_EXPERTS/8 experts -The kernel is written in TileLang, compiled to SM100 (Blackwell) CUBIN. +The kernel uses native NVFP4 block-scaled MMA via tcgen05.mma +kind::mxf8f6f4.block_scale on Blackwell (SM100). + +Native NVFP4 path: + E2M1 (int8, 2 vals/byte) × E2M1 + UE4M3 block-16 scales + → native hardware block-scaled MMA in tensor cores + → float32 accumulator + +This replaces the dequantize-then-BF16-GEMM approach. The native path +performs the E2M1 × E2M1 with UE4M3 block scaling entirely in hardware, +avoiding the costly dequantization step. """ import os import torch +def unpack_ue4m3_u32(x_u32): + """Unpack uint32 packed UE4M3 scales to float8_e4m3fn. + + Each uint32 contains 4 UE4M3 values packed in bits [0:8], [8:16], [16:24], [24:32]. + Must use bit reinterpret (view), NOT value cast (to) — byte 0x3F is the float8 + whose bits are 0x3F (~0.984), NOT the integer 63. + + CUDA doesn't implement bitwise ops on uint32, so we cast to int32 first. + Supports ND tensors — last dim is the packed dim (N words → N*4 float8 values). + """ + # CUDA uint32 lacks bitwise ops — use int32 + x_i32 = x_u32.to(torch.int32) + *prefix, n_words = x_i32.shape + + # Extract 4 bytes, cast to uint8, then bit-reinterpret to float8_e4m3fn + b0 = (x_i32 & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) + b1 = ((x_i32 >> 8) & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) + b2 = ((x_i32 >> 16) & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) + b3 = ((x_i32 >> 24) & 0xFF).to(torch.uint8).view(torch.float8_e4m3fn) + + # Interleave into (*prefix, n_words*4) + out = torch.empty(*prefix, n_words * 4, dtype=torch.float8_e4m3fn, device=x_u32.device) + out[..., 0::4] = b0 + out[..., 1::4] = b1 + out[..., 2::4] = b2 + out[..., 3::4] = b3 + return out + +# CUTLASS native NVFP4 block-scaled GEMM (SM100 Blackwell) +# Primary path: uses CUTLASS MainloopSm100TmaUmmaWarpSpecializedBlockScaled +# which invokes mxf8f6f4.block_scale tensor core instructions directly. +MEGA_MOE_USE_CUTLASS = int(os.environ.get("MEGA_MOE_USE_CUTLASS", "1")) + +try: + from nvfp4_megamoe_kernel.cutlass_nvfp4_gemm.kernel import ( + cutlass_nvfp4_blockscaled_gemm, + cutlass_grouped_nvfp4_gemm, + ) + _CUTLASS_AVAILABLE = True +except ImportError: + _CUTLASS_AVAILABLE = False + # DeepSeek-V4-Pro dimensions HIDDEN = 7168 INTERMEDIATE = 3072 @@ -32,47 +85,153 @@ MEGA_MOE_STATIC = int(os.environ.get("MEGA_MOE_STATIC", "0")) MEGA_MOE_DEBUG = int(os.environ.get("MEGA_MOE_DEBUG", "0")) + +# --------------------------------------------------------------------------- +# Main kernel entry points +# --------------------------------------------------------------------------- + def nvfp4_mega_moe_l1( x_fp4, # (num_tokens, K//2) int8 packed E2M1 - x_sf, # (num_tokens, sf_k_groups) uint32 packed UE4M3 - l1_weights, # (E_per_rank, 2*INTER, K//2) int8 K-major - l1_scales, # (E_per_rank, 2*INTER, sf_k_groups) uint32 packed UE4M3 - topk_ids, # (num_tokens, NUM_TOPK) int32 - topk_weights, # (num_tokens, NUM_TOPK) float32 - num_experts_per_rank, + x_sf, # (num_tokens, sf_k_groups) float8_e4m3fn + l1_weights, # (E_per_rank, K//2, 2*INTER) int8, column-major for CUTLASS + l1_scales, # (E_per_rank, sf_k_groups, 2*INTER) float8_e4m3fn, column-major + slot_expert_ids, # (num_slots,) int32 — per-slot local expert IDs + slot_token, # (num_slots,) int64 — token index per slot + alpha=1.0, # fp32 scalar from stage_activation global scale ): - """L1 GEMM: gate_up_proj — FP4 x FP4 → BF16 with block scaling. - - TODO: TileLang JIT kernel (nvfp4_blockscaled_gemm_2cta_persistent pattern). + """L1 GEMM: gate_up_proj — slot-based, no routing weights. + + Takes pre-built slot mapping (slot_expert_ids, slot_token) from the outer + routing logic. Returns (slot_out, slot_token) where each slot is one + (token, topk) pair. """ - raise NotImplementedError("nvfp4_mega_moe_l1 TileLang kernel not yet implemented") + K_half = x_fp4.shape[1] + K = K_half * 2 + N = l1_weights.shape[2] # 2 * INTERMEDIATE = 6144 + + if MEGA_MOE_DEBUG: + print(f"[nvfp4_moe_l1] tokens={x_fp4.shape[0]} K={K} N={N} slots={slot_expert_ids.shape[0]}") + + x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf + w_sf_fp8 = unpack_ue4m3_u32(l1_scales) if l1_scales.dtype == torch.uint32 else l1_scales + assert w_sf_fp8.dtype == torch.float8_e4m3fn, f"l1_scales after unpack dtype={w_sf_fp8.dtype}" + + slot_out, slot_token = cutlass_grouped_nvfp4_gemm( + x_fp4, x_sf_fp8, + l1_weights, w_sf_fp8, + slot_expert_ids, # 1D per-slot expert IDs + slot_token, # 1D per-slot token indices + alpha=alpha, + ) + print(f"[L1-GEMM-OUT] slots={slot_out.shape[0]} N={N} amax={slot_out.abs().max().item():.4e} mean={slot_out.float().mean().item():.4e}") + return slot_out, slot_token def nvfp4_mega_moe_l2( - x_fp4, # (num_tokens, INTER//2) int8 packed E2M1 - x_sf, # (num_tokens, sf_k_groups) uint32 packed UE4M3 - l2_weights, # (E_per_rank, HIDDEN, INTER//2) int8 K-major - l2_scales, # (E_per_rank, HIDDEN, sf_k_groups) uint32 packed UE4M3 - topk_ids, # (num_tokens, NUM_TOPK) int32 - topk_weights, # (num_tokens, NUM_TOPK) float32 - num_experts_per_rank, + x_fp4, # (num_slots, INTER//2) int8 packed E2M1 + x_sf, # (num_slots, sf_k_groups) float8_e4m3fn + l2_weights, # (E_per_rank, INTER//2, HIDDEN) int8, column-major for CUTLASS + l2_scales, # (E_per_rank, sf_k_groups, HIDDEN) float8_e4m3fn, column-major + slot_expert_ids, # (num_slots,) int32 — per-slot local expert IDs (from L1 routing) + slot_token, # (num_slots,) int64 — token index per slot (from L1) + alpha=1.0, # fp32 scalar from stage_activation global scale ): - """L2 GEMM: down_proj — FP4 x FP4 → BF16 with block scaling. - - TODO: TileLang JIT kernel (same pattern as L1). + """L2 GEMM: down_proj — slot-based, no routing weights. + + Reuses the same slot mapping from L1 (same slot_token and slot_expert_ids). """ - raise NotImplementedError("nvfp4_mega_moe_l2 TileLang kernel not yet implemented") + K_half = x_fp4.shape[1] + K = K_half * 2 + N = l2_weights.shape[2] + + if MEGA_MOE_DEBUG: + print(f"[nvfp4_moe_l2] slots={x_fp4.shape[0]} K={K} N={N} native=1") + + x_sf_fp8 = unpack_ue4m3_u32(x_sf) if x_sf.dtype == torch.uint32 else x_sf + w_sf_fp8 = unpack_ue4m3_u32(l2_scales) if l2_scales.dtype == torch.uint32 else l2_scales + assert w_sf_fp8.dtype == torch.float8_e4m3fn, f"l2_scales after unpack dtype={w_sf_fp8.dtype}" + + slot_out, _ = cutlass_grouped_nvfp4_gemm( + x_fp4, x_sf_fp8, + l2_weights, w_sf_fp8, + slot_expert_ids, # 1D per-slot expert IDs — GEMM handles directly + alpha=alpha, + ) + return slot_out # (num_slots, HIDDEN) bfloat16 + + +# E2M1 (FP4) representable magnitudes: {0, 0.5, 1, 1.5, 2, 3, 4, 6} +# Bit patterns (3-bit, no sign): 000=0, 001=0.5, 010=1, 011=1.5, 100=2, 101=3, 110=4, 111=6 +# Full 4-bit nibble: bit 3 = sign, bits 2:0 = magnitude index +_E2M1_MAGNITUDES = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float32) + + +def _quantize_to_e2m1(x_f32): + """Quantize float32 values to E2M1 (FP4) nibble indices. + + Maps each value to the nearest E2M1 representable magnitude, + then packs as 4-bit sign-magnitude nibbles. + + Returns (nibbles, scales) where: + nibbles: (..., N) uint8 with 4-bit sign-magnitude per value + scales: (..., N//16) float8_e4m3fn block scales + """ + *batch, N = x_f32.shape + assert N % 16 == 0, f"Last dim {N} not divisible by 16 (block size)" + + # Reshape into blocks of 16 for block-wise scaling + x_blocks = x_f32.reshape(*batch, N // 16, 16) + + # Per-block absmax determines the scale + block_max = x_blocks.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8) + + # Scale so that the max maps to 6.0 (largest E2M1 magnitude) + scale_f32 = (block_max / 6.0).clamp(min=1e-8, max=448.0) + x_scaled = x_blocks / scale_f32.clamp(min=1e-8) + + # Find nearest E2M1 magnitude for each value + signs = torch.sign(x_scaled) + abs_scaled = x_scaled.abs() + + mags = _E2M1_MAGNITUDES.to(device=abs_scaled.device) + dists = (abs_scaled.unsqueeze(-1) - mags).abs() + idx = dists.argmin(dim=-1) + + idx = idx.clamp(0, 7).to(torch.uint8) + + sign_bit = (signs < 0).to(torch.uint8) + nibbles = (sign_bit << 3) | idx + + nibbles = nibbles.reshape(*batch, N // 2, 2) + packed = (nibbles[..., 1] << 4) | nibbles[..., 0] + + sf = scale_f32.squeeze(-1).to(torch.float8_e4m3fn) + + return packed.to(torch.int8), sf def stage_activation(x_bf16): """Quantize BF16 activation to FP4 (E2M1) with UE4M3 block16 scales. - This replaces the Triton staging kernel from patches/staging_kernel.py. + Two-level quantization matching the NVFP4 weight format: + 1. Per-tensor global scale: amax / (6.0 * 448.0) + 2. Per-block (16 values) absmax scaling on the normalized values + + Returns (x_fp4, x_sf, input_global_scale) where: + x_fp4: packed E2M1 nibbles + x_sf: UE4M3 block scales (NOT folded with global scale) + input_global_scale: fp32 per-tensor scale, applied as GEMM alpha """ - from vllm.model_executor.layers.quantization.utils.fp4_utils import ( - per_token_cast_to_fp4, - ) - return per_token_cast_to_fp4(x_bf16) + x_f32 = x_bf16.float() + + x_amax = x_f32.abs().amax().to(torch.float32).clamp(min=1e-8) + input_global_scale = x_amax / (6.0 * 448.0) + + x_normalized = x_f32 / input_global_scale + + x_fp4, x_sf = _quantize_to_e2m1(x_normalized) + + return x_fp4, x_sf, input_global_scale def nvfp4_mega_moe_full( @@ -84,71 +243,154 @@ def nvfp4_mega_moe_full( fast_math=False, # fast math flag (unused in NVFP4) ): """Full mega_moe forward pass — replaces deep_gemm.mega.fp8_nvfp4_mega_moe. - - API matches the DeepGEMM fp8_nvfp4_mega_moe call signature used in - the vLLM deepseek_v4.py patch: - - fp8_nvfp4_mega_moe(y, l1_weights, l2_weights, symm_buffer, - activation_clamp=..., fast_math=...) - - Pipeline: - 1. Read staged activation from symm_buffer (already quantized by staging kernel) - 2. L1 GEMM: gate_up_proj (FP4 x FP4 → BF16 with block scaling) - 3. SiLU + Mul (activation) - 4. Quantize L1 output → FP4 + UE4M3 scales - 5. L2 GEMM: down_proj (FP4 x FP4 → BF16 with block scaling) - 6. NVLink sync + reduce across ranks → write to y - - When MEGA_MOE_STATIC=1, returns zeros (bypass) for pipeline testing. + + Slot-based pipeline (routing weights applied ONCE at final scatter): + 1. Read staged activation from symm_buffer + 2. L1 GEMM → slot output (num_slots, 2*INTER) — NO routing weights + 3. SiLU + Mul PER SLOT (nonlinearity before combining expert paths) + 4. Quantize activated slots → FP4 + 5. L2 GEMM → slot output (num_slots, HIDDEN) — NO routing weights + 6. Final scatter: y.index_add_(0, slot_token, slot_weight * l2_slots) + Single routing weight application. """ num_tokens = y.shape[0] device = y.device dtype = y.dtype - + if MEGA_MOE_STATIC: if MEGA_MOE_DEBUG: print(f"[MEGA_MOE_STATIC] Skipping nvfp4_mega_moe, returning zeros " f"shape=({num_tokens}, {y.shape[1]})") y.zero_() return - + # Unpack transformed weights l1_w, l1_sf = transformed_l1_weights l2_w, l2_sf = transformed_l2_weights - + + # Expert sanity check — are experts actually distinct? + if not getattr(nvfp4_mega_moe_full, '_expert_sanity', False): + nvfp4_mega_moe_full._expert_sanity = True + for e in range(min(4, l1_w.shape[0])): + w_sample = l1_w[e].view(torch.uint8)[:8, :8] + sf_sample = l1_sf[e].to(torch.float32)[:4, :4] + print(f"[EXPERT-SANITY e={e}] w_bytes[:8,:8]={w_sample.flatten().tolist()[:16]}") + print(f"[EXPERT-SANITY e={e}] sf[:4,:4]={sf_sample.flatten().tolist()[:8]}") + # Step 1: Read staged activation from symm_buffer - # The staging has already been done by _stage_deepseek_v4_mega_moe_inputs - # and stored in symm_buffer.x, symm_buffer.x_sf x_fp4 = symm_buffer.x[:num_tokens] x_sf = symm_buffer.x_sf[:num_tokens] + l1_global_scale = symm_buffer.input_global_scale topk_ids = symm_buffer.topk_idx[:num_tokens] topk_weights = symm_buffer.topk_weights[:num_tokens] - + + _x_sf_f32 = x_sf.to(torch.float32) + _igs = l1_global_scale if isinstance(l1_global_scale, float) else l1_global_scale.item() if hasattr(l1_global_scale, 'item') else float(l1_global_scale) + if MEGA_MOE_DEBUG: + print(f"[ALPHA L1] alpha={_igs:.4e} x_sf range [{_x_sf_f32.min().item():.4e}, {_x_sf_f32.max().item():.4e}]") + + # Convert global expert IDs to local expert IDs + num_experts_per_rank = l1_w.shape[0] + experts_start_idx = symm_buffer.experts_start_idx + topk_ids_local = topk_ids - experts_start_idx + + # Build slot mapping for this rank + local_topk = (topk_ids >= experts_start_idx) & (topk_ids < experts_start_idx + num_experts_per_rank) + slot_token, slot_k = local_topk.nonzero(as_tuple=True) + slot_expert_local = topk_ids_local[slot_token, slot_k] + slot_weight = topk_weights[slot_token, slot_k] + num_slots = slot_token.shape[0] + + tokens_routed_locally = local_topk.any(dim=-1).sum().item() + print(f"[ROUTING] tokens_routed_local={tokens_routed_locally}/{num_tokens} " + f"num_slots={num_slots}") + if MEGA_MOE_DEBUG: print(f"[nvfp4_mega_moe_full] x_fp4={x_fp4.shape} x_sf={x_sf.shape} " - f"topk_ids={topk_ids.shape} l1_w={l1_w.shape} l2_w={l2_w.shape}") - - # Step 2: L1 GEMM - num_experts_per_rank = l1_w.shape[0] - l1_output = nvfp4_mega_moe_l1( + f"topk_ids range: {topk_ids.min().item()}-{topk_ids.max().item()} " + f"local: {topk_ids_local.min().item()}-{topk_ids_local.max().item()} " + f"slots={num_slots}") + + # Handle no local slots + if num_slots == 0: + y.zero_() + return + + # Ensure alpha is a plain Python float (C extension can't handle torch scalars) + l1_alpha = float(l1_global_scale) if not isinstance(l1_global_scale, float) else l1_global_scale + + # Shape consistency asserts — catch mismatched slot mappings early + assert slot_expert_local.ndim == 1 + assert slot_token.ndim == 1 + assert slot_weight.ndim == 1 + assert slot_expert_local.numel() == num_slots + assert slot_token.numel() == num_slots + assert slot_weight.numel() == num_slots + + # SFB weight scales are remapped per-expert inside CUTLASS on each call. + # ───────────────────────────────────────────────────────────────────── + # NO PREPACK CACHE — see README for rationale. + # DO NOT add a prepack cache. Previous attempts caused: + # - OOM: ~1.75 GiB per prepacked tensor × 61 layers = 214 GiB + # - Peak memory 2× during torch.stack before eviction + # - CUDA graph use-after-free on evicted entries + # - M_for_layout=128 assumption (unverified M-independence) + # The SFB remap is a small scatter kernel (~µs) — not the bottleneck. + # ───────────────────────────────────────────────────────────────────── + + # Step 2: L1 GEMM — slot-based, no routing weights + l1_slots, _ = nvfp4_mega_moe_l1( x_fp4, x_sf, l1_w, l1_sf, - topk_ids, topk_weights, num_experts_per_rank, - ) - - # Step 3: SiLU + Mul - gate, up = l1_output.chunk(2, dim=-1) + slot_expert_local, slot_token, + alpha=l1_alpha, + ) # (num_slots, 2*INTER) bfloat16 + + # Post-L1 shape asserts + assert l1_slots.shape[0] == num_slots + + if MEGA_MOE_DEBUG: + print(f"[L1-out] nan={torch.isnan(l1_slots).any().item()} " + f"abs_max={l1_slots.abs().max().item():.4e}") + + # Step 3: SiLU + Mul PER SLOT — nonlinearity before combining paths + gate, up = l1_slots.chunk(2, dim=-1) activated = torch.nn.functional.silu(gate) * up if activation_clamp is not None: activated = activated.clamp(max=activation_clamp) - - # Step 4: Quantize L1 output → FP4 - l1_fp4, l1_sf_out = stage_activation(activated) - - # Step 5: L2 GEMM - l2_output = nvfp4_mega_moe_l2( + + if MEGA_MOE_DEBUG: + print(f"[silu] nan={torch.isnan(activated).any().item()} " + f"abs_max={activated.abs().max().item():.4e}") + + # Step 4: Quantize activated slots → FP4 + l1_fp4, l1_sf_out, l2_global_scale = stage_activation(activated) + + # Pre-L2 shape asserts + assert activated.shape[0] == num_slots + assert l1_fp4.shape[0] == num_slots + assert l1_sf_out.shape[0] == num_slots + l2_alpha = float(l2_global_scale) if not isinstance(l2_global_scale, float) else l2_global_scale + + if MEGA_MOE_DEBUG: + _l1sf_f32 = l1_sf_out.to(torch.float32) + _l2gs = l2_global_scale if isinstance(l2_global_scale, float) else l2_global_scale.item() + print(f"[ALPHA L2] alpha={_l2gs:.4e} l1_sf range [{_l1sf_f32.min().item():.4e}, {_l1sf_f32.max().item():.4e}]") + + # Step 5: L2 GEMM — slot-based, no routing weights + l2_slots = nvfp4_mega_moe_l2( l1_fp4, l1_sf_out, l2_w, l2_sf, - topk_ids, topk_weights, num_experts_per_rank, + slot_expert_local, slot_token, + alpha=l2_alpha, + ) # (num_slots, HIDDEN) bfloat16 + + if MEGA_MOE_DEBUG: + print(f"[L2-out] nan={torch.isnan(l2_slots).any().item()} " + f"abs_max={l2_slots.abs().max().item():.4e}") + + # Step 6: Final scatter — routing weights applied ONCE + y.zero_() + y.index_add_( + 0, + slot_token, + l2_slots * slot_weight.to(l2_slots.dtype).unsqueeze(1), ) - - # Step 6: Write to output - y.copy_(l2_output) diff --git a/src/nvfp4_megamoe_kernel/symm_buffer.py b/src/nvfp4_megamoe_kernel/symm_buffer.py index a0282dff..d974f1ec 100644 --- a/src/nvfp4_megamoe_kernel/symm_buffer.py +++ b/src/nvfp4_megamoe_kernel/symm_buffer.py @@ -31,21 +31,27 @@ class SymmBuffer: self.top_k = top_k self.hidden_size = hidden_size self.intermediate_size = intermediate_size + self.experts_start_idx = 0 # set by caller before kernel invocation device = torch.cuda.current_device() - # NVFP4: packed E2M1 (2 values per byte), so K//2 - sf_k_groups_hidden = hidden_size // (16 * 4) # UE4M3 block16, 4 packed per uint32 - sf_k_groups_inter = intermediate_size // (16 * 4) - - # Staging buffers (matching DeepGEMM layout) + # NVFP4 packed E2M1: 2 FP4 values per byte → K//2 bytes per token. + # Scales are UE4M3 (float8_e4m3fn), one per 16-element group → K//16 + # bytes per token, UNPACKED. This is what `stage_activation` produces + # and what the CUTLASS NVFP4 block-scaled GEMM consumes directly. + # (The DeepGEMM API packed 4 UE4M3 into one uint32 — we don't, because + # our CUTLASS kernel reads scales as float8_e4m3fn.) + sf_k_groups_hidden = hidden_size // 16 + sf_k_groups_inter = intermediate_size // 16 + + # Staging buffers self.x = torch.empty( max_num_tokens, hidden_size // 2, dtype=torch.int8, device=device, ) self.x_sf = torch.empty( max_num_tokens, sf_k_groups_hidden, - dtype=torch.uint32, device=device, + dtype=torch.float8_e4m3fn, device=device, ) self.topk_idx = torch.empty( max_num_tokens, top_k, @@ -62,6 +68,10 @@ class SymmBuffer: dtype=torch.bfloat16, device=device, ) + # Per-tensor global scale from stage_activation (fp32 scalar) + # Applied as GEMM alpha: D = global_scale * (A_sf * A_fp4) @ (B_sf * B_fp4) + self.input_global_scale = 1.0 + if MEGA_MOE_DEBUG: print(f"[SymmBuffer] x={self.x.shape} x_sf={self.x_sf.shape} " f"topk_idx={self.topk_idx.shape} topk_weights={self.topk_weights.shape} " @@ -83,4 +93,4 @@ def get_symm_buffer_for_nvfp4_mega_moe( return SymmBuffer( group, num_experts, max_num_tokens, top_k, hidden_size, intermediate_size, - ) + ) \ No newline at end of file diff --git a/src/nvfp4_megamoe_kernel/weight_transform.py b/src/nvfp4_megamoe_kernel/weight_transform.py index 8ab2cad9..c56345df 100644 --- a/src/nvfp4_megamoe_kernel/weight_transform.py +++ b/src/nvfp4_megamoe_kernel/weight_transform.py @@ -1,118 +1,162 @@ """ -NVFP4 Weight Transformation for TileLang mega_moe kernel. +NVFP4 Weight Transformation for CUTLASS mega_moe kernel. Converts raw NVFP4 checkpoint weights (uint8 E2M1 + float8_e4m3fn UE4M3 + float32 global scale) -into the TMA-aligned format expected by the block-scaled GEMM kernel: -- Packed FP4 weights (uint8, K-major) -- Packed UE4M3 scales (uint32, TMA-aligned UTCCP layout) +into the format expected by the CUTLASS block-scaled GEMM kernel: +- Packed FP4 weights (int8, K-major) +- UE4M3 block scales (float8_e4m3fn, row-major — CUTLASS SF remap handles interleaving) -This replaces deep_gemm.mega.transform_nvfp4_weights_for_mega_moe. +Weight scales are returned as float8_e4m3fn (NOT packed uint32). The CUTLASS GEMM +consumes float8 scales directly; only activation scales from the staging kernel come +as uint32 and need unpack_ue4m3_u32. + +This replaces deep_gemma.mega.transform_nvfp4_weights_for_mega_moe. + +Call signature matches the nightly vLLM deepseek_v4.py finalize_weights: + transform_nvfp4_weights_for_mega_moe( + (l1_weight, l1_weight_scale), + (l2_weight, l2_weight_scale), + l1_weight_scale_2=..., + l2_weight_scale_2=..., + ) """ import torch -import math -def fold_global_scale( - weight_scale: torch.Tensor, # (N, K//16) float8_e4m3fn - weight_scale_2: torch.Tensor, # (num_logical,) or scalar float32 - logical_widths: list[int] = None, # per-logical-weight row counts +def _fold_global_scale( + weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn + weight_scale_2: torch.Tensor, # (E,) or (E, 2) or scalar float32 ) -> torch.Tensor: - """Fold global scale into block scales: UE4M3 * FP32 → UE4M3 → float32. + """Fold global scale into block scales: UE4M3 * FP32 → float32. - Returns: (N, K//16) float32 folded block scales. + For fused projections (w13 = gate+up), weight_scale_2 is (E, 2): + scale_2[e, 0] applies to gate_proj rows, scale_2[e, 1] applies to up_proj rows. + N is split: gate = weight_scale[:, :N//2, :], up = weight_scale[:, N//2:, :] + For single projections (w2), weight_scale_2 is (E,) or scalar. """ sf_f32 = weight_scale.to(torch.float32) - - if weight_scale_2.numel() == 1: - sf_f32 = sf_f32 * weight_scale_2.to(torch.float32) - elif weight_scale_2.numel() > 1 and logical_widths is not None: - # Per-logical-weight global scale (e.g., gate_up_proj has 2) - expanded = [] - for i, w in enumerate(logical_widths): - if i < len(weight_scale_2): - expanded.append(weight_scale_2[i].flatten()[0].expand(w)) - global_scale = torch.cat(expanded).to(torch.float32).unsqueeze(1) - sf_f32 = sf_f32 * global_scale + gs = weight_scale_2.to(torch.float32) + + if gs.numel() == 1: + sf_f32 = sf_f32 * gs + elif gs.dim() == 2 and gs.shape[1] == 2: + # Fused projection: (E, 2) — gate and up have separate global scales + # weight_scale is (E, N, K//16), N = gate_N + up_N + gate_N = sf_f32.shape[1] // 2 + gs_gate = gs[:, 0].unsqueeze(-1) # (E, 1) + gs_up = gs[:, 1].unsqueeze(-1) # (E, 1) + sf_f32[:, :gate_N, :] = sf_f32[:, :gate_N, :] * gs_gate.unsqueeze(-1) + sf_f32[:, gate_N:, :] = sf_f32[:, gate_N:, :] * gs_up.unsqueeze(-1) else: - sf_f32 = sf_f32 * weight_scale_2.max().to(torch.float32) - + # Per-expert global scale — broadcast multiply + while gs.dim() < sf_f32.dim(): + gs = gs.unsqueeze(-1) + sf_f32 = sf_f32 * gs.expand_as(sf_f32) + return sf_f32 -def pack_ue4m3_to_uint32(sf: torch.Tensor) -> torch.Tensor: - """Pack 4 UE4M3 (float8_e4m3fn) values into one uint32. - - Input: (..., K//16) float8_e4m3fn - Output: (..., K//64) uint32 (4 values packed per word) - """ - # View as raw bytes - sf_u8 = sf.view(torch.uint8) # (..., K//16) uint8 - shape = sf_u8.shape - assert shape[-1] % 4 == 0, f"Last dim {shape[-1]} not divisible by 4" - - # Pack 4 consecutive uint8 values into one uint32 +def _pack_ue4m3_to_uint32(sf: torch.Tensor) -> torch.Tensor: + """Pack 4 UE4M3 (float8_e4m3fn) values into one uint32.""" + sf_u8 = sf.view(torch.uint8) + assert sf_u8.shape[-1] % 4 == 0, f"Last dim {sf_u8.shape[-1]} not divisible by 4" packed = (sf_u8[..., 0::4].to(torch.int32) | (sf_u8[..., 1::4].to(torch.int32) << 8) | (sf_u8[..., 2::4].to(torch.int32) << 16) | (sf_u8[..., 3::4].to(torch.int32) << 24)) - return packed.contiguous() -def interleave_l1_weights( - weight: torch.Tensor, # (E, 2*INTER, K//2) int8, K-major -) -> torch.Tensor: - """Interleave L1 (gate_up) weights for 2CTA UMMA. +def transform_nvfp4_weights_for_mega_moe( + l1_tuple: tuple[torch.Tensor, torch.Tensor], # (weight, weight_scale) + l2_tuple: tuple[torch.Tensor, torch.Tensor], # (weight, weight_scale) + l1_weight_scale_2: torch.Tensor = None, # float32 global scale for L1 + l2_weight_scale_2: torch.Tensor = None, # float32 global scale for L2 +) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]: + """Transform NVFP4 weights for the CUTLASS block-scaled GEMM. - The gate and up projections are interleaved in groups of 8 rows - to match the UMMA 2CTA schedule. - """ - E, N, K_half = weight.shape - assert N % 16 == 0, f"N={N} not divisible by 16" + Matches the call signature from nightly vLLM deepseek_v4.py finalize_weights. - # Reshape to (E, N//16, 16, K_half) → interleave pairs of 8 - w = weight.view(E, N // 16, 16, K_half) - w_gate = w[:, :, :8, :] - w_up = w[:, :, 8:16, :] - interleaved = torch.stack([w_gate, w_up], dim=3).reshape(E, N, K_half) - - return interleaved.contiguous() - - -def transform_nvfp4_weights_for_tilelang( - weight: torch.Tensor, # (E, N, K//2) int8, K-major packed FP4 - weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn UE4M3 - weight_scale_2: torch.Tensor, # (E, 1) or (E, 2) float32 global scale - N: int, # output dimension (2*INTER for L1, HIDDEN for L2) - K: int, # input dimension (HIDDEN for L1, INTER for L2) - gran_k: int = 16, # NVFP4 group_size - is_l1: bool = False, # True for gate_up_proj (needs interleaving) - logical_widths: list[int] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Transform NVFP4 weights for the TileLang block-scaled GEMM. + Args: + l1_tuple: (w13_weight, w13_weight_scale) — gate_up proj + l2_tuple: (w2_weight, w2_weight_scale) — down proj + l1_weight_scale_2: global scale for L1 (float32) + l2_weight_scale_2: global scale for L2 (float32) Returns: - transformed_weight: (E, N, K//2) int8, K-major, contiguously interleaved - transformed_sf: (E, N, K//64) uint32, packed UE4M3 with folded global scale + ((l1_weight, l1_sf_packed), (l2_weight, l2_sf_packed)) """ - # Step 1: Fold global scale into block scales - sf_folded = fold_global_scale(weight_scale, weight_scale_2, logical_widths) - - # Step 2: Clamp and convert back to UE4M3 - sf_clamped = sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn) - - # Step 3: Pack UE4M3 into uint32 - sf_packed = pack_ue4m3_to_uint32(sf_clamped) # (E, N, K//64) - - # Step 4: Interleave L1 weights if needed - if is_l1: - transformed_weight = interleave_l1_weights(weight) + l1_weight, l1_weight_scale = l1_tuple + l2_weight, l2_weight_scale = l2_tuple + + # DEBUG: check raw scales before folding + l1_sf_f32_raw = l1_weight_scale.to(torch.float32) + l1_gs_raw = l1_weight_scale_2.to(torch.float32) if l1_weight_scale_2 is not None else None + if not getattr(transform_nvfp4_weights_for_mega_moe, '_sf_debug', False): + transform_nvfp4_weights_for_mega_moe._sf_debug = True + print(f"[SF-DEBUG] raw l1_sf dtype={l1_weight_scale.dtype} range=[{l1_sf_f32_raw.min().item():.4e}, {l1_sf_f32_raw.max().item():.4e}] " + f"unique_raw={torch.unique(l1_weight_scale.view(torch.uint8)).numel()}") + if l1_gs_raw is not None: + print(f"[SF-DEBUG] l1_gs dtype={l1_weight_scale_2.dtype} shape={tuple(l1_weight_scale_2.shape)} " + f"range=[{l1_gs_raw.min().item():.4e}, {l1_gs_raw.max().item():.4e}] " + f"unique_gs={torch.unique(l1_gs_raw).numel()}") + if l1_gs_raw.dim() == 2 and l1_gs_raw.shape[1] == 2: + print(f"[SF-DEBUG] gate gs unique={torch.unique(l1_gs_raw[:, 0]).numel()} " + f"up gs unique={torch.unique(l1_gs_raw[:, 1]).numel()}") + + # DEBUG: check L2 scales + l2_sf_f32_raw = l2_weight_scale.to(torch.float32) + l2_gs_raw = l2_weight_scale_2.to(torch.float32) if l2_weight_scale_2 is not None else None + if not getattr(transform_nvfp4_weights_for_mega_moe, '_sf_debug_l2', False): + transform_nvfp4_weights_for_mega_moe._sf_debug_l2 = True + print(f"[SF-DEBUG-L2] raw l2_sf dtype={l2_weight_scale.dtype} range=[{l2_sf_f32_raw.min().item():.4e}, {l2_sf_f32_raw.max().item():.4e}] " + f"unique_raw={torch.unique(l2_weight_scale.view(torch.uint8)).numel()}") + if l2_gs_raw is not None: + print(f"[SF-DEBUG-L2] l2_gs dtype={l2_weight_scale_2.dtype} shape={tuple(l2_weight_scale_2.shape)} " + f"range=[{l2_gs_raw.min().item():.4e}, {l2_gs_raw.max().item():.4e}] " + f"unique_gs={torch.unique(l2_gs_raw).numel()}") + + # Post-fold diagnostics — one-time + if not getattr(transform_nvfp4_weights_for_mega_moe, '_sf_debug_fold', False): + transform_nvfp4_weights_for_mega_moe._sf_debug_fold = True + l1_sf_folded = _fold_global_scale(l1_weight_scale, l1_weight_scale_2) if l1_weight_scale_2 is not None else l1_weight_scale.to(torch.float32) + l1_sf_out_check = l1_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn) + l2_sf_folded = _fold_global_scale(l2_weight_scale, l2_weight_scale_2) if l2_weight_scale_2 is not None else l2_weight_scale.to(torch.float32) + l2_sf_out_check = l2_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn) + print(f"[SF-FOLD] l1 pre-fold unique_u8={torch.unique(l1_weight_scale.view(torch.uint8)).numel()} " + f"post-fold unique_u8={torch.unique(l1_sf_out_check.view(torch.uint8)).numel()} " + f"range=[{l1_sf_folded.min().item():.4e}, {l1_sf_folded.max().item():.4e}]") + print(f"[SF-FOLD] l2 pre-fold unique_u8={torch.unique(l2_weight_scale.view(torch.uint8)).numel()} " + f"post-fold unique_u8={torch.unique(l2_sf_out_check.view(torch.uint8)).numel()} " + f"range=[{l2_sf_folded.min().item():.4e}, {l2_sf_folded.max().item():.4e}]") + + # Fold global scales into block scales + # The logical_widths branch was wrong: it treated gs as per-projection + # scalars and only used experts 0 and 1's scales for ALL experts. + # The else branch correctly broadcasts each expert's own global scale. + if l1_weight_scale_2 is not None: + l1_sf_folded = _fold_global_scale(l1_weight_scale, l1_weight_scale_2) else: - transformed_weight = weight.contiguous() - - return transformed_weight, sf_packed + l1_sf_folded = l1_weight_scale.to(torch.float32) + if l2_weight_scale_2 is not None: + l2_sf_folded = _fold_global_scale(l2_weight_scale, l2_weight_scale_2) + else: + l2_sf_folded = l2_weight_scale.to(torch.float32) -# Alias for drop-in replacement of deep_gemm.mega.transform_nvfp4_weights_for_mega_moe -transform_nvfp4_weights_for_mega_moe = transform_nvfp4_weights_for_tilelang + # Clamp and convert back to UE4M3 + l1_sf_out = l1_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn).contiguous() + l2_sf_out = l2_sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn).contiguous() + + # CUTLASS B is declared ColumnMajor — it expects (K, N) in memory. + # Checkpoint weights are (N, K_half) row-major, so we transpose to (K_half, N) + # which is column-major (N, K_half). This is a one-time cost at load time. + l1_weight_out = l1_weight.transpose(-2, -1).contiguous() + l2_weight_out = l2_weight.transpose(-2, -1).contiguous() + + # Same for scale factors: (N, sf_k) row-major → (sf_k, N) column-major + l1_sf_out = l1_sf_out.transpose(-2, -1).contiguous() + l2_sf_out = l2_sf_out.transpose(-2, -1).contiguous() + + return (l1_weight_out, l1_sf_out), (l2_weight_out, l2_sf_out) diff --git a/src/quantize.mojo b/src/quantize.mojo deleted file mode 100644 index 71a7126f..00000000 --- a/src/quantize.mojo +++ /dev/null @@ -1,95 +0,0 @@ -""" -NVFP4 quantization utilities — E2M1 packing and UE4M3 scale handling. -Core math layer for the NVFP4 mega_moe kernel rewrite. -""" - -# E2M1 magnitude lookup table (positive values only) -# Index 0-7 maps to: 0, 0.5, 1, 1.5, 2, 3, 4, 6 -def e2m1_magnitude(index: Int) -> Float64: - if index == 0: return 0.0 - if index == 1: return 0.5 - if index == 2: return 1.0 - if index == 3: return 1.5 - if index == 4: return 2.0 - if index == 5: return 3.0 - if index == 6: return 4.0 - if index == 7: return 6.0 - return 0.0 - - -def quantize_e2m1(value: Float64) -> UInt8: - """Quantize a float64 value to E2M1 (4-bit), returning the 4-bit nibble with sign.""" - var sign = 0 - var abs_val = value - if value < 0.0: - sign = 1 - abs_val = -value - - # Find best E2M1 match - var best_idx = 0 - var best_err = abs_val # error for idx=0 - - for i in range(1, 8): - mag = e2m1_magnitude(i) - err = abs(abs_val - mag) - if err < best_err: - best_err = err - best_idx = i - - return (sign << 3) | best_idx - - -def unpack_e2m1(packed: UInt8, idx: Int) -> Float64: - """Unpack one E2M1 value from a packed byte. - idx=0 -> low nibble, idx=1 -> high nibble. - """ - nibble: UInt8 - if idx == 0: - nibble = packed & 0x0F - else: - nibble = (packed >> 4) & 0x0F # keep sign bit - - sign = (nibble >> 3) & 1 - mag_idx = nibble & 0x07 - magnitude = e2m1_magnitude(Int(mag_idx)) - - if sign: - return -magnitude - return magnitude - - -def dequantize_nvfp4_weight( - packed_weight: UInt8, - block_scale: Float64, - group_offset: Int, -) -> Float64: - """Dequantize a single NVFP4 weight element. - weight = E2M1_magnitude * block_scale - (global_scale is already folded into block_scale) - """ - e2m1_value = unpack_e2m1(packed_weight, group_offset) - return e2m1_value * block_scale - - -def main() raises: - # Test E2M1 quantization round-trip - print("E2M1 quantization test:") - for val in [-6.0, -4.0, -3.0, -2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]: - packed = quantize_e2m1(val) - unpacked = unpack_e2m1(packed, 0) - print(" ", val, " -> E2M1 -> ", unpacked) - - # Test packed byte (two E2M1 values) - print("\nPacked byte test:") - lo = 1.5 - hi = -3.0 - packed = (quantize_e2m1(hi) << 4) | quantize_e2m1(lo) - print(" lo=", lo, " hi=", hi, " packed=", packed) - print(" unpack lo=", unpack_e2m1(packed, 0), " unpack hi=", unpack_e2m1(packed, 1)) - - # Test NVFP4 dequantization - print("\nNVFP4 dequantization test:") - packed_w = UInt8(0x36) # low=6.0, high=3.0 - scale = 0.5 - print(" packed=0x36, scale=0.5, lo=", dequantize_nvfp4_weight(packed_w, scale, 0)) - print(" packed=0x36, scale=0.5, hi=", dequantize_nvfp4_weight(packed_w, scale, 1)) diff --git a/src/staging.py b/src/staging.py deleted file mode 100644 index 42a9ec4d..00000000 --- a/src/staging.py +++ /dev/null @@ -1,65 +0,0 @@ -""" -NVFP4 Activation Quantization — BF16 → packed FP4 + UE4M3 block16 scales - -Port of patches/staging_kernel.py to TileLang. - -This quantizes BF16 hidden states to: -- Packed FP4 (E2M1, 2 values per uint8 byte) -- UE4M3 block16 scales (4 values per uint32 word, group_size=16) -""" - -import torch -import tilelang -import tilelang.language as T - - -@tilelang.jit -def fp4_quant_kernel( - X, # (M, K) bfloat16 input - block_M=128, - block_K=128, - group_size=16, # NVFP4 UE4M3 group_size -): - """Quantize BF16 → packed FP4 (E2M1) + UE4M3 block16 scales. - - Output: - - X_FP4: (M, K//2) uint8 packed E2M1 - - X_SF: (M, K//group_size//4) uint32 packed UE4M3 scales - """ - M, K = T.const("M, K") - X: T.Tensor[[M, K], "bfloat16"] - - X_FP4 = T.empty((M, K // 2), "uint8") - X_SF = T.empty((M, K // (group_size * 4)), "uint32") - - with T.Kernel(T.ceildiv(M, block_M), T.ceildiv(K, block_K), threads=128) as (bx, by): - # Each block quantizes a (block_M, block_K) tile - tx = T.get_thread_binding() - - row = bx * block_M + tx // (block_K // 2) - col = tx % (block_K // 2) - - if row < M and col < K // 2: - # Load 2 BF16 values - val_lo = X[row, col * 2] - val_hi = X[row, col * 2 + 1] - - # Quantize to E2M1 (4-bit each, pack 2 per byte) - # ... (E2M1 quantization logic) - packed = T.uint8(0) # placeholder - X_FP4[row, col] = packed - - # Compute block scale for group of 16 elements - # ... (UE4M3 scale computation) - - return X_FP4, X_SF - - -def fp4_quant_reference(x_bf16, group_size=16): - """Reference FP4 quantization using PyTorch + DeepGEMM. - - Used for correctness verification of the TileLang kernel. - """ - from deep_gemm import per_token_cast_to_fp4 - x_fp4, x_sf = per_token_cast_to_fp4(x_bf16) - return x_fp4, x_sf diff --git a/src/weight_transform.py b/src/weight_transform.py deleted file mode 100644 index 528dbdb5..00000000 --- a/src/weight_transform.py +++ /dev/null @@ -1,147 +0,0 @@ -""" -NVFP4 Weight Transformation for TileLang mega_moe kernel. - -Converts raw NVFP4 checkpoint weights (uint8 E2M1 + float8_e4m3fn UE4M3 + float32 global scale) -into the TMA-aligned format expected by the block-scaled GEMM kernel: -- Packed FP4 weights (uint8, K-major) -- Packed UE4M3 scales (uint32, TMA-aligned UTCCP layout) - -This replaces deep_gemm.mega.transform_nvfp4_weights_for_mega_moe. -""" - -import torch -import math - - -def fold_global_scale( - weight_scale: torch.Tensor, # (N, K//16) float8_e4m3fn - weight_scale_2: torch.Tensor, # (num_logical,) or scalar float32 - logical_widths: list[int] = None, # per-logical-weight row counts -) -> torch.Tensor: - """Fold global scale into block scales: UE4M3 * FP32 → UE4M3 → float32. - - Returns: (N, K//16) float32 folded block scales. - """ - sf_f32 = weight_scale.to(torch.float32) - - if weight_scale_2.numel() == 1: - sf_f32 = sf_f32 * weight_scale_2.to(torch.float32) - elif weight_scale_2.numel() > 1 and logical_widths is not None: - # Per-logical-weight global scale (e.g., gate_up_proj has 2) - expanded = [] - for i, w in enumerate(logical_widths): - if i < len(weight_scale_2): - expanded.append(weight_scale_2[i].flatten()[0].expand(w)) - global_scale = torch.cat(expanded).to(torch.float32).unsqueeze(1) - sf_f32 = sf_f32 * global_scale - else: - sf_f32 = sf_f32 * weight_scale_2.max().to(torch.float32) - - return sf_f32 - - -def pack_ue4m3_to_uint32(sf: torch.Tensor) -> torch.Tensor: - """Pack 4 UE4M3 (float8_e4m3fn) values into one uint32. - - Input: (..., K//16) float8_e4m3fn - Output: (..., K//64) uint32 (4 values packed per word) - """ - # View as raw bytes - sf_u8 = sf.view(torch.uint8) # (..., K//16) uint8 - shape = sf_u8.shape - assert shape[-1] % 4 == 0, f"Last dim {shape[-1]} not divisible by 4" - - # Pack 4 consecutive uint8 values into one uint32 - packed = (sf_u8[..., 0::4].to(torch.int32) | - (sf_u8[..., 1::4].to(torch.int32) << 8) | - (sf_u8[..., 2::4].to(torch.int32) << 16) | - (sf_u8[..., 3::4].to(torch.int32) << 24)) - - return packed.contiguous() - - -def interleave_l1_weights( - weight: torch.Tensor, # (E, 2*INTER, K//2) int8, K-major -) -> torch.Tensor: - """Interleave L1 (gate_up) weights for 2CTA UMMA. - - The gate and up projections are interleaved in groups of 8 rows - to match the UMMA 2CTA schedule. - """ - E, N, K_half = weight.shape - assert N % 16 == 0, f"N={N} not divisible by 16" - - # Reshape to (E, N//16, 16, K_half) → interleave pairs of 8 - w = weight.view(E, N // 16, 16, K_half) - # Interleave: [gate_0..7, up_0..7, gate_8..15, up_8..15, ...] - # Actually: pairs of 8 rows from gate and up - w_gate = w[:, :, :8, :] - w_up = w[:, :, 8:16, :] - interleaved = torch.stack([w_gate, w_up], dim=3).reshape(E, N, K_half) - - return interleaved.contiguous() - - -def transform_nvfp4_weights_for_tilelang( - weight: torch.Tensor, # (E, N, K//2) int8, K-major packed FP4 - weight_scale: torch.Tensor, # (E, N, K//16) float8_e4m3fn UE4M3 - weight_scale_2: torch.Tensor, # (E, 1) or (E, 2) float32 global scale - N: int, # output dimension (2*INTER for L1, HIDDEN for L2) - K: int, # input dimension (HIDDEN for L1, INTER for L2) - gran_k: int = 16, # NVFP4 group_size - is_l1: bool = False, # True for gate_up_proj (needs interleaving) - logical_widths: list[int] = None, -) -> tuple[torch.Tensor, torch.Tensor]: - """Transform NVFP4 weights for the TileLang block-scaled GEMM. - - Returns: - transformed_weight: (E, N, K//2) int8, K-major, contiguously interleaved - transformed_sf: (E, N, K//64) uint32, packed UE4M3 with folded global scale - """ - device = weight.device - - # Step 1: Fold global scale into block scales - sf_folded = fold_global_scale(weight_scale, weight_scale_2, logical_widths) - - # Step 2: Clamp and convert back to UE4M3 - sf_clamped = sf_folded.clamp(0.0, 448.0).to(torch.float8_e4m3fn) - - # Step 3: Pack UE4M3 into uint32 - sf_packed = pack_ue4m3_to_uint32(sf_clamped) # (E, N, K//64) - - # Step 4: Interleave L1 weights if needed - if is_l1: - transformed_weight = interleave_l1_weights(weight) - else: - transformed_weight = weight.contiguous() - - # Step 5: Ensure K-major layout (already K-major from checkpoint) - # The weight should be (E, N, K//2) with stride (N*K//2, K//2, 1) for K-major - - return transformed_weight, sf_packed - - -def main(): - """Test the weight transformation with dummy data.""" - E = 4 # small test - N = 64 - K = 128 - - device = "cpu" # No GPU on sandbox - weight = torch.randint(-128, 127, (E, N, K // 2), dtype=torch.int8, device=device) - weight_scale = torch.randn(E, N, K // 16, device=device).abs().to(torch.float8_e4m3fn) - weight_scale_2 = torch.ones(E, 2, dtype=torch.float32, device=device) - - transformed_w, transformed_sf = transform_nvfp4_weights_for_tilelang( - weight, weight_scale, weight_scale_2, N, K, is_l1=True, - logical_widths=[32, 32], - ) - - print(f"Weight: {transformed_w.shape} {transformed_w.dtype}") - print(f"SF: {transformed_sf.shape} {transformed_sf.dtype}") - print(f"Expected SF: (E={E}, N={N}, K//64={K//64})") - print("Weight transformation test passed!") - - -if __name__ == "__main__": - main() diff --git a/vllm/patches/deepseek_v4.py b/vllm/patches/deepseek_v4.py new file mode 100644 index 00000000..93692955 --- /dev/null +++ b/vllm/patches/deepseek_v4.py @@ -0,0 +1,2304 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import typing +from collections.abc import Callable, Iterable +from itertools import islice + +import regex as re +import os +import torch +import torch.nn as nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig, get_current_vllm_config +from vllm.distributed import ( + get_ep_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp +from vllm.model_executor.layers.deepseek_v4_attention import ( + DeepseekV4Indexer, + DeepseekV4MLAModules, + DeepseekV4MultiHeadLatentAttentionWrapper, +) +from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear +from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod +from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( + fused_topk_bias, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import ( + QuantizationConfig, + QuantizationMethods, +) +from vllm.model_executor.layers.quantization.fp8 import Fp8Config +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + is_layer_skipped, +) +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.utils import set_weight_attrs +from vllm.platforms import current_platform +from vllm.sequence import IntermediateTensors +from vllm.triton_utils import tl, triton +from vllm.utils.torch_utils import direct_register_custom_op + +from .utils import ( + AutoWeightsLoader, + WeightsMapper, + extract_layer_index, + make_layers, + maybe_prefix, +) + +_DEEPSEEK_V4_EXPERT_DTYPES = ("fp4", "fp8") + + +class DeepseekV4MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + swiglu_limit: float | None = None, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + is_sequence_parallel: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + + # If is_sequence_parallel, the input and output tensors are sharded + # across the ranks within the tp_group. In this case the weights are + # replicated and no collective ops are needed. + # Otherwise we use standard TP with an allreduce at the end. + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + disable_tp=is_sequence_parallel, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + if swiglu_limit is not None: + self.act_fn = SiluAndMulWithClamp(swiglu_limit) + else: + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class DeepseekV4FP8Config(Fp8Config): + """FP8 config for DeepSeek V4 with expert-dtype-aware MoE dispatch. + + DeepSeek V4 checkpoints always use FP8 block quantization for + linear/attention layers. The MoE expert weights vary by checkpoint: + - ``expert_dtype="fp4"`` (e.g. DeepSeek-V4-Flash): MXFP4 experts + with ue8m0 (e8m0fnu) FP8 linear scales. + - ``expert_dtype="fp8"`` (e.g. DeepSeek-V4-Flash-Base): FP8 block + experts with float32 FP8 linear scales. + + The dispatch and the linear scale dtype are both keyed off + ``expert_dtype`` from the model's hf_config; missing values default + to ``"fp4"`` so existing FP4 checkpoints stay unchanged. + + NOTE: ``expert_dtype`` is resolved lazily because this config is + constructed during VllmConfig setup, before ``set_current_vllm_config`` + is active. Reading hf_config eagerly in ``__init__`` would always see + the default ``"fp4"`` and silently misroute Flash-Base checkpoints. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._resolved_expert_dtype: str | None = None + # ``is_scale_e8m0`` is a property that resolves on first read, + # by which time the current vllm_config has been set. + + @property + def expert_dtype(self) -> str: + if self._resolved_expert_dtype is None: + try: + hf_config = get_current_vllm_config().model_config.hf_config + except Exception: + # vllm_config not yet set; return safe default but do NOT + # cache — a later call inside set_current_vllm_config may + # resolve differently. + return "fp4" + expert_dtype = getattr(hf_config, "expert_dtype", "fp4") + if expert_dtype not in _DEEPSEEK_V4_EXPERT_DTYPES: + raise ValueError( + f"Unsupported DeepSeek V4 expert_dtype={expert_dtype!r}; " + f"expected one of {_DEEPSEEK_V4_EXPERT_DTYPES}." + ) + self._resolved_expert_dtype = expert_dtype + return self._resolved_expert_dtype + + @property + def is_scale_e8m0(self) -> bool: + # FP4 checkpoints store FP8 linear scales as e8m0fnu; FP8 expert + # checkpoints (Flash-Base) store them as float32. + return self.expert_dtype == "fp4" + + @classmethod + def get_name(cls) -> QuantizationMethods: + return "deepseek_v4_fp8" + + @classmethod + def override_quantization_method( + cls, hf_quant_cfg, user_quant, hf_config=None + ) -> QuantizationMethods | None: + if not ( + isinstance(hf_quant_cfg, dict) + and hf_quant_cfg.get("quant_method") in ("fp8", "deepseek_v4_fp8") + ): + return None + model_type = getattr(hf_config, "model_type", None) + if model_type == "deepseek_v4" or user_quant == "deepseek_v4_fp8": + return "deepseek_v4_fp8" + return None + + def get_quant_method(self, layer, prefix): + if isinstance(layer, FusedMoE): + if is_layer_skipped( + prefix=prefix, + ignored_layers=self.ignored_layers, + fused_mapping=self.packed_modules_mapping, + ): + return UnquantizedFusedMoEMethod(layer.moe_config) + if self.expert_dtype == "fp4": + return Mxfp4MoEMethod(layer.moe_config) + # expert_dtype == "fp8": fall through to Fp8Config which + # returns Fp8MoEMethod with block-wise float32 scales. + return super().get_quant_method(layer, prefix) + + def is_mxfp4_quant(self, prefix, layer): + return isinstance(layer, FusedMoE) and self.expert_dtype == "fp4" + + + + +def make_deepseek_v4_expert_params_mapping( + num_experts: int, +) -> list[tuple[str, str, int, str]]: + # Checkpoint uses gate_proj/up_proj/down_proj, model params use w13_/w2_ + return [ + ( + "experts.w13_" if shard_id in ("w1", "w3") else "experts.w2_", + f"experts.{expert_id}.{ckpt_name}.", + expert_id, + shard_id, + ) + for expert_id in range(num_experts) + for shard_id, ckpt_name in [ + ("w1", "gate_proj"), + ("w2", "down_proj"), + ("w3", "up_proj"), + ] + ] + + +class DeepseekV4MegaMoEExperts(nn.Module): + """MegaMoE experts for DeepSeek V4 with NVFP4 quantization. + + Loads NVFP4 expert weights (E2M1 packed uint8 + float8_e4m3fn block scales + + float32 global scales) and feeds them natively to the DeepGEMM + fp8_nvfp4_mega_moe kernel (kind::mxf4nvf4.scale_vec::4X). + + No conversion to MXFP4. Experts stay NVFP4. The global scale (weight_scale_2) + is folded into the block scales before kernel consumption. + """ + _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} + + # NVFP4 E2M1 lookup table (positive values, sign from bit 3) + E2M1_LUT = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + # MXFP4 E2M1 is the same format + + def __init__( + self, + vllm_config: VllmConfig, + *, + num_experts: int, + num_local_experts: int, + experts_start_idx: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + prefix: str = "", + ): + super().__init__() + self.prefix = prefix + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.experts_start_idx = experts_start_idx + self.experts_end_idx = experts_start_idx + num_local_experts + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + + weight_attrs = {"weight_loader": self.weight_loader} + + # NVFP4 weights: E2M1 packed as uint8, 2 values per byte + self.w13_weight = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight, weight_attrs) + + # NVFP4 block scales: float8_e4m3fn, group_size=16 + # Shape: [num_local_experts, 2*intermediate_size, hidden_size // 16] + self.w13_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // 16, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight_scale, weight_attrs) + self.w13_weight_scale.quant_method = "block" + + # NVFP4 global scales: float32, per-expert, per-projection (gate, up) + # shape (num_local_experts, 2) — one scale for gate_proj, one for up_proj + self.w13_weight_scale_2 = nn.Parameter( + torch.zeros(num_local_experts, 2, dtype=torch.float32), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight_scale_2, weight_attrs) + + # NVFP4 activation scales: float32, per-expert + self.w13_input_scale = nn.Parameter( + torch.zeros(num_local_experts, dtype=torch.float32), + requires_grad=False, + ) + set_weight_attrs(self.w13_input_scale, weight_attrs) + + self.w2_weight = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // 2, + dtype=torch.int8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight, weight_attrs) + + # NVFP4 block scales for w2 + self.w2_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // 16, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight_scale, weight_attrs) + self.w2_weight_scale.quant_method = "block" + + self.w2_weight_scale_2 = nn.Parameter( + torch.zeros(num_local_experts, dtype=torch.float32), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight_scale_2, weight_attrs) + + self.w2_input_scale = nn.Parameter( + torch.zeros(num_local_experts, dtype=torch.float32), + requires_grad=False, + ) + set_weight_attrs(self.w2_input_scale, weight_attrs) + + self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None + self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None + + # Register in the static forward context so the custom-op wrapper + # can look up this module by name from within a torch.compile graph. + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _map_global_expert_id(self, expert_id: int) -> int: + if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx: + return -1 + return expert_id - self.experts_start_idx + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + weight_name: str, + shard_id: str, + expert_id: int, + ) -> bool: + local_expert_id = self._map_global_expert_id(expert_id) + if local_expert_id == -1: + return False + + # DEBUG: log weight loads for expert params (weight only, not scales) + if int(os.environ.get('MEGA_MOE_DEBUG', '0')) and shard_id in ("w1", "w3") and local_expert_id < 2 and loaded_weight.dtype in (torch.uint8, torch.int8): + print(f"[WT-LOAD] {weight_name} expert={expert_id}→local={local_expert_id} " + f"shard={shard_id} loaded_shape={tuple(loaded_weight.shape)} " + f"param_shape={tuple(param.data[local_expert_id].shape)} " + f"loaded_absmax={loaded_weight.view(torch.int8).abs().max().item()}") + + # Scalar params (weight_scale_2, input_scale): per-expert + if "weight_scale_2" in weight_name or "input_scale" in weight_name: + if "w13_" in weight_name and "weight_scale_2" in weight_name: + # w13 is fused gate+up — store gate and up scales separately + # shard_id tells us which projection: w1=gate, w3=up + proj_idx = 0 if shard_id == "w1" else 1 + param.data[local_expert_id, proj_idx].copy_(loaded_weight) + else: + # w2 or input_scale — single scalar per expert + param.data[local_expert_id].copy_(loaded_weight) + return True + + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + if "w13_" not in weight_name: + return False + shard_offset = 0 if shard_id == "w1" else self.intermediate_size + expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size) + elif shard_id == "w2": + if "w2_" not in weight_name: + return False + else: + raise ValueError(f"Unsupported expert shard id: {shard_id}") + + if expert_data.shape != loaded_weight.shape: + raise ValueError( + f"DeepSeek V4 MegaMoE expert weight shape mismatch for " + f"{weight_name}: parameter shard {tuple(expert_data.shape)} " + f"vs checkpoint {tuple(loaded_weight.shape)}" + ) + expert_data.copy_(loaded_weight) + return True + + def _check_runtime_supported(self) -> None: + if not torch.cuda.is_available(): + raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.") + device = self.w13_weight.device + if device.type != "cuda": + raise NotImplementedError( + "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA." + ) + if torch.cuda.get_device_capability(device)[0] < 10: + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + + def finalize_weights(self) -> None: + if self._transformed_l1_weights is not None: + return + + self._check_runtime_supported() + from nvfp4_megamoe_kernel import transform_nvfp4_weights_for_mega_moe + + # === Native NVFP4 path === + # The DeepGEMM nvfp4 mega_moe kernel consumes NVFP4 directly: + # - E2M1 packed uint8 (same as checkpoint) + # - UE4M3 block scales (float8_e4m3fn), group_size=16 + # - float32 global scale folded into block scales + # No conversion to MXFP4. Experts stay NVFP4. + + # Fold global scales into block scales and transform for the kernel + self._transformed_l1_weights, self._transformed_l2_weights = ( + transform_nvfp4_weights_for_mega_moe( + (self.w13_weight.data.contiguous(), + self.w13_weight_scale.data.contiguous()), + (self.w2_weight.data.contiguous(), + self.w2_weight_scale.data.contiguous()), + l1_weight_scale_2=self.w13_weight_scale_2.data.contiguous(), + l2_weight_scale_2=self.w2_weight_scale_2.data.contiguous(), + ) + ) + + # Drop the original loader-side parameters + self.w13_weight = None + self.w13_weight_scale = None + self.w13_weight_scale_2 = None + self.w13_input_scale = None + self.w2_weight = None + self.w2_weight_scale = None + self.w2_weight_scale_2 = None + self.w2_input_scale = None + + @staticmethod + def _ue8m0_to_float32(sf: torch.Tensor) -> torch.Tensor: + """Convert NVFP4 block scales (float8_e4m3fn / UE4M3) to float32. + + Checkpoint stores float8_e4m3fn (standard NVFP4 spec, NOT UE8M0). + Simple .to(float32) is correct — shift-by-23 was wrong (Bug #7 fix). + """ + return sf.to(torch.float32) + + + + def get_symm_buffer(self): + import nvfp4_megamoe_kernel as deep_gemm + from nvfp4_megamoe_kernel import SymmBuffer, get_symm_buffer_for_nvfp4_mega_moe + + group = get_ep_group().device_group + device = torch.accelerator.current_device_index() + key = ( + id(group), + device, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + symm_buffer = self._symm_buffer_cache.get(key) + if symm_buffer is None: + # NVFP4 SymmBuffer: 2x SF size due to group_size=16 + symm_buffer = get_symm_buffer_for_nvfp4_mega_moe( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + self._symm_buffer_cache[key] = symm_buffer + return symm_buffer + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation_clamp: float | None, + fast_math: bool = True, + ) -> torch.Tensor: + if hidden_states.shape[0] > self.max_num_tokens: + raise ValueError( + f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " + f"but the symmetric buffer was sized for {self.max_num_tokens}." + ) + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + torch.ops.vllm.deepseek_v4_mega_moe_experts( + hidden_states, + topk_weights, + topk_ids, + y, + self.prefix, + activation_clamp, + fast_math, + ) + return y + + def _run_mega_moe( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + y: torch.Tensor, + activation_clamp: float | None, + fast_math: bool, + ) -> None: + import os + from nvfp4_megamoe_kernel import stage_activation, nvfp4_mega_moe_full + + symm_buffer = self.get_symm_buffer() + symm_buffer.experts_start_idx = self.experts_start_idx + num_tokens = hidden_states.shape[0] + + # Quantize activation using the kernel's PyTorch stage_activation + # (same code path the kernel uses for L1→L2 requantization). + x_fp4, x_sf, input_global_scale = stage_activation(hidden_states) + symm_buffer.x[:num_tokens].copy_(x_fp4) + symm_buffer.x_sf[:num_tokens].copy_(x_sf) + symm_buffer.input_global_scale = input_global_scale + symm_buffer.topk_idx[:num_tokens].copy_(topk_ids) + symm_buffer.topk_weights[:num_tokens].copy_(topk_weights) + + # This method must have been already called during the weight loading phase. + # We call it again here to cover the dummy weight loading case. + self.finalize_weights() + + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + + nvfp4_mega_moe_full( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + fast_math=fast_math, + ) + if os.environ.get('NVFP4_DEBUG_SYNC', '') == '1': + torch.cuda.synchronize() + + +DeepseekV4MegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined] + + +def _deepseek_v4_mega_moe_experts_op( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + out: torch.Tensor, + layer_name: str, + activation_clamp: float | None, + fast_math: bool, +) -> None: + self = get_forward_context().no_compile_layers[layer_name] + self._run_mega_moe( + hidden_states, + topk_weights, + topk_ids, + out, + activation_clamp, + fast_math, + ) + + +def _deepseek_v4_mega_moe_experts_op_fake( + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + out: torch.Tensor, + layer_name: str, + activation_clamp: float | None, + fast_math: bool, +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_mega_moe_experts", + op_func=_deepseek_v4_mega_moe_experts_op, + mutates_args=["out"], + fake_impl=_deepseek_v4_mega_moe_experts_op_fake, +) + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + self.tp_size = get_tensor_model_parallel_world_size() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.prefix = prefix + self.use_mega_moe = True # Force mega_moe for NVFP4 pipeline + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." + ) + + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.hidden_size = config.hidden_size + + self.n_routed_experts = config.n_routed_experts + self.n_activated_experts = config.num_experts_per_tok + self.moe_intermediate_size = config.moe_intermediate_size + self.swiglu_limit = config.swiglu_limit + self.renormalize = config.norm_topk_prob + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") + if self.use_mega_moe and self.scoring_func != "sqrtsoftplus": + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently supports sqrtsoftplus routing only." + ) + # NVFP4 experts work with mega_moe via NVFP4 weight transformation in finalize_weights + + self.gate = GateLinear( + config.hidden_size, + config.n_routed_experts, + out_dtype=torch.float32, + bias=False, + prefix=f"{prefix}.gate", + ) + self.gate.e_score_correction_bias = None + self.gate.tid2eid = None + is_hash_moe = extract_layer_index(prefix) < config.num_hash_layers + self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 + + if is_hash_moe: + # hash MoE doesn't use e_score_correction_bias + # Use randint instead of empty to avoid garbage values causing + # invalid memory access in dummy mode (--load-format="dummy") + self.gate.tid2eid = nn.Parameter( + torch.randint( + 0, + config.n_routed_experts, + (config.vocab_size, config.num_experts_per_tok), + dtype=self.hash_indices_dtype, + ), + requires_grad=False, + ) + elif getattr(config, "topk_method", None) == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + + if config.n_shared_experts is None: + self.shared_experts = None + else: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + + self.shared_experts = DeepseekV4MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + swiglu_limit=self.swiglu_limit, + quant_config=quant_config, + reduce_results=self.use_mega_moe, + prefix=f"{prefix}.shared_experts", + ) + + if self.use_mega_moe: + self._init_mega_moe_experts(vllm_config, config, prefix) + else: + self._init_fused_moe_experts(config, quant_config, prefix) + + def _init_mega_moe_experts( + self, + vllm_config: VllmConfig, + config, + prefix: str, + ) -> None: + self.ep_group = get_ep_group() + self.ep_size = self.ep_group.world_size + self.ep_rank = self.ep_group.rank_in_group + assert config.n_routed_experts % self.ep_size == 0 + + self.n_local_experts = config.n_routed_experts // self.ep_size + self.experts_start_idx = self.ep_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + + self.experts = DeepseekV4MegaMoEExperts( + vllm_config, + num_experts=config.n_routed_experts, + num_local_experts=self.n_local_experts, + experts_start_idx=self.experts_start_idx, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + prefix=f"{prefix}.experts", + ) + + def _init_fused_moe_experts( + self, + config, + quant_config, + prefix: str, + ) -> None: + self.tp_rank = get_tensor_model_parallel_rank() + assert config.n_routed_experts % self.tp_size == 0 + + self.n_local_experts = config.n_routed_experts // self.tp_size + self.experts_start_idx = self.tp_rank * self.n_local_experts + self.experts_end_idx = self.experts_start_idx + self.n_local_experts + + self.experts = FusedMoE( + shared_experts=self.shared_experts, + gate=self.gate, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + scoring_func=self.scoring_func, + routed_scaling_factor=self.routed_scaling_factor, + e_score_correction_bias=self.gate.e_score_correction_bias, + hash_indices_table=self.gate.tid2eid, + swiglu_limit=self.swiglu_limit, + router_logits_dtype=torch.float32, + ) + + def forward( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + if self.gate.tid2eid is not None and input_ids is None: + raise ValueError("DeepSeek V4 hash MoE routing requires input_ids.") + + if not self.use_mega_moe: + return self._forward_fused_moe(hidden_states, input_ids) + + org_shape = hidden_states.shape + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids = fused_topk_bias( + hidden_states=hidden_states, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None, + topk=self.n_activated_experts, + renormalize=self.renormalize, + indices_type=self.hash_indices_dtype, + input_tokens=input_ids, + hash_indices_table=self.gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + ) + activation_clamp = ( + float(self.swiglu_limit) if self.swiglu_limit is not None else None + ) + final_hidden_states = self.experts( + hidden_states, + topk_weights, + topk_ids, + activation_clamp=activation_clamp, + ) + + # EP all-reduce: each rank only computes its local experts, + # so we must sum across EP ranks to get the full routed output. + torch.distributed.all_reduce( + final_hidden_states, group=self.ep_group.device_group + ) + + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + final_hidden_states += shared_output + + return final_hidden_states.view(org_shape) + + def _forward_fused_moe( + self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None = None + ) -> torch.Tensor: + org_shape = hidden_states.shape + if self.experts.is_internal_router: + # In this case, the gate/router runs inside the FusedMoE class + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=hidden_states, + input_ids=input_ids, + ) + else: + router_logits, _ = self.gate(hidden_states) + final_hidden_states = self.experts( + hidden_states=hidden_states, + router_logits=router_logits, + input_ids=input_ids, + ) + + return final_hidden_states.view(org_shape) + + def finalize_mega_moe_weights(self) -> None: + if self.use_mega_moe: + self.experts.finalize_weights() + + +class DeepseekV4Attention(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + layer_id = extract_layer_index(prefix) + + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.n_heads = config.num_attention_heads + tp_size = get_tensor_model_parallel_world_size() + assert self.n_heads % tp_size == 0 + + self.n_local_heads = self.n_heads // tp_size + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.head_dim = config.head_dim + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = self.head_dim - self.rope_head_dim + self.n_groups = config.o_groups + self.n_local_groups = self.n_groups // tp_size + self.window_size = config.sliding_window + # NOTE(zyongye) Compress ratio can't be 0 + # we do this for because MTP layer is not included + # in the compress ratio list + if layer_id < config.num_hidden_layers: + self.compress_ratio = max(1, config.compress_ratios[layer_id]) + else: + self.compress_ratio = 1 + self.eps = config.rms_norm_eps + self.max_position_embeddings = config.max_position_embeddings + + # Padded to min 64 heads for FlashMLA, initialized to -inf + # (no sink effect). Weight loading fills the first n_local_heads slots. + padded_heads = max(self.n_local_heads, 64) + self.attn_sink = nn.Parameter( + torch.full((padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + + self.fused_wqa_wkv = MergedColumnParallelLinear( + self.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.fused_wqa_wkv", + disable_tp=True, # fused ReplicatedLinear + ) + self.q_norm = RMSNorm(self.q_lora_rank, self.eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.n_heads * self.head_dim, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wq_b", + ) + + self.kv_norm = RMSNorm(self.head_dim, self.eps) + self.wo_a = ColumnParallelLinear( + self.n_heads * self.head_dim // self.n_groups, + self.n_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_a", + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.n_local_groups + self.wo_b = RowParallelLinear( + self.n_groups * self.o_lora_rank, + self.hidden_size, + bias=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.wo_b", + ) + self.softmax_scale = self.head_dim**-0.5 + self.scale_fmt = config.quantization_config["scale_fmt"] + + self.rope_parameters = config.rope_scaling + + # Initialize rotary embedding BEFORE DeepseekV4MLAModules (which needs it) + rope_parameters = dict(config.rope_parameters) + rope_parameters["rope_theta"] = ( + config.compress_rope_theta if self.compress_ratio > 1 else config.rope_theta + ) + if config.rope_parameters["rope_type"] != "default": + config.rope_parameters["rope_type"] = ( + "deepseek_yarn" + if config.rope_parameters.get("apply_yarn_scaling", True) + else "deepseek_llama_scaling" + ) + rope_parameters["mscale"] = 0 # Disable mscale + rope_parameters["mscale_all_dim"] = 0 # Disable mscale + rope_parameters["is_deepseek_v4"] = True + rope_parameters["rope_dim"] = self.rope_head_dim + self.rotary_emb = get_rope( + self.head_dim, + max_position=self.max_position_embeddings, + rope_parameters=rope_parameters, + is_neox_style=False, + ) + + self.indexer = None + if self.compress_ratio == 4: + # Only C4A uses sparse attention and hence has indexer. + self.indexer = DeepseekV4Indexer( + vllm_config, + config=config, + hidden_size=self.hidden_size, + q_lora_rank=self.q_lora_rank, + quant_config=quant_config, + cache_config=vllm_config.cache_config, + topk_indices_buffer=topk_indices_buffer, + compress_ratio=self.compress_ratio, + prefix=f"{prefix}.indexer", + ) + + mla_modules = DeepseekV4MLAModules( + vllm_config=vllm_config, + fused_wqa_wkv=self.fused_wqa_wkv, + q_norm=self.q_norm, + wq_b=self.wq_b, + kv_norm=self.kv_norm, + wo_a=self.wo_a, + wo_b=self.wo_b, + attn_sink=self.attn_sink, + rotary_emb=self.rotary_emb, + indexer=self.indexer, + indexer_rotary_emb=self.rotary_emb, + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.mla_attn = DeepseekV4MultiHeadLatentAttentionWrapper( + hidden_size=self.hidden_size, + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.softmax_scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + v_head_dim=self.head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.head_dim, + o_lora_rank=self.o_lora_rank, + mla_modules=mla_modules, + window_size=self.window_size, + compress_ratio=self.compress_ratio, + cache_config=vllm_config.cache_config, + quant_config=quant_config, + prefix=prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ): + return self.mla_attn(positions, hidden_states, llama_4_scaling) + + +class DeepseekV4DecoderLayer(nn.Module): + def __init__( + self, + vllm_config, + prefix, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream_list: list[torch.cuda.Stream] | None = None, + ): + super().__init__() + + # Lazy import to avoid top-level tilelang dependency. + # Registers both torch.ops.vllm.mhc_pre and mhc_post + import vllm.model_executor.layers.mhc # noqa: F401 + + config = vllm_config.model_config.hf_config + self.hidden_size = config.hidden_size + + self.rms_norm_eps = config.rms_norm_eps + self.attn = DeepseekV4Attention( + vllm_config, + prefix=f"{prefix}.attn", + topk_indices_buffer=topk_indices_buffer, + aux_stream_list=aux_stream_list, + ) + self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") + + self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.hc_post_alpha = 2.0 + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * self.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty( + (mix_hc, hc_dim), + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_base = nn.Parameter( + torch.empty( + mix_hc, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_attn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty( + 3, + dtype=torch.float32, + ), + requires_grad=False, + ) + + def hc_pre( + self, + x: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + ): + post_mix, res_mix, layer_input = torch.ops.vllm.mhc_pre( + residual=x, + fn=hc_fn, + hc_scale=hc_scale, + hc_base=hc_base, + rms_eps=self.rms_norm_eps, + hc_pre_eps=self.hc_eps, + hc_sinkhorn_eps=self.hc_eps, + hc_post_mult_value=self.hc_post_alpha, + sinkhorn_repeat=self.hc_sinkhorn_iters, + ) + return layer_input, post_mix, res_mix + + def hc_post( + self, + x: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, + ): + return torch.ops.vllm.mhc_post(x, residual, post, comb) + + def forward( + self, + x: torch.Tensor, + positions: torch.Tensor, + input_ids: torch.Tensor | None, + ) -> torch.Tensor: + # DEBUG: skip attention entirely, just run FFN on raw input + if int(os.environ.get('SKIP_ATTENTION', '0')): + # Flatten to 2D for ffn, then restore + org_shape = x.shape + x_2d = x.view(-1, x.shape[-1]) + x_2d = self.ffn_norm(x_2d) + x_2d = self.ffn(x_2d, input_ids) + return x_2d.view(org_shape) + + residual = x + x, post, comb = self.hc_pre( + x, self.hc_attn_fn, self.hc_attn_scale, self.hc_attn_base + ) + x = self.attn_norm(x) + x = self.attn(positions, x, None) + x = self.hc_post(x, residual, post, comb) + + residual = x + x, post, comb = self.hc_pre( + x, self.hc_ffn_fn, self.hc_ffn_scale, self.hc_ffn_base + ) + x = self.ffn_norm(x) + x = self.ffn(x, input_ids) + x = self.hc_post(x, residual, post, comb) + return x + + +@support_torch_compile +class DeepseekV4Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.use_mega_moe = True # Force mega_moe for NVFP4 pipeline + if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel: + raise NotImplementedError( + "DeepSeek V4 MegaMoE currently requires expert parallel. " + "Enable it with --enable-expert-parallel, or pick a different " + "moe backend." + ) + self.vocab_size = config.vocab_size + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + self.hc_dim = self.hc_mult * config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + + # Three aux streams: one per non-default input GEMM in + # DeepseekV4MultiHeadLatentAttentionWrapper.attn_gemm_parallel_execute + # (compressor kv_score, indexer.weights_proj, indexer.compressor + # kv_score). fused_wqa_wkv stays on the default stream. + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + + self.device = current_platform.device_type + # Reserved topk indices buffer for all Indexer layers to reuse. + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: DeepseekV4DecoderLayer( + vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + aux_stream_list=aux_stream_list, + ), + prefix=f"{prefix}.layers", + ) + + self.norm = RMSNorm(config.hidden_size, self.rms_norm_eps) + + self.hc_head_fn = nn.Parameter( + torch.empty( + self.hc_mult, + self.hc_dim, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty( + self.hc_mult, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + + # Pre-hc_head residual stream buffer for the MTP draft. Stable + # address (outside the cudagraph pool) so the copy_ in forward() + # refreshes it correctly across captured shapes. + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + self.hc_dim, + dtype=vllm_config.model_config.dtype, + device=self.device, + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states.unsqueeze(-2).repeat(1, self.hc_mult, 1) + if self.use_mega_moe: + input_ids = input_ids.to(torch.int64) + for layer in islice(self.layers, self.start_layer, self.end_layer): + hidden_states = layer( + hidden_states, + positions, + input_ids, + ) + + # Stash pre-hc_head residual for the MTP draft (captured copy_). + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1)) + + hidden_states = hc_head( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + hidden_states = self.norm(hidden_states) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ("compressor.fused_wkv_wgate", "compressor.wkv", 0), + ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + ] + + # Checkpoint key → model param name substitutions. + # Applied to each (name, weight) pair before matching against + # params_dict. Order matters: longer/more-specific patterns first. + CKPT_KEY_SUBST = { + # self_attn projection names → vLLM attn attribute names + ".self_attn.q_a_proj.": ".attn.wq_a.", + ".self_attn.q_b_proj.": ".attn.wq_b.", + ".self_attn.q_a_norm.": ".attn.q_norm.", + ".self_attn.o_a_proj.": ".attn.wo_a.", + ".self_attn.o_b_proj.": ".attn.wo_b.", + ".self_attn.sinks": ".attn.attn_sink", + ".self_attn.kv_proj.": ".attn.wkv.", + ".self_attn.kv_norm.": ".attn.kv_norm.", + # Indexer: self_attn.compressor.indexer → attn.indexer + # MUST come before the generic .self_attn.compressor. rule + ".self_attn.compressor.indexer.q_b_proj.": ".attn.indexer.wq_b.", + ".self_attn.compressor.indexer.kv_norm.": ".attn.indexer.k_norm.", + ".self_attn.compressor.indexer.position_bias": ".attn.indexer.compressor.ape", + ".self_attn.compressor.indexer.gate_proj.": ".attn.indexer.compressor.wgate.", + ".self_attn.compressor.indexer.kv_proj.": ".attn.indexer.compressor.wkv.", + ".self_attn.compressor.indexer.": ".attn.indexer.", + # Compressor: self_attn.compressor → attn.mla_attn.compressor + # Compressor projections for stacking (fused_wkv_wgate) + ".self_attn.compressor.kv_proj.": ".attn.mla_attn.compressor.wkv.", + ".self_attn.compressor.gate_proj.": ".attn.mla_attn.compressor.gate.", + ".self_attn.compressor.kv_norm.": ".attn.kv_norm.", + ".self_attn.compressor.position_bias": ".attn.mla_attn.compressor.ape", + ".self_attn.compressor.": ".attn.mla_attn.compressor.", + # Shared expert projections (stacking into gate_up_proj) + # Must include .mlp. prefix since break prevents .mlp.→.ffn. from + # firing on the same key after these patterns match. + ".mlp.shared_experts.gate_proj.": ".ffn.shared_experts.w1.", + ".mlp.shared_experts.up_proj.": ".ffn.shared_experts.w3.", + ".mlp.shared_experts.down_proj.": ".ffn.shared_experts.down_proj.", + # Hadamard coding params: checkpoint has .attn_hc.base/fn/scale + # and .ffn_hc.base/fn/scale; model has hc_attn_base/fn/scale + # and hc_ffn_base/fn/scale (underscore not dot before base/fn/scale) + ".attn_hc.base": ".hc_attn_base", + ".attn_hc.fn": ".hc_attn_fn", + ".attn_hc.scale": ".hc_attn_scale", + ".ffn_hc.base": ".hc_ffn_base", + ".ffn_hc.fn": ".hc_ffn_fn", + ".ffn_hc.scale": ".hc_ffn_scale", + "hc_head.hc_base": "hc_head_base", + "hc_head.hc_fn": "hc_head_fn", + "hc_head.hc_scale": "hc_head_scale", + # compressor.position_bias → compressor.ape + ".compressor.position_bias": ".compressor.ape", + # modelopt uses mlp, vllm uses ffn internally + ".mlp.": ".ffn.", + } + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + + # TP for attention + tp_size = get_tensor_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() + n_head = self.config.num_attention_heads + n_local_head = n_head // tp_size + head_rank_start = n_local_head * tp_rank + head_rank_end = n_local_head * (tp_rank + 1) + + # Pre-compute expert mapping ONCE. + expert_mapping = self.get_expert_mapping() + + for name, loaded_weight in weights: + # Strip 'model.' prefix from checkpoint keys. + # vLLM's weight iteration yields keys like 'model.layers.0...' + # but named_parameters() on DeepseekV4Model returns 'layers.0...' + if name.startswith("model."): + name = name[len("model."):] + + # Apply checkpoint → model name substitutions + for ckpt_pat, model_pat in CKPT_KEY_SUBST.items(): + if ckpt_pat in name: + name = name.replace(ckpt_pat, model_pat) + break # first match wins (order matters) + + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip MoE routed experts (handled separately below). + # Use .ffn.experts. (not .experts.) to avoid skipping + # shared_experts which also contains ".experts.". + if ".ffn.experts." in name: + continue + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if name_mapped not in params_dict: + continue + name = name_mapped + + param = params_dict[name] + weight_loader = param.weight_loader + + # ModelOpt NVFP4 packed weight fix for MergedColumnParallelLinear. + # + # modelopt exports NVFP4 packed weights as uint8 (2 values/byte + # along the column dim). But MergedColumnParallelLinear creates + # the weight param as bfloat16 (ModelWeightParameter), because + # ModelOptNvFp4Config only patches Linear, not + # MergedColumnParallelLinear. + # + # When loading uint8 packed weights into a bf16 param, we need to + # unpack them. Each uint8 byte contains 2 E2M1 FP4 values. + # We unpack using the LUT and return bf16. + # + # The weight_scale is loaded separately and process_weights_after_loading + # will handle the actual NVFP4 quantization. + if (loaded_weight.dtype == torch.uint8 + and param.data.dtype != torch.uint8 + and loaded_weight.shape[-1] * 2 == param.data.shape[-1]): + # Unpack NVFP4 (E2M1) → BF16 + # E2M1 LUT: 0→0, 1→0.5, 2→1, 3→1.5, 4→2, 5→3, 6→4, 7→6 + # Sign bit in bit 3 (indices 8-15 are negatives) + FP4_LUT = torch.tensor([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, + ], dtype=torch.float32, device=loaded_weight.device) + lower = FP4_LUT[(loaded_weight & 0x0F).long()] # (..., in_packed, ) + upper = FP4_LUT[((loaded_weight >> 4) & 0x0F).long()] + # Interleave: [lower_0, upper_0, lower_1, upper_1, ...] + out = torch.empty( + *loaded_weight.shape[:-1], loaded_weight.shape[-1] * 2, + dtype=torch.float32, device=loaded_weight.device, + ) + out[..., 0::2] = lower + out[..., 1::2] = upper + loaded_weight = out.to(torch.bfloat16) + + try: + weight_loader(param, loaded_weight, shard_id) + except (AssertionError, ValueError, RuntimeError) as e: + raise RuntimeError( + f'Weight load failed: name={name} shard_id={shard_id} ' + f'param.shape={param.shape} param.dtype={param.data.dtype} ' + f'loaded.shape={loaded_weight.shape} loaded.dtype={loaded_weight.dtype}' + ) from e + loaded_params.add(name) + break + else: + if ".ffn.experts." in name: + # E8M0 scales are stored as float8_e8m0fnu in + # MXFP4 checkpoints but NVFP4 uses float8_e4m3fn. + # The uint8 view+copy path is only valid for MXFP4; + # for NVFP4 it would paste raw E8M0 bytes into an + # E4M3 buffer, producing garbage. + if ( + "weight_scale" in name + and loaded_weight.dtype == torch.float8_e8m0fnu + ): + assert False, ( + f"E8M0 weight_scale encountered for NVFP4 experts " + f"({name}) — this is only valid for MXFP4. " + f"Check checkpoint dtype." + ) + for mapping in expert_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name_mapped = name.replace(weight_name, param_name) + if name_mapped not in params_dict: + continue + param = params_dict[name_mapped] + # We should ask the weight loader to return success or not + # here since otherwise we may skip experts with other + # available replicas. + weight_loader = typing.cast( + Callable[..., bool], param.weight_loader + ) + success = weight_loader( + param, + loaded_weight, + name_mapped, + shard_id=shard_id, + expert_id=expert_id, + ) + if success: + name = name_mapped + loaded_params.add(name_mapped) + break + else: + continue + continue + elif "attn_sink" in name: + if name not in params_dict: + continue + narrow_weight = loaded_weight[head_rank_start:head_rank_end] + n = narrow_weight.shape[0] + params_dict[name][:n].copy_(narrow_weight) + loaded_params.add(name) + continue + else: + if name not in params_dict: + # ModelOpt NVFP4 export includes params not in the + # vllm model (e.g., compressor.position_bias). + # Skip them silently. + continue + param = params_dict[name] + + # Handle bf16 → uint8 mismatch for o_a_proj: + # modelopt didn't quantize o_a_proj (bf16, no scales), + # but ModelOptNvFp4Config creates wo_a with NVFP4 quant + # (uint8 weight + scales). We quantize the bf16 weight + # to NVFP4 at load time so the layer runs in NVFP4 path. + if (name.endswith(".weight") + and loaded_weight.dtype != torch.uint8 + and param.data.dtype == torch.uint8): + # Quantize bf16 → NVFP4 (E2M1 packed uint8 + scales) + w_bf16 = loaded_weight + out_dim, in_dim = w_bf16.shape + block_size = 16 + assert in_dim % block_size == 0 + n_blocks = in_dim // block_size + + # Reshape into blocks + w_blocks = w_bf16.reshape(out_dim, n_blocks, block_size) + + # Compute per-block amax + amax = w_blocks.abs().amax(dim=-1) # [out, n_blocks] + + # Global scale (weight_scale_2): max amax / (6.0 * 448.0) + global_amax = amax.max() + # Use 448.0 as the max e4m3 value for scale computation + weight_scale_2_val = global_amax / (6.0 * 448.0) + weight_scale_2 = weight_scale_2_val.to(torch.float32) + + # Per-block scale (weight_scale): UE4M3 format (standard NVFP4) + # block_scale = amax / (6.0 * weight_scale_2) + block_scale = amax / (6.0 * weight_scale_2_val) + weight_scale = block_scale.clamp(0.0, 448.0).to(torch.float8_e4m3fn) + + # Quantize to FP4 (E2M1) + # E2M1 LUT: 0, 0.5, 1, 1.5, 2, 3, 4, 6 (positive) + FP4_POS = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float32, device=w_bf16.device, + ) + # Scale the weight values: normalized = w / (block_scale * weight_scale_2) + block_scale_f32 = block_scale.clamp(0.0, 448.0) + scaled = w_blocks / (block_scale_f32.unsqueeze(-1) * weight_scale_2_val) + # Find nearest FP4 index (0-7 for magnitude) + # Use absolute value for matching, then apply sign + scaled_abs = scaled.abs() + # Find closest FP4 value + diff = (scaled_abs.unsqueeze(-1) - FP4_POS).abs() + fp4_idx = diff.argmin(dim=-1) # [out, n_blocks, block_size] + # Apply sign: negative values get bit 3 set + sign = (scaled < 0).int() + fp4_val = (sign << 3) | fp4_idx.int() + # Pack: 2 FP4 values per uint8 byte + # Even positions → lower nibble, Odd → upper nibble + fp4_flat = fp4_val.reshape(out_dim, -1) # [out, in_dim] + assert fp4_flat.shape[1] % 2 == 0 + even = fp4_flat[:, 0::2] # lower nibble + odd = fp4_flat[:, 1::2] # upper nibble + packed = (odd << 4) | even + weight_packed = packed.to(torch.uint8).view(torch.int8) + + # Reshape weight_scale to [out, n_blocks] + weight_scale_2d = weight_scale.reshape(out_dim, n_blocks) + + # Load the quantized weight into the uint8 param + weight_loader = param.weight_loader + weight_loader(param, weight_packed) + loaded_params.add(name) + + # Load scales into sibling params + base = name.rsplit(".", 1)[0] + # weight_scale + ws_name = f"{base}.weight_scale" + if ws_name in params_dict: + ws_param = params_dict[ws_name] + ws_loader = getattr(ws_param, "weight_loader", default_weight_loader) + ws_loader(ws_param, weight_scale_2d) + loaded_params.add(ws_name) + # weight_scale_2 + ws2_name = f"{base}.weight_scale_2" + if ws2_name in params_dict: + ws2_param = params_dict[ws2_name] + ws2_loader = getattr(ws2_param, "weight_loader", default_weight_loader) + ws2_loader(ws2_param, weight_scale_2.reshape(1)) + loaded_params.add(ws2_name) + # input_scale: use 1.0 default (dynamic quant) + is_name = f"{base}.input_scale" + if is_name in params_dict: + is_param = params_dict[is_name] + is_loader = getattr(is_param, "weight_loader", default_weight_loader) + is_loader(is_param, torch.tensor(1.0, dtype=torch.float32)) + loaded_params.add(is_name) + continue + + # Handle uint8 NVFP4 packed → bf16 unpack for non-stacked + # params (e.g. indexer.weights_proj). Checkpoint stores + # NVFP4 as uint8 (2 values/byte), but model param is bf16. + if (loaded_weight.dtype == torch.uint8 + and param.data.dtype != torch.uint8 + and loaded_weight.shape[-1] * 2 == param.data.shape[-1]): + FP4_LUT = torch.tensor([ + 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, + -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0, + ], dtype=torch.float32, device=loaded_weight.device) + lower = FP4_LUT[(loaded_weight & 0x0F).long()] + upper = FP4_LUT[((loaded_weight >> 4) & 0x0F).long()] + out = torch.empty( + *loaded_weight.shape[:-1], + loaded_weight.shape[-1] * 2, + dtype=torch.float32, device=loaded_weight.device, + ) + out[..., 0::2] = lower + out[..., 1::2] = upper + loaded_weight = out.to(torch.bfloat16) + + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + first_layer = next(iter(islice(self.layers, self.start_layer, self.end_layer))) + if first_layer.ffn.use_mega_moe: + return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts) + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return FusedMoE.make_expert_params_mapping( + self, + ckpt_gate_proj_name="w1", + ckpt_down_proj_name="w2", + ckpt_up_proj_name="w3", + num_experts=self.config.n_routed_experts, + ) + + def finalize_mega_moe_weights(self) -> None: + for layer in islice(self.layers, self.start_layer, self.end_layer): + layer.ffn.finalize_mega_moe_weights() + + def _convert_nvfp4_post_load(self): + """Post-load conversion of NVFP4 weights for vLLM compatibility. + + Strategy: + - wo_a: Convert to FP8 (attention forward reads weight/weight_scale_inv + directly and passes to deepseek_v4_fp8_einsum, bypassing quant_method) + - fused_wqa_wkv, wq_b, wo_b: Dequant NVFP4->bf16 (called via + .forward() which goes through quant_method; FP8 would dtype-mismatch) + - compressor.fused_wkv_wgate: Dequant NVFP4->bf16 (used via direct + torch.mm in attention parallel stream) + - shared_experts (gate_up_proj, down_proj): Stay native NVFP4 via DeepGEMM fp8_fp4_gemm + - MoE experts: Handled by DeepseekV4MegaMoEExperts (NVFP4 native) + """ + E2M1_LUT = torch.tensor( + [0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16 + ) + FP8_MAX = torch.finfo(torch.float8_e4m3fn).max + + # wo_a: attention forward reads .weight and .weight_scale_inv directly + # for fp8_einsum. Only layer that needs FP8 conversion. + fp8_proj_names = {"wo_a"} + # No BF16 dequant paths active — cuBLAS BF16 is broken on Blackwell. + # wo_a goes NVFP4→FP8; compressor gets reconstructed from checkpoint; + # MoE experts stay native NVFP4 via CUTLASS kernel. + + fp8_converted = 0 + fp8_from_bf16 = 0 + compressor_converted = 0 + + # Build shard index once for compressor reconstruction (avoids N×M full-shard loads) + _shard_index = self._build_shard_index("/model") if os.path.isdir("/model") else None + + for layer_idx, layer in enumerate(self.layers): + attn = layer.attn + + # FP8 conversion: only wo_a + for proj_name in fp8_proj_names: + if not hasattr(attn, proj_name): + continue + mod = getattr(attn, proj_name) + if not hasattr(mod, "weight"): + continue + if mod.weight.dtype in (torch.uint8, torch.int8): + # NVFP4 -> dequant to bf16 -> requant to FP8 + self._convert_nvfp4_to_fp8(mod, E2M1_LUT, FP8_MAX) + fp8_converted += 1 + elif mod.weight.dtype == torch.bfloat16: + self._convert_bf16_to_fp8(mod, FP8_MAX) + fp8_from_bf16 += 1 + + # Compressor: fused_wkv_wgate used via direct torch.mm + # Compressor weights were SKIPPED during loading (skip patterns) + # because the stacking weight_loader corrupts NVFP4 uint8 data. + # We reconstruct the bf16 weight from the individual sub-weights + # that were loaded separately before stacking. + # Note: compressor.kv_proj.weight and compressor.gate_proj.weight + # are skipped, so fused_wkv_wgate.weight is zeros (empty tensor). + # We need to manually create it. + mla_attn = getattr(attn, "mla_attn", None) + if mla_attn is not None: + compressor = getattr(mla_attn, "compressor", None) + if compressor is not None and hasattr(compressor, "fused_wkv_wgate"): + compressor_converted += self._reconstruct_compressor_weight( + compressor.fused_wkv_wgate, attn, layer_idx, E2M1_LUT, _shard_index=_shard_index) + # Indexer compressor (C4A layers only) + indexer = getattr(mla_attn, "indexer", None) + if indexer is not None: + idx_compressor = getattr(indexer, "compressor", None) + if idx_compressor is not None and hasattr(idx_compressor, "fused_wkv_wgate"): + compressor_converted += self._reconstruct_compressor_weight( + idx_compressor.fused_wkv_wgate, indexer, layer_idx, E2M1_LUT, sub_path=".indexer", _shard_index=_shard_index) + + total_fp8 = fp8_converted + fp8_from_bf16 + total_bf16 = compressor_converted + if int(os.environ.get('NVFP4_DEBUG', '0')) and (total_fp8 > 0 or total_bf16 > 0): + print(f"NVFP4 post-load: {fp8_converted} NVFP4->FP8, " + f"{fp8_from_bf16} BF16->FP8, " + f"{compressor_converted} compressor->BF16") + + + def _dequant_nvfp4_to_bf16(self, mod, e2m1_lut): + """Dequantize NVFP4 weight to bf16 for normal .forward() path.""" + w_uint8 = mod.weight.data + device = w_uint8.device + w_bf16 = self._unpack_nvfp4_to_bf16(w_uint8, e2m1_lut, device) + + # Dequantize with scales + if hasattr(mod, "weight_scale") and hasattr(mod, "weight_scale_2"): + # NVFP4 block scales are float8_e4m3fn (UE4M3) — standard spec. + # .to(float32) is correct (Bug #7: shift-by-23 was wrong, reverted) + block_scale = self._ue8m0_to_float32(mod.weight_scale.data) + if block_scale.dim() == 2 and w_bf16.dim() == 2: + block_size = w_bf16.shape[1] // block_scale.shape[1] + block_scale_expanded = block_scale.unsqueeze(-1).expand( + -1, -1, block_size + ).reshape(w_bf16.shape) + else: + block_scale_expanded = block_scale + global_scale = mod.weight_scale_2.data.max().item() + input_scale = ( + mod.input_scale.data.max().item() + if hasattr(mod, "input_scale") + else 1.0 + ) + # NOTE: input_scale is for ACTIVATIONS, not weights. + # Weight dequant = e2m1 * block_scale * global_scale (NO input_scale) + w_dequant = w_bf16.float() * block_scale_expanded * global_scale + w_dequant = w_dequant.to(torch.bfloat16) + else: + w_dequant = w_bf16 + + # Free source tensors eagerly to avoid holding uint8+bf16+fp32 simultaneously + del w_uint8, w_bf16 + mod.weight = torch.nn.Parameter(w_dequant, requires_grad=False) + del w_dequant + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + mod.quant_method = UnquantizedLinearMethod() + for attr in ("weight_scale", "weight_scale_2", "input_scale", + "weight_scale_inv"): + if hasattr(mod, attr): + delattr(mod, attr) + + def _convert_nvfp4_to_fp8(self, mod, e2m1_lut, fp8_max): + """Convert NVFP4 weight to FP8 for fp8_einsum path (wo_a only). + + Uses DeepGEMM's deepgemm_post_process_fp8_weight_block to ensure + correct weight and scale format for fp8_einsum with BMM. + """ + w_uint8 = mod.weight.data + device = w_uint8.device + w_bf16 = self._unpack_nvfp4_to_bf16(w_uint8, e2m1_lut, device) + + # Dequantize with scales + if hasattr(mod, "weight_scale") and hasattr(mod, "weight_scale_2"): + # NVFP4 block scales: float8_e4m3fn → .to(float32) (Bug #7 reverted) + block_scale = self._ue8m0_to_float32(mod.weight_scale.data) + if block_scale.dim() == 2 and w_bf16.dim() == 2: + block_size = w_bf16.shape[1] // block_scale.shape[1] + block_scale_expanded = block_scale.unsqueeze(-1).expand( + -1, -1, block_size + ).reshape(w_bf16.shape) + else: + block_scale_expanded = block_scale + global_scale = mod.weight_scale_2.data.max().item() + input_scale = ( + mod.input_scale.data.max().item() + if hasattr(mod, "input_scale") + else 1.0 + ) + # NOTE: input_scale is for ACTIVATIONS, not weights. + # Weight dequant = e2m1 * block_scale * global_scale (NO input_scale) + w_dequant = w_bf16.float() * block_scale_expanded * global_scale + w_dequant = w_dequant.to(torch.bfloat16) + else: + w_dequant = w_bf16 + + # Re-quantize bf16 -> FP8 e4m3 with block quantization + # DeepGEMM expects block-scale format: weight_scale (FP8 e4m3 block scale) + # and weight_scale_inv (per-tensor scale). + # We do per-tensor quantization, so block_scale is all-ones. + w_amax = w_dequant.abs().amax() + if w_amax == 0: + w_amax = torch.tensor(1.0, device=device) + fp8_scale = w_amax / fp8_max + w_fp8 = (w_dequant / fp8_scale).to(torch.float8_e4m3fn) + + # Create block scale filled with the per-tensor fp8_scale value. + # DeepGEMM divides by the block scale, so each block gets fp8_scale. + BLOCK_SIZE = 128 + is_bmm = getattr(mod, "is_bmm", False) + bmm_batch_size = getattr(mod, "bmm_batch_size", 0) + + # Weight is 2D (output_size, input_size) before BMM reshape + # Block scale shape: (output_size / BLOCK_SIZE, input_size / BLOCK_SIZE) + rows = w_fp8.size(0) + cols = w_fp8.size(1) + block_rows = rows // BLOCK_SIZE + block_cols = cols // BLOCK_SIZE + + # Fill block scale with the per-tensor fp8_scale (NOT all-ones!) + # This is correct because we requantized with a single per-tensor scale, + # so every 128x128 block has the same scale = fp8_scale. + ws = torch.full((block_rows, block_cols), fp8_scale.item(), dtype=torch.float32, device=device) + + # Use DeepGEMM's post-processing for proper layout transformation + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + deepgemm_post_process_fp8_weight_block, + ) + w_fp8, ws = deepgemm_post_process_fp8_weight_block( + wq=w_fp8, + ws=ws, + quant_block_shape=(BLOCK_SIZE, BLOCK_SIZE), + use_e8m0=True, # scale_fmt=ue8m0 + is_bmm=is_bmm, + bmm_batch_size=bmm_batch_size, + ) + + # Free source tensors eagerly + del w_uint8, w_bf16, w_dequant + mod.weight = torch.nn.Parameter(w_fp8, requires_grad=False) + del w_fp8 + # weight_scale_inv is what the attention runtime reads as b_scale + # for deepseek_v4_fp8_einsum -> DeepGEMM fp8_einsum. + # It must be the DeepGEMM-formatted block scale (dg_ws), NOT the + # per-tensor scalar. See: deepseek_v4_attention.py line 319. + mod.weight_scale_inv = torch.nn.Parameter(ws, requires_grad=False) + del ws + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + mod.quant_method = UnquantizedLinearMethod() + for attr in ("weight_scale", "weight_scale_2", "input_scale"): + if hasattr(mod, attr): + delattr(mod, attr) + + @staticmethod + def _build_shard_index(ckpt_dir: str) -> dict[str, str]: + """Build key→shard_path index from safetensors metadata (no tensor I/O).""" + import glob + from safetensors import safe_open + index = {} + for shard_file in sorted(glob.glob(os.path.join(ckpt_dir, "model-*.safetensors"))): + try: + with safe_open(shard_file, framework="pt") as f: + for key in f.keys(): + index[key] = shard_file + except Exception: + continue + return index + + def _reconstruct_compressor_weight(self, fused_mod, parent_mod, layer_idx, e2m1_lut, sub_path="", _shard_index=None): + """Reconstruct compressor fused_wkv_wgate from checkpoint. + + Compressor weights are SKIPPED during loading because NVFP4 uint8 data + can't be loaded into bf16 MergedColumnParallelLinear params (shape mismatch). + We read the original uint8 data from the safetensors checkpoint, unpack + E2M1, dequantize, and stack into the fused weight param. + """ + from safetensors import safe_open + + # Find the checkpoint directory + # The model weights are mounted at /model in Docker + ckpt_dir = "/model" + if not os.path.isdir(ckpt_dir): + print(f"WARNING: layer {layer_idx} compressor: checkpoint dir {ckpt_dir} not found") + return 0 + + # Determine the layer's compressor key prefix in the checkpoint + # Before mapper: model.layers.N.self_attn.compressor.{kv_proj,gate_proj} + # After mapper: model.layers.N.attn.mla_attn.compressor.{wkv,wgate} + # We read from checkpoint (before mapper), so use original names + layer_prefix = f"model.layers.{layer_idx}.self_attn.compressor{sub_path}" + + # All keys we need from the checkpoint + keys = { + 'wkv_uint8': f"{layer_prefix}.kv_proj.weight", + 'wgate_uint8': f"{layer_prefix}.gate_proj.weight", + 'wkv_block_scale': f"{layer_prefix}.kv_proj.weight_scale", + 'wgate_block_scale': f"{layer_prefix}.gate_proj.weight_scale", + 'wkv_global_scale': f"{layer_prefix}.kv_proj.weight_scale_2", + 'wgate_global_scale': f"{layer_prefix}.gate_proj.weight_scale_2", + 'wkv_input_scale': f"{layer_prefix}.kv_proj.input_scale", + 'wgate_input_scale': f"{layer_prefix}.gate_proj.input_scale", + } + + # Read tensors using shard index for targeted access (no full-shard loads) + tensors = {} + for name, key in keys.items(): + shard_path = (_shard_index or {}).get(key) + if shard_path is None: + continue + try: + with safe_open(shard_path, framework="pt") as f: + if key in f.keys(): + tensors[name] = f.get_tensor(key) + except Exception: + continue + + wkv_uint8 = tensors.get('wkv_uint8') + wgate_uint8 = tensors.get('wgate_uint8') + + if wkv_uint8 is None or wgate_uint8 is None: + # Layer might not have a compressor (compress_ratio=1 layers) + return 0 + + wkv_block_scale = tensors.get('wkv_block_scale') + wgate_block_scale = tensors.get('wgate_block_scale') + wkv_global_scale = tensors.get('wkv_global_scale') + wgate_global_scale = tensors.get('wgate_global_scale') + wkv_input_scale = tensors.get('wkv_input_scale') + wgate_input_scale = tensors.get('wgate_input_scale') + + device = fused_mod.weight.device + wkv_uint8 = wkv_uint8.to(device) + wgate_uint8 = wgate_uint8.to(device) + + # Unpack E2M1 FP4→bf16 + wkv_bf16 = self._unpack_nvfp4_to_bf16(wkv_uint8, e2m1_lut, device) + wgate_bf16 = self._unpack_nvfp4_to_bf16(wgate_uint8, e2m1_lut, device) + + # Dequantize with scales + def _dequant(w_bf16, block_scale, global_scale, input_scale): + if block_scale is not None and global_scale is not None: + # NVFP4 block scales: float8_e4m3fn → .to(float32) (Bug #7 reverted) + block_scale = self._ue8m0_to_float32(block_scale.to(device)) + if block_scale.dim() == 2 and w_bf16.dim() == 2: + block_size = w_bf16.shape[1] // block_scale.shape[1] + block_scale_exp = block_scale.unsqueeze(-1).expand( + -1, -1, block_size + ).reshape(w_bf16.shape) + else: + block_scale_exp = block_scale + gs = global_scale.to(device).max().item() + # NOTE: input_scale is for activations, not weights. + # Weight dequant = e2m1 * block_scale * global_scale (NO input_scale) + w = w_bf16.float() * block_scale_exp * gs + return w.to(torch.bfloat16) + return w_bf16 + + wkv_dequant = _dequant(wkv_bf16, wkv_block_scale, wkv_global_scale, wkv_input_scale) + wgate_dequant = _dequant(wgate_bf16, wgate_block_scale, wgate_global_scale, wgate_input_scale) + + # Stack: concatenate along output dim (dim 0) + # fused_wkv_wgate.weight = cat([wkv, wgate], dim=0) → (2*head_dim, hidden_size) + w_fused = torch.cat([wkv_dequant, wgate_dequant], dim=0) + + + # Replace the weight + fused_mod.weight = torch.nn.Parameter(w_fused, requires_grad=False) + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + fused_mod.quant_method = UnquantizedLinearMethod() + for attr in ("weight_scale", "weight_scale_2", "input_scale", "weight_scale_inv"): + if hasattr(fused_mod, attr): + delattr(fused_mod, attr) + return 1 + + def _convert_bf16_to_fp8(self, mod, fp8_max): + """Convert BF16 weight to FP8 for fp8_einsum path. + + Used for wo_a which modelopt did NOT quantize (bf16 in checkpoint) + but which the attention forward reads as FP8 for deepseek_v4_fp8_einsum. + Uses DeepGEMM's post-processing for proper BMM + scale format. + """ + w_bf16 = mod.weight.data + device = w_bf16.device + + # Re-quantize bf16 -> FP8 e4m3 with block quantization + w_amax = w_bf16.abs().amax() + if w_amax == 0: + w_amax = torch.tensor(1.0, device=device) + fp8_scale = w_amax / fp8_max + w_fp8 = (w_bf16 / fp8_scale).to(torch.float8_e4m3fn) + + BLOCK_SIZE = 128 + is_bmm = getattr(mod, "is_bmm", False) + bmm_batch_size = getattr(mod, "bmm_batch_size", 0) + + rows = w_fp8.size(0) + cols = w_fp8.size(1) + block_rows = rows // BLOCK_SIZE + block_cols = cols // BLOCK_SIZE + # Fill block scale with per-tensor fp8_scale (NOT all-ones!) + ws = torch.full((block_rows, block_cols), fp8_scale.item(), dtype=torch.float32, device=device) + + from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + deepgemm_post_process_fp8_weight_block, + ) + w_fp8, ws = deepgemm_post_process_fp8_weight_block( + wq=w_fp8, + ws=ws, + quant_block_shape=(BLOCK_SIZE, BLOCK_SIZE), + use_e8m0=True, # scale_fmt=ue8m0 + is_bmm=is_bmm, + bmm_batch_size=bmm_batch_size, + ) + + mod.weight = torch.nn.Parameter(w_fp8, requires_grad=False) + # weight_scale_inv is what the attention runtime reads as b_scale + # for deepseek_v4_fp8_einsum -> DeepGEMM fp8_einsum. + # It must be the DeepGEMM-formatted block scale (dg_ws), NOT the + # per-tensor scalar. See: deepseek_v4_attention.py line 319. + mod.weight_scale_inv = torch.nn.Parameter(ws, requires_grad=False) + # weight_scale is not used at runtime for BMM layers; remove it + # to avoid confusing other code paths. + for attr in ("weight_scale", "weight_scale_2", "input_scale"): + if hasattr(mod, attr): + delattr(mod, attr) + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + mod.quant_method = UnquantizedLinearMethod() + + @staticmethod + def _ue8m0_to_float32(sf: torch.Tensor) -> torch.Tensor: + """Convert NVFP4 block scales (float8_e4m3fn / UE4M3) to float32. + + Checkpoint stores float8_e4m3fn (standard NVFP4 spec, NOT UE8M0). + Simple .to(float32) is correct — shift-by-23 was wrong (Bug #7 fix). + """ + return sf.to(torch.float32) + + def _unpack_nvfp4_to_bf16(self, w_uint8, e2m1_lut, device): + """Unpack NVFP4 uint8 packed weights to bf16 using E2M1 format.""" + # Extract 4-bit FP4 values (0-15, bit 3 = sign) + even_raw = (w_uint8 & 0x0F).int() + odd_raw = ((w_uint8 >> 4) & 0x0F).int() + # Sign: 0-7 = positive, 8-15 = negative + even_sign = torch.where(even_raw >= 8, -1.0, 1.0).to(torch.bfloat16) + odd_sign = torch.where(odd_raw >= 8, -1.0, 1.0).to(torch.bfloat16) + # Magnitude index: lower 3 bits (0-7) + even_vals = even_sign * e2m1_lut.to(device)[even_raw & 0x07] + odd_vals = odd_sign * e2m1_lut.to(device)[odd_raw & 0x07] + # Interleave and flatten + w_bf16 = torch.stack([even_vals, odd_vals], dim=-1) + w_bf16 = w_bf16.reshape(w_uint8.shape[0], -1).to(torch.bfloat16) + return w_bf16 +@torch.compile(backend=current_platform.simple_compile_backend) +def hc_head( + hidden_states: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + hc_mult, hidden_size = hidden_states.shape[-2:] + outer_shape = hidden_states.shape[:-2] + hs_flat = hidden_states.view(-1, hc_mult, hidden_size) + num_tokens = hs_flat.shape[0] + out = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=hidden_states.device + ) + torch.ops.vllm.hc_head_fused_kernel( + hs_flat, + hc_fn, + hc_scale, + hc_base, + out, + hidden_size, + rms_norm_eps, + hc_eps, + hc_mult, + ) + return out.view(*outer_shape, hidden_size) + + +def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: + if expert_dtype == "fp4": + # MXFP4 experts use Mxfp4MoEMethod, which registers scales as + # ``w{1,2,3}_weight_scale`` (no _inv suffix). FP8 linear and + # shared experts use Fp8LinearMethod's block scales, which + # register as ``weight_scale_inv``. + scale_regex = { + re.compile(r"(\.experts\.\d+\.w[123])\.scale$"): r"\1.weight_scale", + re.compile(r"\.scale$"): ".weight_scale_inv", + } + else: + # FP8 experts use Fp8MoEMethod (block_quant=True), which registers + # scales as ``w{13,2}_weight_scale_inv``. Map all ``.scale`` keys + # there. + scale_regex = { + re.compile(r"\.scale$"): ".weight_scale_inv", + } + + # ── ModelOpt NVFP4 export patches ──────────────────────────────── + # modelopt exports with different naming than the original HF ckpt: + # - Expert projections: gate_proj/up_proj/down_proj → w1/w3/w2 + # - Shared expert projections: gate_proj/up_proj → w1/w3 (stacking) + # - Compressor: kv_proj → wkv, gate_proj → wgate (stacking) + # - Attention: self_attn prefix, kv_proj → wkv (stacking) + # - modelopt uses mlp, vllm uses ffn + # Order matters for regex: skip patterns MUST come before renames. + + # Skip NVFP4 scales for compressor+attention fused params. + # After substr renaming, these map to stacked params (fused_wkv_wgate, + # fused_wqa_wkv, gate_up_proj) which don't register NVFP4 scale params + # because ModelOptNvFp4Config only handles Linear, not + # MergedColumnParallelLinear. We unpack weights as bf16 and let + # process_weights_after_loading re-quantize them. + # Must match ORIGINAL checkpoint key names (before substr renaming). + fused_skip_regex = { + # Compressor: SKIP ALL tensors. The compressor uses quant_config=None, + # so MergedColumnParallelLinear creates bf16 weight params. NVFP4 uint8 + # checkpoint data can't be loaded into these params (shape mismatch: + # uint8 (head_dim, hidden_size//2) vs bf16 (head_dim, hidden_size)). + # The stacking weight_loader silently skips the sub-weights, leaving + # random bf16 initialization. We reconstruct the compressor weights + # manually in post-load conversion by reading from the checkpoint. + re.compile(r"\.compressor\.kv_proj\.weight$"): None, + re.compile(r"\.compressor\.gate_proj\.weight$"): None, + re.compile(r"\.compressor\.kv_proj\.weight_scale$"): None, + re.compile(r"\.compressor\.gate_proj\.weight_scale$"): None, + re.compile(r"\.compressor\.kv_proj\.weight_scale_2$"): None, + re.compile(r"\.compressor\.gate_proj\.weight_scale_2$"): None, + re.compile(r"\.compressor\.kv_proj\.input_scale$"): None, + re.compile(r"\.compressor\.gate_proj\.input_scale$"): None, + # Note: attention and shared expert scale tensors are NO LONGER + # skipped. After fixing substr mappings, they correctly map to the + # model's NVFP4 scale parameters (fused_wqa_wkv, wq_b, wo_a, + # wo_b, gate_up_proj). They load via the stacking logic. + } + # Routed expert projections: gate_proj→w1, up_proj→w3, down_proj→w2 + # Regex (not substr) to match ONLY .experts.N. — not .shared_experts. + expert_rename_regex = { + re.compile(r"(\.experts\.\d+\.)gate_proj\."): r"\1w1.", + re.compile(r"(\.experts\.\d+\.)up_proj\."): r"\1w3.", + re.compile(r"(\.experts\.\d+\.)down_proj\."): r"\1w2.", + } + # Merge: skip patterns first, then renames, then original scale_regex + merged_regex = {} + merged_regex.update(fused_skip_regex) + merged_regex.update(expert_rename_regex) + merged_regex.update(scale_regex) + + return WeightsMapper( + orig_to_new_prefix={ + "layers.": "model.layers.", + "embed.": "model.embed.", + "norm.": "model.norm.", + "hc_head": "model.hc_head", + "mtp.": "model.mtp.", + }, + orig_to_new_regex=merged_regex, + orig_to_new_suffix={ + "embed.weight": "embed_tokens.weight", + ".ffn.gate.bias": ".ffn.gate.e_score_correction_bias", + }, + orig_to_new_substr={ + ".attn.compressor.": ".attn.mla_attn.compressor.", + ".shared_experts.w2": ".shared_experts.down_proj", + # ── ModelOpt NVFP4 substr patches ── + # Attention: self_attn → attn (projections at attn level, not mla_attn) + ".self_attn.q_a_proj.": ".attn.wq_a.", + ".self_attn.q_b_proj.": ".attn.wq_b.", + ".self_attn.q_a_norm.": ".attn.q_norm.", + ".self_attn.o_a_proj.": ".attn.wo_a.", + ".self_attn.o_b_proj.": ".attn.wo_b.", + ".self_attn.sinks": ".attn.attn_sink", + # kv_proj → wkv (for stacking into fused_wqa_wkv) + ".self_attn.kv_proj.": ".attn.wkv.", + ".self_attn.kv_norm.": ".attn.kv_norm.", + # kv_norm is at attention level, not compressor/mla_attn level in vllm + # Must come before the general compressor mapping + ".self_attn.compressor.kv_norm.": ".attn.kv_norm.", + # Compressor: self_attn.compressor → attn.mla_attn.compressor + ".self_attn.compressor.": ".attn.mla_attn.compressor.", + # Compressor projections for stacking (fused_wkv_wgate) + ".compressor.kv_proj.": ".compressor.wkv.", + ".compressor.gate_proj.": ".compressor.wgate.", + # Shared expert projections (stacking into gate_up_proj) + # Checkpoint has .shared_experts. but model has .ffn.shared_experts. + ".shared_experts.gate_proj.": ".ffn.shared_experts.w1.", + ".shared_experts.up_proj.": ".ffn.shared_experts.w3.", + # modelopt uses mlp, vllm uses ffn internally + ".mlp.": ".ffn.", + }, + ) + + +class DeepseekV4ForCausalLM(nn.Module): + model_cls = DeepseekV4Model + + # NOTE: We do NOT set hf_to_vllm_mapper here because our custom + # load_weights handles all checkpoint→model name remapping inline. + # If hf_to_vllm_mapper is set, vLLM's AutoWeightsLoader may be invoked + # INSTEAD of our load_weights, silently dropping NVFP4 weight loading. + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config + self.config = config + + self.model = self.model_cls( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + last = logits[-1].float() + top_vals, top_idx = last.topk(10) + print(f"[LOGITS] shape={tuple(logits.shape)} " + f"max={last.max().item():.4e} min={last.min().item():.4e} " + f"mean={last.mean().item():.4e} std={last.std().item():.4e} " + f"finite_frac={torch.isfinite(last).float().mean().item():.4f}") + print(f"[LOGITS] top10 ids: {top_idx.tolist()}") + print(f"[LOGITS] top10 vals: {[f'{v:.3f}' for v in top_vals.tolist()]}") + print(f"[LOGITS] gap top1-top10: {(top_vals[0] - top_vals[-1]).item():.3f}") + + # Probe for "Paris" specifically + try: + from transformers import AutoTokenizer + if not hasattr(self, '_tok'): + self._tok = AutoTokenizer.from_pretrained('/model', trust_remote_code=True) + for v in [' Paris', 'Paris', ' paris', 'paris']: + tid = self._tok.encode(v, add_special_tokens=False)[0] + logit_val = last[tid].item() + rank = (last > logit_val).sum().item() + print(f"[LOGITS-PROBE] {repr(v)}→id={tid} logit={logit_val:.2f} rank={rank}") + except Exception as e: + print(f"[LOGITS-PROBE] failed: {e}") + + return logits + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def get_mtp_target_hidden_states(self) -> torch.Tensor | None: + """Pre-hc_head residual stream buffer (max_num_batched_tokens, + hc_mult * hidden_size) for the MTP draft model. Populated by + forward(); valid after each target step.""" + return getattr(self.model, "_mtp_hidden_buffer", None) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # lm_head lives on this outer model, not on the inner DeepseekV4Model. + # The inner load_weights silently drops lm_head.weight via + # "if name not in params_dict: continue". Extract it here. + rest = [] + for name, loaded_weight in weights: + if name == "lm_head.weight" or name.endswith(".lm_head.weight"): + param = self.lm_head.weight + weight_loader = getattr(param, "weight_loader", + default_weight_loader) + weight_loader(param, loaded_weight) + else: + rest.append((name, loaded_weight)) + # Use the model-level loader which handles NVFP4 expert mapping, + # uint8→bf16 unpacking for MergedColumnParallelLinear, and + # bf16→NVFP4 quantization for unquantized layers. + # AutoWeightsLoader bypasses this logic and would break NVFP4 loading. + loaded_params = self.model.load_weights(rest) + loaded_params.add("lm_head.weight") + self.model.finalize_mega_moe_weights() + self.model._convert_nvfp4_post_load() + if int(os.environ.get('NVFP4_DEBUG', '0')): + # Count loaded expert weights to catch silent load failures + for i, layer in enumerate(self.model.layers): + ffn = layer.ffn + if hasattr(ffn, 'experts') and hasattr(ffn.experts, 'w13_weight'): + w13 = ffn.experts.w13_weight + w13_sf = ffn.experts.w13_weight_scale + w13_sf2 = ffn.experts.w13_weight_scale_2 + w2 = ffn.experts.w2_weight + n_experts = w13.shape[0] + nonzero_w13 = (w13.abs().amax(dim=(1,2)) > 0).sum().item() + nonzero_w2 = (w2.abs().amax(dim=(1,2)) > 0).sum().item() + print(f"[NVFP4_DEBUG] Layer {i}: {nonzero_w13}/{n_experts} w13 nonzero, " + f"{nonzero_w2}/{n_experts} w2 nonzero, " + f"w13_sf shape={w13_sf.shape}, w13_sf2 shape={w13_sf2.shape}") + if os.environ.get('NVFP4_DEBUG_SYNC', '') == '1': + torch.cuda.synchronize() + print("[NVFP4] post-load conversion done, CUDA OK") + + # POST-LOAD: scan for all-zero params (missed renames, failed loads) + zero_attrs = [] + for name, p in self.named_parameters(): + if not torch.is_tensor(p): + continue + sample = p.flatten()[:1024] if p.numel() > 1024 else p.flatten() + if (sample == 0).all().item(): + if (p == 0).all().item(): + zero_attrs.append((name, tuple(p.shape), str(p.dtype))) + print(f"[POST-LOAD] {len(zero_attrs)} all-zero param tensors:") + for n, s, d in zero_attrs[:50]: + print(f" {n} shape={s} dtype={d}") + + return loaded_params + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() + diff --git a/vllm/patches/deepseek_v4_attention.py b/vllm/patches/deepseek_v4_attention.py new file mode 100644 index 00000000..2bde4a53 --- /dev/null +++ b/vllm/patches/deepseek_v4_attention.py @@ -0,0 +1,1155 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +DeepseekV4 MLA Attention Layer +""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import DeepseekV2Config, DeepseekV3Config + +import vllm.envs as envs +from vllm.model_executor.layers.linear import ( + ReplicatedLinear, +) +from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer +from vllm.utils.deep_gemm import fp8_einsum +from vllm.utils.torch_utils import direct_register_custom_op +from vllm.v1.attention.ops.deepseek_v4_ops import ( + combine_topk_swa_indices, + compute_global_topk_indices_and_lens, + dequantize_and_gather_k_cache, + fused_indexer_q_rope_quant, + fused_inv_rope_fp8_quant, + fused_q_kv_rmsnorm, +) + +if TYPE_CHECKING: + from vllm.v1.attention.backends.mla.sparse_swa import ( + DeepseekSparseSWAMetadata, + ) + +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, +) +from vllm.distributed import get_tensor_model_parallel_world_size +from vllm.forward_context import ForwardContext, get_forward_context +from vllm.logger import init_logger +from vllm.model_executor.custom_op import PluggableLayer +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.deepseek_compressor import DeepseekCompressor +from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.input_quant_fp8 import ( + QuantFP8, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, +) +from vllm.utils.multi_stream_utils import ( + execute_in_parallel, + maybe_execute_in_parallel, +) +from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata +from vllm.v1.attention.backends.mla.flashmla_sparse import ( + DeepseekV4FlashMLASparseBackend, + FlashMLASparseBackend, + FlashMLASparseMetadata, +) +from vllm.v1.attention.backends.mla.indexer import ( + DeepseekV4IndexerBackend, + get_max_prefill_buffer_size, +) +from vllm.v1.attention.backends.mla.sparse_swa import DeepseekV4SWACache +from vllm.v1.attention.ops.flashmla import ( + flash_mla_sparse_fwd, + flash_mla_with_kvcache, +) +from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec +from vllm.v1.worker.workspace import current_workspace_manager + +logger = init_logger(__name__) + +# Prefill is processed in fixed-size chunks; this bounds the bf16 kv-gather +# workspace allocated at _forward_prefill (and the matching profile-time +# reservation in attention_impl's dummy-run branch). +PREFILL_CHUNK_SIZE = 4 + + +@dataclass +class DeepseekV4MLAModules: + """Modules used in DeepseekV4 MLA.""" + + vllm_config: VllmConfig + fused_wqa_wkv: torch.nn.Module + q_norm: torch.nn.Module + wq_b: torch.nn.Module + kv_norm: torch.nn.Module + wo_a: torch.nn.Module + wo_b: torch.nn.Module + attn_sink: torch.nn.Module + rotary_emb: torch.nn.Module + indexer: torch.nn.Module | None + indexer_rotary_emb: torch.nn.Module + topk_indices_buffer: torch.Tensor | None + aux_stream_list: list[torch.cuda.Stream] | None = None + + +# --8<-- [start:multi_head_latent_attention] +@PluggableLayer.register("deepseek_v4_multi_head_latent_attention") +class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer): + """Pluggable MLA layer which allows OOT backends to add + custom implementations of the outer MLA layer (including rope & o_proj). + Note that currently oot platforms can still use CustomOp.register_oot to + replace MLA layer entirely, although we use PluggableLayer to register + this layer now. + + This class takes positions and hidden_states as input. + The input tensors can either contain prefill tokens or decode tokens. + The class does the following: + + 1. MLA Preprocess. + 2. Perform multi-head attention to prefill tokens and + multi-query attention to decode tokens separately. + 3. Return the output tensor. + """ + + # --8<-- [end:multi_head_latent_attention] + + def __init__( + self, + hidden_size: int, + num_heads: int, + head_dim: int, + scale: float, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + o_lora_rank: int | None, + mla_modules: DeepseekV4MLAModules, + window_size: int, + compress_ratio: int | None, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.n_local_heads = num_heads + self.head_dim = head_dim + self.scale = scale + + # FlashMLA sparse kernel only supports 64 or 128 heads; pad up to the + # next supported size. Must match DeepseekV4MLAAttention.padded_heads. + if num_heads <= 64: + self.padded_heads = 64 + elif num_heads <= 128: + self.padded_heads = 128 + else: + raise ValueError( + f"DeepseekV4 attention does not support {num_heads} heads " + "(must be <= 128)." + ) + + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.window_size = window_size + self.compress_ratio = compress_ratio if compress_ratio is not None else 1 + self.prefix = prefix + + # Extract config from vllm_config + config = mla_modules.vllm_config.model_config.hf_config + tp_size = get_tensor_model_parallel_world_size() + + # DeepseekV4-specific attributes (num_heads is already TP-adjusted) + self.eps = config.rms_norm_eps + self.rope_head_dim = config.qk_rope_head_dim + self.nope_head_dim = head_dim - self.rope_head_dim + self.n_local_groups = config.o_groups // tp_size + self.o_lora_rank = config.o_lora_rank + + # Store projection modules + self.fused_wqa_wkv = mla_modules.fused_wqa_wkv + self.q_norm = mla_modules.q_norm + self.wq_b = mla_modules.wq_b + + self.kv_norm = mla_modules.kv_norm + self.wo_a = mla_modules.wo_a + + self._wo_a_act_quant = QuantFP8( + static=False, + group_shape=GroupShape(1, 128), + use_ue8m0=True, + ) + # Bypass packed-for-deepgemm path — we need FP32 scales (not packed + # INT32) so fp8_einsum can handle layout transform internally. + self._wo_a_act_quant.use_deep_gemm_supported = False + self.wo_b = mla_modules.wo_b + + # Pick fp8_einsum recipe based on GPU arch: + # SM90: FP32 block scales stay [g, r/128, d/128] → sfb_gran_mn=128 + # SM100: INT32 packed scales become [g, r, ...] → sfb_gran_mn=1 + from vllm.platforms import current_platform + + cap = current_platform.get_device_capability() + assert cap is not None, "DeepseekV4 attention requires a CUDA device" + self._einsum_recipe = (1, 128, 128) if cap.major <= 9 else (1, 1, 128) + self._tma_aligned_scales = cap.major >= 10 + + self.rotary_emb = mla_modules.rotary_emb + self.indexer_rotary_emb = mla_modules.indexer_rotary_emb + self.topk_indices_buffer = mla_modules.topk_indices_buffer + + self.indexer = mla_modules.indexer + + # Per-head RMS normalization for Q (no learnable weights) + self.q_head_norm = RMSNorm(head_dim, eps=self.eps, has_weight=False) + + # TODO(yifan): currently hardcoded for FP8 sparse, make it more generic + head_bytes = ( + self.nope_head_dim # 448 fp8 NoPE + + self.rope_head_dim * 2 # 64 bf16 RoPE + + self.nope_head_dim // 64 # 7B scale factors + + 1 # 1B pad + ) + + self.aux_stream_list = mla_modules.aux_stream_list + # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events; + # [1] doubles as post-GEMM event1. Reuse is safe: GEMM fully joins + # before post-GEMM starts. + self.ln_events = [torch.cuda.Event() for _ in range(4)] + + assert cache_config is not None, "DeepseekV4 attention requires cache_config" + self.swa_cache_layer = DeepseekV4SWACache( + head_dim=self.head_dim, + window_size=self.window_size, + dtype=torch.uint8, + prefix=f"{prefix}.swa_cache", + cache_config=cache_config, + ) + + self.mla_attn = DeepseekV4MLAAttention( + num_heads=self.n_local_heads, + head_dim=self.head_dim, + scale=self.scale, + qk_nope_head_dim=self.nope_head_dim, + qk_rope_head_dim=self.rope_head_dim, + q_lora_rank=self.q_lora_rank, + kv_lora_rank=self.kv_lora_rank, + compress_ratio=self.compress_ratio, + window_size=self.window_size, + head_bytes=head_bytes, + swa_cache_layer=self.swa_cache_layer, + attn_sink=mla_modules.attn_sink, # already padded with -inf + cache_config=cache_config, + quant_config=quant_config, + prefix=prefix, + indexer=self.indexer, + topk_indices_buffer=self.topk_indices_buffer, + ) + # Register this layer in the compilation config's static forward context + # This allows the custom op to retrieve the layer during execution + compilation_config = mla_modules.vllm_config.compilation_config + # HACK + self.layer_name = prefix + ".deepseek_v4_multi_head_latent_attention" + if self.layer_name in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {self.layer_name}") + compilation_config.static_forward_context[self.layer_name] = self + + # Create the compressor for layers with compress_ratio > 1; after + # creating the DeepseekV4MLAAttention layer to get its cache. + self.compressor = None + if self.compress_ratio > 1: + self.compressor = DeepseekCompressor( + vllm_config=mla_modules.vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=self.hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.mla_attn.prefix, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + llama_4_scaling: torch.Tensor | None = None, + ) -> torch.Tensor: + # Pre-allocate attention output with FlashMLA-padded head count. + # The op writes into `o_padded`; we slice to n_local_heads after. + num_tokens = hidden_states.shape[0] + o_padded = torch.empty( + (num_tokens, self.padded_heads, self.head_dim), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + # Attention (inside custom op for torch.compile boundary) + torch.ops.vllm.deepseek_v4_attention( + hidden_states, + positions, + o_padded, + self.layer_name, + ) + o = o_padded[:, : self.n_local_heads, :] + + # O projection: inverse RoPE + FP8 quant + einsum + wo_b + o_fp8, o_scale = fused_inv_rope_fp8_quant( + o, + positions, + self.rotary_emb.cos_sin_cache.to(torch.float32), + n_groups=self.n_local_groups, + heads_per_group=self.n_local_heads // self.n_local_groups, + nope_dim=self.nope_head_dim, + rope_dim=self.rope_head_dim, + tma_aligned_scales=self._tma_aligned_scales, + ) + + wo_a_fp8 = self.wo_a.weight + wo_a_scale = self.wo_a.weight_scale_inv + + z = torch.empty( + (num_tokens, self.n_local_groups, self.o_lora_rank), + device=o.device, + dtype=torch.bfloat16, + ) + torch.ops.vllm.deepseek_v4_fp8_einsum( + o_fp8, + o_scale, + wo_a_fp8, + wo_a_scale, + z, + "bhr,hdr->bhd", + list(self._einsum_recipe), + ) + + return self.wo_b(z.flatten(1)) + + def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: + assert self.aux_stream_list is not None + assert len(self.aux_stream_list) >= 3 + + # fused_wqa_wkv (heaviest) on default; the three lighter input GEMMs + # on aux streams 0..2 when their owning module exists. ln_events[0] + # is the fan-out start event; ln_events[1..3] are per-aux done events. + aux_fns: list[Callable[[], Any] | None] = [None, None, None] + + if self.compressor is not None: + # Local ref so the closure keeps a non-None type for mypy. + compressor = self.compressor + + def compressor_kv_score() -> torch.Tensor: + return torch.mm( + hidden_states, + compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[0] = compressor_kv_score + + if self.indexer is not None: + indexer = self.indexer + + def indexer_weights_proj() -> torch.Tensor: + # ReplicatedLinear returns (output, bias); bias is None. + weights, _ = indexer.weights_proj(hidden_states) + return weights + + def indexer_compressor_kv_score() -> torch.Tensor: + return torch.mm( + hidden_states, + indexer.compressor.fused_wkv_wgate.weight.T, + out_dtype=torch.float32, + ) + + aux_fns[1] = indexer_weights_proj + aux_fns[2] = indexer_compressor_kv_score + + def fused_wqa_wkv() -> torch.Tensor: + # MergedColumnParallelLinear returns (output, bias); bias is None. + qr_kv, _ = self.fused_wqa_wkv(hidden_states) + return qr_kv + + qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel( + fused_wqa_wkv, + aux_fns, + self.ln_events[0], + self.ln_events[1:4], + self.aux_stream_list[:3], + enable=hidden_states.shape[0] + <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD, + ) + + return qr_kv, kv_score, indexer_kv_score, indexer_weights + + def attention_impl( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place + ) -> None: + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + + qr_kv, kv_score, indexer_kv_score, indexer_weights = ( + self.attn_gemm_parallel_execute(hidden_states) + ) + + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + qr, kv = fused_q_kv_rmsnorm( + qr, + kv, + self.q_norm.weight.data, + self.kv_norm.weight.data, + self.eps, + ) + + # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride + # on the default stream so q stays on its consumer stream (mla_attn + # downstream reads q on default). Indexer/compressor go on aux for + # overlap with default's GEMM + cache write. + if self.indexer is not None: + assert self.aux_stream_list is not None + aux_stream = self.aux_stream_list[0] + indexer = self.indexer + # Local ref so the closure keeps a non-None type for mypy. + assert self.compressor is not None + compressor = self.compressor + + def wq_b_kv_insert_and_compress() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + compressor(kv_score, positions, self.rotary_emb) + return q + + q, _ = maybe_execute_in_parallel( + wq_b_kv_insert_and_compress, + lambda: indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + ), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + elif self.compressor is not None: + # wq_b + kv_insert on default, compressor on aux. + assert self.aux_stream_list is not None + aux_stream = self.aux_stream_list[0] + compressor = self.compressor + + def wq_b_kv_insert() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + return q + + q, _ = maybe_execute_in_parallel( + wq_b_kv_insert, + lambda: compressor(kv_score, positions, self.rotary_emb), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + # SWA-only layer: no compressor, no overlap. + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + # Handle dummy run (no metadata). + if not isinstance(attn_metadata, dict): + # Reserve _forward_prefill's bf16-gather workspace; the dummy + # run returns before mla_attn runs, so without this the shared + # workspace locks below the real prefill size. + sub = self.mla_attn + swa_only = sub.compress_ratio <= 1 + N = ( + 0 + if swa_only + else (sub.max_model_len + sub.compress_ratio - 1) // sub.compress_ratio + ) + M = N + sub.window_size + sub.max_num_batched_tokens + current_workspace_manager().get_simultaneous( + ((PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + ) + out.zero_() + return + + # Pad q to FlashMLA-required head count (64 or 128) + if self.n_local_heads < self.padded_heads: + pad_size = self.padded_heads - self.n_local_heads + q = F.pad(q, (0, 0, 0, pad_size), value=0.0) + + # MLA attention writes into the pre-allocated `out` buffer + # ([num_tokens, padded_heads, head_dim]). + self.mla_attn(q, kv, positions, output=out) + + def _fused_qnorm_rope_kv_insert( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + attn_metadata: ( + dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]] | None + ), + ) -> None: + if not isinstance(attn_metadata, dict): + return + + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_kv_cache = self.swa_cache_layer.kv_cache + swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1) + + # Horizontally fused: + # Q side: q_head_norm (per-head RMSNorm, no weight) + GPT-J RoPE + # KV side: GPT-J RoPE + UE8M0 FP8 quant + paged cache insert + # kv is unchanged; mla_attn reads kv solely via swa_kv_cache. + torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert( + q, + kv, + swa_kv_cache_2d, + swa_metadata.slot_mapping, + positions.to(torch.int64), + self.rotary_emb.cos_sin_cache.to(torch.float32), + self.eps, + swa_metadata.block_size, + ) + + +def deepseek_v4_attention( + hidden_states: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, + layer_name: str, +) -> None: + forward_context: ForwardContext = get_forward_context() + self = forward_context.no_compile_layers[layer_name] + self.attention_impl(hidden_states, positions, out) + + +def deepseek_v4_attention_fake( + hidden_states: torch.Tensor, + positions: torch.Tensor, + out: torch.Tensor, + layer_name: str, +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_attention", + op_func=deepseek_v4_attention, + mutates_args=["out"], + fake_impl=deepseek_v4_attention_fake, +) + + +def deepseek_v4_fp8_einsum( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + fp8_einsum(equation, (a, a_scale), (b, b_scale), out, recipe=tuple(recipe)) + + +def deepseek_v4_fp8_einsum_fake( + a: torch.Tensor, + a_scale: torch.Tensor, + b: torch.Tensor, + b_scale: torch.Tensor, + out: torch.Tensor, + equation: str, + recipe: list[int], +) -> None: + return None + + +direct_register_custom_op( + op_name="deepseek_v4_fp8_einsum", + op_func=deepseek_v4_fp8_einsum, + mutates_args=["out"], + fake_impl=deepseek_v4_fp8_einsum_fake, +) + + +class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase): + # FlashMLA FP8 sparse only supports 64 or 128 heads + SUPPORTED_HEAD_COUNTS = (64, 128) + + def __init__( + self, + num_heads: int, + head_dim: int, + scale: float, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + q_lora_rank: int | None, + kv_lora_rank: int, + compress_ratio: int, + window_size: int, + head_bytes: int, + swa_cache_layer: DeepseekV4SWACache, + attn_sink: torch.Tensor, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + # Sparse MLA Args + indexer: object | None = None, + topk_indices_buffer: torch.Tensor | None = None, + aux_stream: torch.cuda.Stream | None = None, + **extra_impl_args, + ) -> None: + super().__init__() + self.num_heads = num_heads + self.num_kv_heads = 1 + self.head_dim = head_dim + self.scale = scale + self.window_size = window_size + self.head_bytes = head_bytes + self.compress_ratio = compress_ratio + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.nope_head_dim = qk_nope_head_dim + self.rope_head_dim = qk_rope_head_dim + self.indexer = indexer + self.topk_indices_buffer = topk_indices_buffer + + self.prefix = prefix # Alias for compatibility with compressor + + self.aux_stream = aux_stream + self.ln_events = [torch.cuda.Event(), torch.cuda.Event()] + + # Determine padded head count for FlashMLA + if num_heads not in self.SUPPORTED_HEAD_COUNTS: + if num_heads < 64: + self.padded_heads = 64 + elif num_heads < 128: + self.padded_heads = 128 + else: + raise ValueError( + f"DeepseekV4MLAAttention does not support {num_heads} heads. " + f"Supported: <= 128 (will be padded to 64 or 128)" + ) + else: + self.padded_heads = num_heads + + # Store attention sink + assert attn_sink is not None + self.attn_sink: torch.Tensor = attn_sink + # Store SWA cache + assert swa_cache_layer is not None + self.swa_cache_layer: DeepseekV4SWACache = swa_cache_layer + + # Get vllm config for cache setup + vllm_config = get_current_vllm_config() + self.max_num_batched_tokens = ( + vllm_config.scheduler_config.max_num_batched_tokens + ) + self.max_model_len = vllm_config.model_config.max_model_len + # DeepseekV4 only supports fp8 kv-cache format for now + kv_cache_dtype = cache_config.cache_dtype if cache_config is not None else "fp8" + + assert kv_cache_dtype.startswith("fp8"), ( + f"DeepseekV4 only supports fp8 kv-cache format for now, " + f"got {kv_cache_dtype}" + ) + assert issubclass(self.get_attn_backend(), FlashMLASparseBackend), ( + "Only FlashMLA Sparse Attention backend is supported for DeepseekV4 for now" + ) + # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format + # Automatically convert fp8 kv-cache format to "fp8_ds_mla" + if ( + issubclass(self.get_attn_backend(), FlashMLASparseBackend) + and kv_cache_dtype.startswith("fp8") + and kv_cache_dtype != "fp8_ds_mla" + ): + assert cache_config is not None + cache_config.cache_dtype = "fp8_ds_mla" + kv_cache_dtype = "fp8_ds_mla" + logger.info_once("Using DeepSeek's fp8_ds_mla KV cache format.") + + self.kv_cache_dtype = kv_cache_dtype + + # Register with compilation context for metadata lookup + compilation_config = vllm_config.compilation_config + if prefix and prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + if prefix: + compilation_config.static_forward_context[prefix] = self + + self.kv_cache = torch.tensor([]) + + def get_attn_backend(self) -> type[AttentionBackend]: + return DeepseekV4FlashMLASparseBackend + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: + if ( + self.compress_ratio <= 1 + ): # SWA part. Allocated separately as DeepseekV4SWACache. + return None + return MLAAttentionSpec( + block_size=vllm_config.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=torch.uint8, + compress_ratio=self.compress_ratio, + cache_dtype_str=self.kv_cache_dtype, + alignment=576, # NOTE: FlashMLA requires 576B alignment + model_version="deepseek_v4", + ) + + def forward( + self, + q: torch.Tensor, + kv: torch.Tensor, + positions: torch.Tensor, + output: torch.Tensor, + ) -> None: + assert output.shape == q.shape, ( + f"output buffer shape {output.shape} must match q shape {q.shape}" + ) + assert output.dtype == q.dtype, ( + f"output buffer dtype {output.dtype} must match q dtype {q.dtype}" + ) + + # Get SWA and indexer metadata from forward context + forward_context = get_forward_context() + attn_metadata = forward_context.attn_metadata + assert isinstance(attn_metadata, dict) + flashmla_metadata = cast( + FlashMLASparseMetadata | None, attn_metadata.get(self.prefix) + ) + swa_metadata = cast( + "DeepseekSparseSWAMetadata | None", + attn_metadata.get(self.swa_cache_layer.prefix), + ) + assert swa_metadata is not None + + swa_only = self.compress_ratio <= 1 + # SWA-only layers (compress_ratio <= 1) don't have their own KV cache + # allocation, so self.kv_cache may be empty after profiling cleanup. + self_kv_cache = self.kv_cache if not swa_only else None + swa_kv_cache = self.swa_cache_layer.kv_cache + + # Split prefill and decode + num_decodes = swa_metadata.num_decodes + num_prefills = swa_metadata.num_prefills + num_decode_tokens = swa_metadata.num_decode_tokens + + if num_prefills > 0: + self._forward_prefill( + q=q[num_decode_tokens:], + positions=positions[num_decode_tokens:], + compressed_k_cache=self_kv_cache, + swa_k_cache=swa_kv_cache, + output=output[num_decode_tokens:], + attn_metadata=flashmla_metadata, + swa_metadata=swa_metadata, + ) + if num_decodes > 0: + self._forward_decode( + q=q[:num_decode_tokens], + kv_cache=self_kv_cache, + swa_metadata=swa_metadata, + attn_metadata=flashmla_metadata, + swa_only=swa_only, + output=output[:num_decode_tokens], + ) + + def _forward_decode( + self, + q: torch.Tensor, + kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_metadata: "DeepseekSparseSWAMetadata", + attn_metadata: FlashMLASparseMetadata | None, + swa_only: bool, + output: torch.Tensor, + ) -> None: + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + topk_indices = None + topk_lens = None + if not swa_only: + assert attn_metadata is not None + assert swa_metadata.is_valid_token is not None + block_size = attn_metadata.block_size // self.compress_ratio + is_valid = swa_metadata.is_valid_token[:num_decode_tokens] + if self.compress_ratio == 4: + # C4A: local indices differ per layer (filled by Indexer). + assert self.topk_indices_buffer is not None + global_indices, topk_lens = compute_global_topk_indices_and_lens( + self.topk_indices_buffer[:num_decode_tokens], + swa_metadata.token_to_req_indices, + attn_metadata.block_table[:num_decodes], + block_size, + is_valid, + ) + topk_indices = global_indices.view(num_decode_tokens, 1, -1) + else: + # C128A: pre-computed during metadata build. + topk_indices = attn_metadata.c128a_global_decode_topk_indices + topk_lens = attn_metadata.c128a_decode_topk_lens + + swa_indices = swa_metadata.decode_swa_indices + swa_lens = swa_metadata.decode_swa_lens + + # We treat queries in the same seq as different queries + # and later we only attend by generated indices. + # q arrives pre-padded to self.padded_heads by the outer wrapper. + q = q.unsqueeze(1) + + # Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes) + # Use unsqueeze to preserve strides (handles padded blocks correctly) + swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2) + # Reshape KV cache to (num_blocks, block_size, 1, head_bytes) + if kv_cache is not None: + kv_cache = kv_cache.unsqueeze(-2) + + # One FlashMLASchedMeta per layer type, shared across all same-type + # layers within this decode step. The first forward call per type + # triggers the in-kernel planner (allocating tile_scheduler_metadata + # and num_splits via PyTorch's graph-aware allocator so CUDA graph + # capture reuses the same addresses on replay); subsequent same-type + # layers see have_initialized=True and skip the planner. + if self.compress_ratio <= 1: + tile_metadata = swa_metadata.tile_sched_swaonly + elif self.compress_ratio == 4: + tile_metadata = swa_metadata.tile_sched_c4a + elif self.compress_ratio == 128: + tile_metadata = swa_metadata.tile_sched_c128a + else: + raise ValueError( + f"Unsupported compress_ratio={self.compress_ratio}; " + "expected 1, 4, or 128." + ) + assert tile_metadata is not None, ( + "swa_metadata missing tile_sched entry for " + f"compress_ratio={self.compress_ratio}; " + "DeepseekSparseSWAMetadataBuilder.build_tile_scheduler did not " + "allocate one for this layer type." + ) + + out, _ = flash_mla_with_kvcache( + q=q, + k_cache=swa_cache, + block_table=None, + head_dim_v=512, + tile_scheduler_metadata=tile_metadata, + cache_seqlens=None, + is_fp8_kvcache=True, + indices=swa_indices, + topk_length=swa_lens, + softmax_scale=self.scale, + attn_sink=self.attn_sink, + extra_k_cache=kv_cache if not swa_only else None, + extra_indices_in_kvcache=topk_indices, + extra_topk_length=topk_lens, + out=output.unsqueeze(1), + ) + + def _forward_prefill( + self, + q: torch.Tensor, + positions: torch.Tensor, + compressed_k_cache: torch.Tensor | None, # Only used when compress_ratio > 1 + swa_k_cache: torch.Tensor, + output: torch.Tensor, + attn_metadata: FlashMLASparseMetadata | None, + swa_metadata: "DeepseekSparseSWAMetadata", + ) -> None: + swa_only = attn_metadata is None + + num_prefills = swa_metadata.num_prefills + num_prefill_tokens = swa_metadata.num_prefill_tokens + num_decodes = swa_metadata.num_decodes + num_decode_tokens = swa_metadata.num_decode_tokens + + # Use pre-computed prefill metadata. + seq_lens = swa_metadata.prefill_seq_lens + gather_lens = swa_metadata.prefill_gather_lens + assert seq_lens is not None + assert gather_lens is not None + + # Derive prefill-local token offsets from the full query_start_loc_cpu. + query_start_loc_cpu = swa_metadata.query_start_loc_cpu + query_start_loc = swa_metadata.query_start_loc + assert query_start_loc_cpu is not None + assert query_start_loc is not None + prefill_token_base = query_start_loc_cpu[num_decodes] + + if not swa_only: + if self.compress_ratio == 4: + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + topk_indices = topk_indices[:num_prefill_tokens] + else: + # C128A: pre-computed during metadata build. + assert attn_metadata is not None + topk_indices = attn_metadata.c128a_prefill_topk_indices + top_k = topk_indices.shape[-1] + # Compressed region must fit the full compressed pool (seq_len // + # compress_ratio), not just top_k. top_k bounds how many indices + # the indexer selects, not the pool size it indexes into. + N = (self.max_model_len + self.compress_ratio - 1) // self.compress_ratio + else: + # NOTE(woosuk): topk_indices will not be used for SWA-only layers. + assert self.topk_indices_buffer is not None + topk_indices = self.topk_indices_buffer[num_decode_tokens:] + top_k = 0 + N = 0 + + M = N + self.window_size + self.max_num_batched_tokens + num_chunks = (num_prefills + PREFILL_CHUNK_SIZE - 1) // PREFILL_CHUNK_SIZE + + workspace_manager = current_workspace_manager() + kv = workspace_manager.get_simultaneous( + ((PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), + )[0] + for chunk_idx in range(num_chunks): + chunk_start = chunk_idx * PREFILL_CHUNK_SIZE + chunk_end = min(chunk_start + PREFILL_CHUNK_SIZE, num_prefills) + chunk_size = chunk_end - chunk_start + if not swa_only: + # Gather compressed KV + assert attn_metadata is not None + block_table = attn_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + compressed_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end] // self.compress_ratio, + gather_lens=None, + block_table=block_table[chunk_start:chunk_end], + block_size=attn_metadata.block_size // self.compress_ratio, + offset=0, + ) + + # Gather SWA KV + swa_block_table = swa_metadata.block_table[num_decodes:] + dequantize_and_gather_k_cache( + kv[:chunk_size], + swa_k_cache, + seq_lens=seq_lens[chunk_start:chunk_end], + gather_lens=gather_lens[chunk_start:chunk_end], + block_table=swa_block_table[chunk_start:chunk_end], + block_size=swa_metadata.block_size, + offset=N, + ) + + # Combine the topk indices and SWA indices for gathered KV cache + query_start = ( + query_start_loc_cpu[num_decodes + chunk_start] - prefill_token_base + ) + query_end = ( + query_start_loc_cpu[num_decodes + chunk_end] - prefill_token_base + ) + + combined_indices, combined_lens = combine_topk_swa_indices( + topk_indices[query_start:query_end], + query_start_loc[ + num_decodes + chunk_start : num_decodes + chunk_end + 1 + ], + seq_lens[chunk_start:chunk_end], + gather_lens[chunk_start:chunk_end], + self.window_size, + self.compress_ratio, + top_k, + M, + N, + ) + + output_chunk, _, _ = flash_mla_sparse_fwd( + q=q[query_start:query_end], + kv=kv.view(-1, 1, q.shape[-1]), + indices=combined_indices.unsqueeze(1), + sm_scale=self.scale, + attn_sink=self.attn_sink, + topk_length=combined_lens, + out=output[query_start:query_end], + ) + + +class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): + def __init__( + self, + head_dim: int, + dtype: torch.dtype, + prefix: str, + cache_config: CacheConfig, + compress_ratio: int = 1, + ): + super().__init__() + self.kv_cache = torch.tensor([]) + self.head_dim = head_dim + self.prefix = prefix + self.cache_config = cache_config + self.dtype = dtype + self.compress_ratio = compress_ratio + compilation_config = get_current_vllm_config().compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: + # head_dim already carries the fp8 scale padding + # compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout. + return MLAAttentionSpec( + block_size=self.cache_config.block_size, + num_kv_heads=1, + head_size=self.head_dim, + dtype=self.dtype, + compress_ratio=self.compress_ratio, + # DeepseekV4 aligns indexer pages to FlashMLA's 576B so they can pack with + # the indexer's compressor state cache. V3.2 keeps the legacy layout. + alignment=576, + ) + + def forward(self): ... + + def get_attn_backend(self) -> type[AttentionBackend]: + return DeepseekV4IndexerBackend + + +class DeepseekV4Indexer(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + config: DeepseekV2Config | DeepseekV3Config, + hidden_size: int, + q_lora_rank: int, + quant_config: QuantizationConfig | None, + cache_config: CacheConfig | None, + topk_indices_buffer: torch.Tensor | None, + compress_ratio: int = 1, + prefix: str = "", + ): + super().__init__() + self.vllm_config = vllm_config + self.config = config + self.quant_config = quant_config + # self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"] + self.topk_tokens = config.index_topk + self.n_head = config.index_n_heads # 64 + self.head_dim = config.index_head_dim # 128 + self.rope_dim = config.qk_rope_head_dim # 64 + self.q_lora_rank = q_lora_rank # 1536 + self.compress_ratio = compress_ratio + self.use_fp4_kv = self.vllm_config.attention_config.use_fp4_indexer_cache + logger.info_once( + "Using %s indexer cache for Lightning Indexer.", + "MXFP4" if self.use_fp4_kv else "FP8", + ) + + # no tensor parallel, just replicated + self.wq_b = ReplicatedLinear( + self.q_lora_rank, + self.head_dim * self.n_head, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.wq_b", + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.n_head, + bias=False, + quant_config=None, + prefix=f"{prefix}.weights_proj", + ) + self.k_norm = LayerNorm(self.head_dim, eps=1e-6) + self.softmax_scale = self.head_dim**-0.5 + + self.scale_fmt = "ue8m0" + self.quant_block_size = 128 # TODO: get from config + self.topk_indices_buffer = topk_indices_buffer + + self.max_model_len = ( + vllm_config.model_config.max_model_len // self.compress_ratio + ) + self.prefix = prefix + + self.max_total_seq_len = ( + get_max_prefill_buffer_size(vllm_config) // self.compress_ratio + ) + + assert cache_config is not None, "Deepseek V4 indexer requires cache_config" + # NOTE(yifan): FP8 indxer cache use the same layout as V3.2: + # head_dim bytes = 128 fp8 + 4 fp32 scale = 132. + # For FP4 indexer cache, we still allocate the same amount of memory as FP8, + # but only use the first half of the memory. + k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4 + self.k_cache = DeepseekV4IndexerCache( + head_dim=k_cache_head_dim, + dtype=torch.uint8, + prefix=f"{prefix}.k_cache", + cache_config=cache_config, + compress_ratio=self.compress_ratio, + ) + self.compressor = DeepseekCompressor( + vllm_config=vllm_config, + compress_ratio=self.compress_ratio, + hidden_size=hidden_size, + head_dim=self.head_dim, + rotate=True, + prefix=f"{prefix}.compressor", + k_cache_prefix=self.k_cache.prefix, + use_fp4_cache=self.use_fp4_kv, + ) + + self.indexer_op = SparseAttnIndexer( + self.k_cache, + self.quant_block_size, + self.scale_fmt, + self.topk_tokens, + self.head_dim, + self.max_model_len, + self.max_total_seq_len, + self.topk_indices_buffer, + skip_k_cache_insert=True, + use_fp4_cache=self.use_fp4_kv, + ) + + def forward( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + compressed_kv_score: torch.Tensor, + indexer_weights: torch.Tensor, + positions: torch.Tensor, + rotary_emb: nn.Module, + ) -> torch.Tensor: + # ReplicatedLinear returns (output, bias); bias is None. + q, _ = self.wq_b(qr) + q = q.view(-1, self.n_head, self.head_dim) + k = self.compressor(compressed_kv_score, positions, rotary_emb) + q_quant, weights = fused_indexer_q_rope_quant( + positions, + q, + rotary_emb.cos_sin_cache, + indexer_weights, + self.softmax_scale, + self.n_head**-0.5, + use_fp4=self.use_fp4_kv, + ) + return self.indexer_op(hidden_states, q_quant, k, weights)