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

@@ -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",

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,
     )

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<typename T, typename AT = float>
 struct AttentionParams {

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)");
 }

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;
+    int sm_count = get_sm_count();
-    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
     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) {
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_mma, HEAD_DIM, p);
-        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);
-    }
 #else
-    if (is_causal) {
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_scalar, HEAD_DIM, p);
-        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);
-    }
 #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) {
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_mma, HEAD_DIM, p, group_size);
-        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);
-    }
 #else
-    if (is_causal) {
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_scalar, HEAD_DIM, p, group_size);
-        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);
-    }
 #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) {
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, group_size);
-        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);
-    }
 #else
-    if (is_causal) {
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size);
-        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);
-    }
 #endif
 
     paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);

csrc/kernels/attn_entry_utils.cuh

@@ -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

@@ -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

@@ -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

@@ -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;
 }