[None][feat] DSv4 follow-up: runtime KV and cache foundations

[jiaganc]

Summary by CodeRabbit

• New Features

  • Added richer KV-cache and performance stats, including peak and evictable counts plus host-drop tracking.
  • Added support for keeping all buffered iteration/request stats with a new “unbounded” setting.
  • Introduced optional SWA scratch reuse and expanded KV-cache configuration options.
  • Added request input-token count access and broader Python serialization support.

• Bug Fixes

  • Improved handling of disaggregated and KV-cache lifecycle flows, including more reliable scheduling, resizing, and transfer checks.
  • Fixed concurrent metrics collection and made server iteration-stat handling safer.

Description

Merge the DeepSeek V4 runtime/KV-cache foundation and DSv4-specific cache
manager work back onto current main.

The PR1 prefix provides the generic foundation:

• KVCacheManagerV2 integration in its current
  kv_cache_manager_v2.py owner, including config-controlled SWA scratch
  reuse and cache-size/statistics plumbing;
• iteration and request performance statistics across executor, scheduler,
  OpenAI serving, and disaggregated serving paths;
• reliable asynchronous iteration-stat draining, lazy buffer creation, and
  benchmark stall diagnostics;
• executor/disaggregation lifecycle hooks, including context-resource updates
  after input-copy synchronization.

The DSv4-CM stack starts after dependency boundary af57e65d8f and applies
nine source commits in feature-branch order:

a7728aad24  Optimize DSV4 block table copy paths
7a4b113c36  Fix MTP by scratch reuse rewind
69e7acc8e3  Add DSV4 KV cache pool ratio config
be4aa75e85  Consolidate disagg gen-init KV setup
1876a39e6e  Cache scratch memory optimization for DSv4
4f9d819a4e  Fix KV cache manager v2 size estimation
c185066bc2  Fix KV cache size estimation limits
0e22041767  Add KV cache block reuse policy
d15f21c8d3  Fix V2 context finalization ordering

This adds DSv4 role-specific block-table/page mapping, compiled scratch-copy
paths, MTP rewind, pool-ratio sizing, V2 estimation, block-reuse policy, and
disaggregated context-finalization ordering. Source V2 changes were ported from
the former resource_manager.py location into kv_cache_manager_v2.py.
Where the first DSv4 commit overlapped generic scratch integration, the current
main implementation from 82f180443d remains authoritative and DSv4 behavior
is connected through the V2 page-table extension hook.

The final branch was compared file by file with
upstream/feat/deepseek_v4@af42658ecb. Every selected source hunk is present,
already equivalent on the dependency base, or documented as a latest-main API
adaptation or moved owner. Feature-only Router/OpenAI policy, sparse custom-op,
MoE/CuTeDSL, and generated FMHA changes remain out of scope.

Test Coverage

A clean build ran on B300 host umb-b300-dp-186 at final SHA fadeb20ae3:

/home/jiaganc/dotfiles/agents/skills/trtllm-docker-compile/scripts/compile.sh \
  --ssh-host umb-b300-dp-186 \
  --trtllm-repo /home/scratch.jiaganc_gpu/worktree/main-dsv4-cache-manager-mergeback-round2 \
  --cuda-architectures 103-real \
  --clean

Result: all 4,543 build targets passed, including bindings and raw-reference
generation.

Changed DSv4 cache-manager, compressor, executor, scheduler, disaggregation,
transfer, V2 runtime, cache-statistics, and LLM-args unit tests then ran in the
worktree-mounted development container:

DSv4 compressor/cache manager       149 passed
Executor/disaggregation/stats       323 passed
Transfer tests                      353 passed, 66 skipped
Standalone KV cache manager V2       63 passed, 12 skipped
LLM args                            171 passed, 3 skipped, 2 failed
Total                              1059 passed, 81 skipped, 2 failed

The two failures are
TestTrtLlmArgs::test_build_config_from_engine and
TestTorchLlmArgs::test_runtime_sizes. Both workers fail while importing the
container package with
ImportError: cannot import name 'bitwidth' from 'triton_kernels.tensor'.
They reproduced before and after the clean build and are an environment package
mismatch rather than failures in the changed code.

