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.

## ⚠️ READ THIS FIRST — THE #1 RULE

**YOU MUST BUILD YOUR OWN KERNELS. ALL OF THEM. DO NOT PATCH vLLM.**

Mike was right — we need our own kernels. Not just for the NVFP4 GEMMs, but for the **ENTIRE attention pipeline**. The current approach of patching individual vLLM functions is a house of cards. Every patch leads to another crash, every workaround reveals three more broken things. FlashMLA, fp8_ds_mla, the fused C++ kernels, the Triton compressor, the indexer — they're all deeply coupled. You cannot swap one piece and expect the rest to work.

**THE ONLY PATH FORWARD:**
1. Build CuTeDSL kernels for EVERYTHING — attention, KV cache, RoPE, the whole stack
2. Test each kernel standalone on the B200 venv BEFORE touching the container
3. Wire them together into a proper vLLM attention backend
4. THEN and ONLY THEN test in the container

**DO NOT:**
- ❌ Try to patch vLLM's FlashMLA code to "work" on Blackwell
- ❌ Use pure PyTorch as a "temporary workaround" — it produces garbage
- ❌ Skip the KV cache write and hope for the best
- ❌ Assume you can mix our kernels with vLLM's existing attention backend
- ❌ Touch the container until ALL kernels pass standalone tests

