From 8a8550184f1153ce45834a783678ef28095a7d31 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 11 Jul 2026 11:03:14 +0800 Subject: [PATCH] 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 | 1 + csrc/kernels/attn_common.cuh | 30 ----------------------- csrc/kernels/attn_common.h | 26 ++++++++++++++++++++ csrc/kernels/attn_decode.cu | 12 ++++----- csrc/kernels/attn_decode_split_kv.cuh | 9 ++++--- csrc/kernels/attn_decode_split_kv_mma.cuh | 8 ++++-- csrc/kernels/attn_entry_utils.cuh | 11 +++++---- csrc/kernels/attn_mma_utils.cuh | 3 +++ csrc/kernels/attn_prefill.cu | 4 +-- csrc/kernels/attn_prefill_split_q.cuh | 8 ++++-- csrc/kernels/attn_prefill_split_q_mma.cuh | 8 ++++-- csrc/tests/attn_decode_test.cu | 16 ++++++------ csrc/tests/attn_prefill_test.cu | 8 +++--- 13 files changed, 80 insertions(+), 64 deletions(-) delete mode 100644 csrc/kernels/attn_common.cuh create mode 100644 csrc/kernels/attn_common.h diff --git a/.gitignore b/.gitignore index 37ca78a..024aee7 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ !csrc/**/*.py !csrc/**/*.cu +!csrc/**/*.h !csrc/**/*.cuh !scripts/**/*.sh diff --git a/csrc/kernels/attn_common.cuh b/csrc/kernels/attn_common.cuh deleted file mode 100644 index a0bdcc8..0000000 --- a/csrc/kernels/attn_common.cuh +++ /dev/null @@ -1,30 +0,0 @@ -#pragma once -#include -#include -#include -#include - -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; -}; diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h new file mode 100644 index 0000000..ea02e69 --- /dev/null +++ b/csrc/kernels/attn_common.h @@ -0,0 +1,26 @@ +#pragma once + + +template +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; +}; diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attn_decode.cu index b01c29e..f37cc81 100644 --- a/csrc/kernels/attn_decode.cu +++ b/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& 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& 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& p) { int G = p.q_head / p.kv_head; return !p.use_mask && G > 1 && G <= 16; } template -static void launch_mma_decode(AttentionParams& p) { +static void launch_mma_decode(AttentionParams& 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 -static void dispatch_decode(AttentionParams& p) { +static void dispatch_decode(AttentionParams& p) { #ifndef ASTRAI_NO_MMA if (decode_use_mma(p)) { launch_mma_decode(p); @@ -80,7 +80,7 @@ torch::Tensor attn_decode( int64_t causal_offset = 0, c10::optional scale = c10::nullopt ) { - AttentionParams p; + AttentionParams 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"); diff --git a/csrc/kernels/attn_decode_split_kv.cuh b/csrc/kernels/attn_decode_split_kv.cuh index 1a610b0..e790313 100644 --- a/csrc/kernels/attn_decode_split_kv.cuh +++ b/csrc/kernels/attn_decode_split_kv.cuh @@ -1,6 +1,9 @@ #pragma once -#include "attn_common.cuh" +#include +#include +#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 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 p) { int bh = blockIdx.x; int d = threadIdx.x; if (d >= p.head_dim) return; diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index 5fd210d..8202669 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -1,7 +1,11 @@ #pragma once -#include "attn_common.cuh" +#include +#include +#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 -__global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { +__global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { constexpr int BR = 16; constexpr int KD = HEAD_DIM / 16; constexpr int NC8 = BC / 8; diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 7fe3e46..7e83730 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -1,7 +1,8 @@ #pragma once #include -#include "attn_common.cuh" +#include "attn_common.h" +template 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 scale, - AttentionParams& p + AttentionParams& 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); diff --git a/csrc/kernels/attn_mma_utils.cuh b/csrc/kernels/attn_mma_utils.cuh index dae9310..1102c21 100644 --- a/csrc/kernels/attn_mma_utils.cuh +++ b/csrc/kernels/attn_mma_utils.cuh @@ -1,4 +1,7 @@ #pragma once +#include +#include +#include // Shared MMA utilities for tensor-core GQA kernels. // mma.sync.m16n8k16 PTX wrappers, ldmatrix helpers, and bf16 packing. diff --git a/csrc/kernels/attn_prefill.cu b/csrc/kernels/attn_prefill.cu index 1ea31f2..3b0d745 100644 --- a/csrc/kernels/attn_prefill.cu +++ b/csrc/kernels/attn_prefill.cu @@ -6,7 +6,7 @@ #endif template -static void dispatch_prefill(AttentionParams& p) { +static void dispatch_prefill(AttentionParams& 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 scale = c10::nullopt ) { - AttentionParams p; + AttentionParams 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"); diff --git a/csrc/kernels/attn_prefill_split_q.cuh b/csrc/kernels/attn_prefill_split_q.cuh index 5df9b43..d84e08b 100644 --- a/csrc/kernels/attn_prefill_split_q.cuh +++ b/csrc/kernels/attn_prefill_split_q.cuh @@ -1,5 +1,9 @@ #pragma once -#include "attn_common.cuh" +#include +#include +#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 -__global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { +__global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { constexpr int DPT = HEAD_DIM / G; int q_tile = blockIdx.x; diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index 783abac..617b6a1 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -1,7 +1,11 @@ #pragma once -#include "attn_common.cuh" +#include +#include +#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 __global__ __launch_bounds__(WARPS * 32, MIN_BLOCKS) -void attn_prefill_split_q_mma_kernel(AttentionParams p) { +void attn_prefill_split_q_mma_kernel(AttentionParams 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) diff --git a/csrc/tests/attn_decode_test.cu b/csrc/tests/attn_decode_test.cu index fdd914a..38a3333 100644 --- a/csrc/tests/attn_decode_test.cu +++ b/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& 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& p) { int G = p.q_head / p.kv_head; return !p.use_mask && G > 1 && G <= 16; } template -static void launch_mma_decode(AttentionParams& p, DecodeScratch& sc) { +static void launch_mma_decode(AttentionParams& 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& 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 -static void dispatch_decode_t(AttentionParams& p, DecodeScratch& sc) { +static void dispatch_decode_t(AttentionParams& p, DecodeScratch& sc) { #ifndef ASTRAI_NO_MMA if (decode_use_mma(p)) { launch_mma_decode(p, sc); return; } #endif launch_scalar_decode(p, sc); } -static void dispatch_decode(AttentionParams& p, DecodeScratch& sc) { +static void dispatch_decode(AttentionParams& 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 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 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); diff --git a/csrc/tests/attn_prefill_test.cu b/csrc/tests/attn_prefill_test.cu index 82e40b3..b80df36 100644 --- a/csrc/tests/attn_prefill_test.cu +++ b/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 -static void launch_prefill(AttentionParams& p) { +static void launch_prefill(AttentionParams& 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& 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 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 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);