Update README: reflect current state, add C128A/C4A topk + warmup fixes
This commit is contained in:
67
README.md
67
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.
|
||||
|
||||
Reference in New Issue
Block a user