fix: fix heterogeneous rdma transport error by XingSong-Sun · Pull Request #1657 · kvcache-ai/Mooncake

XingSong-Sun · 2026-03-14T08:10:16Z

Description

修复异构RDMA Transport对cpu memory的判断错误导致的NPU报错问题

Enhance Store and TransferEngine with health check, metrics, and NUMA support

gemini-code-assist · 2026-03-14T08:10:29Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

此拉取请求旨在解决异构RDMA传输中一个关键的内存识别问题。通过修正 isCpuMemory 函数的判断逻辑，确保系统能够准确区分CPU内存和设备内存，从而避免了因内存类型误判而导致的NPU错误。此外，还更新了示例配置以包含NPU设备，并调整了类之间的友元关系以支持这些更改。

Highlights

修复异构RDMA传输错误: 解决了异构RDMA Transport对CPU内存判断不准确导致NPU报错的问题。
内存判断逻辑优化: 更新了 isCpuMemory 函数的逻辑，使其在 aclrtPointerGetAttributes 调用失败时，或当 location.type 不为1（设备内存）时，能正确地将内存识别为CPU内存。
NPU设备配置: 在示例配置文件中，为 nicPriorityMatrix 添加了 npu:0 条目。
类友元关系: 将 HeterogeneousRdmaTransport 类声明为 RdmaTransport 类的友元。

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

mooncake-transfer-engine/example/transfer_engine_heterogeneous_ascend_perf_initiator.cpp
- 为NIC优先级矩阵添加了NPU设备配置。
mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h
- 将 HeterogeneousRdmaTransport 声明为 RdmaTransport 的友元类。
mooncake-transfer-engine/src/transport/ascend_transport/heterogeneous_rdma_transport/heterogeneous_rdma_transport.cpp
- 修正了 isCpuMemory 函数的逻辑，以正确处理 aclrtPointerGetAttributes 失败的情况，并根据 location.type 准确判断CPU内存。

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request fixes an issue with heterogeneous RDMA transport where CPU memory was being incorrectly identified, leading to NPU errors. The main change is in isCpuMemory to correctly classify memory types. Additionally, a test file is updated to include NPU devices, and a header file is modified to add a friend class declaration. My review focuses on improving code clarity and correctness.

gemini-code-assist · 2026-03-14T08:15:55Z

    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
+                     << ", ret: " << ret


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.

gemini-code-assist · 2026-03-14T08:15:55Z

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


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\": [[" +

codecov-commenter · 2026-03-14T08:43:08Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

alogfans · 2026-03-16T01:38:13Z

-    return (attributes.location.type == 0);
+    // location.type: 0 = ACL host memory, 1 = device memory, 2 = regular
+    // CPU memory (malloc). Only type 1 is device memory; all others are CPU.
+    return (attributes.location.type != 1);


Use constant value if possible.

staryxchen · 2026-03-16T14:52:31Z

    friend class RdmaEndPoint;
    friend class WorkerPool;
-
+    friend class HeterogeneousRdmaTransport;


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.

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.

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?

The corresponding modifications have been submitted.

staryxchen · 2026-03-16T14:53:22Z

-        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


reformat the example file string concatenation

staryxchen

LGTM

--------- Co-authored-by: sunwenhan <sunwenhan@xfusion.com>

XingSong-Sun and others added 4 commits March 13, 2026 15:16

Merge pull request #1 from kvcache-ai/main

a30dfc7

Enhance Store and TransferEngine with health check, metrics, and NUMA support

fix: fix heterogeneous rdma transport bug and add test case

c79a4de

Merge remote-tracking branch 'origin/main'

4b9f292

delete test code

7041ac9

XingSong-Sun requested review from alogfans, ascend-direct-dev, chestnut-Q and doujiang24 as code owners March 14, 2026 08:10

github-actions Bot added run-ci Transfer Engine labels Mar 14, 2026

gemini-code-assist Bot reviewed Mar 14, 2026

View reviewed changes

alogfans reviewed Mar 16, 2026

View reviewed changes

Use constant when judging memory locations

7a68c81

XingSong-Sun force-pushed the main branch from 7275983 to 7a68c81 Compare March 16, 2026 01:58

staryxchen reviewed Mar 16, 2026

View reviewed changes

sunwenhan added 2 commits March 17, 2026 09:43

fix: format string concatenation

7e3f8e5

fix: format the code issues by the code check

eb740f8

XingSong-Sun force-pushed the main branch from 7411292 to eb740f8 Compare March 17, 2026 02:10

XingSong-Sun and others added 2 commits March 17, 2026 11:04

Merge branch 'kvcache-ai:main' into main

63d2021

fix: update main branch and delete friend for HeterogeneousRdmaTransport

1c2c945

staryxchen approved these changes Mar 17, 2026

View reviewed changes

stmatengss merged commit 335de42 into kvcache-ai:main Mar 17, 2026
16 of 17 checks passed

whn09 pushed a commit to whn09/Mooncake that referenced this pull request Apr 4, 2026

[TE] fix: fix heterogeneous rdma transport error (kvcache-ai#1657)

fa94caa

--------- Co-authored-by: sunwenhan <sunwenhan@xfusion.com>

Conversation

XingSong-Sun commented Mar 14, 2026

Description

Uh oh!

gemini-code-assist Bot commented Mar 14, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Mar 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Mar 14, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Mar 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

staryxchen left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

codecov-commenter commented Mar 14, 2026 •

edited

Loading