[TE] Enabling UB Transport on the Kunpeng SuperNode Phase 2 by zchuango · Pull Request #1855 · kvcache-ai/Mooncake

zchuango · 2026-04-09T12:25:52Z

Description

Overview

This PR introduces UbTransport, a new transport implementation for Mooncake's Transfer Engine that enables high-performance data transfer on Kunpeng 950 platforms using the UB (Unified Bus) protocol and URMA (Unified Remote Memory Access) APIs from the openEuler UMDK.

Related Issue: #1773

Phase Contribution Plan

Following suggestion, we propose four phased PRs for incremental review and integration:

Phase 1: Enable UbTransport + UrmaEndpoint Phase 1 PR
Phase 2: CI/Testing Infrastructure With Kunpeng UB Hardware Mock Env ⬅️ Current PR
Phase 3: Enable UbTransport + OBMM Feature
Phase 4: Next TE Tent Module Integration with UbTransport

New Feature

This Phase2 branch has introduced the following new features compared to the main branch:

Mock URMA Implementation : Provided a mock URMA implementation ( mock_urma.cpp ) for testing purposes, allowing tests to run without actual URMA hardware.
Unit Tests : Added comprehensive test cases in ub_transport_test.cpp to verify the functionality of the new transport module.
Build System Integration : Updated CMake configurations to integrate the new Kunpeng transport module into the existing build system.

Module

Type of Change

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have formatted my own code using ./scripts/code_format.sh before submitting.
I have updated the documentation.
I have added tests to prove my changes are effective.

gemini-code-assist

Code Review

This pull request introduces support for Kunpeng UB (Unified Bus) transport using URMA (Unified Remote Memory Access), including a mock implementation for testing and updated topology discovery logic. Feedback identifies a compilation error in topology.cpp caused by a misplaced header, redundant initialization in rdma_transport_test2.cpp, and documentation inaccuracies. It is also suggested to clean up debug messages in the CMake configuration to improve build output.

codecov-commenter · 2026-04-09T13:04:32Z

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

Codecov Report

❌ Patch coverage is 82.48175% with 72 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...src/transport/kunpeng_transport/urma/mock_urma.cpp	79.77%	54 Missing ⚠️
...oncake-transfer-engine/tests/ub_transport_test.cpp	92.07%	8 Missing ⚠️
...transport/kunpeng_transport/urma/urma_endpoint.cpp	82.14%	5 Missing ⚠️
...e/src/transport/kunpeng_transport/ub_transport.cpp	55.55%	4 Missing ⚠️
...ake-transfer-engine/tests/rdma_transport_test2.cpp	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

zchuango · 2026-04-10T13:34:37Z

@alogfans @stmatengss ub trasport feature phase2 pr has been finshed, kindly help review and verify.🙏

stmatengss · 2026-04-12T04:05:35Z

 }

-int UrmaContext::openDevice(const std::string& device_name, uint8_t port,
+int UrmaContext::openDevice(const std::string& device_name, uint8_t por,


por is a typo?

oh，This might be a mistake made during CI testing coding，would fix it in the next phase？

Please fix this typo. I think it will influence the correctness of this method.

Please fix this typo. I think it will influence the correctness of this method.

okay, I have fixed this typo. @stmatengss

stmatengss · 2026-04-12T04:09:10Z

 int Topology::discover(const std::vector<std::string> &filter) {
    matrix_.clear();
-#ifdef USE_UB
-    auto all_hca = listUBDevices(filter);
-    for (auto &ent : discoverCpuTopology(all_hca)) {
-        matrix_[ent.name] = ent;
-    }
-#else
    auto all_hca = listInfiniBandDevices(filter);
    for (auto &ent : discoverCpuTopology(all_hca)) {
        matrix_[ent.name] = ent;
    }
-#endif
 #if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP) || \
    defined(USE_MLU) || defined(USE_MACA)
    for (auto &ent : discoverCudaTopology(all_hca)) {
        matrix_[ent.name] = ent;
    }
+#endif
+#ifdef USE_UB
+    auto ub_all_hca = listUBDevices(filter);
+    for (auto &ent : discoverCpuTopology(ub_all_hca)) {
+        matrix_[ent.name] = ent;
+    }


Consider refactoring this code. What's the difference between CUDA/MUSA/HIP/MLU and UB?

listUBDevices(filter) return UbDevice vector but listInfiniBandDevices(filter) is InfinibandDevice vector, CUDA/MUSA/HIP/MLU method discoverCudaTopology(hca) current using the InfinibandDevice vector， I wonder if we could have a unified device class for this? @stmatengss

Sure. We need a unified device class for different vendors. It can be the next PR

Okay, I'll design it here and submit the reference code in the next PR.

alogfans

LGTM

zchuango added 11 commits April 9, 2026 04:18

add the ub_transport_test UT

e13ec85

modify the mock urma api src

53fe6ad

remove the URMA_MOCK and refine the code for ub_transport_test

33a2c3d

format the code

b2144ef

add the ub transport ci

226d792

fix the ub transport ci bug

986b72c

optimize the ci config and rename the urma api mock source name

a1f6b98

fix the ci config bug

85dc606

fix the memory leakage

b6e2311

fix the code format bug

58a87e0

remove the build-flags with USE_UB

421482f

zchuango requested review from Ann-1024, ShangmingCai, alogfans, chestnut-Q, doujiang24, luketong777, stmatengss and ykwd as code owners April 9, 2026 12:25

github-actions bot added run-ci Transfer Engine labels Apr 9, 2026

gemini-code-assist bot reviewed Apr 9, 2026

View reviewed changes

Comment thread mooncake-transfer-engine/src/topology.cpp

Comment thread docs/source/design/transfer-engine/kunpeng_ub_transport.md

Comment thread mooncake-common/FindUrma.cmake

Comment thread mooncake-common/FindUrma.cmake

Comment thread mooncake-transfer-engine/tests/rdma_transport_test2.cpp

alogfans reviewed Apr 10, 2026

View reviewed changes

Comment thread .github/workflows/ci.yml

optimize the urma port find

841d92c

zchuango mentioned this pull request Apr 10, 2026

[RFC]: Enable UB transport via Kunpeng UMDK URMA apis #1773

Open

1 task

zchuango requested a review from alogfans April 11, 2026 07:50

stmatengss reviewed Apr 12, 2026

View reviewed changes

alogfans approved these changes Apr 13, 2026

View reviewed changes

modify the error code name

6e74a7b

zchuango requested a review from stmatengss April 14, 2026 08:04

stmatengss merged commit 3a69fa4 into kvcache-ai:main Apr 14, 2026
17 checks passed

Conversation

zchuango commented Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Description

Overview

Phase Contribution Plan

New Feature

Module

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

Uh oh!

zchuango commented Apr 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

stmatengss Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

zchuango Apr 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stmatengss Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

zchuango Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

zchuango Apr 12, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss Apr 13, 2026

Choose a reason for hiding this comment

Uh oh!

zchuango Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

alogfans left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

zchuango commented Apr 9, 2026 •

edited

Loading

codecov-commenter commented Apr 9, 2026 •

edited

Loading

zchuango commented Apr 10, 2026 •

edited

Loading

zchuango Apr 12, 2026 •

edited

Loading