**DO:**
- ✅ Build CuTeDSL kernels in `cutedsl/`
- ✅ Test each one in `tests/` on the B200 venv
- ✅ Compare against BF16 reference (cosine >= 0.98 or it's broken)
- ✅ Wire them into a proper attention backend class
- ✅ Only test in the container once everything passes standalone

---

## 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 to 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)
- 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:
- Pure PyTorch SDPA produces garbage in the container
- FlashMLA is broken on Blackwell
- Must build CuTeDSL kernels for Q×K, attn×V, KV cache, RoPE

**KV Cache** — 🔧 NEEDS CuTeDSL:
- The fp8_ds_mla format is FlashMLA-specific (584 bytes per token)
- Must build our own NVFP4 KV cache with our own format

## 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)
```

## Current Status

### ✅ Working (verified on B200 standalone tests)

| Component | Test | Cosine vs BF16 |
|-----------|------|----------------|
| CuTeDSL NVFP4 Linear (q_a, kv, q_b, wo_b) | `test_full_layer_b200.py` | 0.994+ |
| CuTeDSL NVFP4 MoE (L1 gate+up, SiLU, L2 down) | `layertest.py` | 0.988 |
| FP8 KV quantize/dequant | `test_kv_cache_b200.py` | 0.9997 |
| NVFP4 KV quantize/dequant | `test_kv_cache_b200.py` | 0.9943 |
| Paged KV cache read/write | `test_kv_cache_b200.py` | 1.0 |
| FP8 KV → full attention | `test_kv_cache_b200.py` | 0.9997 |
| CSA sparse attention (cr=4) | `test_sparse_attn_b200.py` | works, no NaN |
| HCA sparse attention (cr=128) | `test_sparse_attn_b200.py` | works, no NaN |
| Merged CSA+SWA attention | `test_sparse_attn_b200.py` | works, no NaN |
| Full attention pipeline (all layer types) | `test_v4_attention_b200.py` | 0.981–0.995 |
| RoPE (GPT-J) | `test_v4_attention_b200.py` | works |
| Inverse RoPE + o_a BMM | `test_v4_attention_b200.py` | works |

### 🔧 Needs CuTeDSL Kernels

1. **Attention Q×K^T** — BF16 matmul works standalone, but NVFP4 GEMM too lossy (cosine 0.86). Keep Q×K in BF16.
2. **KV Cache Write** — need CuTeDSL kernel that does: RoPE → fp8 quant → paged cache insert
3. **KV Cache Read** — need CuTeDSL kernel that does: paged cache read → fp8 dequant
4. **Fused Q-norm + RoPE** — currently pure PyTorch (works, slow)
5. **Fused inverse RoPE + o_a BMM** — currently pure PyTorch (works)

### ❌ Does NOT Work

- **NVFP4 Q×K^T GEMM** — cosine 0.86, too lossy for attention scores. Keep attention in BF16.
- **Patching vLLM's FlashMLA path** — house of cards. Don't do it.
- **Pure PyTorch SDPA in the container** — produces garbage because the KV cache isn't written and the pipeline is broken.

## Container Status

The container builds and starts successfully. The server accepts requests and generates tokens. But the output is empty/garbage because the Blackwell attention path is broken. Multiple patches were applied to get this far (KV cache page sizes, FlashMLA alignment, softmax_scale, compressor cache), but the fundamental problem remains: **you cannot half-ass the attention pipeline**.

## Test Files

| Test | What it does | Status |
|------|-------------|--------|
| `tests/test_full_layer_b200.py` | All NVFP4 projections vs BF16 | ✅ 0.994+ |
| `tests/layertest.py` | MoE layer test | ✅ 0.988 |
| `tests/cudagraph_test.py` | CUDAGraph compatibility | ✅ PASS |
| `tests/test_csa_attention_b200.py` | Full attention with SDPA | ✅ 0.988 |
| `tests/test_v4_attention_b200.py` | All 3 layer types (SWA, C128A, C4A) | ✅ 0.981-0.995 |
| `tests/test_kv_cache_b200.py` | FP8/NVFP4 KV cache + paged cache | ✅ 0.9997 |
| `tests/test_sparse_attn_b200.py` | CSA/HCA sparse + SWA merged | ✅ works |
| `tests/test_nvfp4_attn_gemm_b200.py` | NVFP4 Q×K^T GEMM | ❌ 0.86 (too lossy) |

## 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
│   ├── csa_attention.py             # CSA/HCA attention (BF16 SDPA — needs CuTeDSL)
│   ├── custom_ops.py                # torch.autograd wrappers
│   └── kernel/moe/                   # NVIDIA's ScaledGroupedGemmKernel
├── vllm/                             # vLLM integration
│   ├── nvfp4_cutedsl.py             # CuTeDSLMoERunner
│   ├── cutedsl_quant_method.py      # CuTeDSLNvfp4LinearMethod
│   ├── kernels/linear/nvfp4/cutedsl.py  # vLLM kernel registration
│   └── patches/
│       ├── deepseek_v4_attention.py # Attention patch (Blackwell dispatch)
│       ├── patch_kv_cache_utils.py  # KV cache page size fix
│       ├── patch_swa_cache.py       # SWA cache alignment fix
│       ├── patch_indexer_cache.py   # Indexer cache alignment fix
│       ├── patch_compressor_cache.py # Compressor cache alignment fix
│       └── layers/
│           ├── csa_attention.py     # BF16 SDPA (TEMPORARY — needs CuTeDSL)
│           └── ...
├── tests/                            # Standalone tests (run on B200 venv)
└── Dockerfile                        # Container build
```

## Plan

### Phase 1: MoE Kernel ✅ DONE
### Phase 2: NVFP4 Linear Kernels ✅ DONE
### Phase 3: vLLM Integration ✅ DONE (NVFP4 linear + MoE working)

### Phase 4: CuTeDSL Attention Backend 🔧 NEXT — BUILD THE KERNELS

**STOP. READ THIS.**

Do NOT touch the vLLM container until ALL of these kernels pass standalone tests on the B200 venv. The container is a 14-minute build cycle. The venv gives you instant feedback. TEST FIRST.

**Kernels to build (in order):**

1. **KV Cache Write**: BF16 KV → apply RoPE → quantize to fp8 → write to paged cache
   - Test: compare against BF16 reference (cosine >= 0.98 after dequant)
   
2. **KV Cache Read**: paged cache → dequant fp8 → BF16 KV with RoPE
   - Test: write then read back, cosine >= 0.99

3. **BF16 Attention**: Q (with RoPE) × K^T → softmax → attn × V
   - Keep this in BF16 (NVFP4 is too lossy for attention scores)
   - Handle CSA sparse gather (attend to top-k indexed positions)
   - Handle HCA sparse gather (attend to 1/128 positions)
   - Handle SWA (sliding window, full causal within window)
   - Test: compare against PyTorch SDPA reference (cosine >= 0.99)

4. **Full Attention Pipeline**: KV cache read → attention → inverse RoPE → o_a BMM
   - Wire everything together
   - Test: compare against BF16 reference (cosine >= 0.98)

5. **vLLM Backend**: Wrap as a proper AttentionBackend subclass
   - Override `DeepseekSparseSWABackend` on Blackwell
   - Handle the metadata, slot mapping, cache format
   - ONLY THEN test in the container

### Phase 5: Production
- End-to-end benchmarking
- Optimize tile sizes
- Clean up