[None][fix] Filter empty aux buffers from NIXL transfers - #16993
[None][fix] Filter empty aux buffers from NIXL transfers#16993chuangz0 wants to merge 6 commits into
Conversation
cd791e5 to
ca61c38
Compare
|
/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:
WalkthroughAdds validated auxiliary transfer layouts, peer-scoped layout caching, filtering of empty auxiliary buffers during registration and transfer metadata construction, and tests covering validation, caching, registration, and integration execution. ChangesAuxiliary transfer handling
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant TransferWorker
participant PeerRegistrar
participant build_aux_transfer_layout
TransferWorker->>PeerRegistrar: request cached auxiliary transfer layout
TransferWorker->>build_aux_transfer_layout: build layout when cache is empty
build_aux_transfer_layout-->>TransferWorker: return validated non-empty pointers and sizes
TransferWorker->>TransferWorker: build write metadata and register non-empty buffers
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
|
PR_Github #62439 [ run ] triggered by Bot. Commit: |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/disaggregation/native/auxiliary.py (1)
37-45: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winConsider making
AuxTransferLayoutarrays read-only.
@dataclass(frozen=True)only blocks attribute reassignment; the underlying numpy arrays are still mutable in-place. Since this layout is cached and shared across concurrent transfer calls to the same peer (seepeer.py's_aux_transfer_layout_cacheand theis-identity test intest_peer.py), an accidental in-place write anywhere downstream would silently corrupt every subsequent transfer using this cached object.♻️ Suggested hardening
return AuxTransferLayout( - src_base_ptrs=src_meta.ptrs[src_indices], - dst_base_ptrs=dst_meta.ptrs[src_indices], - src_item_sizes=src_item_sizes[src_indices], - dst_item_sizes=dst_item_sizes[src_indices], + src_base_ptrs=_readonly(src_meta.ptrs[src_indices]), + dst_base_ptrs=_readonly(dst_meta.ptrs[src_indices]), + src_item_sizes=_readonly(src_item_sizes[src_indices]), + dst_item_sizes=_readonly(dst_item_sizes[src_indices]), )where
_readonlysetsarr.flags.writeable = Falseand returnsarr.Also applies to: 98-103
🤖 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/native/auxiliary.py` around lines 37 - 45, Make the NumPy arrays stored by AuxTransferLayout read-only at construction time, not merely protected from attribute reassignment by frozen dataclass semantics. Apply the existing _readonly helper to src_base_ptrs, dst_base_ptrs, src_item_sizes, and dst_item_sizes while preserving the cached object’s identity and shared-transfer behavior.tests/unittest/disaggregated/test_aux_buffer_registration.py (1)
1-161: 🧹 Nitpick | 🔵 TrivialTest coverage summary.
- Added test functions:
test_null_pointer_with_non_zero_size_is_rejected,test_zero_size_aux_buffer_is_skipped,test_real_empty_draft_buffer_is_skipped_during_registration,test_all_buffers_registered_when_none_are_empty,test_nothing_registered_when_every_buffer_is_empty,test_real_empty_draft_buffer_is_skipped_during_transfer,test_non_empty_source_requires_non_empty_destination,test_destination_aux_buffer_must_be_large_enough.- Listed in test lists: confirmed present in both
tests/integration/test_lists/test-db/l0_a10.ymlandtests/integration/test_lists/test-db/l0_h100.yml.- Coverage verdict: sufficient — validation errors (null pointer, missing/undersized destination), empty-buffer filtering during both registration and transfer, and the "all empty"/"all non-empty" boundary cases are all exercised, and error messages/indices are asserted precisely against the implementation.
- Minor gap (not blocking):
_build_aux_write_metatests fixtask._slot=0/peer_slot=0, so the per-slot pointer offset (base_ptr + item_size * slot) is never exercised at a non-zero slot.🤖 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_aux_buffer_registration.py` around lines 1 - 161, Add coverage for non-zero slot pointer offsets in the auxiliary transfer path. Update `_build_aux_write_meta` or add a focused test using a task/peer slot greater than zero, then assert `src_ptrs` and `dst_ptrs` include the expected per-slot offset while preserving existing empty-buffer 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.
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/native/auxiliary.py`:
- Around line 37-45: Make the NumPy arrays stored by AuxTransferLayout read-only
at construction time, not merely protected from attribute reassignment by frozen
dataclass semantics. Apply the existing _readonly helper to src_base_ptrs,
dst_base_ptrs, src_item_sizes, and dst_item_sizes while preserving the cached
object’s identity and shared-transfer behavior.
In `@tests/unittest/disaggregated/test_aux_buffer_registration.py`:
- Around line 1-161: Add coverage for non-zero slot pointer offsets in the
auxiliary transfer path. Update `_build_aux_write_meta` or add a focused test
using a task/peer slot greater than zero, then assert `src_ptrs` and `dst_ptrs`
include the expected per-slot offset while preserving existing empty-buffer
behavior.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 159841ad-2e72-4870-b8bc-73faaba4560e
📒 Files selected for processing (7)
tensorrt_llm/_torch/disaggregation/native/auxiliary.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytests/integration/test_lists/test-db/l0_a10.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/disaggregated/test_aux_buffer_registration.pytests/unittest/disaggregated/test_peer.py
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/native/peer.py`:
- Around line 76-78: The peer registration, unregistration, and auxiliary-layout
cache lifecycle are not synchronized, allowing stale layouts or layouts for
unregistered peers to be published. Update the relevant register/unregister
methods and cache_aux_transfer_layout to share one lock or enforce a peer
identity/generation check across layout construction and publication, ensuring
only layouts for the currently registered peer are cached and transfers cannot
combine mismatched RankInfo and pointers. Add a concurrent peer-replacement test
covering these paths.
🪄 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: f2286e4c-a298-4ed0-8f26-e7654575bad2
📒 Files selected for processing (5)
tensorrt_llm/_torch/disaggregation/native/auxiliary.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytests/unittest/disaggregated/test_aux_buffer_registration.pytests/unittest/disaggregated/test_peer.py
🚧 Files skipped from review as they are similar to previous changes (4)
- tests/unittest/disaggregated/test_peer.py
- tests/unittest/disaggregated/test_aux_buffer_registration.py
- tensorrt_llm/_torch/disaggregation/native/transfer.py
- tensorrt_llm/_torch/disaggregation/native/auxiliary.py
|
PR_Github #62439 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #62478 [ run ] triggered by Bot. Commit: |
|
The empty-buffer filtering itself checks out — with a non-empty aux buffer the descriptor list is byte-identical to before (same order, dtype, shape), and registration emits the same The layout cache is the problem.
The stale destination base pointers then persist for every later transfer to that peer, since nothing invalidates again until the next register/unregister. That is the same class of failure this PR is fixing — writes aimed at memory the peer no longer has registered. The GIL doesn't help; it's a multi-step sequence, not a single dict op. Before the change this couldn't happen: Simplest fix is to stamp the layout with the peer Two smaller things, not blocking:
|
|
PR_Github #62478 [ run ] completed with state
|
dc99aaa to
dc86150
Compare
|
/bot run --stage-list "GB200-8_GPUs-2_Nodes-PyTorch-1, GB200-8_GPUs-2_Nodes-PyTorch-2, DGX_B200-PyTorch-6, DGX_H100-PyTorch-4" |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tensorrt_llm/_torch/disaggregation/native/peer.py (1)
134-142: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the new cache API.
Add Google-style docstrings for the peer-layout getter and cache writer, including the peer-key semantics and absent-layout return value.
🤖 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/native/peer.py` around lines 134 - 142, Add Google-style docstrings to get_aux_transfer_layout and cache_aux_transfer_layout describing the peer_name/peer_rank key semantics; document that get_aux_transfer_layout returns None when no layout is cached, and describe the layout stored by the writer.Source: Coding guidelines
tensorrt_llm/_torch/disaggregation/native/auxiliary.py (1)
53-109: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the public layout helpers in Google style.
get_non_empty_aux_indices()andbuild_aux_transfer_layout()are imported by the transfer path, but their docstrings omit argument, return, and raised-error contracts. Document array shapes/dtypes and validation failures.🤖 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/native/auxiliary.py` around lines 53 - 109, Expand the Google-style docstrings for get_non_empty_aux_indices() and build_aux_transfer_layout() with Args, Returns, and Raises sections. Document the expected array shapes and relevant dtypes, the returned index/layout values, and each validation failure raised for mismatched shapes, negative sizes, null non-empty pointers, layout mismatches, missing destination buffers, or insufficient destination capacity.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.
Nitpick comments:
In `@tensorrt_llm/_torch/disaggregation/native/auxiliary.py`:
- Around line 53-109: Expand the Google-style docstrings for
get_non_empty_aux_indices() and build_aux_transfer_layout() with Args, Returns,
and Raises sections. Document the expected array shapes and relevant dtypes, the
returned index/layout values, and each validation failure raised for mismatched
shapes, negative sizes, null non-empty pointers, layout mismatches, missing
destination buffers, or insufficient destination capacity.
In `@tensorrt_llm/_torch/disaggregation/native/peer.py`:
- Around line 134-142: Add Google-style docstrings to get_aux_transfer_layout
and cache_aux_transfer_layout describing the peer_name/peer_rank key semantics;
document that get_aux_transfer_layout returns None when no layout is cached, and
describe the layout stored by the writer.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f01aebe2-0f5e-43cd-ad52-cb641fe32ce2
📒 Files selected for processing (7)
tensorrt_llm/_torch/disaggregation/native/auxiliary.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytests/integration/test_lists/test-db/l0_a10.ymltests/integration/test_lists/test-db/l0_h100.ymltests/unittest/disaggregated/test_aux_buffer_registration.pytests/unittest/disaggregated/test_peer.py
|
/bot run --disable-fail-fast |
|
PR_Github #63391 [ run ] triggered by Bot. Commit: |
|
PR_Github #63391 [ run ] completed with state
|
73afdec to
bb23709
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63632 [ run ] triggered by Bot. Commit: |
|
PR_Github #63632 [ run ] completed with state
|
bb23709 to
1433776
Compare
d05055e to
086b529
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #63742 [ run ] triggered by Bot. Commit: |
|
PR_Github #63742 [ run ] completed with state
|
The draft-token aux buffer is allocated as torch.empty(max_slot_num,
max_draft_len), so with max_draft_len == 0 (the default when speculative
decoding is off) it has numel() == 0 and data_ptr() == 0.
_register_aux_buffer registered all four aux pointers unconditionally,
handing that null address to NIXL. The LIBFABRIC backend rejects it
outright -- libfabric_rail_manager.cpp `if (!buffer)` -> "Invalid buffer
parameter" -- which fails the entire registration and aborts transceiver
setup:
Assertion failed: status == NIXL_SUCCESS (transferAgent.cpp)
NixlTransferAgent::registerMemory(MemoryDescs const&)
UCX tolerates the null descriptor, so this only reproduced with
TRTLLM_NIXL_KVCACHE_BACKEND=LIBFABRIC. The C++ transceiver is unaffected
because it registers only the CacheTransBufferManager staging buffers.
Skip zero-pointer/zero-size buffers, and skip the registration entirely
when nothing is left. Empty buffers carry no data, so nothing is lost.
Verified on 2 nodes x 2 GPU over AWS EFA (16 rails): the
NIXL/PYTHON/LIBFABRIC case goes from TRANSFER_ERROR to PASS at
96.7 GB/s per GPU.
Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
Register the new test file in the l0_h100 and l0_a10 test lists -- without an entry it is never executed by CI. Give the null-pointer case a non-zero size so it fails if only zero-sized buffers are filtered, and assert the surviving descriptors keep their original indices (aux_buffer_ptr_0/2/3) rather than being renumbered. Annotate the new functions, matching the convention in the neighbouring disaggregated unit tests. Baseline-fails check on 1.3.0rc22: 4/4 pass with the fix, 1/4 without it (only the no-empty-buffers control passes). Signed-off-by: Xianjie <5410381+qiaoxj07@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com>
086b529 to
426c4f6
Compare
|
/bot run --stage-list "L40S-PyTorch-1" |
|
PR_Github #63883 [ run ] triggered by Bot. Commit: |
|
PR_Github #63883 [ run ] completed with state |
|
/bot --help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
7 participants
Description
This builds on #16948 and keeps empty auxiliary buffers out of both NIXL memory registration and transfer requests.
When speculative decoding is disabled, the draft-token auxiliary tensor has zero elements and may expose a null data pointer. #16948 filters that tensor during registration, but the generation-first auxiliary transfer path still constructed a descriptor for it. NIXL requires every transfer descriptor to be covered by registered local and remote memory, so the remaining
(address=0, size=0)descriptor could still make transfer request creation fail.The change:
AuxBufferinstances withmax_draft_len=0, mismatched peer layouts, and cache lifecycle checks.User impact
Python-runtime disaggregated serving with NIXL/LIBFABRIC no longer submits an unregistered empty draft-token descriptor during generation-first auxiliary transfers. Incompatible speculative-decoding layouts fail during peer registration with a clear error instead of reaching the transfer backend.
Static layout validation is not repeated for each request. A local microbenchmark of the per-task pointer calculation measured approximately 1.05 microseconds per call, compared with approximately 9.6 microseconds when validation and filtering were performed on every task.
Validation
pre-commit run --files tensorrt_llm/_torch/disaggregation/native/auxiliary.py tensorrt_llm/_torch/disaggregation/native/peer.py tensorrt_llm/_torch/disaggregation/native/transfer.py tests/unittest/disaggregated/test_aux_buffer_registration.py tests/unittest/disaggregated/test_peer.pyPYTHONPATH=$PWD TRTLLM_TEST_PREFETCH_SESSION=0 TRTLLM_TEST_REUSE_SESSION=0 pytest -q tests/unittest/disaggregated/test_aux_buffer_registration.py tests/unittest/disaggregated/region/test_aux.py tests/unittest/disaggregated/test_rank_info.py tests/unittest/disaggregated/test_peer.py tests/unittest/disaggregated/test_pool_matching.pyDev Engineer Review
AuxTransferLayoutand helper functions intensorrt_llm/_torch/disaggregation/native/auxiliary.py:get_non_empty_aux_indices()validates pointer/size descriptor shape matching, rejects negative sizes, rejects “null pointer with non-zero size”, and returns indices wheresize > 0.build_aux_transfer_layout()constructs a validated layout for non-empty aux entries:transfer.py, AUX DRAM descriptor registration for aux pointers is skipped when all aux entries are “real empty”._build_aux_write_metanow prefers a validated/cachedAuxTransferLayoutto compute AUX write metadata (source/destination ptrs and sizes) for non-empty indices, instead of using all slots.tensorrt_llm/_torch/disaggregation/native/peer.py:AuxTransferLayoutper peer key.register()and removes them onunregister().get_aux_transfer_layout()andcache_aux_transfer_layout()APIs for transfer construction to reuse layouts.QA Engineer Review
Test-list files modified
tests/integration/test_lists/test-db/l0_a10.ymlunittest/_torch/sampler/test_penalties.pyunittest/disaggregated/test_aux_buffer_registration.pytests/integration/test_lists/test-db/l0_h100.ymlunittest/_torch/speculative/test_sa_hybrid_state_promotion.py(pre_merge PyTorch speculative)unittest/disaggregated/test_aux_buffer_registration.py(pre_merge disaggregated)Test code added/updated (files outside test lists)
tests/unittest/disaggregated/test_aux_buffer_registration.pytest_null_pointer_with_non_zero_size_is_rejectedtest_zero_size_aux_buffer_is_skippedtest_real_empty_draft_buffer_is_skipped_during_registrationtest_all_buffers_registered_when_none_are_emptytest_nothing_registered_when_every_buffer_is_emptytest_real_empty_draft_buffer_is_skipped_during_transfertest_non_empty_source_requires_non_empty_destinationtest_aux_layout_mismatch_is_checked_lazily_during_transfertest_destination_aux_buffer_must_be_large_enoughtest-db: yes (vial0_a10.ymlandl0_h100.yml)tests/unittest/disaggregated/test_peer.pytest-db/qaby the modified test lists in this PR: no indication.Verdict: needs follow-up (new/updated
test_peer.pycases are not confirmed by the test-db entries changed in this PR.)