[None][feat] Exact multimodal KV blockhashing

[venkywonka]

Description

Today's multimodal KV-cache block hashing assumes each multimodal item occupies one contiguous [pos, pos+length) span in the prompt. Production VLM preprocessors already break that assumption.

Model	Real run shape
Nemotron Nano v2 VL	1 video -> 8 runs x 258 patches, separated by 13-token Frame N sampled at X.XX seconds: gaps
Nemotron Nano v2 VL (EVS-pruned)	1 video -> 8 runs with lengths [1026, 2, 2, 2, 2, 2, 2, 2]
Qwen3-VL (Video-MME)	1 video -> 2 runs [10, 890) and [898, 1778) separated by "<vision_end><0.2 seconds><vision_start>"
Nano-Omni v3	video + audio item, frame separator + appended audio block, reuses v2 VL's preprocessor

What goes wrong under the contiguous claim

The producer emits multimodal_lengths[i] = num_mm_tokens (a count of mm tokens, not a span width). Both consumers then re-derive a span as [pos, pos + length), which is not true for any item whose patches sit in non-contiguous runs:

• C++ block-key extra keys (blockKey.cpp::generateBlockHashExtraKeys) - MmKey.startOffset is computed prompt-relative (block_start - pos), which only equals the item-local offset for contiguous items. Tail blocks past pos + length receive no MmKey at all, even when they hold real mm patches (e.g. positions [2088, 2183) of an 8-frame Nemotron video).
• Python V2 token augmentation (KVCacheManagerV2._augment_tokens_for_block_reuse) - overwrites positions inside the claimed span with digest-derived synthetic IDs. Real text tokens that fall in the gaps (Frame 2 sampled at 2.51 seconds:) are erased before hashing. Tail mm positions outside the span keep their placeholder IDs and carry no digest contribution at all.

Practical impact - correctness debt, not a hot incident

Surface-level cross-prompt collisions are rare in production today:

• The prefix hash chain means any digest mismatch breaks the chain at block 0, so two prompts with different videos diverge before the buggy bookkeeping ever matters.
• Gap text is deterministic for a given (video, sampling config), so same-digest requests have identical gap content - the masking happens to coincide with the truth.
• Tail-block placeholder IDs are reserved tokens that don't appear in normal text, so accidental tail-block collisions across different prompts are structurally unlikely.

What the PR actually closes is structural fragility that might surface as preprocessors evolve:

• The V2 augmentation relies on gap-text determinism never slipping. Tokenizer drift, locale formatting changes, or mixed-modality preprocessors (Nano-Omni v3 with audio appended) would silently serve wrong KV.
• Correct tail-block reuse past the truncated span is accidental - the hash isn't causally tied to the trailing mm content. EVS-pruned and Nano-Omni-style layouts (where the "tail" is most of the runs) need this causation to be enforced.

What this PR adds

Three optional flat arrays alongside the existing fields, so they cross the Python <-> C++ <-> serialization boundary cleanly:

multimodal_item_run_cu_offsets : List[int]   # prefix sum into runs, size: n_item_runs + 1
multimodal_run_positions       : List[int]   # prompt start of each run, flat across items, size: n_item_runs
multimodal_run_lengths         : List[int]   # length of each run, flat across items, size: n_item_runs

End-to-end plumbing:

• Producer (NEW): tensorrt_llm/inputs/multimodal.py::_find_mm_token_runs_from_mask extracts the runs from the per-token mm mask, called from multimodal_hashing_process in inputs/registry.py.
• API surface: MultimodalInput Python dataclass gains the three fields plus _validate_multimodal_runs (all-or-nothing, monotonic cu, sum(run_lengths(item)) == multimodal_lengths[item], runs ordered + non-overlapping, values fit in int32, pos+len <= INT32_MAX). C++ executor::MultimodalInput, serialization.cpp, and the nanobind binding follow; __setstate__ accepts both legacy 4-tuple and new 7-tuple state for backward-compatible deserialization.
• Request object: GenericLlmRequest gains 3 optional shared_ptr<vector<SizeType32>> fields; executor_request_to_llm_request and the nanobind ctor forward them.
• Two consumers refactored into resolve + run-based + legacy-fallback helpers:
  • C++ generateBlockHashExtraKeys -> resolveRunMetadata / generateRunBasedExtraKeys / generateLegacyContiguousExtraKeys.
  • Python _augment_tokens_for_block_reuse -> _resolve_multimodal_run_metadata / _augment_tokens_with_mm_run_metadata / _augment_tokens_with_contiguous_mm_metadata. gen_multimodal_cache_key_tokens (renamed from gen_multi_modal_tokens) gains a token_offset arg so multi-run items don't re-materialize the full per-item synthetic sequence.

