refactor: unify paged and contiguous attention kernels via KVSource policy

- 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:
2026-08-05 14:06:13 +08:00
parent 6dffb0305a
commit 2667b8116d
17 changed files with 415 additions and 890 deletions
+1 -1
View File
@@ -16,7 +16,7 @@ torch::Tensor attn_paged_decode(
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
auto stream = at::cuda::getCurrentCUDAStream();
PagedAttentionParams<bf16> p;
AttentionParams<bf16> p;
attn_pack_paged_decode_params(q, k_cache, v_cache,
req_to_token, req_pool_indices, kv_indptr,
max_seq_len, mask, causal_offset, scale, p);