Repository navigation
[CB] Scale the per-token memory estimate by the TP size - #48566
Open
qgallouedec wants to merge 11 commits into
Open
qgallouedec wants to merge 11 commits into
qgallouedec wants to merge 11 commits into
Conversation
…elism PagedAttentionMemoryHandler estimated bytes per block and the activation peak with the unsharded head counts, while each rank only holds num_heads // tp_size of them, so under TP the pool came out tp x too small: Qwen3-32B at tp4 got 221k tokens and used 36 of 80 GiB per GPU, 14B at tp2 365k. The handler now takes the TP size when the KV projections are in the TP plan and divides both counts by it. Qwen3-32B tp4, 1024 short prompts, hub kernels: 6,523 -> 11,277 tok/s (pool 221k -> 884k tokens, 450 -> 672 concurrent requests); 14B tp2 365k -> 731k tokens.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
# Conflicts: # src/transformers/generation/continuous_batching/cache.py
qgallouedec
changed the title
Contributor
CI recapDashboard: View test results in Grafana |
Member
Author
|
Still needed: PagedAttentionMemoryHandler uses the unsharded head counts. @remi-or @ArthurZucker could you take a look? |
This branch has not been deployed
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.
Rebased on #47809, which already divides the KV heads by
tp_sizefor the page size. What remained: thePagedAttentionMemoryHandlerper-token activation estimate (query projection, new K/V, repeated KV at the attention peak) still used the unsharded head counts, so under TP it reservedtp_sizetimes too much per batch token and the pool came out smaller than it could be.The cache now passes the TP size the KV heads are sharded over (1 when they are not) and the handler scales the head counts with it. A rank holding
1/tp_sizeof the heads predicts the footprint of a model with that many heads; the new test checks exactly that against a handler built on the halved config.