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

refactor: template combine kernel, fix mask bug, unify dispatch 

- Template combine kernel, share macros, extract entry_utils helpers
- Fix mask indexing (pass stride not pre-multiplied base)
- Remove !p.use_mask — MMA handles mask
This commit is contained in:
ViperEkura 2026-07-15 21:43:29 +08:00
parent 1f0be382ad
commit 0654b4b916
7 changed files with 70 additions and 110 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
csrc/kernels/attn_decode.cu
View File

@ -38,7 +38,7 @@ template <int HEAD_DIM>
static void dispatch_decode(AttentionParams<bf16>& p) {
#ifndef ASTRAI_NO_MMA
int G = p.q_head / p.kv_head;
if (!p.use_mask && G >= 1 && G <= 16) {
if (G >= 1 && G <= 16) {
launch_mma_decode<HEAD_DIM, 32>(p);
return;
}
@ -65,7 +65,7 @@ torch::Tensor attn_decode(
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_decode<D>(p); });
DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
return O;
}

3
csrc/kernels/attn_decode_split_kv_mma.cuh
View File

@ -64,7 +64,6 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
// KV: stride-based base — [batch, kv_head, kv_len, head_dim]
const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
const int mask_batch_base = batch * p.mask_b_stride;
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;
@ -125,7 +124,7 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
0, 0,
mask_batch_base, 0,
p.mask_b_stride, 0,
batch,
p.mask, has_mask,
Sacc, Oacc, m0, m1, l0, l1, lane);

161
csrc/kernels/attn_entry_utils.cuh
View File

