DSV4 is **MQA** — 64 or 128 query heads all share the SAME K/V. This means:
- Q shape: `(batch, num_query_heads, T, head_dim)` — each head has its own Q
- K/V shape: `(batch, 1, s_k, head_dim)` — shared across all heads
- O shape: `(batch, num_query_heads, T, head_dim)` — per-head output
---
## Architecture: Two Grid Strategies
### Strategy A: Head-Packed M Dimension (RECOMMENDED for decode)
Fold the head dimension into M. Each CTA processes ALL heads' queries for its M tile.
```
Q reshaped: (batch, T * num_query_heads, head_dim) — heads packed into M
K/V: (batch, s_k, head_dim) — shared, loaded once
Grid: (ceil_div(T * n_h, 128), 1, batch)
At decode T=1, n_h=128: grid = (1, 1, batch) — single CTA per batch!
```
**Pros:**
- K/V loaded once per CTA, shared across all heads in the M tile
- Maximum arithmetic intensity (128 heads × 1 token = 128 rows per tile)
- Simplest code change — just reshape Q and adjust grid
- The current kernel already processes M=128 rows; at T=1+n_h=128, M=128 fits exactly
**Cons:**
- At larger T, M = T * n_h gets large. T=64, n_h=128 → M=8192 → 64 CTA tiles. But each tile still shares K/V.
- The M tile covers multiple heads' rows. The softmax operates per-row, so each row is one head's attention. This is fine — softmax already operates per-row.
**Key constraint:** T * n_h must be a multiple of the M tile size (128). At T=1, n_h=128 → M=128 ✅. At T=2, n_h=128 → M=256 (2 tiles). At T=3, n_h=128 → M=384 (3 tiles). Always a multiple since n_h is a multiple of 128 and M_tile=128.
Wait — T * n_h = 3 * 128 = 384. 384/128 = 3 tiles. Each tile has 128 rows. The first 128 rows are heads 0-127 of token 0. The next 128 rows are heads 0-127 of token 1. The third tile is heads 0-127 of token 2. This works, but the Q TMA needs to know that rows come from different (token, head) pairs. The Q tensor layout must be `(T * n_h, head_dim)` in M-major order.
Actually, for DSV4 decode, T is typically 1 (single token). So M = n_h = 128 (Flash: 64, Pro: 128). At T=1, Flash needs M=64 (half a tile) and Pro needs M=128 (full tile). Flash with M=64 is a problem — the MMA tile is 128 rows. We'd need to handle the edge case or pad.
**For now, start with Strategy B (head as grid dim) and consider Strategy A as an optimization.**
### Strategy B: Head as Grid Dimension (CUTLASS reference approach)
Each CTA handles one query head. Grid = `(num_M_tiles, num_query_heads, batch)`.
The CuTeDSL TMA tensors need to be 4D for Q/O (batch, heads, seq, dim) and 3D for K/V (batch, seq, dim). The head dimension is handled by the grid — each CTA gets one head.
**Key:** K/V have `num_kv_heads=1` in the head dimension. The TMA atom for K/V is created from the full tensor (with head dim=1), and each CTA loads from head 0.
### Step 2: Grid Shape Computation
```python
M_tile = 128 # rows per CTA
num_M_tiles = math.ceil(T / M_tile)
grid = (num_M_tiles, num_query_heads, batch)
```
At decode (T=1): `grid = (1, 128, batch)`.
At prefill (T=64): `grid = (1, 128, batch)` (T=64 < 128, still 1 M tile per head).
### Step 3: Block Coordinate → Tensor Index Mapping
Inside the kernel, map `(block_idx_x, block_idx_y, block_idx_z)` to:
-`m_tile_idx = block_idx_x` → which M tile within this head
-`head_idx = block_idx_y` → which query head
-`batch_idx = block_idx_z` → which batch element
Then Q TMA loads from `Q[batch_idx, head_idx, m_tile_idx*M_tile:(m_tile_idx+1)*M_tile, :]`.
K/V TMA loads from `K[batch_idx, 0, :, :]` (head 0, all tokens).
O TMA stores to `O[batch_idx, head_idx, m_tile_idx*M_tile:(m_tile_idx+1)*M_tile, :]`.
### Step 4: TMA Tensor Construction
The TMA tensors are constructed from the PyTorch tensors before launch. For multi-head, Q needs to be indexed by head. CuTeDSL's `make_tiled_tma_atom_A` creates a TMA descriptor from the tensor shape.
**Key question:** Can the TMA descriptor be created once and reused across heads? Yes — if Q is contiguous in the head dimension, the TMA descriptor covers all heads, and each CTA indexes into its head via the coordinate.
Actually, the CUTLASS reference creates the TMA descriptor from the FULL tensor shape `(batch, n_h, T, head_dim)` and the CTA's block coordinate selects the right head/tile.
### Step 5: Per-CTA Q Offset
Each CTA needs to load Q for its specific head. The TMA load already handles batch and M-tile indexing. For the head dimension, the TMA tensor has a mode for it.
In the kernel:
```python
bidx, bidy, bidz = cute.arch.block_idx()
# Q tile: m_tile=bidx, head=bidy, batch=bidz
```
The TMA copy: `cute.copy(tma_q, mQ[(None, bidy, bidx, bidz)], sQ[(None, qh.index)])` — but this depends on the TMA tensor layout.
### Step 6: K/V Shared Across Heads (MQA Optimization)
For MQA, all CTAs in the head dimension (same `block_idx_x`, same `block_idx_z`, different `block_idx_y`) load the same K/V. Two approaches:
1.**Independent loads (simple):** Each CTA loads its own K/V. Wastes bandwidth but correct.
2.**Cluster-wide load (optimized):** Use `cluster_shape_mn = (1, num_query_heads, 1)` so all heads in a cluster share the same K/V SMEM. Requires cluster barriers.
**Start with independent loads.** The K/V are small at decode (s_k=128 or 512, head_dim=512). Even 128× loads of 128×512 BF16 = 16MB is fine.
### Step 7: Output TMA Store
O has shape `(batch, n_h, T, head_dim)`. Each CTA writes its head's output. The TMA store uses the same block coordinate mapping.
- [ ]**D2.1:** Add `num_query_heads` and `batch_size` to `FmhaKernel.__init__`
- Simple to add, but the grid change is blocked (see below)
- [ ]**D2.3–D2.6:** Multi-CTA grid with runtime block coordinates
- **BLOCKED:** `cute.local_tile` does not support runtime coordinates. Must use `cute.flat_divide` instead.
- **BLOCKED:** `flat_divide` + `epilogue_tma_store` layout mismatch. The epilogue pipeline expects `tCgC` from `local_tile`, but `flat_divide` produces a different coordinate system.
- **Requires:** Full refactor of `tma_partition` + `epilogue_tma_store` to work with `flat_divide`-based GMEM views. This means moving ALL GMEM tensor partitioning into the kernel (like CUTLASS reference does).
- **CUTLASS reference approach:** Uses `flat_divide` + `tma_partition` inside the TMA warp block, and a custom epilogue that handles the flat_divide coordinate system. Estimated 1-2 day effort.
The Q tensor must be laid out so the TMA can efficiently load per-head tiles. Two options:
**Option 1: (batch, n_h, T, head_dim) — head is a TMA mode**
- TMA descriptor covers all heads
- Each CTA selects its head via the block coordinate
- This is how the CUTLASS reference works: `o_shape = (s, d, ((h_r, h_k), b))`
**Option 2: (batch, T * n_h, head_dim) — heads packed into M**
- TMA descriptor is 2D (batch, M, head_dim)
- Each CTA selects its M tile
- No head mode in TMA — simpler but can't distinguish heads in the kernel
**We'll use Option 1** because it matches the CUTLASS reference and allows per-head LSE output.
### Grid vs. Persistent Tile Scheduler
The CUTLASS reference uses a persistent tile scheduler for load balancing. For DSV4 decode (T=1), each CTA has exactly one tile, so a static grid suffices. For prefill with variable sequence lengths, a persistent scheduler would help. **Start with static grid, add persistent scheduler later if needed.**
### MQA K/V Sharing
At decode T=1, n_h=128: 128 CTAs each load the same K/V. Each K/V tile is s_k × head_dim BF16 = 128 × 512 × 2 = 128KB. Total K/V read = 128 × 128KB = 16MB. HBM bandwidth on B200 is ~8TB/s, so 16MB takes ~2μs. The MMA compute per CTA is ~5μs. So K/V redundancy is not a bottleneck at decode.
At prefill T=64, s_k=1024: K/V = 1024 × 512 × 2 = 1MB per CTA. 128 CTAs × 1MB = 128MB. Still ~16μs, comparable to compute. Not a bottleneck.
**Start with independent K/V loads. Optimize later with cluster-wide sharing if profiling shows it's needed.**
---
## Risks and Mitigations
| Risk | Mitigation |
|------|-----------|
| TMA tensor layout mismatch for multi-head Q | Print shapes at trace time. Start with n_h=1 regression. |
