diff --git a/README.md b/README.md index b4932212..6726849e 100644 --- a/README.md +++ b/README.md @@ -1,178 +1,123 @@ -# DSV4 NVFP4 Kernel +# DSV4 Inference Kernel -## Status (May 21, 2026 — 15:40 UTC) +## Status (May 21, 2026 — 17:30 UTC) | Stage | Status | Description | |-------|--------|-------------| | A | ✅ COMPLETE | Q@K^T via tcgen05.mma → TMEM → GMEM | | B | ✅ COMPLETE | QK → identity softmax → P@V pipeline (TMEM alias, KV-tile interleaving) | -| C | 🔨 NEXT | Real softmax: row max, exp, rescale, row sum | +| C | 🔨 IN PROGRESS | Real softmax: row max, exp, rescale, row sum (kernel written, needs test harness) | | D | TODO | Full decode attention: paged KV cache, multi-query, causal mask | -| E | TODO | Production kernel: integrate into `cutedsl/`, PyTorch custom op, vLLM bridge | +| E | TODO | Production kernel: extract into dsv4/kernels/attention/, PyTorch custom op, vLLM bridge | --- -## Where Things Live +## Package Structure -### Canonical Test Files (current working code) +``` +dsv4/ +├── kernels/ Pure GPU code (CuTeDSL @cute.jit, .cu files) +│ ├── gemm/ NVFP4 MoE GEMM kernels (grouped, fused_swiglu, dense, scheduler) +│ ├── attention/ FMHA kernel (stub — extraction is Stage E) +│ ├── compressor/ CSA/HCA token-level compressor +│ ├── decode/ Decode-time attention (sparse, SWA — future) +│ └── cuda/ Raw .cu files (deinterleave_quantize, sparse_topk_metadata) +├── ops/ PyTorch ↔ kernel bridges +│ ├── quantize.py BF16 ↔ NVFP4 conversion, scale factors +│ ├── layouts.py Scale swizzle, gate/up interleave, K-major, offsets +│ ├── gemm_runner.py Warmup, compile, run grouped/fused GEMMs +│ ├── custom_ops.py torch.library.custom_op registrations +│ ├── decode_sparse.py native_sparse_decode dispatcher +│ ├── decode_swa.py native_swa_decode dispatcher +│ ├── rope.py Forward + inverse RoPE +│ └── topk.py Python wrapper for sparse_topk_metadata.cu +├── layers/ nn.Module-style components +│ ├── linear.py Nvfp4Linear +│ ├── grouped_linear.py Nvfp4GroupedLinear +│ ├── moe.py Nvfp4MoE +│ ├── shared_expert.py Nvfp4SharedExpert +│ ├── mhc.py mHCLayer +│ └── (stubs: attention, ffn, router, norm, embedding) +├── model/ Model assembly (stubs — Phase 1) +├── cache/ KV cache infra (stubs — Phase 3) +├── loader/ Checkpoint I/O (stubs — Phase 1) +└── reference/ Slow PyTorch oracles (never imported by production code) + ├── attention.py RoPE, KV cache, causal attention, SWA + ├── csa_attention.py CSA/HCA sparse attention + ├── compressor.py Compressor PyTorch example + └── moe_pipeline.py MoE pipeline reference +``` -| File | Stage | What It Tests | -|------|-------|---------------| -| `tests/test_fmha_v3.py` | A+B | Full QK→softmax→PV with KV-tile interleaving. **This is the canonical Stage B kernel.** | -| `tests/test_pv64_with_softmax.py` | B | Single AB pipeline variant. Simpler, also passing. | -| `tests/test_128_128_vdiag.py` | A+B | (128,128) PV baseline. Square case. | -| `tests/test_qkonly.py` | A | QK with split Q/KV pipelines, no softmax, no PV. | -| `tests/test_qk_softmax.py` | A+partial B | QK + softmax writes P to TMEM. No PV. | - -### Reference Implementations (non-CuTe) - -| File | Description | -|------|-------------| -| `cutedsl/blackwell_attention.py` | Pure PyTorch reference: RoPE, KV cache, causal attention, SWA | -| `cutedsl/csa_attention.py` | CSA/HCA sparse attention reference | -| `cutedsl/native_swa_decode.py` | SWA decode reference | - -### CuTeDSL Kernel Modules (production code) - -| File | Description | -|------|-------------| -| `cutedsl/kernel/moe/` | NVFP4 MoE GEMM kernels (fused SwiGLU, grouped MM) | -| `cutedsl/kernel/blockscaled_gemm/` | Block-scaled GEMM | -| `cutedsl/nvfp4_linear.py` | NVFP4 linear layer wrapper | -| `cutedsl/runner.py` | MoE runner | -| `cutedsl/moe_pipeline.py` | MoE pipeline orchestration | - -### Obsolete Tests (do not use, can delete) - -~100+ test files from debugging stages A and B. Key patterns: -- `test_stage_b_v1.py` through `test_stage_b_v30.py` — incremental Bug 4b debugging -- `test_128_16_*.py` — early (128,16) PV attempts with wrong head_dim -- `test_tmem_*.py`, `test_bf16_*.py` — standalone TMEM/copy debugging -- `test_pv64_no_softmax.py`, `test_pv64_fmha_v.py`, `test_pv64_kmajor_v.py` — Bug 4b root cause isolation - -These can be deleted once the canonical tests are stable and the kernel is extracted. +**Mental model:** `kernels/` → `ops/` → `layers/` → `model/` (dependency flows left to right). `reference/` and `loader/` are sidecars. --- -## Stage C: Real Softmax +## Active Test Files -### What We Have Now +### FMHA (Stages A/B/C) — in `tests/unit/` -Identity softmax: load S FP32 from TMEM, convert to BF16, store P back to TMEM. This proves the TMEM pipeline works but isn't a real softmax. +| File | Stage | Status | +|------|-------|--------| +| `test_fmha_v3.py` | A+B | ✅ Full QK→softmax→PV, cosine 0.999999 | +| `test_fmha_v3_softmax.py` | C | 🔨 Online softmax kernel (needs test harness) | +| `test_pv64_with_softmax.py` | B | ✅ (128,64) PV, single AB pipeline | +| `test_128_128_vdiag.py` | A+B | ✅ (128,128) PV baseline | +| `test_qkonly.py` | A | ✅ QK with split Q/KV pipelines | +| `test_qk_softmax.py` | A+B | ✅ QK + identity softmax, no PV | -### What We Need +### MoE / GEMM — in `tests/unit/` -FMHA-style online softmax per KV-tile: +| File | What | +|------|------| +| `test_cutedsl.py` | NVFP4 grouped GEMM kernel | +| `cudagraph_test.py` | Cudagraph capture + replay | +| `layertest.py` | Per-layer correctness | +| `test_custom_op.py` | torch.library custom ops | +| `test_compile_custom_op.py` | Compile + warmup | +| `test_fp4_roundtrip.py` | BF16 → NVFP4 → BF16 roundtrip | +| `test_interleave.py` | Gate/up weight interleaving | +| `test_interleave_gemm.py` | Interleaved GEMM correctness | +| `test_fused_step1.py` | Fused SwiGLU GEMM | + +### Archived Tests + +`tests/archive/` contains ~190 debug files from Stages A/B. Not maintained. Can be deleted. + +--- + +## Stage C: Online Softmax + +### What We Have + +Identity softmax in `test_fmha_v3.py`: load S FP32 → convert BF16 → store P. Proves TMEM pipeline works. + +### What We Are Building + +Online softmax in `test_fmha_v3_softmax.py` (kernel written, no test runner yet): ``` For each KV tile: 1. QK → S (FP32 in TMEM) - 2. Load S row-max for this tile: tile_max[j] = max(S[j,:]) - 3. Compute new row max: new_max[j] = max(old_max[j], tile_max[j]) - 4. Rescale O: O[j,:] *= exp(old_max[j] - new_max[j]) - 5. Compute P: P[j,i] = exp(S[j,i] - new_max[j]) - 6. Store P to TMEM (BF16, same C-fragment composition store) - 7. Update row sum: row_sum[j] = row_sum[j] * exp(old_max[j] - new_max[j]) + sum(P[j,:]) - 8. PV: O[j,:] += P[j,:] @ V[i,:] + 2. tile_max = max(S[j,:]) + 3. new_max = max(old_max, tile_max) + 4. O *= exp(old_max - new_max) ← TMEM rescale + 5. P = exp2((S - new_max) * scale) ← exp2 with 1/sqrt(d) * log2(e) + 6. Store P to TMEM (FMHA pattern) + 7. row_sum = row_sum * exp(old_max - new_max) + sum(P) + 8. PV: O += P @ V After all tiles: - 9. O[j,:] /= row_sum[j] (final normalization) + 9. O /= row_sum ← final TMEM normalization ``` -### Key Challenges +### Key Implementation Details -1. **Row max across tiles** — Must track per-row maximum across KV-tiles and rescale O when a new max is found. This is the core of online softmax. -2. **Row sum accumulation** — Must accumulate exp(sum) across tiles with proper rescaling. -3. **FP32 precision** — Row max, rescale, and row sum must stay in FP32 for numerical stability. Only P (the exp values) get cast to BF16 for TMEM store. -4. **O rescale in TMEM** — When a new row max is found, the existing O in TMEM must be multiplied by `exp(old_max - new_max)`. This requires loading O, rescaling, and storing back. Same TMEM load/store machinery as softmax P. -5. **Final normalization** — After all KV-tiles, divide O by row_sum. Can be done as part of the epilogue. +- **Row max:** `tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0)` per tile +- **O rescale:** Load O from TMEM, multiply by `exp2(old_max - new_max)`, store back (16-col tiles via `Ld32x32b/St32x32b`) +- **P computation:** `exp2((S - row_max) * scale)` where `scale = 1/sqrt(HEAD_DIM) * log2(e)` +- **Row sum:** Packed `f32x2` reduction using `cute.arch.add_packed_f32x2` (4 unroll, 2-wide) +- **Final norm:** Load O, multiply by `1/row_sum`, store (same TMEM load/store path) -### Expected Structure - -The softmax epilogue warps will expand significantly: -- Currently: load S → convert BF16 → store P (identity softmax) -- After Stage C: load S → compute tile_max → compare with old_max → rescale O → compute exp → store P → update row_sum - -The MMA loop remains the same (QK → softmax → PV per tile). The softmax just does more work between QK completion and PV start. - ---- - -## Stage D: Full Decode Attention - -### What We Have After Stage C - -A working QK → real softmax → PV kernel for a single query sequence against a contiguous KV block. Fixed dimensions (128×128 QK, 128×64 PV), single CTA. - -### What We Need - -1. **Paged KV cache** — KV comes from a paged cache (fp8 or bf8 with per-token inverse scale), not a contiguous tensor. TMA loads must follow page tables. -2. **Multi-query** — Multiple query sequences in flight, each with different KV lengths. Requires grid dimensions > 1, possibly persistent kernel. -3. **Causal masking** — QK must mask future positions. For decode (1 query vs N KVs), this is trivial (no mask needed). For prefill, need a causal mask. -4. **Variable sequence length** — Each CTA handles a different number of KV tiles. The loop bound `n_kv_tiles` becomes dynamic. -5. **Multiple head dimensions** — HEAD_DIM=16, 64, 128 all need to work. Currently only HEAD_DIM=64 is tested. -6. **CSA/HCA sparse attention** — For compress_ratio > 1, KV is read from compressed cache instead of full KV cache. Different TMEM layouts, different attention patterns. - -### Key Question: Do We Need Stage D As A Separate Stage? - -Stage D is really about *scaling* the Stage C kernel, not adding fundamentally new compute. The core pipeline (QK → softmax → PV) doesn't change. What changes is: -- Where the data comes from (paged cache vs contiguous tensor) -- How many CTAs run (grid size) -- Whether we need causal masking - -This could be folded into the production kernel directly rather than being a separate test stage. - ---- - -## Stage E: Production Kernel - -### Goal - -Replace `cutedsl/blackwell_attention.py` (pure PyTorch) with a CuTeDSL kernel that runs on the Blackwell tensor cores. - -### Steps - -1. **Extract kernel from `test_fmha_v3.py`** → `cutedsl/kernel/attention/fmha_kernel.py` - - Class `FmhaKernel` with `@cute.jit __call__` - - Clean parameter interface: Q, K, V, O tensors + config - - No hardcoded dimensions — all derived from MMA shapes - -2. **Add real softmax (Stage C)** to the extracted kernel - -3. **Add paged KV cache support (Stage D)** - - Page table TMA or gather-style loads - - Per-sequence KV length tracking - -4. **Wrap as PyTorch custom op** → `cutedsl/custom_ops.py` - - `blackwell_fmha_forward(q, k, v, ...) -> o` - - Autograd support (or torch.compile integration) - - torch.library custom op registration - -5. **Integrate with vLLM** → `vllm/attention/ops/blackwell_fmha.py` - - Replace the broken FlashMLA Blackwell path - - Hook into vLLM's paged attention interface - - Support both prefill and decode modes - -6. **Benchmark and tune** - - Profile against PyTorch SDPA baseline - - Tune tile sizes, pipeline stages, SMEM usage - - Verify numerical accuracy vs reference across head dims and sequence lengths - -### File Structure (target) - -``` -cutedsl/ - kernel/ - attention/ - __init__.py - fmha_kernel.py ← extracted, clean CuTeDSL kernel - fmha_softmax.py ← real softmax (Stage C) - fmha_epilogue.py ← row sum normalization, O output - blackwell_attention.py ← PyTorch reference (keep for testing) - custom_ops.py ← PyTorch custom op wrappers -``` - ---- - -## TMEM Layout (Current) +### TMEM Layout (Current — Stage B) ``` Col: 0 32 64 96 128 192 256 @@ -181,27 +126,39 @@ Col: 0 32 64 96 128 192 256 | 128 FP32 | 64 FP32 | 32 col | 64 FP32 | ``` -For Stage C, we'll need additional TMEM regions: -- `row_max` — per-row FP32 max (128 rows × 1 col = 128 FP32 values, can use 4 TMEM columns) -- `row_sum` — per-row FP32 sum (128 rows × 1 col, 4 TMEM columns) -- `old_max` — per-row FP32 previous max (4 TMEM columns) - -These are tiny (4-8 TMEM columns each). They can go in the gap at columns 96-128 or after O. +For Stage C, row_max/row_sum are per-thread FP32 scalars (not in TMEM). Future stages may need TMEM-backed state for wider tiles. --- -## Key Lessons From Stages A & B +## Stage E: Production Kernel Extraction -1. **NEVER use `find_tmem_tensor_col_offset()` as a TMEM placement decision.** It returns footprint size, not a safe column offset. The P/O overlap bug cost the entire Bug 4b debugging session. +When ready, extract from `test_fmha_v3.py` → `dsv4/kernels/attention/fmha.py`: +1. Clean `FmhaKernel` class with `@cute.jit __call__`, no hardcoded dimensions +2. Add real softmax (Stage C) +3. Add paged KV cache (Stage D) +5. Wrap as `torch.library.custom_op` in `dsv4/ops/` +6. Integrate with vLLM -2. **FMHA never trusts DLPack tensor layouts.** Reconstruct V as (hd, s_k) MN-major inside CuTe. The DLPack shape (n, hd) has wrong logical modes for PV B-operand. +--- -3. **TMEM allocation must be power of 2.** `TmemAllocator.allocate()` asserts this. +## Key Lessons -4. **P/A alias works.** QK C-fragment composition store + PV A-fragment read alias the same physical TMEM columns. Proven for (128,64) and (128,128). +1. **NEVER use `find_tmem_tensor_col_offset()` as TMEM placement.** It returns footprint size, not a safe offset. +2. **FMHA never trusts DLPack tensor layouts.** Reconstruct V as (hd, s_k) MN-major inside CuTe. +3. **TMEM allocation must be power of 2.** +4. **Square hides bugs.** (128,128) worked for every wrong approach. Always test non-square. +5. **St32x32bOp MUST use Float32**, NOT BFloat16. BFloat16 causes illegal memory access. +6. **First PV ACCUMULATE=False.** Otherwise adds uninitialized TMEM to output. +7. **FMHA P store uses QK C-fragment composition, NOT PV A-fragment.** Two aliases, same TMEM. +8. **Register bridge: FP32 backing (store partition) + BF16 view (QK-load layout).** Do not skip this. -5. **Square hides bugs.** (128,128) PV worked for every wrong approach because both dims are 128. Always test non-square cases. +--- -6. **`St32x32bOp` MUST use Float32, NOT BFloat16.** BFloat16 causes illegal memory access. +## Environment -7. **First PV ACCUMULATE=False.** Otherwise adds uninitialized TMEM to output. +- Server: root@45.76.247.107 (B200, 180 GiB HBM3e per GPU) +- venv: `source /root/dsv4-nvfp4-workspace/venv/bin/activate` +- PYTHONPATH: `/root/dsv4-nvfp4-workspace/kernel` +- Model: `/root/nvidia-meeting/DeepSeek-V4-Pro-NVFP4` +- vLLM repo: `/root/dsv4-nvfp4-workspace/vllm` (modified for Blackwell) +- CUTLASS FMHA reference: `/root/cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py`