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

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
This commit is contained in:
ViperEkura 2026-07-08 13:59:54 +08:00
parent 9ebaea840f
commit fd65b9bc23
7 changed files with 168 additions and 102 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

102
astrai/extension/__init__.py
View File

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

36
astrai/extension/loader.py Normal file
View File

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

86
astrai/extension/ops.py Normal file
View File

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

14
csrc/kernels/gqa_decode_attn.cu
View File

@ -21,13 +21,20 @@ static void dispatch_decode(GQAParams& p) {
gqa_decode_attn_mma_kernel<HEAD_DIM, BC><<<grid, block, smem>>>(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<<<grid, block, smem>>>(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<<<grid, block, smem>>>(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;
}

10
csrc/kernels/gqa_mma_utils.cuh
View File

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

5
csrc/kernels/gqa_prefill_attn.cu
View File

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

17
csrc/kernels/gqa_prefill_attn_mma.cuh
View File

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