Backward compatible: when run buffers are absent on a request, the legacy contiguous path runs exactly as before. Run buffers are only consulted when the producer emits them.

Test Coverage

• C++ unit tests (cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp)
  • GenerateExtraKeysUsesExactMultimodalRuns - non-contiguous runs produce per-run MmKeys with correct item-local offsets; gap-only blocks produce no key; second run's offset reflects accumulated prior run lengths.
  • GenerateExtraKeysDisablesHashingForIncompleteRunMetadata - partial run metadata logs a warning and falls back to no extra keys (request continues).
• C++ serialization round-trip (cpp/tests/unit_tests/executor/serializeUtilsTest.cpp::MultimodalInputWithUuids) extended to cover the 3 new fields.
• Python unit tests (tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py) - exact-runs path, out-of-slice run skipping, incomplete-metadata ValueError, contiguous-fallback preserved, gen_multimodal_cache_key_tokens(token_offset) semantics.
• Validator (tests/unittest/llmapi/test_llm_kv_cache_events.py::test_multimodal_input_exact_run_buffers) - from_components, paired-fields-required, sum-mismatch, int32 overflow.
• Bindings (tests/unittest/bindings/test_executor_bindings.py) - MultimodalInput ctor with run buffers, state-tuple round-trip with both 4-tuple (legacy) and 7-tuple (new) state.

Performance and accuracy validation

Rule for all perf tables: positive delta means feature is slower.

Main perf numbers

Nemotron Omni synthetic run. 4x GB200. tp_size=4. /v1/completions. No profiler. Each cell has 3 runs x 9 measured requests = 27 measured requests per arm.

Chunked prefill	Block reuse	Baseline mean	Feature mean	Delta	95% range
off	off	60.16 ms	60.03 ms	-0.13 ms / -0.21%	-0.39% to -0.04%
off	on	60.39 ms	61.49 ms	+1.10 ms / +1.83%	-0.42% to +4.08%
on	off	61.17 ms	59.52 ms	-1.64 ms / -2.65%	-9.69% to +4.39%
on	on	60.23 ms	60.28 ms	+0.05 ms / +0.08%	-4.88% to +5.05%

Takeaway: no clear perf regression. Only positive deltas are +1.83% and +0.08%, and both ranges cross zero.

Extra perf checks

Check	Baseline	Feature	Result
Qwen2-VL image, C++ KV, block reuse on, warm median	0.203896 s	0.202823 s	feature/baseline 0.994738
Qwen2-VL Video-MME, C++ KV, block reuse off, warm median	1.134069 s	1.143502 s	feature/baseline 1.008318
Qwen2-VL Video-MME, C++ KV, block reuse on, warm median	0.943674 s	0.951920 s	feature/baseline 1.008738
Qwen2-VL Video-MME, Python KVCacheManagerV2, block reuse off, warm median	1.106186 s	1.118415 s	feature/baseline 1.011055
Qwen2-VL Video-MME, Python KVCacheManagerV2, block reuse on, warm median	0.918805 s	0.930309 s	feature/baseline 1.012520
Qwen2.5-VL-7B Video-MME, block reuse off, direct generate time	0.5644 s	0.5706 s	feature/baseline 1.0110
Qwen2.5-VL-7B Video-MME, block reuse on, direct generate time	0.5789 s	0.5557 s	feature/baseline 0.9599

Worst warm repeated ratio observed: 1.012520 (+1.25%).

Accuracy numbers

