Skip to content

fix: fix heterogeneous rdma transport error #1657

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stmatengss merged 9 commits into from
Mar 17, 2026
Merged

fix: fix heterogeneous rdma transport error #1657

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
a30dfc7
Merge pull request #1 from kvcache-ai/main
XingSong-Sun Mar 13, 2026
c79a4de
fix: fix heterogeneous rdma transport bug and add test case
Mar 13, 2026
4b9f292
Merge remote-tracking branch 'origin/main'
Mar 13, 2026
7041ac9
delete test code
Mar 14, 2026
7a68c81
Use constant when judging memory locations
Mar 16, 2026
7e3f8e5
fix: format string concatenation
Mar 17, 2026
eb740f8
fix: format the code issues by the code check
Mar 17, 2026
63d2021
Merge branch 'kvcache-ai:main' into main
XingSong-Sun Mar 17, 2026
1c2c945
fix: update main branch and delete friend for HeterogeneousRdmaTransport
Mar 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions mooncake-transfer-engine/example/transfer_engine_heterogeneous_ascend_perf_initiator.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -159,8 +159,9 @@ std::string loadNicPriorityMatrix() {
"], []], "
" \"cpu:1\": [[" +
device_names +
"], []], "
" \"cuda:0\": [[" +
"], []], " " \"npu:0\": [["
+ device_names +
"], []], " " \"cuda:0\": [[" +
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The string concatenation for adding the npu:0 entry has inconsistent formatting, which makes the code harder to read. Please reformat this section for better clarity and consistency with the surrounding code.

Suggested change
"], []], " " \"npu:0\": [["
+ device_names +
"], []], " " \"cuda:0\": [[" +
"], []], "
" \"npu:0\": [[" +
device_names +
"], []], "
" \"cuda:0\": [[" +

device_names +
"], []], "
" \"musa:0\": [[" +
Expand Down
3 changes: 2 additions & 1 deletion mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,12 +37,13 @@ class RdmaContext;
class RdmaEndPoint;
class TransferMetadata;
class WorkerPool;
class HeterogeneousRdmaTransport;

class RdmaTransport : public Transport {
friend class RdmaContext;
friend class RdmaEndPoint;
friend class WorkerPool;

friend class HeterogeneousRdmaTransport;
Copy link
Copy Markdown
Collaborator

@staryxchen staryxchen Mar 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you clarify which private member(s) this is needed for? If none currently, this should be removed — friend breaks encapsulation and shouldn't be added speculatively.

Copy link
Copy Markdown
Contributor Author

@XingSong-Sun XingSong-Sun Mar 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HeterogeneousRdmaTransport needs to access the private members submitTransfer, submitTransferTask, and getTransferStatus in RdmaTransport. If friend is not added, it will result in a compilation error.

Copy link
Copy Markdown
Collaborator

@staryxchen staryxchen Mar 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

HeterogeneousRdmaTransport needs to access the private members submitTransfer, submitTransferTask, and getTransferStatus in RdmaTransport. If friend is not added, it will result in a compilation error.

Thanks for the explanation — that makes sense for the original code. However, PR #1663 has since been merged, which moved submitTransfer and related methods from private to public. So the friend class declaration is no longer needed. Could you rebase onto the latest main and drop that change?

Copy link
Copy Markdown
Contributor Author

@XingSong-Sun XingSong-Sun Mar 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The corresponding modifications have been submitted.

public:
using BufferDesc = TransferMetadata::BufferDesc;
using SegmentDesc = TransferMetadata::SegmentDesc;
Expand Down
14 changes: 11 additions & 3 deletions .../transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,21 @@
namespace mooncake {

namespace {
// ACL memory location types for aclrtPtrAttributes::location.type:
// 0 = ACL host memory, 1 = device memory, 2 = regular CPU memory (malloc)
static constexpr int kDeviceMemoryLocationType = 1;

bool isCpuMemory(void *addr) {
aclrtPtrAttributes attributes{};
if (int ret = aclrtPointerGetAttributes(addr, &attributes)) {
LOG(ERROR) << "aclrtPointrtGetAttributes error, ret: " << ret;
return false;
// If ACL cannot identify the pointer, it is not ACL-managed device
// memory, so treat it as regular CPU memory.
LOG(WARNING) << "aclrtPointerGetAttributes failed for addr " << addr
Copy link
Copy Markdown
Collaborator

@staryxchen staryxchen Mar 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

reformat the example file string concatenation

Copy link
Copy Markdown
Contributor Author

@XingSong-Sun XingSong-Sun Mar 17, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

done.

<< ", ret: " << ret
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The logic for determining if a memory address belongs to the CPU has been updated. The previous implementation only checked for ACL_HOST memory (type == 0) and incorrectly returned false on aclrtPointerGetAttributes failure. The new logic correctly treats any memory that is not explicitly device memory (type != 1) as CPU memory. This includes ACL_HOST memory (type == 0), regular CPU memory (type == 2), and cases where aclrtPointerGetAttributes fails. This change is crucial for correctness.

<< ". Treating as CPU memory.";
return true;
}
return (attributes.location.type == 0);
return (attributes.location.type != kDeviceMemoryLocationType);
}
} // namespace

Expand Down
Loading