From 7feeb0b93e7a86a3fedfe354bf7afcc8c9b93137 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 1 Aug 2026 11:04:45 +0800 Subject: [PATCH] refactor: replace magic layout ints with TensorLayout enum - 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] --- astrai/extension/__init__.py | 2 + astrai/extension/attention_ops.py | 25 ++++++--- csrc/kernels/attn_common.h | 8 +++ csrc/kernels/attn_decode.cu | 5 +- csrc/kernels/attn_dispatchers.cuh | 86 +++++++++++++++---------------- csrc/kernels/attn_entry_utils.cuh | 4 +- csrc/kernels/attn_paged_decode.cu | 5 +- csrc/kernels/attn_prefill.cu | 5 +- csrc/kernels/rotary_emb.cu | 10 ++++ 9 files changed, 91 insertions(+), 59 deletions(-) diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index b39aaaf..7fba00d 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -25,6 +25,7 @@ from astrai.extension.attention_backend import ( get_backend, ) from astrai.extension.attention_ops import ( + TensorLayout, attn_decode, attn_paged_decode, attn_prefill, @@ -37,6 +38,7 @@ __all__ = [ "AttentionBackend", "CudaBackend", "TorchNativeBackend", + "TensorLayout", "attention", "attn_backend", "get_backend", diff --git a/astrai/extension/attention_ops.py b/astrai/extension/attention_ops.py index ce9dbaa..e11257a 100644 --- a/astrai/extension/attention_ops.py +++ b/astrai/extension/attention_ops.py @@ -12,11 +12,24 @@ Interface (all functions): mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep) """ +import enum +from typing import Optional + import torch from astrai.extension.loader import _available, _modules +class TensorLayout(enum.IntEnum): + """Q/K/V tensor layout, mirrors the C++ ``TensorLayout`` enum in ``attn_common.h``. + + Kernels internally operate on BHLD; BLHD inputs are transposed at entry. + """ + + BHLD = 0 # [batch, n_heads, seq_len, head_dim] + BLHD = 1 # [batch, seq_len, n_heads, head_dim] + + def _check_available(name: str): if not _available.get(name): raise RuntimeError( @@ -29,7 +42,7 @@ def attn_decode( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - mask: torch.Tensor | None = None, + mask: Optional[torch.Tensor] = None, is_causal: bool = False, ) -> torch.Tensor: """GQA decode attention (q_len == 1). @@ -47,7 +60,7 @@ def attn_decode( _check_available("attn_decode") causal_offset = (k.size(1) - 1) if is_causal else -1 return _modules["attn_decode"].attn_decode( - q, k, v, mask=mask, causal_offset=causal_offset, layout=1 + q, k, v, mask=mask, causal_offset=causal_offset, layout=TensorLayout.BLHD ) @@ -55,7 +68,7 @@ def attn_prefill( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - mask: torch.Tensor | None = None, + mask: Optional[torch.Tensor] = None, is_causal: bool = False, ) -> torch.Tensor: """GQA prefill attention (q_len > 1). @@ -73,7 +86,7 @@ def attn_prefill( _check_available("attn_prefill") causal_offset = (k.size(1) - q.size(1)) if is_causal else -1 return _modules["attn_prefill"].attn_prefill( - q, k, v, mask=mask, causal_offset=causal_offset, layout=1 + q, k, v, mask=mask, causal_offset=causal_offset, layout=TensorLayout.BLHD ) @@ -84,7 +97,7 @@ def attn_paged_decode( v_cache: torch.Tensor, page_size: int, kv_len: int, - mask: torch.Tensor | None = None, + mask: Optional[torch.Tensor] = None, is_causal: bool = False, ) -> torch.Tensor: """Paged GQA decode attention (q_len == 1, direct page-table access). @@ -113,5 +126,5 @@ def attn_paged_decode( kv_len, mask=mask, causal_offset=causal_offset, - layout=1, + layout=TensorLayout.BLHD, ) diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h index 9904ba7..20e57ec 100644 --- a/csrc/kernels/attn_common.h +++ b/csrc/kernels/attn_common.h @@ -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 struct AttentionParams { diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attn_decode.cu index a9fec7a..a303a65 100644 --- a/csrc/kernels/attn_decode.cu +++ b/csrc/kernels/attn_decode.cu @@ -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)"); } diff --git a/csrc/kernels/attn_dispatchers.cuh b/csrc/kernels/attn_dispatchers.cuh index 68060e5..5a5c25f 100644 --- a/csrc/kernels/attn_dispatchers.cuh +++ b/csrc/kernels/attn_dispatchers.cuh @@ -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 ; 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(__VA_ARGS__); \ + else FN(__VA_ARGS__); \ + } else { \ + if (has_mask) FN(__VA_ARGS__); \ + else FN(__VA_ARGS__); \ + } \ + } while (0) + // ====================================================================== // Prefill // ====================================================================== @@ -56,21 +88,9 @@ static inline void dispatch_prefill(AttentionParams& p) { bool has_mask = (p.use_mask && p.mask); #ifndef ASTRAI_NO_MMA - if (is_causal) { - if (has_mask) launch_prefill_mma(p); - else launch_prefill_mma(p); - } else { - if (has_mask) launch_prefill_mma(p); - else launch_prefill_mma(p); - } + DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_mma, HEAD_DIM, p); #else - if (is_causal) { - if (has_mask) launch_prefill_scalar(p); - else launch_prefill_scalar(p); - } else { - if (has_mask) launch_prefill_scalar(p); - else launch_prefill_scalar(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& p) { int group_size = p.q_head / p.kv_head; #ifndef ASTRAI_NO_MMA - if (is_causal) { - if (has_mask) launch_decode_mma(p, group_size); - else launch_decode_mma(p, group_size); - } else { - if (has_mask) launch_decode_mma(p, group_size); - else launch_decode_mma(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(p, group_size); - else launch_decode_scalar(p, group_size); - } else { - if (has_mask) launch_decode_scalar(p, group_size); - else launch_decode_scalar(p, group_size); - } + DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_scalar, HEAD_DIM, p, group_size); #endif attn_decode_combine_kernel<<>>(p); @@ -174,21 +182,9 @@ static inline void dispatch_paged_decode(PagedAttentionParams& p) { int group_size = p.q_head / p.kv_head; #ifndef ASTRAI_NO_MMA - if (is_causal) { - if (has_mask) launch_paged_decode_mma(p, group_size); - else launch_paged_decode_mma(p, group_size); - } else { - if (has_mask) launch_paged_decode_mma(p, group_size); - else launch_paged_decode_mma(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(p, group_size); - else launch_paged_decode_scalar(p, group_size); - } else { - if (has_mask) launch_paged_decode_scalar(p, group_size); - else launch_paged_decode_scalar(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); diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 50cb34e..495be47 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -37,7 +37,7 @@ inline void alloc_split_partials(P& p) { // ---- Shared Q-dims + strides extraction ---- template 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); diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu index b6f8d67..1398f95 100644 --- a/csrc/kernels/attn_paged_decode.cu +++ b/csrc/kernels/attn_paged_decode.cu @@ -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."); } diff --git a/csrc/kernels/attn_prefill.cu b/csrc/kernels/attn_prefill.cu index d4f1c0a..a0a0cb8 100644 --- a/csrc/kernels/attn_prefill.cu +++ b/csrc/kernels/attn_prefill.cu @@ -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)"); } diff --git a/csrc/kernels/rotary_emb.cu b/csrc/kernels/rotary_emb.cu index c9bb356..3e717c9 100644 --- a/csrc/kernels/rotary_emb.cu +++ b/csrc/kernels/rotary_emb.cu @@ -1,4 +1,6 @@ #include +#include +#include #include __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; }