[https://nvbugs/5989923][fix] fix test_py_cache_transceiver_mp

[chuangz0]

Summary by CodeRabbit

Release Notes

• New Features

  • Added array-based batch construction API for memory descriptors with from_arrays and from_arrays_uniform_device methods.

• Performance Improvements

  • Vectorized pointer and memory operations using NumPy for faster execution in transfer and extraction paths.
  • Added NVTX profiling instrumentation for better performance visibility across transfer, extraction, and mapping operations.

• Bug Fixes

  • Improved test reliability by fixing port allocation collisions in multi-process test scenarios.

• Tests

  • Removed previously skipped test cases; updated test suite for new NumPy-based data structures.

Description

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)

• 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.

[coderabbitai]

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 67ec05a1-18a5-4eb9-8aa1-38cbf5c02dcf

📥 Commits

Reviewing files that changed from the base of the PR and between f0b336e and 7d19624.

📒 Files selected for processing (17)

• cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp
• cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp
• tensorrt_llm/_torch/disaggregation/base/agent.py
• tensorrt_llm/_torch/disaggregation/base/region.py
• tensorrt_llm/_torch/disaggregation/base/transfer.py
• tensorrt_llm/_torch/disaggregation/native/auxiliary.py
• tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py
• tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py
• tensorrt_llm/_torch/disaggregation/native/transfer.py
• tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py
• tensorrt_llm/_torch/disaggregation/transceiver.py
• tests/integration/test_lists/waives.txt
• tests/unittest/disaggregated/region/test_block.py
• tests/unittest/disaggregated/test_extractor.py
• tests/unittest/disaggregated/test_kv_transfer.py
• tests/unittest/disaggregated/test_kv_transfer_mp.py
• tests/unittest/disaggregated/test_py_cache_transceiver_mp.py

---

📝 Walkthrough

Walkthrough

The changes convert the TensorRT-LLM cache transmission disaggregation infrastructure from Python lists to NumPy arrays for vectorized operations, add NumPy-backed batch construction helper methods to C++ MemoryDescs bindings, update TransferRequest Python binding construction semantics via move-based ownership transfer, introduce NVTX3 performance instrumentation across transfer and mapper operations, and update tests and test configuration accordingly.

Changes

Cohort / File(s)	Summary
C++ Bindings and Transfer Agent cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/agentBindings.cpp, cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp	Added MemoryDescs.from_arrays and from_arrays_uniform_device classmethods for NumPy array conversion; updated TransferRequest.__init__ to use lambda-based construction with move semantics. Added NVTX3 instrumentation to transfer phases (createXferReq, postXferReq, splitTransferDescsFromRegistry).
Base Data Structure Conversions tensorrt_llm/_torch/disaggregation/base/agent.py, tensorrt_llm/_torch/disaggregation/base/region.py, tensorrt_llm/_torch/disaggregation/base/transfer.py	Converted MemoryDescs from dataclass to custom class with from_arrays/from_arrays_uniform_device classmethods; changed MemRegionGroup.ptrs from List[int] to np.ndarray; updated RegionExtractorBase.extract signature to accept np.ndarray for region_ids; changed KVSlice.block_ids_per_layer_groups from List[List[int]] to List[np.ndarray].
Auxiliary Buffer Metadata tensorrt_llm/_torch/disaggregation/native/auxiliary.py	Changed AuxBufferMeta fields (ptrs, size, item_sizes) from list[int] to np.ndarray; updated serialization/deserialization via .tolist()/np.array(..., dtype=np.int64) conversions.
Peer Mapper Vectorization tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py, tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py	Added NVTX instrumentation to mapper methods; replaced len(ptrs) assertions with ptrs.size; converted pointer list arithmetic to NumPy vectorized operations (np.add.outer, array broadcasting); precomputed flat offset arrays in HeadMismatchMapper and ConvStateMismatchMapper during initialization.
Native Transfer Implementation tensorrt_llm/_torch/disaggregation/native/transfer.py	Changed RecvReqInfo.block_ids_per_layer_groups to list[np.ndarray] with custom serialization; converted WriteMeta pointer and size fields to np.ndarray; updated _make_agent_request to use MemoryDescs.from_arrays_uniform_device; replaced list-based construction with vectorized NumPy operations in _build_kv_write_meta and AUX pointer computation.
Resource Extraction tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py	Updated KVRegionExtractorV1.extract signature to accept np.ndarray for region_ids; added NVTX instrumentation; replaced list comprehension filtering with NumPy boolean masking (region_ids[valid]).
Transceiver tensorrt_llm/_torch/disaggregation/transceiver.py	Modified _get_block_ids to return np.ndarray via np.fromiter or np.asarray; updated _create_kv_slice to use ndarray semantics (.size, direct ndarray appending); added NVTX instrumentation decorators to async methods.
Test Updates tests/unittest/disaggregated/region/test_block.py, tests/unittest/disaggregated/test_extractor.py, tests/unittest/disaggregated/test_kv_transfer.py, tests/unittest/disaggregated/test_kv_transfer_mp.py, tests/unittest/disaggregated/test_py_cache_transceiver_mp.py	Updated test inputs from Python lists to np.ndarray for ptrs and region_ids; replaced list equality assertions with np.testing.assert_array_equal; added type assertions for numpy arrays; updated get_block_ids_per_layer_groups to convert results to np.asarray(..., dtype=np.int64); added retry logic with port re-selection in transceiver MP test.
Test Configuration tests/integration/test_lists/waives.txt	Removed four SKIP waivers for test_v2_transceiver_mp parameterizations (nvbugs/5989923, nvbugs/5996642, nvbugs/6018051).

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Possibly related PRs

• [TRTLLM-9523][chore] PyTransceiver code consolidation #12342: Modifies MemoryDescs/MemoryDesc and TransferRequest construction in Python/C++ bindings and call sites, directly related at code level.

Suggested reviewers

• schetlur-nv
• pcastonguay

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}