nvfp4-megamoe-kernel/README.md

# DeepSeek-V4 NVFP4 Kernel Suite

CuTeDSL kernels for DeepSeek-V4 (Blackwell B200, SM100). All kernels use `cutlass.cute` (CuTeDSL) with Blackwell tensor cores.

## File Map

```
cutedsl/
├── native_swa_decode.py         # SWA decode attention — IN PROGRESS (v3 tcgen05 rewrite)
├── native_sparse_decode.py      # Sparse (CSA/HCA) decode — NOT YET REWRITTEN
├── nvfp4_cutedsl.py             # NVFP4 MoE runner (CuTeDSL) — WORKING
├── moe_pipeline.py              # MoE fused SwiGLU pipeline — WORKING
├── blackwell_attention.py       # vLLM bridge for Blackwell attention path
├── csa_attention.py             # CSA/HCA sparse attention bridge
├── custom_ops.py                # Custom CUDA ops registration
└── kernel/
    └── blockscaled_gemm/
        └── dense_blockscaled_gemm_persistent.py  # REFERENCE: Blackwell TMEM/tcgen05 GEMM

tests/
├── test_stage_a_v2.py           # ✅ Stage A: bare Q@K^T via tcgen05.mma → TMEM → GMEM
├── test_stage_b_v7.py           # 🔨 Stage B: two MMAs + identity softmax (runs, wrong output)
├── test_stage_b_minimal.py      # ✅ Stage B minimal: two MMAs, no softmax (runs, NaN expected)
├── test_stage_b_pipeline_only.py # ✅ Stage B pipeline-only: PipelineUmmaAsync, no ld/st (runs, NaN expected)
├── diag_tmem.py                 # Diagnostic: TMEM layout inspection
├── test_stage_b_v6.py           # ❌ Stage B v6 (hardcoded offsets, crashes)
├── test_stage_a_qk.py           # ❌ Stage A v1 (broken, superseded by v2)
├── test_stage_a_minimal.py      # ❌ Stage A minimal (broken, superseded by v2)
├── test_attention_path_b200.py  # Full attention path test (uses naive BF16 attn)
└── ...
```

## Current Status

### ✅ Stage A: Bare Q@K^T via tcgen05.mma — COMPLETE (May 20)

**File**: `tests/test_stage_a_v2.py`
**Result**: Q(128,128) @ K^T(128,128) → S(128,128), cosine 0.999999

Validates the full tcgen05.mma → TMEM → epilogue → GMEM path:
- tcgen05.mma with BF16 inputs, FP32 TMEM accumulator
- TMA load for A and B (cute.nvgpu.make_tiled_tma_atom_A/B)
- TMA store for C (cpasync.CopyBulkTensorTileS2GOp)
- Warp specialization: 4 epilogue warps + 1 MMA warp + 1 TMA warp = 192 threads
- PipelineTmaUmma for AB pipeline, PipelineUmmaAsync for acc pipeline
- TmemAllocator for TMEM allocation/deallocation
- utils.gemm.sm100.epilogue_tma_store for the TMEM→reg→SMEM→TMA→GMEM epilogue

### 🔨 Stage B: Two MMAs + Identity Softmax — IN PROGRESS (May 20)

**Latest**: `tests/test_stage_b_v7.py`
**Status**: Kernel compiles and runs without crashing. Identity softmax produces wrong output (cosine ≈ -0.02).

**What was fixed today:**

1. **PipelineUmmaAsync consumer group size crash (THE bug):**
   `PipelineUmmaAsync` with `Agent.Thread` requires **thread count** (128), NOT warp count (4), for the consumer group. fmha.py uses `32 * len(softmax_warp_ids) = 128`. Using 4 caused `CUDA_ERROR_LAUNCH_FAILED` (not a deadlock — the barrier reached wrong threshold causing illegal TMEM access).
   
   ```python
   # WRONG (caused CUDA_ERROR_LAUNCH_FAILED):
   consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 4)
   # CORRECT (matches fmha.py):
   consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 128)
   ```

2. **TMEM offset computation (no more hardcoding):**
   - `s_cols = find_tmem_tensor_col_offset(tStS) = 128` — QK accumulator physical TMEM columns
   - `o_cols = find_tmem_tensor_col_offset(tOtO) = 128` — PV accumulator physical TMEM columns
   - `tmem_s0_offset = 0, tmem_p0_offset = 32, tmem_o0_offset = 128` — matches fmha.py
   - `find_tmem_tensor_col_offset(tOrP_sliced) = 32800 = 0x8020` — 0x8000 is TMEM space tag, column offset = 32
   - Total: 256 TMEM cols (verified by `get_num_tmem_alloc_cols`)