Check	Baseline	Feature	Result
Qwen2-VL MMMU image sample	accuracy 40.0, MM keys 654, stored blocks 741	accuracy 40.0, MM keys 654, stored blocks 741	pass: output text and input IDs matched
Qwen2-VL Video-MME, C++ KV, block reuse off	answer A, token IDs [32, 151645]	answer A, token IDs [32, 151645]	pass: gold was A
Qwen2-VL Video-MME, C++ KV, block reuse on	answer A, token IDs [32, 151645]	answer A, token IDs [32, 151645]	pass: gold was A
Qwen2-VL Video-MME, Python KVCacheManagerV2, block reuse off	answer C	answer C	parity only: gold was A
Qwen2-VL Video-MME, Python KVCacheManagerV2, block reuse on	answer C	answer C	parity only: gold was A
Qwen2.5-VL-7B Video-MME, block reuse off	answer A, token IDs [32, 151645]	answer A, token IDs [32, 151645]	parity only: gold was D
Qwen2.5-VL-7B Video-MME, block reuse on	answer A, token IDs [32, 151645]	answer A, token IDs [32, 151645]	parity only: gold was D

Claim boundary: perf numbers above show no clear regression. Accuracy numbers above show one passing local Video-MME sample plus branch-vs-baseline parity in the other completed video lanes. No scorer-backed Nemotron Omni Video-MME accuracy number yet.

PR Checklist

• PR description clearly explains what and why.
• PR follows TRT-LLM CODING GUIDELINES.
• Test cases provided for new code paths.
• No new third-party dependencies.
• CODEOWNERS updated if ownership changes - N/A.
• Documentation updated as needed - pending review (no public-API docs reference the contiguous-span assumption).
• Update tava architecture diagram - N/A (no significant design change to the KV-cache layer).
• The reviewers assigned automatically/manually are appropriate for the PR.
• Please check this after reviewing the above items as appropriate for this PR.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR extends TensorRT-LLM with multimodal run metadata support, adding optional fields describing flat token run layouts (per-item offsets, positions, lengths) to enable run-based cache key generation. The implementation includes backward-compatible serialization, refactored block-key extra-key generation supporting both run-based and legacy paths, Python validation and run derivation helpers, and comprehensive test coverage across C++ and Python layers.

Changes

Multimodal Run Metadata Support

Layer / File(s)	Summary
Data Contracts & Public APIs cpp/include/tensorrt_llm/executor/executor.h, cpp/include/tensorrt_llm/batch_manager/llmRequest.h, tensorrt_llm/inputs/multimodal.py	MultimodalInput and GenericLlmRequest gain optional fields for run metadata: multimodalItemRunCuOffsets, multimodalRunPositions, multimodalRunLengths with corresponding getter methods. MultimodalInput dataclass adds fields and validation via _validate_multimodal_runs().
C++ Core Implementation cpp/tensorrt_llm/executor/multimodalInput.cpp, cpp/tensorrt_llm/executor/serialization.cpp, cpp/tensorrt_llm/batch_manager/blockKey.cpp	Constructor stores new fields via move-initialization; serialization conditionally encodes/decodes extension tag for backward compatibility; extra-key generation refactored into generateRunBasedExtraKeys() and generateLegacyContiguousExtraKeys() branches; executor::Request constructor extracts run metadata.
Python Core & Validation tensorrt_llm/inputs/multimodal.py, tensorrt_llm/inputs/registry.py	Run-derivation helper _find_mm_token_runs_from_mask() converts token masks to exact run metadata; input processor integration replaces per-token start offsets with run-based representation.
Token Generation Refactoring 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/_block_radix_tree.py	gen_multi_modal_tokens() renamed to gen_multimodal_cache_key_tokens() with added token_offset parameter for flexible synthetic token placement.
KV Cache Block Reuse Integration tensorrt_llm/_torch/pyexecutor/resource_manager.py	Token augmentation resolves run metadata via _resolve_multimodal_run_metadata() and branches to run-based or contiguous augmentation; digest-based token generation replaces per-item iteration.
Batch Manager & Worker Integration cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp, cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h, cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp, cpp/tensorrt_llm/nanobind/executor/request.cpp, tensorrt_llm/_torch/pyexecutor/llm_request.py, tensorrt_llm/executor/base_worker.py, tensorrt_llm/_torch/pyexecutor/_util.py	Nanobind bindings expose run metadata properties and extended constructors; executor-to-LLM request conversion extracts and forwards run fields; worker passes full multimodal input including dummy request support with run metadata.
C++ Unit Tests cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp, cpp/tests/unit_tests/executor/serializeUtilsTest.cpp	Tests verify run-based extra-key generation with correct offsets/UUIDs, incomplete metadata disablement, and legacy wire format alignment/backward compatibility.
Python Unit Tests tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py, tests/unittest/bindings/test_executor_bindings.py, tests/unittest/llmapi/test_llm_kv_cache_events.py	Tests validate token_offset application, run-based and contiguous augmentation paths, slicing behavior, metadata validation, int32 overflow detection, and pickle round-trips.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~50 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 27.03% 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	Title clearly describes the main feature (exact multimodal KV blockhashing) with proper formatting including ticket type and category.
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 pull request contains a comprehensive description explaining the issue (non-contiguous multimodal runs in production VLM preprocessors), the problem it solves (incorrect KV-cache block hashing and token augmentation), and the solution (three optional flat arrays for run metadata). The description includes practical impact analysis, detailed implementation details across producer/API/consumer layers, test coverage summary, and performance/accuracy validation numbers.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (4)