Earlier PR1-focused remote validation also covered the disaggregated perf
metrics collector and OpenAI iteration-stat tests; that batch completed with
283 passed, 3 skipped before the final autosquash.

Explicitly skipped by request:

• tests/unittest/bindings/test_executor_bindings.py
• tests/unittest/_torch/modeling/test_modeling_deepseekv4.py
• tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py
• tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py
• tests/unittest/_torch/speculative/test_eagle3.py

Changed accuracy integration tests were not run because their model data and
full integration environment were unavailable in the unit-test container.

Full-range pre-commit, git diff --check, strict conflict-marker search,
source-order/trailer audit, and forbidden-scope checks passed. The final branch
tree is byte-identical to the tree used for the clean build and final tests.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR updates KV-cache stats, executor stats limits and buffering, request token transport, SWA scratch reuse, disaggregation mapping, and serving telemetry.

Changes

KV-cache Stats Reporting

Layer / File(s)	Summary
KV stats schema and exports cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h, cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp, tensorrt_llm/runtime/kv_cache_manager_v2/{__init__.py,__init__.pyi,_core/__init__.py,_stats.py}, tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py, tests/unittest/executor/test_stats_serializer.py	KvCacheIterationStats adds evictable and peak counters, the runtime exports PoolGroupPeakBlockStats, and serializers/tests include the new iteration keys.
KV peak and drop tracking tensorrt_llm/runtime/kv_cache_manager_v2/{_core/_kv_cache.py,_core/_kv_cache_manager.py,_page.py,_storage_manager.py}, tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py	KVCacheManagerV2 tracks per-interval peak blocks, records host-tier drops through allocation and eviction paths, and exposes resettable peak stats.
KV stats behavior tests tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py	Context-allocation stats expectations are updated to match the new peak/drop and disaggregated-init behavior.

Executor Stats Limits and Buffering

Layer / File(s)	Summary
Stats limit contract cpp/include/tensorrt_llm/executor/executor.h, cpp/tensorrt_llm/executor/executorConfig.cpp, tensorrt_llm/llmapi/llm_args.py, cpp/tests/unit_tests/executor/executorConfigTest.cpp, tests/unittest/bindings/test_executor_bindings.py	ExecutorConfig and the Python args accept -1 as the unlimited stats sentinel, while 0 still disables stats collection.
C++ buffering semantics cpp/tensorrt_llm/executor/executorImpl.h, cpp/tensorrt_llm/executor/executorImpl.cpp	The C++ executor stores stats caps as SizeType32 and uses enabled/bounded helpers for trimming and collection gating.
Python executor buffering and scheduling tensorrt_llm/_torch/pyexecutor/py_executor.py, tensorrt_llm/_torch/pyexecutor/request_utils.py, tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py, tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py	Request counts are broadcast separately, disaggregated fill stalls use a timeout, and disagg init passes history_length into resize_context.

Request Token Serialization

Layer / File(s)	Summary
Lazy prompt-token state tensorrt_llm/executor/request.py	GenerationRequest lazily materializes prompt token ids from an int32 byte buffer and preserves that encoding across pickling.
Nanobind request packing cpp/tensorrt_llm/nanobind/executor/request.cpp, tensorrt_llm/executor/base_worker.py	The request binding accepts contiguous int32 ndarrays or lists, serializes the token buffer as raw bytes, and BaseWorker reuses the packed prompt-token buffer.

SWA Scratch Reuse and DeepSeek V4

Layer / File(s)	Summary
SWA config and pool setup tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/_torch/models/modeling_deepseekv4.py, tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py	KvCacheConfig adds enable_swa_scratch_reuse, and KVCacheManagerV2 builds the matching scratch-reuse config and attention-op pool layout.
Scratch offset preparation tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py, tensorrt_llm/_torch/attention_backend/{interface.py,trtllm.py,sparse/deepseek_v4/cache_manager.py,sparse/deepseek_v4/compressor.py,sparse/deepseek_v4/deepseek_v4.py}	SWA block-offset helpers add scratch metadata staging, shared-page addressing, and updated sparse attention block-table wiring for the reuse path.
Lifecycle cleanup and validation tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py, tests/unittest/llmapi/test_llm_args.py	Context prep and cleanup disable reuse, release paths clear base page-index buffers, and the config test covers the new flag default.

Disaggregation Mapping and Transceiver

