nvfp4-megamoe-kernel/README.md

# NVFP4 MegaMoE Kernel

Full NVFP4 inference pipeline for DeepSeek-V4 on NVIDIA Blackwell (SM100). The entire model — MoE experts, shared experts, attention projections, and attention compute — runs in native NVFP4 with zero dequantization overhead.

## What This Is

A native NVFP4 inference stack for DeepSeek-V4:

**MoE Experts** — CuTeDSL ScaledGroupedGemmKernel ✅:
```
BF16 input → quantize to NVFP4
  L1 GEMM: NVFP4 × NVFP4 → BF16 (gate + up)
  SiLU(gate) * up → BF16 (only nonlinear — can't avoid BF16 here)
  Re-quantize → NVFP4
  L2 GEMM: NVFP4 × NVFP4 → BF16 (down_proj)
  Scatter with routing weights → BF16 output
```

**Attention Projections** — CuTeDSL NVFP4 GEMM ✅:
- `q_a_proj`, `q_b_proj`, `kv_proj`, `wo_b_proj` — native NVFP4, cosine 0.995 vs BF16
- `wo_a` — BF16 BMM (o_a_proj weights are BF16 in checkpoint)
- `compressor.kv_proj`, `compressor.gate_proj` — native NVFP4, cosine 0.995 vs BF16
- All verified with `tests/test_full_layer_b200.py`

**Shared Experts** — CuTeDSL NVFP4 GEMM ✅:
- `gate_up_proj`, `down_proj` — native NVFP4, cosine 0.990 vs BF16

