diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index 4a3b3a4..5454f85 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -1,19 +1,19 @@ """CUDA attention kernel wrappers with torch fallback. Public API: - - ``gqa_decode_attn`` — single-query decode attention - - ``gqa_prefill_attn`` — multi-query prefill attention + - ``attn_decode`` — single-query decode attention + - ``attn_prefill`` — multi-query prefill attention -Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.gqa_*``) +Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``) when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``. """ from astrai.extension.loader import KERNEL_NAMES, is_available -from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn +from astrai.extension.ops import attn_decode, attn_prefill __all__ = [ - "gqa_decode_attn", - "gqa_prefill_attn", + "attn_decode", + "attn_prefill", "is_available", "KERNEL_NAMES", ] diff --git a/astrai/extension/loader.py b/astrai/extension/loader.py index 21afbeb..5633e9b 100644 --- a/astrai/extension/loader.py +++ b/astrai/extension/loader.py @@ -11,7 +11,7 @@ import logging logger = logging.getLogger(__name__) -KERNEL_NAMES = ["gqa_decode_attn", "gqa_prefill_attn"] +KERNEL_NAMES = ["attn_decode", "attn_prefill"] _available: dict[str, bool] = {} _modules: dict[str, object] = {} diff --git a/astrai/extension/ops.py b/astrai/extension/ops.py index 1be3464..39f8faa 100644 --- a/astrai/extension/ops.py +++ b/astrai/extension/ops.py @@ -42,7 +42,7 @@ def _torch_fallback( ) -def gqa_decode_attn( +def attn_decode( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -51,8 +51,8 @@ def gqa_decode_attn( causal_offset: int = 0, scale: float | None = None, ) -> torch.Tensor: - if _available["gqa_decode_attn"]: - return _modules["gqa_decode_attn"].gqa_decode_attn( + if _available["attn_decode"]: + return _modules["attn_decode"].attn_decode( q, k, v, @@ -64,7 +64,7 @@ def gqa_decode_attn( return _torch_fallback(q, k, v, mask, is_causal, scale) -def gqa_prefill_attn( +def attn_prefill( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, @@ -73,8 +73,8 @@ def gqa_prefill_attn( causal_offset: int = 0, scale: float | None = None, ) -> torch.Tensor: - if _available["gqa_prefill_attn"]: - return _modules["gqa_prefill_attn"].gqa_prefill_attn( + if _available["attn_prefill"]: + return _modules["attn_prefill"].attn_prefill( q, k, v, diff --git a/csrc/build.py b/csrc/build.py index a445ba2..8605830 100644 --- a/csrc/build.py +++ b/csrc/build.py @@ -42,5 +42,5 @@ def register(name: str, sources: list[str] | None = None, **kwargs): } -register("gqa_decode_attn") -register("gqa_prefill_attn") +register("attn_decode") +register("attn_prefill") diff --git a/csrc/kernels/gqa_common.cuh b/csrc/kernels/attn_common.cuh similarity index 80% rename from csrc/kernels/gqa_common.cuh rename to csrc/kernels/attn_common.cuh index 981d601..a0bdcc8 100644 --- a/csrc/kernels/gqa_common.cuh +++ b/csrc/kernels/attn_common.cuh @@ -7,7 +7,7 @@ using bf16 = __nv_bfloat16; using std::min; -struct GQAParams { +struct AttentionParams { int batch; int q_head; int kv_head; @@ -22,5 +22,9 @@ struct GQAParams { const bf16* __restrict__ k; const bf16* __restrict__ v; const bool* __restrict__ mask; + bf16* __restrict__ o; + float* __restrict__ o_part; + float* __restrict__ ml_part; + int num_splits; }; diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/attn_decode.cu similarity index 53% rename from csrc/kernels/gqa_decode_attn.cu rename to csrc/kernels/attn_decode.cu index f88aa40..b01c29e 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/attn_decode.cu @@ -1,70 +1,67 @@ -#include "gqa_decode_attn.cuh" -#include "gqa_entry_utils.cuh" +#include "attn_decode_split_kv.cuh" +#include "attn_entry_utils.cuh" #ifndef ASTRAI_NO_MMA -#include "gqa_decode_attn_mma.cuh" +#include "attn_decode_split_kv_mma.cuh" #endif -// Scalar fallback: one warp per query head, per (batch, kv_head) block. -static void launch_scalar_decode(const GQAParams& p) { - int group_size = p.q_head / p.kv_head; - size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); - gqa_decode_attn_kernel<<>>(p); -} - -#ifndef ASTRAI_NO_MMA -// Tensor-core head-packing requires 1 < G <= 16 (the MMA M dim) and no mask. -static bool decode_use_mma(const GQAParams& p) { - int G = p.q_head / p.kv_head; - return !p.use_mask && G > 1 && G <= 16; -} - // Decode has only batch*kv_head independent tasks; without split-K the grid is // tiny (e.g. 16 blocks) and leaves most SMs idle. Pick the smallest split count -// that fills the device (~2 blocks/SM), capped by the tile count and 32. -static int decode_num_splits(const GQAParams& p, int tiles_total) { +// that fills the device (~2 blocks/SM), capped by the tile count, min work per +// split (at least 8 tiles), and 32. +static int decode_num_splits(const AttentionParams& p, int tiles_total) { int sm_count = 0; cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0); int base_blocks = p.kv_head * p.batch; int desired = 2 * (sm_count > 0 ? sm_count : 64); int n = (desired + base_blocks - 1) / base_blocks; - return std::max(1, std::min(n, std::min(tiles_total, 32))); + int max_by_work = tiles_total / 8; + return std::max(1, std::min({n, tiles_total, 32, max_by_work})); +} + +// Scalar fallback: one warp per query head, split-KV across grid.z. +static void launch_scalar_decode(AttentionParams& p) { + int group_size = p.q_head / p.kv_head; + int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK; + p.num_splits = decode_num_splits(p, chunks_total); + + auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt); + auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt); + p.o_part = o_part.data_ptr(); + p.ml_part = ml_part.data_ptr(); + + size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); + attn_decode_split_kv_kernel<<>>(p); + attn_decode_combine_kernel<<>>(p); +} + +#ifndef ASTRAI_NO_MMA +// Tensor-core head-packing requires 1 < G <= 16 (the MMA M dim) and no mask. +static bool decode_use_mma(const AttentionParams& p) { + int G = p.q_head / p.kv_head; + return !p.use_mask && G > 1 && G <= 16; } template -static void launch_mma_decode(GQAParams& p) { - constexpr int BR = 16, LD = HEAD_DIM; // XOR swizzle → no padding - int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16); +static void launch_mma_decode(AttentionParams& p) { int tiles_total = (p.kv_len + BC - 1) / BC; - int num_splits = decode_num_splits(p, tiles_total); + p.num_splits = decode_num_splits(p, tiles_total); - // Enough (batch, kv_head) work to fill the SMs → single pass, direct write. - if (num_splits <= 1) { - cudaFuncSetAttribute(gqa_decode_attn_mma_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - gqa_decode_attn_mma_kernel - <<>>(p); - return; - } - - // Split-K (FlashDecoding): partition kv across blocks, then reduce. auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto o_part = torch::empty({p.batch, p.q_head, num_splits, p.head_dim}, fopt); - auto ml_part = torch::empty({p.batch, p.q_head, num_splits, 2}, fopt); + auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt); + auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt); + p.o_part = o_part.data_ptr(); + p.ml_part = ml_part.data_ptr(); - cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - gqa_decode_attn_mma_splitk_kernel - <<>>( - p, o_part.data_ptr(), ml_part.data_ptr(), num_splits); - gqa_decode_combine_kernel<<>>( - o_part.data_ptr(), ml_part.data_ptr(), p.o, - num_splits, p.head_dim); + attn_decode_split_kv_mma_kernel + <<>>(p); + attn_decode_combine_kernel<<>>(p); } #endif template -static void dispatch_decode(GQAParams& p) { +static void dispatch_decode(AttentionParams& p) { #ifndef ASTRAI_NO_MMA if (decode_use_mma(p)) { launch_mma_decode(p); @@ -74,7 +71,7 @@ static void dispatch_decode(GQAParams& p) { launch_scalar_decode(p); } -torch::Tensor gqa_decode_attn( +torch::Tensor attn_decode( torch::Tensor q, torch::Tensor k, torch::Tensor v, @@ -83,8 +80,8 @@ torch::Tensor gqa_decode_attn( int64_t causal_offset = 0, c10::optional scale = c10::nullopt ) { - GQAParams p; - gqa_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); + AttentionParams p; + attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1"); TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32"); @@ -112,7 +109,7 @@ torch::Tensor gqa_decode_attn( } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("gqa_decode_attn", &gqa_decode_attn, + m.def("attn_decode", &attn_decode, py::arg("q"), py::arg("k"), py::arg("v"), diff --git a/csrc/kernels/gqa_decode_attn.cuh b/csrc/kernels/attn_decode_split_kv.cuh similarity index 51% rename from csrc/kernels/gqa_decode_attn.cuh rename to csrc/kernels/attn_decode_split_kv.cuh index 3a9826e..9730cb2 100644 --- a/csrc/kernels/gqa_decode_attn.cuh +++ b/csrc/kernels/attn_decode_split_kv.cuh @@ -1,5 +1,5 @@ #pragma once -#include "gqa_common.cuh" +#include "attn_common.cuh" constexpr int DC_CHUNK = 64; @@ -9,9 +9,10 @@ __device__ inline float warp_reduce_sum(float val) { return val; } -__global__ void gqa_decode_attn_kernel(GQAParams p) { +__global__ void attn_decode_split_kv_kernel(AttentionParams p) { int batch = blockIdx.x / p.kv_head; int kv_head = blockIdx.x % p.kv_head; + int split = blockIdx.z; int group_size = blockDim.y; int q_head = kv_head * group_size + threadIdx.y; int lane = threadIdx.x; @@ -29,7 +30,14 @@ __global__ void gqa_decode_attn_kernel(GQAParams p) { extern __shared__ __align__(16) bf16 k_smem[]; - for (int chunk_start = 0; chunk_start < p.kv_len; chunk_start += DC_CHUNK) { + // Split-KV: each split processes a contiguous subset of chunks + int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK; + int chunks_per_split = (chunks_total + p.num_splits - 1) / p.num_splits; + int ch_begin = split * chunks_per_split; + int ch_end = min(chunks_total, ch_begin + chunks_per_split); + + for (int ci = ch_begin; ci < ch_end; ci++) { + int chunk_start = ci * DC_CHUNK; int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start); int total = this_chunk * p.head_dim; @@ -61,7 +69,47 @@ __global__ void gqa_decode_attn_kernel(GQAParams p) { __syncthreads(); } - int out_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread; - for (int i = 0; i < hd_per_thread; i++) - p.o[out_off + i] = __float2bfloat16(acc_reg[i] / d); + // ---- write UN-normalised partials for this split ---- + size_t bh = (size_t)batch * p.q_head + q_head; + size_t slot = bh * p.num_splits + split; + int d0 = lane * hd_per_thread; + for (int i = 0; i < hd_per_thread; i++) { + int dd = d0 + i; + p.o_part[slot * p.head_dim + dd] = acc_reg[i]; + } + if (lane == 0) { + p.ml_part[slot * 2] = m; + p.ml_part[slot * 2 + 1] = d; + } +} + +// Reduce split-K partials into the final bf16 output. One block per (batch, +// q_head); each thread owns one head_dim element and folds across all splits +// with a numerically-stable online rescale. +__global__ void attn_decode_combine_kernel(AttentionParams p) { + int bh = blockIdx.x; + int d = threadIdx.x; + if (d >= p.head_dim) return; + + size_t split_base = (size_t)bh * p.num_splits; + + const float* mlp = p.ml_part + split_base * 2; + float mstar = -FLT_MAX; + for (int s = 0; s < p.num_splits; s++) + mstar = fmaxf(mstar, mlp[s * 2]); + + float lstar = 0.0f; + for (int s = 0; s < p.num_splits; s++) { + float mi = mlp[s * 2]; + if (mi > -FLT_MAX) lstar += mlp[s * 2 + 1] * __expf(mi - mstar); + } + + const float* op = p.o_part + split_base * p.head_dim; + float acc = 0.0f; + for (int s = 0; s < p.num_splits; s++) { + float mi = mlp[s * 2]; + if (mi > -FLT_MAX) acc += op[s * p.head_dim + d] * __expf(mi - mstar); + } + float inv = (lstar > 1e-20f) ? (1.0f / lstar) : 0.0f; + p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv); } diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh new file mode 100644 index 0000000..76ac025 --- /dev/null +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -0,0 +1,164 @@ +#pragma once +#include "attn_common.cuh" +#include "attn_mma_utils.cuh" + +// Split-K (FlashDecoding) tensor-core decode via GQA head-packing. +// +// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work +// on its own. But GQA gives us G = q_head / kv_head query heads that all share +// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16, +// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile +// across all G heads (K/V load is the decode bottleneck, so the reuse is the win, +// not the flops). The KV sequence is partitioned across gridDim.z blocks so that +// a decode with only batch*kv_head independent tasks can fill all SMs. Each +// (batch, kv_head, split) block computes an UN-normalised partial (Oacc, m, l) +// over its KV slice; the combine kernel below reduces across splits. Fixes the +// "grid too small" bottleneck (0.04 waves/SM → many blocks) for long-context, +// small-batch decode. +// +// Partial layout (float, contiguous): +// o_part : [batch, q_head, num_splits, HEAD_DIM] +// ml_part: [batch, q_head, num_splits, 2] (m, l) +// +// Optimizations: +// - cp.async global→shared for K/V (bypasses registers, cuts instruction count) +// - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts +// - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply +// - single-buffer: keeps smem small for high occupancy +template +__global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { + constexpr int BR = 16; + constexpr int KD = HEAD_DIM / 16; + constexpr int NC8 = BC / 8; + constexpr int KT2 = BC / 16; + constexpr int DN8 = HEAD_DIM / 8; + constexpr int LD = HEAD_DIM; + constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1); + + const int lane = threadIdx.x; + const int gid = lane >> 2; + const int tid4 = lane & 3; + + const int kv_head = blockIdx.x; + const int batch = blockIdx.y; + const int split = blockIdx.z; + const int G = p.q_head / p.kv_head; + const int q_head0 = kv_head * G; + + __shared__ __align__(16) bf16 sK[BC * HEAD_DIM]; + __shared__ __align__(16) bf16 sV[BC * HEAD_DIM]; + __shared__ __align__(16) bf16 sQ[BR * HEAD_DIM]; + + bf16 scale_bf16 = __float2bfloat16(p.scale); + for (int i = lane; i < BR * HEAD_DIM; i += 32) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + bf16 val = __float2bfloat16(0.0f); + if (r < G) { + int qh = q_head0 + r; + val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d]; + } + sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16); + } + __syncwarp(); + + unsigned Qa[KD][4]; + int qrow_l = (lane & 7) + (lane & 8); + int qcol_l = (lane & 16) ? 8 : 0; +#pragma unroll + for (int kt = 0; kt < KD; kt++) + ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]); + + float Oacc[DN8][4]; +#pragma unroll + for (int j = 0; j < DN8; j++) + Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; + float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; + + const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM; + const int mask_base = batch * p.kv_len; + const int tiles_total = (p.kv_len + BC - 1) / BC; + const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits; + const int ti_begin = split * tiles_per_split; + const int ti_end = min(tiles_total, ti_begin + tiles_per_split); + const int has_mask = p.use_mask && p.mask; + + for (int ti = ti_begin; ti < ti_end; ti++) { + int kv0 = ti * BC; + + bool full_tile = (kv0 + BC <= p.kv_len); + if (full_tile) { + constexpr int VEC = 8; + int total = BC * HEAD_DIM; +#pragma unroll + for (int i = lane * VEC; i < total; i += 32 * VEC) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + int kc = kv0 + r; + cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)], + &p.k[kv_base + kc * HEAD_DIM + d]); + cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)], + &p.v[kv_base + kc * HEAD_DIM + d]); + } + cp_async_commit(); + cp_async_wait_all(); + } else { + for (int i = lane; i < BC * HEAD_DIM; i += 32) { + int r = i / HEAD_DIM, d = i % HEAD_DIM; + int kc = kv0 + r; + bf16 z = __float2bfloat16(0.0f); + sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = + (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z; + sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = + (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z; + } + } + __syncwarp(); + + float Sacc[NC8][4]; + mma_compute_scores(Qa, sK, LD, SWIZ_MASK, lane, Sacc); + + int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; + mma_softmax_tile(kv0, maxc, maxc, + mask_base, p.mask, has_mask, + Sacc, Oacc, m0, m1, l0, l1, lane); + + mma_pv_accumulate(Sacc, sV, LD, SWIZ_MASK, lane, Oacc); + __syncwarp(); + } + + // ---- write UN-normalised partials for this split ---- + auto split_slot = [&](int h) -> size_t { + size_t bh = (size_t)batch * p.q_head + h; + return bh * p.num_splits + split; + }; +#pragma unroll + for (int dn8 = 0; dn8 < DN8; dn8++) { + int d = dn8 * 8 + 2 * tid4; + int r0 = gid, r1 = gid + 8; + if (r0 < G) { + int h = q_head0 + r0; + float* op = p.o_part + split_slot(h) * HEAD_DIM; + op[d] = Oacc[dn8][0]; + op[d + 1] = Oacc[dn8][1]; + } + if (r1 < G) { + int h = q_head0 + r1; + float* op = p.o_part + split_slot(h) * HEAD_DIM; + op[d] = Oacc[dn8][2]; + op[d + 1] = Oacc[dn8][3]; + } + } + if (tid4 == 0) { + int r0 = gid, r1 = gid + 8; + if (r0 < G) { + int h = q_head0 + r0; + float* mp = p.ml_part + split_slot(h) * 2; + mp[0] = m0; mp[1] = l0; + } + if (r1 < G) { + int h = q_head0 + r1; + float* mp = p.ml_part + split_slot(h) * 2; + mp[0] = m1; mp[1] = l1; + } + } +} + diff --git a/csrc/kernels/gqa_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh similarity index 94% rename from csrc/kernels/gqa_entry_utils.cuh rename to csrc/kernels/attn_entry_utils.cuh index 19b0605..7fe3e46 100644 --- a/csrc/kernels/gqa_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -1,8 +1,8 @@ #pragma once #include -#include "gqa_common.cuh" +#include "attn_common.cuh" -inline void gqa_pack_params( +inline void attn_pack_params( torch::Tensor q, torch::Tensor k, torch::Tensor v, @@ -10,7 +10,7 @@ inline void gqa_pack_params( bool is_causal, int64_t causal_offset, c10::optional scale, - GQAParams& p + AttentionParams& p ) { TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda()); TORCH_CHECK(q.dtype() == torch::kBFloat16); diff --git a/csrc/kernels/gqa_mma_utils.cuh b/csrc/kernels/attn_mma_utils.cuh similarity index 100% rename from csrc/kernels/gqa_mma_utils.cuh rename to csrc/kernels/attn_mma_utils.cuh diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/attn_prefill.cu similarity index 83% rename from csrc/kernels/gqa_prefill_attn.cu rename to csrc/kernels/attn_prefill.cu index 8d2f335..1ea31f2 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/attn_prefill.cu @@ -1,12 +1,12 @@ -#include "gqa_prefill_attn.cuh" -#include "gqa_entry_utils.cuh" +#include "attn_prefill_split_q.cuh" +#include "attn_entry_utils.cuh" #ifndef ASTRAI_NO_MMA -#include "gqa_prefill_attn_mma.cuh" +#include "attn_prefill_split_q_mma.cuh" #endif template -static void dispatch_prefill(GQAParams& p) { +static void dispatch_prefill(AttentionParams& p) { #ifndef ASTRAI_NO_MMA constexpr int WARPS = 4, BR = 16; // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier + @@ -23,16 +23,16 @@ static void dispatch_prefill(GQAParams& p) { dim3 block(WARPS * 32, 1, 1); // Static shared memory — no dynamic smem or cudaFuncSetAttribute needed. // sK[BC*LD] + sV[BC*LD] + sQ[BR*LD], all sized by template params. - gqa_prefill_attn_mma_kernel<<>>(p); + attn_prefill_split_q_mma_kernel<<>>(p); #else constexpr int G = 8, ROWS = 32, P_BC = 32; dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch); dim3 block(G, ROWS, 1); - gqa_prefill_attn_kernel_t<<>>(p); + attn_prefill_split_q_kernel_t<<>>(p); #endif } -torch::Tensor gqa_prefill_attn( +torch::Tensor attn_prefill( torch::Tensor q, torch::Tensor k, torch::Tensor v, @@ -41,8 +41,8 @@ torch::Tensor gqa_prefill_attn( int64_t causal_offset = 0, c10::optional scale = c10::nullopt ) { - GQAParams p; - gqa_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); + AttentionParams p; + attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16"); auto O = torch::empty_like(q); @@ -69,7 +69,7 @@ torch::Tensor gqa_prefill_attn( } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("gqa_prefill_attn", &gqa_prefill_attn, + m.def("attn_prefill", &attn_prefill, py::arg("q"), py::arg("k"), py::arg("v"), diff --git a/csrc/kernels/gqa_prefill_attn.cuh b/csrc/kernels/attn_prefill_split_q.cuh similarity index 98% rename from csrc/kernels/gqa_prefill_attn.cuh rename to csrc/kernels/attn_prefill_split_q.cuh index c360e9b..5df9b43 100644 --- a/csrc/kernels/gqa_prefill_attn.cuh +++ b/csrc/kernels/attn_prefill_split_q.cuh @@ -1,5 +1,5 @@ #pragma once -#include "gqa_common.cuh" +#include "attn_common.cuh" // v9: group-split register blocking. G threads cooperate on one query row, // each owning HEAD_DIM/G dims of qreg[]/acc[]. Small per-thread footprint keeps @@ -31,7 +31,7 @@ __device__ __forceinline__ void ld8(const bf16* p, float* o) { } template -__global__ void gqa_prefill_attn_kernel_t(GQAParams p) { +__global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { constexpr int DPT = HEAD_DIM / G; int q_tile = blockIdx.x; diff --git a/csrc/kernels/gqa_prefill_attn_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh similarity index 98% rename from csrc/kernels/gqa_prefill_attn_mma.cuh rename to csrc/kernels/attn_prefill_split_q_mma.cuh index b9bbbd6..9badf81 100644 --- a/csrc/kernels/gqa_prefill_attn_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -1,6 +1,6 @@ #pragma once -#include "gqa_common.cuh" -#include "gqa_mma_utils.cuh" +#include "attn_common.cuh" +#include "attn_mma_utils.cuh" // Tensor-core prefill flash attention (raw mma.sync PTX). // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor @@ -38,7 +38,7 @@ template __global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS) -void gqa_prefill_attn_mma_kernel(GQAParams p) { +void attn_prefill_split_q_mma_kernel(AttentionParams p) { constexpr int BR = 16; constexpr int KD = HEAD_DIM / 16; // Q/K k-tiles constexpr int NC8 = BC / 8; // S n-tiles (N=8 each) diff --git a/csrc/kernels/gqa_decode_attn_mma.cuh b/csrc/kernels/gqa_decode_attn_mma.cuh deleted file mode 100644 index 84daf85..0000000 --- a/csrc/kernels/gqa_decode_attn_mma.cuh +++ /dev/null @@ -1,321 +0,0 @@ -#pragma once -#include "gqa_common.cuh" -#include "gqa_mma_utils.cuh" - -// Tensor-core decode via GQA head-packing with cp.async loads. -// -// Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work -// on its own. But GQA gives us G = q_head / kv_head query heads that all share -// one kv_head. We pack those G heads into the M=16 rows of mma.sync.m16n8k16, -// turning G independent GEMVs into a single GEMM that reuses each loaded K/V tile -// across all G heads (K/V load is the decode bottleneck, so the reuse is the win, -// not the flops). Fragment layout is identical to the prefill mma kernel; the -// only differences are (1) the M rows come from different heads at position 0 -// instead of different sequence positions of one head, and (2) causal masking is -// a single scalar bound shared by every row. One warp owns one (batch, kv_head); -// requires G <= 16. -// -// Optimizations: -// - cp.async global→shared for K/V (bypasses registers, cuts instruction count) -// - XOR swizzle (swiz_col): LD=HEAD_DIM, zero waste, no bank conflicts -// - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply -// - single-buffer: keeps smem small for high occupancy - -template -__global__ void gqa_decode_attn_mma_kernel(GQAParams p) { - constexpr int BR = 16; - constexpr int KD = HEAD_DIM / 16; // Q/K k-tiles - constexpr int NC8 = BC / 8; // S n-tiles (N=8 each) - constexpr int KT2 = BC / 16; // P k-tiles (K=16 each) - constexpr int DN8 = HEAD_DIM / 8; // O n-tiles (N=8 each) - constexpr int LD = HEAD_DIM; // XOR swizzle handles bank conflicts, zero waste - constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1); - - const int lane = threadIdx.x; // single warp - const int gid = lane >> 2; // 0..7 → rows gid, gid+8 - const int tid4 = lane & 3; - - const int kv_head = blockIdx.x; - const int batch = blockIdx.y; - const int G = p.q_head / p.kv_head; - const int q_head0 = kv_head * G; - - extern __shared__ __align__(16) bf16 smem[]; - bf16* sK = smem; // [BC][LD] - bf16* sV = sK + BC * LD; // [BC][LD] - bf16* sQ = sV + BC * LD; // [BR][LD] - - // ---- stage Q into shared (pre-scaled, swizzled) ---- - bf16 scale_bf16 = __float2bfloat16(p.scale); - for (int i = lane; i < BR * HEAD_DIM; i += 32) { - int r = i / HEAD_DIM, d = i % HEAD_DIM; - bf16 val = __float2bfloat16(0.0f); - if (r < G) { - int qh = q_head0 + r; - val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d]; // q_len == 1 - } - sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16); - } - __syncwarp(); - - // Q resident A-fragments - unsigned Qa[KD][4]; - int qrow_l = (lane & 7) + (lane & 8); - int qcol_l = (lane & 16) ? 8 : 0; -#pragma unroll - for (int kt = 0; kt < KD; kt++) - ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]); - - float Oacc[DN8][4]; -#pragma unroll - for (int j = 0; j < DN8; j++) - Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; - float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; - - const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM; - const int mask_base = batch * p.kv_len; - const int tiles = (p.kv_len + BC - 1) / BC; - const int has_mask = p.use_mask && p.mask; - - for (int ti = 0; ti < tiles; ti++) { - int kv0 = ti * BC; - - // ---- load K/V tile to shared (cp.async on full tiles) ---- - bool full_tile = (kv0 + BC <= p.kv_len); - if (full_tile) { - constexpr int VEC = 8; // 8 bf16 = 16 bytes per cp.async - int total = BC * HEAD_DIM; -#pragma unroll - for (int i = lane * VEC; i < total; i += 32 * VEC) { - int r = i / HEAD_DIM, d = i % HEAD_DIM; - int kc = kv0 + r; - cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)], - &p.k[kv_base + kc * HEAD_DIM + d]); - cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)], - &p.v[kv_base + kc * HEAD_DIM + d]); - } - cp_async_commit(); - cp_async_wait_all(); - } else { - for (int i = lane; i < BC * HEAD_DIM; i += 32) { - int r = i / HEAD_DIM, d = i % HEAD_DIM; - int kc = kv0 + r; - bf16 z = __float2bfloat16(0.0f); - sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = - (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z; - sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = - (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z; - } - } - __syncwarp(); - - // S = Q @ K^T + online softmax + O += P @ V (shared MMA functions) - float Sacc[NC8][4]; - mma_compute_scores(Qa, sK, LD, SWIZ_MASK, lane, Sacc); - - int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; - mma_softmax_tile(kv0, maxc, maxc, - mask_base, p.mask, has_mask, - Sacc, Oacc, m0, m1, l0, l1, lane); - - mma_pv_accumulate(Sacc, sV, LD, SWIZ_MASK, lane, Oacc); - __syncwarp(); // sK/sV reused next tile - } - - // ---- write output ---- - float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f; - float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f; -#pragma unroll - for (int dn8 = 0; dn8 < DN8; dn8++) { - int d = dn8 * 8 + 2 * tid4; - int r0 = gid, r1 = gid + 8; - if (r0 < G) { - int o_off = (batch * p.q_head + q_head0 + r0) * HEAD_DIM + d; - p.o[o_off] = __float2bfloat16(Oacc[dn8][0] * rl0); - p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][1] * rl0); - } - if (r1 < G) { - int o_off = (batch * p.q_head + q_head0 + r1) * HEAD_DIM + d; - p.o[o_off] = __float2bfloat16(Oacc[dn8][2] * rl1); - p.o[o_off + 1] = __float2bfloat16(Oacc[dn8][3] * rl1); - } - } -} - -// --------------------------------------------------------------------------- -// Split-K (FlashDecoding) decode: identical math to the kernel above, but the -// KV sequence is partitioned across gridDim.z blocks so that a decode with only -// batch*kv_head independent tasks can fill all SMs. Each (batch, kv_head, split) -// block computes an UN-normalised partial (Oacc, m, l) over its KV slice; the -// combine kernel below reduces across splits. Fixes the "grid too small" -// bottleneck (0.04 waves/SM → many blocks) for long-context, small-batch decode. -// -// Partial layout (float, contiguous): -// o_part : [batch, q_head, num_splits, HEAD_DIM] -// ml_part: [batch, q_head, num_splits, 2] (m, l) -template -__global__ void gqa_decode_attn_mma_splitk_kernel(GQAParams p, - float* __restrict__ o_part, - float* __restrict__ ml_part, - int num_splits) { - constexpr int BR = 16; - constexpr int KD = HEAD_DIM / 16; - constexpr int NC8 = BC / 8; - constexpr int KT2 = BC / 16; - constexpr int DN8 = HEAD_DIM / 8; - constexpr int LD = HEAD_DIM; - constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1); - - const int lane = threadIdx.x; - const int gid = lane >> 2; - const int tid4 = lane & 3; - - const int kv_head = blockIdx.x; - const int batch = blockIdx.y; - const int split = blockIdx.z; - const int G = p.q_head / p.kv_head; - const int q_head0 = kv_head * G; - - extern __shared__ __align__(16) bf16 smem[]; - bf16* sK = smem; - bf16* sV = sK + BC * LD; - bf16* sQ = sV + BC * LD; - - bf16 scale_bf16 = __float2bfloat16(p.scale); - for (int i = lane; i < BR * HEAD_DIM; i += 32) { - int r = i / HEAD_DIM, d = i % HEAD_DIM; - bf16 val = __float2bfloat16(0.0f); - if (r < G) { - int qh = q_head0 + r; - val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d]; - } - sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16); - } - __syncwarp(); - - unsigned Qa[KD][4]; - int qrow_l = (lane & 7) + (lane & 8); - int qcol_l = (lane & 16) ? 8 : 0; -#pragma unroll - for (int kt = 0; kt < KD; kt++) - ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]); - - float Oacc[DN8][4]; -#pragma unroll - for (int j = 0; j < DN8; j++) - Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; - float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; - - const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM; - const int mask_base = batch * p.kv_len; - const int tiles_total = (p.kv_len + BC - 1) / BC; - const int tiles_per_split = (tiles_total + num_splits - 1) / num_splits; - const int ti_begin = split * tiles_per_split; - const int ti_end = min(tiles_total, ti_begin + tiles_per_split); - const int has_mask = p.use_mask && p.mask; - - for (int ti = ti_begin; ti < ti_end; ti++) { - int kv0 = ti * BC; - - bool full_tile = (kv0 + BC <= p.kv_len); - if (full_tile) { - constexpr int VEC = 8; - int total = BC * HEAD_DIM; -#pragma unroll - for (int i = lane * VEC; i < total; i += 32 * VEC) { - int r = i / HEAD_DIM, d = i % HEAD_DIM; - int kc = kv0 + r; - cp_async_16(&sK[r * LD + swiz_col(d, r, SWIZ_MASK)], - &p.k[kv_base + kc * HEAD_DIM + d]); - cp_async_16(&sV[r * LD + swiz_col(d, r, SWIZ_MASK)], - &p.v[kv_base + kc * HEAD_DIM + d]); - } - cp_async_commit(); - cp_async_wait_all(); - } else { - for (int i = lane; i < BC * HEAD_DIM; i += 32) { - int r = i / HEAD_DIM, d = i % HEAD_DIM; - int kc = kv0 + r; - bf16 z = __float2bfloat16(0.0f); - sK[r * LD + swiz_col(d, r, SWIZ_MASK)] = - (kc < p.kv_len) ? p.k[kv_base + kc * HEAD_DIM + d] : z; - sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = - (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z; - } - } - __syncwarp(); - - float Sacc[NC8][4]; - mma_compute_scores(Qa, sK, LD, SWIZ_MASK, lane, Sacc); - - int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; - mma_softmax_tile(kv0, maxc, maxc, - mask_base, p.mask, has_mask, - Sacc, Oacc, m0, m1, l0, l1, lane); - - mma_pv_accumulate(Sacc, sV, LD, SWIZ_MASK, lane, Oacc); - __syncwarp(); - } - - // ---- write UN-normalised partials for this split ---- -#pragma unroll - for (int dn8 = 0; dn8 < DN8; dn8++) { - int d = dn8 * 8 + 2 * tid4; - int r0 = gid, r1 = gid + 8; - if (r0 < G) { - int hh = q_head0 + r0; - float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM; - op[d] = Oacc[dn8][0]; - op[d + 1] = Oacc[dn8][1]; - } - if (r1 < G) { - int hh = q_head0 + r1; - float* op = o_part + ((size_t)(batch * p.q_head + hh) * num_splits + split) * HEAD_DIM; - op[d] = Oacc[dn8][2]; - op[d + 1] = Oacc[dn8][3]; - } - } - if (tid4 == 0) { - int r0 = gid, r1 = gid + 8; - if (r0 < G) { - float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r0) * num_splits + split) * 2; - mp[0] = m0; mp[1] = l0; - } - if (r1 < G) { - float* mp = ml_part + ((size_t)(batch * p.q_head + q_head0 + r1) * num_splits + split) * 2; - mp[0] = m1; mp[1] = l1; - } - } -} - -// Reduce split-K partials into the final bf16 output. One block per (batch, -// q_head); each thread owns one head_dim element and folds across all splits -// with a numerically-stable online rescale. -__global__ void gqa_decode_combine_kernel(const float* __restrict__ o_part, - const float* __restrict__ ml_part, - bf16* __restrict__ out, - int num_splits, int head_dim) { - int bh = blockIdx.x; // batch * q_head + head - int d = threadIdx.x; - if (d >= head_dim) return; - - const float* mlp = ml_part + (size_t)bh * num_splits * 2; - float mstar = -FLT_MAX; - for (int s = 0; s < num_splits; s++) - mstar = fmaxf(mstar, mlp[s * 2]); - - float lstar = 0.0f; - for (int s = 0; s < num_splits; s++) { - float mi = mlp[s * 2]; - if (mi > -FLT_MAX) lstar += mlp[s * 2 + 1] * __expf(mi - mstar); - } - - const float* op = o_part + (size_t)bh * num_splits * head_dim; - float acc = 0.0f; - for (int s = 0; s < num_splits; s++) { - float mi = mlp[s * 2]; - if (mi > -FLT_MAX) acc += op[s * head_dim + d] * __expf(mi - mstar); - } - float inv = (lstar > 1e-20f) ? (1.0f / lstar) : 0.0f; - out[(size_t)bh * head_dim + d] = __float2bfloat16(acc * inv); -} \ No newline at end of file diff --git a/csrc/tests/gqa_decode_test.cu b/csrc/tests/attn_decode_test.cu similarity index 84% rename from csrc/tests/gqa_decode_test.cu rename to csrc/tests/attn_decode_test.cu index 0dfb910..fdd914a 100644 --- a/csrc/tests/gqa_decode_test.cu +++ b/csrc/tests/attn_decode_test.cu @@ -2,16 +2,16 @@ Pure-C test: nvcc -I csrc -arch=sm_89 -O3 \ --use_fast_math --ptxas-options=-O3 --extra-device-vectorization \ - csrc/tests/gqa_decode_test.cu -o test && ./test + csrc/tests/attn_decode_test.cu -o test && ./test */ #include #include #include #include -#include "../kernels/gqa_decode_attn.cuh" +#include "../kernels/attn_decode_split_kv.cuh" #ifndef ASTRAI_NO_MMA -#include "../kernels/gqa_decode_attn_mma.cuh" +#include "../kernels/attn_decode_split_kv_mma.cuh" #endif static double now_ms() { @@ -28,63 +28,58 @@ struct DecodeScratch { float* ml_part = nullptr; }; -// Launch the production decode path (tensor-core head-packing MMA on sm_80+, -// scalar fallback otherwise), mirroring dispatch_decode() in gqa_decode_attn.cu. -#ifndef ASTRAI_NO_MMA -static bool decode_use_mma(const GQAParams& p) { - int G = p.q_head / p.kv_head; - return !p.use_mask && G > 1 && G <= 16; -} - -static int decode_num_splits(const GQAParams& p, int tiles_total) { +static int decode_num_splits(const AttentionParams& p, int tiles_total) { int sm_count = 0; cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0); int base_blocks = p.kv_head * p.batch; int desired = 2 * (sm_count > 0 ? sm_count : 64); int n = (desired + base_blocks - 1) / base_blocks; - return max(1, min(n, min(tiles_total, 32))); + int max_by_work = tiles_total / 8; + return max(1, min(min(min(n, tiles_total), 32), max_by_work)); +} + +// Launch the production decode path (tensor-core head-packing MMA on sm_80+, +// scalar fallback otherwise), mirroring dispatch_decode() in attn_decode.cu. +#ifndef ASTRAI_NO_MMA +static bool decode_use_mma(const AttentionParams& p) { + int G = p.q_head / p.kv_head; + return !p.use_mask && G > 1 && G <= 16; } template -static void launch_mma_decode(GQAParams& p, DecodeScratch& sc) { - constexpr int BR = 16, LD = HEAD_DIM; - int smem = (2 * BC * LD + BR * LD) * (int)sizeof(bf16); +static void launch_mma_decode(AttentionParams& p, DecodeScratch& sc) { int tiles_total = (p.kv_len + BC - 1) / BC; - int num_splits = decode_num_splits(p, tiles_total); + p.num_splits = decode_num_splits(p, tiles_total); + p.o_part = sc.o_part; + p.ml_part = sc.ml_part; - if (num_splits <= 1) { - cudaFuncSetAttribute(gqa_decode_attn_mma_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - gqa_decode_attn_mma_kernel - <<>>(p); - return; - } - cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem); - gqa_decode_attn_mma_splitk_kernel - <<>>( - p, sc.o_part, sc.ml_part, num_splits); - gqa_decode_combine_kernel<<>>( - sc.o_part, sc.ml_part, p.o, num_splits, p.head_dim); + attn_decode_split_kv_mma_kernel + <<>>(p); + attn_decode_combine_kernel<<>>(p); } #endif -static void launch_scalar_decode(const GQAParams& p) { +static void launch_scalar_decode(AttentionParams& p, DecodeScratch& sc) { int gs = p.q_head / p.kv_head; + int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK; + p.num_splits = decode_num_splits(p, chunks_total); + p.o_part = sc.o_part; + p.ml_part = sc.ml_part; + size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); - gqa_decode_attn_kernel<<>>(p); + attn_decode_split_kv_kernel<<>>(p); + attn_decode_combine_kernel<<>>(p); } template -static void dispatch_decode_t(GQAParams& p, DecodeScratch& sc) { +static void dispatch_decode_t(AttentionParams& p, DecodeScratch& sc) { #ifndef ASTRAI_NO_MMA if (decode_use_mma(p)) { launch_mma_decode(p, sc); return; } #endif - (void)sc; - launch_scalar_decode(p); + launch_scalar_decode(p, sc); } -static void dispatch_decode(GQAParams& p, DecodeScratch& sc) { +static void dispatch_decode(AttentionParams& p, DecodeScratch& sc) { switch (p.head_dim) { case 32: dispatch_decode_t<32>(p, sc); break; case 64: dispatch_decode_t<64>(p, sc); break; @@ -169,7 +164,7 @@ static void bench() { for (size_t i=0;i #include #include #include -#include "../kernels/gqa_prefill_attn.cuh" +#include "../kernels/attn_prefill_split_q.cuh" #ifndef ASTRAI_NO_MMA -#include "../kernels/gqa_prefill_attn_mma.cuh" +#include "../kernels/attn_prefill_split_q_mma.cuh" #endif static double now_ms() { @@ -21,9 +21,9 @@ static double now_ms() { } // Launch the production prefill path (tensor-core MMA on sm_80+, else the -// scalar fallback), mirroring dispatch_prefill() in gqa_prefill_attn.cu. +// scalar fallback), mirroring dispatch_prefill() in attn_prefill.cu. template -static void launch_prefill(GQAParams& p) { +static void launch_prefill(AttentionParams& p) { #ifndef ASTRAI_NO_MMA constexpr int WARPS = 4, BR = 16; constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16; @@ -31,16 +31,16 @@ static void launch_prefill(GQAParams& p) { : (HEAD_DIM <= 128) ? 3 : 2; dim3 grid((p.q_len + BR * WARPS - 1) / (BR * WARPS), p.q_head, p.batch); dim3 block(WARPS * 32, 1, 1); - gqa_prefill_attn_mma_kernel<<>>(p); + attn_prefill_split_q_mma_kernel<<>>(p); #else constexpr int G = 8, ROWS = 32, P_BC = 32; dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch); dim3 block(G, ROWS, 1); - gqa_prefill_attn_kernel_t<<>>(p); + attn_prefill_split_q_kernel_t<<>>(p); #endif } -static void dispatch_prefill(GQAParams& p) { +static void dispatch_prefill(AttentionParams& p) { switch (p.head_dim) { case 64: launch_prefill<64>(p); break; case 128: launch_prefill<128>(p); break; @@ -123,7 +123,7 @@ static void bench() { for (size_t i=0;i