Layer / File(s)	Summary
Pool and mapper metadata tensorrt_llm/_torch/disaggregation/base/region.py, tensorrt_llm/_torch/disaggregation/resource/{page.py,utils.py,kv_extractor.py}, tensorrt_llm/_torch/disaggregation/native/{peer.py,mixers/attention/peer.py}, tensorrt_llm/_torch/disaggregation/resource/cache_reuse.py	Disaggregation pool views now carry pool_role and mapper_kind, and the peer matching and KV-mapper construction logic use the new metadata.
Transfer-complete checks tensorrt_llm/_torch/disaggregation/transceiver.py	The transfer-complete boundary now asserts declared prompt history instead of trimming KV state.

Serving Telemetry

Layer / File(s)	Summary
Tracing helpers tensorrt_llm/_torch/pyexecutor/trace_log_utils.py, tensorrt_llm/_torch/attention_backend/interface.py	Trace helpers log memory and tensor sizes when profiling is enabled, and RoPE table creation logs the new tensor size.
Warmup snapshots tensorrt_llm/_torch/pyexecutor/model_engine.py	Warmup logs memory snapshots at multiple milestones, and _prepare_tp_inputs() records cached KV token counts in iteration state.
OpenAI iteration stats tensorrt_llm/serve/openai_server.py, tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py	OpenAIServer drains iteration stats under a lock, and tests cover buffered and unbuffered endpoint behavior.
Perf metrics locking tensorrt_llm/serve/perf_metrics.py, tests/unittest/llmapi/apps/test_disagg_perf_metrics_collector.py	DisaggPerfMetricsCollector splits collection and aggregation across locks, and the concurrency test asserts serialized client collection.

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#12413: Expands the same KV-cache iteration-stat schema and serialization paths.
• NVIDIA/TensorRT-LLM#15378: Touches the same KV-cache V2 stats plumbing and peak-stat reporting code paths.

Suggested labels

api-compatible

Suggested reviewers

• schetlur-nv
• yizhang-nv
• hchings
• lowsfer
• Superjomn
• chang-l

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.85% 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 is concise and clearly points to the DSv4 runtime KV/cache foundation follow-up.
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.
Description check	✅ Passed	The PR description includes the required Description, Test Coverage, and PR Checklist sections and is sufficiently detailed.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 9

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

tensorrt_llm/serve/openai_server.py (1)

940-940: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Return an empty metrics payload for VisualGen before draining stats.

register_visual_gen_routes() exposes /metrics, and _get_iteration_stats_buffer_maxlen() handles VisualGen, but get_iteration_stats() still eventually calls self.generator.get_stats_async(...). The provided VisualGen class does not expose that method, so VisualGen /metrics will 500.

Suggested change

 async def get_iteration_stats(self) -> JSONResponse:
+        if isinstance(self.generator, VisualGen):
+            return JSONResponse(content=[])
+
         # When the background collector loop is active it is the sole

🤖 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/serve/openai_server.py` at line 940, The get_iteration_stats
path still falls through to self.generator.get_stats_async even for VisualGen,
which causes /metrics to fail because VisualGen lacks that method. Update
get_iteration_stats in openai_server.py to detect the VisualGen generator type,
return an empty metrics payload for that case, and only drain or fetch stats for
generators that actually implement get_stats_async; use the existing VisualGen
handling in register_visual_gen_routes and _get_iteration_stats_buffer_maxlen as
the reference points.

tensorrt_llm/_torch/pyexecutor/py_executor.py (1)

1815-1839: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Preserve the max_stats_len == 0 disabled contract and exact cap.

max_stats_len == 0 still appends one stats record in the legacy path and one per-rank iteration in the ADP path. For positive caps, the legacy path also trims before append with >, so it retains max_stats_len + 1 records. This breaks the stated stats limit contract.

Proposed fix

             if self.dist.tp_rank == 0:
                 with self.stats_lock:
+                    if self.max_stats_len == 0:
+                        return
                     # Wrap as ("per_rank_dict", dict) so the serializer can
                     # distinguish from the legacy (stats, req_stats, kv) tuple.
@@
         #   [7] gpu_forward_time_ms: Optional[float]
         with self.stats_lock:
-            if (not _stats_buffer_is_unbounded(self.max_stats_len)
-                    and len(self.stats) > self.max_stats_len):
-                self.stats.pop(0)
+            if self.max_stats_len == 0:
+                return
             self.stats.append(
                 (stats, req_stats, kv_iter_stats, attention_dp_rank,
                  host_step_time_ms, prev_device_step_time_ms, scheduler_mode,
                  gpu_forward_time_ms))
+            if not _stats_buffer_is_unbounded(self.max_stats_len):
+                overflow = max(0, len(self.stats) - self.max_stats_len)
+                if overflow:
+                    del self.stats[:overflow]

🤖 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/py_executor.py` around lines 1815 - 1839, The
stats buffering logic in py_executor.py is violating the max_stats_len contract
in both the legacy path and the per-rank iteration path. Update the checks
around self.stats handling so max_stats_len == 0 truly disables storage and no
records are appended, and for positive limits enforce an exact cap before
appending (use the existing _stats_buffer_is_unbounded helper, the
self.stats_lock section, and the gathered/per-rank append loop as the key places
to adjust). Make the legacy trim condition and the ADP overflow trimming
consistent so the buffer never exceeds max_stats_len.

