docs: revised Stage D/E plan — indexer removes paged TMA, one kernel for CSA/HCA/SWA, sink merge

2026-05-23 03:10:41 +00:00
parent bf0bb8241c
commit 3be9d6ed8c
1 changed files with 195 additions and 29 deletions
--- a/README.md
+++ b/README.md
@@ -145,9 +145,13 @@ 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 | ⚠️ SINGLE-TILE OK, MULTI-TILE 3% ERROR | n=128 cos 0.973. n=256 cos 0.793. TMEM layout mismatch between MMA C-fragment and get_tmem_load_op. See below. |
-| D | TODO | Full decode attention: paged KV cache, multi-query, causal mask |
-| E | TODO | Production kernel: extract into dsv4/kernels/attention/, PyTorch custom op, vLLM bridge |
+| C | ⚠️ SINGLE-TILE OK, MULTI-TILE 3% ERROR | n=128 cos 0.973. n=256 cos 0.793. TMEM layout mismatch. See below. |
+| D1 | TODO | Parameterize HEAD_DIM (64 → 512) |
+| D2 | TODO | Multi-query grid with head packing (128 Q heads, MQA) |
+| D3 | TODO | SWA sequence length mask (swa_lens per batch) |
+| D4 | TODO | Causal mask on SWA branch only |
+| D5 | TODO | SWA + sink merge (two MMA loops, log-sum-exp merge) |
+| E1-E7 | TODO | Production extraction (class, custom op, cache, cleanup) |

 ---

@@ -389,50 +393,212 @@ Col 128+:   O (PV acc, 64 FP32, rescale via Ld32x32bOp Repetition(16))

 ---

-## Stage D: Full Decode Attention (TODO)
+## Stage D: Full Decode Attention (revised May 23)

-### What Stage D Needs
+### Key Insight: The Indexer Solves Paging Upstream

-The FMHA pipeline (Stages A-C) processes a single query row against K/V. Stage D makes it a **real attention kernel** by adding:
+The indexer now hands the kernel `selected_kv: [T, top_k, head_dim] BF16` — a **dense, materialized, dequantized** K/V tile. FMHA sees a dense `[T, top_k, 512]` tile, exactly like Stage A/B's existing `k` and `v` inputs. **The kernel doesn't need to know it's sparse.** Paged TMA, scattered HBM reads, FP8 dequantization — all handled by `gather_selected_kv` upstream.

-1. **Paged KV cache** — read K/V from a paged cache (not contiguous GMEM). Each page is a fixed-size block (e.g., 128 tokens). Page tables map logical → physical addresses. This is how vLLM handles variable sequence lengths without reallocation.
+The SWA branch is the only "irregular" thing: it reads from the state cache's ring buffer with a position mask. SWA is small (`n_win=128` per query), so it's a separate fused branch with a sink-weighted merge.

-2. **Multi-query attention (MQA) / Grouped-query attention (GQA)** — DSV4 has 128 heads sharing a (T, 512) KV latent. All heads share the same K/V, so Q is (128, 512) but K/V are (T, 512). One K/V load, 128 Q heads. The current single-head kernel needs a Q-head loop.
+**One FMHA kernel serves all three DSV4 attention types:**
+- **CSA:** `compressed_kv` = top-k from indexer, `swa_kv` from cache → sink merge
+- **HCA:** `compressed_kv` = all classical pool entries (gather-all mode), `swa_kv` from cache → sink merge
+- **SWA-only (Flash layers 0-1):** `compressed_kv` = empty (`top_k=0`), only SWA runs. Sink merge degenerates to just `o_swa` after renormalization.

-3. **Causal mask** — for autoregressive decode, each query position can only attend to K/V positions ≤ its own. This is a simple upper-triangular mask on the QK scores (or just limiting the K/V range loaded per query row).
+### Build Order

-4. **CSA/HCA sparse attention** — DSV4 uses compressed sparse attention (m=4) and heavily compressed attention (m=128). The indexer selects top-k blocks, and attention only runs on those blocks. This changes the KV tile iteration from sequential to sparse.
+**D1 — Parameterize HEAD_DIM** (~½ day)

