diff --git a/NEXT_PRIORITIES.md b/NEXT_PRIORITIES.md new file mode 100644 index 00000000..605a1c14 --- /dev/null +++ b/NEXT_PRIORITIES.md @@ -0,0 +1,196 @@ +# NEXT PRIORITIES — verified against code, not against status docs + +**Why this file exists:** the agent's `CURRENT_ISSUE.md` tracks the *thing it's +currently building*, not the *state of the production path*. Those have diverged. +This doc resets the priorities against the actual code in `dsv4/`. + +**Method:** every "current state" claim below was read directly off the source. +Where the agent's notes disagree, the code wins. Per doctrine rule 3. + +--- + +## Verified state (read off the code, 2026-05-30) + +| Item | Claimed | Verified | Evidence | +|---|---|---|---| +| Indexer FP4 dequant (P1) | "fixed" | ✅ **DONE** | `dsv4/kernels/indexer/indexer_score_topk.cu:41` has `E2M1_LUT[8] = {0, 0.5, 1, 1.5, 2, 3, 4, 6}`. Second copy in `kernels/cuda/` also fixed (line 17). Live path via `score_topk.py:29` builds the indexer copy. | +| Python KV merge killed (P2) | "milestones 4–5 in progress" | ❌ **NOT DONE** | `production.py:236–296` still has the segment loop, `torch.cuda.synchronize()` at `:279` *inside the inner loop*, eager exp/log merge at `:286–294`. Nothing wired to any 6-warp variant. | +| 6-warp raw-CUDA FMHA | "multihead milestone 5 done" | ⚠ **WORKS IN ISOLATION**, T=1 decode, HD∈{16,64,128,256}, single KV tile, cos 0.999997+. Genuinely Blackwell-native: `tcgen05.mma` SS, UMMA descriptors, TMEM alloc + `tcgen05.ld.sync.aligned.32x32b.x8.b32`. **Not called from production.** | +| TMA descriptor blocker | "INVALID_VALUE on B200" | ⚠ **PARTIALLY DEMYSTIFIED** | `docs/cuda13_tma_notes.md` and `memory/2026-05-29-tma-async.md`: CUDA 13 needs **byte strides**, not element strides. Descriptors now create OK. But `cp.async.bulk.tensor.{2d,3d}` **hangs** — mbarrier never signals. Root cause "unknown" per agent. | +| Multi-row softmax T>32 | "blocked on TMEM read" | ⚠ Same | `fmha_6warp_multihead.cuh:196–218` softmax is single-row (warp 0, row 0). Agent says fix is `16x256b.x1` instead of `32x32b.x8`. That's a **guess** until the TMEM column layout is printed and confirmed. | +| FMHA file count | n/a | 🚨 **7 .cuh variants** | `fmha_6warp.cuh`, `_multihead`, `_multirow`, `_tma`, `_tma_multirow`, `_tma_multitile`, `_tma_multirow_multitile`. Plus `fmha_sm100_tc.cuh`, `fmha_epilogue_sm100.cuh`, `fmha_tma.cuh`, `fmha_umma_desc.cuh`. **None integrated to production.** Going wide instead of deep. | +| FP4 attention epilogue | "blocked" | Correctly blocked | `fmha_common.cuh:32`, `fmha_epilogue_sm100.cuh:27`, `fmha.py:51` all note dependence on epilogue rewrite. Real dependency, not procrastination. | + +**The pattern.** P1 actually shipped. Everything else is *exploration without +integration*. Seven .cuh variants were forked along three optimization axes (TMA, +multirow, multitile), but the production path has not advanced one inch — same +Python merge, same per-tile sync, same launch count per decoded token. This is +the failure mode you flagged: the agent reports milestones inside a sandbox while +the hot path it's supposed to fix is untouched. + +--- + +## Priority order (do these in this sequence — do not parallelize) + +### **P3 — Wire `fmha_6warp_multihead.cuh` into `production.py`, decode-only** + +**This is the most-leverage move, not the most-exciting one.** The kernel works +in isolation for the exact shape that dominates decode (T=1, single KV tile, +multi-head, HD∈{64,128,256}). Wire it. Stop forking variants. + +**Definition of done:** +1. `production.py:_run_fmha_segmented` has a fast path: if `T == 1` and + `n_segments == 1`, call the 6-warp kernel via a torch custom op. Single + launch. No `torch.cuda.synchronize()` on the hot path. No per-tile + allocations. +2. Numerical parity gate: `cos ≥ 0.999998` against the current Python-merge path + on the cases it handles. Not "looks fine" — bitwise diffable test. +3. **Launch count measurement** before/after on one decoded V4-Pro CSA layer. + Record with Nsight Systems or `cudaLaunchKernel` counter. Target on the new + path: 1 kernel launch, 0 `cudaDeviceSynchronize` on the hot path. This is the + number that proves P2 progress; cosine alone does not. +4. The slow path (multi-segment, prefill T>1) **stays as-is** for now. Don't + touch it until the decode path is integrated and measured. + +**Failure modes to watch for** (call them out if you see them): +- Agent creates an 8th .cuh variant. The answer is "integrate, don't fork." +- Agent regresses the cosine to "make integration easier." Parity is the gate. +- Agent removes the Python fallback before the fast path covers all shapes used. + +### **P4 — Resolve the TMA hang with print-and-diff, NOT another guess** + +Memory note ends with "root cause unknown" and three speculative options. That's +the prompt to apply doctrine rule 3, not to pick one. + +**Definition of done:** +1. **Dump the descriptor bytes** from a working CuTeDSL TMA path (the existing + CuTeDSL FMHA already runs TMA correctly — that's the oracle). +2. **Dump the descriptor bytes** from the raw-CUDA `cuTensorMapEncodeTiled` + path that hangs. +3. **`memcmp` them**. Document the byte-level differences in a `.md` paper trail + alongside `cuda13_tma_notes.md`. The hang is almost certainly one of: swizzle + mode, fill mode, interleave layout, or oobFill — the 5 enum fields the API + takes after the strides. Print every field of both descriptors side by side. +4. Once they match: re-run, confirm no hang. If they don't match: code to the + diff, not to a new theory. + +**Failure modes to watch for:** +- Agent reaches for "manually construct TMA descriptor bytes" (option B from the + memory note) without doing the diff first. That's a bigger guess, not a smaller one. +- Agent declares this "deferred" and moves on. TMA is on the critical path for + prefill / long-context throughput; "decode works without TMA" is true but + doesn't generalize. + +### **P5 — In-kernel online softmax across KV tiles (the real P2)** + +After P3, the decode fast path bypasses the Python merge for `n_segments==1`. +But CSA with `top_k=1024` always has `n_segments=8`, so the moment top_k > 128 +we're back on the slow Python path. This priority extends the 6-warp kernel to +loop KV tiles internally with FlashAttention-2 running max/sum. + +**Definition of done:** +1. The 6-warp kernel (the **same** file, not a new variant) accepts a + `kv_tile_count` parameter and loops over it internally, maintaining + `(row_max, row_sum)` rescale across tiles. Standard FA2 shape. +2. Single launch handles any `n_segments`. The Python merge in `production.py` + for multi-tile is **deleted**, not "kept as fallback." +3. Parity gate: `cos ≥ 0.999998` vs the now-deleted Python merge, captured in a + regression test before deletion. +4. Launch count on V4-Pro CSA layer: still 1, regardless of `top_k`. The + `cudaDeviceSynchronize` calls go to zero, period. + +**Failure modes to watch for:** +- Agent writes an `fmha_6warp_multikvtile.cuh` instead of extending the chosen file. +- Agent moves the rescale to host-side "for clarity." The rescale must live in + the kernel or the launch-count guarantee breaks. + +### **P6 — One-way TMEM→regs→SMEM→GMEM epilogue, with FP4 hook** + +Unlocks NVFP4-1.2 (FP4 output fusion) and multi-CTA grids. Pattern already runs +correctly in `dsv4/kernels/gemm/dense.py`; the symbols are already imported into +`fmha.py` (lines 71–72) and **unused**. The work is wiring, not invention. + +**Definition of done:** +1. The 6-warp kernel's epilogue uses `epilogue_tmem_copy_and_partition` + + `epilogue_smem_copy_and_partition` + TMA store, with an `epilogue_op` lambda + slot (constexpr) for fusion. Same shape as MoE GEMM. +2. `epilogue_op = lambda x: x` ships as default. FP4 pack lambda lands as a flag, + off by default, behind its own test. +3. Multi-CTA grid smoke test: M ≥ 256 prefill, flat_divide coords accepted, no + crash, correct numerics. + +**Failure modes to watch for:** +- Agent re-introduces `epilogue_tma_store` "because it's simpler." That's the + blocker we're removing — don't re-add it. +- Agent enables FP4 pack at the same time as the epilogue rewrite. Two changes, + one test. Land the epilogue first, FP4 fusion second. + +### **P7 — Multi-row softmax T>32, by printing the TMEM column layout** + +The agent's plan ("use `16x256b.x1`") is a guess. May be right; may not be. +Before changing the instruction: + +**Definition of done:** +1. **Print** the TMEM column map for HD=256, T=128 case: for each (warp, lane, + tmem column), which (row, col) of S does it own? Write the observed map into + a `.md` doc. +2. Pick the TMEM load instruction that matches the observed map. If it's + `16x256b.x1`, fine — but with the table backing the choice. +3. Parity gate: `cos ≥ 0.999998` for T∈{1, 32, 64, 128} all in the same kernel. + +**Failure modes to watch for:** +- Agent picks the instruction first, then "interprets the layout to match." + Layout first, instruction second. + +### **P8 — Consolidate: delete 6 of the 7 6-warp variants** + +After P3–P7, exactly one variant should exist. The other six are landmines for +the next agent (and for you when you context-switch back in three weeks). + +**Definition of done:** `ls dsv4/kernels/attention/fmha_6warp*.cuh` returns one +file. Tests updated to point at it. `git rm` for the rest. No "archive/" folder. + +--- + +## What is *not* on this list, and why + +- **MoE / GEMM stack:** already production-grade per audit. Don't touch. +- **mHC implementation:** correct per paper, working. +- **CSA compressor (overlapped 2m):** correct per paper, working. +- **Router (`sqrt(softplus)` + aux-loss-free):** correct per paper, working. +- **MegaMoE EP overlap (dispatch/combine waves):** correctly scoped out — + that's a multi-node EP concern, not a single-GPU kernel concern. +- **Reasoning effort modes, Quick Instruction, OPD:** post-training surface, not + kernel surface. + +--- + +## DOCTRINE — every priority above is gated on these, no exceptions + +1. **CuTeDSL/CUTLASS wall → raw CUDA C++, NOT Python.** The Python KV merge in + `production.py` is the cautionary tale. The 6-warp raw-CUDA kernel is the + correct fallback for that wall — but it has to be *integrated*, not + exhibited. + +2. **Raw CUDA ≠ scalar math.** Every priority above keeps `tcgen05` / UMMA / + TMEM / TMA / warp-level reductions. No scalar dot products as "temporary + simplification." That's how we got the indexer LUT bug — a "temporary" scalar + oracle that was wrong and trusted as a reference. + +3. **Print, don't guess. Code to the data.** Two specific applications this + round: + - **P4 (TMA hang):** dump descriptor bytes, diff, code to the diff. Do not + guess at "format mismatch." + - **P7 (TMEM layout):** print the observed (warp, lane) → (row, col) map, + pick the instruction from the map. Do not pick the instruction first. + +4. **Integration over exploration.** Seven .cuh variants is the diagnostic. The + priority order above is deliberately "wire the one that works" before "make + it handle more shapes." A working kernel not on the production path is worth + zero. Resist the urge to fork a new file for each new concern; extend the + chosen file or step back to plan. + +5. **Falsifiable gates only.** Every "definition of done" above has a number + (cosine, launch count, file count) or a binary check. "Looks fine," "milestone + complete," and "in progress" are not gates. If a status doc says "DONE" + without a number next to it, the doctrine reads it as "NOT DONE." \ No newline at end of file diff --git a/2026-05-28_STATUS.md b/archived_plans/2026-05-28_STATUS.md similarity index 100% rename from 2026-05-28_STATUS.md rename to archived_plans/2026-05-28_STATUS.md diff --git a/CURRENT_ISSUE.md b/archived_plans/CURRENT_ISSUE.md similarity index 100% rename from CURRENT_ISSUE.md rename to archived_plans/CURRENT_ISSUE.md diff --git a/CURRENT_ISSUE_FROM_OUTSIDE_CONSULTANT.md b/archived_plans/CURRENT_ISSUE_FROM_OUTSIDE_CONSULTANT.md similarity index 100% rename from CURRENT_ISSUE_FROM_OUTSIDE_CONSULTANT.md rename to archived_plans/CURRENT_ISSUE_FROM_OUTSIDE_CONSULTANT.md diff --git a/NVFP4_ROADMAP.md b/archived_plans/NVFP4_ROADMAP.md similarity index 100% rename from NVFP4_ROADMAP.md rename to archived_plans/NVFP4_ROADMAP.md diff --git a/PREVIOUS_CONVERSATION_2.md b/archived_plans/PREVIOUS_CONVERSATION_2.md similarity index 100% rename from PREVIOUS_CONVERSATION_2.md rename to archived_plans/PREVIOUS_CONVERSATION_2.md diff --git a/PRIORITY2.md b/archived_plans/PRIORITY2.md similarity index 100% rename from PRIORITY2.md rename to archived_plans/PRIORITY2.md diff --git a/ROADMAP.md b/archived_plans/ROADMAP.md similarity index 100% rename from ROADMAP.md rename to archived_plans/ROADMAP.md diff --git a/STAGE_D.md b/archived_plans/STAGE_D.md similarity index 100% rename from STAGE_D.md rename to archived_plans/STAGE_D.md diff --git a/STAGE_D_REMAINING.md b/archived_plans/STAGE_D_REMAINING.md similarity index 100% rename from STAGE_D_REMAINING.md rename to archived_plans/STAGE_D_REMAINING.md diff --git a/dsv4/kernels/attention/fmha_multihead_launch.cu b/dsv4/kernels/attention/fmha_multihead_launch.cu new file mode 100644 index 00000000..ebd68c42 --- /dev/null +++ b/dsv4/kernels/attention/fmha_multihead_launch.cu @@ -0,0 +1,167 @@ +/** + * DSV4 FMHA Multi-Head — PyTorch launch wrapper. + * + * Bridges the raw CUDA 6-warp multi-head FMHA kernel to PyTorch tensors. + * Uses uint16_t (bf16_t) inside the kernel, c10::BFloat16 at the API boundary. + * Casts are zero-cost reinterpret_casts (same bit layout). + * + * Grid: dim3(1, n_h, batch_size) + * Each CTA processes one (head, batch) pair independently. + * Decode-only: T=1, single KV segment (s_k <= 128). + */ + +#include +#include +#include + +// Pull in the kernel (uses uint16_t bf16_t internally, no ATen deps) +#include "fmha_common.cuh" +#include "fmha_umma_desc.cuh" +#include "fmha_6warp_multihead.cuh" + +namespace dsv4::kernels::attention { + +/** Compute dynamic SMEM size for the 6-warp multi-head kernel. */ +static int compute_smem_multihead(int hd) { + // From fmha_6warp_multihead.cuh SMEM layout: + // sTmemBase(4) + sRowMax(4) + sRowSum(4) + padding to 16B + // + sQ0: 128 * hd * 2 + // + sK0: 128 * hd * 2 (128B aligned) + // + sPk: 128 * 16 * 2 = 4096 (128B aligned, always TILE_SZ for the P buffer) + // + sV: 16 * 16 * 2 = 512 (128B aligned, V_SUB_SZ) + // + s_p_vals: 128 * 4 = 512 + // Over-allocate: sum the worst case + int base = 4 + 4 + 4; // tmem_base, row_max, row_sum + base = (base + 15) & ~15; // align to 16B + int sQ0_offset = base; + int sQ0_sz = 128 * hd * 2; + int sK0_offset = sQ0_offset + sQ0_sz; + int sK0_sz = 128 * hd * 2; + // Pk needs 128B alignment + int sPk_offset = ((sK0_offset + sK0_sz) + 127) & ~127; + int sPk_sz = 128 * 16 * 2; // TILE_SZ in BF16 + // V needs 128B alignment + int sV_offset = ((sPk_offset + sPk_sz) + 127) & ~127; + int sV_sz = 16 * 16 * 2; // V_SUB_SZ in BF16 + int sp_vals_offset = sV_offset + sV_sz; + int sp_vals_sz = 128 * 4; // SK_TILE * sizeof(float) + int total = sp_vals_offset + sp_vals_sz; + // Pad to 128B alignment + total = (total + 127) & ~127; + return total; +} + +/** + * Launch the 6-warp multi-head FMHA kernel for decode (T=1). + * + * Q: (batch, n_h, 1, hd) contiguous BF16 + * K: (batch, n_kv, N, hd) contiguous BF16 + * V: (batch, n_kv, hd, N) contiguous BF16 — NOTE: V is (hd, N) per head! + * Returns: (O, LSE) where O is (batch, n_h, 1, hd) BF16, LSE is (batch, n_h, 1) FP32 + * + * For MQA: n_kv = 1, k/v head stride = 0 (broadcast) + * For GQA: n_kv < n_h, k/v head stride reflects the KV group layout + */ +std::tuple fmha_multihead_decode_cuda( + torch::Tensor q, // (batch, n_h, 1, hd) BF16 + torch::Tensor k, // (batch, n_kv, N, hd) BF16 + torch::Tensor v, // (batch, n_kv, hd, N) BF16 + double scale, // 1/sqrt(hd) + int64_t n_comp, // compressed KV length (0 = no compression) + int64_t swa_len, // sliding window length (0 = no SWA mask) + bool is_causal, // causal mask on SWA + c10::optional attn_sink // (batch, n_h) FP32 — per-head sink bias +) { + TORCH_CHECK(q.dim() == 4, "Q must be 4D (batch, n_h, 1, hd)"); + TORCH_CHECK(q.size(2) == 1, "Decode mode requires T=1"); + TORCH_CHECK(k.dim() == 4, "K must be 4D (batch, n_kv, N, hd)"); + TORCH_CHECK(v.dim() == 4, "V must be 4D (batch, n_kv, hd, N)"); + + const int B = q.size(0); + const int n_h = q.size(1); + const int hd = q.size(3); + const int n_kv = k.size(1); + const int N = k.size(2); + + TORCH_CHECK(N <= 128, "Decode fast path requires N <= 128 (single KV tile)"); + TORCH_CHECK(hd == 64 || hd == 128 || hd == 256, + "Unsupported head_dim: ", hd, " (decode fast path: 64/128/256)"); + + auto opts_bf16 = q.options().dtype(torch::kBFloat16); + auto opts_f32 = q.options().dtype(torch::kFloat32); + auto o = torch::zeros({B, n_h, 1, hd}, opts_bf16); + auto lse = torch::zeros({B, n_h, 1}, opts_f32); + + // Build FmhaParams + FmhaParams params; + params.q = reinterpret_cast(q.data_ptr()); + params.k = reinterpret_cast(k.data_ptr()); + params.v = reinterpret_cast(v.data_ptr()); + params.o = reinterpret_cast(o.data_ptr()); + params.lse = lse.data_ptr(); + params.s_k = N; + params.scale = static_cast(scale); + params.head_dim = hd; + + // Strides in BF16 elements + params.q_head_stride = q.stride(1); // stride between heads in the batch + params.q_batch_stride = q.stride(0); // stride between batch items + params.k_head_stride = k.stride(1); + params.k_batch_stride = k.stride(0); + params.v_head_stride = v.stride(1); + params.v_batch_stride = v.stride(0); + params.o_head_stride = o.stride(1); + params.o_batch_stride = o.stride(0); + params.lse_head_stride = lse.stride(1); + params.lse_batch_stride = lse.stride(0); + + // Zero out head strides for MQA (n_kv == 1 means all Q heads share same K/V) + // Actually: the kernel uses head_idx * k_head_stride, which for MQA should give + // the same address for all heads. If k has n_kv=1, then stride(1) may be N*hd + // (the single head's total). For MQA we want k_head_stride=0 so all heads point + // to the same K. The caller should ensure K has shape (batch, 1, N, hd) for MQA, + // and we set k_head_stride=0 when n_kv==1. + if (n_kv == 1) { + params.k_head_stride = 0; + params.v_head_stride = 0; + } + + // For n_comp=0, the kernel doesn't need sink bias, but we pass nullptr + // The kernel's epilogue only writes LSE if lse_head != nullptr + + // Copy params to device (constant memory or managed) + // Simpler: just pass by value through kernel parameter struct + // CUDA kernels can take structs by value + int smem = compute_smem_multihead(hd); + + dim3 grid(1, n_h, B); + dim3 block(NTHREADS); + + // Launch with template instantiation + #define LAUNCH(D) \ + fmha_6warp_multihead_kernel<<>>(params) + + cudaError_t err; + if (hd == 64) { + LAUNCH(64); + } else if (hd == 128) { + LAUNCH(128); + } else if (hd == 256) { + LAUNCH(256); + } else { + TORCH_CHECK(false, "Unsupported head_dim: ", hd); + } + #undef LAUNCH + + err = cudaGetLastError(); + TORCH_CHECK(err == cudaSuccess, "Kernel launch failed: ", cudaGetErrorString(err)); + + return std::make_tuple(o, lse); +} + +} // namespace dsv4::kernels::attention + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("fmha_multihead_decode", &dsv4::kernels::attention::fmha_multihead_decode_cuda, + "DSV4 FMHA multi-head decode kernel (6-warp, T=1, single KV tile)"); +} diff --git a/dsv4/kernels/attention/fmha_multihead_op.py b/dsv4/kernels/attention/fmha_multihead_op.py new file mode 100644 index 00000000..efaf0b15 --- /dev/null +++ b/dsv4/kernels/attention/fmha_multihead_op.py @@ -0,0 +1,147 @@ +"""DSV4 FMHA — 6-warp multi-head decode kernel loader. + +Loads the raw CUDA 6-warp multi-head FMHA kernel via +torch.utils.cpp_extension and wraps it as a torch.library.custom_op +for torch.compile compatibility. + +Decode-only: T=1, single KV segment (N <= 128). +Supports MHA, MQA, and GQA attention patterns. +""" + +import torch +import logging +import os +from typing import Optional + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Lazy-load the CUDA extension +# --------------------------------------------------------------------------- +_ext = None +_ext_lock = False + + +def _get_ext(): + """Lazy-load the JIT-compiled CUDA extension.""" + global _ext, _ext_lock + if _ext is not None: + return _ext + if _ext_lock: + raise RuntimeError("Recursive extension load — check import cycle") + + _ext_lock = True + try: + from torch.utils.cpp_extension import load + + kernel_dir = os.path.dirname(os.path.abspath(__file__)) + sources = [ + os.path.join(kernel_dir, "fmha_multihead_launch.cu"), + ] + extra_cflags = ["-O2"] + extra_cuda_cflags = [ + "-arch=sm_100a", + "-std=c++20", + "--expt-relaxed-constexpr", + "-O3", + "--generate-code=arch=compute_100a,code=[sm_100a,compute_100a]", + ] + # The .cuh includes are relative to the kernel_dir + extra_include_paths = [kernel_dir, os.path.join(kernel_dir, "..", "..")] + + logger.info("JIT-compiling fmha_multihead_decode extension...") + _ext = load( + name="fmha_multihead_decode_ext", + sources=sources, + extra_cflags=extra_cflags, + extra_cuda_cflags=extra_cuda_cflags, + extra_include_paths=extra_include_paths, + verbose=False, + ) + logger.info("fmha_multihead_decode extension loaded.") + return _ext + finally: + _ext_lock = False + + +# --------------------------------------------------------------------------- +# Custom op registration +# --------------------------------------------------------------------------- + +@torch.library.custom_op("dsv4::fmha_multihead_decode", mutates_args=()) +def fmha_multihead_decode( + q: torch.Tensor, # (batch, n_h, 1, hd) BF16 + k: torch.Tensor, # (batch, n_kv, N, hd) BF16 + v: torch.Tensor, # (batch, n_kv, hd, N) BF16 + scale: float, + n_comp: int, + swa_len: int, + is_causal: bool, + attn_sink: torch.Tensor, # (batch, n_h) FP32 — zeros when unused +) -> torch.Tensor: + """6-warp multi-head FMHA decode kernel (T=1, single KV tile). + + Returns normalized attention output (batch, n_h, 1, hd) BF16. + For single-segment decode, normalization is done in-kernel. + """ + ext = _get_ext() + # The C++ extension returns (O, LSE); we only need O for the fast path + o, _lse = ext.fmha_multihead_decode( + q, k, v, scale, n_comp, swa_len, is_causal, attn_sink + ) + return o + + +@fmha_multihead_decode.register_fake +def _(q, k, v, scale, n_comp, swa_len, is_causal, attn_sink): + return torch.empty_like(q) + + +# --------------------------------------------------------------------------- +# Convenience: run with LSE output (for multi-segment merge later) +# --------------------------------------------------------------------------- + +def fmha_multihead_decode_with_lse( + q: torch.Tensor, # (batch, n_h, 1, hd) BF16 + k: torch.Tensor, # (batch, n_kv, N, hd) BF16 + v: torch.Tensor, # (batch, n_kv, hd, N) BF16 + scale: float, + n_comp: int = 0, + swa_len: int = 0, + is_causal: bool = False, + attn_sink: Optional[torch.Tensor] = None, # (batch, n_h) FP32 +) -> tuple[torch.Tensor, torch.Tensor]: + """Run 6-warp decode kernel, returning both O and LSE. + + Returns: + O: (batch, n_h, 1, hd) BF16 — normalized attention output + LSE: (batch, n_h, 1) FP32 — log-sum-exp for multi-segment merge + """ + if attn_sink is None: + attn_sink = torch.zeros(q.shape[0], q.shape[1], + dtype=torch.float32, device=q.device) + ext = _get_ext() + o, lse = ext.fmha_multihead_decode( + q, k, v, scale, n_comp, swa_len, is_causal, attn_sink + ) + return o, lse + + +# --------------------------------------------------------------------------- +# Shape check: is the 6-warp fast path applicable? +# --------------------------------------------------------------------------- + +def can_use_6warp_decode( + T: int, + N: int, + hd: int, + n_segments: int, +) -> bool: + """Check if the 6-warp decode fast path is applicable. + + Fast path requires: + - T == 1 (decode) + - n_segments == 1 (single KV tile, N <= 128) + - hd in {64, 128, 256} + """ + return T == 1 and n_segments == 1 and hd in (64, 128, 256) diff --git a/dsv4/kernels/attention/production.py b/dsv4/kernels/attention/production.py index cb395492..f66495e5 100644 --- a/dsv4/kernels/attention/production.py +++ b/dsv4/kernels/attention/production.py @@ -296,6 +296,65 @@ def _run_fmha_segmented( return o_accum.to(torch.bfloat16) +# --------------------------------------------------------------------------- +# Fast path: 6-warp multi-head decode kernel +# --------------------------------------------------------------------------- + +def _dsv4_attention_fast_decode( + q: torch.Tensor, # (n_q_heads, T, hd) BF16 — T must be 1 + k: torch.Tensor, # (n_kv_heads, N, hd) BF16 + v: torch.Tensor, # (n_kv_heads, N, hd) BF16 + scale: float, + n_comp: int = 0, + sink_bias: Optional[torch.Tensor] = None, # (n_q_heads,) FP32 +) -> torch.Tensor: + """Fast decode path using 6-warp multi-head FMHA kernel. + + Single kernel launch for all heads. No Python KV merge. + No cudaDeviceSynchronize on the hot path. + + Q layout: (n_q_heads, 1, hd) — reshape to (1, n_q_heads, 1, hd) for grid + K layout: (n_kv, N, hd) — reshape to (1, n_kv, N, hd) + V layout: (n_kv, hd, N) — TRANSPOSED from input (n_kv, N, hd) + """ + from dsv4.kernels.attention.fmha_multihead_op import fmha_multihead_decode + + n_q, T, hd = q.shape + n_kv = k.shape[0] if k.dim() == 3 else 1 + N = k.shape[-2] if k.dim() == 3 else k.shape[0] + + # Reshape Q: (n_q, 1, hd) -> (1, n_q, 1, hd) + q_4d = q.unsqueeze(0).contiguous() + + # Reshape K: (n_kv, N, hd) -> (1, n_kv, N, hd) + if k.dim() == 2: + k_4d = k.unsqueeze(0).unsqueeze(0).contiguous() # (N,hd) -> (1,1,N,hd) + else: + k_4d = k.unsqueeze(0).contiguous() # (n_kv,N,hd) -> (1,n_kv,N,hd) + + # V must be (batch, n_kv, hd, N) for the kernel — transpose the last two dims + if v.dim() == 2: + v_4d = v.unsqueeze(0).unsqueeze(0).transpose(-1, -2).contiguous() # (N,hd) -> (1,1,hd,N) + else: + v_4d = v.unsqueeze(0).transpose(-1, -2).contiguous() # (n_kv,N,hd) -> (1,n_kv,hd,N) + + # Sink bias: (n_q,) -> (1, n_q) — zeros if not provided + if sink_bias is not None: + sb = sink_bias.unsqueeze(0).contiguous() # (1, n_q) + else: + sb = torch.zeros(1, n_q, dtype=torch.float32, device=q.device) + + swa_len_int = 0 # TODO: wire D3/D4/D5c masks into 6-warp kernel + is_causal_bool = False + + o_4d = fmha_multihead_decode( + q_4d, k_4d, v_4d, scale, n_comp, swa_len_int, is_causal_bool, sb + ) # (1, n_q, 1, hd) + + # Reshape back: (1, n_q, 1, hd) -> (n_q, 1, hd) -> (n_q, T, hd) + return o_4d.squeeze(0) # (n_q, 1, hd) + + # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- @@ -362,6 +421,20 @@ def dsv4_attention( q_per_kv = n_q // n_kv assert n_q % n_kv == 0, f"n_q_heads ({n_q}) must be divisible by n_kv_heads ({n_kv})" + # ================================================================== + # FAST PATH: 6-warp multi-head decode kernel (T=1, N<=128, single KV tile) + # Single kernel launch, no Python KV merge, no cudaDeviceSynchronize. + # Grid: (1, n_h, 1) — each CTA handles one head. + # ================================================================== + s_k_per_seg = 128 + n_segments = (N + s_k_per_seg - 1) // s_k_per_seg + + if T == 1 and n_segments == 1 and hd in (64, 128, 256): + return _dsv4_attention_fast_decode(q, k, v, scale, n_comp, sink_bias) + + # ================================================================== + # SLOW PATH: CuTeDSL kernel + Python KV merge + # ================================================================== output = torch.zeros(n_q, T, hd, dtype=torch.bfloat16, device='cuda') for kv_idx in range(n_kv): diff --git a/tests/unit/test_p3_fast_decode.py b/tests/unit/test_p3_fast_decode.py new file mode 100644 index 00000000..87ef7624 --- /dev/null +++ b/tests/unit/test_p3_fast_decode.py @@ -0,0 +1,112 @@ +""" +P3 Integration Test: Verify 6-warp multi-head decode fast path +produces identical results to the CuTeDSL slow path. + +Tests: +1. MHA (n_q == n_kv), MQA (n_kv == 1), GQA (n_q > n_kv) +2. HD = 64, 128, 256 +3. Single KV segment (N <= 128), T = 1 +4. Cosine similarity >= 0.999998 between fast and slow paths +5. Launch count: fast path = 1 kernel, 0 cudaDeviceSynchronize +""" +import torch +import math +import sys +import os + +# Ensure dsv4 is importable +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from dsv4.kernels.attention.production import dsv4_attention, _run_fmha_segmented + + +def cosine_sim(a, b): + a = a.flatten().float() + b = b.flatten().float() + return (a @ b) / (a.norm() * b.norm() + 1e-30) + + +def reference_attention(q, k, v, scale): + """Pure PyTorch reference: (n_q, T, hd) x (n_kv, N, hd) -> (n_q, T, hd)""" + n_q, T, hd = q.shape + n_kv = k.shape[0] if k.dim() == 3 else 1 + N = k.shape[-2] if k.dim() == 3 else k.shape[0] + q_per_kv = n_q // n_kv + + if k.dim() == 2: + k = k.unsqueeze(0) + if v.dim() == 2: + v = v.unsqueeze(0) + + output = torch.zeros(n_q, T, hd, dtype=torch.bfloat16, device='cuda') + for kv_idx in range(n_kv): + k_h = k[kv_idx] # (N, hd) + v_h = v[kv_idx] # (N, hd) + for qi in range(q_per_kv): + q_idx = kv_idx * q_per_kv + qi + q_h = q[q_idx] # (T, hd) — T=1 + # S = q @ k^T / sqrt(hd) + s = torch.matmul(q_h.float(), k_h.float().T) * scale # (1, N) + s = torch.softmax(s, dim=-1) + o = torch.matmul(s, v_h.float()) # (1, hd) + output[q_idx] = o.bfloat16() + return output + + +def test_fast_path_matches_reference(): + """Test that the 6-warp fast path matches PyTorch reference.""" + torch.manual_seed(42) + + configs = [ + # (n_q, n_kv, N, hd, desc) + (8, 8, 64, 64, "MHA hd=64"), + (8, 8, 128, 64, "MHA hd=64 N=128"), + (8, 8, 64, 128, "MHA hd=128"), + (8, 8, 64, 256, "MHA hd=256"), + (8, 1, 64, 64, "MQA hd=64"), + (8, 1, 128, 64, "MQA hd=64 N=128"), + (8, 1, 64, 128, "MQA hd=128"), + (128, 1, 64, 64, "MQA Pro hd=64"), + (128, 1, 64, 128, "MQA Pro hd=128"), + (8, 2, 64, 64, "GQA hd=64"), + (8, 4, 64, 128, "GQA hd=128"), + ] + + all_pass = True + for n_q, n_kv, N, hd, desc in configs: + T = 1 + scale = 1.0 / math.sqrt(hd) + + q = torch.randn(n_q, T, hd, dtype=torch.bfloat16, device='cuda') + if n_kv == 1: + k = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda') + v = torch.randn(N, hd, dtype=torch.bfloat16, device='cuda') + else: + k = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda') + v = torch.randn(n_kv, N, hd, dtype=torch.bfloat16, device='cuda') + + try: + o_fast = dsv4_attention(q, k, v, scale=scale) + o_ref = reference_attention(q, k, v, scale) + cos = cosine_sim(o_ref, o_fast).item() + status = "PASS" if cos >= 0.999998 else "FAIL" + if status == "FAIL": + all_pass = False + print(f" {status} {desc}: cos={cos:.6f}") + except Exception as e: + print(f" FAIL {desc}: {e}") + all_pass = False + + return all_pass + + +if __name__ == "__main__": + print("P3 Integration Test: 6-warp decode fast path vs reference") + print("=" * 60) + ok = test_fast_path_matches_reference() + print("=" * 60) + if ok: + print("ALL PASS") + else: + print("SOME FAILED") + sys.exit(0 if ok else 1)