[perf] vulkan: SDPA KV Cache update - #6801
Conversation
Summary: Replace the generic Concat-based SDPA KV cache path with dedicated Vulkan cache update and append shaders. This keeps decode cache updates on GPU and reduces repeated full-cache copies during autoregressive generation. Changes: 1. Add dedicated append and grow-copy shaders for SDPA KV cache updates 2. Reuse existing VkMat cache capacity when appending new decode tokens 3. Wire SDPA Vulkan key and value cache outputs through the new cache helper
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eb3bc8449f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| top_blob = past_blob; | ||
| top_blob.h = dst_seqlen; | ||
| record_kvcache_append(cur_blob, top_blob, width, past_seqlen, cur_seqlen, num_group, pipeline_append, cmd); |
There was a problem hiding this comment.
When the previous cache has spare capacity, this aliases top_blob to past_blob and then records an append into that same buffer, so a normal forward call mutates one of its input blobs. This breaks callers that reuse a prefix cache for multiple continuations/beam branches or keep the first returned cache while running another step from the same input cache, because later appends overwrite the shared storage that earlier outputs also reference; the old concat path always produced an independent output cache.
Useful? React with 👍 / 👎.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #6801 +/- ##
==========================================
+ Coverage 95.13% 95.49% +0.36%
==========================================
Files 947 878 -69
Lines 416089 392899 -23190
==========================================
- Hits 395853 375207 -20646
+ Misses 20236 17692 -2544 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Summary: Return compact K/V cache outputs from Vulkan SDPA so generic download, packing, and layer tests interpret the cache layout consistently. This fixes the SDPA KV cache mismatch caused by encoding hidden capacity in VkMat::cstep. Changes: 1. Remove hidden KV cache capacity growth from the public SDPA Vulkan output path. 2. Allocate updated K/V cache tensors with the logical sequence length. 3. Rename the cache helper to reflect compact update behavior.
crafcat7
changed the title
Summary: Restore the empty past-cache fast path so SDPA can alias the current K/V tensors without allocating or dispatching extra work. This keeps the Vulkan KV cache path aligned with the intended zero-cost initialization behavior while removing an unnecessary append shader dependency. Changes: 1. Skip cache allocation and copy work when past_seqlen is zero by reusing the current K/V tensors directly. 2. Remove the append-only KV cache shader and the related pipeline plumbing from the Vulkan SDPA implementation. 3. Keep the cache update path focused on the non-empty history case so future growth still uses the dedicated copy logic.
2 participants
Summary
This PR updates the Vulkan
SDPAKV cache path so K/V cache materialization is handled insideSDPA_vulkanwith dedicated Vulkan shaders instead of a generic cache concat path.The final implementation keeps public K/V cache outputs compact. This is required because
VkMatis the shared Vulkan tensor representation used by ncnn layers, packing conversion, download/extract, and tests. A cache tensor returned fromSDPA_vulkan::forward()must therefore expose a normal, self-consistentw/h/c/csteplayout.Committed code changes are limited to Vulkan SDPA implementation files and shaders. This PR does not modify benchmark tools or external tooling.
What Changed
1. Dedicated KV cache update shaders
The Vulkan SDPA KV cache path now uses SDPA-specific shaders for cache materialization:
The new shader roles are:
sdpa_kvcache_update.comppastand currentcurinto a compactpast + curcache output.sdpa_kvcache_append.compcurinto the output atpast_seqlen + y.This removes the dependency on a generic concat layer for the SDPA KV cache update step and makes the cache path explicit inside
SDPA_vulkan.2. Compact public cache outputs
The final cache output layout is:
The public
VkMatoutput does not encode hidden cache capacity incstep.3. Existing attention compute paths remain unchanged
The cache update happens before attention kernels read K/V. The existing SDPA compute paths are preserved:
qk -> softmax -> qkvPerformance
Comparison basis
All results compare clean worktrees built with the same compiler, CMake options, and benchmark parameters:
81e8cbff39cc7404Illegal instructionissueImprovement is calculated as:
End-to-end LLM benchmark
The normal committed benchmark path was used. It keeps K/V cache as host
Matbetween decoder calls; no GPU-resident cache mode or modified benchmark path is involved.Parameters and sampling:
loop_count = 3num_threads = 1powersave = 0gpu_device = 0cooling_down = 0The prefill suite total changed from
3330.81 msto3384.11 ms(-1.60%). Prefill has no existing cache to concatenate, so it is a control case rather than the target of this optimization.The end-to-end result is small and noisy because model compute and host
Matcache upload/download dominate the dedicated KV update dispatch. It is still the comparable normal benchmark result: across the complete decode suite, the optimized revision reduces measured time by0.65%.Isolated Vulkan SDPA decode benchmark
perf_sdpa_decodeisolatesSDPAforward and uploads inputs before timing, so it exposes the cost removed by replacing genericConcatwith the SDPA-specific KV update pipeline.Both revisions used the same measurement cases on
gpu-0:Each value below is the mean of two runs. Each run uses the median of 20 timed command-buffer batches, divided by the calibrated batch size to obtain per-forward latency.
The dedicated KV update path improves isolated FP32 SDPA decode by
5.26%to10.73%. The FP16 arithmetic (fp16psa) path improves by up to2.48%at shorter cache lengths and becomes neutral as attention computation dominates at longer lengths.These two benchmark levels describe different effects: