From fd65b9bc23b34cdda7052b3262717e960dfa9ff7 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 8 Jul 2026 13:59:54 +0800 Subject: [PATCH] feat: support HEAD_DIM=32 and split extension into loader/ops - add case 32 to decode/prefill dispatch switch - fix swiz_col out-of-bounds for HEAD_DIM=32: XOR mask now limited to chunk count (3 for 32, 7 for >=64) instead of always 7, which produced column offsets >= LD=32 and corrupted shared memory - restructure decode dispatch to #ifndef/#else/#endif matching prefill - split astrai/extension/__init__.py into loader.py (kernel .so discovery) and ops.py (wrapper functions + torch SDPA fallback); __init__.py now re-exports the public API --- astrai/extension/__init__.py | 102 ++++---------------------- astrai/extension/loader.py | 36 +++++++++ astrai/extension/ops.py | 86 ++++++++++++++++++++++ csrc/kernels/gqa_decode_attn.cu | 14 +++- csrc/kernels/gqa_mma_utils.cuh | 10 ++- csrc/kernels/gqa_prefill_attn.cu | 5 +- csrc/kernels/gqa_prefill_attn_mma.cuh | 17 +++-- 7 files changed, 168 insertions(+), 102 deletions(-) create mode 100644 astrai/extension/loader.py create mode 100644 astrai/extension/ops.py diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index a1efda1..4a3b3a4 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -1,91 +1,19 @@ -import importlib -import logging +"""CUDA attention kernel wrappers with torch fallback. -import torch -import torch.nn.functional as F +Public API: + - ``gqa_decode_attn`` — single-query decode attention + - ``gqa_prefill_attn`` — multi-query prefill attention -logger = logging.getLogger(__name__) +Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.gqa_*``) +when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``. +""" -_available: dict[str, bool] = {} -_modules: dict[str, object] = {} +from astrai.extension.loader import KERNEL_NAMES, is_available +from astrai.extension.ops import gqa_decode_attn, gqa_prefill_attn -for _name in ["gqa_decode_attn", "gqa_prefill_attn"]: - try: - _mod = importlib.import_module(f".{_name}", package=__package__) - _available[_name] = True - _modules[_name] = _mod - except ImportError: - _available[_name] = False - _modules[_name] = None - - -def _expand_kv_heads( - k: torch.Tensor, v: torch.Tensor, q_head: int -) -> tuple[torch.Tensor, torch.Tensor]: - """Expand K/V heads to match Q heads for GQA fallback.""" - kv_head = k.size(1) - if kv_head == q_head: - return k, v - group = q_head // kv_head - k = k.repeat_interleave(group, dim=1) - v = v.repeat_interleave(group, dim=1) - return k, v - - -def _torch_fallback( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - mask: torch.Tensor | None, - is_causal: bool, - scale: float | None, -) -> torch.Tensor: - k, v = _expand_kv_heads(k, v, q.size(1)) - attn_mask = mask[:, None, None, :] if mask is not None else None - return F.scaled_dot_product_attention( - q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale - ) - - -def gqa_decode_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - mask: torch.Tensor | None = None, - is_causal: bool = False, - causal_offset: int = 0, - scale: float | None = None, -) -> torch.Tensor: - if _available["gqa_decode_attn"]: - return _modules["gqa_decode_attn"].gqa_decode_attn( - q, - k, - v, - mask=mask, - is_causal=is_causal, - causal_offset=causal_offset, - scale=scale, - ) - return _torch_fallback(q, k, v, mask, is_causal, scale) - - -def gqa_prefill_attn( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - mask: torch.Tensor | None = None, - is_causal: bool = False, - causal_offset: int = 0, - scale: float | None = None, -) -> torch.Tensor: - if _available["gqa_prefill_attn"]: - return _modules["gqa_prefill_attn"].gqa_prefill_attn( - q, - k, - v, - mask=mask, - is_causal=is_causal, - causal_offset=causal_offset, - scale=scale, - ) - return _torch_fallback(q, k, v, mask, is_causal, scale) +__all__ = [ + "gqa_decode_attn", + "gqa_prefill_attn", + "is_available", + "KERNEL_NAMES", +] diff --git a/astrai/extension/loader.py b/astrai/extension/loader.py new file mode 100644 index 0000000..21afbeb --- /dev/null +++ b/astrai/extension/loader.py @@ -0,0 +1,36 @@ +"""Dynamic discovery and loading of compiled CUDA kernel modules. + +Each kernel is registered in ``csrc/build.py`` and built into a ``.so`` placed +in this package directory. On import we try to load each one; kernels that +failed to build (or are running on a CPU-only machine) are marked unavailable +so the wrapper functions can fall back to ``torch`` SDPA. +""" + +import importlib +import logging + +logger = logging.getLogger(__name__) + +KERNEL_NAMES = ["gqa_decode_attn", "gqa_prefill_attn"] + +_available: dict[str, bool] = {} +_modules: dict[str, object] = {} + +for _name in KERNEL_NAMES: + try: + _mod = importlib.import_module(f".{_name}", package=__package__) + _available[_name] = True + _modules[_name] = _mod + except ImportError: + _available[_name] = False + _modules[_name] = None + + +def is_available(name: str) -> bool: + """Return ``True`` if the compiled kernel ``name`` was loaded.""" + return _available.get(name, False) + + +def get_module(name: str) -> object: + """Return the loaded kernel module for ``name``, or ``None`` if unavailable.""" + return _modules.get(name) diff --git a/astrai/extension/ops.py b/astrai/extension/ops.py new file mode 100644 index 0000000..1be3464 --- /dev/null +++ b/astrai/extension/ops.py @@ -0,0 +1,86 @@ +"""GQA attention wrapper functions — one entry point per compiled kernel. + +Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when +available, otherwise falls back to ``torch`` SDPA. + +Add new kernel wrappers here; split into per-variant files only if this file +grows large. +""" + +import torch +import torch.nn.functional as F + +from astrai.extension.loader import _available, _modules + + +def _expand_kv_heads( + k: torch.Tensor, v: torch.Tensor, q_head: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Expand K/V heads to match Q heads for GQA fallback.""" + kv_head = k.size(1) + if kv_head == q_head: + return k, v + group = q_head // kv_head + k = k.repeat_interleave(group, dim=1) + v = v.repeat_interleave(group, dim=1) + return k, v + + +def _torch_fallback( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor | None, + is_causal: bool, + scale: float | None, +) -> torch.Tensor: + """Reference attention via ``scaled_dot_product_attention``.""" + k, v = _expand_kv_heads(k, v, q.size(1)) + attn_mask = mask[:, None, None, :] if mask is not None else None + return F.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale + ) + + +def gqa_decode_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor | None = None, + is_causal: bool = False, + causal_offset: int = 0, + scale: float | None = None, +) -> torch.Tensor: + if _available["gqa_decode_attn"]: + return _modules["gqa_decode_attn"].gqa_decode_attn( + q, + k, + v, + mask=mask, + is_causal=is_causal, + causal_offset=causal_offset, + scale=scale, + ) + return _torch_fallback(q, k, v, mask, is_causal, scale) + + +def gqa_prefill_attn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor | None = None, + is_causal: bool = False, + causal_offset: int = 0, + scale: float | None = None, +) -> torch.Tensor: + if _available["gqa_prefill_attn"]: + return _modules["gqa_prefill_attn"].gqa_prefill_attn( + q, + k, + v, + mask=mask, + is_causal=is_causal, + causal_offset=causal_offset, + scale=scale, + ) + return _torch_fallback(q, k, v, mask, is_causal, scale) diff --git a/csrc/kernels/gqa_decode_attn.cu b/csrc/kernels/gqa_decode_attn.cu index 3f2aba8..72286a1 100644 --- a/csrc/kernels/gqa_decode_attn.cu +++ b/csrc/kernels/gqa_decode_attn.cu @@ -21,13 +21,20 @@ static void dispatch_decode(GQAParams& p) { gqa_decode_attn_mma_kernel<<>>(p); return; } -#endif // scalar fallback (per-KV-head, one warp per query head) int group_size = p.q_head / p.kv_head; size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); dim3 block(32, group_size); dim3 grid(p.batch * p.kv_head); gqa_decode_attn_kernel<<>>(p); +#else + // scalar fallback (per-KV-head, one warp per query head) + int group_size = p.q_head / p.kv_head; + size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); + dim3 block(32, group_size); + dim3 grid(p.batch * p.kv_head); + gqa_decode_attn_kernel<<>>(p); +#endif } torch::Tensor gqa_decode_attn( @@ -74,6 +81,9 @@ torch::Tensor gqa_decode_attn( p.o = (bf16*)O.data_ptr(); switch (p.head_dim) { + case 32: + dispatch_decode<32>(p); + break; case 64: dispatch_decode<64>(p); break; @@ -85,7 +95,7 @@ torch::Tensor gqa_decode_attn( break; default: TORCH_CHECK(false, "decode: unsupported head_dim ", p.head_dim, - " (supported: 64, 128, 256)"); + " (supported: 32, 64, 128, 256)"); } return O; } diff --git a/csrc/kernels/gqa_mma_utils.cuh b/csrc/kernels/gqa_mma_utils.cuh index 5cdd6be..f8932dc 100644 --- a/csrc/kernels/gqa_mma_utils.cuh +++ b/csrc/kernels/gqa_mma_utils.cuh @@ -58,10 +58,12 @@ __device__ __forceinline__ void ldmatrix_x2_trans(unsigned* r, const bf16* p) { // XOR swizzle for shared-memory column at 8-bf16 chunk granularity. // Eliminates ldmatrix bank conflicts without LD padding: consecutive rows -// land in distinct bank groups. swiz_col(d, r) = ((d>>3)^(r&7))<<3 | (d&7). -// Works for any d; aligned (d%8==0) simplifies to d ^ ((r&7)<<3). -__device__ __forceinline__ int swiz_col(int d, int r) { - return ((d >> 3) ^ (r & 7)) << 3 | (d & 7); +// land in distinct bank groups. swiz_col(d, r, mask) = ((d>>3)^(r&mask))<<3 | (d&7). +// mask must cover log2(HEAD_DIM/8) chunk bits but stay within LD: use 7 for +// HEAD_DIM>=64 (8+ chunks), 3 for HEAD_DIM=32 (4 chunks). Default 7 keeps +// existing HEAD_DIM>=64 call sites working unchanged. +__device__ __forceinline__ int swiz_col(int d, int r, int mask = 7) { + return ((d >> 3) ^ (r & mask)) << 3 | (d & 7); } // cp.async: copy 16 bytes (8 bf16) from global to shared memory directly, diff --git a/csrc/kernels/gqa_prefill_attn.cu b/csrc/kernels/gqa_prefill_attn.cu index d9dced0..0936573 100644 --- a/csrc/kernels/gqa_prefill_attn.cu +++ b/csrc/kernels/gqa_prefill_attn.cu @@ -68,6 +68,9 @@ torch::Tensor gqa_prefill_attn( p.o = (bf16*)O.data_ptr(); switch (p.head_dim) { + case 32: + dispatch_prefill<32>(p); + break; case 64: dispatch_prefill<64>(p); break; @@ -79,7 +82,7 @@ torch::Tensor gqa_prefill_attn( break; default: TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim, - " (supported: 64,128,256)"); + " (supported: 32,64,128,256)"); } return O; } diff --git a/csrc/kernels/gqa_prefill_attn_mma.cuh b/csrc/kernels/gqa_prefill_attn_mma.cuh index a1ebee1..ccced71 100644 --- a/csrc/kernels/gqa_prefill_attn_mma.cuh +++ b/csrc/kernels/gqa_prefill_attn_mma.cuh @@ -30,6 +30,7 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { 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 (swiz_col) handles bank conflicts + constexpr int SWIZ_MASK = (HEAD_DIM >= 64) ? 7 : (HEAD_DIM / 8 - 1); // chunk bits, stay within LD const int warp = threadIdx.x / 32; const int lane = threadIdx.x % 32; @@ -61,12 +62,12 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { int qr = qrow0 + r; bf16 qv = (qr < p.q_len) ? p.q[q_base + qr * HEAD_DIM + d] : __float2bfloat16(0.0f); - sQ[r * LD + swiz_col(d, r)] = __hmul(qv, scale_bf16); + sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(qv, scale_bf16); } __syncwarp(); #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)]); + ldmatrix_x4(Qa[kt], &sQ[qrow_l * LD + swiz_col(kt * 16 + qcol_l, qrow_l, SWIZ_MASK)]); } __syncthreads(); // prevent next warp from overwriting sQ prematurely } @@ -106,8 +107,8 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { int r = i / HEAD_DIM; int d = i % HEAD_DIM; int kc = kv0 + r; - cp_async_16(&sK[r * LD + swiz_col(d, r)], &p.k[kv_base + kc * HEAD_DIM + d]); - cp_async_16(&sV[r * LD + swiz_col(d, r)], &p.v[kv_base + kc * HEAD_DIM + d]); + 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(); @@ -116,9 +117,9 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { 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)] = (kc < p.kv_len) + 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)] = (kc < p.kv_len) + sV[r * LD + swiz_col(d, r, SWIZ_MASK)] = (kc < p.kv_len) ? p.v[kv_base + kc * HEAD_DIM + d] : z; } } @@ -137,7 +138,7 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { #pragma unroll for (int kt = 0; kt < KD; kt++) { unsigned b[2]; - ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l)]); + ldmatrix_x2(b, &sK[krow_l * LD + swiz_col(kt * 16 + kcol_h, krow_l, SWIZ_MASK)]); mma16816(Sacc[n8], Qa[kt], b, Sacc[n8]); } } @@ -218,7 +219,7 @@ __global__ void gqa_prefill_attn_mma_kernel(GQAParams p) { #pragma unroll for (int dn8 = 0; dn8 < DN8; dn8++) { unsigned b[2]; - ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l)]); + ldmatrix_x2_trans(b, &sV[vrow_l * LD + swiz_col(dn8 * 8, vrow_l, SWIZ_MASK)]); mma16816(Oacc[dn8], Pa, b, Oacc[dn8]); } }