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

docs/source/features/kvcache.md

@@ -74,6 +74,31 @@ 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, which does not split the KV
+cache budget between the target and draft managers. Under `auto`, a model
+default of V2 falls back to V1 for that combination; setting
+`use_kv_cache_manager_v2: true` explicitly is not supported there.
+
 ### Mamba Snapshot Boundaries
 
 Hybrid Mamba models must retain the recurrent Mamba state together with the
@@ -111,8 +136,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,25 @@ 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`` always wins over this
+        default. Two-model Eagle3 is a known exception -- V2 does not split
+        the KV cache budget between the target and draft managers -- so that
+        combination should be run with the flag set to False explicitly.
+        """
+        return {"kv_cache_config": {"use_kv_cache_manager_v2": True}}
+
     @classmethod
     def get_preferred_transceiver_runtime(
         cls,

tensorrt_llm/llmapi/llm_utils.py

@@ -563,10 +563,16 @@ def _resolve_kv_cache_manager_v2_auto(
         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 always 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, and V2 cannot split the KV cache budget
+      between the two -- both managers would size their pools from the full
+      budget and double-allocate it -- so it falls back to V1.
     """
     setting = (llm_args.kv_cache_config.use_kv_cache_manager_v2
                if original_setting is None else original_setting)
@@ -595,6 +601,22 @@ def _resolve_kv_cache_manager_v2_auto(
                 "falling back to V1.", runtime, effective_backend)
             model_default = False
 
+    spec_config = llm_args.speculative_config
+    if model_default and spec_config is not None:
+        # ``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.
+        spec_dec_mode = getattr(spec_config, "spec_dec_mode", None)
+        if spec_dec_mode is not None and spec_dec_mode.has_draft_model():
+            decoding_type = getattr(spec_config, "decoding_type",
+                                    "speculative decoding")
+            logger.info(
+                "KV cache manager V2 is the model default, but %s runs the "
+                "draft model in a separate engine and V2 cannot split the KV "
+                "cache budget between the target and draft managers; 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/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,49 @@ 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 cannot split the budget between the two, 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_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)