Repository navigation
[Bugfix] Engage indexer TP row sharding on total rows, not per-rank rows - #61647
Merged
ZJY0516 merged 7 commits intoOct 15, 2026
Merged
Conversation
The row-sharding gate introduced in #54951 required MIN_TP_SHARD_ROWS_PER_RANK * tp_size scheduled prefill rows, so the total threshold grew linearly with tp_size: TP8 needs 128K rows, i.e. --max-num-batched-tokens >= 128K, which 80 GB Hopper nodes cannot even start with (the KV-cache budget goes negative during memory profiling). Per rank, the saved replicated indexer work (~ (1 - 1/tp) * rows) and the all_gatherv index traffic (~ (tp - 1)/tp * rows * topk * 4B) scale identically with tp, so the engagement crossover in TOTAL rows is tp-independent to first order. Replace the per-rank floor with the TP4-calibrated total floor of 64K rows, which keeps TP2/TP4 behavior unchanged and lets TP8 engage with --max-num-batched-tokens 65536. Also add VLLM_INDEXER_TP_SHARD_MIN_ROWS to override the floor and an info log on first engagement. Measured on 8xH200 with an 80 GB-equivalent budget (TP8, --max-num-batched-tokens 65536, 256K-context requests, vllm bench serve): mean TTFT 8246 ms -> 7318 ms (-11.3%), zero engagements before the fix vs. engaged at num_rows=65536 after it. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: Zheng Cai <caizheng1993@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Keep the change minimal: the TP-independent total floor already carries the auto-tune behavior, and no workload has asked for an override yet. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Signed-off-by: Zheng Cai <caizheng1993@gmail.com>
The 64K total floor is calibrated on NVSwitch; on slower interconnects (e.g. PCIe) the per-row all_gatherv cost is an order of magnitude higher and the fixed floor would engage a losing path, so deployments need an escape hatch without a code change. Drop the info_once engage log: lru_cache keys on (msg, *args), so every distinct num_rows logs again — noisy over a long-serving run, and not needed in this PR. Also pin the TP2 behavior explicitly: 32K-64K total row batches engaged under the old per-rank floor and now stay replicated; the crossover's latency term is largest at small tp, so the later engagement is the right direction, but it is a deliberate change, not an invariant. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Contributor
Author
|
/ci run |
|
❌ @zigzagcai, A reviewer with write access must run |
ZJY0516
reviewed
Oct 15, 2026
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
ZJY0516
requested review from
ProExpertProg,
WoosukKwon,
houseroad,
mgoin,
robertgshaw2-redhat,
tlrmchlsmth,
yewentao256 and
youkaichao
as code owners
October 15, 2026 12:09
ZJY0516
changed the title
|
✅ @zigzagcai, CI is now available for this PR.
|
ZJY0516
approved these changes
Oct 15, 2026
ZJY0516
left a comment
ZJY0516
left a comment
Member
There was a problem hiding this comment.
LGTM now, after some cleanup. BTW, we prefer config instead of env variable
Member
|
/ci run |
|
✅ Triggered Buildkite CI #95122 for commit |
ZJY0516
enabled auto-merge (squash)
October 15, 2026 12:24
Contributor
Author
Thanks for the review and changes! Sure, I also think config is betther. |
michael-gpu
pushed a commit
to practical-inference/vllm-ampere-extended
that referenced
this pull request
Oct 20, 2026
…ows (vllm-project#61647) Signed-off-by: Zheng Cai <caizheng1993@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>
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.
Overview
Follow-up to #54951, addressing #54951 (comment): gate the sparse-indexer TP row sharding on total scheduled prefill rows instead of per-rank rows, so the 8xH100 TP8 recipe can actually engage it. The floor becomes
AttentionConfig.indexer_tp_shard_min_rows(default 65536).Not a duplicate: no other open PR changes how this floor scales with TP. #60601 lowers the same floor on ROCm only; whichever lands second needs a rebase.
Claims
--max-num-batched-tokens 65536; before, it needed 131072, which OOMs on 80 GB Hopper.-ac '{"indexer_tp_shard_min_rows": 131072}'.Validation
Problem reproduced on 8xH200 with an 80 GB-equivalent budget (
--gpu-memory-utilization 0.56, TP8, BF16 KV, merged-commit build3e8e1364):--max-num-batched-tokens 131072Available KV cache memory: -6.94 GiB--max-num-batched-tokens 65536num_rows=65536 < threshold=131072on every chunk → 0 engagementsTTFT at the H100 budget (8xH200,
--gpu-memory-utilization 0.568≈ 81,662 MiB/card vs H100's 81,559 MiB), TP8,--max-num-batched-tokens 65536,--max-model-len 1048576, FP8 KV,FLASH_ATTN_MLA_SPARSE, streaming completions with one request in flight; only the floor differs between arms:The replicated arm was measured with an earlier revision's
VLLM_INDEXER_TP_SHARD_MIN_ROWS=2000000; the equivalent now is-ac '{"indexer_tp_shard_min_rows": 2000000}'. The gain grows with context depth (0.48 s saved at 255K vs 16.35 s at 1.02M): per-row indexer cost ramps with causal depth while the per-row exchange cost stays constant.Tests
pytest tests/v1/attention/test_indexer_tp_row_shard.py: 12 passed (TP8 parametrization, total-not-per-rank floor, TP2 32K–64K window stays replicated, floor never belowtp_size).Details
Per rank, the saved replicated indexer work (
~ (1 - 1/tp) * rows) and theall_gathervindex traffic (~ (tp - 1)/tp * rows * topk * 4B) both scale with(tp - 1)/tp, so the engagement crossover in total rows is tp-independent to first order. The old per-rank floor grew the total requirement linearly withtp. 65536 is the TP4-calibrated total floor from #54951.tp.balanced_prefill_row_shardtakes the floor as an argument and clamps it to at leasttp_size, so every rank always gets at least one row.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code