biondizzle/nvfp4-megamoe-kernel
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

README: update for new dsv4/ package structure

Browse Source
This commit is contained in:
biondizzle
2026-05-21 17:34:40 +00:00
parent 9cbdc92744
commit 364d9edcd3
1 changed files with 118 additions and 161 deletions
Download Patch File Download Diff File Expand all files Collapse all files

279
README.md
View File

@@ -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`