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
@@ -57,8 +57,8 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
            int s = i / p.head_dim;
            int d_dim = i % p.head_dim;
            int kc = chunk_start + s;
-            int token = KV::resolve_token(p, kctx, kc, true);
-            KVAddr a = KV::kv_addr_from_token(p, kctx, token, d_dim);
+            KVAddr a = KV::template decode_addr<1>(
+                p, kctx, batch, kv_head, kc, d_dim, true, true);
            k_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.k) : (bf16)0.f;
            v_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.v) : (bf16)0.f;
        }