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

perf: remove split partials memset and overlap decode tile loads

Browse Source 
- alloc_split_partials now uses torch::empty: the split kernel writes every slot it owns, so the per-call zeros/full memset was pure overhead (2 kernels per layer per step)
- decode split-KV MMA kernels now run a true multi-stage cp.async pipeline (wait_group<STAGES-1> instead of wait_group<0>), keeping STAGES-1 tile loads in flight; the old wait_group<0> serialized load and compute so deeper STAGES made no difference
- add a fallback path when ntiles < STAGES to avoid a race on the last tile
This commit is contained in:
ViperEkura
2026-07-31 22:37:44 +08:00
parent 21ddead238
commit 530d280e33
4 changed files with 64 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
-1
View File
@@ -30,7 +30,6 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
import contextvars
import enum
import math
from abc import ABC, abstractmethod
from contextlib import contextmanager
from typing import Optional, Union
csrc/kernels/attn_decode_split_kv_mma.cuh
+29 -21
View File
@@ -76,26 +76,17 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
cp_async_commit();
};
constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0;
// Prologue
if (ti_begin < ti_end) {
load_tile(ti_begin, 0);
}
for (int ti = ti_begin; ti < ti_end; ti++) {
int buf = (ti - ti_begin) & BUF_MASK;
cp_async_wait_group<0>();
__syncwarp();
if constexpr (Traits::STAGES > 1) {
if (ti + 1 < ti_end)
load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
}
// ---- Multi-stage cp.async pipeline ----
// Prologue loads STAGES tiles; each loop iteration waits only for the
// oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
// tile loads stay in flight and overlap with the current tile's compute.
constexpr int STAGES = Traits::STAGES;
const int ntiles = ti_end - ti_begin;
auto process_tile = [&](int it, int buf) {
const bf16* bK = sK + buf * Traits::BC * Traits::LD;
const bf16* bV = sV + buf * Traits::BC * Traits::LD;
int kv0 = ti * Traits::BC;
int kv0 = (ti_begin + it) * Traits::BC;
float Sacc[Traits::NC8][4];
mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
@@ -115,12 +106,29 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
__syncwarp();
};
if constexpr (Traits::STAGES == 1) {
if (ti + 1 < ti_end)
load_tile(ti + 1, 0);
if (ntiles >= STAGES) {
#pragma unroll
for (int i = 0; i < STAGES; i++)
load_tile(ti_begin + i, i);
for (int it = 0; it < ntiles; it++) {
cp_async_wait_group<STAGES - 1>();
__syncwarp();
process_tile(it, it & (STAGES - 1));
__syncwarp();
if (it + STAGES < ntiles)
load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
}
} else {
// Fewer tiles than stages: load all, wait for all, process.
for (int i = 0; i < ntiles; i++)
load_tile(ti_begin + i, i);
cp_async_wait_group<0>();
__syncwarp();
for (int it = 0; it < ntiles; it++)
process_tile(it, it);
}
// ---- write UN-normalised partials for this split ----
csrc/kernels/attn_entry_utils.cuh
+6 -2
View File
@@ -21,11 +21,15 @@ using bf16 = __nv_bfloat16;
" (supported: 32, 64, 128, 256)"); \
}
// The split kernel unconditionally writes every (batch, q_head, split) slot it
// owns — including empty split ranges, which store m = -FLT_MAX so the combine
// skips them. Allocators are therefore left uninitialized (torch::empty); the
// per-call memset (torch::zeros / torch::full) was pure overhead.
template<typename P>
inline void alloc_split_partials(P& p) {
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
auto o_part = torch::zeros(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
auto ml_part = torch::full(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, -FLT_MAX, fopt);
auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
p.o_part = (float*)o_part.data_ptr();
p.ml_part = (float*)ml_part.data_ptr();
}
csrc/kernels/attn_paged_decode_split_kv_mma.cuh
+29 -20
View File
@@ -91,25 +91,17 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
cp_async_commit();
};
constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0;
if (ti_begin < ti_end) {
load_tile(ti_begin, 0);
}
for (int ti = ti_begin; ti < ti_end; ti++) {
int buf = (ti - ti_begin) & BUF_MASK;
cp_async_wait_group<0>();
__syncwarp();
if constexpr (Traits::STAGES > 1) {
if (ti + 1 < ti_end)
load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
}
// ---- Multi-stage cp.async pipeline ----
// Prologue loads STAGES tiles; each loop iteration waits only for the
// oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
// tile loads stay in flight and overlap with the current tile's compute.
constexpr int STAGES = Traits::STAGES;
const int ntiles = ti_end - ti_begin;
auto process_tile = [&](int it, int buf) {
const bf16* bK = sK + buf * Traits::BC * Traits::LD;
const bf16* bV = sV + buf * Traits::BC * Traits::LD;
int kv0 = ti * Traits::BC;
int kv0 = (ti_begin + it) * Traits::BC;
float Sacc[Traits::NC8][4];
mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
@@ -128,12 +120,29 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
Sacc, Oacc, m0, m1, l0, l1, lane);
mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
__syncwarp();
};
if constexpr (Traits::STAGES == 1) {
if (ti + 1 < ti_end)
load_tile(ti + 1, 0);
if (ntiles >= STAGES) {
#pragma unroll
for (int i = 0; i < STAGES; i++)
load_tile(ti_begin + i, i);
for (int it = 0; it < ntiles; it++) {
cp_async_wait_group<STAGES - 1>();
__syncwarp();
process_tile(it, it & (STAGES - 1));
__syncwarp();
if (it + STAGES < ntiles)
load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
}
} else {
// Fewer tiles than stages: load all, wait for all, process.
for (int i = 0; i < ntiles; i++)
load_tile(ti_begin + i, i);
cp_async_wait_group<0>();
__syncwarp();
for (int it = 0; it < ntiles; it++)
process_tile(it, it);
}
auto split_slot = [&](int h) -> size_t {