From b98a35f5c3856030a391fbb791a79cbe6cca8702 Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Tue, 28 Jul 2026 17:29:24 +0900 Subject: [PATCH] [None][feat] Opt GPT-OSS in to KV cache manager V2 by default GPT-OSS applies a sliding window to every other layer, so its KV cache is VSWA with two distinct attention window sizes. KVCacheManagerV2 groups layers by lifecycle and coalesces buffers per pool group, sizing the sliding-window and full-attention pools independently instead of statically dividing memory between them. Add GptOssForCausalLM.get_model_defaults() returning kv_cache_config.use_kv_cache_manager_v2=True. This is consumed by the per-model auto-selection path added in #15823: the model default applies only when the user leaves the flag at the "auto" sentinel, and an explicit user value otherwise wins. Two-model speculative decoding is a known exception. The draft model runs in a separate engine with its own KV cache manager, and the budget split is skipped for that route under V2 (_util.py:1854), so both managers would size their pools from the full max_gpu_total_bytes and double-allocate it. Add a compatibility arm to _resolve_kv_cache_manager_v2_auto, symmetric with the existing disaggregated one: when the model default is V2 and the speculative config reports spec_dec_mode.has_draft_model() -- the same predicate py_executor_creator uses to decide whether to build a draft engine -- log and fall back to V1. Keying the guard on the resolver rather than on individual call sites keeps GPT-OSS + Eagle3 with eagle3_one_model=False working on the default "auto" setting. The arm is keyed on the draft engine rather than on the model class, so it also covers NemotronH, Qwen3-Next and DeepSeek-V4, which default to V2 as well. It does not reach models routed to a V2 manager unconditionally: the sparse-attention path picks its class from the algorithm alone and _is_kv_cache_manager_v2 is derived from the resolved class, not the flag. That gap predates this change and is recorded in the resolver docstring. Since the resolver returns early for any non-"auto" setting, an explicit use_kv_cache_manager_v2=True previously bypassed the arm entirely and proceeded into the allocation the arm exists to avoid. Reject that combination with a ValueError instead, consistent with the other explicit use_kv_cache_manager_v2 conflicts raised from get_kv_cache_manager_cls, and state the rule in kvcache.md as enforced behavior rather than guidance. test_eagle3_4gpus skips all v2_kv_cache + two_model cells accordingly, matching the guard test_eagle3_vswa_reuse_4gpus already used; no CI or QA list enrolls such a cell. Also documents KV cache manager selection in docs/source/features/kvcache.md, which previously described the V2 default only in the hybrid Mamba section. Co-Authored-By: Yueh-Ting Chen Signed-off-by: Yueh-Ting Chen --- docs/source/features/kvcache.md | 33 +++++++++- .../_torch/models/modeling_gpt_oss.py | 26 +++++++- tensorrt_llm/llmapi/llm_utils.py | 59 +++++++++++++++-- .../defs/accuracy/test_llm_api_pytorch.py | 6 +- .../_torch/modeling/test_modeling_gpt_oss.py | 63 ++++++++++++++++++- 5 files changed, 176 insertions(+), 11 deletions(-) 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)