[#18849][fix] Pass vocab_size to hybrid-Mamba KV cache managers

[dibyo10]

Description

Fixes #18849.

_create_kv_cache_manager passed vocab_size=config.vocab_size on the MLA branch and the plain-attention branch only. The three hybrid-Mamba branches (Kimi K3 / KDA, nemotron_hybrid, and qwen3_next / Qwen3.5 / Qwen4-Exp) constructed the manager without it, and MambaHybridCacheManagerV2 absorbs unknown keywords through **kwargs, so KVCacheManagerV2.vocab_size silently stayed None.

That field is read only when building multimodal block-reuse cache keys (_augment_tokens_with_mm_run_metadata / _augment_tokens_with_contiguous_mm_metadata), which is why a hybrid deployment serves text-only traffic indefinitely and then dies on its first image_url request: gen_multimodal_cache_key_tokens() receives None as id_offset, the resulting TypeError escapes into the PyExecutor event loop, and every rank is hard-killed.

Two changes:

1. vocab_size moves into manager_extra_kwargs, which is already gated on issubclass(kv_cache_manager_cls, KVCacheManagerV2), and the two per-branch vocab_size=config.vocab_size arguments are removed. Every V2 branch now gets the value from one place, so a branch added later cannot omit it — the same reasoning _mamba_conv_layout_kwargs already applies to conv_state_layout. The V1 managers are unaffected: KVCacheManager swallowed the argument into **kwargs without reading it, and MixedMambaHybridCacheManager (which has no **kwargs) must not receive it.
2. The new _resolve_vocab_size helper falls back to text_config.vocab_size, because composite VLM configs keep the field on the nested text config.

Additionally, KVCacheManagerV2._augment_tokens_for_block_reuse now raises a RuntimeError naming the missing construction argument when it is asked to build multimodal cache keys without a vocab_size, rather than letting None reach the binding. Text-only requests return before that check, so they are unaffected.

Not addressed here, to keep this PR to one concern: suggestion 3 in the issue — failing the individual request with a 4xx instead of terminating the executor when multimodal metadata cannot be built (which also covers the disable_mm_encoder: true case). That is executor-level error classification and is worth its own PR, so it is now tracked separately in #18971.

I do not have Hopper hardware, so the end-to-end reproduction from the issue was not re-run; the reporter's verified workaround fills in the same value this change now passes.

Test Coverage

tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py

• test_hybrid_v2_manager_receives_vocab_size — routes a Qwen3 hybrid config through _create_kv_cache_manager and asserts the V2 manager is constructed with vocab_size. Fails on main.
• test_hybrid_v2_manager_reads_vocab_size_from_text_config — same, with vocab_size only on the nested text_config.
• test_hybrid_v1_manager_does_not_receive_vocab_size — guards the V1 MixedMambaHybridCacheManager, whose signature has no **kwargs.

tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py

• test_augment_tokens_for_block_reuse_reports_missing_vocab_size — the guard names the missing argument.
• test_augment_tokens_for_block_reuse_ignores_missing_vocab_size_for_text — text-only requests still work without a vocab_size.

Both files are CPU-only. pre-commit's formatters and linters (ruff 0.9.4 check/format on the Group A files, yapf 0.43.0 / isort 5.12.0 / ruff-legacy on _util.py, codespell) pass on the changed files. The test cases themselves have not been executed — TensorRT-LLM does not build on the macOS host I am working from — so please run the two files in CI.

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.

Dev Engineer Review

• _resolve_vocab_size reads from the main config or nested text_config.
• V2 managers, including hybrid-Mamba managers, now receive the resolved value.
• V1 construction remains compatible.
• Multimodal cache-key generation raises RuntimeError when vocab_size is missing. Text-only requests remain unchanged.
• Executor-level request isolation remains out of scope.

QA Engineer Review

• test_kv_cache_v2_multimodal_runs.py adds coverage for missing vocab_size errors and text-only behavior.
• test_mamba_cache_manager.py adds coverage for direct and nested configuration resolution and V1 compatibility.
• No entries for these unit tests were found in tests/integration/test_lists/test-db/ or tests/integration/test_lists/qa/.
• Tests were not run locally because TensorRT-LLM does not build on the contributor’s macOS host.
• Coverage verdict: sufficient for the implemented scope. Executor-level error handling needs follow-up.

Per-File QA Perspective

• tensorrt_llm/_torch/pyexecutor/_util.py: Verify vocab_size propagation for hybrid-Mamba and composite VLM configurations. Verify that V1 managers do not receive the V2-only argument.
• tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py: Verify that missing vocab_size raises RuntimeError only during multimodal cache-key generation.
• tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py: Covers missing-value errors and unchanged text-only behavior. It is not listed in the CI or manual-QA test lists.
• tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py: Covers direct and nested configuration resolution and V1 compatibility. It is not listed in the CI or manual-QA test lists.

[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: e2365d0b-71dc-4ebb-a1f5-c3ba924e2f2a

📥 Commits

Reviewing files that changed from the base of the PR and between 7bad4c5 and 94a6cc6.

📒 Files selected for processing (1)

• tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py

🚧 Files skipped from review as they are similar to previous changes (1)

• tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py

Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.

---

Walkthrough

The executor resolves vocabulary size from top-level or nested text configuration and passes it to V2 KV-cache managers. Multimodal metadata generation reports missing vocabulary size explicitly. Tests cover hybrid propagation and text-only behavior.

Changes

KV-cache vocabulary handling

Layer / File(s)	Summary
Resolve and propagate vocabulary size tensorrt_llm/_torch/pyexecutor/_util.py, tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py	The executor resolves vocab_size from the model configuration or nested text_config and passes it to V2 cache managers. Tests cover constructor inputs and confirm that the V1 mixed manager does not receive the V2-only argument.
Validate multimodal metadata tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py, tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py	KVCacheManagerV2 accepts an omitted vocab_size and raises a descriptive RuntimeError for multimodal metadata when it is missing. Text-only augmentation remains unchanged.

Priority: ⬆️ High

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

Severity of issue fixed: High

Merge Risk: ⚪ Minimal · up to d281a

This change propagates vocabulary size to V2 hybrid-Mamba cache managers and adds a clear multimodal validation error. No concrete merge-blocking risk remains in the supplied evidence.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the primary fix: passing vocab_size to hybrid-Mamba KV cache managers.
Description check	✅ Passed	The description is complete and explains the problem, solution, test coverage, limitations, and deferred executor-level work.
Linked Issues check	✅ Passed	The changes address issue #18849 by propagating vocab_size to V2 hybrid-Mamba managers, resolving nested text_config values, preserving V1 compatibility, and guarding missing values during multimodal …
Out of Scope Changes check	✅ Passed	The code and tests remain focused on vocab_size propagation, configuration resolution, multimodal validation, and V1 compatibility. No unrelated changes are evident.
Docstring Coverage	✅ Passed	Docstring coverage is 82.35% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 17 functions across 4 files.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[mikeiovine]

Stamp on behalf of runtime devs, delegating proper review to @NVIDIA/trt-llm-kv-cache-manager-devs; please ping me if you think this is not accurate