Conversation
The CUDA DSA prefill indexer recomputes identical logits and Top-K on every tensor-parallel rank. Assign each prefill query row to exactly one rank and publish the finished rows with one all_gatherv per indexer layer. Rows are independent, so this is a layout-preserving concatenation rather than a Top-K merge. The partition is contiguous and balances the exact number of scored keys, derived from CPU scheduler metadata with no device sync. Assisted-by: OpenAI Codex Assisted-by: Claude Code (Opus) Signed-off-by: lanqinghuan <qinghuan_lan@163.com>
|
Updated the PR body with matched vendor-runtime A/B data and explicit scope/non-regression evidence. Headline result: at 1M context / C1, p50 TTFT improved from 63,898.9 ms to 49,482.6 ms (-22.56%), p50 TPOT from 102.70 ms to 92.75 ms (-9.69%), and total token throughput from 14,212 to 18,076 tok/s (+27.19%). At 128K / C4, the corresponding changes are TTFT -3.79%, TPOT -2.43%, and throughput +3.03%. Control is the unmodified vendor vLLM package; treatment overlays only this PR's three implementation files. Both arms use TP4 on H200, BF16 KV, FP8 weights, chunked prefill, MTP=0, no prefix cache, enforce-eager, temperature 0, ignore-eos, and one warmup. All 7 optimized requests and all control requests completed successfully. Focused row-sharding tests pass (50/50), and related existing indexer/MLA regression tests pass (61/61). |
Performance improvement rationale and resultsThe main bottleneck addressed by this PR is the replicated sparse indexer prefill computation across tensor-parallel ranks. For GLM-5.3-Flash, the indexer projections and compressed K cache are replicated on every TP rank. During a long prefill, every rank therefore performs the same MQA-logits and top-k computation for all query rows, even though the final top-k indices are the only result needed by the subsequent attention computation. This redundant work becomes increasingly expensive as the context length grows because each query row scans a longer causal compressed-K range. This PR introduces cost-balanced contiguous row sharding across TP ranks:
The communication overhead is small compared with the eliminated replicated computation. For Matched A/B benchmarkBoth arms used the same configuration:
All requests completed successfully.
The improvement becomes substantially larger at 1M context because the eliminated replicated MQA/top-k work scales with the causal context length, while the additional communication remains limited to the final top-k indices. Scope and non-regressionThe functional change is limited to:
No scheduler, model weights, KV-cache layout, decode kernel, or public API is modified. The optimization is explicitly disabled for short prefills and unsupported execution modes, including DCP/PCP, ROCm, batch-invariant execution, symm-mem/PynCCL-disabled configurations, and FULL CUDA graph execution. When the row threshold is not met, the existing replicated indexer path is retained. Validation completed:
The A/B test intentionally uses the matching vendor runtime for both arms because the editable checkout currently has an unrelated DeepGEMM/FlashAttention ABI mismatch. This keeps the comparison focused on the PR's Python implementation changes rather than mixing in a binary compatibility issue. We would appreciate a review of the implementation and benchmark methodology. Given the clear long-context gains, especially the 1M-context result, we hope this optimization can be considered for merging. |
Signed-off-by: Zheng Cai <8370601+zigzagcai@users.noreply.github.com>
Inference-accuracy A/B (updated)I also validated inference correctness because this optimization changes the prefill execution/communication path. The first version exposed an important GLM-5.3-specific detail: On 4x NVIDIA H200 (TP=4), with native FP8 weights, BF16 KV cache, chunked prefill (
The change is limited to long CUDA prefill MQA/indexer row partitioning and one TP all-gather of the complete token-index rows. Decode, DCP/PCP, full CUDA-graph, short-prefill, and unsupported communication configurations remain on the original path. The focused row-sharding test passes ( The 1M-token case is not included in this accuracy run because this checkpoint/runtime is capped at |
|
Close this PR and re-open a PR targeted for vLLM upstream, since maintainer suggest a separate PR would be better |
2 participants
Summary
all_gathervindex_kpool=4pool expansion, tail-token handling, and mixed decode/prefill semanticsMotivation
The indexer projections and K cache are replicated across TP ranks, so long prefill currently repeats the same MQA/top-k work on every rank. At 128K and 1M contexts the planner partitions rows by causal compressed-key cost; the final communication is only
prefill_rows * index_topk * sizeof(int32)per indexer layer (1 MiB and 8 MiB respectively forindex_topk=2048).Matched optimized A/B
Both arms used the same 4x NVIDIA H200 TP4 process, native FP8 weights, BF16 KV, chunked prefill with
max_num_batched_tokens=8192, no prefix cache,enforce_eager, MTP=0, temperature 0,ignore_eos, and one warmup request. The control is the unmodified vendor vLLM package; the treatment overlays only this PR's Python changes. Values are fromvllm bench servepercentile output; every request succeeded.Relative to baseline, optimized changes are:
The 1M test is intentionally C1 to isolate the long-prefill critical path; C4 at 128K covers queueing/throughput behavior. Raw JSON results and the full reproduction note are in the serving workspace at
results/long_context_row_sharding.md.Scope and non-regression
The functional change is limited to four files: the generic sparse indexer, the GLM kpool indexer, the MLA metadata/planner, and a focused TP row-sharding test. No scheduler, model weights, KV-cache layout, decode kernel, or public API is changed. For short prefills the planner returns
Nonebelow 1024 rows per TP rank, so the existing replicated indexer path is used byte-for-byte. DCP/PCP, symm-mem/PynCCL-disabled, batch-invariant, ROCm, and FULL CUDA-graph configurations are explicitly gated off and retain their prior path.pytest -q tests/v1/attention/test_indexer_tp_row_shard.py: 50 passedpython -m compileallover changed modules: passedgit diff --check: passedPlanner evidence
[65536, 27146, 20829, 17561]rows on TP0..TP3[524288, 217167, 166638, 140483]rows on TP0..TP3The source editable checkout cannot currently be used for an apples-to-apples runtime because its DeepGEMM/FlashAttention binaries are ABI-incompatible with the model image and fail in the pre-existing
fp8_fp4_mqa_logitspath even withVLLM_DISABLE_PYNCCL=1. Therefore the A/B above deliberately uses the matching vendor runtime for both arms; this isolates the PR change instead of attributing an unrelated extension failure to the optimization.Test environment
4x NVIDIA H200, TP4, native FP8 weights, BF16 KV, chunked prefill, model snapshot from the GLM-5.3-Flash release.