[PR #1399] [MERGED] [TE] Enable ubshmem transport via ascend vmm apis #2221

New issue

Closed

opened 2026-10-29 21:33:46 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:33:46 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/Mooncake/pull/1399
Author: @VNightMare
Created: 1/18/2026
Status: ✅ Merged
Merged: 1/26/2026
Merged by: @alogfans

Base: main ← Head: ubshmem_transport_dev

📝 Commits (1)

0f677e7 ubshmem transport with fabric mem without conflicts

📊 Changes

16 files changed (+1103 additions, -23 deletions)

View changed files

📝 mooncake-common/common.cmake (+7 -1)
📝 mooncake-transfer-engine/example/CMakeLists.txt (+4 -0)
📝 mooncake-transfer-engine/example/transfer_engine_bench.cpp (+79 -11)
📝 mooncake-transfer-engine/include/cuda_alike.h (+3 -1)
📝 mooncake-transfer-engine/include/gpu_vendor/mnnvl.h (+7 -1)
➕ mooncake-transfer-engine/include/gpu_vendor/ubshmem.h (+4 -0)
➕ mooncake-transfer-engine/include/transport/ascend_transport/ubshmem_transport/ubshmem_transport.h (+99 -0)
📝 mooncake-transfer-engine/include/transport/transport.h (+3 -0)
📝 mooncake-transfer-engine/src/CMakeLists.txt (+9 -7)
📝 mooncake-transfer-engine/src/multi_transport.cpp (+8 -0)
📝 mooncake-transfer-engine/src/transfer_metadata.cpp (+3 -2)
📝 mooncake-transfer-engine/src/transport/CMakeLists.txt (+5 -0)
📝 mooncake-transfer-engine/src/transport/ascend_transport/CMakeLists.txt (+6 -0)
➕ mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp (+541 -0)
📝 mooncake-transfer-engine/tests/CMakeLists.txt (+6 -0)
➕ mooncake-transfer-engine/tests/ubshmem_transport_test.cpp (+319 -0)

📄 Description

Description

CANN provides VMM (Virutal Memory Management) APIs to support Fabric Memory for Ascend chips, which allows for building NVLink-like Transport as an additional supplement to Ascend Direct Transport.
Related issue: https://github.com/kvcache-ai/Mooncake/issues/1312

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?

We run mooncake-transfer-engine/example/transfer-engine-bench, with 12 intiator threads and batch_size set to 1 to process D2RD transfer tasks, get the performance of ubshmem as below

block_size	operation	throughput（GB/s）
64K	write/read	7.76 / 8.34
128K	write/read	15.77 / 15.73
256K	write/read	31.86 / 30.27
512K	write/read	60.49 / 58.31
1M	write/read	108.83 / 108.57
2M	write/read	118.57 / 148.85
4M	write/read	131.45 / 158.17
8M	write/read	134.60 / 162.70
12M	write/read	135.69 / 164.20

Add a test case mooncake-transfer-engine/tests/ubshmem_transport_test to test transfer functionality and accuracy.

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.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/Mooncake/pull/1399 **Author:** [@VNightMare](https://github.com/VNightMare) **Created:** 1/18/2026 **Status:** ✅ Merged **Merged:** 1/26/2026 **Merged by:** [@alogfans](https://github.com/alogfans) **Base:** `main` ← **Head:** `ubshmem_transport_dev` --- ### 📝 Commits (1) - [`0f677e7`](https://github.com/kvcache-ai/Mooncake/commit/0f677e72bcd375865be25387691163004b995fde) ubshmem transport with fabric mem without conflicts ### 📊 Changes **16 files changed** (+1103 additions, -23 deletions) <details> <summary>View changed files</summary> 📝 `mooncake-common/common.cmake` (+7 -1) 📝 `mooncake-transfer-engine/example/CMakeLists.txt` (+4 -0) 📝 `mooncake-transfer-engine/example/transfer_engine_bench.cpp` (+79 -11) 📝 `mooncake-transfer-engine/include/cuda_alike.h` (+3 -1) 📝 `mooncake-transfer-engine/include/gpu_vendor/mnnvl.h` (+7 -1) ➕ `mooncake-transfer-engine/include/gpu_vendor/ubshmem.h` (+4 -0) ➕ `mooncake-transfer-engine/include/transport/ascend_transport/ubshmem_transport/ubshmem_transport.h` (+99 -0) 📝 `mooncake-transfer-engine/include/transport/transport.h` (+3 -0) 📝 `mooncake-transfer-engine/src/CMakeLists.txt` (+9 -7) 📝 `mooncake-transfer-engine/src/multi_transport.cpp` (+8 -0) 📝 `mooncake-transfer-engine/src/transfer_metadata.cpp` (+3 -2) 📝 `mooncake-transfer-engine/src/transport/CMakeLists.txt` (+5 -0) 📝 `mooncake-transfer-engine/src/transport/ascend_transport/CMakeLists.txt` (+6 -0) ➕ `mooncake-transfer-engine/src/transport/ascend_transport/ubshmem_transport/ubshmem_transport.cpp` (+541 -0) 📝 `mooncake-transfer-engine/tests/CMakeLists.txt` (+6 -0) ➕ `mooncake-transfer-engine/tests/ubshmem_transport_test.cpp` (+319 -0) </details> ### 📄 Description ## Description  CANN provides VMM (Virutal Memory Management) APIs to support Fabric Memory for Ascend chips, which allows for building NVLink-like Transport as an additional supplement to Ascend Direct Transport. Related issue: https://github.com/kvcache-ai/Mooncake/issues/1312 ## Type of Change * Types - [ ] Bug fix - [x] New feature - [x] Transfer Engine - [ ] Mooncake Store - [ ] Mooncake EP - [ ] Integration - [ ] P2P Store - [ ] Python Wheel - [ ] Breaking change - [ ] CI/CD - [ ] Documentation update - [ ] Other ## How Has This Been Tested? We run `mooncake-transfer-engine/example/transfer-engine-bench`, with 12 intiator threads and batch_size set to 1 to process D2RD transfer tasks, get the performance of ubshmem as below | block_size | operation | throughput（GB/s） | | ---------- | ---------- | ------------------ | | 64K | write/read | 7.76 / 8.34 | | 128K | write/read | 15.77 / 15.73 | | 256K | write/read | 31.86 / 30.27 | | 512K | write/read | 60.49 / 58.31 | | 1M | write/read | 108.83 / 108.57 | | 2M | write/read | 118.57 / 148.85 | | 4M | write/read | 131.45 / 158.17 | | 8M | write/read | 134.60 / 162.70 | | 12M | write/read | 135.69 / 164.20 | Add a test case `mooncake-transfer-engine/tests/ubshmem_transport_test` to test transfer functionality and accuracy.  ## 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. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-29 21:33:46 +00:00