Update README.md and CURRENT_BUG.md: eliminate stale issues, document NaN investigation, clarify our kernels are clean
This commit is contained in:
117
README.md
117
README.md
@@ -52,18 +52,16 @@ BF16 input → quantize to NVFP4
|
||||
**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
|
||||
**Attention Pipeline** — ✅ Verified standalone, 🔧 vLLM integration blocked by NaN:
|
||||
- KV cache write (RoPE → fp8 quant → paged cache) — cosine 0.999
|
||||
- KV cache read (paged cache → fp8 dequant → BF16) — cosine 0.999
|
||||
- Decode attention (1 query vs N cached KVs) — cosine 0.9998
|
||||
- Full pipeline (inv RoPE + o_a BMM + o_b) — cosine 0.996–0.999
|
||||
- All 5 layer types (C128A, C4A, SWA) — cosine ≥0.996
|
||||
|
||||
## Architecture: DeepSeek-V4-Pro
|
||||
|
||||
**CSA + HCA + mHC** (NOT MLA — vLLM misnames it "MLA" in code):
|
||||
**MegaMoE (384 experts, top-6) with CSA + HCA + mHC:**
|
||||
|
||||
- **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.
|
||||
@@ -76,42 +74,44 @@ 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)
|
||||
```
|
||||
|
||||
**Expert intermediate size: 3072** (NOT 18432 — that's 6×3072 for top-6)
|
||||
|
||||
**DeepGEMM MegaMoE**: DeepSeek's persistent grouped GEMM for MoE uses TMA tensormap updates per expert with variable block_m (16-192) based on expected tokens per expert. Our CuTeDSL runner uses `run_nvfp4_grouped_gemm` (simpler, but proven correct in standalone tests).
|
||||
|
||||
## Current Status
|
||||
|
||||
### ✅ Working (verified on B200 standalone tests)
|
||||
### ✅ Verified (B200 venv, real weights, zero NaN)
|
||||
|
||||
| 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 |
|
||||
| KV cache write + decode attention | `test_decode_attention_b200.py` | 0.9998 |
|
||||
| Decode vs prefill consistency (5 layers) | `test_decode_vs_prefill_b200.py` | 0.996–0.999 |
|
||||
| E2E 61-layer model (shared experts) | `test_e2e_decode_b200.py` | healthy logits |
|
||||
| MoE runner (grouped GEMM, 16 experts) | `test_moe_runner_nan_b200.py` | no NaN, all sizes |
|
||||
| Full layer (attention + MoE) | `test_full_layer_nan_b200.py` | no NaN |
|
||||
| Multi-layer chain (3 layers) | `test_full_layer_nan_b200.py` | no NaN |
|
||||
|
||||
### 🔧 Needs CuTeDSL Kernels
|
||||
### ❌ Container — NaN in vLLM compiled execution
|
||||
|
||||
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)
|
||||
The container produces empty/garbage output. Debug logs show NaN in `hidden_states` from the first forward pass. **The NaN is NOT from our kernels** — it comes from vLLM's compiled execution infrastructure (see CURRENT_BUG.md for full investigation).
|
||||
|
||||
Most likely sources:
|
||||
1. `attn_gemm_parallel_execute` — fused parallel GEMM (NOT our CuTeDSL kernel)
|
||||
2. `fused_q_kv_rmsnorm` — CUDA kernel that may produce NaN on Blackwell
|
||||
3. Weight packing during model loading
|
||||
4. `torch.compile` + cudagraph interaction with CuTeDSL buffers
|
||||
|
||||
### ❌ 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
|
||||
|
||||
@@ -120,11 +120,15 @@ The container builds and starts successfully. The server accepts requests and ge
|
||||
| `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) |
|
||||
| `tests/test_decode_attention_b200.py` | Prefill + decode with KV cache | ✅ 0.9998 |
|
||||
| `tests/test_decode_vs_prefill_b200.py` | Decode vs prefill consistency | ✅ 0.996-0.999 |
|
||||
| `tests/test_e2e_decode_b200.py` | 61-layer E2E (shared experts) | ✅ healthy logits |
|
||||
| `tests/test_moe_nan_b200.py` | Single expert NaN check | ✅ no NaN |
|
||||
| `tests/test_moe_runner_nan_b200.py` | MoE grouped GEMM NaN check | ✅ no NaN |
|
||||
| `tests/test_full_layer_nan_b200.py` | Full layer + multi-layer NaN check | ✅ no NaN |
|
||||
|
||||
## Project Structure
|
||||
|
||||
@@ -133,63 +137,42 @@ 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)
|
||||
│ ├── runner.py # CuTeDSLMoERunner — grouped GEMM MoE
|
||||
│ ├── blackwell_attention.py # KV cache + attention (standalone, works)
|
||||
│ ├── csa_attention.py # CSA/HCA attention (BF16 SDPA)
|
||||
│ ├── custom_ops.py # torch.autograd wrappers
|
||||
│ └── kernel/moe/ # NVIDIA's ScaledGroupedGemmKernel
|
||||
├── vllm/ # vLLM integration
|
||||
│ ├── nvfp4_cutedsl.py # CuTeDSLMoERunner
|
||||
│ ├── nvfp4_cutedsl.py # CuTeDSLMoERunner (vLLM wrapper)
|
||||
│ ├── cutedsl_quant_method.py # CuTeDSLNvfp4LinearMethod
|
||||
│ ├── kernels/linear/nvfp4/cutedsl.py # vLLM kernel registration
|
||||
│ └── patches/
|
||||
│ ├── deepseek_v4_attention.py # Attention patch (Blackwell dispatch)
|
||||
│ ├── deepseek_compressor.py # Compressor patch (skip fused kernel on Blackwell)
|
||||
│ ├── 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)
|
||||
│ └── ...
|
||||
│ ├── csa_attention.py # BF16 SDPA + KV cache (our Blackwell path)
|
||||
│ └── deepseek_compressor.py # Skip fused kernel on Blackwell
|
||||
├── tests/ # Standalone tests (run on B200 venv)
|
||||
└── Dockerfile # Container build
|
||||
├── Dockerfile # Container build
|
||||
├── README.md # This file
|
||||
└── CURRENT_BUG.md # Current bug investigation
|
||||
```
|
||||
|
||||
## Plan
|
||||
|
||||
### Phase 1: MoE Kernel ✅ DONE
|
||||
### Phase 2: NVFP4 Linear Kernels ✅ DONE
|
||||
### Phase 3: vLLM Integration ✅ DONE (NVFP4 linear + MoE working)
|
||||
### Phase 3: Attention Pipeline ✅ DONE (standalone, all tests pass)
|
||||
### Phase 4: vLLM Integration 🔧 IN PROGRESS — blocked by NaN from vLLM infrastructure
|
||||
|
||||
### Phase 4: CuTeDSL Attention Backend 🔧 NEXT — BUILD THE KERNELS
|
||||
**Current blocker:** NaN in the vLLM container's compiled execution. Our kernels produce zero NaN standalone. The NaN comes from vLLM's `attn_gemm_parallel_execute` or `fused_q_kv_rmsnorm` CUDA kernels, weight packing, or torch.compile interaction.
|
||||
|
||||
**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
|
||||
**Next:**
|
||||
1. Install vllm in B200 venv, test the exact parallel GEMM path
|
||||
2. Test with torch.compile disabled in the container
|
||||
3. Add NaN checks inside the parallel GEMM wrapper
|
||||
4. If the parallel GEMM is the source, replace it with our CuTeDSL kernels (path of least resistance)
|
||||
|
||||
### Phase 5: Production
|
||||
- End-to-end benchmarking
|
||||
|
||||
Reference in New Issue
Block a user