diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index faa8b064caf6..9651b78a3b63 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index cc598cbf67ca..66bd65aaef16 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/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. diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py index aed1ddfe94fb..0be8abb2fa8a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_multimodal_runs.py +++ b/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 diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 49fef0c8d50e..609e1eea4327 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/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, ):