cuda: add opt-in ordered KV gathering for indexed decode - #993
Open
JordiPosthumus wants to merge 1 commit into
Open
cuda: add opt-in ordered KV gathering for indexed decode#993JordiPosthumus wants to merge 1 commit into
JordiPosthumus wants to merge 1 commit into
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add
DS4_CUDA_INDEXED_DECODE_GATHER=1, an opt-in CUDA path that gathers Flash's selected F32 KV rows once before running the existing single-token attention kernel. Default behavior is unchanged.Why
Indexed decode makes each of 64 heads read the same scattered KV rows. The selected rows are shared, but their scattered memory access is repeated. This patch copies the raw window and selected compressed rows into contiguous scratch once, then calls the original attention kernel with remapped indices.
The important invariant is same rows, same slot order, same F32 values, same attention arithmetic. Duplicate indices stay duplicated; invalid and not-yet-visible rows stay masked. There is no sorting, precision conversion, new reduction or change to model math.
Compatibility
0and unrecognized values retain the original path. Only1enables gathering.top_k <= 512. Other shapes and graph capture retain the original path. Metal and ROCm are unchanged.ds4_image.oto resolveds4_deepseek4_attention_bounds; the missing prerequisite was reproduced on pristine main and is corrected here as test-link wiring.This is independent of the indexer optimization in #763 and token-tile prefill fix in #869. The top-k expansion in #478 must retain this path's 512-slot guard; #566's stream override and #628's backend split would need normal integration/rebase work if merged first.
Validation
Based directly on main
9ab705347c1775e7599ede7eb81a6255ec7dccb5, as one commit with no dependency on other PRs.make clean;make -j2 all ds4_test ds4_agent_test tests/test_session_state;./ds4_test --server;./ds4_agent_test;./tests/test_session_state;make -j2 cpu— pass. The build retains 27 Metal deprecation warnings; no Mac model was loaded.sm_121a: all five CUDA frontends build on pristine base and this PR; the same three host unit groups pass.make -j2 CUDA_ARCH=sm_121a cuda-regressionpasses both the existing smoke test and the new 48-case API test.compute-sanitizer --tool memcheck --error-exitcode 99 ./tests/test_cuda_indexed_gatherand the correspondingsynccheckrun report zero errors.nsys profile --trace=cuda --sample=none --cpuctxsw=none ./tests/test_cuda_indexed_gatherrecords 148 gather launches, exactly the test's expected count.git diff --checkpasses.Full-model evidence was collected before extraction into this clean upstream commit, using the same gather/attention implementation in an existing integration. These are supporting measurements, not a claim that a clean-head
ds4-benchsweep was run:All six paired blocks were positive at these frontiers. Excluding the first balanced block changes the 131045-token result to +8.05%. A separate warmed 2K dispatch trace recorded zero gather launches and effectively flat speed (19.69388 off / 19.69315 on); no 2K gain is claimed. The existing engine first uses indexed attention after 1024 compressed rows. The separately linked off-only baseline had only two observations per frontier: off-vs-off differences at 8K+ were −0.25% to −2.82%, so this is not proof of unchanged default speed. Warmed 2K prefill was 773.28 vs 771.68 tok/s (−0.21%). Startup samples and outliers were retained.
All 30,768,640 saved F32 values matched bit-for-bit across reference/candidate builds, including text and image-conditioned continuations. Separate private-server checks passed long-prefix restart/restore, image restore, cancellation, tool continuation and shared-prefix/history isolation. These full-model/server results apply to the pre-extraction integration. No
ds4_test --all, other-GPU, other-quantization or multi-GPU result is claimed.