P3: wire 6-warp multi-head FMHA decode fast path into production.py
- fmha_multihead_launch.cu: PyTorch launch wrapper for fmha_6warp_multihead_kernel
(c10::BFloat16 boundary, uint16_t bf16_t inside kernel, zero-cost casts)
- fmha_multihead_op.py: torch.utils.cpp_extension JIT loader + custom_op registration
(dsv4::fmha_multihead_decode for torch.compile)
- production.py: fast path dispatch for T=1, n_segments==1, hd in {64,128,256}
Falls through to CuTeDSL slow path for multi-segment/prefill
- test_p3_fast_decode.py: integration test (MHA/MQA/GQA, cosine >= 0.999998)
Architecture:
Grid: dim3(1, n_h, batch_size) — one CTA per (head, batch)
MQA: k_head_stride=0 so all Q heads share same K/V
Single kernel launch, zero cudaDeviceSynchronize on hot path
Normalized output for single-segment decode
This commit is contained in:
196
NEXT_PRIORITIES.md
Normal file
196
NEXT_PRIORITIES.md
Normal file
@@ -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."
|
||||
167
dsv4/kernels/attention/fmha_multihead_launch.cu
Normal file
167
dsv4/kernels/attention/fmha_multihead_launch.cu
Normal file
@@ -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 <ATen/ATen.h>
|
||||
#include <torch/extension.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// 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<torch::Tensor, torch::Tensor> 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<torch::Tensor> 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<const bf16_t*>(q.data_ptr<at::BFloat16>());
|
||||
params.k = reinterpret_cast<const bf16_t*>(k.data_ptr<at::BFloat16>());
|
||||
params.v = reinterpret_cast<const bf16_t*>(v.data_ptr<at::BFloat16>());
|
||||
params.o = reinterpret_cast<bf16_t*>(o.data_ptr<at::BFloat16>());
|
||||
params.lse = lse.data_ptr<float>();
|
||||
params.s_k = N;
|
||||
params.scale = static_cast<float>(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<D, 128><<<grid, block, smem>>>(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)");
|
||||
}
|
||||
147
dsv4/kernels/attention/fmha_multihead_op.py
Normal file
147
dsv4/kernels/attention/fmha_multihead_op.py
Normal file
@@ -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)
|
||||
@@ -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):
|
||||
|
||||
112
tests/unit/test_p3_fast_decode.py
Normal file
112
tests/unit/test_p3_fast_decode.py
Normal file
@@ -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)
|
||||
Reference in New Issue
Block a user