refactor: template AttentionParams, rename .cuh to .h

- Convert AttentionParams to a template struct supporting arbitrary types
- Rename attn_common.cuh -> attn_common.h (no CUDA-specific code remains)
- Include standard headers explicitly in each .cuh instead of via attn_common.cuh
- Allow .h files in csrc/ via .gitignore

.gitignore

@@ -11,6 +11,7 @@
 !csrc/**/*.py
 
 !csrc/**/*.cu
+!csrc/**/*.h
 !csrc/**/*.cuh
 
 !scripts/**/*.sh

csrc/kernels/attn_common.cuh

@@ -1,30 +0,0 @@
-#pragma once
-#include <cuda_bf16.h>
-#include <cuda_runtime.h>
-#include <cfloat>
-#include <algorithm>
-
-using bf16 = __nv_bfloat16;
-using std::min;
-
-struct AttentionParams {
-    int batch;
-    int q_head;
-    int kv_head;
-    int q_len;
-    int kv_len;
-    int head_dim;
-    int use_mask;
-    int is_causal;
-    int causal_offset;
-    float scale;
-    const bf16* __restrict__ q;
-    const bf16* __restrict__ k;
-    const bf16* __restrict__ v;
-    const bool* __restrict__ mask;
-    
-    bf16* __restrict__ o;
-    float* __restrict__ o_part;
-    float* __restrict__ ml_part;
-    int num_splits;
-};

csrc/kernels/attn_common.h

@@ -0,0 +1,26 @@
+#pragma once
+
+
+template<typename T, typename AT = float>
+struct AttentionParams {
+    int batch;
+    int q_head;
+    int kv_head;
+    int q_len;
+    int kv_len;
+    int head_dim;
+    int use_mask;
+    int is_causal;
+    int causal_offset;
+    int num_splits;
+    float scale;
+    
+    const T* __restrict__ q;
+    const T* __restrict__ k;
+    const T* __restrict__ v;
+    const bool* __restrict__ mask;
+    
+    T* __restrict__ o;
+    AT* __restrict__ o_part;
+    AT* __restrict__ ml_part;
+};

csrc/kernels/attn_decode.cu

@@ -9,7 +9,7 @@
 // tiny (e.g. 16 blocks) and leaves most SMs idle. Pick the smallest split count
 // that fills the device (~2 blocks/SM), capped by the tile count, min work per
 // split (at least 8 tiles), and 32.
