ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: unify paged and contiguous attention kernels via KVSource policy

Browse Source 
- merge AttentionParams and PagedAttentionParams into one struct
- add attn_kv_source.cuh with ContigKV/PagedKV addressing policies
- template prefill/decode kernels (MMA + scalar) on the KV policy, deleting the four duplicated attn_paged_*.cuh variants
- template dispatcher launchers on KV; single combine kernel
- verify: all correctness tests pass and SASS matches baseline (no perf regression)
This commit is contained in:
ViperEkura
2026-08-05 14:06:13 +08:00
parent 6dffb0305a
commit 2667b8116d
17 changed files with 415 additions and 890 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/tests/attn_paged_test.cu
+6 -6
View File
@@ -212,7 +212,7 @@ static int run_decode_test(int B, int Hq, int Hkv, int max_seq,
B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
// Kernel launch
PagedAttentionParams<bf16> p;
AttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = B;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
@@ -347,7 +347,7 @@ static int run_decode_mask_test(int B, int Hq, int Hkv, int max_seq,
h_mask, max_sl,
B, Hq, Hkv, HEAD_DIM, max_ctx, h_o_ref);
PagedAttentionParams<bf16> p;
AttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = B;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
@@ -480,7 +480,7 @@ static int run_prefill_test(int B, int Hq, int Hkv,
B, Hq, Hkv, HEAD_DIM, max_ctx, causal, h_o_ref);
// Kernel launch
PagedAttentionParams<bf16> p;
AttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = total_q;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
@@ -617,7 +617,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) {
h_mask, q_len, q_len,
B, Hq, Hkv, HEAD_DIM, max_ctx, 0, h_o_ref);
PagedAttentionParams<bf16> p;
AttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = total_q;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
@@ -707,7 +707,7 @@ static void bench_decode(int B, int Hq, int Hkv, int seq_len) {
for (int b = 0; b < B; b++) h_kvi[b + 1] = h_kvi[b] + seq_len;
cudaMemcpy(d_kvi, h_kvi, sz_kvi, cudaMemcpyHostToDevice);
PagedAttentionParams<bf16> p;
AttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = B;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
@@ -784,7 +784,7 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau
for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_len;
cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice);
PagedAttentionParams<bf16> p;
AttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv;
p.head_dim = HEAD_DIM; p.total_q = total_q;
p.q_stride_l = Hq * HEAD_DIM; p.q_stride_h = HEAD_DIM; p.q_stride_d = 1;
csrc/tests/test_utils.cuh
+1 -1
View File
@@ -120,7 +120,7 @@ inline void set_default_strides(P& p) {
p.mask_q_stride = 0;
}
// Set default Q strides for contiguous b h l d layout on PagedAttentionParams.
// Set default Q strides for a paged decode params struct.
template<typename P>
inline void set_default_paged_strides(P& p) {
p.q_stride_b = p.q_head * p.q_len * p.head_dim;