refactor: stride-based attn interface with layout and causal mask
- Replace is_causal + causal_offset with unified causal_offset (-1 = off, >=0 = first Q pos)
- Causal and mask can now coexist (was mutually exclusive)
- Add stride-based addressing for Q/KV/O (layout-agnostic, zero-copy)
- Add layout param ("bhld"/"blhd") parsed in Python, passed as int to C++
- Support 2D [batch, kv_len] and 3D [batch, q_len, kv_len] mask
- Vectorize paged KV gather in Python fallback (was per-token Python loop)
- Extract shared helpers: compute_num_splits, alloc_split_partials, dispatch_head_dim
- Unify paged_decode entry via attn_pack_paged_params
- Update mma_softmax_tile for 3D mask with per-row qrow indexing
This commit is contained in:
@@ -10,11 +10,19 @@ struct AttentionParams {
|
||||
int kv_len;
|
||||
int head_dim;
|
||||
int use_mask;
|
||||
int is_causal;
|
||||
int causal_offset;
|
||||
int causal_offset; // -1 = non-causal; >=0 = absolute position of first Q token
|
||||
int num_splits;
|
||||
float scale;
|
||||
|
||||
// Q strides (element offsets for each dim — layout-agnostic)
|
||||
int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
|
||||
// KV strides (K and V share the same layout — only base pointers differ)
|
||||
int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
|
||||
|
||||
// Mask: 2D [batch, kv_len] (mask_q_stride=0) or 3D [batch, q_len, kv_len]
|
||||
int mask_b_stride; // = kv_len (both 2D and 3D)
|
||||
int mask_q_stride; // 2D: 0 (all q rows share); 3D: kv_len
|
||||
|
||||
const T* __restrict__ q;
|
||||
const T* __restrict__ k;
|
||||
const T* __restrict__ v;
|
||||
@@ -34,7 +42,6 @@ struct PagedAttentionParams {
|
||||
int kv_len;
|
||||
int head_dim;
|
||||
int use_mask;
|
||||
int is_causal;
|
||||
int causal_offset;
|
||||
float scale;
|
||||
|
||||
@@ -42,12 +49,19 @@ struct PagedAttentionParams {
|
||||
int page_size;
|
||||
int max_pages;
|
||||
|
||||
// Q strides (layout-agnostic)
|
||||
int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
|
||||
|
||||
// Mask strides (2D or 3D)
|
||||
int mask_b_stride;
|
||||
int mask_q_stride;
|
||||
|
||||
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;
|
||||
AT* __restrict__ ml_part;
|
||||
|
||||
Reference in New Issue
Block a user