From 3be9d6ed8c58a0ff3d6b34b9b4cbdae09dac7a0b Mon Sep 17 00:00:00 2001 From: biondizzle Date: Sat, 23 May 2026 03:10:41 +0000 Subject: [PATCH] =?UTF-8?q?docs:=20revised=20Stage=20D/E=20plan=20?= =?UTF-8?q?=E2=80=94=20indexer=20removes=20paged=20TMA,=20one=20kernel=20f?= =?UTF-8?q?or=20CSA/HCA/SWA,=20sink=20merge?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 224 +++++++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 195 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index 0ddbe710..eb5b6dd9 100644 --- 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. ---