@ -5,10 +5,6 @@
using bf16 = __nv_bfloat16;
// ---------------------------------------------------------------------------
// Shared dispatch helpers — eliminates duplication across .cu entry files.
// ---------------------------------------------------------------------------
inline int compute_num_splits(int base_blocks, int tiles_total) {
int sm_count = 0;
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
@ -16,21 +12,20 @@ inline int compute_num_splits(int base_blocks, int tiles_total) {
return std::max(1, std::min(n, std::min(tiles_total, 32)));
}
// Dispatch head_dim to a generic lambda FN (callable as FN.operator()<D>()).
template <typename Fn>
inline void dispatch_head_dim(int hd, Fn&& fn) {
switch (hd) {
case 32: fn.template operator()<32>(); break;
case 64: fn.template operator()<64>(); break;
case 128: fn.template operator()<128>(); break;
case 256: fn.template operator()<256>(); break;
default:
TORCH_CHECK(false, "unsupported head_dim ", hd,
" (supported: 32, 64, 128, 256)");
}
// Dispatch head_dim: shared macro — avoids C++20 lambda template syntax.
// Usage: DISPATCH_HEAD_DIM(hd, fn, arg)
// Expands to: fn<32>(arg); fn<64>(arg); etc.
#define DISPATCH_HEAD_DIM(hd, fn, arg) \
switch (hd) { \
case 32: fn<32>(arg); break; \
case 64: fn<64>(arg); break; \
case 128: fn<128>(arg); break; \
case 256: fn<256>(arg); break; \
default: \
TORCH_CHECK(false, "unsupported head_dim ", hd, \
" (supported: 32, 64, 128, 256)"); \
}
// Allocate split-KV partial buffers and wire into params.
template<typename P>
inline void alloc_split_partials(P& p) {
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
@ -40,10 +35,48 @@ inline void alloc_split_partials(P& p) {
p.ml_part = (float*)ml_part.data_ptr();
}
// ---------------------------------------------------------------------------
// Param packing — fills AttentionParams from torch tensors with validation.
// ---------------------------------------------------------------------------
// ---- Shared Q-dims + strides extraction ----
template <typename P>
inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
if (layout == 1) q = q.transpose(1, 2);
p.batch = (int)q.size(0);
p.q_head = (int)q.size(1);
p.q_len = (int)q.size(2);
p.head_dim = (int)q.size(3);
p.q_stride_b = (int)q.stride(0);
p.q_stride_h = (int)q.stride(1);
p.q_stride_l = (int)q.stride(2);
p.q_stride_d = (int)q.stride(3);
}
// ---- Shared mask packing ----
template <typename P>
inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
if (p.use_mask) {
auto m = mask.value();
TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
if (m.dim() == 2) {
p.mask_b_stride = (int)m.stride(0);
p.mask_q_stride = 0;
} else if (m.dim() == 3) {
TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_q_stride = (int)m.stride(1);
} else {
TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
}
p.mask = m.data_ptr<bool>();
} else {
p.mask = nullptr;
p.mask_b_stride = 0;
p.mask_q_stride = 0;
}
}
// ---- attn_pack_params (contiguous KV) ----
template<typename T>
inline void attn_pack_params(
torch::Tensor q,
@ -52,7 +85,7 @@ inline void attn_pack_params(
c10::optional<torch::Tensor> mask,
int64_t causal_offset,
double scale,
int64_t layout, // 0 = b h l d, 1 = b l h d
int64_t layout,
AttentionParams<T>& p
) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
@ -64,26 +97,14 @@ inline void attn_pack_params(
TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes");
TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D");
// Normalize to b h l d view (zero-copy transpose if user passed b l h d)
if (layout == 1) {
q = q.transpose(1, 2);
k = k.transpose(1, 2);
v = v.transpose(1, 2);
}
extract_q_dims_and_strides(q, layout, p);
if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2);
p.batch = (int)q.size(0);
p.q_head = (int)q.size(1);
p.q_len = (int)q.size(2);
p.head_dim = (int)q.size(3);
p.kv_head = (int)k.size(1);
p.kv_len = (int)k.size(2);
TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q");
// Strides (layout-agnostic: works for b h l d and b l h d)
p.q_stride_b = (int)q.stride(0);
p.q_stride_h = (int)q.stride(1);
p.q_stride_l = (int)q.stride(2);
p.q_stride_d = (int)q.stride(3);
p.kv_stride_b = (int)k.stride(0);
p.kv_stride_h = (int)k.stride(1);
p.kv_stride_l = (int)k.stride(2);
@ -100,34 +121,10 @@ inline void attn_pack_params(
p.o_part = nullptr;
p.ml_part = nullptr;
if (p.use_mask) {
auto m = mask.value();
TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
if (m.dim() == 2) {
p.mask_b_stride = (int)m.stride(0);
p.mask_q_stride = 0;
} else if (m.dim() == 3) {
TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_q_stride = (int)m.stride(1);
} else {
TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
}
p.mask = m.data_ptr<bool>();
} else {
p.mask = nullptr;
p.mask_b_stride = 0;
p.mask_q_stride = 0;
}
pack_mask(mask, p);
}
// ---------------------------------------------------------------------------
// Param packing for paged attention.
// ---------------------------------------------------------------------------
// ---- attn_pack_paged_params ----
template<typename T>
inline void attn_pack_paged_params(
torch::Tensor q,
@ -139,7 +136,7 @@ inline void attn_pack_paged_params(
c10::optional<torch::Tensor> mask,
int64_t causal_offset,
double scale,
int64_t layout, // 0 = b h l d, 1 = b l h d
int64_t layout,
PagedAttentionParams<T>& p
) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
@ -151,15 +148,8 @@ inline void attn_pack_paged_params(
TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes");
// Normalize Q to b h l d view if user passed b l h d
if (layout == 1) {
q = q.transpose(1, 2);
}
extract_q_dims_and_strides(q, layout, p);
p.batch = (int)q.size(0);
p.q_head = (int)q.size(1);
p.q_len = (int)q.size(2);
p.head_dim = (int)q.size(3);
p.kv_head = (int)k_cache.size(2);
p.kv_len = (int)kv_len;
p.page_size = (int)page_size;
@ -171,12 +161,6 @@ inline void attn_pack_paged_params(
"k_cache dim 1 must equal page_size, got ",
k_cache.size(1), " vs ", page_size);
// Q strides
p.q_stride_b = (int)q.stride(0);
p.q_stride_h = (int)q.stride(1);
p.q_stride_l = (int)q.stride(2);
p.q_stride_d = (int)q.stride(3);
p.causal_offset = (int)causal_offset;
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim);
@ -189,26 +173,5 @@ inline void attn_pack_paged_params(
p.o_part = nullptr;
p.ml_part = nullptr;
if (p.use_mask) {
auto m = mask.value();
TORCH_CHECK(m.is_cuda(), "mask must be on CUDA");
TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool");
TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch");
TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch");
if (m.dim() == 2) {
p.mask_b_stride = (int)m.stride(0);
p.mask_q_stride = 0;
} else if (m.dim() == 3) {
TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_q_stride = (int)m.stride(1);
} else {
TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
}
p.mask = m.data_ptr<bool>();
} else {
p.mask = nullptr;
p.mask_b_stride = 0;
p.mask_q_stride = 0;
}
pack_mask(mask, p);
}

4
csrc/kernels/attn_paged_decode.cu
View File

@ -34,7 +34,7 @@ template <int HEAD_DIM>
static void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
#ifndef ASTRAI_NO_MMA
int G = p.q_head / p.kv_head;
if (!p.use_mask && G >= 1 && G <= 16 && p.page_size >= 32) {
if (G >= 1 && G <= 16 && p.page_size >= 32) {
launch_paged_mma_decode<HEAD_DIM, 32>(p);
return;
}
@ -62,7 +62,7 @@ torch::Tensor attn_paged_decode(
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_paged_decode<D>(p); });
DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
return O;
}

3
csrc/kernels/attn_paged_decode_split_kv_mma.cuh
View File

@ -51,7 +51,6 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
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 mask_batch_base = batch * p.mask_b_stride;
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;
@ -119,7 +118,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
0, 0,
mask_batch_base, 0,
p.mask_b_stride, 0,
batch,
p.mask, has_mask,
Sacc, Oacc, m0, m1, l0, l1, lane);

2
csrc/kernels/attn_prefill.cu
View File

@ -47,7 +47,7 @@ torch::Tensor attn_prefill(
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_prefill<D>(p); });
DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
return O;
}

3
csrc/kernels/attn_prefill_split_q_mma.cuh
View File

@ -94,7 +94,6 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
const int block_max_kv =
blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset;
const int has_mask = p.use_mask && p.mask;
const int mask_batch_base = batch * p.mask_b_stride;
// Last active tile: block-level causal bound (all warps in the block share
// the K/V load, so the prefetch range is the block max, not per-warp).
@ -165,7 +164,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
: p.kv_len;
mma_softmax_tile<NC8, DN8>(kv0, maxc0, maxc1,
qr0, qr1,
mask_batch_base, p.mask_q_stride,
p.mask_b_stride, p.mask_q_stride,
batch,
p.mask, has_mask,
Sacc, Oacc, m0, m1, l0, l1, lane);