perf: launch CUDA kernels on torch's current stream

- Thread a cudaStream_t through attn dispatchers onto torch's current stream
- Scope the device guard to the entry function so kernels run on tensor device
- DISPATCH_HEAD_DIM now forwards varargs so stream reaches each dispatch
- Parallelize CPU reference kernels with OpenMP (paged test 31s -> 7s)
- Merge decode/prefill standalone tests into attn_test.cu with correctness tables
- Drop bench error column (CPU ref too slow at large sizes)
- Update cuda_kernels.md for the merged test layout

csrc/kernels/attn_decode.cu

@@ -10,6 +10,9 @@ torch::Tensor attn_decode(
     double scale,
     int64_t layout
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     AttentionParams<bf16> p;
     attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
     TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
@@ -20,7 +23,7 @@ torch::Tensor attn_decode(
     p.o = (bf16*)O_view.data_ptr();
 
     alloc_split_partials(p);
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/attn_dispatchers.cuh

@@ -66,33 +66,33 @@ inline int compute_num_splits(int base_blocks, int tiles_total,
 
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_prefill_mma(AttentionParams<bf16>& p) {
+static inline void launch_prefill_mma(AttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int WARPS = 4;
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
     using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
     dim3 grid((p.q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS), p.q_head, p.batch);
     dim3 block(Traits::NUM_THREADS);
-    attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block>>>(p);
+    attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_prefill_scalar(AttentionParams<bf16>& p) {
+static inline void launch_prefill_scalar(AttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int G = 8, ROWS = 32, P_BC = 32;
     dim3 grid((p.q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
     dim3 block(G, ROWS);
-    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask><<<grid, block>>>(p);
+    attn_prefill_split_q_kernel_t<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask><<<grid, block, 0, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_prefill(AttentionParams<bf16>& p) {
+static inline void dispatch_prefill(AttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
 
 #ifndef ASTRAI_NO_MMA
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_mma, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_mma, HEAD_DIM, p, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_scalar, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_prefill_scalar, HEAD_DIM, p, stream);
 #endif
 }
 
@@ -106,7 +106,7 @@ static inline void dispatch_prefill(AttentionParams<bf16>& p) {
 // enabling STAGES=2 (double-buffer) within the 32KB smem budget — eliminates
 // the 176-byte spill that STAGES=1+BC=32 suffered.
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_decode_mma(AttentionParams<bf16>& p, int group_size) {
+static inline void launch_decode_mma(AttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
     int G = p.q_head / p.kv_head;
     constexpr int MAX_G = 16;
     int num_passes = (G + MAX_G - 1) / MAX_G;
@@ -116,34 +116,34 @@ static inline void launch_decode_mma(AttentionParams<bf16>& p, int group_size) {
     constexpr int STAGES = 2;
     using Traits = KernelTraits<HEAD_DIM, BC, 1, STAGES>;
     dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
-    attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask><<<grid, 32>>>(p);
+    attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask><<<grid, 32, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_decode_scalar(AttentionParams<bf16>& p, int group_size) {
+static inline void launch_decode_scalar(AttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
     int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
     p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
     size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
     int g = min(group_size, 32);  // cap at 32 to respect 1024-thread limit
     dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
     dim3 block(32, g);
-    attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
+    attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_decode(AttentionParams<bf16>& p) {
+static inline void dispatch_decode(AttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
     int group_size = p.q_head / p.kv_head;
 
 #ifndef ASTRAI_NO_MMA
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_mma, HEAD_DIM, p, group_size);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_mma, HEAD_DIM, p, group_size, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_scalar, HEAD_DIM, p, group_size);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_decode_scalar, HEAD_DIM, p, group_size, stream);
 #endif
 
-    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+    attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }
 
 // ======================================================================
@@ -152,7 +152,7 @@ static inline void dispatch_decode(AttentionParams<bf16>& p) {
 
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int) {
+static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     int G = p.q_head / p.kv_head;
     constexpr int MAX_G = 16;
     constexpr int BC = 16;
@@ -162,34 +162,34 @@ static inline void launch_paged_decode_mma(PagedAttentionParams<bf16>& p, int) {
     constexpr int STAGES = 2;
     using Traits = KernelTraits<HEAD_DIM, BC, 1, STAGES>;
     dim3 grid(p.kv_head * num_passes, p.batch, p.num_splits);
-    paged_attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask> <<<grid, 32>>>(p);
+    paged_attn_decode_split_kv_mma_kernel<Traits, IsCausal, HasMask> <<<grid, 32, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_decode_scalar(PagedAttentionParams<bf16>& p, int group_size) {
+static inline void launch_paged_decode_scalar(PagedAttentionParams<bf16>& p, int group_size, cudaStream_t stream) {
     int chunks_total = (p.max_seq_len + PDC_CHUNK - 1) / PDC_CHUNK;
     p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
     size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
     int g = min(group_size, 32);
     dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
     dim3 block(32, g);
-    paged_attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem>>>(p);
+    paged_attn_decode_split_kv_kernel<HEAD_DIM, IsCausal, HasMask><<<grid, block, smem, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
+static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
     int group_size = p.q_head / p.kv_head;
 
 #ifndef ASTRAI_NO_MMA
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, 0);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_mma, HEAD_DIM, p, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_decode_scalar, HEAD_DIM, p, group_size, stream);
 #endif
 
-    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
+    paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim, 0, stream>>>(p);
 }
 
 // ======================================================================
@@ -198,35 +198,35 @@ static inline void dispatch_paged_decode(PagedAttentionParams<bf16>& p) {
 
 #ifndef ASTRAI_NO_MMA
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_prefill_mma(PagedAttentionParams<bf16>& p) {
+static inline void launch_paged_prefill_mma(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int WARPS = 4;
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
     using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
     int max_q_tiles = (p.max_q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS);
     dim3 grid(max_q_tiles, p.q_head, p.batch);
     dim3 block(Traits::NUM_THREADS);
-    paged_attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block>>>(p);
+    paged_attn_prefill_split_q_mma_kernel<Traits, IsCausal, HasMask><<<grid, block, 0, stream>>>(p);
 }
 #endif
 
 template <int HEAD_DIM, bool IsCausal, bool HasMask>
-static inline void launch_paged_prefill_scalar(PagedAttentionParams<bf16>& p) {
+static inline void launch_paged_prefill_scalar(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     constexpr int G = 8, ROWS = 32, P_BC = 32;
     int max_q_tiles = (p.max_q_len + ROWS - 1) / ROWS;
     dim3 grid(max_q_tiles, p.q_head, p.batch);
     dim3 block(G, ROWS);
     paged_attn_prefill_split_q_kernel<HEAD_DIM, G, ROWS, P_BC, IsCausal, HasMask>
-        <<<grid, block>>>(p);
+        <<<grid, block, 0, stream>>>(p);
 }
 
 template <int HEAD_DIM>
-static inline void dispatch_paged_prefill(PagedAttentionParams<bf16>& p) {
+static inline void dispatch_paged_prefill(PagedAttentionParams<bf16>& p, cudaStream_t stream) {
     bool is_causal = (p.causal_offset >= 0);
     bool has_mask = (p.use_mask && p.mask);
 
 #ifndef ASTRAI_NO_MMA
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_mma, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_mma, HEAD_DIM, p, stream);
 #else
-    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_scalar, HEAD_DIM, p);
+    DISPATCH_CAUSAL_MASK(is_causal, has_mask, launch_paged_prefill_scalar, HEAD_DIM, p, stream);
 #endif
 }

csrc/kernels/attn_entry_utils.cuh

@@ -8,14 +8,14 @@
 using bf16 = __nv_bfloat16;
 
 // 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) \
+// Usage: DISPATCH_HEAD_DIM(hd, fn, args...)
+//   Expands to: fn<32>(args...); fn<64>(args...); etc.
+#define DISPATCH_HEAD_DIM(hd, fn, ...) \
     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; \
+        case 32:  fn<32>(__VA_ARGS__); break; \
+        case 64:  fn<64>(__VA_ARGS__); break; \
+        case 128: fn<128>(__VA_ARGS__); break; \
+        case 256: fn<256>(__VA_ARGS__); break; \
         default: \
             TORCH_CHECK(false, "unsupported head_dim ", hd, \
                          " (supported: 32, 64, 128, 256)"); \

csrc/kernels/attn_paged_decode.cu

@@ -13,6 +13,9 @@ torch::Tensor attn_paged_decode(
     int64_t causal_offset,
     double scale
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     PagedAttentionParams<bf16> p;
     attn_pack_paged_decode_params(q, k_cache, v_cache,
                                    req_to_token, req_pool_indices, kv_indptr,
@@ -22,7 +25,7 @@ torch::Tensor attn_paged_decode(
     p.o = (bf16*)O.data_ptr();
 
     alloc_split_partials(p);
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/attn_paged_prefill.cu

@@ -14,6 +14,9 @@ torch::Tensor attn_paged_prefill(
     int64_t causal_offset,
     double scale
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     PagedAttentionParams<bf16> p;
     attn_pack_paged_prefill_params(q, k_cache, v_cache,
                                     req_to_token, req_pool_indices,
@@ -23,7 +26,7 @@ torch::Tensor attn_paged_prefill(
     auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
     p.o = (bf16*)O.data_ptr();
 
-    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_prefill, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/attn_prefill.cu

@@ -10,6 +10,9 @@ torch::Tensor attn_prefill(
     double scale,
     int64_t layout
 ) {
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
     AttentionParams<bf16> p;
     attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
     TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
@@ -18,7 +21,7 @@ torch::Tensor attn_prefill(
     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);
+    DISPATCH_HEAD_DIM(p.head_dim, dispatch_prefill, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;
 }

csrc/kernels/rotary_emb.cu

@@ -49,6 +49,7 @@ torch::Tensor rotary_emb(
     torch::Tensor freqs_cis
 ) {
     const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
+    auto stream = at::cuda::getCurrentCUDAStream();
 
     TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
     TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
@@ -77,7 +78,7 @@ torch::Tensor rotary_emb(
     int block = 256;
     int grid = std::min((total + block - 1) / block, 1024);
 
-    rotary_emb_kernel<<<grid, block>>>(
+    rotary_emb_kernel<<<grid, block, 0, stream>>>(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
         freqs_cis.data_ptr<float>(),
         reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),