ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: rename gqa_* to attn_*, split-KV for all decode paths 

- Rename all csrc/kernels/gqa_*.cuh/cu to attn_*, with _split_q / _split_kv
  strategy suffix and optional _mma compute suffix
- Remove non-split MMA decode kernel, keep only split-KV path
- Convert scalar decode fallback to split-KV (o_part/ml_part + combine)
- Move combine kernel to attn_decode_split_kv.cuh (shared by both paths)
- Rename GQAParams to AttentionParams
- Update all C++ #include, PYBIND11, and Python extension references
This commit is contained in:
ViperEkura 2026-07-10 23:35:14 +08:00
parent 29b0423c4e
commit d923ebe38d
16 changed files with 346 additions and 459 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
astrai/extension/__init__.py
View File

@ -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",
]

2
astrai/extension/loader.py
View File

@ -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] = {}

12
astrai/extension/ops.py
View File

@ -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,

4
csrc/build.py
View File

@ -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")

6
csrc/kernels/gqa_common.cuh → csrc/kernels/attn_common.cuh
View File

@ -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;
};

95
csrc/kernels/gqa_decode_attn.cu → csrc/kernels/attn_decode.cu
View File

@ -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.batch * p.kv_head, dim3(32, group_size), smem>>>(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<float>();
p.ml_part = ml_part.data_ptr<float>();
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, group_size), smem>>>(p);
attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(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 <int HEAD_DIM, int BC>
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<HEAD_DIM, BC>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
<<<dim3(p.kv_head, p.batch), 32, smem>>>(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<float>();
p.ml_part = ml_part.data_ptr<float>();
cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
<<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
p, o_part.data_ptr<float>(), ml_part.data_ptr<float>(), num_splits);
gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
o_part.data_ptr<float>(), ml_part.data_ptr<float>(), p.o,
num_splits, p.head_dim);
attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
}
#endif
template <int HEAD_DIM>
static void dispatch_decode(GQAParams& p) {
static void dispatch_decode(AttentionParams& p) {
#ifndef ASTRAI_NO_MMA
if (decode_use_mma(p)) {
launch_mma_decode<HEAD_DIM, 32>(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<double> 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"),

60
csrc/kernels/gqa_decode_attn.cuh → csrc/kernels/attn_decode_split_kv.cuh
View File

@ -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);
}

164
csrc/kernels/attn_decode_split_kv_mma.cuh Normal file
View File

@ -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 <int HEAD_DIM, int BC>
__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<KD, NC8>(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<NC8, DN8>(kv0, maxc, maxc,
mask_base, p.mask, has_mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<DN8, KT2>(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;
}
}
}

6
csrc/kernels/gqa_entry_utils.cuh → csrc/kernels/attn_entry_utils.cuh
View File

@ -1,8 +1,8 @@
#pragma once
#include <torch/extension.h>
#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<double> scale,
GQAParams& p
AttentionParams& p
) {
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
TORCH_CHECK(q.dtype() == torch::kBFloat16);

0
csrc/kernels/gqa_mma_utils.cuh → csrc/kernels/attn_mma_utils.cuh
View File

20
csrc/kernels/gqa_prefill_attn.cu → csrc/kernels/attn_prefill.cu
View File

@ -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 <int HEAD_DIM>
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<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(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<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(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<double> 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"),

4
csrc/kernels/gqa_prefill_attn.cuh → csrc/kernels/attn_prefill_split_q.cuh
View File

@ -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 <int HEAD_DIM, int G, int ROWS, int P_BC>
__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;

6
csrc/kernels/gqa_prefill_attn_mma.cuh → csrc/kernels/attn_prefill_split_q_mma.cuh
View File

@ -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 <int HEAD_DIM, int WARPS, int BC, int MIN_BLOCKS>
__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)

321
csrc/kernels/gqa_decode_attn_mma.cuh
View File

@ -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 <int HEAD_DIM, int BC>
__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<KD, NC8>(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<NC8, DN8>(kv0, maxc, maxc,
mask_base, p.mask, has_mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<DN8, KT2>(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 <int HEAD_DIM, int BC>
__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<KD, NC8>(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<NC8, DN8>(kv0, maxc, maxc,
mask_base, p.mask, has_mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<DN8, KT2>(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);
}

73
csrc/tests/gqa_decode_test.cu → csrc/tests/attn_decode_test.cu
View File

@ -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 <cstdio>
#include <cstdlib>
#include <cmath>
#include <sys/time.h>
#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 <int HEAD_DIM, int BC>
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<HEAD_DIM, BC>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
gqa_decode_attn_mma_kernel<HEAD_DIM, BC>
<<<dim3(p.kv_head, p.batch), 32, smem>>>(p);
return;
}
cudaFuncSetAttribute(gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
gqa_decode_attn_mma_splitk_kernel<HEAD_DIM, BC>
<<<dim3(p.kv_head, p.batch, num_splits), 32, smem>>>(
p, sc.o_part, sc.ml_part, num_splits);
gqa_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(
sc.o_part, sc.ml_part, p.o, num_splits, p.head_dim);
attn_decode_split_kv_mma_kernel<HEAD_DIM, BC>
<<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(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.batch * p.kv_head, dim3(32, gs), smem>>>(p);
attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, gs), smem>>>(p);
attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
}
template <int HEAD_DIM>
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<HEAD_DIM, 32>(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<nKV;i++) tmp[i]=f2bf(randf());
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
GQAParams p;
AttentionParams p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
p.use_mask=0; p.is_causal=0; p.causal_offset=0;
p.scale=1.0f/sqrtf((float)D);
@ -245,7 +240,7 @@ int main() {
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
GQAParams p;
AttentionParams p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
p.use_mask=0; p.is_causal=0; p.causal_offset=0;
p.scale=1.0f/sqrtf((float)D);

20
csrc/tests/gqa_prefill_test.cu → csrc/tests/attn_prefill_test.cu
View File

@ -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_prefill_test.cu -o test && ./test
csrc/tests/attn_prefill_test.cu -o test && ./test
*/
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <sys/time.h>
#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 <int HEAD_DIM>
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<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(p);
attn_prefill_split_q_mma_kernel<HEAD_DIM, WARPS, BC, MIN_BLOCKS><<<grid, block>>>(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<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(p);
attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC><<<grid, block>>>(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<nKV;i++) tmp[i]=f2bf(randf());
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
GQAParams p;
AttentionParams p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
p.scale=1.0f/sqrtf((float)D);
@ -191,7 +191,7 @@ int main() {
for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
GQAParams p;
AttentionParams p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
p.scale=1.0f/sqrtf((float)D);