3. **P fragment construction (matching fmha.py):**
   ```python
   tP = cute.make_tensor(tStS.iterator, p_tmem_s.outer)  # A-layout from PV MMA
   tOrP = pv_thr.make_fragment_A(tP)[None, None, None, 0]
   tOrP0 = cute.make_tensor(tOrP.iterator + 2 * tmem_p0_offset, tOrP.layout)
   ```
   Previously used `cute.composition` on C-layout — wrong, must use PV MMA's A-layout.

4. **V SMEM aliasing:**
   V shares the same SMEM as K with a different layout interpretation:
   ```python
   sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner)
   sV = cute.make_tensor(sV_ptr, v_smem_s.outer)
   tCrV = pv_mma.make_fragment_B(sV)  # Uses MN-major V layout
   ```

**What's still broken:**

The identity softmax C→A layout transform produces garbage output (cosine ≈ -0.02). The kernel runs, Stage A (Q@K^T) gives cosine 0.999999, but the full (Q@K^T)@V pipeline is wrong. The issue is in the tcgen05.ld/st identity softmax path — either the ld/st copy atoms, the register conversion, or the A-layout write positions are incorrect.

**Bisection results:**
- ✅ Stage B minimal (no pipeline, no softmax): runs, NaN (expected — no C→A transform)
- ✅ Stage B pipeline-only (PipelineUmmaAsync, no ld/st): runs, NaN (expected)
- 🔨 Stage B full (identity softmax): runs, cosine -0.02 (wrong — softmax transform is broken)
- All three crash with consumer_group=4, all run with consumer_group=128

**TMEM layout diagnostic data:**
```
QK accumulator C fragment:
  tStS.layout = ((128,128),1,1):((65536,1),0,0)
  cute.size = 16384, cute.cosize = 8323200
  find_tmem_tensor_col_offset = 128

PV A-fragment (P operand):
  tOrP_sliced.layout = ((128,16),1,4):((65536,1),0,16)
  cute.size = 8192, cute.cosize = 8323136
  find_tmem_tensor_col_offset = 32800 = 0x8020 (0x8000 tag + col 32)
```

### 🔨 Stage C: Online Softmax — AFTER B

The hard part. Per the pseudocode:
- Epilogue warps tcgen05.ld scores from TMEM into register fragments
- Compute per-row: tile_max, new_max, rescale = exp(old_max - new_max)
- Apply rescale to tmem_output in place (tmem_output *= rescale)
- Compute exp(scores - new_max), tcgen05.st back to TMEM as P operand for MMA2
- Update row_sum = row_sum * rescale + new_tile_sum

**The register fragment layout from tcgen05.ld is NOT (row, col).** It's determined by the MMA instruction's partition of the accumulator. Need to figure out the mapping from fragment indices to logical (head, kv_pos) positions for per-row softmax operations. fmha.py uses `tTMEM_LOADrS.load().reduce(cute.ReductionOp.MAX, row_max, 0)` for the row max — a built-in reduction that handles the layout.

### 🔨 Stage D: FP8 Paged KV Gather — AFTER C

Replace BF16 TMA load of KV with:
- Indexed cp.async gather from paged KV cache (fp8)
- Per-position dequant scale (inv_scale) applied during or after gather
- Keep KV in fp8 in SMEM, let the MMA's per-row scale handle dequant (like blockscaled GEMM)

### Architecture: Per-Tile Flow (from /root/fragile-kernel-example/README.md)

```
For each KV tile:
  1. Load warp writes sKV[stage] (paged FP8 gather via indexed cp.async)
  2. MMA warp issues MMA1: sQ @ sKV[stage]^T → tmem_scores (accumulate=False)
     Signals scores_full_mbar (via PipelineUmmaAsync commit)
  3. Epilogue warps wait on mma_si consumer (scores ready), then:
     a. tcgen05.ld scores from TMEM → register fragments
     b. Compute tile_max, new_max, rescale = exp(old_max - new_max)
     c. Apply rescale to tmem_output IN PLACE (tmem_output *= rescale)
     d. tcgen05.st exp(scores - new_max) back to TMEM → now it's the P operand
     e. Release mma_si (softmax_done — MMA warp can re-acquire and issue PV MMA)
  4. MMA warp waits on mma_si acquire (softmax done), then MMA2: P @ sKV[stage] → tmem_output (accumulate=True)
  5. Stage released, load warp can refill it

After all tiles: epilogue warps tcgen05.ld tmem_output, divide by row_sum, cast to BF16, store to GMEM
```

### ✅ NVFP4 MoE (CuTeDSL) — WORKING
- `nvfp4_cutedsl.py` + `moe_pipeline.py`
- CuTeDSL NVFP4 Linear (q_a, kv, q_b, o_b) — cosine 0.994+
- CuTeDSL NVFP4 MoE (L1 gate+up, SiLU, L2 down) — cosine 0.988
- Fused SwiGLU epilogue (granularity-8 weight interleave) — cosine 0.988