**Attention Compute** — **NEEDS CuTeDSL NVFP4** 🔧:
- Currently using pure PyTorch SDPA as a TEMPORARY workaround
- Q×K and attn×V are activation×activation matmuls that CAN be NVFP4
- FlashMLA (vLLM's CUDA kernel) is broken on Blackwell
- **Plan: CuTeDSL NVFP4 attention kernel** — quantize Q/K to NVFP4, use CuTeDSL GEMMs

**KV Cache Write** — **NEEDS CuTeDSL NVFP4** 🔧:
- The SWA KV cache uses `fp8_ds_mla` packed format (37376 bytes per slot, not 512)
- C++ kernel `fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert` is broken on Blackwell
- Currently skipped in Blackwell path (works for prefill, breaks decode)
- **Plan: NVFP4 quant + paged cache insert in CuTeDSL**

## Architecture: DeepSeek-V4-Pro

**CSA + HCA + mHC** (NOT MLA — vLLM misnames it "MLA" in code):

- **CSA (Compress Ratio 4)**: Compressed Sparse Attention — KV compressed 4x with overlap (coff=2). Indexer finds per-layer top-k.
- **HCA (Compress Ratio 128)**: Heavily Compressed Attention — KV compressed 128x. Top-k indices pre-computed during metadata build.
- **mHC**: Manifold-Constrained Hyper-Connections — replaces standard residual connections. Learned mixing with Sinkhorn normalization.
- **SWA**: Sliding Window Attention — local window (compress_ratio=0, last layer only)

**Compress Ratios (from config.json):**
```
Layer 0: 128 (HCA)   Layer 1: 128 (HCA)   Layer 2: 4 (CSA)   Layer 3: 128 (HCA)
Layer 4: 4 (CSA)     ...alternating 4/128...                  Layer 60: 0 (SWA)
```

**Checkpoint Key Names** (different from vLLM's internal names):
```
q_a_proj, q_b_proj, kv_proj       (NOT fused_wqa_wkv, wq_b)
q_a_norm                           (NOT q_norm)
attn_hc.fn/base/scale              (MHC attention)
ffn_hc.fn/base/scale               (MHC FFN)
compressor.kv_proj, compressor.gate_proj  (CSA/HCA compressor)
compressor.position_bias
sinks                              (attn_sink)
```

## Current Status: Attention + KV Cache Need CuTeDSL 🔧

**What works (verified on B200):**
- CuTeDSL NVFP4 linear kernels: cosine 0.989–0.999 vs BF16 ✅
- CuTeDSL NVFP4 MoE: cosine 0.988 ✅
- Full attention path with PyTorch SDPA: cosine 0.988 vs BF16 ✅
- MHC, RMS norm, RoPE (BF16), wo_a BMM, shared experts ✅
- Compressor + indexer (Triton, works on SM100) ✅

**What's broken:**
- FlashMLA CUDA kernel → garbage on Blackwell → model outputs immediate EOS
- `fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert` C++ kernel → crashes on Blackwell
- Pure PyTorch SDPA is a TEMPORARY workaround — must replace with CuTeDSL NVFP4

**What needs to be built:**

### 1. CuTeDSL NVFP4 Attention Kernel
- Quantize Q and K to NVFP4 per-head
- Use CuTeDSL GEMM for Q×K and attn×V
- Support prefill (batched) and decode (single-token) paths
- Handle CSA sparse gather (attend to top-k positions only)
- This is exactly what FlashMLA does with FP8 — we just use NVFP4 instead
- **Test first**: build standalone test in `tests/` with real weights

### 2. CuTeDSL NVFP4 KV Cache Insert
- Replace C++ `fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert`
- Per-head RMS norm on Q + GPT-J RoPE on Q + RoPE on KV + NVFP4 quant + paged cache write
- The SWA cache uses `fp8_ds_mla` packed format: row width = 37376 bytes
  - Layout: [nope_dim FP8 values | rope_dim FP8 values | UE8M0 scale blocks]
  - NOT just [head_dim] — it's a packed FP8 format with interleaved scales
- **Option A**: Understand and write the fp8_ds_mla format from CuTeDSL
- **Option B**: Use our own NVFP4 cache format (simpler, more efficient, but diverges from vLLM)

### 3. CuTeDSL Fused RoPE + Norm Kernel
- Currently pure PyTorch (works, but slow)
- Fuse: Q norm → RoPE → NVFP4 quant → all in one pass
- Same for KV side: RoPE → NVFP4 quant → cache write

### 4. CuTeDSL CSA Sparse Gather
- Currently `torch.gather` (slow, not GPU-optimal)
- CuTeDSL can do the gather + GEMM in one fused operation
- The whole point of CSA is sparse KV access — we should do it right

## vLLM Integration

The Blackwell detection and dispatch is in `vllm/patches/deepseek_v4_attention.py`:
- `attention_impl()` detects SM100+ → `_attention_impl_blackwell()`
- Currently uses pure PyTorch (SDPA) — **must replace with CuTeDSL**
- The dispatch is INSIDE the `torch.ops.vllm.deepseek_v4_attention` custom op boundary (important for torch.compile)

**Config issues:**
- `quant_method: modelopt` → vLLM uses ModelOpt's NVFP4 handler
- Our CuTeDSL IS registered (via `register_cutedsl_kernel.py`) and forced with `VLLM_NVFP4_GEMM_BACKEND=cutedsl`
- FlashMLA hard assertion in `DeepseekV4MLAAttention.__init__` — patched with `_is_blackwell` flag
- `kv_cache_scheme: {"num_bits": 8, "type": "float"}` → FP8 KV cache → FlashMLA (broken on Blackwell)

**Key discovery: warmup gs is irrelevant.** CuTeDSL runner recomputes activation global scale per-call internally. Changing it 10x has zero effect on output (cosine 0.9993). The `input_scale` from the checkpoint is NOT the activation global scale — it's a calibration constant.

## Test Files

| Test | What it does | Status |
|------|-------------|--------|
| `tests/test_full_layer_b200.py` | All NVFP4 projections vs BF16 (layer 0) | ✅ All pass (0.989–0.999) |
| `tests/test_model_forward_b200.py` | Warmup gs vs dynamic gs diagnostic | ✅ Warmup gs irrelevant |
| `tests/test_csa_attention_b200.py` | Full attention path with SDPA | ✅ cosine 0.988 |
| `tests/layertest.py` | MoE layer test | ✅ cosine 0.988 |
| `tests/cudagraph_test.py` | CUDAGraph compatibility | ✅ PASS |
| `tests/test_shared_expert.py` | Shared expert standalone | ✅ cosine 0.990 |

## Project Structure

```
nvfp4-megamoe-kernel/
├── cutedsl/                          # CuTeDSL kernel + bridge layer
│   ├── bridge.py                     # Tensor layout conversion, quantization, kernel launch
│   ├── nvfp4_linear.py              # CuTeDSLNvfp4Linear — NVFP4 GEMM runner
│   ├── moe_pipeline.py              # Full MoE pipeline (L1→SiLU→L2→scatter)
│   ├── shared_expert_pipeline.py    # Shared expert pipeline (1-expert MoE variant)
│   ├── csa_attention.py             # CSA/HCA attention (currently SDPA, needs CuTeDSL)
│   ├── custom_ops.py                # torch.autograd wrappers for compile boundary
│   └── kernel/moe/                   # NVIDIA's ScaledGroupedGemmKernel
├── vllm/                             # vLLM integration
│   ├── nvfp4_cutedsl.py             # CuTeDSLMoERunner — cudagraph-safe MoE kernel
│   ├── cutedsl_quant_method.py      # CuTeDSLNvfp4LinearMethod — vLLM quant method
│   ├── kernels/linear/nvfp4/cutedsl.py  # CuTeDSLNvFp4LinearKernel — vLLM kernel registration
│   └── patches/
│       ├── deepseek_v4.py           # Model patch (NVFP4 native, MHC, MoE)
│       ├── deepseek_v4_attention.py # Attention patch (Blackwell dispatch)
│       ├── layers/
│       │   ├── mhc.py               # MHC pure PyTorch (replaces TileLang)
│       │   ├── csa_attention.py     # CSA attention (TEMPORARY — needs CuTeDSL)
│       │   └── deepseek_compressor.py  # Compressor (Triton, works on SM100)
│       └── fused_moe/experts/cutedsl_moe.py  # MoE CuTeDSL integration
├── tests/                            # Standalone tests (run on B200 outside container)
└── Dockerfile                        # Container build
```

## Plan

### Phase 1: MoE Kernel ✅ DONE
- CuTeDSL ScaledGroupedGemmKernel with NVFP4
- Full pipeline: cosine 0.988, cudagraph-safe

### Phase 2: NVFP4 Linear Kernels ✅ DONE
- All attention projections: cosine 0.995
- Shared experts: cosine 0.990
- Compressor projections: cosine 0.995

### Phase 3: vLLM Integration ✅ DONE (with PyTorch fallback)
- CuTeDSL kernels registered and working for all NVFP4 linear layers
- Blackwell dispatch in attention_impl
- MHC pure PyTorch
- MoE CuTeDSL

### Phase 4: CuTeDSL NVFP4 Attention 🔧 NEXT
- Replace pure PyTorch SDPA with CuTeDSL NVFP4 GEMMs for Q×K and attn×V
- NVFP4 KV cache insert (replace C++ kernel)
- Fused RoPE + norm + quant kernel
- CSA sparse gather in CuTeDSL
- **Test each component standalone before integrating into vLLM**

### Phase 5: Production
- End-to-end benchmarking
- Optimize tile sizes for occupancy
- Clean up old C++ kernel code