[TransferEngine][WIP]heterogeneous-tcp #1070

AscendTransport · 2025-11-18T01:52:05Z

Description

Type of Change

Types
- Bug fix
- New feature
  - Transfer Engine
  - Mooncake Store
  - Mooncake EP
  - Integration
  - P2P Store
  - Python Wheel
- Breaking change
- CI/CD
- Documentation update
- Other

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have updated the documentation.
I have added tests to prove my changes are effective.

gemini-code-assist · 2025-11-18T01:54:10Z

Summary of Changes

Hello @AscendTransport, 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!

This pull request introduces a new HeterogeneousTcpTransport to enable TCP-based data transfers between Ascend NPUs and GPUs. This new transport mechanism works by managing dedicated host and device memory buffers and using asynchronous memory copy operations to move data between the NPU and host, before leveraging the existing TcpTransport for host-to-host communication. The changes include updates to the build system and core transfer engine logic to support this new heterogeneous transfer path.

Highlights

New TCP Transport for Heterogeneous Ascend/GPU: Introduced a new CMake option USE_ASCEND_HETEROGENEOUS_TCP and a corresponding HeterogeneousTcpTransport class to enable TCP-based data transfers between Ascend NPUs and GPUs.
Memory Management for Device-Host Transfers: The HeterogeneousTcpTransport manages dedicated host (aligned_alloc) and device (aclrtMalloc) memory buffers, utilizing asynchronous memory copy operations (aclrtMemcpyAsync) to efficiently move data between the NPU and host before TCP transmission.
Integration with Existing Transfer Engine: The new transport is seamlessly integrated into the MultiTransport and TransferEngine components, allowing it to be selected and used based on the build configuration.
Build System Updates: CMake files across mooncake-common, mooncake-transfer-engine, and mooncake-integration have been updated to support the new build option, ensuring proper compilation and linking of the new transport.
TcpTransport::install Visibility: The install method in TcpTransport was changed from private to public, enabling HeterogeneousTcpTransport to properly utilize its functionality for underlying TCP communication.

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 introduces a new TCP-based transport for heterogeneous communication involving Ascend NPUs. The changes are extensive, touching CMake build scripts, integration code, and the core transfer engine. While the feature addition is significant, the current implementation has several critical issues that must be addressed. These include compilation errors from incorrect include paths and non-standard function usage, potential data corruption and race conditions due to incorrect types and synchronization, and performance bottlenecks from busy-waiting. I've provided specific comments and suggestions to resolve these problems.

gemini-code-assist · 2025-11-18T01:56:17Z

mooncake-transfer-engine/src/multi_transport.cpp

+#include "transport/ascend_transport/heterogeneous_transport/heterogeneous_rdma_transport.h"
+#endif
+#ifdef USE_ASCEND_HETEROGENEOUS_TCP
+#include "transport/ascend_transport/heterogeneous_transport/heterogeneous_tcp_transport.h"


The include path for heterogeneous_tcp_transport.h appears to be incorrect. Based on the file's location (mooncake-transfer-engine/include/transport/ascend_transport/heterogeneous_tcp_transport.h) and the project's include paths, the extra heterogeneous_transport/ directory in the path is likely a copy-paste error and will cause a compilation failure.

Suggested change

#include "transport/ascend_transport/heterogeneous_transport/heterogeneous_tcp_transport.h"

#include "transport/ascend_transport/heterogeneous_tcp_transport.h"

The include path for heterogeneous_tcp_transport.h appears to be incorrect. Based on the file's location (mooncake-transfer-engine/include/transport/ascend_transport/heterogeneous_tcp_transport.h) and the project's include paths, the extra heterogeneous_transport/ directory in the path is likely a copy-paste error and will cause a compilation failure.

fixed

gemini-code-assist · 2025-11-18T01:56:18Z

mooncake-transfer-engine/src/multi_transport.cpp

+#ifdef USE_ASCEND_HETEROGENEOUS_TCP
+    else if (std::string(proto) == "ascend") {
+        transport = new HeterogeneousTcpTransport();
+    }
+#endif


This else if block is conditional on USE_ASCEND_HETEROGENEOUS_TCP, but it's syntactically attached to the preceding if-else if chain. If USE_ASCEND_HETEROGENEOUS is disabled, this else if will be dangling and cause a compilation error. The entire block for selecting an ascend transport should be restructured to handle the mutually exclusive options correctly, for example by using a single if-else if chain with #if/#elif/#endif preprocessor directives to ensure correctness regardless of which flags are enabled.

