diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index b8b395abdc4d..0eb5085b3daf 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -74,6 +74,34 @@ scheduler_config: enable_prefix_aware_scheduling: false ``` +### Selecting the KV Cache Manager + +TensorRT LLM ships two KV cache manager implementations. `use_kv_cache_manager_v2` +selects between them and defaults to `auto`, which adopts the model's own +preference and falls back to the V1 C++ manager for models that do not declare +one. Set it to `true` or `false` to override the model default. + +Models that select the V2 manager by default: + +| Model | Reason | +| --- | --- | +| Hybrid Mamba (NemotronH, Qwen3-Next) | Attention KV and Mamba state pools must be sized together | +| DeepSeek-V4 | Sparse attention attaches auxiliary per-layer buffers | +| GPT-OSS | Sliding window on every other layer (VSWA), so the sliding-window and full-attention pools are sized independently | + +Separately, Gemma4 hybrid attention and sparse-attention models are routed to +V2 unconditionally: their per-layer buffer layouts cannot be represented by V1's +unified pool, so `use_kv_cache_manager_v2` does not apply to them. + +Two-model speculative decoding (for example Eagle3 with +`eagle3_one_model=False`) is not supported by V2: the draft model runs in a +separate engine with its own KV cache manager, and V2 sizes both managers from +the full `max_gpu_total_bytes` budget instead of partitioning it between them. +Under `auto`, a model default of V2 falls back to V1 for that combination; +setting `use_kv_cache_manager_v2: true` explicitly raises an error. This +applies to every model that selects its manager through +`use_kv_cache_manager_v2`, not just GPT-OSS. + ### Mamba Snapshot Boundaries Hybrid Mamba models must retain the recurrent Mamba state together with the @@ -111,8 +139,9 @@ If neither `avg_seq_len` nor an explicit `pool_ratio` is configured, hybrid Mamba models warn and fall back to half of `max_seq_len`, which can produce a suboptimal pool split. Exact explicit boundaries currently require `MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid -Mamba models select V2 by default when -`use_kv_cache_manager_v2: auto`; set it to `false` to select the V1 C++ +Mamba models select V2 by default (see +[Selecting the KV Cache Manager](#selecting-the-kv-cache-manager)); set +`use_kv_cache_manager_v2` to `false` to select the V1 C++ compatibility manager. In disaggregated serving, V2 Mamba requires the Python NIXL transceiver (`transceiver_runtime: PYTHON`); V1 routes support periodic snapshots only. diff --git a/tensorrt_llm/_torch/models/modeling_gpt_oss.py b/tensorrt_llm/_torch/models/modeling_gpt_oss.py index d48d454715e1..6ba492929cd5 100644 --- a/tensorrt_llm/_torch/models/modeling_gpt_oss.py +++ b/tensorrt_llm/_torch/models/modeling_gpt_oss.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Literal, Optional +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional import torch from torch import nn @@ -32,6 +32,9 @@ from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, filter_weights, register_auto_model +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + # Use TinyGEMM when the number of tokens is not larger than this threshold MIN_LATENCY_TINYGEMM_NUM_TOKENS = 128 @@ -552,6 +555,27 @@ def forward( @register_auto_model("GptOssForCausalLM") class GptOssForCausalLM(SpecDecOneEngineForCausalLM[Transformer, GptOssConfig]): + @classmethod + def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: + """Select KV cache manager V2 by default. + + GPT-OSS applies a sliding window to every other layer + (see ``AttentionBlock.__init__``), so the KV cache is VSWA: two + distinct attention window sizes. V2 groups layers by lifecycle and + coalesces buffers within each pool group, which sizes the + sliding-window and full-attention pools independently instead of + statically dividing memory between them. + + Users keep full control: an explicit + ``kv_cache_config.use_kv_cache_manager_v2`` otherwise wins over this + default. Two-model speculative decoding is the exception, since V2 + sizes both the target and draft KV cache managers from the full + budget: ``auto`` demotes this default to V1 there and an explicit + ``True`` is rejected, both in + ``llm_utils._resolve_kv_cache_manager_v2_auto``. + """ + return {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + @classmethod def get_preferred_transceiver_runtime( cls, diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 63b360f5584b..2c7b275732cd 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -557,20 +557,61 @@ def _compute_applied(defaults: Dict[str, Any], return _compute_applied(model_defaults_dict, user_overrides) +def _two_model_spec_dec_decoding_type( + llm_args: 'TorchLlmArgs') -> Optional[str]: + """Return the decoding type when a separate draft engine is configured. + + ``has_draft_model()`` is the same predicate py_executor_creator uses to + decide whether to build a separate draft model engine, which is what forces + the second KV cache manager. Returns ``None`` for single-engine runs. + """ + spec_config = llm_args.speculative_config + if spec_config is None: + return None + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + if spec_dec_mode is None or not spec_dec_mode.has_draft_model(): + return None + return getattr(spec_config, "decoding_type", "speculative decoding") + + def _resolve_kv_cache_manager_v2_auto( llm_args: 'TorchLlmArgs', model_defaults_dict: Dict[str, Any], original_setting: Optional[Union[bool, str]] = None) -> bool: """Resolve the KV cache manager auto setting after model defaults are applied. - The transceiver runtime auto setting must be resolved first. In - disaggregated serving, hybrid Mamba V2 requires the Python transceiver with - NIXL, so an incompatible route falls back to V1 unless the user explicitly - selected V2. + A model default of V2 is demoted to V1 for routes V2 cannot serve; an + explicit user value otherwise wins. The compatibility arms are: + + - Disaggregated serving: hybrid Mamba V2 requires the Python transceiver + with NIXL, so any other route falls back to V1. The transceiver runtime + auto setting must be resolved first. + - Two-model speculative decoding: the draft model runs in a separate engine + with its own KV cache manager. ``build_managers`` hands that manager the + target's ``kv_cache_config`` unsplit, and V2 capacity is governed solely + by ``max_gpu_total_bytes``, so both managers size their pools from the + full budget. The model default falls back to V1, and an explicit ``True`` + is rejected rather than deferred to that allocation. + + The fallback only reaches models whose manager class is selected by + ``use_kv_cache_manager_v2``. Models routed to a V2 manager unconditionally + -- sparse attention picks its class from the algorithm alone -- keep a V2 + manager after the demotion, so the arm does not protect them. """ setting = (llm_args.kv_cache_config.use_kv_cache_manager_v2 if original_setting is None else original_setting) if setting != "auto": + if setting is True: + decoding_type = _two_model_spec_dec_decoding_type(llm_args) + if decoding_type is not None: + raise ValueError( + "kv_cache_config.use_kv_cache_manager_v2=True is not " + f"supported with {decoding_type}: the draft model runs in " + "a separate engine and V2 sizes both KV cache managers " + "from the full max_gpu_total_bytes budget instead of " + "partitioning it between them. Set " + "use_kv_cache_manager_v2 to False or 'auto', or use the " + "one-model variant of this decoding mode.") return setting kv_cache_defaults = model_defaults_dict.get("kv_cache_config", {}) @@ -595,6 +636,16 @@ def _resolve_kv_cache_manager_v2_auto( "falling back to V1.", runtime, effective_backend) model_default = False + if model_default: + decoding_type = _two_model_spec_dec_decoding_type(llm_args) + if decoding_type is not None: + logger.info( + "KV cache manager V2 is the model default, but %s runs the " + "draft model in a separate engine and V2 sizes both KV cache " + "managers from the full max_gpu_total_bytes budget; falling " + "back to V1.", decoding_type) + model_default = False + llm_args.kv_cache_config.use_kv_cache_manager_v2 = model_default return model_default diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 234274da3559..5389c8defde4 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -5557,10 +5557,10 @@ def test_eagle3_4gpus(self, v2_kv_cache, moe_backend, one_model, mocker.patch.object(GPQADiamond, "MAX_OUTPUT_LEN", MAX_OUTPUT_LEN) mocker.patch.object(GPQADiamond, "MAX_INPUT_LEN", MAX_INPUT_LEN) - if v2_kv_cache and not one_model and overlap_scheduler: + if v2_kv_cache and not one_model: pytest.skip( - "KVCacheManagerV2 not compatible with two-model overlap scheduling" - ) + "KVCacheManagerV2 sizes the target and draft managers from the " + "same budget, so two-model Eagle3 is rejected") # https://nvbugs/5590408: 2-Model overlap scheduling has accuracy issue pytorch_config = dict(disable_overlap_scheduler=not overlap_scheduler, diff --git a/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py b/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py index 410c50189192..00e3812ed0fb 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py +++ b/tests/unittest/_torch/modeling/test_modeling_gpt_oss.py @@ -18,7 +18,11 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import \ KvCacheConfig as BindingsKvCacheConfig -from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MoeConfig +from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, + KvCacheConfig, MoeConfig) +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs +from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, + apply_model_defaults_to_llm_args) from tensorrt_llm.mapping import Mapping configs = """ @@ -51,6 +55,63 @@ def test_gpt_oss_prefers_python_transceiver() -> None: assert GptOssForCausalLM.get_preferred_transceiver_runtime() == "PYTHON" +def _resolve_gpt_oss_kv_cache_manager_v2(**llm_args_kwargs) -> bool: + """Run GPT-OSS model defaults through the same path model loading uses.""" + llm_args = TorchLlmArgs(model="/tmp/dummy_model", **llm_args_kwargs) + original_setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 + model_defaults = GptOssForCausalLM.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, model_defaults) + return _resolve_kv_cache_manager_v2_auto(llm_args, + model_defaults, + original_setting=original_setting) + + +def test_gpt_oss_model_defaults_select_v2(): + """GPT-OSS is VSWA, so "auto" resolves to KVCacheManagerV2.""" + assert _resolve_gpt_oss_kv_cache_manager_v2() is True + + +@pytest.mark.parametrize("user_setting", [False, True]) +def test_gpt_oss_explicit_setting_wins(user_setting): + """An explicit user value is never overridden by the model default.""" + assert _resolve_gpt_oss_kv_cache_manager_v2(kv_cache_config=KvCacheConfig( + use_kv_cache_manager_v2=user_setting)) is user_setting + + +def test_gpt_oss_two_model_eagle3_falls_back_to_v1(): + """Two-model Eagle3 builds a separate draft engine with its own KV cache + manager, and V2 sizes both from the full budget, so the model default is + demoted to V1.""" + assert _resolve_gpt_oss_kv_cache_manager_v2( + speculative_config=Eagle3DecodingConfig( + max_draft_len=3, + speculative_model="/tmp/dummy_eagle_model", + eagle3_one_model=False)) is False + + +def test_gpt_oss_explicit_v2_rejects_two_model_eagle3(): + """The demotion above only applies to the model default. An explicit + request for the same unsupported combination is rejected rather than + silently honored.""" + with pytest.raises(ValueError, + match="use_kv_cache_manager_v2=True is not supported"): + _resolve_gpt_oss_kv_cache_manager_v2( + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=True), + speculative_config=Eagle3DecodingConfig( + max_draft_len=3, + speculative_model="/tmp/dummy_eagle_model", + eagle3_one_model=False)) + + +def test_gpt_oss_one_model_eagle3_keeps_v2(): + """One-model Eagle3 shares the target engine, so V2 still applies.""" + assert _resolve_gpt_oss_kv_cache_manager_v2( + speculative_config=Eagle3DecodingConfig( + max_draft_len=3, + speculative_model="/tmp/dummy_eagle_model", + eagle3_one_model=True)) is True + + def dump_config_json(dst_dir): if os.path.exists(dst_dir): shutil.rmtree(dst_dir)