Skip to content

cuda: accelerate long-context B1 indexer and 1M HCA - #763

Open
wangshang23 wants to merge 1 commit into
antirez:mainfrom
wangshang23:perf/cuda-long-context-indexer-hca
Open

cuda: accelerate long-context B1 indexer and 1M HCA#763
wangshang23 wants to merge 1 commit into
antirez:mainfrom
wangshang23:perf/cuda-long-context-indexer-hca

Conversation

@wangshang23

Copy link
Copy Markdown

Fixes #762.

What changed

  • Route non-quality B1 indexer scoring to the existing WMMA scorer at
    n_comp >= 8192; retain the direct scorer for shorter contexts, quality
    mode, and the existing DS4_CUDA_NO_INDEXER_WMMA escape hatch.
  • Raise the attention score capacity from 8192 to 8448 rows so the 1M HCA
    shape fits 8192 compressed rows plus the existing 256-row raw reserve.
  • Extend the CUDA long-context smoke test to cover 8192 compressed HCA rows
    and direct-vs-WMMA B1 score/top-k equivalence on QAT-format values.

The HCA change does not add a new attention implementation. It keeps the
existing exact score-split path active instead of falling back to the slow
online kernel for the last 256 rows at 1M.

DGX Spark results

Base: 84cc882, NVIDIA GB10 (sm_121), DeepSeek-V4-Flash IQ2XXS Q2 model,
Q8-to-F16 weight cache disabled. A/B paths alternate in one process/session,
with 3 warmups and 30 samples each.

Context Base tok/s PR tok/s Speedup
128K 12.168 12.468 1.025x
256K 10.642 11.290 1.061x
512K 8.647 9.574 1.107x
1M 4.372 6.981 1.597x

At 1M, Nsight Systems attributes the main change to:

  • CSA: direct scorer 78.036 ms -> WMMA scorer 48.565 ms across 21 layers.
  • HCA: online kernel 97.286 ms -> exact score-split score+finalize
    25.400 ms across 20 layers.

Full benchmark details and the existing-work audit are in #762. Related #258
uses approximate HISA pruning and is complementary to this full-top-512 flat
scorer optimization.

Correctness

  • B1 QAT score max delta: 6.10352e-05.
  • B1 top-512 differences: 0/512.
  • Isolated 8192-row HCA regression output delta: 0.
  • Full synthetic-frontier logits are bit-identical through 512K.
  • At 1M, legacy online vs existing exact score-split has max logit delta
    0.0469046, RMSE 0.00876944, and identical argmax.
  • Real 30,474-token story fact-recall regression passes.

Validation

make cuda-regression CUDA_ARCH=sm_121

Passes on DGX Spark / GB10.

--logprob-vectors currently fails short_code_completion identically on
both base and PR builds with this GGUF, so it is not a regression from this
change.

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.

CUDA: B1 indexer misses WMMA and 1M HCA falls off the 7936-row fast path

1 participant