This else if block is conditional on USE_ASCEND_HETEROGENEOUS_TCP, but it's syntactically attached to the preceding if-else if chain. If USE_ASCEND_HETEROGENEOUS is disabled, this else if will be dangling and cause a compilation error. The entire block for selecting an ascend transport should be restructured to handle the mutually exclusive options correctly, for example by using a single if-else if chain with #if/#elif/#endif preprocessor directives to ensure correctness regardless of which flags are enabled.

fixed

gemini-code-assist · 2025-11-18T01:56:18Z

...e/src/transport/ascend_transport/heterogeneous_tcp_transport/heterogeneous_tcp_transport.cpp

+        return ret;
+    }
+
+    hostAddr_ = std::aligned_alloc(64, HUGE_HOST_SIZE);


std::aligned_alloc is not part of the C++ standard; it's a C11 function. Using std::aligned_alloc will likely cause compilation errors on standard-compliant compilers. You should include <cstdlib> and use aligned_alloc. The memory is correctly freed with free() in the destructor.

Suggested change

hostAddr_ = std::aligned_alloc(64, HUGE_HOST_SIZE);

hostAddr_ = static_cast<char*>(aligned_alloc(64, HUGE_HOST_SIZE));

std::aligned_alloc is not part of the C++ standard; it's a C11 function. Using std::aligned_alloc will likely cause compilation errors on standard-compliant compilers. You should include <cstdlib> and use aligned_alloc. The memory is correctly freed with free() in the destructor.

fixed

gemini-code-assist · 2025-11-18T01:56:18Z

...e/src/transport/ascend_transport/heterogeneous_tcp_transport/heterogeneous_tcp_transport.cpp

+        if (transferQueues_.empty()) {
+            transfer_cond_.wait(lock);
+        }


The wait on the condition variable should be done in a loop or with a predicate to handle spurious wakeups. The current implementation might cause the thread to wake up, find the queue empty, and then crash when trying to access transferQueues_.front(). A predicate in the wait call is the standard way to solve this.

Suggested change

if (transferQueues_.empty()) {

transfer_cond_.wait(lock);

}

transfer_cond_.wait(lock, [this] { return !transferQueues_.empty() || !running_; });

if (transferQueues_.empty()) {

continue;

}

The wait on the condition variable should be done in a loop or with a predicate to handle spurious wakeups. The current implementation might cause the thread to wake up, find the queue empty, and then crash when trying to access transferQueues_.front(). A predicate in the wait call is the standard way to solve this.

fixed

gemini-code-assist · 2025-11-18T01:56:18Z

...e/src/transport/ascend_transport/heterogeneous_tcp_transport/heterogeneous_tcp_transport.cpp

+                << "HeterogeneousTcpTransport: empty transfer task batch";
+            continue;
+        }
+        int total_length = pkg.total_length;


There is a potential data truncation here. pkg.total_length is a uint64_t, but it's being assigned to an int. On a typical 64-bit system, int is 32 bits, which can lead to incorrect length if the transfer size exceeds 2GB. The total_length variable should be uint64_t to match the type in the TransferTaskTCP struct.

Suggested change

int total_length = pkg.total_length;

uint64_t total_length = pkg.total_length;

There is a potential data truncation here. pkg.total_length is a uint64_t, but it's being assigned to an int. On a typical 64-bit system, int is 32 bits, which can lead to incorrect length if the transfer size exceeds 2GB. The total_length variable should be uint64_t to match the type in the TransferTaskTCP struct.

fixed

gemini-code-assist · 2025-11-18T01:56:19Z

mooncake-transfer-engine/include/transport/ascend_transport/heterogeneous_tcp_transport.h

+#define HUGE_HOST_SIZE 3ULL * 1024 * 1024 * 1024
+#define HUGE_DEVICE_SIZE 8 * 1024 * 1024


It's a good practice to wrap macro definitions containing expressions in parentheses to avoid potential operator precedence issues when they are used in more complex expressions. While it may not be causing a bug here, it improves code robustness.

Suggested change

#define HUGE_HOST_SIZE 3ULL * 1024 * 1024 * 1024

#define HUGE_DEVICE_SIZE 8 * 1024 * 1024

#define HUGE_HOST_SIZE (3ULL * 1024 * 1024 * 1024)

#define HUGE_DEVICE_SIZE (8 * 1024 * 1024)

It's a good practice to wrap macro definitions containing expressions in parentheses to avoid potential operator precedence issues when they are used in more complex expressions. While it may not be causing a bug here, it improves code robustness.

fixed

gemini-code-assist · 2025-11-18T01:56:19Z

mooncake-transfer-engine/include/transport/ascend_transport/heterogeneous_tcp_transport.h

