diff --git a/CURRENT_BUG.md b/CURRENT_BUG.md index 7045a7c7..9452582a 100644 --- a/CURRENT_BUG.md +++ b/CURRENT_BUG.md @@ -1,41 +1,116 @@ # CURRENT_BUG.md -## Status: CSA/HCA kernel works. Need vLLM integration. +## Status: vLLM server starts but returns empty output (immediate EOS) -### What We Know -- **CuTeDSL NVFP4 kernels**: All pass (cosine 0.988-0.999 vs BF16) -- **Warmup gs**: Irrelevant (runner recomputes per-call) -- **CSA attention kernel** (`cutedsl/csa_attention.py`): Works with PyTorch SDPA -- **Full layer 0 forward**: CuTeDSL + SDPA = cosine 0.988 vs BF16 ✅ -- **Logits**: std=2.98, reasonable top-5 tokens ✅ +### Root Cause -### Root Cause of vLLM Empty Output -vLLM uses two compiled CUDA kernels that DON'T work on Blackwell (SM100): -1. `torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert` — fused RoPE + KV cache -2. `FlashMLA sparse attention` — the actual attention computation +vLLM's compiled CUDA kernels don't work on Blackwell (SM100): +1. **`torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_kv_rope_quant_insert`** — fused RoPE + FP8 KV cache write (C++ kernel) +2. **FlashMLA sparse attention** — the attention computation (CUDA kernel) -The model uses **CSA (Compressed Sparse Attention) + HCA (Heavily Compressed Attention)**, NOT MLA. -vLLM misnames it "MLA" in code but the architecture is CSA/HCA with mHC. +Both crash or produce garbage on SM100. The model outputs immediate EOS → empty chat completions. -### Integration Plan -Replace vLLM's broken CUDA kernels in `DeepseekV4MLAAttention.forward`: -1. Replace `fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert` → pure PyTorch RoPE + FP8 quant + cache insert -2. Replace FlashMLA → our CSA/HCA kernel using PyTorch SDPA -3. Keep the compressor (it's mostly Triton which may work on SM100) -4. Keep the indexer (it calls into sparse_attn_indexer which is also Triton) +### What Works (verified with standalone tests on B200) -### Test Results +**CuTeDSL NVFP4 kernels — ALL PASS (cosine 0.989–0.999 vs BF16):** ``` -test_full_layer_b200.py: - q_a_proj: 0.995 ✅ kv_proj: 0.995 ✅ q_b_proj: 0.995 ✅ - wo_b_proj: 0.995 ✅ comp.kv_proj: 0.994 ✅ comp.gate: 0.995 ✅ - shared_expert: 0.990 ✅ - -test_model_forward_b200.py: - Warmup gs is IRRELEVANT (10x change → cosine 0.9993) - CuTeDSL cosine vs BF16: 0.999 - -test_csa_attention_b200.py: - Full path CuTeDSL + SDPA vs BF16: 0.988 ✅ - Logit std: 2.98 ✅ +q_a_proj: 0.995 ✅ kv_proj: 0.995 ✅ q_b_proj: 0.995 ✅ +wo_b_proj: 0.995 ✅ comp.kv_proj: 0.994 ✅ comp.gate: 0.995 ✅ +shared_expert: 0.990 ✅ ``` + +**Full attention path with SDPA — cosine 0.988 vs BF16, logit std 2.98 ✅** + +**Warmup gs is IRRELEVANT** — CuTeDSL runner recomputes activation global scale per-call internally. Changing it 10x has zero effect on output (cosine 0.9993). + +### Current Workaround (TEMPORARY — needs replacement) + +`_attention_impl_blackwell()` in `vllm/patches/deepseek_v4_attention.py`: +- Replaces FlashMLA with `full_sdpa_attention()` — **pure PyTorch matmuls, NO CuTeDSL** +- Replaces C++ fused kernel with `fused_qnorm_rope_kv_insert_py()` — pure PyTorch RoPE +- Skips SWA KV cache write (cache uses fp8_ds_mla packed format, shape [slot, 37376] not [slot, 512]) + +### THE PLAN: Replace all pure PyTorch with CuTeDSL/NVFP4 + +**Mike's directive: NO "optimize later" BS. Build the full thing with NVFP4/CuTeDSL.** + +The attention Q×K and attn×V are activation×activation matmuls. They CAN be done in NVFP4: +- Quantize Q and K to NVFP4 (4-bit activation + block scales) +- Use CuTeDSL grouped GEMM for the sparse attention pattern +- This is exactly what FlashMLA does with FP8 — we just use NVFP4 instead + +#### Specific replacements needed: + +1. **Attention (Q×K, attn×V)** → CuTeDSL NVFP4 GEMM + - Quantize Q and K to NVFP4 per-head + - Use `CuTeDSLNvfp4Linear` or raw CuTeDSL GEMM for the matmuls + - Support both prefill (batched) and decode (single-token) paths + - Handle CSA sparse gather pattern (only attend to top-k positions) + +2. **KV cache write** → NVFP4 quant + paged cache insert + - The SWA cache uses `fp8_ds_mla` format: packed FP8 values + UE8M0 scales + - Row width = 37376 bytes (not just head_dim=512) + - Layout: [nope_dim FP8 values | rope_dim FP8 values | scale blocks] + - Need to understand the exact fp8_ds_mla layout and replicate in CuTeDSL + - OR: skip SWA cache entirely and use our own NVFP4 cache format + +3. **RoPE** → CuTeDSL fused kernel + - Currently pure PyTorch (works, but slow) + - Could fuse with Q norm + KV quant into a single CuTeDSL kernel + - Pattern: Q norm → RoPE → NVFP4 quant → all in one pass + +4. **CSA sparse gather** → CuTeDSL indexed access + - Currently uses `torch.gather` (slow, not GPU-optimal) + - CuTeDSL can do the gather + GEMM in one fused operation + - This is the whole point of CSA — sparse KV access + +5. **Compressor** → already Triton (works on SM100) ✅ +6. **Indexer** → already Triton (works on SM100) ✅ +7. **MHC** → already pure PyTorch ✅ +8. **MoE** → already CuTeDSL ✅ + +### Config Issues (from config.json) + +- `quant_method: modelopt` → vLLM uses ModelOpt's NVFP4 handler +- Our CuTeDSL IS registered in `_POSSIBLE_NVFP4_KERNELS` (via register_cutedsl_kernel.py) +- Added `VLLM_NVFP4_GEMM_BACKEND=cutedsl` env var to force it +- `kv_cache_scheme: {"num_bits": 8, "type": "float"}` → FP8 KV cache → FlashMLA +- Hard assertion `issubclass(get_attn_backend(), FlashMLASparseBackend)` — patched with `_is_blackwell` flag + +### Checkpoint Key Names (different from vLLM 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.position_bias +sinks (attn_sink) +``` + +### Compress Ratios (from config.json compress_ratios) + +``` +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-only) +``` + +### Architecture: CSA + HCA + mHC (NOT MLA!) + +- **CSA (Compress Ratio 4)**: Compressed Sparse Attention — KV compressed 4x with overlap (coff=2) +- **HCA (Compress Ratio 128)**: Heavily Compressed Attention — KV compressed 128x +- **mHC**: Manifold-Constrained Hyper-Connections — replaces standard residual connections +- **SWA**: Sliding Window Attention (compress_ratio=0, last layer only) + +### Files + +- **Kernel**: `cutedsl/csa_attention.py` — CSA/HCA attention (currently SDPA, needs CuTeDSL) +- **vLLM patch**: `vllm/patches/deepseek_v4_attention.py` — `_attention_impl_blackwell()` +- **vLLM patch**: `vllm/patches/layers/csa_attention.py` — `fused_qnorm_rope_kv_insert_py()`, `full_sdpa_attention()` +- **Standalone tests**: `tests/test_full_layer_b200.py`, `tests/test_csa_attention_b200.py`, `tests/test_model_forward_b200.py` +- **CuTeDSL NVFP4 linear**: `cutedsl/nvfp4_linear.py` — CuTeDSLNvfp4Linear runner +- **CuTeDSL bridge**: `cutedsl/bridge.py` — quantize_activation_nvfp4, NVFP4 GEMM wrappers diff --git a/README.md b/README.md index 924fcfa3..da07c94c 100644 --- a/README.md +++ b/README.md @@ -1,12 +1,12 @@ # NVFP4 MegaMoE Kernel -Full NVFP4 inference pipeline for DeepSeek-V4 on NVIDIA Blackwell (SM100). The entire model — MoE experts, shared experts, and attention projections — runs in native NVFP4 with zero dequantization overhead. +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 (our work): +**MoE Experts** — CuTeDSL ScaledGroupedGemmKernel ✅: ``` BF16 input → quantize to NVFP4 L1 GEMM: NVFP4 × NVFP4 → BF16 (gate + up) @@ -16,156 +16,121 @@ BF16 input → quantize to NVFP4 Scatter with routing weights → BF16 output ``` -**Attention Projections** — CuTeDSL NVFP4 GEMM (our work, in progress): -- `wq_b`, `wo_b`, `fused_wqa_wkv` — native NVFP4, no conversion -- `wo_a` — NVFP4→FP8 for `fp8_einsum` (only attention weight that needs conversion) -- Compressor — BF16 (weight_loader stacking issue, small matmul) +**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 (our work, in progress): -- `gate_up_proj`, `down_proj` — native NVFP4 +**Shared Experts** — CuTeDSL NVFP4 GEMM ✅: +- `gate_up_proj`, `down_proj` — native NVFP4, cosine 0.990 vs BF16 -Both GEMM types use `float4_e2m1fn_x2` for weights, `float8_e4m3fn` for block scales, `float32` for global scales. BF16 is used only for SiLU activation, the final MoE scatter, and the compressor — the minimum possible. +**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 -## Current Status: Building Our Own Kernels 🔧 +**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** -**vLLM's built-in `FlashInferCutlassNvFp4LinearKernel` is broken on B200.** The same class of C++ CUTLASS FP4 bugs we hit with MoE (documented in "How We Got Here") affects the attention and shared expert paths. After a full day of debugging (broken `input_scale`, `process_weights_after_loading` timing, forward hook failures, BF16 dequant workarounds), we're replacing ALL vLLM NVFP4 kernels with our own CuTeDSL implementations. +## Architecture: DeepSeek-V4-Pro -**Test Results:** -- `tests/layertest.py`: cosine 0.988 vs BF16 reference ✅ -- `tests/cudagraph_test.py`: capture + replay PASS ✅ -- vLLM inference: produces empty/garbage output — vLLM's pipeline is broken, not our kernel +**CSA + HCA + mHC** (NOT MLA — vLLM misnames it "MLA" in code): -**What works:** -- MoE expert CuTeDSL kernel — production-ready, cosine 0.988, cudagraph-safe -- All NVFP4 weight dequantization — valid BF16 output confirmed in standalone tests +- **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) -**What's in progress:** -- Shared expert CuTeDSL kernel — runner WIP, scale assembly for num_groups=1 -- Attention projection CuTeDSL kernel — planned after shared experts - -**Why we're building our own:** vLLM's `FlashInferCutlassNvFp4LinearKernel` uses the same C++ CUTLASS FP4 path that was broken for MoE. The CuTeDSL approach (Python-based CUTLASS via MLML→PTX) is what NVIDIA's CUTLASS team recommends for Blackwell. Our MoE kernel proves it works. Time to apply the same approach to the rest of the model. - -**vLLM serves DeepSeek-V4-Pro NVFP4 with cudagraph enabled.** The model loads, cudagraph captures successfully, and inference runs. Output quality is still being tuned (garbage tokens currently), but this is the first time the entire pipeline — model loading, kernel compilation, cudagraph capture, and inference — works end-to-end. - -**Test Results:** -- `tests/layertest.py`: cosine 0.988 vs BF16 reference ✅ -- `tests/cudagraph_test.py`: capture + replay PASS ✅ -- vLLM inference: running with cudagraph, output quality in progress - -## How We Got Here - -### The C++ CUTLASS Kernel Was Broken - -The original kernel was a C++ `.cu` file using CUTLASS's C++ API directly. It passed all the simple tests (uniform data → exact output, SF remap verifier → 0 errors) but produced **cosine 0.05** with real random data. After weeks of debugging the SF remap (8+ iterations, all producing the same 0.2 cosine against a wrong reference), we discovered: - -1. **The BF16 reference comparison was wrong** — our Python dequantization didn't match CUTLASS's internal FP4 handling. A wrong reference is worse than no reference. We chased ghosts through 8+ SF remap rewrites because the 0.2 cosine was never about the remap. - -2. **The C++ CUTLASS kernel misinterpreted FP4 data** — even with SF remap verified correct (0 byte errors), the GEMM produced garbage with non-uniform data. The issue was in how CUTLASS's C++ API handles FP4 packing/tiling internally — something we couldn't easily debug or fix. - -3. **The checkpoint `input_scale` was a red herring** — we tried using the checkpoint's calibration scale as the activation normalization scale. It saturated all block scales to 448.0 (max float8). The `input_scale` is a calibration constant for alpha computation, not a normalization scale. - -### The CuTeDSL Kernel Works - -NVIDIA's CuTeDSL approach (Python-based CUTLASS kernels compiled via MLML → PTX) is what the CUTLASS team recommends for Blackwell. Their official MoE scaled grouped GEMM example (`torch_scaled_grouped_mm.py`) supports NVFP4 out of the box. We adapted it. - -**Results with real DeepSeek-V4 layer 0 weights:** -- L1 GEMM alone: cosine 0.995 -- Full MoE pipeline (L1→SiLU→L2→scatter): cosine 0.989 -- Weight loading: **0% loss** — direct uint8→float4_e2m1fn_x2 view-cast, bit-identical to checkpoint -- Activation quantization: ~1.1% cosine loss (dynamic BF16→NVFP4 — inherent to the format, unavoidable) -- GEMM kernel: 0% loss (CuTeDSL is correct) - -The 0.989 cosine is entirely from activation quantization. The weights are bit-identical to the checkpoint — no BF16 round-trip, no precision loss. - -### The Dequant→Requant Anti-Pattern - -Early versions dequantized all NVFP4 weights to BF16, then let vLLM's `FlashInferCutlassNvFp4LinearKernel` requantize them back to NVFP4 at inference time. This: -- Wasted 5 minutes on load doing NVFP4→BF16 conversion -- Lost precision on the double round-trip -- Caused vLLM to hang — the NVFP4 attention kernel expects native NVFP4 weights, not BF16 weights with an NVFP4 quant_method attached - -The fix: **keep everything in NVFP4**. The checkpoint stores NVFP4. The kernels consume NVFP4. No conversion needed. - -### CUDAGraph Compatibility - -vLLM uses CUDA graphs to eliminate kernel launch overhead in the decode path. CUDA graphs record the entire forward pass once, then replay it — but they require **fixed tensor shapes, fixed memory addresses, and zero CPU-GPU syncs**. - -Our original runner was not cudagraph-safe. We had to fix several classes of issues: - -#### 1. CPU↔CUDA Tensor Copies - -`torch.tensor([0,1,...], device=x.device)` creates the tensor on **CPU first**, then copies to CUDA. This copy is forbidden during graph capture. The fix: **cache tensors per device** on first use, outside the graph. - -```python -# BAD — CPU→CUDA copy inside graph -step_to_idx = torch.tensor([0,1,2,3,4,4,5,5,6,6,6,7,7], device=x.device) - -# GOOD — cached on first use, reused in graph -step_to_idx = _get_step_to_idx_lut(x.device) # returns cached CUDA tensor +**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) ``` -Similarly, `torch.zeros` and `torch.rand` don't support `float4_e2m1fn_x2` or `float8_e4m3fn` dtypes. The fix: create as `uint8` or `float16`, then `.view()` or `.to()` the target dtype. - -#### 2. GPU Scalar Slicing - -`buf[:gpu_scalar, :]` requires the runtime to query the GPU scalar's value to determine the output shape. This triggers an implicit CPU-GPU sync, which invalidates the graph. The fix: **always use full pre-allocated buffers**. Extra rows are zeros that contribute nothing to the computation. - -```python -# BAD — GPU scalar as slice index (implicit sync) -total_padded_rows = padded_expert_offsets[-1] # GPU scalar -padded_scales = buf[:total_padded_rows, :padded_cols] # sync! - -# GOOD — full pre-allocated buffer, zero out before use -padded_scales = self._padded_scales_buf # always max size -padded_scales.zero_() +**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) ``` -**Design decision:** Padding to max size wastes a few rows of compute on zero data, but: -- The extra rows are zeros → zero GEMM output → no accuracy impact -- GEMMs are memory-bandwidth bound → multiplying zeros is nearly free -- VRAM cost is negligible (~350KB for activation intermediates across all MoE layers) -- vLLM already does this everywhere (attention, FFN, etc.) +## Current Status: Attention + KV Cache Need CuTeDSL 🔧 -#### 3. Kernel Compilation in the Forward Path +**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) ✅ -`cute.compile()` is a host-side JIT operation that generates PTX and compiles a CUDA kernel. It cannot be called inside cudagraph capture. The fix: **compile once during warmup, cache the compiled kernel, then only invoke `compiled()` on subsequent calls**. +**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 -The compiled kernel uses `separate_tensormap_init=True`, which handles TMA descriptor re-initialization for new tensor data. We create new `mark_layout_dynamic` CuTe tensor views for each forward call, and the compiled kernel accepts them. +**What needs to be built:** -**Critical lesson:** Caching the compiled kernel across different tensor allocations initially produced wrong results (cosine 0.5 instead of 0.99). The issue was NOT that caching is fundamentally broken — it was that our bridge had other bugs (wrong `make_b_k_major` stride check, `quantize_weight_to_nvfp4` packing N instead of K). Once those were fixed, cached compilation works correctly. +### 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 -#### 4. Weight Quantization: K is the Packed Dimension +### 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) -`quantize_weight_to_nvfp4` packs K (dim 0) differently from `quantize_to_nvfp4` which packs the last dimension. For a weight matrix `(K, N)`: -- K=7168 is the packed dimension (7168 → 3584 in float4) -- N=6144 stays as-is -- Block scales are computed along K blocks: `(K//16, N)` not `(K//2, N//16)` -- The nibble packing uses `[:, ::2, :]` and `[:, 1::2, :]` (along the K block dim) +### 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 -Confusing the two quantization functions produces wrong tensor shapes that crash or produce garbage. +### 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 -#### 5. B Tensor K-Major Layout +## vLLM Integration -The CuTeDSL kernel expects B tensors in K-major memory layout (K elements contiguous in memory). `torch.stack` produces N-major layout. The fix: **double-permute trick** — transpose, make contiguous, transpose back. Same shape, different strides. +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) -```python -# Double-permute: (E,K,N) → (E,N,K) → contiguous → (E,K,N) -# Same shape, but K-contiguous memory layout -return b_tensor.permute(0, 2, 1).contiguous().permute(0, 2, 1) -``` +**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) -A single permute changes the tensor SHAPE (swapping K and N), which breaks everything downstream. +**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. -### Key Lessons +## Test Files -1. **A wrong reference is worse than no reference** — the 0.2 cosine against a broken BF16 dequant sent us chasing SF remap bugs for weeks -2. **The C++ CUTLASS API is a footgun for FP4** — CuTeDSL handles tensor layouts, tiling, and SF construction correctly by construction -3. **Test with real data early** — uniform tests pass even with broken kernels; random data reveals real bugs -4. **Separate the GEMM from the pipeline** — our `layertest.py` runs without vLLM, Docker, or tensor parallelism. It caught the kernel bug that vLLM's integration layers masked. -5. **Don't dequant what's already quantized** — if the kernel expects NVFP4 and the checkpoint is NVFP4, leave it alone. No BF16 round-trips. -6. **GPU scalar slicing is a silent cudagraph killer** — no error, no warning, just `cudaErrorStreamCaptureInvalidated` with no pointer to the cause. The test harness catches it. -7. **Weight vs activation quantization are different** — K-packed (weights) vs last-dim-packed (activations). Mixing them up produces wrong shapes and garbage output. -8. **Double-permute for memory layout changes** — single permute changes shape, double-permute changes layout. The kernel cares about both. +| 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 @@ -173,128 +138,53 @@ A single permute changes the tensor SHAPE (swapping K and N), which breaks every 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) -│ └── kernel/moe/ # NVIDIA's ScaledGroupedGemmKernel (untouched) -│ ├── torch_scaled_grouped_mm.py # The working kernel (3900 lines) -│ ├── moe_utils.py -│ moe_persistent_scheduler.py -│ └── moe_sched_extension.py +│ ├── 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 interface +│ ├── 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 # DeepSeek-V4 model patch (NVFP4 native) -│ └── deepseek_v4_attention.py # Attention patch (NVFP4 native) -├── tests/ -│ ├── cudagraph_test.py # CUDAGraph compatibility test (✅ PASS) -│ ├── layertest.py # Layer 0 comparison: CuTeDSL vs BF16 (✅ cosine 0.988) -│ ├── test_cutedsl.py # Small standalone CuTeDSL test (✅ cosine 0.991) -│ ├── test_uniform_fp4.py # Uniform data GEMM test -│ ├── test_b_layout.py # B matrix column layout test -│ └── test_quick_rand.py # Quick random GEMM sanity check -└── reference/ # Reference files for study +│ ├── 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 ``` -## The Bridge Layer (`cutedsl/bridge.py`) - -Handles all tensor layout conversion from our pipeline to what the CuTeDSL kernel expects: - -| Function | What it does | -|----------|--------------| -| `quantize_to_nvfp4()` | BF16 → float4_e2m1fn_x2 + float8_e4m3fn block scales (NOT cudagraph-safe — uses .max()) | -| `quantize_activation_nvfp4()` | Same, but cudagraph-safe (fixed global_scale, no .max()) | -| `quantize_weight_to_nvfp4()` | Same, but K is the packed dimension (different block scale shape) | -| `assemble_scales_2d_side()` | Pad and swizzle activation scale factors (2Dx3D A side) | -| `assemble_scales_3d_side()` | Pad and swizzle weight scale factors (2Dx3D B side) | -| `make_b_k_major()` | Convert B tensor from N-major to K-major strides (double-permute trick) | -| `run_nvfp4_grouped_gemm()` | Kernel launch with cached compilation (cudagraph-safe) | - -## Running Tests - -On the B200 (host venv, no container): - -```bash -cd /root/nvfp4-megamoe-kernel -source tests/.venv/bin/activate -export CUDA_TOOLKIT_PATH=/usr/local/cuda - -# CUDAGraph compatibility test -python3 tests/cudagraph_test.py - -# Small standalone test -python3 tests/test_cutedsl.py - -# Full layer 0 comparison with real weights -python3 tests/layertest.py -``` - -## NVFP4 Coverage - -| Component | Format | Kernel | Status | -|-----------|--------|--------|--------| -| MoE experts (L1+L2) | NVFP4 native | CuTeDSL ScaledGroupedGemm | ✅ Working, cosine 0.988 | -| Shared experts | NVFP4 native | CuTeDSL GEMM (1 group) | 🔧 In progress | -| wq_b, wo_b, fused_wqa_wkv | NVFP4 native | CuTeDSL GEMM | 📋 Planned | -| wo_a | NVFP4 → FP8 | fp8_einsum | 📋 May stay FP8 or go native NVFP4 | -| Compressor | NVFP4 → BF16 | torch.mm | ✅ Done (weight_loader stacking issue) | -| KV cache | FP8 | FlashInfer MLA | ✅ Works | - ## Plan -### Phase 1: Kernel ✅ DONE -- CuTeDSL ScaledGroupedGemmKernel works with NVFP4 -- Bridge layer handles all tensor layout conversion -- Full MoE pipeline (L1→SiLU→L2→scatter) produces cosine 0.988 vs BF16 +### Phase 1: MoE Kernel ✅ DONE +- CuTeDSL ScaledGroupedGemmKernel with NVFP4 +- Full pipeline: cosine 0.988, cudagraph-safe -### Phase 2: vLLM Integration ✅ DONE -- CuTeDSLMoERunner wires CuTeDSL kernel into vLLM -- Weight loading: checkpoint uint8 → float4_e2m1fn_x2 view-cast (bit-preserving) -- Block scales (float8_e4m3fn) and global scales (float32) pass through directly -- L1 dual global scale handling: normalize to max(gate_gs, up_gs), fold ratio into block scales -- Attention projections stay native NVFP4 (FlashInferCutlassNvFp4LinearKernel) -- CuTeDSL kernel warmup during model load (prevents RPC timeout) -- Removed all debug prints and env var gates from vLLM serving path +### Phase 2: NVFP4 Linear Kernels ✅ DONE +- All attention projections: cosine 0.995 +- Shared experts: cosine 0.990 +- Compressor projections: cosine 0.995 -### Phase 2.5: CUDAGraph Compatibility ✅ DONE -- CuTeDSLMoERunner is fully cudagraph-safe -- Zero CPU-GPU syncs, zero dynamic shapes, zero GPU scalar slicing -- All intermediate buffers pre-allocated at max_num_tokens * top_k -- `quantize_activation_nvfp4` uses cached LUT (no CPU→CUDA copy) -- `torch.zeros/rand` for float4/float8 → uint8→view or float16→cast -- Test harness validates capture + replay -- VRAM overhead: ~350KB (negligible) -- Compute overhead: zero rows through GEMM on padding (memory-bound, free) -- Kernel compilation cached: `cute.compile()` once during warmup, `compiled()` on forward calls +### 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 3: Output Quality 🔧 IN PROGRESS -- vLLM serves the model with cudagraph, but output is garbage tokens -- Layer 0 cosine is 0.988 in isolation, so the GEMM math is correct -- Root cause: vLLM's `FlashInferCutlassNvFp4LinearKernel` is broken on B200 - - Same class of C++ CUTLASS FP4 bugs we hit with MoE - - `input_scale` from checkpoint causes NaN during activation quantization - - BF16 dequant workaround doesn't fix the underlying pipeline issues -- Solution: Replace ALL vLLM NVFP4 kernels with our own CuTeDSL implementations - -### Phase 3.5: Shared Expert CuTeDSL Kernel 🔧 IN PROGRESS -- Replacing `FlashInferCutlassNvFp4LinearKernel` with CuTeDSL GEMM -- Shared expert = MoE with 1 expert, no routing -- `cutedsl/shared_expert_pipeline.py` — dedicated runner (scale assembly needs fixing) -- `tests/test_shared_expert.py` — standalone test (ready) -- Target: cosine ≥ 0.98 vs BF16 reference - -### Phase 3.6: Attention CuTeDSL Kernel 📋 PLANNED -- Replace attention NVFP4 path with CuTeDSL GEMMs -- Each projection (fused_wqa_wkv, wq_b, wo_a, wo_b) = standard NVFP4 GEMM -- `fused_wqa_wkv` has dual weight_scale_2 (same pattern as MoE gate+up) -- Test each projection individually, then integrate - -### Phase 4: Optimization -- Replace wo_a FP8 conversion with native NVFP4 GEMM (eliminate last dequant) -- Fix compressor weight_loader so it stays NVFP4 native -- Explore larger tile sizes for better occupancy -- Profile end-to-end inference on full model -- Proper TMA descriptor management for kernel caching across different tensor pools +### 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 -- Clean up old C++ kernel code (tagged `the-last-of-cutlass`) -- Add proper error handling and logging -- Benchmark vs BF16 baseline +- End-to-end benchmarking +- Optimize tile sizes for occupancy +- Clean up old C++ kernel code