[None][feat] Opt GPT-OSS in to KV cache manager V2 by default

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.

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,

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
 

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,

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)