diff --git a/README.md b/README.md index 19794507..c7981519 100644 --- a/README.md +++ b/README.md @@ -74,12 +74,28 @@ vLLM's internal kernels (FlashMLA, fp8_ds_mla, fused compressor, Triton indexer) - **Cosine 0.9999+** vs combined SDPA reference - Supports both CSA (cr=4) and HCA (cr=128) layers +### ✅ Sparse Topk Metadata Kernels (C128A + C4A) + +`cutedsl/kernels/sparse_topk_metadata.cu` + `cutedsl/sparse_topk_metadata.py`: +- **`build_c128a_topk_metadata`**: position-based compressed KV slot lookup via block table for C128A (cr=128) decode tokens. Maps `(position, block_table) → global compressed KV slot IDs + lengths` +- **`compute_c4a_global_topk`**: local topk index → global KV cache slot mapping via block table for C4A (cr=4) decode tokens +- Both tested: correct block table lookups, proper padding, valid length counts +- **No FlashMLA, no vLLM Triton dependency** — own CUDA kernels + ### ✅ Blackwell Attention (standalone tests) - `cutedsl/blackwell_attention.py` — KV cache write/read, full attention pipeline - `cutedsl/csa_attention.py` — CSA (cr=4) and HCA (cr=128) sparse attention - All standalone tests pass: KV cache (0.9997), CSA/HCA, prefill+decode (0.9998) +### ✅ CuTeDSL Warmup Compilation + +`warmup_compilation()` and `warmup_fused_swiglu_compilation()` in `bridge.py`: +- Eagerly JIT-compiles GEMM kernels before model forward pass +- Uses **quantized random BF16** (via `quantize_to_nvfp4`) for warmup data +- Zero-filled FP4/FP8 causes `cudaErrorIllegalInstruction` — random bytes produce NaN in MMA dequant +- All three shapes compile successfully: L1 (48 experts, 3584×3072), L2 (48 experts, 3072×3584), Fused L1 + --- ## Bridge Layer (`cutedsl/bridge.py`) @@ -96,8 +112,8 @@ Quantization, layout, kernel launch utilities: | `deinterleave_quantize_nvfp4_cuda()` | Custom CUDA: de-interleave + quantize in one kernel | | `make_b_k_major()` | B tensor stride conversion | | `assemble_scales_2d_side()` / `assemble_scales_3d_side()` | Scale assembly + swizzle | -| `warmup_compilation()` | Eager JIT compilation (base GEMM) | -| `warmup_fused_swiglu_compilation()` | Eager JIT compilation (fused SwiGLU) | +| `warmup_compilation()` | Eager JIT compilation with quantized random data (base GEMM) | +| `warmup_fused_swiglu_compilation()` | Eager JIT compilation with quantized random data (fused SwiGLU) | | `run_nvfp4_grouped_gemm()` | Base GEMM entry point | | `run_fused_swiglu_grouped_gemm()` | Fused SwiGLU GEMM entry point | @@ -129,15 +145,39 @@ Quantization, layout, kernel launch utilities: --- +## Blackwell Decode Path (vLLM Integration) + +The Blackwell decode path in `attention.py` routes through our own kernels: + +**SWA-only layers (cr=0):** `native_swa_decode_attention` — CuTeDSL kernel + +**CSA layers (cr=4):** `native_sparse_decode_attention` with topk indices from `compute_c4a_global_topk` — our CUDA kernel maps indexer local topk → global KV cache slots + +**HCA layers (cr=128):** `native_sparse_decode_attention` with topk indices from `build_c128a_topk_metadata` — our CUDA kernel maps positions → compressed KV slot IDs via block table lookup + +**Metadata flow:** +- `DeepseekSparseSWAMetadataBuilder` builds SWA indices + C128A buffers +- `attention.py` detects FlashMLA vs Indexer metadata at runtime +- Blackwell path reads `indexer_metadata.decode.block_table` for block table access +- No FlashMLA dependency on Blackwell + +--- + ## Correctness Bugs Fixed (May 20, 2026) | Bug | Issue | Fix | |-----|-------|-----| -| 1 | `_needs_token_refill` myth — cute.compile doesn't corrupt GPU memory | Removed hack, pre-allocated workspace per cache entry | -| 2 | Dequantize→requantize supposedly lossy | Verified 100% byte-identical round-trip. Deprecated `prepare_weights_from_dequantized` | -| 3 | `clamp(min=1e-8)` on zero blocks gives nonzero FP8 scale | Detect zero blocks, force FP8 scale to exact 0 | -| 4 | Underflow blocks (amax < 6×2⁻⁹) get nonzero FP4 from div-by-tiny-number | Detect underflow blocks, zero x_norm before division | -| 5 | Expert counting materializes 18M bool tensor | `torch.bincount` replaces O(n×E) comparison | +| C128A topk missing | `DeepseekSparseSWAMetadataBuilder` returned None for C128A topk → SWA-only fallback | `build_c128a_topk_metadata` CUDA kernel computes global slot IDs from positions + block table | +| C4A topk missing | Relied on vLLM's Triton `compute_global_topk_indices_and_lens` (not ours) | `compute_c4a_global_topk` CUDA kernel replaces it on Blackwell | +| Warmup crash | Zero-filled FP4/FP8 → `cudaErrorIllegalInstruction` in MMA hardware | Quantize random BF16 through `quantize_to_nvfp4` for mathematically consistent warmup data | +| Warmup disabled | Was commented out → lazy JIT on first forward → OOM competing with model | Re-enabled in runner.py; L1/L2/fused all compile eagerly | +| `_fused_swiglu` not initialized | `CuTeDSLMoERunner.__init__` missing `self._fused_swiglu = False` | Added initialization | +| FlashMLA assert crash | `assert flashmla_metadata is not None` crashes on Blackwell where indexer_metadata is used instead | Fixed assert to accept either | +| `_needs_token_refill` myth | cute.compile doesn't corrupt GPU memory | Removed hack | +| Zero block FP8 scale | `clamp(min=1e-8)` gives nonzero scale for zero blocks | Detect zero blocks, force FP8 scale to exact 0 | +| Underflow blocks | amax < 6×2⁻⁹ gets nonzero FP4 | Detect underflow, zero x_norm before division | +| Expert counting | Materializes 18M bool tensor | `torch.bincount` replaces O(n×E) comparison | +| Dequantize→requantize | "Supposedly lossy" | Verified 100% byte-identical round-trip | --- @@ -198,9 +238,6 @@ The custom CUDA quantize kernel needs the **L2 activation global scale** (from t | What | Status | Notes | |------|--------|-------| | In-epilogue NVFP4 quantize (replace BF16 TMA with FP4 TMA) | 🔨 Future | Saves ~0.14ms/layer; requires register→GMEM mapping for FP4 output | -| GPU-native SWA decode attention | ✅ Done | CuTeDSL kernel, cosine 0.9999+ | -| GPU-native sparse + SWA decode attention | ✅ Done | CuTeDSL kernel, cosine 0.9999+ | -| vLLM Blackwell decode path | ✅ Done | _attention_impl_blackwell uses native SWA + sparse kernels | | Fuse fp8→bf16 dequant into CuTeDSL kernel | 🔨 Future | Currently pre-dequantized on host; need vectorized fp8 loads | | CSA/HCA sink weight merge in CuTeDSL | 🔨 Future | Applied on host for now; fuse into kernel for perf | @@ -230,8 +267,12 @@ cutedsl/ ├── csa_attention.py # CSA/HCA attention ├── custom_ops.py # torch.autograd wrappers ├── moe_pipeline.py # Standalone test pipeline (fused + non-fused) +├── sparse_topk_metadata.py # C128A + C4A topk metadata (Python wrapper) +├── native_swa_decode.py # GPU-native SWA decode (CuTeDSL) +├── native_sparse_decode.py # GPU-native sparse+SWA decode (CuTeDSL) ├── kernels/ -│ └── deinterleave_quantize.cu # Custom CUDA: de-interleave + NVFP4 quantize +│ ├── deinterleave_quantize.cu # Custom CUDA: de-interleave + NVFP4 quantize +│ └── sparse_topk_metadata.cu # Custom CUDA: C128A + C4A topk metadata └── kernel/moe/ ├── torch_scaled_grouped_mm.py # ScaledGroupedGemmKernel (the GEMM) └── fused_swiglu_grouped_mm.py # FusedSwiGLUScaledGroupedGemmKernel @@ -272,3 +313,7 @@ tests/ 10. **⛔ "SiLU on all positions" is NOT SwiGLU.** SwiGLU pairs silu(gate)*up. Applying SiLU to the full (M, 2×intermediate) output is just SiLU. The pairing must be explicit. 11. **⛔ The global scale must match the data being quantized.** Passing the L1 input gs to the SwiGLU quantize causes FP8 overflow → NaN. The gs must come from the SwiGLU output's magnitude. + +12. **⛔ NEVER use zero-filled or random-byte data for CuTeDSL warmup.** Zeros cause division-by-zero in scale dequant. Random uint8 bytes as FP4 produce NaN/Inf in MMA → `cudaErrorIllegalInstruction`. Always quantize random BF16 through `quantize_to_nvfp4` for mathematically consistent warmup data. + +13. **⛔ NEVER borrow kernels from vLLM or FlashMLA.** We own all our kernels. If we need a kernel that exists in vLLM's Triton or FlashMLA's C++, we build our own CUDA/CuTeDSL equivalent from scratch.