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

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2325,6 +2325,22 @@ def _mamba_conv_layout_kwargs(kv_cache_manager_cls: type,
     return {"model_type": model_type}
 
 
+def _resolve_vocab_size(config) -> Optional[int]:
+    """Vocabulary size the V2 manager needs to build multimodal cache keys.
+
+    ``KVCacheManagerV2`` reserves token ids at and above ``vocab_size`` for the
+    synthetic tokens that carry a multimodal item's content digest into the
+    radix tree, so every V2 manager needs the value before the first image
+    request; text-only traffic never reads it.  Composite VLM configs keep the
+    field on the nested ``text_config``, so fall back there rather than
+    reading ``config.vocab_size`` directly.
+    """
+    vocab_size = getattr(config, "vocab_size", None)
+    if vocab_size is not None:
+        return vocab_size
+    return getattr(getattr(config, "text_config", None), "vocab_size", None)
+
+
 def _get_qwen4_exp_ple_cache_params(config, *, total_layers: int,
                                     is_draft: bool):
     """Align target-only PLE state with a target/draft cache layout."""
@@ -2511,6 +2527,10 @@ def _create_kv_cache_manager(
             draft_config_for_kv)
     manager_extra_kwargs = {}
     if issubclass(kv_cache_manager_cls, KVCacheManagerV2):
+        # Set here rather than per branch: the hybrid-Mamba branches below used
+        # to omit vocab_size, and a manager built without it fails inside the
+        # first multimodal request instead of at start-up.
+        manager_extra_kwargs["vocab_size"] = _resolve_vocab_size(config)
         manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats
         manager_extra_kwargs[
             "cold_page_codec_provider"] = cold_page_codec_provider
@@ -2616,7 +2636,6 @@ def _create_kv_cache_manager(
             mapping=mapping,
             dtype=kv_cache_dtype,
             spec_config=spec_config,
-            vocab_size=config.vocab_size,
             max_num_tokens=max_num_tokens,
             max_beam_width=max_beam_width,
             is_draft=is_draft,
@@ -2900,7 +2919,6 @@ def _create_kv_cache_manager(
             mapping=mapping,
             dtype=kv_cache_dtype,
             spec_config=spec_config,
-            vocab_size=config.vocab_size,
             max_num_tokens=max_num_tokens,
             model_config=binding_model_config,
             max_beam_width=max_beam_width,

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -921,7 +921,7 @@ def __init__(
         dtype: DataType = DataType.HALF,
         spec_config=None,
         layer_mask: Optional[List[bool]] = None,
-        vocab_size: int = None,
+        vocab_size: Optional[int] = None,
         max_num_tokens: int = 8192,
         model_config: Optional[ModelConfigCpp] = None,
         max_beam_width: int = 1,
@@ -3096,6 +3096,16 @@ def _augment_tokens_for_block_reuse(
         ):
             return tokens[chunk_start:chunk_end] if is_sliced else tokens
 
+        if self.vocab_size is None:
+            # Only multimodal requests read vocab_size, so a manager built
+            # without it serves text traffic indefinitely and then fails on the
+            # first image. Name the missing argument here; the alternative is a
+            # binding TypeError reporting a NoneType id_offset.
+            raise RuntimeError(
+                f"{type(self).__name__} was constructed without vocab_size, "
+                "which is required to build multimodal block-reuse cache keys"
+            )
+
         # Multimodal path: materialize a Python-int list (digest bytes get spliced in below),
         # which flows through the per-element binding fallback. tokens may be a zero-copy numpy
         # int32 view (get_tokens_view) — use tolist() so elements are Python ints, not np.int32.

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

@@ -180,3 +180,33 @@ def test_augment_tokens_for_block_reuse_keeps_contiguous_metadata_path():
 
     sliced = KVCacheManagerV2._augment_tokens_for_block_reuse(manager, tokens, req, start=1, end=5)
     assert sliced == [tokens[1], *mm_tokens]
+
+
+def test_augment_tokens_for_block_reuse_reports_missing_vocab_size():
+    """https://github.com/NVIDIA/TensorRT-LLM/issues/18849: a manager built
+    without vocab_size must name the missing construction argument instead of
+    reaching the cache-key binding with a None id_offset."""
+    tokens = list(range(8))
+    manager = _make_manager(None)
+    req = _make_request(
+        tokens,
+        multimodal_hashes=[_HASH_INTS],
+        multimodal_positions=[2],
+        multimodal_lengths=[3],
+        multimodal_item_run_cu_offsets=None,
+        multimodal_run_positions=None,
+        multimodal_run_lengths=None,
+    )
+
+    with pytest.raises(RuntimeError, match="without vocab_size"):
+        KVCacheManagerV2._augment_tokens_for_block_reuse(manager, tokens, req)
+
+
+def test_augment_tokens_for_block_reuse_ignores_missing_vocab_size_for_text():
+    """Text-only requests never build multimodal cache keys, so they must keep
+    working on a manager that has no vocab_size."""
+    tokens = list(range(8))
+    manager = _make_manager(None)
+    req = _make_request(tokens)
+
+    assert KVCacheManagerV2._augment_tokens_for_block_reuse(manager, tokens, req) == tokens

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

@@ -544,6 +544,104 @@ def __init__(self, *args, **kwargs):
     assert not any("RecordingV2Manager was selected" in log for log in fallback_logs)
 
 
+def _capture_qwen3_hybrid_manager_ctor(monkeypatch, manager_base, pretrained_config):
+    """Route a Qwen3 hybrid config through _create_kv_cache_manager and capture
+    the constructor kwargs."""
+    captured = {}
+
+    class RecordingManager(manager_base):
+        def __init__(self, *args, **kwargs):
+            """Record the constructor kwargs without building a real manager.
+
+            The base constructor is deliberately not called: it would allocate
+            cache pools and need a GPU, and these tests only assert on what the
+            factory passes in.
+            """
+            captured.update(kwargs)
+
+    monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 90)
+    _create_kv_cache_manager(
+        model_engine=None,
+        kv_cache_manager_cls=RecordingManager,
+        mapping=Mapping(world_size=1, tp_size=1, pp_size=1),
+        kv_cache_config=KvCacheConfig(
+            use_kv_cache_manager_v2=issubclass(manager_base, KVCacheManagerV2)
+        ),
+        tokens_per_block=32,
+        max_seq_len=2048,
+        max_batch_size=4,
+        spec_config=None,
+        sparse_attention_config=None,
+        max_num_tokens=256,
+        max_beam_width=1,
+        kv_connector_manager=None,
+        model_config=SimpleNamespace(pretrained_config=pretrained_config, quant_config=None),
+        dtype=torch.bfloat16,
+        is_draft=False,
+    )
+    return captured
+
+
+def _qwen3_hybrid_pretrained_config(*, vocab_size=None, text_vocab_size=None):
+    """Minimal Qwen3 hybrid text config, with vocab_size placed either on the
+    config itself or on a nested text_config as composite VLM configs do.
+
+    Both keywords default to None so the field can also be left off entirely.
+    """
+    config = SimpleNamespace(
+        architectures=["Qwen3_5ForCausalLM"],
+        num_hidden_layers=2,
+        layer_types=["linear_attention", "full_attention"],
+        linear_key_head_dim=8,
+        linear_conv_kernel_dim=4,
+        linear_num_value_heads=4,
+        linear_num_key_heads=1,
+        linear_value_head_dim=8,
+        num_key_value_heads=2,
+        num_attention_heads=4,
+        hidden_size=32,
+        torch_dtype=torch.bfloat16,
+    )
+    if vocab_size is not None:
+        config.vocab_size = vocab_size
+    if text_vocab_size is not None:
+        config.text_config = SimpleNamespace(vocab_size=text_vocab_size)
+    return config
+
+
+def test_hybrid_v2_manager_receives_vocab_size(monkeypatch):
+    """https://github.com/NVIDIA/TensorRT-LLM/issues/18849: the hybrid-Mamba
+    branch omitted vocab_size, so the V2 manager could not build multimodal
+    cache keys and the first image request killed the executor."""
+    kwargs = _capture_qwen3_hybrid_manager_ctor(
+        monkeypatch,
+        MambaHybridCacheManagerV2,
+        _qwen3_hybrid_pretrained_config(vocab_size=151936),
+    )
+    assert kwargs["vocab_size"] == 151936
+
+
+def test_hybrid_v2_manager_reads_vocab_size_from_text_config(monkeypatch):
+    """Composite VLM configs keep vocab_size on the nested text_config."""
+    kwargs = _capture_qwen3_hybrid_manager_ctor(
+        monkeypatch,
+        MambaHybridCacheManagerV2,
+        _qwen3_hybrid_pretrained_config(text_vocab_size=151936),
+    )
+    assert kwargs["vocab_size"] == 151936
+
+
+def test_hybrid_v1_manager_does_not_receive_vocab_size(monkeypatch):
+    """vocab_size is a V2-only parameter; MixedMambaHybridCacheManager has no
+    **kwargs, so passing it there would be a TypeError at start-up."""
+    kwargs = _capture_qwen3_hybrid_manager_ctor(
+        monkeypatch,
+        MixedMambaHybridCacheManager,
+        _qwen3_hybrid_pretrained_config(vocab_size=151936),
+    )
+    assert "vocab_size" not in kwargs
+
+
 def test_hybrid_cache_manager_factory_rejects_cpp_preference_with_explicit_v2(
     monkeypatch,
 ):