[https://nvbugs/5977180][fix] size KV cache and admit requests per beam width - #16802
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
WalkthroughKV-cache estimation now accounts for maximum beam width. Resource managers expose request-level KV-cache budgets across supported pools. PyExecutor validates these budgets after sampler-specific validation. Tests cover estimation, delegation, capacity bounds, and admission outcomes. ChangesKV-cache budgeting and admission
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🔵 Low · up to The change updates KV-cache sizing and request admission for large beam widths, but the VSWA and linear-attention opt-out paths remain untested, leaving a bounded regression risk in those configurations. The PR is mergeable with explicit owner awareness and follow-up testing. Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant ResourceManager
participant KVCacheManager
PyExecutor->>ResourceManager: request KV block budgets
ResourceManager->>KVCacheManager: query target, draft, and cross managers
KVCacheManager-->>ResourceManager: return required and primary-capacity blocks
ResourceManager-->>PyExecutor: return resource-tagged budgets
PyExecutor->>PyExecutor: reject request when required blocks exceed capacity
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)
1261-1272: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAttention-window clamping still needs beam-width headroom.
get_max_atten_window_upper_bound()now reduces toblocks_in_primary_pool * tokens_per_block, somax_beam_width > 1no longer shrinks the clamp._validate_and_adjust_attention_windows()can still accept windows that exhaust KV cache during beam search; the- tokens_per_blockfallback is not enough for wider beams.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 1261 - 1272, The get_max_atten_window_upper_bound() calculation must reserve sufficient KV-cache capacity for the requested max_beam_width, rather than only subtracting one tokens_per_block when beams exceed one. Update the upper-bound computation so _validate_and_adjust_attention_windows() cannot accept a window that exhausts the primary pool during beam search, while preserving the existing positive-capacity assertions.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_resource_manager.py (1)
1017-1112: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winGood coverage of the beam-sharing math and delegation; two gaps worth a follow-up.
TestRequestBudgetsolidly covers: beam-width=1 parity with the historical value, prompt-sharing across beams, block-aligned prompts,ResourceManagerdelegation/no-op, and the_validate_request_budgetaccept/reject boundary. Coverage is currently insufficient for:
get_max_atten_window_upper_bound's beam-width behavior change (no test exercisesmax_beam_width > 1for that method in this file)._validate_request_budgetagainst aDRAFT_KV_CACHE_MANAGER/CROSS_KV_CACHE_MANAGER-only-constrained scenario (onlyKV_CACHE_MANAGERis exercised).Suggest adding these to
tests/unittest/_torch/executor/test_resource_manager.py, either in this PR or a fast follow-up, based on the resolution of the related correctness comments inresource_manager.pyandpy_executor.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/executor/test_resource_manager.py` around lines 1017 - 1112, Extend TestRequestBudget with coverage for get_max_atten_window_upper_bound using max_beam_width greater than one, verifying the updated beam-aware result. Also add _validate_request_budget coverage where capacity is constrained only by DRAFT_KV_CACHE_MANAGER or CROSS_KV_CACHE_MANAGER, confirming oversized requests are rejected and fitting requests are accepted without relying on KV_CACHE_MANAGER.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 706-710: Update _get_token_num_for_estimation so num_cache_blocks
is multiplied by the configured maximum beam width before calculating
max_num_tokens_for_estimation. Preserve the existing block and pool-group
scaling, and use the same max_beam_width value applied by
_create_dummy_context_requests to keep the warm-up KV-cache estimate consistent
with beam search.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 4736-4750: Extend _validate_request_budget to account for every
active KV pool, including DRAFT_KV_CACHE_MANAGER and CROSS_KV_CACHE_MANAGER,
rather than checking only the primary KV cache capacity. Compare the request’s
required blocks against each applicable pool’s available capacity and raise the
existing ValueError with the relevant pool context when any pool cannot satisfy
the request.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 1261-1272: The get_max_atten_window_upper_bound() calculation must
reserve sufficient KV-cache capacity for the requested max_beam_width, rather
than only subtracting one tokens_per_block when beams exceed one. Update the
upper-bound computation so _validate_and_adjust_attention_windows() cannot
accept a window that exhausts the primary pool during beam search, while
preserving the existing positive-capacity assertions.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_resource_manager.py`:
- Around line 1017-1112: Extend TestRequestBudget with coverage for
get_max_atten_window_upper_bound using max_beam_width greater than one,
verifying the updated beam-aware result. Also add _validate_request_budget
coverage where capacity is constrained only by DRAFT_KV_CACHE_MANAGER or
CROSS_KV_CACHE_MANAGER, confirming oversized requests are rejected and fitting
requests are accepted without relying on KV_CACHE_MANAGER.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 70fb17e8-709f-4519-b207-9d5ce33005c4
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytests/unittest/_torch/executor/test_resource_manager.py
|
Thanks for the detailed write-up — the diagnosis is right, and I agree with the direction: beams Line numbers are PR head 🔴 Blocker 1 — the budget accessors don't hold across KV cache manager implementations1a.
|
| Class | Base | get_num_blocks |
|---|---|---|
KVCacheManager |
BaseResourceManager |
✅ resource_manager.py:711 |
KVCacheManagerV2 |
BaseResourceManager |
❌ |
DeepseekV4CacheManager (main sparse/deepseek_v4/cache_manager.py:165) |
KVCacheManagerV2 |
❌ |
MiniMaxM3KVCacheManagerV2 (main sparse/minimax_m3/cache_manager.py:132) |
KVCacheManagerV2 |
❌ |
MambaHybridCacheManagerV2 (main mamba_cache_manager.py:2469) |
KVCacheManagerV2, … |
❌ |
MixedMambaHybridCacheManager, CppMambaHybridCacheManager, DSACacheManager, RocketKVCacheManager |
KVCacheManager, … |
✅ inherited |
BaseResourceManager (resource_manager.py:117-141) declares neither get_num_blocks nor
__getattr__, so there is no fallback.
DeepSeek-V4 hits this on both possible routes: with a sparse config,
get_kv_cache_manager_cls returns DeepseekV4CacheManager at _util.py:162-163; without one,
_non_hybrid_kv_cache_manager_cls returns KVCacheManagerV2, since
DeepseekV4ForCausalLM.get_model_defaults sets use_kv_cache_manager_v2: True
(modeling_deepseekv4.py:2459), resolved through llm_utils.py:547-566. Gemma4-hybrid is forced
to V2 unconditionally at _util.py:84-85.
The failure is silent. The AttributeError is caught by the bare except Exception in
_respond_if_invalid (py_executor.py:4982-4986), which calls _handle_errors(..., charge_budget=False) — documented at py_executor.py:6066-6070 as a per-request failure that
never triggers shutdown. So the server starts cleanly, accepts connections, and fails every
request with "'DeepseekV4CacheManager' object has no attribute 'get_num_blocks'" as the
client-facing error text. No traceback, no startup failure.
That's every request, not most: py_executor.py:4996 is the only path by which new requests enter
active_requests (:5471 is ADP dummy padding, :6462 re-filters existing ones).
py_executor.py:3828 _executor_loop
→ :3851 _prepare_and_schedule_batch
→ :3496 _fetch_and_activate_new_requests
→ :4980 _respond_if_invalid → _validate_request
→ :4768 _validate_request_budget ← added here
→ :4741 resource_manager.get_num_blocks(request)
→ resource_manager.py:2510 kv_cache_manager.get_num_blocks(request) 💥
Same fault via _executor_loop_overlap (:4293 → :4323) and _executor_loop_pp
(:2478 → :2496).
1b. get_max_resource_count() is not a block count
Adding the missing method wouldn't be enough, because the other accessor rests on an assumption the
interface doesn't make. BaseResourceManager documents it as:
# resource_manager.py:119-121
@abstractmethod
def get_max_resource_count(self) -> int:
"""Return the maximum number of real requests this manager can admit."""Implementations disagree on units:
KVCacheManager→impl.max_num_blocks— blocks ✔KVCacheManagerV2→ literal1(kv_cache_manager_v2.py:2999-3001,# TODO: implement this)DeepseekV4CacheManager→int(self.impl.get_quota(GPU_LEVEL))— bytes (main
sparse/deepseek_v4/cache_manager.py:999-1001; the same call populates
kv_cache_stats.allocated_bytesatkv_cache_manager_v2.py:2513)
So get_num_blocks_available() compares a block count against blocks, 1, or a byte count
depending on the model — rejecting everything in one case and nothing in another.
Even on the V1 path the number is loose: impl.max_num_blocks → mAllBlocksById.size(), reserved
and filled from both pools (main kvCacheManager.cpp:856), so with host_cache_size > 0 the
ceiling counts offloaded blocks that can't be GPU-resident — a request can pass admission and still
stall the scheduler, which is the failure mode this check exists to prevent.
blocks_in_primary_pool would be the honest ceiling.
Suggested shape: don't reuse get_max_resource_count for this. Add an explicit, documented
block-capacity method to the KV managers that need it, and gate the check on capability rather than
None — isinstance(kv_cache_manager, KVCacheManager), or a supports_budget_check property on
BaseResourceManager defaulting to False. Note the signature already hints at this:
get_resource_manager returns Optional[BaseResourceManager] (resource_manager.py:2498-2499),
and neither budget method is on that interface.
🔴 Blocker 2 — the window clamp from the cited cost model is missing
_num_blocks_for (resource_manager.py:700) correctly reproduces the shared/unshared split from
getNeededBlocksOneStep — but not that function's clamp. In main kvCacheManager.cpp:3504-3517:
auto const promptCacheLen
= std::min((isCrossKv() ? req.getEncoderOutputLen() : req.mPromptLen) + maxDraftTokensToAdd,
windowSize + mChunkSize) // ← clamp
+ mSinkBubbleLength;
auto const numSharedBlocks = promptCacheLen / getTokensPerBlock();
auto const numUnSharedTokens = promptCacheLen % getTokensPerBlock();
auto const numUnSharedBlocks = tc::ceilDiv(numUnSharedTokens, getTokensPerBlock()) * beamWidth;getRemainingBlocksToCompletion clamps the same way (kvCacheManager.cpp:3622-3626). Without it
the estimate grows without bound in prompt_len + max_new_tokens, while the true requirement
saturates at the window.
Whether that crosses the capacity threshold is config-dependent — pool block count is
memory-derived, so on a large GPU the headroom often absorbs it. So: a latent false-reject, most
reachable in exactly the regime this PR targets (high beam width, long generation on a windowed
model), rather than a guaranteed break.
One part is unit-invalid regardless of config: for VSWA, get_max_resource_count() →
getMaxNumBlocks() → sumWindows(...) (main kvCacheManager.h:1842-1845) — the sum across all
window pools — compared against a single-window full-attention demand.
The inputs are already on the object: self.max_attention_window_vec, self.is_vswa,
self.is_linear_attention are set at resource_manager.py:388-406. Clamp by window, or skip the
check when is_vswa or is_linear_attention.
🟡 Smaller items
Vestigial beam logic in get_max_atten_window_upper_bound (resource_manager.py:1261-1272).
After the divisor removal, math.floor(token_capacity / tokens_per_block) is exact division
(token_capacity is blocks_in_primary_pool * tokens_per_block), so max_blocks_per_seq == blocks_in_primary_pool and the function collapses to token_capacity. What survives is a
beam-conditional -= tokens_per_block at :1269 whose rationale — headroom against the
beam-scaled bound — no longer exists, and a max_beam_width parameter used nowhere else in the
function. The description defers this, but this PR is what invalidates it; leaving it half-done is
more confusing than finishing or reverting it.
Naming.
get_num_blocks(request)sits besideget_num_kv_blocks(num_tokens)
(resource_manager.py:1437) — near-identical names, different argument kinds.
get_num_blocks_to_completiondisambiguates.get_num_blocks_available()returns total capacity, not available blocks. The comment says so;
the name still misleads at the call site.
Implicit no-op coupling. The no-KV-manager case works only because 0 > 0 is False — two
methods silently conspiring. An explicit early return in _validate_request_budget is clearer and
survives future edits.
Unaccounted (all under-estimates, so safe-direction, but worth a comment): draft/spec-decode
tokens and num_extra_kv_tokens (resource_manager.py:367); the draft and cross-attention pools,
which get their own managers (_util.py:1602-1605) and aren't counted at all.
✅ Verified correct
_num_blocks_for(P, G, 1)is algebraically identical to the code it replaces:
ceil((P+G)/tpb) − P//tpb == ceil((P%tpb + G)/tpb). Theget_needed_resource_to_completion
refactor genuinely preserves behavior.- For a fresh request with no reuse and no sliding window it matches
getRemainingBlocksToCompletion(mainkvCacheManager.cpp:3689-3695):
contextBlocks + genBlocksPerBeam × beamWidth. - The shared/unshared decomposition and the
getNeededBlocksOneStepcitation are accurate. - The error message at
py_executor.py:4744-4749carries required, available, and all three
inputs — good for debuggability.
Tests
Good coverage of the pure arithmetic. The gaps track the risks:
- Neither blocker is covered — nothing exercises a V2-derived manager or a windowed/VSWA manager.
A parametrized test over the manager classes that can occupy theKV_CACHE_MANAGERslot would
have caught 1a directly. get_max_atten_window_upper_boundand_get_token_num_for_estimation(_util.py:629) are the
load-bearing changes and have no direct tests; the beam-divisor removal is untested.- Request stubs use
SimpleNamespace, so a rename oforig_prompt_len/py_beam_width/
max_new_tokensonLlmRequestfails silently.LlmRequestis already imported in this file. test_validate_request_budget_rejects_oversized_requestpasses aMagicMockasselfwith both
accessors stubbed, so it verifies only the>comparison — not the wiring through
_validate_request, and not that rejection yields a clean client-facing error.- The behavioral test promised in the description ("clean error response, not a hang") isn't
present. That's the one demonstrating the fix; worth an integration test at largebeam_width. - Move
from types import SimpleNamespaceand thePyExecutorimport to module scope.
thorjohnsen
left a comment
There was a problem hiding this comment.
The basic idea of this PR is correct and I think you are nearly there for KVCacheManager V1. However, because of inconsistencies between V1 and V2 APIs, an implementation that is functionally correct for V1 will not work without additional changes for V2. I left a more detailed comment about what those inconsistencies are. It was produced by Claude, which can hallucinate, but I looked through it and verified the major findings.
15a36e6 to
6a38c0d
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_resource_manager.py (1)
1323-1387: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd request-admission coverage for KV-cache budget validation.
_validate_requestcalls_validate_request_budgetafter sampler validation. The three tests call the helper directly. Add an_validate_requesttest with a sampler-valid oversized request.Coverage verdict: insufficient. The added tests are not listed in
test-db/orqa/. Addunittest/_torch/executor/test_resource_manager.pytotests/integration/test_lists/test-db/l0_a10.yml.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/executor/test_resource_manager.py` around lines 1323 - 1387, The added tests only exercise _validate_request_budget directly and are not registered in the test database. Add coverage that invokes PyExecutor._validate_request with a sampler-valid oversized request, while retaining the existing budget assertions, and register unittest/_torch/executor/test_resource_manager.py in tests/integration/test_lists/test-db/l0_a10.yml.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/executor/test_resource_manager.py`:
- Around line 1323-1387: The added tests only exercise _validate_request_budget
directly and are not registered in the test database. Add coverage that invokes
PyExecutor._validate_request with a sampler-valid oversized request, while
retaining the existing budget assertions, and register
unittest/_torch/executor/test_resource_manager.py in
tests/integration/test_lists/test-db/l0_a10.yml.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: d43af79b-c79e-4a54-911a-65df85b2c90c
📒 Files selected for processing (5)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/executor/test_resource_manager.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/pyexecutor/py_executor.py
|
/bot run |
|
PR_Github #65088 [ run ] triggered by Bot. Commit: |
|
PR_Github #65088 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65152 [ run ] triggered by Bot. Commit: |
|
PR_Github #65152 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
6a38c0d to
ba972d8
Compare
|
PR_Github #66884 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #67143 [ run ] triggered by Bot. Commit: |
|
PR_Github #67143 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
/bot run --disable-fail-fast |
|
PR_Github #67511 [ run ] triggered by Bot. Commit: |
|
PR_Github #67511 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68335 [ run ] triggered by Bot. Commit: |
|
PR_Github #68335 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68851 [ run ] triggered by Bot. Commit: |
…am width Torch flow derived the maximum attention window by dividing pool capacity by max_beam_width, so a large beam width shrank the sequence length the resource manager believed it could support. Startup then asserted against a configured max_seq_len it should have accepted, and lowering max_seq_len to get past that drove the deduced token budget negative at runtime. Beams share the prompt context rather than each holding a private copy, so the divisor was never the right model. Drop it and gate admission on a beam-aware per-request estimate instead. - return the full token capacity from get_max_atten_window_upper_bound and drop its now-unused max_beam_width and max_seq_len arguments - add _num_blocks_to_completion, mirroring the C++ shared/unshared cost model: full prompt blocks are shared across beams, the partial last prompt block and generation are per beam. get_needed_resource_to_completion delegates to it with beam_width=1 so the scheduler's value is unchanged - reject in PyExecutor._validate_request any request whose worst-case demand exceeds the GPU-primary capacity of the self, draft or cross pool, so it fails cleanly instead of stalling the scheduler - scope that check to the exact KVCacheManager; sparse/compressed and mamba hybrid subclasses opt in through _dense_kv_block_budget once verified - take the warm-up estimation beam factor from the configured max_beam_width rather than the dummy request's sampling config Signed-off-by: Athena Cai <athenac@nvidia.com>
8fb7c3f to
3a9c68d
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #68867 [ run ] triggered by Bot. Commit: |
|
PR_Github #68851 [ run ] completed with state |
|
PR_Github #68867 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #68962 [ run ] triggered by Bot. Commit: |
|
PR_Github #68962 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #69151 [ run ] triggered by Bot. Commit: |
3894f74 to
3a9c68d
Compare
|
PR_Github #69151 [ run ] completed with state
|
Ticket: 【TRTLLM】TRT‑LLM Torch Flow has incorrect KVcache resource estimation and sequence length calculation when using large beam_width values, causing max_seq_len–related assertion failures.
https://nvbugspro.nvidia.com/bug/5977180/6
Ticket Description:
Background:
We’re trying to benchmark TRT‑LLM Torch Flow with very large beam widths (for example, beam_width = 128 or 256), but we’ve found that the current Torch Flow implementation has issues with GPU memory estimation and sequence length estimation in this setting. Concretely, we’re seeing two problems:
During server startup, the resource manager first computes the “maximum supported sequence length” based on the configuration. As far as I can tell, this happens in the get_max_atten_window_upper_bound function, where a key step is:
max_blocks_per_seq = math.floor(token_capacity / (max_beam_width * tokens_per_block))
In other words, the current design derives the maximum number of blocks by taking the total KV‑cache token capacity and dividing it by beam_width. When beam_width is very large, this effectively caps the maximum sequence length that the resource manager believes it can support. Then, during warmup, if the max_seq_len value in my YAML config is set to a large number, it triggers an assertion because the configured value exceeds what the resource manager computed. Conceptually, this means the resource manager is using the following formula to allocate KV‑cache memory:
max_seq_len × beam × layer × dim × nhead × 2 (K+V) × 2 bytes.
However, we believe it’s more reasonable for different beams to share the prompt context length, rather than treating them as fully independent for KV‑cache sizing.
At runtime, even if I set a smaller max_seq_len in the YAML so that startup can get past the assertion, I still hit another assertion when sending a request with a long context:
default_max_tokens (-5200) must be greater than 0, default_max_tokens (-5200) = max_seq_len (4800)
This shows that the impact of issue (1) carries over into runtime; there isn’t a separate or divergent handling path at runtime that would avoid this problem.
What is happening?
get_max_atten_window_upper_boundruns once at server startup, inside the KVCacheManager constructor — long before any request arrives. The call chain:Cost model (original, pre-fix):
tensorrt_llm/_torch/pyexecutor/resource_manager.pyThis follows the cost model:
How does this get used?
get_max_atten_window_upper_boundreturns a single number — the largest attention window (in tokens) that the KV-cache pool can physically hold for one sequence. It's used as a safety clamp on the configured window / max_seq_len during KVCacheManager construction. Here's what its return value drives:tensorrt_llm/_torch/pyexecutor/resource_manager.pySolution (implemented)
The fix has three parts, all in commit
6a38c0df7a([https://nvbugs/5977180][fix] size KV cache and admit requests per beam width):max_beam_widthdivisor fromget_max_atten_window_upper_boundso the attention-window upper bound is just the raw token capacity of the pool. Beams no longer shrink the sequence length the manager believes it can support.KVCacheCreator._get_token_num_for_estimation, but take it from the configuredmax_beam_widthinstead of the dummy request's sampling config. (This differs from the earlier plan, which proposed removing the factor entirely — see "Warm-up estimation" below for why it has to stay.)PyExecutor._validate_request, compute the worst-case beam-aware block demand for the request and reject it up front if it exceeds the GPU-primary capacity of any KV pool, instead of letting the scheduler spin on a request that can never be served.Downstream effects of removing the
max_beam_widthdivisor fromget_max_atten_window_upper_boundfree_fractionandtpbtell you how many blocks exist and how big each is, but the engine still needs a fixed per-sequence block-table width and a positional upper bound to (a) allocate the fixed-shape host_kv_cache_block_offsets, (b) build the C++ pool metadata, and (c) capture CUDA graphs. Those are laid out once and can't grow per request — so max_seq_len (or an equivalent explicit cap) is unavoidable at startup. You can change how it's computed, but you can't drop it.Linking fact: max_blocks_per_seq is derived from max_seq_len
The Python manager reads it straight from the C++ impl:
resource_manager.py
Line 715
self.max_blocks_per_seq = self.impl.max_blocks_per_seq
and the C++ computes that value purely from maxSequenceLength (via maxTokensHeld):
So max_seq_len → maxBlocksPerSeq → max_blocks_per_seq. Everything below is dimensioned by it.
(a) The fixed-shape host block-offset buffer
Allocated once in the constructor, with max_blocks_per_seq as its last dimension:
This is a single persistent buffer sized at init — it cannot grow when a request arrives.
Consequence of changing max_seq_len: Not a bug — it's the intended size. A beam that reaches the full configured length genuinely needs a block table that long. The current shrunk value is what's wrong; it would be too small for a beam to ever reach the config length (the clamp just forbids that instead). No correctness issue here.
(b) The C++ pool metadata
The WindowSizeMetadata built in the C++ constructor (snippet above, kvCacheManager.cpp:730-732) stores maxBlocksPerSeq for the pool's block-table geometry. It's computed once at construction from maxSequenceLength, alongside allottedPrimaryBlocks/allottedSecondaryBlocks (which come from the free_gpu_memory_fraction-derived pool bytes and tokensPerBlock).
Consequence of changing max_seq_len:
Pool block count is unchanged. allottedPrimaryBlocks comes from free_gpu_memory_fraction-derived bytes ÷ block bytes, independent of max_seq_len.
A single beam-search request is now permitted to demand up
contextBlocks + genBlocks × beam_width blocks, which can exceed the entire pool. getRemainingBlocksToCompletion will then return more than total free blocks.This is ok because the divisor removal is paired with an explicit beam-aware admission check that rejects (errors) a request whose worst-case completion exceeds pool capacity, instead of letting the scheduler spin — implemented as
_validate_request_budget/get_request_kv_block_budget(see below).(c) CUDA-graph-captured device buffer
The device-side block-offset tensor is also shaped by max_blocks_per_seq, and is explicitly a graph-capture buffer (capture_graph=capture_graph):
And the graph runner captures the attn_metadata (which owns that buffer) and cannot reallocate inputs across replays — i.e., the shape is frozen at capture time:
Consequence of changing
max_seq_len: Capture correctness is unaffected. The captured device buffer just gets a larger (padded) last dim; generation graphs pad to max_blocks_per_seq. No shape mismatch, only the modest memory from (a).Implementation
Reused existing primitives (minimized duplication)
KVCacheManager.get_num_kv_blocks(num_tokens)— ceil-div bytokens_per_block.KVCacheManager.get_needed_resource_to_completion()— the beam-unaware (beam=1) version of the same math; refactored to share a helper rather than duplicating the ceil math.orig_prompt_len,max_new_tokens,py_beam_width,encoder_output_len.blocks_in_primary_poolattribute (set both on the estimation dry-run path and the standard path inKVCacheManager.__init__).1. Startup: beam-independent attention-window upper bound
get_max_atten_window_upper_boundis now a pure capacity query — the beam divisor, the-= tokens_per_blockbeam fudge branch, and themax_beam_width/max_seq_lenparameters are all gone (resource_manager.py:1411):_validate_and_adjust_attention_windowslost itsmax_beam_widthparameter accordingly, and theKVCacheManager.__init__call site no longer passes it. Everything else about the clamp is unchanged: windows larger than the bound are still reduced,blocks_per_windowkeys are still rewritten, andadjusted_max_seq_len = max(adjusted_window_vec)is still how the effectivemax_seq_lenis derived. The net effect is that a largemax_beam_widthno longer shrinks the startup sequence-length ceiling — which is exactly the startup assertion in the ticket.2. Warm-up estimation: beam factor kept, sourced from config
KvCacheCreator._get_token_num_for_estimation(_util.py:896) still scales the block budget by beam width, but reads it fromself._max_beam_widthrather thanself._dummy_reqs[0].sampling_config.beam_width:Why the factor stays (contrary to the earlier proposal):
_create_dummy_context_requestsbuilds its requests withbeam_width=max_beam_width(_util.py:766), so the temporary KV cache allocated for the estimation dry run genuinely has to holdmax_beam_widthbeams per dummy request. Removing the factor would make warm-up itself fail to allocate. The two forms are numerically identical today; sourcing it from the config removes the dependency on_dummy_reqsbeing populated and makes the intent explicit. The startup over-restriction is fixed in (1), not here.3. Runtime: per-pool, beam-aware admission check
Shared block-count helper (resource_manager.py:781) mirroring the C++ shared/unshared model (kvCacheManager.cpp:3557-3561): full prompt blocks are shared across beams; the partial-last-prompt block plus generation are per-beam.
get_needed_resource_to_completionnow delegates to it withbeam_width=1, so the Python scheduler keeps its historical beam-unaware value bit-for-bit while the duplicated ceil math is gone.Opt-in budget hook.
BaseResourceManager.get_request_kv_block_budget(request)returnsOptional[Tuple[int, int]]—(required_blocks, primary_capacity)— and defaults toNone(resource_manager.py:150), meaning "this manager does not support the request-level feasibility check". BecauseKVCacheManagerV2derives fromBaseResourceManagerand does not override it, the V2 path opts out automatically and the check is a no-op there.Scoped to the exact base manager.
KVCacheManager.get_request_kv_block_budget(resource_manager.py:792) returnsNonefor any subclass and otherwise delegates to_dense_kv_block_budget(resource_manager.py:808):The dense model assumes every token of the sequence occupies a block of this pool, and subclasses break that assumption in both directions: sparse/compressed KV managers (
RocketKVCacheManager,DSACacheManager) retain less than the full sequence and would be over-estimated into false rejections, while the mamba hybrids split capacity across an extra state cache that a single(required, capacity)pair cannot express. A subclass opts in deliberately, by overriding the method and delegating to_dense_kv_block_budgetonce its own cost model has been checked against it._dense_kv_block_budgetitself:get_num_kv_blocks(request.encoder_output_len)(no beam factor — encoder KV is shared).None) for anything that isn't plain full-attention self cache: non-SELFcache types,is_vswa,is_linear_attention, or any window inmax_attention_window_vecsmaller thanmax_seq_len. Those need per-pool / window-aware accounting that a single(required, capacity)pair can't express._num_blocks_to_completion(orig_prompt_len, max_new_tokens + num_extra_kv_tokens + _kv_reserve_draft_tokens, py_beam_width), so speculative-decoding reservations are included.blocks_in_primary_pool(GPU-primary only, notget_max_resource_count()/impl.max_num_blocks): secondary/offloaded blocks cannot make an otherwise-impossible request schedulable on the GPU. It is total pool capacity rather than transient free blocks, so the check rejects only requests that can never fit and does not false-reject requests that would fit once others drain.Container fan-out (resource_manager.py:2707):
ResourceManager.get_request_kv_block_budgets(request)pollsKV_CACHE_MANAGER,DRAFT_KV_CACHE_MANAGER, andCROSS_KV_CACHE_MANAGER, skipping absent managers and those returningNone, and returns a list of(ResourceManagerType, required_blocks, primary_capacity). A model with no KV manager yields[], making the check an explicit no-op.The gate (py_executor.py:4927), called at the end of
_validate_requestafter the sampler-specific validation:Because it raises from
_validate_request, the request fails cleanly with a diagnostic error instead of stalling the scheduler forever.Naming/shape differences from the earlier plan
_num_blocks_for→_num_blocks_to_completion.get_num_blocks/get_num_blocks_availableaccessor pair was replaced by a singleget_request_kv_block_budgetreturning(required, capacity), plus the plural container method — this is what makes per-pool (self/draft/cross) checking and per-manager opt-out possible.blocks_in_primary_pool, notget_max_resource_count().Nonesentinel replaces the "return 0 so the check is a no-op" convention, which is explicit rather than relying on a magic zero.Notes / scope boundaries
enable_block_reuseis off in the target config). Unlike the original plan, draft/spec extra tokens are accounted for vianum_extra_kv_tokens + _kv_reserve_draft_tokens.KVCacheManagersubclass explicitly opt out of the guard rather than being approximated.max_beam_widthargument and the-= tokens_per_blockbeam branch inget_max_atten_window_upper_bound) is done as part of change (1).Testing
tests/unittest/_torch/executor/test_resource_manager.py::TestRequestBudget(GPU-free, builds a bareKVCacheManagervia__new__):test_num_blocks_beam_one_matches_needed_resource— beam=1 equalsget_needed_resource_to_completionacross several prompt/decode lengths.test_num_blocks_shares_prompt_across_beams— beam width multiplies only generation + partial-last-prompt block, not the full prompt.test_num_blocks_block_aligned_prompt— a block-aligned prompt contributes no per-beam partial block.test_container_returns_v1_primary_pool_budget,test_container_no_kv_cache_manager_is_explicit_noop,test_v2_manager_does_not_run_v1_budget_check,test_windowed_v1_manager_does_not_use_single_pool_budget— fan-out and opt-out behavior.test_subclass_manager_opts_out_by_default— a plainKVCacheManagersubclass inherits no budget, so the container returns[].test_subclass_manager_can_opt_in— a subclass that overrides the hook and delegates to_dense_kv_block_budgetgets the dense estimate.test_attention_window_upper_bound_does_not_scale_with_beam_width— the bound is now the full token capacity.test_validate_request_budget_rejects_oversized_request/..._checks_draft_manager/..._checks_cross_manager— oversized requests raise aValueErrornaming the pool and block counts; fitting requests pass, for the self, draft, and cross pools respectively.tests/unittest/_torch/executor/test_kv_cache_estimation.py:test_max_beam_width_scales_estimation_blocks— warm-up capacity uses the configured beam width (V1 path).Dev Engineer Review
The implementation adds beam-aware KV-cache estimation and request admission checks.
Review blockers remain:
get_num_blocksis not universally implemented.get_max_resource_count()has inconsistent units and may include offloaded or secondary blocks. The implementation needs an explicit GPU-primary capacity interface with capability-based gating.LlmRequestinputs,_validate_request()end-to-end coverage, and an integration test for clean oversized-request rejection.No configuration or test-list changes are reported.
QA Engineer Review
Test changes cover:
_validate_request_budget().Modified test files:
tests/unittest/_torch/executor/test_resource_manager.pytests/unittest/_torch/executor/test_kv_cache_estimation.pytests/unittest/_torch/sampler/test_beam_search.pyNo corresponding
test-db/,qa/, ortests/integration/test_lists/coverage is reported. The verdict is needs follow-up until CI or manual-QA coverage is confirmed.