Skip to content

[https://nvbugs/5977180][fix] size KV cache and admit requests per beam width - #16802

Merged
pcastonguay merged 1 commit into
NVIDIA:mainfrom
athena-nv:trtllm-bug-5977180
Aug 25, 2026
Merged

[https://nvbugs/5977180][fix] size KV cache and admit requests per beam width#16802
pcastonguay merged 1 commit into
NVIDIA:mainfrom
athena-nv:trtllm-bug-5977180

Conversation

@athena-nv

@athena-nv athena-nv commented Jul 23, 2026

Copy link
Copy Markdown
Collaborator

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_bound runs once at server startup, inside the KVCacheManager constructor — long before any request arrives. The call chain:

KVCacheManager.__init__  (resource_manager.py:281)
   └─ _validate_and_adjust_attention_windows(...)   (line 571 → def at 2237)
        └─ get_max_atten_window_upper_bound(...)     (line 2268 → def at 1411)

Cost model (original, pre-fix):

tensorrt_llm/_torch/pyexecutor/resource_manager.py

    def get_max_atten_window_upper_bound(self, blocks_in_primary_pool,
                                         tokens_per_block, max_beam_width,
                                         max_seq_len: Optional[int]):
        token_capacity = blocks_in_primary_pool * tokens_per_block
        max_blocks_per_seq = math.floor(token_capacity /
                                        (max_beam_width * tokens_per_block))
        assert max_blocks_per_seq > 0, "Impossible to fit in any sequence in kvCache"

        max_atten_window_upper_bound = max_blocks_per_seq * tokens_per_block
        if max_seq_len is not None and max_seq_len > max_atten_window_upper_bound and max_beam_width > 1:
            max_atten_window_upper_bound -= tokens_per_block
        assert max_atten_window_upper_bound > 0, "Impossible to fit in any sequence in kvCache"
        return max_atten_window_upper_bound

This follows the cost model:

blocks_needed = ceil((P + G) / tpb) * B
# beam width B, prompt length P, and generated length G (so total sequence = P + G):

How does this get used?

get_max_atten_window_upper_bound returns 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:

  1. Setting the effective max_seq_len TensorRT-LLM pre-allocates the entire KV cache and its bookkeeping structures once at startup (that's the whole point of free_gpu_memory_fraction — grab a fixed pool up front), and several of those fixed-shape structures are dimensioned from max_seq_len. So it has to be finalized before any request arrives. — line 2297 tensorrt_llm/_torch/pyexecutor/resource_manager.py
        adjusted_max_seq_len = max(adjusted_window_vec)
        # self._max_seq_len = kv_cache_manager.max_seq_len -> base_worker.max_seq_len → _deduce_max_tokens
  1. Clamping the attention window(s) — resource_manager.py:2268-2275 Each window_size in blocks_per_window is compared against the bound; anything larger is reduced to the bound and recorded in window_adjustments. These clamped values become the actual per-layer attention window sizes (adjusted_window_vec) used by the attention kernels and the KV-cache pool layout.
  2. Rewriting the pool block map — line 2285-2295 The blocks_per_window dict keys are rewritten to the clamped window sizes so the C++ pool is built with matching window keys.

Solution (implemented)

The fix has three parts, all in commit 6a38c0df7a ([https://nvbugs/5977180][fix] size KV cache and admit requests per beam width):

  1. Startup: drop the max_beam_width divisor from get_max_atten_window_upper_bound so 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.
  2. Warm-up estimation: keep a beam-width factor in KVCacheCreator._get_token_num_for_estimation, but take it from the configured max_beam_width instead 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.)
  3. Runtime: in 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_width divisor from get_max_atten_window_upper_bound

free_fraction and tpb tell 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):

kvCacheManager.cpp
Lines 723-732
        auto const maxTokenNum = std::max(windowSize, maxSequenceLength) + sinkBubbleLength;
        // For sliding window attention, we at most hold chunk_size + window_size number of tokens.
        // For full attention, all tokens are held.
        auto const maxTokensHeld
            = windowSize < maxSequenceLength ? std::min(maxSequenceLength, maxTokenNum + chunkSize) : maxTokenNum;
        auto const maxBlocksPerSeq = tc::ceilDiv(maxTokensHeld, tokensPerBlock);
        ...
        mWindowSizeToMetadata[windowSize]
            = WindowSizeMetadata{..., maxBlocksPerSeq, ...};

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:

resource_manager.py
Lines 719-726
        self.host_kv_cache_block_offsets = torch.zeros(
            self.num_pools,
            max_batch_size * max_beam_width,
            2,
            self.max_blocks_per_seq,
            dtype=torch.int32,
            pin_memory=prefer_pinned(),
            device='cpu')

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):

trtllm.py
Lines 324-333
            self.kv_cache_block_offsets = self.get_empty(
                buffers,
                [
                    num_attention_op_pools, self.max_num_sequences, 2,
                    self.kv_cache_manager.max_blocks_per_seq
                ],
                cache_name="kv_cache_block_offsets",
                dtype=torch.int32,
                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:

cuda_graph_runner.py
Lines 468-500
        # [CUDA graph spec decode padding]
        # We pad input IDs/position IDs to the maximum draft length (token per request).
        # We're forced to do this because we cannot reallocate inputs over many graph runs.
        ...
        self.graph_metadata[key] = {
            "attn_metadata": initial_inputs["attn_metadata"],
            "spec_metadata": initial_inputs.get("spec_metadata", None),
        }

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 by tokens_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.
  • Request fields already exist: orig_prompt_len, max_new_tokens, py_beam_width, encoder_output_len.
  • Pool capacity comes from the existing blocks_in_primary_pool attribute (set both on the estimation dry-run path and the standard path in KVCacheManager.__init__).

1. Startup: beam-independent attention-window upper bound

get_max_atten_window_upper_bound is now a pure capacity query — the beam divisor, the -= tokens_per_block beam fudge branch, and the max_beam_width / max_seq_len parameters are all gone (resource_manager.py:1411):

def get_max_atten_window_upper_bound(self, blocks_in_primary_pool,
                                     tokens_per_block):
    token_capacity = blocks_in_primary_pool * tokens_per_block
    assert token_capacity > 0, "Impossible to fit in any sequence in kvCache"
    return token_capacity

_validate_and_adjust_attention_windows lost its max_beam_width parameter accordingly, and the KVCacheManager.__init__ call site no longer passes it. Everything else about the clamp is unchanged: windows larger than the bound are still reduced, blocks_per_window keys are still rewritten, and adjusted_max_seq_len = max(adjusted_window_vec) is still how the effective max_seq_len is derived. The net effect is that a large max_beam_width no 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 from self._max_beam_width rather than self._dummy_reqs[0].sampling_config.beam_width:

# Dummy context requests use the configured maximum beam width. Scale
# their block budget by the same value so the temporary KV cache used
# during warm-up can accommodate those requests.
num_cache_blocks *= self._max_beam_width

max_num_tokens_for_estimation = (num_cache_blocks * self._tokens_per_block)

Why the factor stays (contrary to the earlier proposal): _create_dummy_context_requests builds its requests with beam_width=max_beam_width (_util.py:766), so the temporary KV cache allocated for the estimation dry run genuinely has to hold max_beam_width beams 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_reqs being 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.

def _num_blocks_to_completion(self, prompt_len: int, max_new_tokens: int,
                              beam_width: int) -> int:
    shared_context_blocks = prompt_len // self.tokens_per_block
    per_beam_tokens = prompt_len % self.tokens_per_block + max_new_tokens
    per_beam_blocks = self.get_num_kv_blocks(per_beam_tokens) * beam_width
    return shared_context_blocks + per_beam_blocks

get_needed_resource_to_completion now delegates to it with beam_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) returns Optional[Tuple[int, int]](required_blocks, primary_capacity) — and defaults to None (resource_manager.py:150), meaning "this manager does not support the request-level feasibility check". Because KVCacheManagerV2 derives from BaseResourceManager and 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) returns None for any subclass and otherwise delegates to _dense_kv_block_budget (resource_manager.py:808):

if type(self) is not KVCacheManager:
    return None
return self._dense_kv_block_budget(request)

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_budget once its own cost model has been checked against it.

_dense_kv_block_budget itself:

  • Cross cache: required = get_num_kv_blocks(request.encoder_output_len) (no beam factor — encoder KV is shared).
  • Opts out (None) for anything that isn't plain full-attention self cache: non-SELF cache types, is_vswa, is_linear_attention, or any window in max_attention_window_vec smaller than max_seq_len. Those need per-pool / window-aware accounting that a single (required, capacity) pair can't express.
  • Self cache: required = _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.
  • Capacity = blocks_in_primary_pool (GPU-primary only, not get_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) polls KV_CACHE_MANAGER, DRAFT_KV_CACHE_MANAGER, and CROSS_KV_CACHE_MANAGER, skipping absent managers and those returning None, 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_request after the sampler-specific validation:

def _validate_request_budget(self, request: LlmRequest) -> None:
    for resource_type, required_blocks, primary_capacity in (
            self.resource_manager.get_request_kv_block_budgets(request)):
        if required_blocks > primary_capacity:
            raise ValueError(...)  # names the pool, both block counts,
                                   # prompt_len, max_new_tokens, beam_width

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.
  • The get_num_blocks / get_num_blocks_available accessor pair was replaced by a single get_request_kv_block_budget returning (required, capacity), plus the plural container method — this is what makes per-pool (self/draft/cross) checking and per-manager opt-out possible.
  • Capacity is blocks_in_primary_pool, not get_max_resource_count().
  • The None sentinel 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

  • Coarse by design: ignores KV block reuse (conservative overestimate; enable_block_reuse is off in the target config). Unlike the original plan, draft/spec extra tokens are accounted for via num_extra_kv_tokens + _kv_reserve_draft_tokens.
  • VSWA / linear-attention / sliding-window, KV-cache-manager V2, and every KVCacheManager subclass explicitly opt out of the guard rather than being approximated.
  • The follow-up cleanup previously listed as out of scope (vestigial max_beam_width argument and the -= tokens_per_block beam branch in get_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 bare KVCacheManager via __new__):

  • test_num_blocks_beam_one_matches_needed_resource — beam=1 equals get_needed_resource_to_completion across 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 plain KVCacheManager subclass 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_budget gets 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 a ValueError naming 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.

  • Startup and warm-up estimation use the configured beam width.
  • Request budgeting accounts for shared prompt blocks, per-beam partial and generation blocks, speculative tokens, and cross-attention encoder length.
  • Budget checks cover target, draft, and cross KV-cache pools.
  • Unsupported managers and attention types can opt out.
  • Completion estimation reuses shared block calculations.
  • Attention-window validation removes obsolete beam-width and sequence-length inputs.

Review blockers remain:

  • The budget interface is not compatible with all V2-derived managers. get_num_blocks is 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.
  • The block estimate does not apply the sliding-window clamp used by the C++ implementation. This can cause false request rejection.
  • VSWA and linear-attention managers need window-aware handling or explicit opt-out.
  • Add direct startup and warm-up tests, V2-derived and windowed-manager coverage, real LlmRequest inputs, _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:

  • Beam-aware KV-cache block arithmetic.
  • Shared prompt blocks and block-aligned prompts.
  • V1 delegation.
  • No-manager, V2, subclass, and windowed-cache behavior.
  • Attention-window bounds.
  • Target, draft, and cross KV-cache budget validation.
  • Warm-up estimation through the V1 memory-cap path.
  • Test-stub compatibility with _validate_request_budget().

Modified test files:

  • tests/unittest/_torch/executor/test_resource_manager.py
  • tests/unittest/_torch/executor/test_kv_cache_estimation.py
  • tests/unittest/_torch/sampler/test_beam_search.py

No corresponding test-db/, qa/, or tests/integration/test_lists/ coverage is reported. The verdict is needs follow-up until CI or manual-QA coverage is confirmed.

@coderabbitai

coderabbitai Bot commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 10f38171-5d5a-4447-97e5-f6f40c823afc

📥 Commits

Reviewing files that changed from the base of the PR and between 997549d and 5fdf20a.

📒 Files selected for processing (1)
  • tests/unittest/_torch/sampler/test_beam_search.py

Walkthrough

KV-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.

Changes

KV-cache budgeting and admission

Layer / File(s) Summary
Beam-aware block estimation
tensorrt_llm/_torch/pyexecutor/resource_manager.py, tensorrt_llm/_torch/pyexecutor/_util.py, tests/unittest/_torch/executor/test_kv_cache_estimation.py, tests/unittest/_torch/executor/test_resource_manager.py
KV-cache calculations share full prompt blocks and multiply partial or generated blocks per beam. V1 token estimation uses the configured maximum beam width.
Resource capacity and window integration
tensorrt_llm/_torch/pyexecutor/resource_manager.py, tests/unittest/_torch/executor/test_resource_manager.py
ResourceManager collects applicable target, draft, and cross KV-cache budgets. Attention-window bounds use full primary block capacity. Tests cover unsupported managers, windowed managers, and capacity bounds.
Request budget admission
tensorrt_llm/_torch/pyexecutor/py_executor.py, tests/unittest/_torch/executor/test_resource_manager.py, tests/unittest/_torch/sampler/test_beam_search.py
PyExecutor rejects requests when required KV blocks exceed a resource pool’s primary capacity. Tests cover fitting and oversized requests for primary, draft, and cross KV-cache managers. Beam-width validation mocks provide the new budget-validation callback.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: 🔵 Low · up to 5fdf2

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
Loading

Suggested reviewers: chienchunhung, cascade812, nv-xtf

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 39.53% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the bug fix and its focus on KV-cache sizing and request admission for beam width.
Description check ✅ Passed The description clearly explains the issue, implementation, scope boundaries, and extensive test coverage, although it omits the explicit checklist section.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 win

Attention-window clamping still needs beam-width headroom. get_max_atten_window_upper_bound() now reduces to blocks_in_primary_pool * tokens_per_block, so max_beam_width > 1 no longer shrinks the clamp. _validate_and_adjust_attention_windows() can still accept windows that exhaust KV cache during beam search; the - tokens_per_block fallback 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 win

Good coverage of the beam-sharing math and delegation; two gaps worth a follow-up.

TestRequestBudget solidly covers: beam-width=1 parity with the historical value, prompt-sharing across beams, block-aligned prompts, ResourceManager delegation/no-op, and the _validate_request_budget accept/reject boundary. Coverage is currently insufficient for:

  • get_max_atten_window_upper_bound's beam-width behavior change (no test exercises max_beam_width > 1 for that method in this file).
  • _validate_request_budget against a DRAFT_KV_CACHE_MANAGER/CROSS_KV_CACHE_MANAGER-only-constrained scenario (only KV_CACHE_MANAGER is 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 in resource_manager.py and py_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

📥 Commits

Reviewing files that changed from the base of the PR and between f4e692d and aed74f7.

📒 Files selected for processing (4)
  • tensorrt_llm/_torch/pyexecutor/_util.py
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/_torch/pyexecutor/resource_manager.py
  • tests/unittest/_torch/executor/test_resource_manager.py

Comment thread tensorrt_llm/_torch/pyexecutor/_util.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated
@thorjohnsen

Copy link
Copy Markdown
Collaborator

Thanks for the detailed write-up — the diagnosis is right, and I agree with the direction: beams
share the prompt prefix, so gating max_seq_len on capacity / beam_width was structurally wrong,
and moving feasibility to per-request admission is the correct replacement. The block math checks
out (details at the bottom). Two issues need resolving before this lands.

Line numbers are PR head aed74f74fe42 unless marked main.


🔴 Blocker 1 — the budget accessors don't hold across KV cache manager implementations

1a. get_num_blocks is missing on four of the nine possible managers

resource_manager.py:2506-2510 guards on is None, then calls a method that only exists on
KVCacheManager:

kv_cache_manager = self.get_resource_manager(ResourceManagerType.KV_CACHE_MANAGER)
if kv_cache_manager is None:
    return 0
return kv_cache_manager.get_num_blocks(request)   # AttributeError on any V2-derived manager

Every class that can occupy that slot:

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:

  • KVCacheManagerimpl.max_num_blocks — blocks ✔
  • KVCacheManagerV2 → literal 1 (kv_cache_manager_v2.py:2999-3001, # TODO: implement this)
  • DeepseekV4CacheManagerint(self.impl.get_quota(GPU_LEVEL))bytes (main
    sparse/deepseek_v4/cache_manager.py:999-1001; the same call populates
    kv_cache_stats.allocated_bytes at kv_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_blocksmAllBlocksById.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
Noneisinstance(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 beside get_num_kv_blocks(num_tokens)
    (resource_manager.py:1437) — near-identical names, different argument kinds.
    get_num_blocks_to_completion disambiguates.
  • 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). The get_needed_resource_to_completion
    refactor genuinely preserves behavior.
  • For a fresh request with no reuse and no sliding window it matches
    getRemainingBlocksToCompletion (main kvCacheManager.cpp:3689-3695):
    contextBlocks + genBlocksPerBeam × beamWidth.
  • The shared/unshared decomposition and the getNeededBlocksOneStep citation are accurate.
  • The error message at py_executor.py:4744-4749 carries 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 the KV_CACHE_MANAGER slot would
    have caught 1a directly.
  • get_max_atten_window_upper_bound and _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 of orig_prompt_len / py_beam_width /
    max_new_tokens on LlmRequest fails silently. LlmRequest is already imported in this file.
  • test_validate_request_budget_rejects_oversized_request passes a MagicMock as self with 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 large beam_width.
  • Move from types import SimpleNamespace and the PyExecutor import to module scope.

@thorjohnsen thorjohnsen left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@athena-nv
athena-nv force-pushed the trtllm-bug-5977180 branch 2 times, most recently from 15a36e6 to 6a38c0d Compare August 7, 2026 22:07
@coderabbitai

coderabbitai Bot commented Aug 7, 2026

Copy link
Copy Markdown
Contributor

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.

@athena-nv athena-nv changed the title [DRAFT][feat][5977180] runtime evaluation of blocks required by requests using beam search decode [https://nvbugs/5977180][fix] size KV cache and admit requests per beam width Aug 7, 2026

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_resource_manager.py (1)

1323-1387: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add request-admission coverage for KV-cache budget validation.

_validate_request calls _validate_request_budget after sampler validation. The three tests call the helper directly. Add an _validate_request test with a sampler-valid oversized request.

Coverage verdict: insufficient. The added tests are not listed in test-db/ or qa/. Add unittest/_torch/executor/test_resource_manager.py to tests/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

📥 Commits

Reviewing files that changed from the base of the PR and between 4b69c59 and 6a38c0d.

📒 Files selected for processing (5)
  • tensorrt_llm/_torch/pyexecutor/_util.py
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/_torch/pyexecutor/resource_manager.py
  • tests/unittest/_torch/executor/test_kv_cache_estimation.py
  • tests/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

@athena-nv
athena-nv requested a review from thorjohnsen August 10, 2026 17:27
@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65088 [ run ] triggered by Bot. Commit: 6a38c0d Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65088 [ run ] completed with state SUCCESS. Commit: 6a38c0d
/LLM/main/L0_MergeRequest_PR pipeline #52892 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65152 [ run ] triggered by Bot. Commit: 6a38c0d Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #65152 [ run ] completed with state FAILURE. Commit: 6a38c0d
/LLM/main/L0_MergeRequest_PR pipeline #52946 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Ask a member of NVIDIA/trt-llm-ci-approvers to add the label, then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #66884 [ run ] completed with state FAILURE. Commit: 22e5df2
/LLM/main/L0_MergeRequest_PR pipeline #54441 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Tabrizian

Copy link
Copy Markdown
Member

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #67143 [ run ] triggered by Bot. Commit: b60a5d5 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #67143 [ run ] completed with state FAILURE. Commit: b60a5d5
/LLM/main/L0_MergeRequest_PR pipeline #54677 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #67511 [ run ] triggered by Bot. Commit: 8fb7c3f Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #67511 [ run ] completed with state FAILURE. Commit: 8fb7c3f
/LLM/main/L0_MergeRequest_PR pipeline #55008 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68335 [ run ] triggered by Bot. Commit: 8fb7c3f Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68335 [ run ] completed with state SUCCESS. Commit: 8fb7c3f
/LLM/main/L0_MergeRequest_PR pipeline #55776 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Tabrizian

Copy link
Copy Markdown
Member

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68851 [ run ] triggered by Bot. Commit: 8fb7c3f Link to invocation

…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>
@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68867 [ run ] triggered by Bot. Commit: 3a9c68d Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68851 [ run ] completed with state ABORTED. Commit: 8fb7c3f

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68867 [ run ] completed with state FAILURE. Commit: 3a9c68d
/LLM/main/L0_MergeRequest_PR pipeline #56254 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68962 [ run ] triggered by Bot. Commit: 3a9c68d Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #68962 [ run ] completed with state SUCCESS. Commit: 3a9c68d
/LLM/main/L0_MergeRequest_PR pipeline #56343 completed with status: 'SUCCESS'

CI Report

Link to invocation

@athena-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #69151 [ run ] triggered by Bot. Commit: 3894f74 Link to invocation

@athena-nv athena-nv closed this Aug 25, 2026
@athena-nv athena-nv reopened this Aug 25, 2026
@pcastonguay
pcastonguay merged commit 675e17d into NVIDIA:main Aug 25, 2026
19 checks passed
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #69151 [ run ] completed with state FAILURE. Commit: 3894f74
/LLM/main/L0_MergeRequest_PR pipeline #56520 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.