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:
2026-07-30 20:38:34 +08:00
parent 97114b95a4
commit 3067a8e1a6
19 changed files with 438 additions and 81 deletions
+1 -1
View File
@@ -24,7 +24,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
// KV: [batch, kv_head, kv_len, head_dim] — stride-based base
int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h;
int mask_base = batch * p.mask_b_stride;
int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};