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

perf: enable MMA decode path for G=1 full attention 

- inline decode_use_mma() into dispatch_decode()
- drop G>1 guard, MMA works correctly for G>=1
- decode is memory-bound, tensor cores + cp.async still win at G=1
This commit is contained in:
ViperEkura 2026-07-11 11:45:13 +08:00
parent a4ae7d17fb
commit 2c0b5d0b5e
1 changed files with 4 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
csrc/kernels/attn_decode.cu
View File

@ -30,12 +30,8 @@ static void launch_scalar_decode(AttentionParams<bf16>& 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<bf16>& p) {
int G = p.q_head / p.kv_head;
return !p.use_mask && G > 1 && G <= 16;
}
// MMA head-packing requires G <= 16 (sQ has BR=16 rows). sm_80+ tensor-core
// + cp.async wins even at G=1 (decode is memory-bound, not compute-bound).
template <int HEAD_DIM, int BC>
static void launch_mma_decode(AttentionParams<bf16>& p) {
int tiles_total = (p.kv_len + BC - 1) / BC;
@ -56,7 +52,8 @@ static void launch_mma_decode(AttentionParams<bf16>& p) {
template <int HEAD_DIM>
static void dispatch_decode(AttentionParams<bf16>& p) {
#ifndef ASTRAI_NO_MMA
if (decode_use_mma(p)) {
int G = p.q_head / p.kv_head;
if (!p.use_mask && G >= 1 && G <= 16) {
launch_mma_decode<HEAD_DIM, 32>(p);
return;
}