Skip to content

cuda: add opt-in ordered KV gathering for indexed decode - #993

Open
JordiPosthumus wants to merge 1 commit into
antirez:mainfrom
JordiPosthumus:codex/cuda-ordered-indexed-decode
Open

cuda: add opt-in ordered KV gathering for indexed decode#993
JordiPosthumus wants to merge 1 commit into
antirez:mainfrom
JordiPosthumus:codex/cuda-ordered-indexed-decode

Conversation

@JordiPosthumus

Copy link
Copy Markdown
Contributor

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

  • Unset, 0 and unrecognized values retain the original path. Only 1 enables gathering.
  • Limited to one GPU, single-token decode, 64 heads, dimension 512, at most 256 raw rows and top_k <= 512. Other shapes and graph capture retain the original path. Metal and ROCm are unchanged.
  • Reuses the existing ordered scratch allocator, with input/output overlap guards and allocation-failure fallback. Maximum staging request is about 1.50 MiB; no per-session registry, cache-format change or retained activation pointer is added.
  • The regression target gains a 48-case API test. Its existing smoke target also needs ds4_image.o to resolve ds4_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.

  • M3 Ultra, macOS 27.0: 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.
  • GB10, CUDA 13.0.88 / 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-regression passes both the existing smoke test and the new 48-case API test.
  • The API test checks bit-exact full outputs, a double-precision scalar oracle, wrapped raw rings, duplicate/invalid/invisible indices, shape fallbacks, guards, scratch growth/reuse and pending A/B calls with an intervening scratch user.
  • compute-sanitizer --tool memcheck --error-exitcode 99 ./tests/test_cuda_indexed_gather and the corresponding synccheck run report zero errors. nsys profile --trace=cuda --sample=none --cpuctxsw=none ./tests/test_cuda_indexed_gather records 148 gather launches, exactly the test's expected count.
  • git diff --check passes.

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-bench sweep was run:

  • NVIDIA GB10, CUDA 13.0.88, driver 580.173.02; Vision-Exp IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8 and its vision encoder.
  • One fixed text corpus; same 128 teacher-forced tokens restored from the same snapshot. Six balanced observations per mode/frontier; first-token latency reported separately. Rates below are harmonic steady-decode rates within one executable, gathering off/on.
Context Off tok/s On tok/s Change
8192 15.879 17.555 +10.55%
16384 16.007 17.535 +9.55%
32768 15.181 16.584 +9.24%
65536 14.577 15.860 +8.80%
131045 13.446 14.628 +8.79%

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant