[None][fix] Skip empty aux buffers when registering NIXL memory - #16948
[None][fix] Skip empty aux buffers when registering NIXL memory#16948qiaoxj07 wants to merge 2 commits into
Conversation
|
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 (3)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughChangesAuxiliary-buffer registration now skips null-pointer or zero-sized descriptors and avoids registration when all entries are empty. New unit tests cover filtering, complete registration, descriptor index preservation, the all-empty case, and pre-merge test-list inclusion. Auxiliary buffer registration
Estimated code review effort: 2 (Simple) | ~10 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 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
🤖 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_aux_buffer_registration.py`:
- Around line 47-56: The test test_null_pointer_aux_buffer_is_skipped should
isolate null-pointer filtering from zero-size filtering. Give the skipped buffer
a non-zero size, update the expected registered descriptors accordingly, and
assert their names are the sparse sequence aux_buffer_ptr_0, aux_buffer_ptr_2,
and aux_buffer_ptr_3 alongside the existing pointer and registration assertions.
🪄 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: 05061793-ad1c-407e-816b-b1b2408d65c9
📒 Files selected for processing (2)
tensorrt_llm/_torch/disaggregation/native/transfer.pytests/unittest/disaggregated/test_aux_buffer_registration.py
|
/bot run |
|
PR_Github #62181 [ run ] triggered by Bot. Commit: |
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>
f73628b to
7158463
Compare
|
/bot run |
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 `@tests/unittest/disaggregated/test_aux_buffer_registration.py`:
- Around line 33-43: Add complete parameter and return type annotations to
_fake_worker, _registered_descs, and each of the four test functions in this
diff. Use precise types for helper inputs and outputs, and annotate every test
function with -> None without changing their 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: f780af62-52de-4bd9-95fe-ee12bceaf409
📒 Files selected for processing (2)
tensorrt_llm/_torch/disaggregation/native/transfer.pytests/unittest/disaggregated/test_aux_buffer_registration.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/disaggregation/native/transfer.py
|
PR_Github #62184 [ run ] triggered by Bot. Commit: |
|
PR_Github #62181 [ run ] completed with state |
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>
|
Thanks @coderabbitai — all three points were valid and are addressed in 1. Test list registration (the important one). The file had no entry under 2. Null-pointer vs zero-size independence. Correct — the skipped entry previously had both 3. Type annotations. Added. For the record, Validated with a baseline-fails check on 1.3.0rc22 (real container, not mocks):
|
|
Tip For best results, initiate chat on the files or code changes.
🐇 ✅ |
|
/bot run |
|
PR_Github #62187 [ run ] triggered by Bot. Commit: |
|
PR_Github #62184 [ run ] completed with state |
|
PR_Github #62187 [ run ] completed with state
|
2 participants
Description
The draft-token aux buffer is allocated as
torch.empty(max_slot_num, max_draft_len), so whenmax_draft_len == 0(the default with speculative decoding off) it hasnumel() == 0anddata_ptr() == 0.TransferWorker._register_aux_bufferregistered all four aux pointers unconditionally, handing thatnull address to NIXL. The LIBFABRIC backend rejects it outright —
libfabric_rail_manager.cppif (!buffer) { NIXL_ERROR << "Invalid buffer parameter"; }— which fails the entire registrationand aborts transceiver setup:
That string occurs at exactly one site in NIXL and its only precondition is
buffer != nullptr, so thefailure is purely the null descriptor — no alignment, size, or memory-type constraint is involved.
Two reasons this went unnoticed:
TRTLLM_NIXL_KVCACHE_BACKEND=LIBFABRIC.registerMemorycall site(
agent_utils/connection.cpp) registers just theCacheTransBufferManagerstaging buffers, never theaux buffer. Only the Python transceiver registers aux memory.
This blocks
cache_transceiver_config.transceiver_runtime: PYTHON+ LIBFABRIC entirely, which in turnblocks EFA for any model pinned to the V2 KV-cache manager.
Fix
Skip zero-pointer / zero-size buffers, and skip the registration call entirely if nothing is left.
Empty buffers carry no data, so nothing is lost. Descriptor names stay tied to the original buffer
index so peer matching is unchanged.
Test Coverage
New
tests/unittest/disaggregated/test_aux_buffer_registration.py(4 cases, no GPU required): nullpointer skipped, zero size skipped, all-non-empty unchanged (including desc names), and nothing
registered when every buffer is empty.
Validated end-to-end on 2 nodes x 2 GPU over AWS EFA (16 rails), NIXL + LIBFABRIC:
TRANSFER_ERRORFor reference the C++ transceiver reaches 93.6 GB/s per GPU on the same setup; the Python path is
slightly faster since it is zero-copy and skips the staging-buffer hop.
PR Checklist
[JIRA/NVBUG/None][type] SummaryDev Engineer Review
TransferWorker._register_aux_bufferto skip auxiliary-buffer slots where the aux pointer (data_ptr()) is0or the aux size is0, preventing NIXL/LIBFABRIC registration failures when speculative decoding leaves auxiliary buffers empty._agent.register_memoryis not called with no valid(ptr, size)descriptors.aux_buffer_ptr_{i}) for surviving entries.QA Engineer Review
Touched:
tests/unittest/disaggregated/test_aux_buffer_registration.pytests/integration/test_lists/test-db/l0_a10.ymltests/integration/test_lists/test-db/l0_h100.ymlTest functions added:
test_null_pointer_aux_buffer_is_skipped()test_zero_size_aux_buffer_is_skipped()test_all_buffers_registered_when_none_are_empty()test_nothing_registered_when_every_buffer_is_empty()Integration coverage (
test-db/):unittest/disaggregated/test_aux_buffer_registration.pyto:tests/integration/test_lists/test-db/l0_a10.ymltests/integration/test_lists/test-db/l0_h100.ymlVerdict: needs follow-up