[None][perf] Memoize multimodal run metadata and drop identity block-offset gather

[eopXD]

Description

Two host-side overhead reductions in the KV cache managers. Both are performance-only with no external behavior change.

1. Memoize per-request multimodal run metadata (kv_cache_manager_v2.py). _resolve_multimodal_run_metadata is resolved on the first context chunk, again on every commit during chunked prefill, and on prefetch. Its result is a pure function of the request's immutable multimodal run layout (multimodal_item_run_cu_offsets / multimodal_run_positions / multimodal_run_lengths), so the CPU tensors it builds (arange / repeat_interleave / cumsum / cat / int64 casts) were being rebuilt identically on every call. The result is now cached on the request, whose lifetime bounds the cache, so no separate teardown is needed. This resolves the existing TODO(perf). The non-None path is exercised by the VLMs that populate run metadata (Qwen2VL, Qwen3VL, LLaVA-Next, Kimi-K2.5, Nemotron-Nano).

2. Remove an identity gather (resource_manager.py, BlockManager.copy_block_offsets). base_block_offsets was a torch.arange, so base_block_offsets[block_ids] returns block_ids unchanged. The block ids are now converted directly and the unused base_block_offsets buffer is removed (also dropping its per-manager allocation). This BlockManager backs the RocketKV KT cache.

Test Coverage

No new tests are added: both changes are performance-only with no external behavior change. Existing coverage exercises the touched paths:

• tests/unittest/_torch/executor/test_kv_cache_v2_multimodal_runs.py — calls _augment_tokens_for_block_reuse multiple times on the same request, exercising the memoized _resolve_multimodal_run_metadata on both the cache-miss (build) and cache-hit (return) paths, as well as the contiguous None path.
• tests/unittest/_torch/executor/test_kv_cache_manager_v2.py — KV cache manager v2 behavior.
• tests/unittest/_torch/attention/sparse/rocketkv/test_rocketkv.py — RocketKV path that calls BlockManager.copy_block_offsets.

PR Checklist

• No API changes; no new dependencies; no CODEOWNERS or documentation changes required.

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

Summary by CodeRabbit

• Performance Improvements
  • Improved multimodal request processing by reusing computed metadata during repeated operations.
  • Streamlined block offset preparation to reduce unnecessary intermediate data handling.

[eopXD]

/bot run

[coderabbitai]

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: 0f9dec33-4caa-4af2-a457-83c09b74665a

📥 Commits

Reviewing files that changed from the base of the PR and between e2ce7a6 and 6810969.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py

---

📝 Walkthrough

Walkthrough

The PyExecutor now caches multimodal run metadata per request and copies block IDs directly into destination tensors without a precomputed CPU offset tensor.

Changes

Multimodal metadata caching

Layer / File(s)	Summary
Request-scoped metadata memoization tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py	The metadata resolver reuses cached request metadata and stores either None or computed _MmRunMetadata values on the request.

Block offset copying

Layer / File(s)	Summary
Direct block offset copying tensorrt_llm/_torch/pyexecutor/resource_manager.py	BlockManager removes the precomputed CPU offset tensor and creates CPU tensors directly from request block IDs during copying.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise and matches the two performance optimizations described in the PR.
Description check	✅ Passed	The PR description follows the template with Description, Test Coverage, and PR Checklist sections filled in.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.}

[tensorrt-cicd]

PR_Github #58648 [ run ] triggered by Bot. Commit: 6810969 Link to invocation