[None][feat] Opt GPT-OSS in to KV cache manager V2 by default - #16942
Conversation
a7d7c85 to
30b9064
Compare
|
/bot run |
WalkthroughGPT-OSS now enables KV cache manager V2 by default. Explicit settings remain effective. Two-model Eagle3 speculative decoding falls back to V1, while explicit V2 raises an error. Documentation describes these rules and updates hybrid Mamba guidance. ChangesGPT-OSS KV cache manager selection
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant TorchLlmArgs
participant GptOssForCausalLM
participant KVCacheResolver
participant KVCacheManager
TorchLlmArgs->>GptOssForCausalLM: provide GPT-OSS model arguments
GptOssForCausalLM->>KVCacheResolver: provide V2 model default
KVCacheResolver->>KVCacheResolver: inspect speculative decoding configuration
KVCacheResolver->>KVCacheManager: select V1 for two-model Eagle3
KVCacheResolver->>KVCacheManager: retain V2 otherwise
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/modeling/test_modeling_gpt_oss.py (1)
54-54: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the return annotation to the new test.
Change
test_gpt_oss_model_defaultsto returnNone, consistent with the Python guidelines and the neighboring test.Proposed fix
-def test_gpt_oss_model_defaults(): +def test_gpt_oss_model_defaults() -> None:As per coding guidelines,
**/*.pyrequires annotations for every function.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_modeling_gpt_oss.py` at line 54, Update the test_gpt_oss_model_defaults function signature to include a None return annotation, matching the neighboring test and the repository’s Python annotation guidelines.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/modeling/test_modeling_gpt_oss.py`:
- Around line 54-68: Register test_eagle3_guided_decoding_4gpus and
test_eagle3_1gpu from the GPT-OSS modeling test suite in the appropriate test-db
CI lists, matching the existing registration for test_eagle3_2gpus;
alternatively, explicitly document that both tests are intentionally QA-only.
Ensure their CI coverage is represented consistently across the relevant test
selectors.
---
Nitpick comments:
In `@tests/unittest/_torch/modeling/test_modeling_gpt_oss.py`:
- Line 54: Update the test_gpt_oss_model_defaults function signature to include
a None return annotation, matching the neighboring test and the repository’s
Python annotation guidelines.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 872c1991-45d1-4e88-9df2-f804bb80113b
📒 Files selected for processing (4)
docs/source/features/kvcache.mdtensorrt_llm/_torch/models/modeling_gpt_oss.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/unittest/_torch/modeling/test_modeling_gpt_oss.py
|
PR_Github #63370 [ run ] triggered by Bot. Commit: |
|
PR_Github #63370 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63445 [ run ] triggered by Bot. Commit: |
|
PR_Github #63445 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #63603 [ run ] triggered by Bot. Commit: |
|
PR_Github #63603 [ run ] completed with state
|
30b9064 to
ffdb7af
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63739 [ run ] triggered by Bot. Commit: |
ffdb7af to
e26073e
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/unittest/_torch/modeling/test_modeling_gpt_oss.py (2)
58-66: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAssert the field consumed by model loading.
_resolve_kv_cache_manager_v2_autoreturns a boolean and updatesllm_args.kv_cache_config.use_kv_cache_manager_v2at Line 620. This helper returns only the boolean. A regression could keep the return value correct while leaving the configuration field stale. Assert both values.Proposed test assertion
- return _resolve_kv_cache_manager_v2_auto(llm_args, - model_defaults, - original_setting=original_setting) + resolved = _resolve_kv_cache_manager_v2_auto( + llm_args, model_defaults, original_setting=original_setting) + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 == resolved + return resolved🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_modeling_gpt_oss.py` around lines 58 - 66, Update _resolve_gpt_oss_kv_cache_manager_v2 to assert that llm_args.kv_cache_config.use_kv_cache_manager_v2 matches the boolean returned by _resolve_kv_cache_manager_v2_auto, while preserving the helper’s return value so tests validate both the result and the field consumed by model loading.
74-78: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover explicit overrides with two-model Eagle3.
test_gpt_oss_explicit_setting_winsdoes not setspeculative_config. The Eagle3 tests cover only the default"auto"setting. Add explicitTrueandFalsecases witheagle3_one_model=False. This protects the contract that the compatibility fallback applies only to model defaults.Proposed coverage
+@pytest.mark.parametrize("user_setting", [False, True]) +def test_gpt_oss_explicit_setting_wins_for_two_model_eagle3(user_setting): + assert _resolve_gpt_oss_kv_cache_manager_v2( + kv_cache_config=KvCacheConfig( + use_kv_cache_manager_v2=user_setting), + speculative_config=Eagle3DecodingConfig( + max_draft_len=3, + speculative_model="/tmp/dummy_eagle_model", + eagle3_one_model=False)) is user_setting🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_modeling_gpt_oss.py` around lines 74 - 78, Extend test_gpt_oss_explicit_setting_wins to parameterize explicit True and False KV-cache settings together with speculative_config configured for Eagle3 with eagle3_one_model=False. Assert _resolve_gpt_oss_kv_cache_manager_v2 preserves each explicit user value, ensuring the compatibility fallback only affects model defaults.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/modeling/test_modeling_gpt_oss.py`:
- Around line 58-66: Update _resolve_gpt_oss_kv_cache_manager_v2 to assert that
llm_args.kv_cache_config.use_kv_cache_manager_v2 matches the boolean returned by
_resolve_kv_cache_manager_v2_auto, while preserving the helper’s return value so
tests validate both the result and the field consumed by model loading.
- Around line 74-78: Extend test_gpt_oss_explicit_setting_wins to parameterize
explicit True and False KV-cache settings together with speculative_config
configured for Eagle3 with eagle3_one_model=False. Assert
_resolve_gpt_oss_kv_cache_manager_v2 preserves each explicit user value,
ensuring the compatibility fallback only affects model defaults.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: abd1da48-53f4-4df9-8a86-587a5128dff2
📒 Files selected for processing (4)
docs/source/features/kvcache.mdtensorrt_llm/_torch/models/modeling_gpt_oss.pytensorrt_llm/llmapi/llm_utils.pytests/unittest/_torch/modeling/test_modeling_gpt_oss.py
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/source/features/kvcache.md
- tensorrt_llm/_torch/models/modeling_gpt_oss.py
|
Re-reviewed at spec_config = llm_args.speculative_config
if model_default and spec_config is not None:
spec_dec_mode = getattr(spec_config, "spec_dec_mode", None)
if spec_dec_mode is not None and spec_dec_mode.has_draft_model():
...
model_default = FalseGating on That closes my objection: the unsupported combination is now handled in the resolver instead of by pinning Not approving, only because this is still a default flip on a shipped model and it has no approvals yet — that wants someone with KV-cache-manager ownership on it, not me first. No outstanding findings from my side. One small thing, non-blocking: |
|
/bot run --disable-fail-fast |
|
PR_Github #63988 [ run ] triggered by Bot. Commit: |
|
PR_Github #63988 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #64178 [ run ] triggered by Bot. Commit: |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
The GPT-OSS opt-in itself is a one-liner into existing infrastructure, and the validation campaign in the description is more than I'd normally ask for on a default flip.
Two things on the description: it doesn't mention the new two-model-spec-dec guard in llm_utils.py at all, even though that's the larger half of the diff and changes resolution for NemotronH, Qwen3-Next, and DeepSeek-V4 as well as GPT-OSS. Worth calling out explicitly so those model owners see it. Second, a default-manager flip with a disclosed throughput trade-off is a feature, not a chore — this should carry a JIRA ticket rather than [None].
|
PR_Github #64178 [ run ] completed with state |
yechank-nvidia
left a comment
yechank-nvidia
left a comment
There was a problem hiding this comment.
LGTM on modeling-side
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 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 <yueh.ting.chen@gmail.com> Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
e26073e to
b98a35f
Compare
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/llmapi/llm_utils.py`:
- Around line 642-646: Update the logger.info call in the KV cache manager
fallback path to preformat the message with decoding_type before passing it to
the logger, ensuring no literal "%s" remains and no separate formatting argument
is supplied.
In `@tests/unittest/_torch/modeling/test_modeling_gpt_oss.py`:
- Line 58: Annotate the changed functions in the GPT-OSS KV-cache tests: give
_resolve_gpt_oss_kv_cache_manager_v2’s llm_args_kwargs a precise
keyword-argument type, add -> None to the test functions, and annotate the
parametrized test’s user_setting parameter as bool. Apply these annotations
consistently across the functions in the referenced range.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: df21e209-a898-4310-976c-4ae2f3bc87be
📒 Files selected for processing (5)
docs/source/features/kvcache.mdtensorrt_llm/_torch/models/modeling_gpt_oss.pytensorrt_llm/llmapi/llm_utils.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/unittest/_torch/modeling/test_modeling_gpt_oss.py
🚧 Files skipped from review as they are similar to previous changes (2)
- docs/source/features/kvcache.md
- tensorrt_llm/_torch/models/modeling_gpt_oss.py
|
/bot skip --comment "The MR previously passed the full pre-merge CI before. The latest update is regards to guarding two-model runs safely and message logging them explicitly. We should not need to run the full CI again and can land the MR." |
|
PR_Github #64998 [ skip ] triggered by Bot. Commit: |
|
PR_Github #64998 [ skip ] completed with state |
…#16942) Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com> Signed-off-by: Xin He (SW-GPU) <200704525+xinhe-nv@users.noreply.github.com>
9 participants
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()(tensorrt_llm/_torch/models/modeling_gpt_oss.py:558) 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:3799).ModelLoadermerges model defaults (tensorrt_llm/_torch/pyexecutor/model_loader.py:432-440), 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 —model_loader.py:455-456passes the pre-merge value asoriginal_setting.llm_utils.py:586-596) falls back to V1 unless the route isbackend=NIXL+transceiver_runtime=PYTHON. GPT-OSS already declaresget_preferred_transceiver_runtime() -> "PYTHON"(modeling_gpt_oss.py:577-582), so the NIXL route keeps V2.GPT-OSS joins hybrid Mamba (NemotronH, Qwen3-Next, Qwen3.5) and DeepSeek-V4 as models that select V2 by default.
Relationship to #14047
#14047 (KVCacheManagerV2 C++ translation) has landed, so
TLLM_KV_CACHE_MANAGER_V2_BACKENDnow defaults tocpp(tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py:23). This PR therefore moves GPT-OSS onto the C++ V2 core. There is no code dependency on #14047 — this is purely the manager-selection flip — but the validation below was run with bothcppandpythonV2 backends so a regression can be attributed to the manager switch versus the language switch.Validation
Campaign on GPT-OSS-120B (MXFP4), B200 + GB200, three arms per cell (
v2cpp/v2py/v1), paired within-node:Accuracy — 50/50 cells PASS, 0 FAIL. The highest-risk cell,
test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model](VSWA + block reuse + Eagle3 draft tokens together), passed at 67.172 vs a 55.734 threshold with provenance confirming the C++ backend was active and no silent fallback toKVCacheManager.test_w4_1gpu,test_w4_4gpus,test_w4_2gpus,test_eagle3_4gpuspassed across arms; the V2 nanobind suite and thekvCacheManagerV2C++ gtests passed on both platforms.Performance. Of the GPT-OSS cases that ship in perf-sanity CI:
dep4_1k8k(TP4+ADP): +15.5% ± 0.88 throughput, mean TTFT −91% (72.3 s → 6.5 s), n=3.tep2_1k8k: −3.4% ± 3.80 over n=8 (95% CI [−6.6%, −0.2%]) — within the 5% post-merge gate at the mean, but the interval reaches past it. This is the noisiest operating point measured (v1 run-to-run CV 3.10% there versus 0.30% one rung up); an earlier n=1 reading of −8.58% did not reproduce.tp1_mtp0_8k1k,tp2_1k8k) landed within ±0.51%.Known perf trade-off, disclosed for reviewers. V1 and V2 differ in how aggressively they admit requests when the batch ceiling binds. When both converge on the same effective batch, all arms land within ~1% and the manager is off the critical path. When the ceiling binds, V2 fills the batch and V1 self-limits, which buys V2 a 10–20× better TTFT — but whether the larger batch pays in throughput flips on attention-DP. On a synthetic ladder rung not in CI (
max_batch_size=1024, attention-DP off), v2cpp measured −11.3% ± 5.01 and v2py −22.2% ± 4.55 (n=3, wide interval). This is a behavioural difference between the managers, not port overhead — v2py regresses further than v2cpp — and it predates this PR; the flip is what exposes GPT-OSS to it. Reuse (prefix caching) performance is not covered: all perf arms ranenable_block_reuse=False.Known limitation: two-model Eagle3
V2 does not split the KV cache budget between the target and draft managers.
_needs_gpu_kv_cache_budget_splitreturns the one-model predicate under V2 (tensorrt_llm/_torch/pyexecutor/_util.py:1809-1818), the host-budget split is skipped forv2_two_model(_util.py:1853-1861), and the two-model path asserts the draft config is unsplit (_util.py:1886-1890) before handing the draft manager the target's config.Three GPT-OSS Eagle3 accuracy tests parametrize
one_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_4gpus(tests/integration/defs/accuracy/test_llm_api_pytorch.py:5531) andtest_eagle3_vswa_reuse_4gpus, both of which already parametrizev2_kv_cacheexplicitly and are unaffected.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 is covered generically intests/unittest/llmapi/test_llm_args.py:488-563, over exactly the dict GPT-OSS returns: model default adopted for both implicit and explicit"auto", unknown model falls back toFalse, explicit user value wins in both directions, and the disagg route falls back to V1 forNIXL/CPP,UCX,MPIwhile keeping V2 forNIXL/PYTHON.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. The remaining GPT-OSS tests that set no flag (test_w4_1gpu_piecewise_cuda_graph,test_dflash,test_w4a16,test_w4_2gpus,test_w4_chunked_prefill,test_w4_4gpus_online_eplb, the disagg/serve suites and perf-sanity cases) flip to V2 with this PR — that is the intended blast radius and what CI on this PR measures.Docs:
docs/source/features/kvcache.mdgains a "Selecting the KV Cache Manager" section, since the V2 default was previously described only in the hybrid Mamba section. It also notes that Gemma4 hybrid and sparse-attention models are routed to V2 unconditionally (_util.py:83-88), so readers do not infer that an unlisted model always gets V1.Local validation:
pre-commitclean against the PR diff (isort / yapf / autoflake / codespell / baseline-gated ruff / test-list AST validation).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.Dev Engineer Review
KVCacheManagerV2by default throughget_model_defaults.QA Engineer Review
test_gpt_oss_model_defaults_select_v2test_gpt_oss_explicit_setting_winstest_gpt_oss_two_model_eagle3_falls_back_to_v1test_gpt_oss_explicit_v2_rejects_two_model_eagle3test_gpt_oss_one_model_eagle3_keeps_v2test-db/orqa/entries were found.