🧹 Nitpick comments (2)

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py (1)

463-562: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Consolidate the duplicated PoolGroupPeakBlockStats-from-statistics construction.

The same make_typed(lambda pool_group_index: PoolGroupPeakBlockStats(available=..., unavailable=..., evictable=...), self._storage.num_pool_groups) block is repeated three times: in _current_block_stats_by_cache_level.collect (469-477), in the non-None branch of _reset_iteration_peak_num_blocks (487-495), and in get_and_reset_iteration_peak_block_stats (546-559). Extracting a single helper that maps a cache_level to its per-pool-group PoolGroupPeakBlockStats list would remove the duplication and keep the three call sites in sync if the stat fields ever change.

♻️ Sketch

+    def _pool_group_block_stats(
+        self, cache_level: CacheLevel
+    ) -> TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats]:
+        stats_by_pool_group = self._storage.get_statistics(cache_level)
+        return make_typed(
+            lambda pool_group_index: PoolGroupPeakBlockStats(
+                available=stats_by_pool_group[pool_group_index].available,
+                unavailable=stats_by_pool_group[pool_group_index].unavailable,
+                evictable=stats_by_pool_group[pool_group_index].evictable,
+            ),
+            self._storage.num_pool_groups,
+        )

_current_block_stats_by_cache_level, _reset_iteration_peak_num_blocks, and get_and_reset_iteration_peak_block_stats can then reuse it.

As per coding guidelines (DRY in CODING_GUIDELINES.md).

🤖 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/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py` around
lines 463 - 562, The same PoolGroupPeakBlockStats construction from
get_statistics is duplicated in _current_block_stats_by_cache_level,
_reset_iteration_peak_num_blocks, and get_and_reset_iteration_peak_block_stats.
Add a small helper on KVCacheManagerV2 to build the per-cache-level
TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats] from a cache_level, then
have those three methods call it so the mapping logic stays in one place and
remains consistent if the tracked fields change.

Source: Coding guidelines

tensorrt_llm/executor/request.py (1)

203-203: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add return/type annotations to the new helpers.

These new methods are unannotated; add annotations such as -> list[int] | list[list[int]] | None, -> None, -> object, and -> dict[str, object] as appropriate. As per coding guidelines, “Always annotate functions with return types.” Based on learnings, Python 3.10+ features can be used throughout the codebase.

Also applies to: 213-213, 217-217, 225-225, 239-239

🤖 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/executor/request.py` at line 203, The new helper
methods/properties in Request are missing return annotations, so add explicit
return types to each one to match the coding guidelines. Update prompt_token_ids
and the other newly added helpers to use precise Python 3.10+ annotations such
as list[int] | list[list[int]] | None, None, object, or dict[str, object]
depending on what each accessor returns. Keep the signatures consistent across
the new accessors so the Request API remains fully typed.

Sources: Coding guidelines, Learnings