### ✅ FP8 KV Quantize/Dequant — WORKING
- FP8 KV: cosine 0.9997
- NVFP4 KV: cosine 0.9943 (2x smaller than FP8)
- Paged KV cache read/write: cosine 1.0

### ❌ Sparse Decode Attention — NOT YET REWRITTEN
`native_sparse_decode.py` still has the scalar FMA bug. Needs the same tcgen05.mma rewrite.

### ✅ Full Attention Pipeline (standalone tests) — WORKING
- FP8 KV → full attention: cosine 0.9997
- CSA sparse attention (cr=4): works
- HCA sparse attention (cr=128): works
- Merged CSA+SWA attention: works

## Critical APIs & Lessons

### PipelineUmmaAsync consumer group size — THE MAY 20 BUG

**For `Agent.Thread` groups in `PipelineUmmaAsync`: use thread count, NOT warp count.**

```python
# WRONG (caused CUDA_ERROR_LAUNCH_FAILED):
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 4)  # warp count

# CORRECT (matches fmha.py):
consumer_group=pipeline.CooperativeGroup(pipeline.Agent.Thread, 32 * len(softmax_warp_ids))  # thread count
```

This applies to ALL PipelineUmmaAsync consumers where the consumer is multiple warps. fmha.py line 671: `self.threads_per_warp * len(self.softmax0_warp_ids) = 32 * 4 = 128`.

**Note:** The earlier README incorrectly stated that warp count was correct. That was wrong. The `Agent.Thread` agent type measures group size in threads.

### TMEM offset arithmetic

- `find_tmem_tensor_col_offset(fragment)` — returns physical TMEM column count (with 0x8000 tag for A-fragments)
- QK accumulator C fragment: 128 TMEM columns
- PV A-fragment: offset 0x8020 = tag(0x8000) + col(32) — the 0x8000 is a TMEM memory-space identifier
- P OVERLAPS S in TMEM — P is written at column 32 within the S region (C-layout columns 0..127)
- `tOrP0 = cute.make_tensor(tOrP.iterator + acc_dtype.width // q_dtype.width * tmem_p0_offset, tOrP.layout)` — A-fragment offset scaled by dtype width ratio (F32/BF16 = 2)

### `make_trivial_tiled_mma` has two overloads

```python
# New (preferred):
make_trivial_tiled_mma(a_dtype, b_dtype, a_leading_mode, b_leading_mode,
                        acc_dtype, cta_group, mma_tiler_mn, a_source=SMEM)

# Deprecated (still works, used by Stage A):
make_trivial_tiled_mma(ab_dtype, a_leading_mode, b_leading_mode,
                        acc_dtype, cta_group, mma_tiler_mn, a_source=SMEM)
```

### V SMEM aliasing (K and V share SMEM)

```python
# K and V share the same SMEM buffer, but with different layouts:
v_smem_s = utils.sm100.make_smem_layout_b(pv_mma, pv_mma_tiler, b_dtype, 1)
sV_ptr = cute.recast_ptr(sB.iterator, v_smem_s.inner)
sV = cute.make_tensor(sV_ptr, v_smem_s.outer)
tCrV = pv_mma.make_fragment_B(sV)
```

### Other APIs discovered from Stage A

1. **`cute.Tensor` API** — `cutlass_torch.from_dlpack(t).mark_layout_dynamic(leading_dim=...)`
2. **3D tensors** — Tensors must be 3D (M, K, L) for `cute.local_tile` — add L=1 dimension
3. **`PipelineTmaUmma.create(...).make_participants()`** — returns `(producer, consumer)` pair
4. **`utils.gemm.sm100.epilogue_tma_store`** — handles transform + partition/dcopy. DO NOT hand-roll.
5. **`get_num_tmem_alloc_cols`** — correct TMEM allocation (accepts list of fragments, sums cols, rounds to power of 2)
6. **`smem.allocate_tensor()`** — for SMEM tensors (not SharedStorage struct for A/B/C)
7. **`LayoutEnum.from_tensor(a).mma_major_mode()`** — major mode from cute tensor
8. **Minimum valid N tile for tcgen05.mma BF16**: 32 (step 32, range 32-256)

## Environment

- **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)
- **Pseudocode**: `/root/fragile-kernel-example/README.md` — authoritative per-tile attention flow
- **fmha.py reference**: `/root/cutlass/examples/python/CuTeDSL/cute/blackwell/kernel/attention/fmha/fmha.py`

## 4-Stage Build Plan

| Stage | Goal | Status |
|-------|------|--------|
| A | Bare Q@K^T via tcgen05.mma → TMEM → GMEM | ✅ COMPLETE |
| B | Two MMAs + identity softmax (validates TMEM A operand, shared KV, layout transform, barrier ordering) | 🔨 Runs without crash, identity softmax produces wrong output |
| C | Online softmax between MMA1 and MMA2 (the hard part) | ⬜ TODO |
| D | FP8 paged KV gather + dequant (replace BF16 TMA load) | ⬜ TODO |