[https://nvbugs/6567554][fix] Mirror KVCacheManagerV2.max_blocks_per_seq arithmetic in the harness budget… - #17373
Conversation
|
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 (3)
🚧 Files skipped from review as they are similar to previous changes (3)
WalkthroughThe benchmark CLI now supports automatic MoE prefill backend selection. Attention metadata uses attention-DP sizing. KV-cache capacity and dummy-request admission checks now account for benchmark runtime requirements. ChangesLayer-wise benchmark updates
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@examples/layer_wise_benchmarks/run.py`:
- Line 57: Update the post-argument-parsing logic around the assignment to
args.moe_backend_for_prefill so the derived CUTLASS value is applied only when
the parsed prefill backend is None. Preserve explicitly selected DEEPGEMM,
WIDEEP, or other valid choices, and keep the argument parser default as None.
In `@tensorrt_llm/tools/layer_wise_benchmarks/runner.py`:
- Around line 797-808: Remove the max_util_for_resume division from the
blocks_per_seq calculation in the benchmark setup, keeping the KVCacheManagerV2
quota derivation unchanged. Add a regression test that registers max_batch_size
dummy requests and verifies the resulting cache allocation fits the intended
quota without duplicate scaling.
🪄 Autofix
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: 7361dac3-863f-49f6-8f80-8fc1a8bfc607
📒 Files selected for processing (3)
examples/layer_wise_benchmarks/README.mdexamples/layer_wise_benchmarks/run.pytensorrt_llm/tools/layer_wise_benchmarks/runner.py
| kv_cache_config = KvCacheConfig(enable_block_reuse=False) | ||
| # `max_tokens` must cover what the manager reserves per request, not just the | ||
| # rounded sequence length, or the pools cannot hold `max_batch_size` requests. | ||
| # Mirror `KVCacheManagerV2.max_blocks_per_seq` -- room for one extra decode | ||
| # token, padded to a multiple of 4 blocks by the copy_block_offsets kernel -- | ||
| # then scale by 1/max_util_for_resume, since only that fraction of the derived | ||
| # quota is resumable. | ||
| blocks_per_seq = math.ceil( | ||
| round_up(ceil_div(max_seq_len + 1, tokens_per_block), 4) | ||
| / kv_cache_config.max_util_for_resume | ||
| ) | ||
| kv_cache_config.max_tokens = max_batch_size * blocks_per_seq * tokens_per_block |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -t f -g 'llm_args.py' -g 'kv_cache_manager_v2.py' tensorrt_llm \
-x rg -n -C 8 \
'max_util_for_resume|max_blocks_per_seq|_get_quota_from_max_tokens|_gpu_max_tokens' {}Repository: NVIDIA/TensorRT-LLM
Length of output: 931
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- candidate files ---'
fd -t f . tensorrt_llm | rg '(^|/)(llm_args\.py|.*kv.*cache.*manager.*\.py|runner\.py)$'
printf '%s\n' '--- relevant definitions and call sites ---'
rg -n -C 12 \
'max_util_for_resume|max_blocks_per_seq|_get_quota_from_max_tokens|_gpu_max_tokens|create_kv_cache_manager|add_dummy_requests' \
tensorrt_llm/tools/layer_wise_benchmarks/runner.py \
tensorrt_llm/llmapi \
tensorrt_llm | head -n 1200Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- KvCacheConfig definition and runner imports ---'
rg -n -C 20 'class KvCacheConfig|KvCacheConfig' tensorrt_llm/llmapi/llm_args.py tensorrt_llm/tools/layer_wise_benchmarks/runner.py
printf '%s\n' '--- V2 quota and block calculations ---'
rg -n -C 18 \
'max_blocks_per_seq|_get_quota_from_max_tokens|gpu_max_tokens|max_util_for_resume|max_tokens' \
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py \
tensorrt_llm/runtime/kv_cache_manager_v2 \
tensorrt_llm/_torch/pyexecutor/_util.py \
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py \
| head -n 700
printf '%s\n' '--- manager construction arguments ---'
sed -n '805,915p' tensorrt_llm/tools/layer_wise_benchmarks/runner.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- KvCacheConfig fields ---'
sed -n '3757,3940p' tensorrt_llm/llmapi/llm_args.py | cat -n
printf '%s\n' '--- V2 initialization and quota calculation ---'
sed -n '900,985p' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py | cat -n
printf '%s\n' '--- V2 block-count properties and max-token conversion ---'
rg -n -C 12 \
'def max_blocks_per_seq|`@property`.*max_blocks_per_seq|def _get_quota_from_max_tokens|_get_quota_from_max_tokens' \
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
printf '%s\n' '--- all max_util_for_resume consumers ---'
rg -n -C 6 'max_util_for_resume' tensorrt_llm/_torch/pyexecutor tensorrt_llm/runtime/kv_cache_manager_v2 tensorrt_llm/llmapi/llm_args.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 41690
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- exact max_blocks_per_seq definitions and assignments ---'
rg -n -C 10 \
'max_blocks_per_seq|blocks_per_seq|copy_block_offsets' \
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py \
tensorrt_llm/_torch/pyexecutor \
tensorrt_llm/runtime/kv_cache_manager_v2 \
| head -n 500
printf '%s\n' '--- max-token conversion body ---'
sed -n '1440,1495p' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py | cat -n
printf '%s\n' '--- deterministic formula probe ---'
python3 - <<'PY'
import math
max_seq_len = 4096
tokens_per_block = 32
max_batch_size = 8
util = 0.95
base_blocks = math.ceil((max_seq_len + 1) / tokens_per_block)
aligned_blocks = math.ceil(base_blocks / 4) * 4
runner_blocks = math.ceil(aligned_blocks / util)
print({
"base_blocks": base_blocks,
"aligned_blocks": aligned_blocks,
"runner_blocks": runner_blocks,
"runner_max_tokens": max_batch_size * runner_blocks * tokens_per_block,
"manager_effective_quota_in_base_token_units":
max_batch_size * runner_blocks / util,
"single_util_application_target":
max_batch_size * aligned_blocks / util,
"double_util_application_ratio":
1 / (util * util),
})
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50375
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
import math
import re
runner = Path("tensorrt_llm/tools/layer_wise_benchmarks/runner.py").read_text()
manager = Path("tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py").read_text()
args = Path("tensorrt_llm/llmapi/llm_args.py").read_text()
assert re.search(r"class KvCacheConfig\b[\s\S]*?max_util_for_resume:\s*float\s*=\s*Field", args)
assert "kv_cache_config.max_util_for_resume" in manager
assert re.search(
r"_get_quota_from_max_tokens\(kv_cache_config\.max_tokens\)\s*/\s*max_util_for_resume",
manager,
)
assert "round_up(ceil_div(max_seq_len + 1, tokens_per_block), 4)" in runner
assert " / kv_cache_config.max_util_for_resume" in runner
util = 0.95
for max_seq_len, tokens_per_block, max_batch_size in (
(4096, 32, 8),
(4097, 32, 8),
(8192, 64, 16),
):
base = math.ceil((max_seq_len + 1) / tokens_per_block)
aligned = math.ceil(base / 4) * 4
runner_blocks = math.ceil(aligned / util)
runner_tokens = max_batch_size * runner_blocks * tokens_per_block
manager_scale = 1 / util
intended_tokens = max_batch_size * aligned * tokens_per_block
print({
"inputs": (max_seq_len, tokens_per_block, max_batch_size),
"aligned_blocks": aligned,
"runner_blocks": runner_blocks,
"runner_tokens": runner_tokens,
"manager_division": manager_scale,
"effective_scale_vs_aligned_capacity":
runner_blocks / aligned / util,
"intended_tokens_before_manager_division": intended_tokens,
})
PY
printf '%s\n' '--- V2 reserve-token initialization ---'
sed -n '840,910p' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py | cat -n
printf '%s\n' '--- V2 dummy-request implementation ---'
rg -n -C 16 'def add_dummy_requests' tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 6221
Remove the duplicate max_util_for_resume scaling.
KVCacheManagerV2 already divides the quota derived from max_tokens by max_util_for_resume. Remove the division at Line 806. Otherwise, the cache is over-allocated by approximately 10.8% at the default value of 0.95, which can increase GPU memory pressure. Add a regression test that registers max_batch_size dummy requests.
🤖 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/tools/layer_wise_benchmarks/runner.py` around lines 797 - 808,
Remove the max_util_for_resume division from the blocks_per_seq calculation in
the benchmark setup, keeping the KVCacheManagerV2 quota derivation unchanged.
Add a regression test that registers max_batch_size dummy requests and verifies
the resulting cache allocation fits the intended quota without duplicate
scaling.
Source: MCP tools
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
The three fixes look right and the inline rationale is clear. Two things before merge:
No guard against the same failure recurring. The budget fix removes today's trigger, but add_dummy_requests still returns None on any partial-allocation failure and [runner.py:910](https://github.com/NVIDIA/TensorRT-LLM/pull/17373/files#diff-db99631da16085596489c77e071e5505b83ed2f768b087d0681367aa5951b862R910) still discards it, so the next time the budget is off (different tokens_per_block, a spec-dec config, a new cache manager) the symptom is again a C++ IndexMapper assertion far from the cause rather than a Python error at the allocation site. A one-line assert on the return value is cheap and is what makes this class of bug self-diagnosing.
No regression test. The GEN + sparse-attention combination that failed here isn't in tests/unittest/tools/test_layer_wise_benchmarks.py (the DSA case is CTX-only), and the world_size→dp_size fix is invisible to CI because only world_size=1 params are scheduled. Adding a GEN case with tokens_per_block=128 would cover both.
Also worth calling out in the description: sizing all_rank_num_tokens by dp_size changes MoE chunking/communication selection for the existing ADP-off (TEP) benchmarks, so their measured numbers will shift. That's a correction, not a regression, but people tracking those numbers should know.
| # quota is resumable. | ||
| blocks_per_seq = math.ceil( | ||
| round_up(ceil_div(max_seq_len + 1, tokens_per_block), 4) | ||
| / kv_cache_config.max_util_for_resume |
There was a problem hiding this comment.
The / max_util_for_resume looks like double-counting. KVCacheManagerV2 already applies it when turning max_tokens into a quota:
# kv_cache_manager_v2.py:951-957
quota_from_max_tokens = int(math.ceil(
self._get_quota_from_max_tokens(kv_cache_config.max_tokens) / max_util_for_resume))so the resumable capacity derived from max_tokens is already max_tokens worth of tokens; pre-dividing here inflates the allocation another ~5%. Harmless for a benchmark, but the comment asserts a rationale the manager handles, which will mislead the next reader.
Related: the comment says it mirrors max_blocks_per_seq, but the mirror drops num_extra_kv_tokens and _kv_reserve_draft_tokens from that expression. That's fine today only because this call site passes spec_config=None; worth saying so explicitly, since the harness does run MTP-shaped GEN batches (seq_len_q > 1).
| / kv_cache_config.max_util_for_resume | ||
| ) | ||
| kv_cache_config.max_tokens = max_batch_size * blocks_per_seq * tokens_per_block | ||
| kv_cache_manager_cls = get_kv_cache_manager_cls(model_config, kv_cache_config) |
There was a problem hiding this comment.
Consider asserting the add_dummy_requests result at [runner.py:910](https://github.com/NVIDIA/TensorRT-LLM/pull/17373/files#diff-db99631da16085596489c77e071e5505b83ed2f768b087d0681367aa5951b862R910):
assert kv_cache_manager.add_dummy_requests(...) is not None, (
f"KV cache budget too small: {max_batch_size} requests x {max_seq_len} tokens "
f"do not fit in max_tokens={kv_cache_config.max_tokens}")It returns None after releasing the partially built batch, which is exactly how this bug turned into an IndexMapper assertion two call layers away. Without the check, the next budget mismatch produces the same undiagnosable failure.
| args.moe_backend_for_prefill = "CUTLASS" | ||
| # Let prefill resolve like decode when decode is AUTO; the "CUTLASS" default cannot | ||
| # serve every checkpoint (see README limitations). | ||
| args.moe_backend_for_prefill = "AUTO" if args.moe_backend == "AUTO" else "CUTLASS" |
There was a problem hiding this comment.
This only propagates the AUTO case. --moe-backend DEEPGEMM (or WIDEEP) alone still leaves prefill on CUTLASS, which is the exact FP8-on-Blackwell failure the README limitation describes — so the documented workaround still requires passing both flags. Since --moe-backend-for-prefill already restricts choices, defaulting to args.moe_backend when it's one of the supported prefill backends, and CUTLASS otherwise, would cover all of them:
PREFILL_BACKENDS = {"AUTO", "CUTLASS", "DEEPGEMM", "WIDEEP"}
args.moe_backend_for_prefill = (
args.moe_backend if args.moe_backend in PREFILL_BACKENDS else "CUTLASS")If mirroring more broadly is deliberately out of scope, say so in the comment.
| max_num_tokens=batch_size * seq_len_q, | ||
| # One entry per attention-DP rank, not per world rank: with attention DP | ||
| # off `dp_size` is 1 and the MoE non-DP path requires a single-element list. | ||
| all_rank_num_tokens=[batch_size * seq_len_q] * self.model_config.mapping.dp_size, |
There was a problem hiding this comment.
Matches _get_all_rank_num_tokens (model_engine.py:3033), which allgathers only under attention DP and returns None otherwise — and moe_scheduler.py:143 expands None to [local_num_tokens], so the 1-element list is equivalent. Good catch that moving this into the constructor also lets __post_init__ see it.
One gap: the engine uses tp_cp_allgather under CP helix, giving a list of length tp*cp, which dp_size won't reproduce. The harness hardcodes cp_size=1 in create_mapping, so it can't hit that today — a note to that effect would keep a future CP-enabled harness from inheriting a subtly wrong list.
The GEN phase aborted in attn_metadata.prepare() with a C++ assertion "Request ID not found in IndexMapper". Three independent defects in the layer-wise benchmark harness contributed: 1. create_kv_cache_manager budgeted max_tokens from round_up(max_seq_len, tokens_per_block), which under-counts what KVCacheManagerV2 reserves per request: max_blocks_per_seq leaves room for one extra decode token and pads to a multiple of 4 blocks for the copy_block_offsets kernel, and the derived quota is divided by max_util_for_resume. For the reported shape that is 13 blocks/request against 9 budgeted, so add_dummy_requests could not place the whole batch. Its failure path releases every already-registered request and returns None, leaving all request IDs unregistered, which surfaced two layers later as the IndexMapper assertion. Mirror the manager's arithmetic instead. 2. run.py defaulted --moe-backend-for-prefill to the literal "CUTLASS", so --moe-backend AUTO resolved a backend for decode but never for prefill. On an FP8_BLOCK_SCALES checkpoint CUTLASS dispatches to a Hopper-only kernel and prefill aborted on SM100. Default prefill to AUTO when decode is AUTO. 3. create_run_pack sized all_rank_num_tokens by world_size, but the field holds one entry per attention-DP rank. With attention DP disabled the MoE non-DP path requires a single-element list, so any world_size > 1 run tripped an assertion in calculate_num_chunks. Use mapping.dp_size. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
d0b7c2c to
43be01b
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. |
3 participants
Summary
max_tokensmadeadd_dummy_requestsrelease the batch and returnNone(discarded, leaving all IDs unregistered); a hardcoded"CUTLASS"prefill MoE backend bypassedAUTOresolution; andall_rank_num_tokenswas sized byworld_sizeinstead ofdp_size.KVCacheManagerV2.max_blocks_per_seqarithmetic in the harness budget, default prefill toAUTOwhen decode isAUTO, and sizeall_rank_num_tokensbymapping.dp_size.python .repair-bot/repro.pyTest plan
Links
Dev Engineer Review
max_tokenssizing to matchKVCacheManagerV2.max_blocks_per_seqarithmetic.AUTOwhen the decode backend usesAUTO.all_rank_num_tokenswithmapping.dp_sizeinstead ofworld_size.AUTO.QA Engineer Review
No test changes.