🤖 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 `@cpp/tensorrt_llm/executor/executorImpl.cpp`:
- Around line 1935-1945: The bounded stats trimming in `appendMultipleIterStats`
can call `pop_front()` more times than `mIterationStats` currently holds when
`currentIterStatsVec` alone exceeds the cap. Update the removal logic to clamp
the number of pops to the existing deque size before draining `mIterationStats`,
and apply the same fix in `appendMultipleRequestStats` so both batch-append
paths avoid empty-deque access.

In `@cpp/tensorrt_llm/nanobind/executor/request.cpp`:
- Around line 634-637: The request deserialization in request.cpp around the
inputTokenIdsBytes to VecTokens conversion should validate that the byte buffer
size is evenly divisible by sizeof(VecTokens::value_type) before calling
std::memcpy. If the size is invalid, reject the state or throw an error in the
same requestGetstate/request restore path; otherwise compute the element count
and copy only the validated number of bytes into inputTokenIds.

In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 3898-3899: The cached_kv_tokens stat is only updated in
_prepare_tp_inputs(), so incremental-update batches can leave it stale. Update
_prepare_incremental_update_metadata() in model_engine.py to assign
self.iter_states['cached_kv_tokens'] from the current num_cached_tokens_per_seq,
keeping the exported value in sync for both incremental-update flows.

In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2735-2789: Make the benchmark stall abort decision collective
across ADP ranks in the py_executor logic that checks benchmark fill progress.
The current timeout in the disagg benchmark path uses per-rank `time.time()` and
`_bench_disagg_last_gen_count_time`, which can make some ranks call
`_handle_errors()` and enter different collectives earlier than others. Update
the `is_benchmark_disagg` / `_benchmark_fill_phase_active` branch so all ranks
agree on whether the stall timeout has expired, and reset or suppress the stall
timer whenever any rank still reports fitting work via the `tp_allgather`
result. Keep the decision synchronized around `tp_allgather`, `total_ready_gen`,
`any_rank_has_fitting`, and `_handle_errors()` so no rank can diverge into abort
handling alone.

In `@tensorrt_llm/_torch/pyexecutor/request_utils.py`:
- Around line 83-105: The diagnostic helper in request_utils is using broad
Exception handlers in the state counting and kv_cache_manager info path, which
should be narrowed to the expected failure types instead of silently swallowing
everything. Update the logic around LlmRequestState(req.state_value).name,
idx_mapper.num_free_slots()/size(), and kv_cache_manager.get_num_free_blocks()
to catch only the specific enum/accessor errors you expect, and when a field is
unavailable, record a clear fallback marker in the output rather than ignoring
it. Keep the best-effort behavior, but avoid generic Exception in these branches
to satisfy the repository’s exception-handling guidelines and Ruff findings.

In `@tensorrt_llm/executor/base_worker.py`:
- Around line 457-461: The prompt token selection in base_worker.py is
incorrectly preferring _prompt_token_ids_i32 even after request.prompt_token_ids
has been materialized and may have been mutated in place. Update the logic
around the prompt_token_ids assignment in the request handling path so it only
uses the packed buffer when request._prompt_token_ids is still None, and
otherwise always uses the current list value from request.prompt_token_ids. Keep
the behavior consistent with prompt_adapter_request checks and the existing
request field names.

In `@tensorrt_llm/serve/openai_server.py`:
- Around line 948-950: The `/metrics` path in the background stats flow can
block while holding `_iteration_stats_lock` because
`_drain_iteration_stats_to_sinks_unlocked` is called with a 2-second timeout
even when the collector is already running. Update the metrics handling in
`openai_server.py` so the on-demand drain is non-blocking in background mode, or
skip draining entirely and read directly from `_iteration_stats_buffer`/the tee
buffer instead. Keep the change localized around `_iteration_stats_lock`,
`_drain_iteration_stats_to_sinks_unlocked`, and the `/metrics` stats collection
path.

In `@tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py`:
- Line 168: The resize callback setup in KV cache scheduler tests needs to
accept the new history_length keyword argument consistently. Update the custom
failure-path callback in test_kv_cache_v2_scheduler alongside the existing
resize_context_fn fallback so both signatures can receive history_length without
TypeError. Make sure the test cases that exercise skip behavior still use
resize_context on the same mocked manager path and validate both the positive
and failure paths after the callback signature change.

