fix(rdma): restore task.request association in submitTransfer#2610
Conversation
RdmaTransport::submitTransfer(BatchID, entries) was simplified in kvcache-ai#772 to delegate slice construction to submitTransferTask(), but the rewrite dropped the per-entry assignment of task.request. submitTransferTask() then runs `assert(task.request); auto &request = *task.request;`, so every task it receives from this entry point carries a null request pointer. The overload is reached through HeterogeneousRdmaTransport::submitTransfer, which forwards CPU-source and staged-host batches straight to transport_->submitTransfer(batch_id, entries). With the missing assignment that dereferences a null pointer (assert failure in debug, UB/segfault in release). Repopulate batch_id and request for each newly created task before delegating, matching how MultiTransport::submitTransfer and TcpTransport associate each task with its originating request. Signed-off-by: Yufeng He <[email protected]>
There was a problem hiding this comment.
Code Review
This pull request modifies submitTransfer in rdma_transport.cpp to iterate over the input entries and populate TransferTask requests and batch IDs. Feedback highlights a critical memory safety issue where resizing batch_desc.task_list can trigger reallocation and cause a double-free or use-after-free due to the lack of a custom move constructor in TransferTask. Additionally, there is a risk of undefined behavior from casting away constness on entries using const_cast.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| std::vector<TransferTask *> task_list; | ||
| for (auto &task : batch_desc.task_list) task_list.push_back(&task); | ||
| for (auto &request : entries) { | ||
| auto &task = batch_desc.task_list[task_id]; |
There was a problem hiding this comment.
Critical Memory Safety Issue: Potential Double-Free / Use-After-Free of Slices
batch_desc.task_list is a std::vector<TransferTask>. When batch_desc.task_list.resize() is called, if the new size exceeds the vector's current capacity, the vector will reallocate its internal storage.
During reallocation, the existing TransferTask elements are copied or moved to the new memory location, and the old elements are destroyed.
However, TransferTask (defined in transport.h) has a custom destructor ~TransferTask() that deallocates all slices in its slice_list:
~TransferTask() {
for (auto &slice : slice_list)
Transport::getSliceCache().deallocate(slice);
}But TransferTask does not define a custom move constructor or copy constructor (violating the Rule of Three/Five).
- Since a custom destructor is defined, the compiler will not implicitly generate a move constructor.
- It will instead fall back to the default copy constructor (which is deprecated but generated), performing a shallow copy of the
slice_listpointer vector. - When the old
TransferTaskelements are destroyed during reallocation, their destructors will run and deallocate all active slices of any pre-existing tasks in the batch. - The new
TransferTaskelements in the reallocated vector will now hold dangling pointers to these deallocated slices, leading to use-after-free and subsequent double-free when they are eventually destroyed.
Recommendation
To fix this, you should implement a proper noexcept move constructor for TransferTask in transport.h that clears the slice_list of the moved-from object:
TransferTask(TransferTask&& other) noexcept
: slice_count(other.slice_count),
success_slice_count(other.success_slice_count),
failed_slice_count(other.failed_slice_count),
transferred_bytes(other.transferred_bytes),
is_finished(other.is_finished),
total_bytes(other.total_bytes),
batch_id(other.batch_id),
transport_(other.transport_),
slice_list(std::move(other.slice_list)) {
other.slice_list.clear();
}And delete or properly implement the copy constructor/assignment operator to prevent accidental copies.
| ++task_id; | ||
| task.batch_id = batch_id; | ||
| #ifdef USE_ASCEND_HETEROGENEOUS | ||
| task.request = const_cast<TransferRequest *>(&request); |
There was a problem hiding this comment.
Undefined Behavior Risk with const_cast
entries is passed as a const std::vector<TransferRequest> &. Casting away constness using const_cast<TransferRequest *>(&request) and subsequently modifying the request (as implied by the comment in transport.h: // need to modify the request's source address...) leads to undefined behavior if the caller passed a truly const object or a temporary.
To ensure safety and maintain const-correctness, the transport APIs should either:
- Accept
std::vector<TransferRequest> &(non-const) if modification is required. - Avoid modifying the original request object directly, and instead store the modified state (like the CPU source address) in the
TransferTaskorSliceitself.
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
EfaTransport::submitTransfer(BatchID, entries) and UbTransport::submitTransfer(BatchID, entries) resize the batch task list and then hand every slot to submitTransferTask() without setting task.request on the newly created tasks. submitTransferTask() runs `assert(task.request); auto& request = *task.request;`, so each task from this entry point carries a null request pointer (assert failure in debug, UB / segfault in release). This is the same defect kvcache-ai#2610 fixed in RdmaTransport. After that change, EFA and Kunpeng UB are the only two transports whose submitTransfer override still omits the per-entry association that MultiTransport, TcpTransport, and (post-kvcache-ai#2610) RdmaTransport already perform. Repopulate batch_id and request for each newly added entry before delegating, and submit only the new tasks instead of re-pushing the whole list, matching kvcache-ai#2610. Signed-off-by: Yufeng He <[email protected]>
…2617) EfaTransport::submitTransfer(BatchID, entries) and UbTransport::submitTransfer(BatchID, entries) resize the batch task list and then hand every slot to submitTransferTask() without setting task.request on the newly created tasks. submitTransferTask() runs `assert(task.request); auto& request = *task.request;`, so each task from this entry point carries a null request pointer (assert failure in debug, UB / segfault in release). This is the same defect #2610 fixed in RdmaTransport. After that change, EFA and Kunpeng UB are the only two transports whose submitTransfer override still omits the per-entry association that MultiTransport, TcpTransport, and (post-#2610) RdmaTransport already perform. Repopulate batch_id and request for each newly added entry before delegating, and submit only the new tasks instead of re-pushing the whole list, matching #2610. Signed-off-by: Yufeng He <[email protected]>
…e-ai#2610) RdmaTransport::submitTransfer(BatchID, entries) was simplified in kvcache-ai#772 to delegate slice construction to submitTransferTask(), but the rewrite dropped the per-entry assignment of task.request. submitTransferTask() then runs `assert(task.request); auto &request = *task.request;`, so every task it receives from this entry point carries a null request pointer. The overload is reached through HeterogeneousRdmaTransport::submitTransfer, which forwards CPU-source and staged-host batches straight to transport_->submitTransfer(batch_id, entries). With the missing assignment that dereferences a null pointer (assert failure in debug, UB/segfault in release). Repopulate batch_id and request for each newly created task before delegating, matching how MultiTransport::submitTransfer and TcpTransport associate each task with its originating request. Signed-off-by: Yufeng He <[email protected]>
…vcache-ai#2617) EfaTransport::submitTransfer(BatchID, entries) and UbTransport::submitTransfer(BatchID, entries) resize the batch task list and then hand every slot to submitTransferTask() without setting task.request on the newly created tasks. submitTransferTask() runs `assert(task.request); auto& request = *task.request;`, so each task from this entry point carries a null request pointer (assert failure in debug, UB / segfault in release). This is the same defect kvcache-ai#2610 fixed in RdmaTransport. After that change, EFA and Kunpeng UB are the only two transports whose submitTransfer override still omits the per-entry association that MultiTransport, TcpTransport, and (post-kvcache-ai#2610) RdmaTransport already perform. Repopulate batch_id and request for each newly added entry before delegating, and submit only the new tasks instead of re-pushing the whole list, matching kvcache-ai#2610. Signed-off-by: Yufeng He <[email protected]>
Description
RdmaTransport::submitTransfer(BatchID, entries)stopped associating each task with its originating request, so the tasks it hands tosubmitTransferTask()carry a nullrequestpointer.In #772 this overload was simplified to delegate slice construction to
submitTransferTask()instead of building slices inline. The inline loop used to read eachentries[i]directly; the rewrite removed that loop but did not add thetask.request = &entries[i]assignment thatsubmitTransferTask()now depends on:The leftover
task_idvariable in the current code (computed, then only used byresize) is what remains of that deleted per-entry loop.This overload is not on the normal engine path —
MultiTransport::submitTransfersetstask.requestitself and callssubmitTransferTaskdirectly. It is reached throughHeterogeneousRdmaTransport::submitTransfer, which forwards both the CPU-source fast path and the staged-host path totransport_->submitTransfer(batch_id, entries)(wheretransport_is astd::unique_ptr<RdmaTransport>). On that path every submitted task hasrequest == nullptr: an assertion failure in debug builds, and a null-pointer dereference (request.length,request.source, …) in release builds.The fix repopulates
batch_idandrequestfor each newly created task before delegating, mirroring whatMultiTransport::submitTransferdoes (including theUSE_ASCEND_HETEROGENEOUSconst handling, sinceTransferTask::requestis non-const in that build) and whatTcpTransport::submitTransferdoes inline. As a side effect it also stops re-submitting pre-existing tasks in the batch, since it now iterates the new entries rather than the wholetask_list.Module
mooncake-transfer-engine)Type of Change
How Has This Been Tested?
I have not been able to exercise this path at runtime: the only caller is
HeterogeneousRdmaTransport, which is built underUSE_ASCEND_HETEROGENEOUSand needs Ascend NPU hardware, andRdmaTransportitself needs an RDMA-capable NIC to instantiate. There is no hardware-free unit harness for it in the tree.So this is verified by code inspection rather than execution:
requestconsumption was removed without a replacement assignment);submitTransferTask()unconditionally requirestask.request;MultiTransport::submitTransfer/TcpTransport::submitTransfersiblings;I'd appreciate a check from someone with the Ascend heterogeneous setup to confirm the runtime path end to end.
Test results: