[TRTLLM-13969][feat] Support MiniMax M3 for Disaggregated Serving - #16017
Merged
chuangz0 merged 21 commits intoJul 23, 2026
Conversation
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #57871 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57871 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #58028 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58028 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58076 [ run ] triggered by Bot. Commit: |
Contributor
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
WalkthroughAdds mapper-aware NHD and replicated KV pool views, MiniMax-M3 disaggregation support, geometry-based peer mapping, region-aware NIXL descriptor splitting/coalescing, and extensive unit and integration coverage. ChangesDisaggregation and transfer updates
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py (1)
600-602: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider precise
Callablesignatures for the injectable hooks.Bare
Callableloses the parameter/return contract thatmanager_factory,init_fn, andverify_fnmust satisfy, so signature drift between the DeepSeek and MiniMax injections won't be caught by the type checker. As per coding guidelines ("Prefer specifying argument types inCallabletype hints").🤖 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/disaggregated/test_deepseek_v4_kv_transfer.py` around lines 600 - 602, The injectable hook parameters in the DeepSeek KV transfer test helpers currently use bare Callable types, which hides the expected function contracts. Update the type hints for manager_factory, init_fn, and verify_fn in the relevant test setup helpers to use explicit Callable signatures that match _create_managers_for_instance, _init_pool_data, and verify_all_requests so the type checker can catch mismatches.Source: Coding guidelines
tests/unittest/disaggregated/test_rank_info.py (1)
84-109: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd companion test for whole-byte dtype path.
Coverage gap: this file only exercises
from_kv_cache_manager's sub-byte branch (NVFP4). No test asserts the whole-byte branch (e.g. HALF) still produces anintelement_bytes. A regression collapsing both branches to always return a float wouldn't be caught here.Suggest adding
test_rank_info_represents_whole_byte_cache(or extending this test) withdtype=DataType.HALF, assertingelement_bytes == 2andisinstance(element_bytes, int).🤖 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/disaggregated/test_rank_info.py` around lines 84 - 109, Add a companion test for the whole-byte path in RankInfo.from_kv_cache_manager so the dtype handling is covered beyond NVFP4; either extend test_rank_info_represents_subbyte_nvfp4_cache or add test_rank_info_represents_whole_byte_cache using DataType.HALF on the same manager setup, then assert RankInfo.attention.element_bytes equals 2 and is an int, and verify the serialized round-trip via RankInfo.to_bytes/from_bytes preserves the same integer element_bytes value.Source: Path instructions
tests/unittest/disaggregated/test_cache_reuse_adapter.py (1)
654-717: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winQA coverage note: this suite thoroughly covers
_prepare_kv_blocks_for_transfer's documented behaviors (draft-block trimming, oversized-diff rejection,dst_start_tokentrimming, SWA requirements/clamping). One gap: no case exercisessrc_block_idslonger thandst_block_ids(negativeblock_diff), which is currently unvalidated in the implementation (see companion comment intensorrt_llm/_torch/disaggregation/native/transfer.py). Consider adding a test once that behavior is confirmed/fixed.🤖 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/disaggregated/test_cache_reuse_adapter.py` around lines 654 - 717, The KV block transfer coverage is missing a case where src_block_ids is longer than dst_block_ids, so negative block_diff is not exercised. Add a test in TestPrepareKvBlocksForTransfer that calls Sender._prepare_kv_blocks_for_transfer with a shorter dst input and asserts the expected behavior once the implementation in _prepare_kv_blocks_for_transfer is confirmed. Use the existing helper _prepare and align the assertion with the transfer.py block_diff handling so this edge case is explicitly covered.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 `@tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py`:
- Around line 82-90: The region size selection in kv_extractor.py incorrectly
uses a truthy fallback, so an explicit bytes_per_region of 0 is treated the same
as unset. Update the logic in the KV cache extraction path around region_size to
check bytes_per_region with an explicit None test instead of using or, so only
None falls back to pool.slot_bytes. Keep the rest of the SpecRegion and
MemRegionGroup construction unchanged.
In `@tests/unittest/disaggregated/test_cache_reuse_adapter.py`:
- Around line 694-701: The pytest.raises assertion in
test_swa_requires_prompt_length uses a regex match string with an unescaped
period, so update the match pattern to treat the literal dot in
session.prompt_len as a regex escape. Adjust the pytest.raises call in
test_swa_requires_prompt_length to use a properly escaped match value while
keeping the same ValueError expectation from _prepare.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_cache_reuse_adapter.py`:
- Around line 654-717: The KV block transfer coverage is missing a case where
src_block_ids is longer than dst_block_ids, so negative block_diff is not
exercised. Add a test in TestPrepareKvBlocksForTransfer that calls
Sender._prepare_kv_blocks_for_transfer with a shorter dst input and asserts the
expected behavior once the implementation in _prepare_kv_blocks_for_transfer is
confirmed. Use the existing helper _prepare and align the assertion with the
transfer.py block_diff handling so this edge case is explicitly covered.
In `@tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py`:
- Around line 600-602: The injectable hook parameters in the DeepSeek KV
transfer test helpers currently use bare Callable types, which hides the
expected function contracts. Update the type hints for manager_factory, init_fn,
and verify_fn in the relevant test setup helpers to use explicit Callable
signatures that match _create_managers_for_instance, _init_pool_data, and
verify_all_requests so the type checker can catch mismatches.
In `@tests/unittest/disaggregated/test_rank_info.py`:
- Around line 84-109: Add a companion test for the whole-byte path in
RankInfo.from_kv_cache_manager so the dtype handling is covered beyond NVFP4;
either extend test_rank_info_represents_subbyte_nvfp4_cache or add
test_rank_info_represents_whole_byte_cache using DataType.HALF on the same
manager setup, then assert RankInfo.attention.element_bytes equals 2 and is an
int, and verify the serialized round-trip via RankInfo.to_bytes/from_bytes
preserves the same integer element_bytes value.
🪄 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: b9803a29-b4c0-4a6e-b73f-ffee4feb6e13
📒 Files selected for processing (19)
tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.pytensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/rank_info.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/disaggregation/resource/page.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/disaggregated/region/test_page.pytests/unittest/disaggregated/test_cache_reuse_adapter.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/disaggregated/test_extractor.pytests/unittest/disaggregated/test_minimax_m3_kv_transfer.pytests/unittest/disaggregated/test_peer.pytests/unittest/disaggregated/test_pool_matching.pytests/unittest/disaggregated/test_rank_info.py
Collaborator
|
PR_Github #58076 [ run ] completed with state
|
chuangz0
reviewed
Jul 8, 2026
chuangz0
reviewed
Jul 8, 2026
chuangz0
reviewed
Jul 8, 2026
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #58301 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58301 [ run ] completed with state
|
chuangz0
reviewed
Jul 9, 2026
Collaborator
|
[by Codex] @lowsfer Could you review this PR? Thanks! |
…sagg resource layer Introduce MapperKind and per-cache-class pool views in the page table so heterogeneous pools (e.g. MiniMax M3 full-attention vs SWA classes) are addressed without storage-layout changes; PoolView is constructed with bytes_per_layer instead of mutating it. Extend the KV extractor to build page tables from the manager per class and default Role.INDEX_KEY to REPLICATED in V2 role mapping. Covered by region/page, extractor and pool-matching tests, including DSA indexer-K coverage. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
… MiniMax M3 Surface per-cache-class pool metadata from KVCacheManagerV2 and the MiniMax M3 sparse cache manager so the disagg resource layer can build per-class views over the existing storage layout (no storage changes). Covered by kv_cache_manager_v2 unit tests. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…M3 disagg transfer Unify all KV mappers on entries-driven layer addressing, organized as IntactMapper/ReplicatedMapper and the HND/NHD head-mismatch mappers (dead V1-era mappers dropped), so the attention mixer can map heterogeneous per-class pools between ctx and gen peers. RankInfo carries the per-class page tables; transfer sessions consume the mapped entries. Covered by peer and rank-info tests. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…ansceiverV2 Exchange the cancelled/failed/completed id lists in a single batched allgather (packed list-of-lists) instead of three collectives, for both ctx (TP then PP) and gen consensus; update bounce gate test fakes accordingly. A ctx idle fast-path that skips the variable-length allgathers on idle polls is left as a TODO for a follow-up PR — its reduction must mirror _ctx_consensus()'s communicator scope (see the skipped test test_ctx_consensus_fastpath_skips_when_idle). Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
… entries Extract the shared KV-transfer harness from the DSv4 test module and port DSv4 onto it; add the MiniMax M3 KV-transfer suite (per-class pool views, coalesced vs separate schemes) and DSA indexer-K transceiver coverage; register the disagg KV-transfer suites in the l0_h100 test list. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…apper refactor The refactoring stack rewrites the pool-view mapper layer and supersedes some later review-feedback changes in PR 16017: buffer_mapper_kinds helpers, list-returning get_pool_mapping golden tests, the MapperKind.MIXED scenarios, and the extracted _prepare_kv_blocks_for_transfer helper (equivalent trimming logic lives inline in _build_kv_write_meta). Align these files with the refactored design. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
A region-map miss returned the {0,0} sentinel, so two unknown regions
compared equal in the merge guard and contiguity alone could merge
pieces across chunk or registration boundaries (e.g. when the remote
peer sent no region metadata). Lookups now report a found flag and a
piece whose lookup misses on either side is never merged, degrading to
split-only behavior. Tests updated to provide region metadata and new
cases lock the no-merge-on-miss behavior.
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Three silent-failure paths in peer pool matching now raise instead: - get_pool_mapping raises when a self pool view's layers span multiple peer layer groups (mismatched grouping between peers); previously the first hit won and the remaining layers were silently skipped. - get_layer_to_layer_group raises on a duplicate global_layer_id across layer groups; previously last-write-wins. - _get_buffers_per_layer counts entries per layer and requires them uniform; total-count divisibility let skewed layouts (e.g. 1+3 over two layers) pass and mis-derive the per-layer buffer count. All three are defensive: current page-table builders uphold these invariants by construction. Unit tests added for each guard. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…group_descs API KVCacheManagerV2 now exposes its slot layout through the public pool_group_descs API (added for the upcoming C++ backend, where the Python storage internals no longer exist). Rebuild _build_page_table_v2 on that API instead of reaching into storage._buffer_attr, _levels and _life_cycles: iterate pool groups and their slot-desc variants (one per layer group), derive buffer offsets from buffer_ids order times single_buffer_size (identical to the storage layer's own offset assignment), and keep the per-(pool, mapper-kind) view bucketing, bytes_per_layer stamping and canonical ordering unchanged. Window sizes now come from init_config.layers. The duck-typed test manager fakes pool_group_descs accordingly. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Commit 9115e48 (NVIDIA#16099) added an index_scale keyword argument to KVCacheManagerV2._get_batch_cache_indices_by_pool_id and passes it explicitly from get_batch_cache_indices, so the MiniMax M3 override must accept it. M3 bypasses the V1 block-id conversion by design (its forward path indexes paged views directly by slot id), so the caller-supplied scale is ignored alongside index_scales. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
…tion Commit f84b6da (NVIDIA#15245) made _CacheReuseAdapterV1.get_block_ids translate block ids through get_memory_pool_block_indices with the layer group's window size, but test_v1_adapter_uses_request_py_beam_width kept a windowless fake layer group and a fake manager without the translation hook, tripping the window_size assertion. The gap went unnoticed because test_cache_reuse_adapter.py is not registered in any CI test list. Give the fake layer group a window, add an identity get_memory_pool_block_indices to the fake manager, and assert the adapter passes the layer group's window through to the translation. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
get_kv_map iterated the overlapping layers sorted by global id, which is correct (per-layer byte offsets are looked up explicitly) but can split a physically contiguous copy into multiple fragments when the slot layout is not monotonic in global id. Iterate the overlap in self's physical slot order instead so byte-contiguous layers stay adjacent in the offset arrays and the mappers merge them into a single fragment. Add a unit test covering a non-monotonic but contiguous layout merging into one fragment. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
_kv_pool_mapping_offset derived the per-layer offset from layer_grouping's iteration order, which is an implementation detail of the Python KVCacheManagerV2 and not an API contract (per review on PR 16017, the ordering is not guaranteed and may change with the upcoming C++ backend). Rank the group's layers by their Role.KEY pool base address instead: layer_grouping is now used only for membership, and the offset always reflects the physical slot layout, keeping the NVFP4 block_scale_offset == offset cross-check in the base pool-mapping loop meaningful under any future grouping order. Add a unit test where layer_grouping order deliberately disagrees with the physical addresses. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
chuangz0
force-pushed
the
peihengh/trtllm-13969-minimax-m3-disagg
branch
from
July 23, 2026 02:23
f6aa318 to
0cb44ef
Compare
Collaborator
|
/bot run --stage-list "A100X-PyTorch-1, DGX_B200-PyTorch-1, DGX_B200-PyTorch-5" |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/disaggregated/region/test_block.py (1)
124-125: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winMake the fixture regions large enough for their offsets.
bytes_per_regiondescribes the base memory block, but these tests provide blocks smaller than the mapped layers. The partial-layer case needs at least3 * bytes_per_layer; the peer case needs at least2 * peer_bytes_per_layerto contain the layer at offsetpeer_bytes_per_layer.Suggested fixture correction
- src_group = MemRegionGroup(ptrs=np.array([10, 20], dtype=np.int64), bytes_per_region=1) - dst_group = MemRegionGroup(ptrs=np.array([30, 40], dtype=np.int64), bytes_per_region=1) + src_group = MemRegionGroup( + ptrs=np.array([10, 20], dtype=np.int64), + bytes_per_region=3 * bytes_per_layer, + ) + dst_group = MemRegionGroup( + ptrs=np.array([30, 40], dtype=np.int64), + bytes_per_region=3 * bytes_per_layer, + )Also applies to: 149-150
🤖 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/disaggregated/region/test_block.py` around lines 124 - 125, Update the MemRegionGroup fixtures in the affected tests to allocate base blocks large enough for their mapped offsets: use at least 3 * bytes_per_layer for the partial-layer case and 2 * peer_bytes_per_layer for the peer case. Adjust the bytes_per_region values near the src_group and dst_group definitions while preserving the existing pointer arrays and test behavior.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/disaggregation/resource/utils.py (1)
44-49: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the
buffer_entriesparameter.Every other parameter and the return are typed, but
buffer_entries(the structuredBUFFER_ENTRY_DTYPEarray / iterable of entries) is not. As per coding guidelines, "Annotate every function ... document public function arguments."🤖 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 `@tensorrt_llm/_torch/disaggregation/resource/utils.py` around lines 44 - 49, Annotate the buffer_entries parameter in compute_layer_byte_ranges with the appropriate type for the structured BUFFER_ENTRY_DTYPE array or iterable of entries, matching the project’s existing typing conventions and preserving the current function behavior.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/disaggregated/kv_transfer_harness.py`:
- Around line 41-43: Update the environment setup in the KV transfer harness to
deterministically assign TRTLLM_NIXL_NUM_THREADS to "1" instead of using
setdefault, set UCX_TLS to "^ib,gdr_copy", and remove UCX_NET_DEVICES before
constructing real NIXL agents so inherited environment values cannot override
the test configuration.
In `@tests/unittest/disaggregated/region/test_block.py`:
- Around line 128-140: Strengthen the partial-layer and HND mapper tests around
IntactMapper.map by using independent source and destination offsets rather than
identical values. Assert the exact result.src.memory.ptrs and
result.dst.memory.ptrs, including each peer layer offset and computed
head/buffer offsets, while retaining the existing fragment-count and
bytes-per-region assertions.
---
Outside diff comments:
In `@tests/unittest/disaggregated/region/test_block.py`:
- Around line 124-125: Update the MemRegionGroup fixtures in the affected tests
to allocate base blocks large enough for their mapped offsets: use at least 3 *
bytes_per_layer for the partial-layer case and 2 * peer_bytes_per_layer for the
peer case. Adjust the bytes_per_region values near the src_group and dst_group
definitions while preserving the existing pointer arrays and test behavior.
---
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/resource/utils.py`:
- Around line 44-49: Annotate the buffer_entries parameter in
compute_layer_byte_ranges with the appropriate type for the structured
BUFFER_ENTRY_DTYPE array or iterable of entries, matching the project’s existing
typing conventions and preserving the current function behavior.
🪄 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: 667de667-49b9-4faa-b46a-5e0f0d1c6c27
📒 Files selected for processing (36)
cpp/include/tensorrt_llm/executor/transferAgent.hcpp/tensorrt_llm/common/envUtils.cppcpp/tensorrt_llm/common/envUtils.hcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cppcpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.hcpp/tensorrt_llm/executor/cache_transmission/transferAgent.cppcpp/tests/unit_tests/executor/CMakeLists.txtcpp/tests/unit_tests/executor/coalesceTest.cppcpp/tests/unit_tests/executor/transferAgentTest.cpptensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.pytensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.pytensorrt_llm/_torch/disaggregation/native/mixers/attention/spec.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/rank_info.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/disaggregation/resource/page.pytensorrt_llm/_torch/disaggregation/resource/utils.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytests/integration/test_lists/test-db/l0_a10.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/disaggregated/kv_transfer_harness.pytests/unittest/disaggregated/region/test_block.pytests/unittest/disaggregated/region/test_page.pytests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_cache_reuse_adapter.pytests/unittest/disaggregated/test_cache_transceiver_single_process.pytests/unittest/disaggregated/test_deepseek_v4_kv_transfer.pytests/unittest/disaggregated/test_extractor.pytests/unittest/disaggregated/test_minimax_m3_kv_transfer.pytests/unittest/disaggregated/test_peer.pytests/unittest/disaggregated/test_pool_matching.pytests/unittest/disaggregated/test_rank_info.pytests/unittest/disaggregated/test_transceiver_bounded_polling.py
💤 Files with no reviewable changes (1)
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h
🚧 Files skipped from review as they are similar to previous changes (27)
- cpp/tensorrt_llm/common/envUtils.h
- tensorrt_llm/_torch/disaggregation/native/rank_info.py
- cpp/tensorrt_llm/common/envUtils.cpp
- tests/unittest/disaggregated/test_cache_reuse_adapter.py
- tests/unittest/disaggregated/test_bounce.py
- cpp/tests/unit_tests/executor/transferAgentTest.cpp
- cpp/tests/unit_tests/executor/CMakeLists.txt
- cpp/include/tensorrt_llm/executor/transferAgent.h
- tests/unittest/disaggregated/test_rank_info.py
- tests/unittest/disaggregated/test_transceiver_bounded_polling.py
- tests/unittest/_torch/executor/test_kv_cache_manager_v2.py
- tensorrt_llm/_torch/pyexecutor/resource_manager.py
- cpp/tests/unit_tests/executor/coalesceTest.cpp
- tensorrt_llm/_torch/disaggregation/resource/page.py
- tensorrt_llm/_torch/disaggregation/native/transfer.py
- tests/unittest/disaggregated/region/test_page.py
- tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
- cpp/tensorrt_llm/executor/cache_transmission/transferAgent.cpp
- cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
- tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py
- tests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py
- tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py
- tensorrt_llm/_torch/disaggregation/native/peer.py
- tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
- tests/unittest/disaggregated/test_pool_matching.py
- tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py
- tests/unittest/disaggregated/test_peer.py
Collaborator
|
PR_Github #61157 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61157 [ run ] completed with state
|
…assertions Address PR review comments: - kv_transfer_harness.py builds real NIXL agents, so set TRTLLM_NIXL_NUM_THREADS=1 and UCX_TLS=^ib,gdr_copy explicitly instead of setdefault(..., "0"): a single worker thread avoids CPU oversubscription and an inherited/poisoned dev/CI value can no longer destabilize the transport. - region/test_block.py now asserts the exact mapper-produced src/dst addresses. The IntactMapper partial-layer test uses distinct src/dst layer offsets so it fails if the two are ever crossed. The HND head-mismatch test uses peer tp_rank=1 to force a nonzero source head offset and checks exact pointers derived from geometry, so incorrect layer, buffer, or head offsets can no longer pass silently. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Collaborator
|
/bot run --stage-list "A100X-PyTorch-1, DGX_B200-PyTorch-1, DGX_B200-PyTorch-5" |
Collaborator
|
PR_Github #61166 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61166 [ run ] completed with state |
chuangz0
disabled auto-merge
July 23, 2026 05:05
Collaborator
|
/bot skip --comment "all tests have passed" |
Collaborator
|
PR_Github #61192 [ skip ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61192 [ skip ] completed with state |
1 task
hhzhang16
added a commit
to hhzhang16/TensorRT-LLM
that referenced
this pull request
Jul 23, 2026
…nnahz/dep-1083-port-flashinfer-stable-va-lifecycle-for-native-all-reduce * 'main' of https://github.com/NVIDIA/TensorRT-LLM: (83 commits) [None][feat] Bind SourceIdentity to checkpoint artifacts (NVIDIA#16159) [None][infra] Waive 4 failed cases for main in pre-merge 49550 (NVIDIA#16798) [None][fix] Make FlashInfer sampling op wrappers opaque to Dynamo (NVIDIA#16732) [None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e (NVIDIA#16420) [None][feat] Default GLM-5 to the Python KV-cache transceiver (NVIDIA#16524) [None][chore] Add NVTX ranges to per-iteration ADP sync points in PyExecutor (NVIDIA#16422) [https://nvbugs/6426850][test] Unwaive Qwen3.5 397B NVFP4 ADP4 TRTLLM test (NVIDIA#16348) [https://nvbugs/6445456][fix] Restore inplace ops for functionalization v2 (NVIDIA#16410) [None][infra] Waive 1 failed cases for main in pre-merge 49229 (NVIDIA#16786) [None][fix] Load DeepSeek V4 mixed-precision NVFP4 checkpoints (NVIDIA#16433) [None][feat] ADP conversation router: configurable least-queued placement for new conversations (NVIDIA#16294) [None][infra] Waive 1 failed cases for main in pre-merge 49424 (NVIDIA#16780) [None][infra] Waive 1 failed cases for main in pre-merge 49424 (NVIDIA#16781) [TRTLLMINF-188][infra] Require approval for PerfSanity wildcard runs (NVIDIA#16777) [TRTLLM-13969][feat] Support MiniMax M3 for Disaggregated Serving (NVIDIA#16017) [NVIDIA#11932][fix] Filter CUTLASS MoE GEMM tile configs by device shared memory on SM121 (NVIDIA#12704) [None][chore] Remove attention backend test waivers (NVIDIA#16723) [TRTLLM-14540][perf] Skip fp32 state round-trip in FlashInfer GDN pre… (NVIDIA#16703) [TRTLLM-13694][feat] Add IBDB recipe provenance and refresh configs (NVIDIA#16254) [None][infra] Preview/bump/main (NVIDIA#16758) ... Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
yuanjingx87
pushed a commit
to yuanjingx87/TensorRT-LLM
that referenced
this pull request
Jul 26, 2026
…IDIA#16017) Signed-off-by: peihengh <259410613+peihu-nv@users.noreply.github.com> Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Co-authored-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.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
14 participants
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.
Summary by CodeRabbit
Dev Engineer Review
Correctness / API consistency
get_disagg_role_mapper_kinds(bothKVCacheManagerV2and MiniMax M3 sparse manager V2).disable_index_value=Truefor all sparse layers in disagg mode).MapperKindtoINDEXED,REPLICATED,NHD(removedFLAT) and extendedPoolViewwith optionalbytes_per_layerto support geometry-aware views._get_batch_cache_indices_by_pool_idnow supportsnum_blocks_per_seqby directly returning base page indices fromkv_cache.get_base_page_indices(pool_id)and truncating per-request indices (preservingBAD_PAGE_INDEXentries)._kv_pool_mapping_offsetnow computes consistent per-layer offsets based on physical K base address ordering rather thanlayer_groupingiteration order.IntactMapper(fragments only the selected contiguous per-layer regions) and routed mapper dispatch byMapperKind:ReplicatedMapper(specializedIntactMapper) forMapperKind.REPLICATEDHNDHeadMismatchMapperand newNHDHeadMismatchMapperfor head-mismatch cases with NHD’s token-major layout invariants.ValueErrors for region/byte/alignment/buffer-count invariants (including sub-byte head alignment requirements).PeerRegistrar:get_pool_view_global_layer_idson both sides.get_kv_mapto compute overlap layers and offsets from explicit per-layer byte ranges and validated per-layer buffer entry distribution._owns_tp_fan_in).VmmDescSplitterAPI tosplitAndCoalesceTransferDescs(..., enableCoalesce=true)and updated call sites.getEnvNixlEnableCoalesce()→getEnvNixlDisableCoalesce()TRTLLM_NIXL_ENABLE_COALESCE→TRTLLM_NIXL_DISABLE_COALESCEelement_bytesas anint | float-typed quantity, with unit tests added to lock expected typing/value behavior.Tests / validation additions
PoolViewroundtrip ofmapper_kindandbytes_per_layer.compute_layer_byte_ranges/get_layer_byte_rangesvalidation and failure modes.IntactMapper,ReplicatedMapper,HNDHeadMismatchMapper,NHDHeadMismatchMapper.KVCacheManagerV2.get_disagg_role_mapper_kinds()defaults.QA Engineer Review
Test list changes (tests/integration/test_lists)
tests/integration/test_lists/test-db/l0_h100.ymlunittest/disaggregated/test_transceiver_bounded_polling.pyunittest/disaggregated/test_pool_matching.pyunittest/disaggregated/test_deepseek_v4_kv_transfer.pyunittest/disaggregated/test_minimax_m3_kv_transfer.pyunittest/disaggregated/test_cache_transceiver_single_process.py::test_cache_transceiver_v1_dsa_indexertests/integration/test_lists/test-db/l0_a10.ymlunittest/disaggregated/test_cache_reuse_adapter.pyTest code changes (files under tests/, outside test-list files)
tests/unittest/_torch/executor/test_kv_cache_manager_v2.py—test_disagg_role_mapper_kinds_default_to_indexedtests/unittest/disaggregated/region/test_page.py—PoolViewkind/bytes_per_layerroundtrip + multipleget_layer_byte_rangesvalidation teststests/unittest/disaggregated/test_cache_reuse_adapter.py— extended request window propagation assertiontests/unittest/disaggregated/test_deepseek_v4_kv_transfer.py— refactor to usekv_transfer_harness.run_kv_transfer_testand tightened typingtests/unittest/disaggregated/test_extractor.py— added V1 DSA indexer-K and V2 page-table builder tests + validation casestests/unittest/disaggregated/test_minimax_m3_kv_transfer.py— new MiniMax M3 disaggregated transfer matrix + multiple unit validationstests/unittest/disaggregated/test_peer.py— expanded mapper/registrar dispatch + geometry/alignment/buffer validation teststests/unittest/disaggregated/test_pool_matching.py— added/updated golden/unit coverage for pool mapping rules and validation errorstests/unittest/disaggregated/test_rank_info.py— dtype-dependentelement_bytesnumeric/type teststests/unittest/disaggregated/test_cache_transceiver_single_process.py— new V1 DSA indexer-K transfer test and harness plumbingtests/unittest/disaggregated/test_transceiver_bounded_polling.py— new skipped unit testtest_ctx_consensus_fastpath_skips_when_idletests/unittest/disaggregated/kv_transfer_harness.py— added threaded single-process NIXL/V2 KV transfer harnesstests/unittest/disaggregated/test_bounce.py,tests/unittest/disaggregated/region/test_block.py— updated mapper/region coverage and fan-in safety assertionstests/integration/test_lists/test-db/l0_h100.ymlandl0_a10.yml(where applicable), but full per-test CBTS linkage/coverage data is not provided.Description
Enables MiniMax M3 with native NIXL disaggregated serving using KV Cache Manager V2.
MiniMax M3 requires model-specific KV-cache handling because it uses token-major NHD K/V storage and a replicated sparse-attention index-key cache. This PR:
Known limitation: long-context head-mismatched TEP↔DEP transfers remain functionally correct but require a future gather/staging/scatter path for production performance.
Test Coverage
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.