[TRTLLM-15264][fix] Reject non-Python transceiver routes for Kimi K3 disaggregated serving - #17480
[TRTLLM-15264][fix] Reject non-Python transceiver routes for Kimi K3 disaggregated serving#17480brnguyen2 wants to merge 7 commits into
Conversation
|
/bot run |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughKimi K3 now declares Python as its preferred transceiver runtime. Disaggregated serving validates the Python NIXL route and rejects unsupported configurations. Tests cover preference resolution, invalid routes, valid cache-manager selection, and an integration test waiver. ChangesKimi K3 transceiver routing
Estimated code review effort: 2 (Simple) | ~10 minutes Sequence Diagram(s)sequenceDiagram
participant KimiLinearForCausalLM
participant DisaggregatedServing
participant MixedMambaHybridCacheManager
KimiLinearForCausalLM->>DisaggregatedServing: provide preferred runtime "PYTHON"
DisaggregatedServing->>DisaggregatedServing: validate runtime "PYTHON" and backend "NIXL"
DisaggregatedServing->>MixedMambaHybridCacheManager: select valid route
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 (2)
tests/unittest/_torch/executor/test_mamba_cache_manager.py (2)
796-798: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
cache_transceiver_config.Use
CacheTransceiverConfig | Nonefor the parameter type.Proposed type annotation
def test_kimi_disagg_rejects_non_python_transceiver_route( - monkeypatch: pytest.MonkeyPatch, cache_transceiver_config + monkeypatch: pytest.MonkeyPatch, + cache_transceiver_config: CacheTransceiverConfig | None, ) -> None:As per coding guidelines, “Annotate 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/executor/test_mamba_cache_manager.py` around lines 796 - 798, Update the test_kimi_disagg_rejects_non_python_transceiver_route function signature to annotate cache_transceiver_config as CacheTransceiverConfig | None, using the existing CacheTransceiverConfig symbol and preserving the current test behavior.Source: Coding guidelines
778-832: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
cache_transceiver_config.
test_kimi_disagg_rejects_non_python_transceiver_routeleaves this parameter untyped. UseCacheTransceiverConfig | None.Test coverage summary: sufficient. Added tests cover Kimi defaults, runtime preference, rejected routes, and the Python NIXL route. The test file is included through
unittest/_torch/executorentries in the test-db lists. No tests were modified or removed.🤖 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/executor/test_mamba_cache_manager.py` around lines 778 - 832, Annotate the cache_transceiver_config parameter in test_kimi_disagg_rejects_non_python_transceiver_route with the union type CacheTransceiverConfig | None, preserving the existing parametrized values and test behavior.Source: Path instructions
🤖 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/executor/test_mamba_cache_manager.py`:
- Around line 786-795: Add a parameterized case in the existing
cache_transceiver_config matrix for backend="UCX" with
transceiver_runtime="PYTHON", and assign it a distinct id. Ensure the test
exercises the rejection path in the relevant executor configuration validation
while retaining the existing NIXL and C++ cases.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_mamba_cache_manager.py`:
- Around line 796-798: Update the
test_kimi_disagg_rejects_non_python_transceiver_route function signature to
annotate cache_transceiver_config as CacheTransceiverConfig | None, using the
existing CacheTransceiverConfig symbol and preserving the current test behavior.
- Around line 778-832: Annotate the cache_transceiver_config parameter in
test_kimi_disagg_rejects_non_python_transceiver_route with the union type
CacheTransceiverConfig | None, preserving the existing parametrized values and
test behavior.
🪄 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: dab9c8ab-cbe8-4e4b-b824-26e4423c54e2
📒 Files selected for processing (3)
tensorrt_llm/_torch/models/modeling_kimi_linear.pytensorrt_llm/_torch/pyexecutor/_util.pytests/unittest/_torch/executor/test_mamba_cache_manager.py
|
PR_Github #65138 [ run ] triggered by Bot. Commit: |
|
/bot run |
21d1124 to
4ef04d0
Compare
|
PR_Github #65148 [ run ] triggered by Bot. Commit: |
|
PR_Github #65138 [ run ] completed with state |
|
PR_Github #65148 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65290 [ run ] triggered by Bot. Commit: |
|
PR_Github #65290 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65353 [ run ] triggered by Bot. Commit: |
|
PR_Github #65353 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65613 [ run ] completed with state |
|
/bot run |
|
PR_Github #65997 [ run ] triggered by Bot. Commit: |
|
PR_Github #65997 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66042 [ run ] triggered by Bot. Commit: |
|
PR_Github #66042 [ run ] completed with state
|
|
/bot run |
5f2ff40 to
7f57f00
Compare
|
PR_Github #66080 [ run ] triggered by Bot. Commit: |
|
PR_Github #66080 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66123 [ run ] triggered by Bot. Commit: |
…disaggregated serving
A K3 disaggregated deployment that left transceiver_runtime at its
default ('auto') silently resolved to the C++ transceiver: the hybrid
manager routing returned CppMambaHybridCacheManager with no error, even
though only the Python NIXL transceiver (KvCacheTransceiverV2) can move
the KDA recurrent state. The C++ path is unvalidated for K3 and would
serve wrong results rather than fail.
Two layers of fix:
- KimiLinearForCausalLM now overrides get_preferred_transceiver_runtime
to 'PYTHON' (the Qwen3-Next pattern), so model loading resolves
'auto' to the Python transceiver when the backend is NIXL.
- get_kv_cache_manager_cls rejects any K3 disagg route that is not
PYTHON+NIXL with an actionable error, covering explicit non-Python
runtimes and paths that skip model-default resolution (e.g.
AutoDeploy, where 'auto' falls back to the C++ runtime).
Adds routing unit tests: the preference override, rejection of
no-config / unresolved-auto / explicit-CPP / UCX routes, and the
positive PYTHON+NIXL route to MixedMambaHybridCacheManager.
Addresses review feedback on NVIDIA#17334.
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…er-route test Review follow-up: backend=UCX with transceiver_runtime=PYTHON passes the runtime check and is rejected only by the NIXL-backend check; without this case the matrix cannot detect removal of that check. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…_Instruct_2507::test_skip_softmax_attention[target_sparsity_0.9-fp8kv=True] (pre-existing failure, tracking bug pending) Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…r does not truncate it The test-list parser treats everything after a # as a comment, so a # inside a waive reason truncates the line and leaves unbalanced parentheses, breaking collection. Drop the # from the PR reference. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…ention waive Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
7f57f00 to
6d2d615
Compare
|
PR_Github #66201 [ run ] triggered by Bot. Commit: |
|
PR_Github #66123 [ run ] completed with state |
|
PR_Github #66201 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66239 [ run ] triggered by Bot. Commit: |
…_nvfp4[adp4_cutedsl] waive The waive cited nvbugs/6535767, which tracks a different, already-fixed (V2C) test (TestQwen3NextInstruct::test_nvfp4). The correct open tracking bug for the TestQwen3_5_397B_A17B::test_nvfp4 family is nvbugs/6535779 (Open - To fix), under which its sibling params (adp4_trtllm, tep4_trtllm) are already waived on main. Transposed digit (767 -> 779). Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
PR_Github #66239 [ run ] completed with state |
5 participants
Description
Follow-up to review feedback on #17334 (#17334 (comment)).
A Kimi K3 disaggregated deployment that left
cache_transceiver_config.transceiver_runtimeat its default (auto) silently resolved to the C++ transceiver: the hybrid manager routing returnedCppMambaHybridCacheManagerwith no error, even though only the Python NIXL transceiver (KvCacheTransceiverV2) can move the KDA recurrent state. The C++ path is unvalidated for K3 and would serve wrong results rather than fail. (The shipped K3 example configs all pintransceiver_runtime: PYTHON, so this affects hand-rolled configs only.)Two layers of fix:
KimiLinearForCausalLMnow overridesget_preferred_transceiver_runtimeto"PYTHON"(the Qwen3-Next pattern), so model loading resolvesautoto the Python transceiver when the effective backend is NIXL.get_kv_cache_manager_clsrejects any K3 disagg route that is not PYTHON+NIXL with an actionable error, covering explicit non-Python runtimes and paths that skip model-default resolution (e.g. AutoDeploy, whereautofalls back to the C++ runtime).Test Coverage
New cases in
tests/unittest/_torch/executor/test_mamba_cache_manager.py:test_kimi_preferred_transceiver_runtimetest_kimi_disagg_rejects_non_python_transceiver_route(no config / unresolved auto / explicit CPP / UCX+CPP)test_kimi_disagg_python_nixl_routes_to_mixed_managerFull file run on Blackwell hardware: 125 passed, 4 xfailed.
PR Checklist
[TRTLLM-15264][fix]conventionDev Engineer Review
PYTHONtransceiver.PYTHONtransceiver with the NIXL backend.auto, C++ runtime, and UCX routes.QA Engineer Review
Added:
test_kimi_preferred_transceiver_runtimetest_kimi_disagg_rejects_non_python_transceiver_routetest_kimi_disagg_python_nixl_routes_to_mixed_managerThese are test-code changes outside
tests/integration/test_lists/. Notest-db/orqa/coverage entries are reported. The full test file passed 125 tests, with 4 expected Blackwell hardware failures.Verdict: needs follow-up
pr-babysitter waive set-bug NVIDIA/TensorRT-LLM#17480 <bug-url>— the waives entry carries a placeholder until then