-5. **SWA branch + sink merge** — Sliding window attention (window=128) runs alongside sparse attention. Sink weights combine the two branches: `O = sink * O_sparse + (1-sink) * O_swa`.
+Currently hardcoded at 64. Promote to constructor arg, thread through `_setup`. Test at 64, then 512 (DSV4's real value).

-### Architecture Plan
+Risk at HEAD_DIM=512: TMEM column budget. `_setup` already does `find_tmem_tensor_col_offset(tOtO)` dynamically. Verify the total fits in 512 TMEM columns. If not, reduce `kv_stage` from 2 to 1 (lose K/V double-buffering) before sacrificing math.
+
+Done when: identical result at HEAD_DIM=64 (regression), passes at HEAD_DIM=512 against FP32 oracle.
+
+**D2 — Multi-query grid with head packing** (~1 day)
+
+Grid changes from `(1, 1, 1)` to `(num_q_blocks, 1, batch)`. DSV4 is MQA — all `n_h=128` query heads share the same K/V. The query-head axis is folded into the M dimension of the Q tile: `M_tile = 128` covers `M = T * n_h` rows. At decode T is small (1-16), so packing heads into M fills the MMA. At prefill T=64, M is already 8192 with heads packed.
+
+Done when: batch=4, T=64, n_h=128, num_kv_heads=1 produces correct attention against FP32 oracle.
+
+**D3 — SWA sequence length mask** (~½ day)
+
+The indexer's `top_k` is fixed (512 for Flash, 1024 for Pro). Compressed-K input is always `[T, top_k, head_dim]` with the same `top_k` at compile time.
+
+What varies: the SWA window holds up to `n_win=128` tokens but starts with fewer. Add `swa_lens: [batch] int32` as kernel input. Mask SWA-branch logits to `-inf` where `swa_idx >= swa_lens[b]`.
+
+Done when: batched input with varying SWA fill levels (some requests at position 50, some at 5000) produces correct masked output.
+
+**D4 — Causal mask on SWA branch** (~½ day)
+
+The compressed K the indexer selects is already from `s < floor(t/m)` (paper eq. 17). The indexer enforces causality at selection time. FMHA sees only causally-valid candidates. **The main path has no mask.**
+
+The SWA branch needs a causal mask within the window. Add `is_causal: bool` constructor flag, apply `swa_idx > q_pos` masking to `-inf` in the SWA pass.
+
+Done when: prefill mode produces correct output with the causal mask applied to SWA.
+
+**D5 — SWA + sink merge** (~2-3 days) ← the real new work
+
+Per `dsv4/ops/decode_sparse.py`:
+```
+o = (exp(lse_sparse) * o_sparse + exp(attn_sink) * exp(lse_swa) * o_swa)
+    / (exp(lse_sparse) + exp(attn_sink) * exp(lse_swa))
+```
+
+Both branches must produce **un-normalized `o`** AND **`lse = log(row_sum) + row_max`**. Stage C currently normalizes by `1/row_sum` in the epilogue — that must change.
+
+Three structural changes:
+1. **TMEM grows:** two O accumulators + two row_max/row_sum per row. Roughly doubles TMEM column usage. Verify it fits.
+2. **Q is loaded once, used twice.** Free win — largest input stays in SMEM.
+3. **K and V have two sources.** Compressed K/V from contiguous BF16 input. SWA K/V from state cache — for now, dequantize SWA in a small prep kernel before FMHA, let FMHA see two contiguous BF16 sources. Optimize later.
+
+Sub-steps:
+- **5a:** Modify Stage C to emit un-normalized `o` + `lse` instead of normalized `o`. (Keep normalize as a flag so standalone tests still work.)
+- **5b:** Run kernel twice externally (once with compressed_kv, once with swa_kv), merge results in Python. **End-to-end correctness without touching kernel structure.** Hours of work.
+- **5c:** Fuse the two passes into one kernel launch. Q stays in SMEM, two MMA loops sequentially.
+- **5d:** Fuse the merge into the kernel epilogue.
+
+Done when: end-to-end kernel produces correct attention against FP32 oracle that does sparse+SWA+sink merge.
+
+**~~D5 (old) paged TMA~~ — REMOVED.** The indexer + gather handles all paging upstream.
+
+### Kernel Architecture (after D5)

 ```
-Input: Q (num_heads, 512, 1), K/V from paged cache
+Input: Q [T, n_h, 512], compressed_kv [T, top_k, 512], swa_kv [batch, n_win, 512]
+       swa_lens [batch], sink_logits [n_h], request_ids [T]
  │
-  ├─ Q-head loop (128 heads, shared K/V)
-  │   │
-  │   ├─ Causal mask: only load K/V up to current position
-  │   │
-  │   ├─ CSA path: load top-k compressed KV blocks → sparse FMHA
-  │   │
-  │   ├─ SWA path: load last 128 tokens → windowed FMHA
-  │   │
-  │   └─ Sink merge: O = sink * O_csa + (1-sink) * O_swa
+  ├─ Load Q to SMEM (once)
  │
-  └─ Output: O (num_heads, 512, 1) → inv RoPE → o_a BMM → o_b projection
+  ├─ Loop 1: compressed KV (top_k tokens)
+  │   QK → online softmax → PV → O_sparse, lse_sparse in TMEM
+  │
+  ├─ Loop 2: SWA window (n_win tokens, masked by swa_lens)
+  │   QK → online softmax → PV → O_swa, lse_swa in TMEM
+  │
+  └─ Sink merge epilogue:
+      O = (exp(lse_sparse) * O_sparse + exp(sink) * exp(lse_swa) * O_swa)
+          / (exp(lse_sparse) + exp(sink) * exp(lse_swa))
 ```

 ### Reference Files

- KV cache spec: `dsv4/cache/` (stubs)
- CSA/HCA attention: `dsv4/reference/csa_attention.py`
- vLLM Blackwell attention backend: `/root/dsv4-nvfp4-workspace/vllm/vllm/attention/backends/`
+- Sink merge spec: `dsv4/ops/decode_sparse.py` (formula)
+- SWA decode: `dsv4/ops/decode_swa.py`
+- Attention reference: `dsv4/reference/attention.py`
+- CSA attention: `dsv4/reference/csa_attention.py`

-### Dependencies
+### Stage C Note

- Stage C precision fix (correction_epilog) should be resolved before production, but doesn't block Stage D development
- Paged KV cache read/write kernels (Stage D.1) are the first concrete deliverable
+When implementing D5a, Stage C's epilogue changes from "multiply by 1/row_sum" to "emit un-normalized o + lse". Defer this until D5. Through D1-D4, keep Stage C normalize as-is and test as standalone dense FMHA.
+
+---
+
+## Stage E: Production Extraction (revised May 23)
+
+### E1 — File placement
+
+`dsv4/kernels/attention/fmha.py`. Class: `FmhaKernel`. Constructor takes all dimensions and dtypes, no module-level constants. Drop `if __name__ == "__main__"` test block.
+
+### E2 — Constructor signature
+
+```python
+class FmhaKernel:
+    def __init__(
+        self,
+        head_dim: int,           # 512 for DSV4
+        num_query_heads: int,     # 128 for Pro, 64 for Flash
+        sliding_window: int,      # 128
+        top_k: int,               # 512 (Flash) or 1024 (Pro)
+        q_dtype=BFloat16,
+        kv_dtype=BFloat16,
+        o_dtype=BFloat16,
+        qk_acc_dtype=Float32,
+        pv_acc_dtype=Float32,
+        is_causal: bool = False,  # affects SWA mask only
+        cta_group: tcgen05.CtaGroup = tcgen05.CtaGroup.ONE,
+        cluster_shape_mn: tuple = (1, 1),
+    ):
+```
+
+All architecture-level shapes from config flow into the constructor. No FMHA-internal magic numbers.
+
+### E3 — Call signature
+
+```python
+def __call__(
+    self,
+    q: torch.Tensor,            # [T, n_h, head_dim] BF16
+    compressed_kv: torch.Tensor, # [T, top_k, head_dim] BF16 — from indexer gather
+    swa_kv: torch.Tensor,        # [batch, n_win, head_dim] BF16 — from cache prep
+    swa_lens: torch.Tensor,      # [batch] int32
+    sink_logits: torch.Tensor,   # [n_h] FP32
+    request_ids: torch.Tensor,   # [T] int32 — maps query to its SWA slot
+    o: torch.Tensor,             # [T, n_h, head_dim] BF16 — preallocated
+    stream: cuda.CUstream,
+):
+```
+
+Notably absent: block_table, paged KV, inv_scale, FP8 dequant. All handled upstream.
+
+### E4 — Kernel cache + warmup
+
+Mirror `dsv4/ops/gemm_runner.py`'s `_compiled_kernel_cache`. Key on `(head_dim, num_query_heads, top_k, is_causal, ...)`. Pre-allocate at warmup, reuse at call. For DSV4, the cache has at most ~2 entries (Flash/Pro × causal/non).
+
+### E5 — torch.library custom op
+
+```python
+@torch.library.custom_op("dsv4::sparse_fmha_with_swa", mutates_args=("o",))
+def sparse_fmha_with_swa_op(
+    q: torch.Tensor,
+    compressed_kv: torch.Tensor,
+    swa_kv: torch.Tensor,
+    swa_lens: torch.Tensor,
+    sink_logits: torch.Tensor,
+    request_ids: torch.Tensor,
+    o: torch.Tensor,
+    runner_id: int,
+) -> None:
+    runner = get_runner(runner_id)
+    runner._run_impl(q, compressed_kv, swa_kv, swa_lens, sink_logits, request_ids, o)
+```
+
+Mutates `o` (preallocated buffer). Consistent with cudagraphs.
+
+### E6 — Reference parity hook
+
+`dsv4/reference/attention.py` stays as the FP32 oracle. New test: `tests/unit/test_fmha_kernel.py`.
+
+```python
+def test_sparse_fmha_matches_spec(T=64, n_h=128, top_k=1024, n_win=128, hd=512):
+    q = torch.randn(T, n_h, hd, dtype=torch.bfloat16, device='cuda')
+    ck = torch.randn(T, top_k, hd, dtype=torch.bfloat16, device='cuda')
+    swa = torch.randn(4, n_win, hd, dtype=torch.bf16, device='cuda')
+    swa_lens = torch.tensor([128, 50, 128, 75], dtype=torch.int32)
+    sink = torch.randn(n_h, device='cuda')
+    req_ids = torch.randint(0, 4, (T,), dtype=torch.int32)
+
+    # Oracle: pure FP32 spec
+    o_sparse, lse_sparse = attention_with_lse_f32(q, ck, ck)
+    o_swa, lse_swa = attention_swa_with_lse_f32(q, swa, swa, swa_lens, req_ids)
+    e_sink = sink.exp()
+    num = lse_sparse.exp().unsqueeze(-1) * o_sparse \
+        + e_sink[None, :, None] * lse_swa.exp().unsqueeze(-1) * o_swa
+    den = lse_sparse.exp() + e_sink[None, :] * lse_swa.exp()
+    expected = num / den.unsqueeze(-1)
+
+    # Kernel
+    o = torch.empty_like(expected, dtype=torch.bfloat16)
+    fmha = FmhaKernel(head_dim=hd, num_query_heads=n_h, sliding_window=n_win, top_k=top_k)
+    fmha(q, ck, swa, swa_lens, sink, req_ids, o, stream=...)
+
+    torch.testing.assert_close(o.float(), expected, atol=5e-3, rtol=5e-3)
+```
+
+### E7 — Cleanup
+
+Delete all debug test files. `test_fmha_v3.py` becomes `dsv4/kernels/attention/fmha.py`. Only `tests/unit/test_fmha_kernel.py` remains as the attention test.

 ---