feat: paged decode attention with split-KV (scalar + MMA)
- PagedAttentionParams merged into attn_common.h - Scalar variant: warp-per-query-head split-KV, resolves page table per-position for the K/V shared-memory tile load - MMA variant (sm_80+): tensor-core head-packing with cp.async, single page-table lookup per tile (BC=32 fits within page_size>=32) - Standalone test: 14 cases across head_dim 32/64/128/256, GQA, multi-batch, both paths verified against CPU reference
This commit is contained in:
@@ -14,11 +14,39 @@ struct AttentionParams {
|
||||
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;
|
||||
};
|
||||
|
||||
template<typename T, typename AT = float>
|
||||
struct PagedAttentionParams {
|
||||
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;
|
||||
|
||||
int num_splits;
|
||||
int page_size;
|
||||
int max_pages;
|
||||
|
||||
const T* __restrict__ q;
|
||||
const T* __restrict__ k_cache;
|
||||
const T* __restrict__ v_cache;
|
||||
const bool* __restrict__ mask;
|
||||
const int64_t* __restrict__ page_table;
|
||||
|
||||
T* __restrict__ o;
|
||||
AT* __restrict__ o_part;
|
||||
|
||||
Reference in New Issue
Block a user