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
2026-07-30 20:38:34 +08:00
parent 97114b95a4
commit 3067a8e1a6
19 changed files with 438 additions and 81 deletions
@@ -109,8 +109,8 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
        int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
        mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
                                           0, 0,
-                                           p.mask_b_stride, 0,
-                                           batch,
+                                           p.mask_b_stride, 0, 0,
+                                           batch, 0,
                                           p.mask,
                                           Sacc, Oacc, m0, m1, l0, l1, lane);