feat: unify attention backend with multi-dim mask support
- Add attention() functional entry delegating to active backend - GQA/MLA forward calls attention() instead of inline cache/SDPA - CUDA kernels support 2D/3D/4D mask via mask_h_stride field - CudaBackend.fwd_decode builds 2D padding mask for mixed seq_lens - KVCache.max_len precomputed in bind_tasks to avoid GPU sync - batch==1 decode short-circuits mask=None - Split tests into conftest, test_backend, test_backend_equivalence, test_kernel_mask - 440 tests pass, L20 decode 1.44-1.60x speedup vs torch native
This commit is contained in:
@@ -64,7 +64,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
|
||||
|
||||
// KV: stride-based base
|
||||
int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
|
||||
int mask_batch_base = batch * p.mask_b_stride;
|
||||
int mask_batch_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
|
||||
int tiles = (p.kv_len + P_BC - 1) / P_BC;
|
||||
int tt = G * ROWS;
|
||||
int lid = row * G + gpos;
|
||||
|
||||
Reference in New Issue
Block a user