perf: preload V in decode split-kv shared mem and cache partial tensors

- Preload V into shared memory alongside K to eliminate per-element KV address lookups in the inner softmax/accum loop (doubles smem)
- Cache split-KV partial tensors (o_part, ml_part) with static tensors instead of per-call allocation in both decode and paged-decode paths
- Force is_causal=True in CUDA decode backend (decode is always causal)
This commit is contained in:
2026-08-06 18:27:00 +08:00
parent 4f2e03880b
commit d0c5debbab
5 changed files with 31 additions and 14 deletions
+1 -1
View File
@@ -157,7 +157,7 @@ struct DecodeLauncherScalar {
int kv_len = KV::host_kv_len(p);
int chunks_total = (kv_len + DC_CHUNK - 1) / DC_CHUNK;
p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
size_t smem = 2 * DC_CHUNK * p.head_dim * sizeof(bf16);
int g = min(group_size, 32); // cap at 32 to respect 1024-thread limit
dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
dim3 block(32, g);