tests/unittest/llmapi/test_llm_kv_cache_events.py (1)

412-458: Strong unit coverage for exact-run buffer validation and int32 limits.

This adds both functional and failure-mode checks for the new multimodal run metadata, including overflow boundary behavior.

QA list updates are unnecessary here since the PR changes in this file are unit-test scoped (no new/changed tests/integration/defs/... coverage to schedule).

Also applies to: 460-482

🤖 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/llmapi/test_llm_kv_cache_events.py` around lines 412 - 458,
Add the same exact-run buffer validation and int32 overflow boundary tests for
multimodal run metadata in the sibling test block (lines referenced 460-482):
use _find_mm_token_runs_from_mask to compute
item_run_cu_offsets/run_positions/run_lengths, assert those values, create
MultimodalInput via from_components and assert the stored multimodal_*
attributes, then add the two failure-mode assertions that instantiate
MultimodalInput with mismatched provided run metadata and with a run length that
overflows (expecting ValueError with the same "must be provided together" and
"sum to 3, expected 2" messages); reference the functions/classes
MultimodalInput, MultimodalInput.from_components, and
_find_mm_token_runs_from_mask to locate where to add the duplicated checks.

cpp/tensorrt_llm/executor/serialization.cpp (1)

392-431: 💤 Low value

LGTM — serialize/serializedSize are mutually consistent.

serialize and serializedSize produce a coherent byte layout. One minor inconsistency worth noting: serialize caches hasMultimodalInputRunExtension(multimodalInput) in a const local (line 394), while serializedSize evaluates it twice (lines 415 and 424). Storing it in a local variable would harmonise the two functions.

♻️ Suggested consistency improvement

 size_t Serialization::serializedSize(MultimodalInput const& multimodalInput)
 {
     size_t totalSize = 0;
+    auto const hasRunExtension = hasMultimodalInputRunExtension(multimodalInput);
-    if (hasMultimodalInputRunExtension(multimodalInput))
+    if (hasRunExtension)
     {
         totalSize += su::serializedSize(kMultimodalInputRunExtensionTag);
         totalSize += su::serializedSize(kMultimodalInputRunExtensionVersion);
     }
     totalSize += su::serializedSize(multimodalInput.mMultimodalHashes);
     totalSize += su::serializedSize(multimodalInput.mMultimodalPositions);
     totalSize += su::serializedSize(multimodalInput.mMultimodalLengths);
     totalSize += su::serializedSize(multimodalInput.mMultimodalUuids);
-    if (hasMultimodalInputRunExtension(multimodalInput))
+    if (hasRunExtension)
     {

🤖 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 `@cpp/tensorrt_llm/executor/serialization.cpp` around lines 392 - 431,
serializedSize recomputes hasMultimodalInputRunExtension(multimodalInput) twice;
to match serialize's approach, call hasMultimodalInputRunExtension once and
cache the result in a local const (e.g. hasRunExtension) at the top of
Serialization::serializedSize, then use that variable in both conditional blocks
when adding kMultimodalInputRunExtensionTag, kMultimodalInputRunExtensionVersion
and the multimodal run fields (mMultimodalItemRunCuOffsets,
mMultimodalRunPositions, mMultimodalRunLengths).

tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py (1)

28-78: ⚡ Quick win

Add a two-item exact-run case here.

These exact-run assertions only cover a single multimodal item, so they never exercise the new run_item_indices / run_item_offsets mapping across multiple items. A small two-item fixture would catch regressions in the per-item offset math much earlier.

As per coding guidelines, "Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix."

🤖 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_kv_cache_v2_multimodal_runs.py` around
lines 28 - 78, Add a new unit test that exercises a two-item multimodal-run case
to validate per-item offset math in
KVCacheManagerV2._augment_tokens_for_block_reuse: create a fixture with two
multimodal items (e.g., multimodal_hashes containing two hashes,
multimodal_positions, multimodal_lengths, multimodal_item_run_cu_offsets
containing offsets for both items, and multimodal_run_positions/run_lengths
covering runs that span both items), generate expected mm_tokens for each item
via gen_multimodal_cache_key_tokens, call _augment_tokens_for_block_reuse with
full and sliced ranges, and assert that tokens are replaced at the correct
indices for each item (and that runs outside the slice are skipped); this will
exercise the run_item_indices / run_item_offsets mapping across multiple items.

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

