[None][feat] Opt GPT-OSS in to KV cache manager V2 by default - #16942
Draft
eopXD wants to merge 1 commit into
Draft
[None][feat] Opt GPT-OSS in to KV cache manager V2 by default#16942eopXD wants to merge 1 commit into
eopXD wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Opts GPT-OSS in to
KVCacheManagerV2by 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()returningkv_cache_config.use_kv_cache_manager_v2=True. This plugs into the per-model auto-selection infrastructure added in #15823:use_kv_cache_manager_v2defaults to the"auto"sentinel (tensorrt_llm/llmapi/llm_args.py:3766).ModelLoadermerges model defaults, then_resolve_kv_cache_manager_v2_auto(tensorrt_llm/llmapi/llm_utils.py:560) collapses"auto"to a concrete bool.use_kv_cache_manager_v2always wins over the model default.backend=NIXL+transceiver_runtime=PYTHON. GPT-OSS already declaresget_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
mainon its own. On top of #14047 the flip additionally moves GPT-OSS onto the C++ V2 core, sinceTLLM_KV_CACHE_MANAGER_V2_BACKENDdefaults tocppthere — 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 parametrizeone_model=[True, False]without setting the flag, so theirtwo_modelvariants would have silently picked up V2:TestGPTOSS::test_eagle3_guided_decoding_4gpusTestGPTOSS::test_eagle3_2gpusTestGPTOSS::test_eagle3_1gpuThese now pass
use_kv_cache_manager_v2="auto" if one_model else False, which pins V1 fortwo_model(preserving current behavior) while leavingone_modelon"auto"so it exercises the new default. This mirrors the existingpytest.skip("KVCacheManagerV2 not compatible with two-model overlap scheduling")intest_eagle3_4gpusandtest_eagle3_vswa_reuse_4gpus.DFlash is unaffected:
SpeculativeDecodingMode.has_draft_model()excludesDFLASH, 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, mirroringtest_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 toFalse) is covered generically intests/unittest/llmapi/test_llm_args.py:451-487.TestGPTOSS::test_w4_1gpu,test_w4_4gpus,test_eagle3_4gpus,test_eagle3_vswa_reuse_4gpusalready parametrizev2_kv_cacheexplicitly; they pass the flag, so they are unaffected by the default change and continue to cover both managers.one_modelvariants 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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.