-static int decode_num_splits(const AttentionParams& p, int tiles_total) {
+static int decode_num_splits(const AttentionParams<bf16>& p, int tiles_total) {
     int sm_count = 0;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
     int base_blocks = p.kv_head * p.batch;
@@ -20,7 +20,7 @@ static int decode_num_splits(const AttentionParams& p, int tiles_total) {
 }
 
 // Scalar fallback: one warp per query head, split-KV across grid.z.
-static void launch_scalar_decode(AttentionParams& p) {
+static void launch_scalar_decode(AttentionParams<bf16>& p) {
     int group_size = p.q_head / p.kv_head;
     int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
     p.num_splits = decode_num_splits(p, chunks_total);
@@ -38,13 +38,13 @@ static void launch_scalar_decode(AttentionParams& p) {
 
 #ifndef ASTRAI_NO_MMA
 // Tensor-core head-packing requires 1 < G <= 16 (the MMA M dim) and no mask.
-static bool decode_use_mma(const AttentionParams& p) {
+static bool decode_use_mma(const AttentionParams<bf16>& p) {
     int G = p.q_head / p.kv_head;
     return !p.use_mask && G > 1 && G <= 16;
 }
 
 template <int HEAD_DIM, int BC>
-static void launch_mma_decode(AttentionParams& p) {
+static void launch_mma_decode(AttentionParams<bf16>& p) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
     p.num_splits = decode_num_splits(p, tiles_total);
 
@@ -61,7 +61,7 @@ static void launch_mma_decode(AttentionParams& p) {
 #endif
 
 template <int HEAD_DIM>
-static void dispatch_decode(AttentionParams& p) {
+static void dispatch_decode(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
     if (decode_use_mma(p)) {
         launch_mma_decode<HEAD_DIM, 32>(p);
@@ -80,7 +80,7 @@ torch::Tensor attn_decode(
     int64_t causal_offset = 0,
     c10::optional<double> scale = c10::nullopt
 ) {
-    AttentionParams p;
+    AttentionParams<bf16> p;
     attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
     TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
     TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");

csrc/kernels/attn_decode_split_kv.cuh

@@ -1,6 +1,9 @@
 #pragma once
-#include "attn_common.cuh"
+#include <cuda_bf16.h>
+#include <float.h>
+#include "attn_common.h"
 
+using bf16 = __nv_bfloat16;
 constexpr int DC_CHUNK = 64;
 
 __device__ inline float warp_reduce_sum(float val) {
@@ -9,7 +12,7 @@ __device__ inline float warp_reduce_sum(float val) {
     return val;
 }
 
-__global__ void attn_decode_split_kv_kernel(AttentionParams p) {
+__global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
     int batch = blockIdx.x / p.kv_head;
     int kv_head = blockIdx.x % p.kv_head;
     int split = blockIdx.z;
@@ -86,7 +89,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams p) {
 // Reduce split-K partials into the final bf16 output. One block per (batch,
 // q_head); each thread folds across all splits with a single-pass
 // online-rescale reduction (expf + FMA counts halved vs 3-pass original).
-__global__ void attn_decode_combine_kernel(AttentionParams p) {
+__global__ void attn_decode_combine_kernel(AttentionParams<bf16> p) {
     int bh = blockIdx.x;
     int d = threadIdx.x;
     if (d >= p.head_dim) return;

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -1,7 +1,11 @@
 #pragma once
-#include "attn_common.cuh"
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
 #include "attn_mma_utils.cuh"
 
+using bf16 = __nv_bfloat16;
+
 // Split-K (FlashDecoding) tensor-core decode via GQA head-packing.
 //
 // Decode has q_len == 1, so S = q @ K^T is a GEMV per head — no tensor-core work
@@ -26,7 +30,7 @@
 //   - pre-scaled Q: Q scaled during load, softmax skips per-tile multiply
 //   - single-buffer: keeps smem small for high occupancy
 template <int HEAD_DIM, int BC>
-__global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) {
+__global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
     constexpr int BR = 16;
     constexpr int KD = HEAD_DIM / 16;
     constexpr int NC8 = BC / 8;

csrc/kernels/attn_entry_utils.cuh

@@ -1,7 +1,8 @@
 #pragma once
 #include <torch/extension.h>
-#include "attn_common.cuh"
+#include "attn_common.h"
 
+template<typename T>
 inline void attn_pack_params(
     torch::Tensor q,
     torch::Tensor k,
@@ -10,7 +11,7 @@ inline void attn_pack_params(
     bool is_causal,
     int64_t causal_offset,
     c10::optional<double> scale,
-    AttentionParams& p
+    AttentionParams<T>& p
 ) {
     TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda());
     TORCH_CHECK(q.dtype() == torch::kBFloat16);
@@ -27,9 +28,9 @@ inline void attn_pack_params(
     p.is_causal = is_causal ? 1 : 0;
     p.causal_offset = (int)causal_offset;
     p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim);
-    p.q = (const bf16*)q.data_ptr();
-    p.k = (const bf16*)k.data_ptr();
-    p.v = (const bf16*)v.data_ptr();
+    p.q = (const T*)q.data_ptr();
+    p.k = (const T*)k.data_ptr();
+    p.v = (const T*)v.data_ptr();
     if (p.use_mask) {
         TORCH_CHECK(mask.value().dtype() == torch::kBool);
         TORCH_CHECK(mask.value().dim() == 2);

csrc/kernels/attn_mma_utils.cuh

@@ -1,4 +1,7 @@
 #pragma once
+#include <cfloat>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
 
 // Shared MMA utilities for tensor-core GQA kernels.
 // mma.sync.m16n8k16 PTX wrappers, ldmatrix helpers, and bf16 packing.

csrc/kernels/attn_prefill.cu

@@ -6,7 +6,7 @@
 #endif
 
 template <int HEAD_DIM>
-static void dispatch_prefill(AttentionParams& p) {
+static void dispatch_prefill(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
     constexpr int WARPS = 4, BR = 16;
     // KV tile: bigger tiles amortize the per-tile cp.async wait + barrier +
@@ -41,7 +41,7 @@ torch::Tensor attn_prefill(
     int64_t causal_offset = 0,
     c10::optional<double> scale = c10::nullopt
 ) {
-    AttentionParams p;
+    AttentionParams<bf16> p;
     attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
     TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16");
 

csrc/kernels/attn_prefill_split_q.cuh

@@ -1,5 +1,9 @@
 #pragma once
-#include "attn_common.cuh"
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
+
+using bf16 = __nv_bfloat16;
 
 // v9: group-split register blocking. G threads cooperate on one query row,
 // each owning HEAD_DIM/G dims of qreg[]/acc[]. Small per-thread footprint keeps
@@ -31,7 +35,7 @@ __device__ __forceinline__ void ld8(const bf16* p, float* o) {
 }
 
 template <int HEAD_DIM, int G, int ROWS, int P_BC>
-__global__ void attn_prefill_split_q_kernel_t(AttentionParams p) {
+__global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
     constexpr int DPT = HEAD_DIM / G;
 
     int q_tile = blockIdx.x;

csrc/kernels/attn_prefill_split_q_mma.cuh

@@ -1,7 +1,11 @@
 #pragma once
-#include "attn_common.cuh"
+#include <cfloat>
+#include <cuda_bf16.h>
+#include "attn_common.h"
 #include "attn_mma_utils.cuh"
 
+using bf16 = __nv_bfloat16;
+
 // Tensor-core prefill flash attention (raw mma.sync PTX).
 // One warp owns BR=16 query rows. S = Q@K^T and O = P@V run on bf16 tensor
 // cores via mma.sync.m16n8k16 (f32 accumulate). Q fragments are loaded once
@@ -38,7 +42,7 @@
 
 template <int HEAD_DIM, int WARPS, int BC, int MIN_BLOCKS>
 __global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS)
-void attn_prefill_split_q_mma_kernel(AttentionParams p) {
+void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
     constexpr int BR = 16;
     constexpr int KD = HEAD_DIM / 16;  // Q/K k-tiles
     constexpr int NC8 = BC / 8;        // S n-tiles (N=8 each)

csrc/tests/attn_decode_test.cu

@@ -28,7 +28,7 @@ struct DecodeScratch {
     float* ml_part = nullptr;
 };
 
-static int decode_num_splits(const AttentionParams& p, int tiles_total) {
+static int decode_num_splits(const AttentionParams<bf16>& p, int tiles_total) {
     int sm_count = 0;
     cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
     int base_blocks = p.kv_head * p.batch;
@@ -41,13 +41,13 @@ static int decode_num_splits(const AttentionParams& p, int tiles_total) {
 // Launch the production decode path (tensor-core head-packing MMA on sm_80+,
 // scalar fallback otherwise), mirroring dispatch_decode() in attn_decode.cu.
 #ifndef ASTRAI_NO_MMA
-static bool decode_use_mma(const AttentionParams& p) {
+static bool decode_use_mma(const AttentionParams<bf16>& p) {
     int G = p.q_head / p.kv_head;
     return !p.use_mask && G > 1 && G <= 16;
 }
 
 template <int HEAD_DIM, int BC>
-static void launch_mma_decode(AttentionParams& p, DecodeScratch& sc) {
+static void launch_mma_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
     p.num_splits = decode_num_splits(p, tiles_total);
     p.o_part = sc.o_part;
@@ -59,7 +59,7 @@ static void launch_mma_decode(AttentionParams& p, DecodeScratch& sc) {
 }
 #endif
 
-static void launch_scalar_decode(AttentionParams& p, DecodeScratch& sc) {
+static void launch_scalar_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
     int gs = p.q_head / p.kv_head;
     int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
     p.num_splits = decode_num_splits(p, chunks_total);
@@ -72,14 +72,14 @@ static void launch_scalar_decode(AttentionParams& p, DecodeScratch& sc) {
 }
 
 template <int HEAD_DIM>
-static void dispatch_decode_t(AttentionParams& p, DecodeScratch& sc) {
+static void dispatch_decode_t(AttentionParams<bf16>& p, DecodeScratch& sc) {
 #ifndef ASTRAI_NO_MMA
     if (decode_use_mma(p)) { launch_mma_decode<HEAD_DIM, 32>(p, sc); return; }
 #endif
     launch_scalar_decode(p, sc);
 }
 
-static void dispatch_decode(AttentionParams& p, DecodeScratch& sc) {
+static void dispatch_decode(AttentionParams<bf16>& p, DecodeScratch& sc) {
     switch (p.head_dim) {
         case 32:  dispatch_decode_t<32>(p, sc);  break;
         case 64:  dispatch_decode_t<64>(p, sc);  break;
@@ -164,7 +164,7 @@ static void bench() {
         for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
 
-        AttentionParams p;
+        AttentionParams<bf16> p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
         p.use_mask=0; p.is_causal=0; p.causal_offset=0;
         p.scale=1.0f/sqrtf((float)D);
@@ -240,7 +240,7 @@ int main() {
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
         cudaMemcpy(dMask,hMask,B*sl,cudaMemcpyHostToDevice);
 
-        AttentionParams p;
+        AttentionParams<bf16> p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
         p.use_mask=0; p.is_causal=0; p.causal_offset=0;
         p.scale=1.0f/sqrtf((float)D);

csrc/tests/attn_prefill_test.cu

@@ -23,7 +23,7 @@ static double now_ms() {
 // Launch the production prefill path (tensor-core MMA on sm_80+, else the
 // scalar fallback), mirroring dispatch_prefill() in attn_prefill.cu.
 template <int HEAD_DIM>
-static void launch_prefill(AttentionParams& p) {
+static void launch_prefill(AttentionParams<bf16>& p) {
 #ifndef ASTRAI_NO_MMA
     constexpr int WARPS = 4, BR = 16;
     constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
@@ -40,7 +40,7 @@ static void launch_prefill(AttentionParams& p) {
 #endif
 }
 
-static void dispatch_prefill(AttentionParams& p) {
+static void dispatch_prefill(AttentionParams<bf16>& p) {
     switch (p.head_dim) {
         case 64:  launch_prefill<64>(p);  break;
         case 128: launch_prefill<128>(p); break;
@@ -123,7 +123,7 @@ static void bench() {
         for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(randf());
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
 
-        AttentionParams p;
+        AttentionParams<bf16> p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
         p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
         p.scale=1.0f/sqrtf((float)D);
@@ -191,7 +191,7 @@ int main() {
         for (size_t i=0;i<nKV;i++) tmp[i]=f2bf(hV[i]);
         cudaMemcpy(dV,tmp,nKV*2,cudaMemcpyHostToDevice);
 
-        AttentionParams p;
+        AttentionParams<bf16> p;
         p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
         p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
         p.scale=1.0f/sqrtf((float)D);