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

refactor: unify attention param field names

Browse Source 
- rename q_stride_* to q_*_stride to match mask stride convention
- rename mask_q_stride to mask_l_stride for consistent l-dim naming
- merge k/v and k_cache/v_cache into k_ptr/v_ptr; rename q to q_ptr
- KVSource policy selects contiguous vs paged mode at compile time
This commit is contained in:
ViperEkura
2026-08-09 20:23:58 +08:00
parent d565d44c43
commit a5a3cc1fc2
11 changed files with 124 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_common.h
+25 -19
View File
@@ -23,22 +23,35 @@ struct AttentionParams {
int q_head;
int kv_head;
int head_dim;
int use_mask;
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;
// -1 = non-causal; >=0 = absolute position of first Q token
int causal_offset;
int use_mask;
int num_splits;
// Mask: 2D [batch, kv_len], 3D [batch, q_len, kv_len],
// or 4D [batch, n_heads, q_len, kv_len] (head dim broadcasts when stride=0)
int mask_b_stride; // batch stride
int mask_h_stride; // head stride (0 = broadcast across heads)
int mask_q_stride; // q stride (0 = all q rows share)
// Q strides
int q_b_stride;
int q_h_stride;
int q_l_stride;
int q_d_stride;
// K/V strides
int kv_b_stride;
int kv_h_stride;
int kv_l_stride;
int kv_d_stride;
// Mask strides
int mask_b_stride;
int mask_h_stride;
int mask_l_stride;
const T* __restrict__ q_ptr;
const T* __restrict__ k_ptr;
const T* __restrict__ v_ptr;
const bool* __restrict__ mask;
const T* __restrict__ q;
T* __restrict__ o;
AT* __restrict__ o_part;
AT* __restrict__ ml_part;
@@ -46,13 +59,6 @@ struct AttentionParams {
// ---- contiguous K/V mode ----
int q_len;
int kv_len;
int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d;
const T* __restrict__ k;
const T* __restrict__ v;
// ---- paged (SGLang flat pool) mode ----
const T* __restrict__ k_cache;
const T* __restrict__ v_cache;
// Indexing
const int64_t* __restrict__ req_to_token; // [num_reqs, max_context_len]