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

refactor: replace magic layout ints with TensorLayout enum

Browse Source 
- Add TensorLayout enum (C++ + Python) to replace magic layout ints
- Add C10_CUDA_CHECK post-launch error checking to all kernel entries
- Add CUDAGuard + freqs_cis shape validation to rotary_emb.cu
- Cache SM count to eliminate per-call cudaDeviceGetAttribute
- Add DISPATCH_CAUSAL_MASK macro to deduplicate dispatcher if/else
- Convert mask type hints from X|None to Optional[X]
This commit is contained in:
ViperEkura
2026-08-01 11:05:52 +08:00
parent 3639b50b4a
commit 7feeb0b93e
9 changed files with 91 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_common.h
+8
View File
@@ -1,5 +1,13 @@
#pragma once
// Tensor layout for Q/K/V tensors passed to attention kernels.
// Internally, kernels always operate on BHLD [batch, n_heads, seq_len, head_dim].
// When the caller passes BLHD, dims 1 and 2 are transposed at entry.
enum TensorLayout : int {
BHLD = 0, // [batch, n_heads, seq_len, head_dim]
BLHD = 1, // [batch, seq_len, n_heads, head_dim]
};
template<typename T, typename AT = float>
struct AttentionParams {
csrc/kernels/attn_decode.cu
+3 -2
View File
@@ -16,11 +16,12 @@ torch::Tensor attn_decode(
TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
alloc_split_partials(p);
DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
C10_CUDA_CHECK(cudaGetLastError());
return O;
}
@@ -32,6 +33,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("mask") = py::none(),
py::arg("causal_offset") = -1,
py::arg("scale") = 0.0,
py::arg("layout") = 0,
py::arg("layout") = (int64_t)BHLD,
"GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
}
csrc/kernels/attn_dispatchers.cuh
+41 -45
View File
@@ -14,18 +14,50 @@
#include "attn_paged_decode_split_kv_mma.cuh"
#endif
// Cached SM count — cudaDeviceGetAttribute is a host-side call that was
// invoked on every decode/paged-decode launch. Cache per-device so multi-GPU
// setups with heterogeneous GPUs still get the right count, while the common
// single-GPU path hits the cache after the first call.
inline int get_sm_count() {
int dev = 0;
cudaGetDevice(&dev);
static int cached_dev = -1;
static int cached_count = 0;
if (dev != cached_dev) {
cudaDeviceGetAttribute(&cached_count, cudaDevAttrMultiProcessorCount, dev);
cached_dev = dev;
}
return cached_count;
}
// Split-KV: compute number of splits to fill all SMs for small-batch decode.
// Caps splits so each split processes at least `min_tiles_per_split` tiles,
// avoiding excessive loop/prologue overhead when tiles are small.
inline int compute_num_splits(int base_blocks, int tiles_total,
int min_tiles_per_split = 1) {
int sm_count = 0;
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
int min_tiles_per_split = 1) {
int sm_count = get_sm_count();
int n = (2 * sm_count + base_blocks - 1) / base_blocks;
int max_by_work = tiles_total / min_tiles_per_split;
return std::max(1, std::min(n, std::min(max_by_work, MAX_SPLITS)));
}
// Dispatch IsCausal × HasMask — eliminates the duplicated 4-way if/else
// ladder that appeared in each dispatch_* function. FN must be a function
// template <int HEAD_DIM, bool IsCausal, bool HasMask>; HEAD_DIM is forwarded
// as the first template argument so callers only spell it once.
//
// Usage: DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_mma, HEAD_DIM, p, group_size);
#define DISPATCH_CAUSAL_MASK(is_causal, has_mask, FN, HEAD_DIM, ...) \
do { \
if (is_causal) { \
if (has_mask) FN<HEAD_DIM, true, true>(__VA_ARGS__); \
else FN<HEAD_DIM, true, false>(__VA_ARGS__); \
} else { \
if (has_mask) FN<HEAD_DIM, false, true>(__VA_ARGS__); \
else FN<HEAD_DIM, false, false>(__VA_ARGS__); \
} \
} while (0)
// ======================================================================
// Prefill
// ======================================================================
@@ -56,21 +88,9 @@ static inline void dispatch_prefill(AttentionParams<bf16>& p) {
bool has_mask = (p.use_mask && p.mask);
#ifndef ASTRAI_NO_MMA
if (is_causal) {
if (has_mask) launch_prefill_mma<HEAD_DIM, true, true>(p);
else launch_prefill_mma<HEAD_DIM, true, false>(p);
} else {
if (has_mask) launch_prefill_mma<HEAD_DIM, false, true>(p);
else launch_prefill_mma<HEAD_DIM, false, false>(p);
}
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_mma, HEAD_DIM, p);
#else
if (is_causal) {
if (has_mask) launch_prefill_scalar<HEAD_DIM, true, true>(p);
else launch_prefill_scalar<HEAD_DIM, true, false>(p);
} else {
if (has_mask) launch_prefill_scalar<HEAD_DIM, false, true>(p);
else launch_prefill_scalar<HEAD_DIM, false, false>(p);
}
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_scalar, HEAD_DIM, p);
#endif
}
@@ -116,21 +136,9 @@ static inline void dispatch_decode(AttentionParams<bf16>& p) {
int group_size = p.q_head / p.kv_head;
#ifndef ASTRAI_NO_MMA
if (is_causal) {
if (has_mask) launch_decode_mma<HEAD_DIM, true, true>(p, group_size);
else launch_decode_mma<HEAD_DIM, true, false>(p, group_size);
} else {
if (has_mask) launch_decode_mma<HEAD_DIM, false, true>(p, group_size);
else launch_decode_mma<HEAD_DIM, false, false>(p, group_size);
}
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_mma, HEAD_DIM, p, group_size);
#else
if (is_causal) {
if (has_mask) launch_decode_scalar<HEAD_DIM, true, true>(p, group_size);
else launch_decode_scalar<HEAD_DIM, true, false>(p, group_size);
} else {
if (has_mask) launch_decode_scalar<HEAD_DIM, false, true>(p, group_size);
else launch_decode_scalar<HEAD_DIM, false, false>(p, group_size);
}
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_scalar, HEAD_DIM, p, group_size);
#endif
attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
@@ -174,21 +182,9 @@ static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
int group_size = p.q_head / p.kv_head;
#ifndef ASTRAI_NO_MMA
if (is_causal) {
if (has_mask) launch_paged_decode_mma<HEAD_DIM, true, true>(p, group_size);
else launch_paged_decode_mma<HEAD_DIM, true, false>(p, group_size);
} else {
if (has_mask) launch_paged_decode_mma<HEAD_DIM, false, true>(p, group_size);
else launch_paged_decode_mma<HEAD_DIM, false, false>(p, group_size);
}
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, group_size);
#else
if (is_causal) {
if (has_mask) launch_paged_decode_scalar<HEAD_DIM, true, true>(p, group_size);
else launch_paged_decode_scalar<HEAD_DIM, true, false>(p, group_size);
} else {
if (has_mask) launch_paged_decode_scalar<HEAD_DIM, false, true>(p, group_size);
else launch_paged_decode_scalar<HEAD_DIM, false, false>(p, group_size);
}
DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size);
#endif
paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
csrc/kernels/attn_entry_utils.cuh
+2 -2
View File
@@ -37,7 +37,7 @@ inline void alloc_split_partials(P& p) {
// ---- 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);
if (layout == BLHD) 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);
@@ -109,7 +109,7 @@ inline void attn_pack_params(
extract_q_dims_and_strides(q, layout, p);
if (layout == 1) k = k.transpose(1, 2), v = v.transpose(1, 2);
if (layout == BLHD) k = k.transpose(1, 2), v = v.transpose(1, 2);
p.kv_head = (int)k.size(1);
p.kv_len = (int)k.size(2);
csrc/kernels/attn_paged_decode.cu
+3 -2
View File
@@ -18,11 +18,12 @@ torch::Tensor attn_paged_decode(
page_size, kv_len, mask, causal_offset, scale, layout, p);
auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
alloc_split_partials(p);
DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
C10_CUDA_CHECK(cudaGetLastError());
return O;
}
@@ -37,6 +38,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("mask") = py::none(),
py::arg("causal_offset") = -1,
py::arg("scale") = 0.0,
py::arg("layout") = 0,
py::arg("layout") = (int64_t)BHLD,
"Paged GQA decode — split-KV with direct page-table access.");
}
csrc/kernels/attn_prefill.cu
+3 -2
View File
@@ -15,10 +15,11 @@ torch::Tensor attn_prefill(
TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
p.o = (bf16*)O_view.data_ptr();
DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p);
C10_CUDA_CHECK(cudaGetLastError());
return O;
}
@@ -30,6 +31,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("mask") = py::none(),
py::arg("causal_offset") = -1,
py::arg("scale") = 0.0,
py::arg("layout") = 0,
py::arg("layout") = (int64_t)BHLD,
"GQA prefill (tensor-core mma on sm_80+, scalar fallback)");
}
csrc/kernels/rotary_emb.cu
+10
View File
@@ -1,4 +1,6 @@
#include <torch/extension.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_bf16.h>
__global__ void rotary_emb_kernel(
@@ -46,6 +48,8 @@ torch::Tensor rotary_emb(
torch::Tensor x,
torch::Tensor freqs_cis
) {
const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
@@ -53,6 +57,7 @@ torch::Tensor rotary_emb(
TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
TORCH_CHECK(freqs_cis.dim() == 4, "freqs_cis must be 4D [batch, seq_len, dim/2, 2]");
TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous");
TORCH_CHECK(freqs_cis.scalar_type() == torch::kFloat32, "freqs_cis must be f32");
int batch = x.size(0);
int seq_len = x.size(1);
@@ -60,6 +65,10 @@ torch::Tensor rotary_emb(
int head_dim = x.size(3);
TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
TORCH_CHECK(freqs_cis.size(0) == batch, "freqs_cis batch mismatch");
TORCH_CHECK(freqs_cis.size(1) == seq_len, "freqs_cis seq_len mismatch");
TORCH_CHECK(freqs_cis.size(2) == head_dim / 2, "freqs_cis dim/2 mismatch");
TORCH_CHECK(freqs_cis.size(3) == 2, "freqs_cis last dim must be 2 [cos, sin]");
auto out = torch::empty_like(x);
@@ -74,6 +83,7 @@ torch::Tensor rotary_emb(
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
batch, seq_len, n_heads, head_dim
);
C10_CUDA_CHECK(cudaGetLastError());
return out;
}