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

refactor: stride-based attn interface with layout and causal mask

Browse Source 
- 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:
ViperEkura
2026-07-14 21:34:42 +08:00
parent 2c7a71a9c0
commit 57729fd92d
18 changed files with 558 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/tests/attn_decode_test.cu
+4 -3
View File
@@ -98,8 +98,9 @@ static void bench() {
AttentionParams<bf16> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl;
p.head_dim = D; p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
p.head_dim = D; p.use_mask = 0; p.causal_offset = -1;
p.scale = 1.0f / sqrtf((float)D);
set_default_strides(p);
p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO;
DecodeScratch sc;
@@ -160,7 +161,7 @@ int main() {
AttentionParams<bf16> p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D;
p.use_mask=0; p.is_causal=0; p.causal_offset=0;
p.use_mask=0; p.causal_offset=-1;
p.scale=1.0f/sqrtf((float)D);
p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
@@ -180,7 +181,7 @@ int main() {
cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
float* ref=new float[nQ];
cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, 0, 0);
cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, -1);
float max_err=0;
for (size_t i=0;i<nQ;i++){
csrc/tests/attn_paged_decode_test.cu
+5 -3
View File
@@ -138,13 +138,14 @@ static int run_test(int B, int Hq, int Hkv, int kv_len, int page_size, int seed)
}
float* h_o_ref = (float*)calloc(B * Hq * HEAD_DIM, sizeof(float));
cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv, 1, kv_len, HEAD_DIM, 0, 0);
cpu_attention_ref(h_q_f, h_k_f, h_v_f, nullptr, h_o_ref, B, Hq, Hkv, 1, kv_len, HEAD_DIM, -1);
float scale_val = 1.0f / sqrtf((float)HEAD_DIM);
PagedAttentionParams<bf16, float> p;
p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1;
p.kv_len = kv_len; p.head_dim = HEAD_DIM;
p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0;
p.use_mask = 0; p.causal_offset = -1;
set_default_strides(p);
p.num_splits = 1; p.scale = scale_val;
p.page_size = page_size; p.max_pages = max_pages;
p.page_table = d_pt;
@@ -272,7 +273,8 @@ static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) {
PagedAttentionParams<bf16, float> pa;
pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1;
pa.kv_len = kv_len; pa.head_dim = HEAD_DIM;
pa.use_mask = 0; pa.is_causal = 0; pa.causal_offset = 0;
pa.use_mask = 0; pa.causal_offset = -1;
set_default_paged_strides(pa);
pa.num_splits = 1; pa.scale = scale_val;
pa.page_size = page_size; pa.max_pages = max_pages;
pa.page_table = d_pt;
csrc/tests/attn_prefill_test.cu
+5 -3
View File
@@ -75,7 +75,8 @@ static void bench() {
AttentionParams<bf16> p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
p.use_mask=0; p.causal_offset=causal?0:-1;
set_default_strides(p);
p.scale=1.0f/sqrtf((float)D);
p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
@@ -143,7 +144,8 @@ int main() {
AttentionParams<bf16> p;
p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D;
p.use_mask=0; p.is_causal=causal; p.causal_offset=0;
p.use_mask=0; p.causal_offset=causal?0:-1;
set_default_strides(p);
p.scale=1.0f/sqrtf((float)D);
p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO;
@@ -158,7 +160,7 @@ int main() {
cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost);
float* ref=new float[nQ];
cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal, 0);
cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal ? 0 : -1);
float max_err=0;
for (size_t i=0;i<nQ;i++) {
csrc/tests/test_utils.cuh
+29 -2
View File
@@ -101,6 +101,32 @@ void dispatch_by_head_dim(int head_dim, Fn&& fn) {
_HeadSwitch<32, 64, 128, 256>::call(head_dim, fn);
}
// Set default strides for contiguous b h l d layout on AttentionParams.
template<typename P>
inline void set_default_strides(P& p) {
p.q_stride_b = p.q_head * p.q_len * p.head_dim;
p.q_stride_h = p.q_len * p.head_dim;
p.q_stride_l = p.head_dim;
p.q_stride_d = 1;
p.kv_stride_b = p.kv_head * p.kv_len * p.head_dim;
p.kv_stride_h = p.kv_len * p.head_dim;
p.kv_stride_l = p.head_dim;
p.kv_stride_d = 1;
p.mask_b_stride = p.kv_len;
p.mask_q_stride = 0;
}
// Set default Q strides for contiguous b h l d layout on PagedAttentionParams.
template<typename P>
inline void set_default_paged_strides(P& p) {
p.q_stride_b = p.q_head * p.q_len * p.head_dim;
p.q_stride_h = p.q_len * p.head_dim;
p.q_stride_l = p.head_dim;
p.q_stride_d = 1;
p.mask_b_stride = p.kv_len;
p.mask_q_stride = 0;
}
// Generic CPU reference for multi-query / grouped-query attention.
// Tensor shapes (all float*):
// Q : [B, Hq, q_len, D]
@@ -108,10 +134,11 @@ void dispatch_by_head_dim(int head_dim, Fn&& fn) {
// V : [B, Hk, kv_len, D]
// O : [B, Hq, q_len, D]
// mask: if q_len == 1, shape is [B, kv_len]; otherwise mask is not supported.
// causal_offset: -1 = non-causal; >=0 = absolute position of first Q token.
static void cpu_attention_ref(
const float* Q, const float* K, const float* V, const bool* mask,
float* O, int B, int Hq, int Hk, int q_len, int kv_len, int D,
int is_causal, int causal_offset
int causal_offset
) {
float scale = 1.0f / sqrtf((float)D);
int n_rep = Hq / Hk;
@@ -122,7 +149,7 @@ static void cpu_attention_ref(
float mv = -INFINITY, sv = 0.0f;
float accum[256] = {0.0f};
int lim = kv_len;
if (is_causal) {
if (causal_offset >= 0) {
int c = qi + causal_offset + 1;
lim = (c < kv_len) ? c : kv_len;
}