[TRTLLM-15264][test] Kimi K3 disagg review fixups: KDA test geometry, gate docs, example cleanup - #17483
[TRTLLM-15264][test] Kimi K3 disagg review fixups: KDA test geometry, gate docs, example cleanup#17483brnguyen2 wants to merge 6 commits into
Conversation
|
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 KDA disaggregated-serving documentation now describes sharded and replicated layouts. Transfer tests derive production cache parameters and cover matched or heterogeneous TP with attention-DP enabled or disabled. ChangesKDA transfer validation
Estimated code review effort: 3 (Moderate) | ~25 minutes 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: 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 `@tests/unittest/disaggregated/test_kda_mamba_transfer.py`:
- Around line 691-716: Add unittest/disaggregated/test_kda_mamba_transfer.py to
the repository’s integration test registration list so test_kda_transfer and the
other module tests are executed by CI and QA; update the existing test-list
configuration rather than changing the test implementation.
- Around line 527-544: Unset UCX_NET_DEVICES during setup of the KDA loopback
transfer tests, before creating any NIXL transceivers, while preserving the
existing UCX_TLS and TRTLLM_NIXL_NUM_THREADS configuration. Apply this to the
shared setup covering test_kda_hetero_tp_sharded_accepted and test_kda_transfer
so CI-provided device selection cannot affect either test.
🪄 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: 1a1dfd1b-a513-412a-b4d7-7402c910563f
📒 Files selected for processing (8)
examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yamlexamples/kimi_k3/disagg/README.mdexamples/kimi_k3/disagg/ctx_config.yamltensorrt_llm/_torch/disaggregation/native/bounce/config.pytensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.pytensorrt_llm/_torch/disaggregation/transceiver.pytests/unittest/disaggregated/test_kda_mamba_transfer.py
💤 Files with no reviewable changes (1)
- examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml
|
/bot run |
|
PR_Github #65160 [ run ] triggered by Bot. Commit: |
|
PR_Github #65160 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65196 [ run ] triggered by Bot. Commit: |
|
PR_Github #65196 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65220 [ run ] triggered by Bot. Commit: |
|
PR_Github #65220 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65280 [ run ] triggered by Bot. Commit: |
|
PR_Github #65280 [ run ] completed with state
|
|
/bot run |
1da35af to
f6098c7
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. |
|
PR_Github #65322 [ run ] triggered by Bot. Commit: |
|
PR_Github #65322 [ run ] completed with state
|
|
/bot run |
3 similar comments
|
/bot run |
|
/bot run |
|
/bot run |
10c2c7f to
3d7619f
Compare
|
/bot run |
|
PR_Github #66035 [ run ] triggered by Bot. Commit: |
|
PR_Github #66035 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66041 [ run ] triggered by Bot. Commit: |
|
PR_Github #66041 [ run ] completed with state
|
… gate docs, example cleanup Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Review follow-ups: - _create_transceivers now takes enable_attention_dp and sets it on the transceiver Mapping, so the ADP-on KDA cases exercise the production attention-DP control flow (_dp_rank, sync policy, _kv_size_rank_factor) instead of the non-ADP path. The transfer-byte assertion uses rank factor 1 under ADP, matching _kv_size_rank_factor. - Note why UCX_NET_DEVICES must be left untouched: clearing the cluster-provided pin makes UCX open every interface and fail worker creation on unbindable ones. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…test_gpt_oss_two_model_eagle3_falls_back_to_v1 (pre-existing failure, tracking bug pending) Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
3d7619f to
8981a76
Compare
|
PR_Github #66098 [ run ] triggered by Bot. Commit: |
|
PR_Github #66098 [ run ] completed with state
|
…r exists) The waive entry pointed at unittest/_torch/modeling/test_modeling_gpt_oss.py::test_gpt_oss_two_model_eagle3_falls_back_to_v1, which was removed from main as obsolete in NVIDIA#17597 (commit 8398196). The 2-model Eagle3 fallback scenario it asserted was invalidated when NVIDIA#17366 forced 2-model spec-dec to fall back to 1-model, so the test was deleted rather than fixed. Waiving a nonexistent test fails the AST-based validate-test-lists pre-commit hook (check_test_list.py --validate) and gates the whole pipeline. NVBug 6606544 is moot now that the test is gone. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
|
PR_Github #66275 [ run ] triggered by Bot. Commit: |
|
PR_Github #66275 [ run ] completed with state
|
|
/bot run |
|
PR_Github #66292 [ run ] triggered by Bot. Commit: |
…DIA#17584) test_kda_verify_matches_sequential_decode[2-1] was waived under nvbugs/6599150, but that bug is closed-fixed: PR NVIDIA#17584 (merged) fixed the test by initializing dt_bias in the parity test and removed the need for the waive. Main carries no such waive. Drop the re-added line so the now-passing test runs. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
PR_Github #66292 [ run ] completed with state |
|
/bot run |
|
PR_Github #66481 [ run ] triggered by Bot. Commit: |
|
PR_Github #66481 [ run ] completed with state
|
5 participants
Description
Follow-ups from review of the Kimi K3 disaggregated-serving work, bundled into one change:
_create_kda_managersintest_kda_mamba_transfer.pybuilds its Mamba cache params viaextract_mamba_kv_cache_paramswith unscaled global head counts, letting the manager's own attention-DP gate do the sharding — matching how the model runtime head-shards KDA state across TP when attention-DP is off.validate_peer_compatiblecorrectly accepts hetero-TP peers;test_kda_hetero_tp_rejectedbecametest_kda_hetero_tp_sharded_accepted, asserting concrete per-rank slot sizes before acceptance. New transfer parametrizations cover matched TP2 (attention-DP on and off) and hetero 2->4 / 4->2 with bitwise comparison.peer.pydocstring/error text, the K3 disagg example README, andctx_config.yaml— the claim that hetero ctx/gen TP is rejected for K3 predated the head-sharded layout. Matched geometry remains the only configuration validated end-to-end.bounce/config.pyexplaining what each admission control covers.transceiver.py_slice_num_bytes(sums all pool views) andbounce/impl.pyblock_bytes_per_group(pool 0 only), pointing at the TRTLLM-15194 sizing gap without changing behavior.Test Coverage
On Blackwell hardware:
tests/unittest/disaggregated/test_kda_mamba_transfer.py18 passed;tests/unittest/disaggregated/test_bounce.py75 passed. Product-code changes are comment/docstring/error-text only.PR Checklist
[TRTLLM-15264][test]conventionDev Engineer Review
extract_mamba_kv_cache_params.QA Engineer Review
Modified test code:
tests/unittest/disaggregated/test_kda_mamba_transfer.pytest_kda_hetero_tp_rejected()withtest_kda_hetero_tp_sharded_accepted().test_kda_transfer(ctx_tp, gen_tp)to acceptenable_attention_dp.tests/unittest/disaggregated/test_mamba_transfer.py_create_transceiverswithenable_attention_dp.Mapping.These tests are not individually listed in
tests/integration/test_lists/. The change adds one waiver entry intests/integration/test_lists/waives.txt, but it does not provide CI or manual-QA entries for the modified test functions.Verdict: needs follow-up. Confirm CBTS coverage and add or verify the required CI and manual-QA test-list coverage.
pr-babysitter waive set-bug NVIDIA/TensorRT-LLM#17483 <bug-url>— the waives entry carries a placeholder until then