In `@tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py`:
- Around line 88-96: Add a test in this file to cover the VisualGen-specific
branch of `get_iteration_stats` in the OpenAI server path: assert that a
VisualGen-style server returns an empty list from `/metrics` and does not invoke
`get_stats_async()`. Use the existing
`test_metrics_endpoint_reads_queue_without_background_buffer` style and the
`server.get_iteration_stats()` / `metrics_collector` setup to locate the right
branch, and keep the assertion focused on VisualGen behavior rather than the
buffered/unbuffered LLM drain cases.

---

Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 1815-1839: The stats buffering logic in py_executor.py is
violating the max_stats_len contract in both the legacy path and the per-rank
iteration path. Update the checks around self.stats handling so max_stats_len ==
0 truly disables storage and no records are appended, and for positive limits
enforce an exact cap before appending (use the existing
_stats_buffer_is_unbounded helper, the self.stats_lock section, and the
gathered/per-rank append loop as the key places to adjust). Make the legacy trim
condition and the ADP overflow trimming consistent so the buffer never exceeds
max_stats_len.

In `@tensorrt_llm/serve/openai_server.py`:
- Line 940: The get_iteration_stats path still falls through to
self.generator.get_stats_async even for VisualGen, which causes /metrics to fail
because VisualGen lacks that method. Update get_iteration_stats in
openai_server.py to detect the VisualGen generator type, return an empty metrics
payload for that case, and only drain or fetch stats for generators that
actually implement get_stats_async; use the existing VisualGen handling in
register_visual_gen_routes and _get_iteration_stats_buffer_maxlen as the
reference points.

---

Nitpick comments:
In `@tensorrt_llm/executor/request.py`:
- Line 203: The new helper methods/properties in Request are missing return
annotations, so add explicit return types to each one to match the coding
guidelines. Update prompt_token_ids and the other newly added helpers to use
precise Python 3.10+ annotations such as list[int] | list[list[int]] | None,
None, object, or dict[str, object] depending on what each accessor returns. Keep
the signatures consistent across the new accessors so the Request API remains
fully typed.

In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py`:
- Around line 463-562: The same PoolGroupPeakBlockStats construction from
get_statistics is duplicated in _current_block_stats_by_cache_level,
_reset_iteration_peak_num_blocks, and get_and_reset_iteration_peak_block_stats.
Add a small helper on KVCacheManagerV2 to build the per-cache-level
TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats] from a cache_level, then
have those three methods call it so the mapping logic stays in one place and
remains consistent if the tracked fields change.

🪄 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: f27f2760-20a5-4b3c-8bf1-3bd41e89a424

📥 Commits

Reviewing files that changed from the base of the PR and between a6f29f5 and 6746a1f.

📒 Files selected for processing (38)

• cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
• cpp/include/tensorrt_llm/executor/executor.h
• cpp/tensorrt_llm/executor/executorConfig.cpp
• cpp/tensorrt_llm/executor/executorImpl.cpp
• cpp/tensorrt_llm/executor/executorImpl.h
• cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp
• cpp/tensorrt_llm/nanobind/executor/request.cpp
• cpp/tests/unit_tests/executor/executorConfigTest.cpp
• tensorrt_llm/_torch/attention_backend/interface.py
• tensorrt_llm/_torch/attention_backend/trtllm.py
• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py
• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tensorrt_llm/_torch/pyexecutor/py_executor.py
• tensorrt_llm/_torch/pyexecutor/request_utils.py
• tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
• tensorrt_llm/_torch/pyexecutor/trace_log_utils.py
• tensorrt_llm/executor/base_worker.py
• tensorrt_llm/executor/request.py
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py
• tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py
• tensorrt_llm/serve/openai_server.py
• tensorrt_llm/serve/perf_metrics.py
• tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py
• tests/unittest/bindings/test_executor_bindings.py
• tests/unittest/executor/test_stats_serializer.py
• tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py
• tests/unittest/llmapi/apps/test_disagg_perf_metrics_collector.py
• tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py
• tests/unittest/llmapi/test_llm_args.py

[jiaganc]

/bot run

[tensorrt-cicd]

PR_Github #56155 [ run ] triggered by Bot. Commit: 1ee7248 Link to invocation

[tensorrt-cicd]

PR_Github #56155 [ run ] completed with state FAILURE. Commit: 1ee7248
/LLM/main/L0_MergeRequest_PR pipeline #45018 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