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

Rename FmhaV3StageC → FmhaKernel — no dev stage artifacts in production API

Browse Source
This commit is contained in:
biondizzle
2026-05-23 05:45:58 +00:00
parent 787a25516d
commit bfacfeca7b
2 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
README.md
View File

@@ -144,7 +144,7 @@ Summary
|-------|--------|-------------|
| A | ✅ COMPLETE | Q@K^T via tcgen05.mma → TMEM → GMEM |
| B | ✅ COMPLETE | QK → identity softmax → P@V pipeline (TMEM alias, KV-tile interleaving) |
| C | ✅ MIGRATED TO MODULE | Real online softmax + normalize. n=128 cos 0.973. Migrated to `dsv4/kernels/attention/fmha.py` as `FmhaV3StageC`. TMEM layout mismatch still present (3% error). |
| C | ✅ MIGRATED TO MODULE | Real online softmax + normalize. n=128 cos 0.973. Migrated to `dsv4/kernels/attention/fmha.py` as `FmhaKernel`. TMEM layout mismatch still present (3% error). |
| D1 | 🔨 IN PROGRESS | Parameterize HEAD_DIM (64 → 512). SMEM-P path for hd>64 (register→SMEM copy TODO). |
| D2 | TODO | Multi-query grid with head packing (128 Q heads, MQA) |
| D3 | TODO | SWA sequence length mask (swa_lens per batch) |
@@ -160,7 +160,7 @@ Summary
dsv4/
├── kernels/ Pure GPU code (CuTeDSL @cute.jit, .cu files)
│ ├── gemm/ NVFP4 MoE GEMM kernels (grouped, fused_swiglu, dense, scheduler)
│ ├── attention/ FMHA kernel — FmhaV3StageC (migrated from tests), SMEM-P stub
│ ├── attention/ FMHA kernel — FmhaKernel (hd=64, TMEM-P proven; SMEM-P stub for hd>64)
│ ├── compressor/ CSA/HCA token-level compressor (CuTeDSL, 419 lines)
│ ├── indexer/ CSA indexer — score+topk (FP32 dot products, top-k selection)
│ ├── router/ Dense router decode kernel (warp-specialized persistent GEMM)
@@ -228,7 +228,7 @@ dsv4/
|------|-------|--------|
| `test_fmha_v3.py` | A+B | ✅ Full QK→identity softmax→PV, cosine 0.999999 |
| `test_fmha_v3_12w.py` | A+B | ✅ 12-warp QK→PV, cosine 0.999999 |
| `test_fmha_v3_stage_c.py` | C | ✅ Real online softmax + normalize, n=128 cos 0.973. **Also migrated to `dsv4/kernels/attention/fmha.py` as `FmhaV3StageC`.** |
| `test_fmha_v3_stage_c.py` | C | ✅ Real online softmax + normalize, n=128 cos 0.973. **Also in module as `FmhaKernel`.** |
| `test_fmha_v3_stage_d1.py` | D1 | 🔨 Parameterized hd + SMEM-P path (WIP) |
| `test_d1_*.py` | D1 | 🔨 Debug/diagnostic variants (hd512, regression, sweep, raw, debug) |
| `test_paired_epilog.py` | C | ✅ Paired atom epilogue experiments |
@@ -532,7 +532,7 @@ When implementing D5a, Stage C's epilogue changes from "multiply by 1/row_sum" t
### E1 — File placement
`dsv4/kernels/attention/fmha.py`. Currently contains `FmhaV3StageC` (exact migration from test). Will become `FmhaKernel` once D1 parameterization is complete and the SMEM-P path is working. Constructor takes all dimensions and dtypes, no module-level constants.
`dsv4/kernels/attention/fmha.py`. Currently contains `FmhaKernel` (migrated from test, hd=64 TMEM-P). Will gain parameterized `head_dim` and SMEM-P path in D1. Constructor takes all dimensions and dtypes, no module-level constants.
### E2 — Constructor signature
@@ -557,8 +557,6 @@ class FmhaKernel:
All architecture-level shapes from config flow into the constructor. No FMHA-internal magic numbers.
**Naming convention:** The class will be `FmhaKernel` once D1 is complete (replacing the current `FmhaV3StageC`). The progression: `FmhaV3StageC` (hd=64, TMEM-P only) → `FmhaKernel` (parameterized hd, TMEM-P + SMEM-P). The old name stays in the test file for regression.
### E3 — Call signature
```python

2
dsv4/kernels/attention/fmha.py
View File

@@ -16,7 +16,7 @@ import math
HEAD_DIM = 64
class FmhaV3StageC:
class FmhaKernel:
def __init__(self, s_k=128, scale_softmax=None):
self.s_k = s_k
self.n_kv_tiles = s_k // 128