+    Status getTransferStatus(BatchID batch_id,
+                             std::vector<TransferStatus> &status);


This overload of getTransferStatus is declared but not defined in the corresponding .cpp file. It also doesn't override any base class method. If it's not used, it should be removed to avoid confusion and dead code.

This overload of getTransferStatus is declared but not defined in the corresponding .cpp file. It also doesn't override any base class method. If it's not used, it should be removed to avoid confusion and dead code.
fixed

gemini-code-assist · 2025-11-18T01:56:19Z

...e/src/transport/ascend_transport/heterogeneous_tcp_transport/heterogeneous_tcp_transport.cpp

+        s = transport_->submitTransferTask(task_list);
+        if (!s.ok()) {
+            LOG(ERROR)
+                << "HeterogeneousTcpTransport: Rdma submitTransferTask error";


The error message mentions Rdma, which seems to be a copy-paste error from another transport implementation. This should be corrected to Tcp or HeterogeneousTcp for accurate logging.

Suggested change

<< "HeterogeneousTcpTransport: Rdma submitTransferTask error";

<< "HeterogeneousTcpTransport: Tcp submitTransferTask error";

gemini-code-assist · 2025-11-18T01:56:19Z

...e/src/transport/ascend_transport/heterogeneous_tcp_transport/heterogeneous_tcp_transport.cpp

+    if (transport_ == nullptr) {
+        LOG(ERROR) << "HeterogeneousTcpTransport:transport is null";
+        return ret;
+    }


If transport_ is null, this function returns ret, which holds the value from aclrtGetDevice. This is likely not the intended error code. It would be better to return a specific error code, like -1, to indicate the failure due to the null transport.

Suggested change

if (transport_ == nullptr) {

LOG(ERROR) << "HeterogeneousTcpTransport:transport is null";

return ret;

}

if (transport_ == nullptr) {

LOG(ERROR) << "HeterogeneousTcpTransport:transport is null";

return -1;

}

gemini-code-assist · 2025-11-18T01:56:19Z

...e/src/transport/ascend_transport/heterogeneous_tcp_transport/heterogeneous_tcp_transport.cpp

+
+    ret = transport_->install(local_server_name_, meta, topo);
+    if (ret) {
+        LOG(ERROR) << "HeterogeneousTcpTransport::RdmaTransport install error, ret: " << ret;


The error message mentions RdmaTransport, which seems to be a copy-paste error from another transport implementation. This should be corrected to TcpTransport for accurate logging.

Suggested change

LOG(ERROR) << "HeterogeneousTcpTransport::RdmaTransport install error, ret: " << ret;

LOG(ERROR) << "HeterogeneousTcpTransport::TcpTransport install error, ret: " << ret;

AscendTransport changed the title ~~heterogeneous-tcp~~ [TransferEngine][WIP]heterogeneous-tcp Nov 18, 2025

gemini-code-assist bot reviewed Nov 18, 2025

View reviewed changes

heterogeneous-tcp

f0daa9e

AscendTransport force-pushed the ascend-tcp-agg branch from fe45a11 to f0daa9e Compare November 19, 2025 07:03

	#include "transport/ascend_transport/heterogeneous_transport/heterogeneous_tcp_transport.h"
	#include "transport/ascend_transport/heterogeneous_tcp_transport.h"

	hostAddr_ = std::aligned_alloc(64, HUGE_HOST_SIZE);
	hostAddr_ = static_cast<char*>(aligned_alloc(64, HUGE_HOST_SIZE));

	int total_length = pkg.total_length;
	uint64_t total_length = pkg.total_length;

		#define HUGE_HOST_SIZE 3ULL * 1024 * 1024 * 1024
		#define HUGE_DEVICE_SIZE 8 * 1024 * 1024

		Status getTransferStatus(BatchID batch_id,
		std::vector<TransferStatus> &status);

	<< "HeterogeneousTcpTransport: Rdma submitTransferTask error";
	<< "HeterogeneousTcpTransport: Tcp submitTransferTask error";

	LOG(ERROR) << "HeterogeneousTcpTransport::RdmaTransport install error, ret: " << ret;
	LOG(ERROR) << "HeterogeneousTcpTransport::TcpTransport install error, ret: " << ret;

[TransferEngine][WIP]heterogeneous-tcp #1070

Are you sure you want to change the base?

[TransferEngine][WIP]heterogeneous-tcp #1070

Uh oh!

Conversation

AscendTransport commented Nov 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist bot commented Nov 18, 2025

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 Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Nov 18, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

AscendTransport commented Nov 18, 2025 •

edited

Loading