164-183: ⚡ Quick win

Materialize the overlap tensors before the Python loop.

This runs on the pyexecutor block-reuse path, but the loop still does tensor masking/indexing per multimodal item. Converting the overlap metadata to Python lists once before iterating will avoid repeated tensor-bound work in a hot path.

Based on learnings: In files under tensorrt_llm/_torch/pyexecutor, avoid accessing torch.Tensor objects inside for-loops when iterating over requests. Convert batched tensors to Python lists beforehand using tensor.tolist(), and then iterate over those lists.

🤖 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 164 - 183,
The loop is repeatedly indexing and operating on torch.Tensors
(overlap_run_item_indices, overlap_item_indices, overlap_starts, overlap_ends,
token_offsets, result_offsets, lengths, etc.), which is expensive on the hot
pyexecutor path; before the Python for-loop that iterates multimodal items (the
block around overlap_item_indices and the for result_offset,... zip(...,
strict=True)), convert the per-item tensors to native Python lists (e.g., use
.tolist() or .cpu().tolist()) and materialize any small per-item arrays used
inside the loop (e.g., overlap_run_item_indices, overlap_starts, overlap_ends,
token_offsets, result_offsets, lengths, and multimodal_hashes entries used by
_hash_to_digest) so the inner loop uses only Python scalars/lists rather than
indexing torch.Tensors. Ensure you still compute digest = _hash_to_digest(...)
with the materialized multimodal hash value, and preserve the existing semantics
and ordering.

