Skip to content

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

Draft
eopXD wants to merge 1 commit into
NVIDIA:mainfrom
eopXD:feat/gpt-oss-kvcache-v2-default-main
Draft

[None][feat] Opt GPT-OSS in to KV cache manager V2 by default#16942
eopXD wants to merge 1 commit into
NVIDIA:mainfrom
eopXD:feat/gpt-oss-kvcache-v2-default-main

Conversation

@eopXD

@eopXD eopXD commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

Description

Opts GPT-OSS in to KVCacheManagerV2 by default.

GPT-OSS applies a sliding window to every other layer (AttentionBlock.__init__, tensorrt_llm/_torch/models/modeling_gpt_oss.py:96), so its KV cache is VSWA with two distinct attention window sizes. V2 groups layers by lifecycle and coalesces buffers within each pool group, sizing the sliding-window and full-attention pools independently instead of statically dividing memory between them.

The change adds GptOssForCausalLM.get_model_defaults() returning kv_cache_config.use_kv_cache_manager_v2=True. This plugs into the per-model auto-selection infrastructure added in #15823:

  • use_kv_cache_manager_v2 defaults to the "auto" sentinel (tensorrt_llm/llmapi/llm_args.py:3766).
  • ModelLoader merges model defaults, then _resolve_kv_cache_manager_v2_auto (tensorrt_llm/llmapi/llm_utils.py:560) collapses "auto" to a concrete bool.
  • An explicit user use_kv_cache_manager_v2 always wins over the model default.
  • In disaggregated serving, the existing guard falls back to V1 unless the route is backend=NIXL + transceiver_runtime=PYTHON. GPT-OSS already declares get_preferred_transceiver_runtime() -> "PYTHON", so the NIXL route keeps V2.

GPT-OSS joins hybrid Mamba (NemotronH, Qwen3-Next) and DeepSeek-V4 as models that select V2 by default.

Relationship to #14047

This is being validated alongside #14047 (KVCacheManagerV2 C++ translation), but it has no code dependency on it and is proposed against main on its own. On top of #14047 the flip additionally moves GPT-OSS onto the C++ V2 core, since TLLM_KV_CACHE_MANAGER_V2_BACKEND defaults to cpp there — accuracy/perf runs are being done with both backend settings so a regression can be attributed to the manager switch versus the language switch.

Known limitation: two-model Eagle3

V2 does not split the KV cache budget between the target and draft managers (tensorrt_llm/_torch/pyexecutor/_util.py:1681-1689, assert at :1708). Three GPT-OSS Eagle3 accuracy tests parametrize one_model=[True, False] without setting the flag, so their two_model variants would have silently picked up V2:

  • TestGPTOSS::test_eagle3_guided_decoding_4gpus
  • TestGPTOSS::test_eagle3_2gpus
  • TestGPTOSS::test_eagle3_1gpu

These now pass use_kv_cache_manager_v2="auto" if one_model else False, which pins V1 for two_model (preserving current behavior) while leaving one_model on "auto" so it exercises the new default. This mirrors the existing pytest.skip("KVCacheManagerV2 not compatible with two-model overlap scheduling") in test_eagle3_4gpus and test_eagle3_vswa_reuse_4gpus.

DFlash is unaffected: SpeculativeDecodingMode.has_draft_model() excludes DFLASH, so no separate draft engine is created.

Scope is intentionally limited to the manager selection. Other V2 tuning knobs GPT-OSS could take (tokens_per_block, enable_swa_scratch_reuse — both set by DeepSeek-V4) are deliberately left out pending perf measurement.

Test Coverage

New:

  • tests/unittest/_torch/modeling/test_modeling_gpt_oss.py::test_gpt_oss_model_defaults — asserts the defaults dict, mirroring test_modeling_deepseekv4.py::test_deepseek_v4_model_defaults.

Existing coverage that exercises this path:

  • "auto"-resolution precedence (explicit user value wins, unknown model falls back to False) is covered generically in tests/unittest/llmapi/test_llm_args.py:451-487.
  • TestGPTOSS::test_w4_1gpu, test_w4_4gpus, test_eagle3_4gpus, test_eagle3_vswa_reuse_4gpus already parametrize v2_kv_cache explicitly; they pass the flag, so they are unaffected by the default change and continue to cover both managers.
  • The one_model variants of the three tests modified here newly run through the "auto" resolution path.

Validation status: formatting (yapf/isort) and baseline-gated lint pass locally. Draft until the GPT-OSS accuracy (GSM8K / GPQA) and perf runs complete; results will be posted here.

PR Checklist

Please review the following before submitting your PR:

  • PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

  • PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

  • Test cases are provided for new code paths (see test instructions)

  • If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

  • Any new dependencies have been scanned for license and vulnerabilities

  • CODEOWNERS updated if ownership changes

  • Documentation updated as needed

  • Update tava architecture diagram if there is a significant design change in PR.

  • The reviewers assigned automatically/manually are appropriate for the PR.

  • Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

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 NVIDIA#15823: the model default applies
only when the user leaves the flag at the "auto" sentinel, and an explicit
user value still wins.

Two-model Eagle3 is a known exception -- V2 does not split the KV cache
budget between the target and draft managers -- so the three GPT-OSS Eagle3
accuracy tests that parametrize one_model now pin V1 for their two_model
variants and leave one_model on "auto" to exercise the new default.

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 <yueh.ting.chen@gmail.com>
Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant