refactor: standardize packed 3d inference

- keep training attention on dense 4d tensors - use packed 3d tensors with KV cache for inference - extend CUDA rotary embedding to packed 3d inputs - adapt torch, CUDA and FlashAttention backend dispatch Benchmark: NVIDIA L20, BF16, 1B model, paged KV cache, CUDA Graph, prompt 512, generation 128 (median of 3 alternating runs) - batch 1: 234.5 -> 242.6 tok/s (1.034x, +3.4%) - batch 8: 1243.1 -> 1286.6 tok/s (1.035x, +3.5%)
2026-08-19 00:36:53 +08:00
parent f7f14d0e5f
commit 3d3ea47d37
9 changed files with 151 additions and 30 deletions
@@ -73,8 +73,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
            int r = i / Traits::HEAD_DIM, d = i % Traits::HEAD_DIM;
            int kc = kv0 + r;
            bool valid = kc < seq_len;
-            int token = KV::resolve_token(p, kctx, kc, valid);
-            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d);
+            // All GQA passes consume new K/V directly. Only the first pass
+            // persists it, so no cross-block synchronization is required.
+            KVAddr a = KV::template decode_addr<Traits::VEC>(
+                p, kctx, batch, kv_head, kc, d, valid, pass == 0);
            int off = r * Traits::LD + swiz_col(d, r, Traits::SWIZ_MASK);
            cp_async_16_pred(&dK[off], a.k, a.valid);
            cp_async_16_pred(&dV[off], a.v, a.valid);