🤖 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/batch_manager/blockKey.cpp`:
- Around line 211-215: The overlap check uses signed addition startPos + length
which can overflow; change the legacy branch to compute a safe run end using
getValidRunEnd(startPos, length) and use that value in the if-condition and when
computing mmStartInBlock so the logic mirrors the run-based path; update the
condition if (endTokenIdx > startPos && startTokenIdx < startPos + length) to
use the safe end (e.g. auto const runEnd = getValidRunEnd(startPos, length); if
(endTokenIdx > startPos && startTokenIdx < runEnd) { auto const mmStartInBlock =
startPos >= startTokenIdx ? 0 : startTokenIdx - startPos; ... }) ensuring you
reference makeHashArray(multimodalHashes[itemIdx]) and getUuid(multimodalUuids,
itemIdx) unchanged when emplacing into extraKeys.

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 100-105: Replace the debug-only assert in _hash_to_digest with
explicit runtime validation that raises a ValueError when the input is
malformed: check that hash_ints is a Sequence of length exactly 8 (and
optionally that each element is an int) and raise ValueError with a clear
message if not, before converting values to bytes; keep the existing conversion
logic using v.to_bytes('big', signed=True) unchanged so the function still
returns the 32-byte digest when validation passes.

In `@tensorrt_llm/inputs/multimodal.py`:
- Line 149: Before calling _validate_multimodal_runs(), add an explicit
cardinality check that len(self.multimodal_hashes) ==
len(self.multimodal_positions) == len(self.multimodal_lengths) and raise a clear
ValueError if not; this prevents an IndexError later when building
multimodal_item_run_cu_offsets (and the per-item loop indexed by
self.multimodal_lengths) — apply the same pre-check in both places where
_validate_multimodal_runs() is invoked so the three multimodal lists are
guaranteed aligned before further processing.

---

Nitpick comments:
In `@cpp/tensorrt_llm/executor/serialization.cpp`:
- Around line 392-431: serializedSize recomputes
hasMultimodalInputRunExtension(multimodalInput) twice; to match serialize's
approach, call hasMultimodalInputRunExtension once and cache the result in a
local const (e.g. hasRunExtension) at the top of Serialization::serializedSize,
then use that variable in both conditional blocks when adding
kMultimodalInputRunExtensionTag, kMultimodalInputRunExtensionVersion and the
multimodal run fields (mMultimodalItemRunCuOffsets, mMultimodalRunPositions,
mMultimodalRunLengths).

In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 164-183: The loop is repeatedly indexing and operating on
torch.Tensors (overlap_run_item_indices, overlap_item_indices, overlap_starts,
overlap_ends, token_offsets, result_offsets, lengths, etc.), which is expensive
on the hot pyexecutor path; before the Python for-loop that iterates multimodal
items (the block around overlap_item_indices and the for result_offset,...
zip(..., strict=True)), convert the per-item tensors to native Python lists
(e.g., use .tolist() or .cpu().tolist()) and materialize any small per-item
arrays used inside the loop (e.g., overlap_run_item_indices, overlap_starts,
overlap_ends, token_offsets, result_offsets, lengths, and multimodal_hashes
entries used by _hash_to_digest) so the inner loop uses only Python
scalars/lists rather than indexing torch.Tensors. Ensure you still compute
digest = _hash_to_digest(...) with the materialized multimodal hash value, and
preserve the existing semantics and ordering.

In `@tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py`:
- Around line 28-78: Add a new unit test that exercises a two-item
multimodal-run case to validate per-item offset math in
KVCacheManagerV2._augment_tokens_for_block_reuse: create a fixture with two
multimodal items (e.g., multimodal_hashes containing two hashes,
multimodal_positions, multimodal_lengths, multimodal_item_run_cu_offsets
containing offsets for both items, and multimodal_run_positions/run_lengths
covering runs that span both items), generate expected mm_tokens for each item
via gen_multimodal_cache_key_tokens, call _augment_tokens_for_block_reuse with
full and sliced ranges, and assert that tokens are replaced at the correct
indices for each item (and that runs outside the slice are skipped); this will
exercise the run_item_indices / run_item_offsets mapping across multiple items.

In `@tests/unittest/llmapi/test_llm_kv_cache_events.py`:
- Around line 412-458: Add the same exact-run buffer validation and int32
overflow boundary tests for multimodal run metadata in the sibling test block
(lines referenced 460-482): use _find_mm_token_runs_from_mask to compute
item_run_cu_offsets/run_positions/run_lengths, assert those values, create
MultimodalInput via from_components and assert the stored multimodal_*
attributes, then add the two failure-mode assertions that instantiate
MultimodalInput with mismatched provided run metadata and with a run length that
overflows (expecting ValueError with the same "must be provided together" and
"sum to 3, expected 2" messages); reference the functions/classes
MultimodalInput, MultimodalInput.from_components, and
_find_mm_token_runs_from_mask to locate where to add the duplicated checks.

🪄 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: 1e8117b1-fdce-4914-93db-84eeb3df59be

📥 Commits

Reviewing files that changed from the base of the PR and between 2899e92 and 46f13a3.

📒 Files selected for processing (23)

• cpp/include/tensorrt_llm/batch_manager/llmRequest.h
• cpp/include/tensorrt_llm/executor/executor.h
• cpp/tensorrt_llm/batch_manager/blockKey.cpp
• cpp/tensorrt_llm/executor/multimodalInput.cpp
• cpp/tensorrt_llm/executor/serialization.cpp
• cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
• cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.cpp
• cpp/tensorrt_llm/nanobind/batch_manager/llmRequest.h
• cpp/tensorrt_llm/nanobind/executor/request.cpp
• cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp
• cpp/tests/unit_tests/executor/serializeUtilsTest.cpp
• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/pyexecutor/llm_request.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py
• tensorrt_llm/executor/base_worker.py
• tensorrt_llm/inputs/multimodal.py
• tensorrt_llm/inputs/registry.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/_block_radix_tree.py
• tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py
• tests/unittest/bindings/test_executor_bindings.py
• tests/unittest/llmapi/test_llm_kv_cache_events.py

[venkywonka]

/bot run

[tensorrt-cicd]

PR_Github #47792 [ run ] triggered by Bot. Commit: a4b0e83 Link to invocation

[SimengLiu-nv]

When any exact-run field is present but the run metadata is incomplete or malformed, this returns an empty extra-key vector. That does not disable KV reuse; it builds ordinary text/placeholder-only BlockKeys, so two requests with the same prompt placeholder IDs but different multimodal content can collide and reuse the wrong KV. Since the legacy multimodal hashes/positions/lengths are still available, this path should either fall back to the legacy extra-key generation or make the request uncacheable instead of returning reusable empty keys.

Assume two requests differ only by image content.

Setup

• tokens_per_block = 4
• Prompt tokens for both requests: [101, IMG, IMG, 102, 103]
• IMG is the same placeholder token ID in both requests.
• Request A image hash: hashA
• Request B image hash: hashB
• Legacy metadata is present:
  • multimodal_hashes = [hashA/hashB]
  • multimodal_positions = [1]
  • multimodal_lengths = [2]
• Exact-run metadata is malformed:
  • multimodal_item_run_cu_offsets = [0, 1]
  • multimodal_run_positions = null
  • multimodal_run_lengths = null

What Happens Today

1. generateBlockHashExtraKeys() sees at least one exact-run field, so it calls resolveRunMetadata().
2. resolveRunMetadata() returns kInvalid because the three run fields are not all present.
3. The current code returns {} from blockKey.cpp.
4. buildBlockKeys() still creates reusable BlockKeys with extraKeys = {}.
5. For block [101, IMG, IMG, 102], Request A and Request B now have identical block keys:
   • same token IDs: [101, IMG, IMG, 102]
   • same loraTaskId
   • same cacheSaltID
   • same extraKeys = {}
6. The KV cache can reuse Request A’s cached block for Request B, even though hashA != hashB.

Why This Is Wrong

extraKeys is the only thing that makes identical multimodal placeholder tokens distinguish different media content. Returning an empty vector does not mean “disable reuse”; it means “hash this like ordinary text.” For multimodal placeholders, that is unsafe.

Safe Behavior

Preferred behavior:

• If exact-run metadata is invalid but legacy multimodal_positions, multimodal_lengths, and multimodal_hashes are valid, fall back to generateLegacyContiguousExtraKeys(...).

Stricter behavior:

• Mark the request or block as uncacheable instead of producing normal text-only block keys.

Unsafe behavior:

• Returning {} while still committing/reusing the block.

[tensorrt-cicd]

PR_Github #47792 [ run ] completed with state SUCCESS. Commit: a4b0e83
/LLM/main/L0_MergeRequest_PR pipeline #37683 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

[SimengLiu-nv]

LGTM.

[venkywonka]

/bot run

[tensorrt-cicd]

PR_Github #48032 [ run ] triggered by Bot. Commit: 55e6a1c Link to invocation

[moraxu]

Thanks for the comments in resource_manager.py

[tensorrt-cicd]

PR_Github #48032 [ run ] completed with state SUCCESS. Commit: 55e6a1c
/LLM/main/L0_MergeRequest_PR pipeline #37866 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

[venkywonka]

/bot run

[tensorrt-cicd]

PR_Github #49158 [ run ] completed with state ABORTED. Commit: b62ac47

Link to invocation

[venkywonka]

/bot run

[tensorrt-cicd]

PR_Github #49204 [ run ] triggered by Bot. Commit: 538536e Link to invocation

[tensorrt-cicd]

PR_Github #49193 [ run ] completed with state ABORTED. Commit: b62ac47

Link to invocation

[tensorrt-cicd]

PR_Github #49204 [ run ] completed with state SUCCESS. Commit: 538536e
/LLM/main/L0_MergeRequest_PR pipeline #38879 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

[venkywonka]

/bot run

[tensorrt-cicd]

PR_Github #49249 [ run ] triggered by Bot. Commit: 71e4c3f Link to invocation

[tensorrt-cicd]

PR_Github #49249 [ run ] completed with state SUCCESS. Commit: 71e4c3f
/LLM/main/L0_MergeRequest_PR pipeline #38919 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

[venkywonka]

/bot run

[tensorrt-cicd]

PR_Github #49273 [ run ] triggered by Bot. Commit: 71e4c3f Link to invocation

[tensorrt-cicd]

PR_Github #49273 [ run ] completed with state SUCCESS. Commit: 71e4c3f
/LLM/main/L0_MergeRequest_PR pipeline #38939 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation