[PR #675] [MERGED] feat(transfer_engine_bench): Add multi-GPU support #1738

New issue

Closed

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

bsakaguchi commented

2026-10-29 21:29:56 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/Mooncake/pull/675
Author: @staryxchen
Created: 7/26/2025
Status: ✅ Merged
Merged: 7/28/2025
Merged by: @alogfans

Base: main ← Head: opt/bench

📝 Commits (4)

c94f335 feat(transfer_engine_bench): add multi-GPU support and improve memory allocation logic
3fac6ad replace thread array with vector
d77bc7f move DRAM log messages to initialization blocks
8c6a218 fix typo

📊 Changes

3 files changed (+82 additions, -31 deletions)

View changed files

📝 mooncake-transfer-engine/example/transfer_engine_bench.cpp (+78 -27)
📝 mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp (+2 -2)
📝 mooncake-transfer-engine/tests/rdma_transport_test.cpp (+2 -2)

📄 Description

Summary

This PR enhances the mooncake transfer engine benchmark to support multi-GPU configurations and improves memory allocation logic for better performance demonstration.

Change

Added support for using all available GPUs: When gpu_id=-1, the benchmark now automatically detects and utilizes all available GPUs
Dynamic buffer allocation: The number of buffers now scales with the number of GPUs when using VRAM
Enhanced logging: Added detailed logging for GPU memory allocation to aid in debugging and monitoring

Usage

Single GPU: --gpu_id=0 (existing behavior)
Multi-GPU: --gpu_id=-1 (new feature - uses all available GPUs)

Testing Result

The changes maintain full backward compatibility with existing single-GPU configurations while enabling multi-GPU scenarios for enhanced performance testing. I have test the performance with the below command:

target:

./transfer_engine_bench --mode=target --metadata_server=redis://${redis_ip}:${redis_port} --local_server_name=${local_ip} --use_vram=true --block_size=$((1*1024*1024)) --threads=16 --auto_discovery --gpu_id=-1

initiator:

./transfer_engine_bench --metadata_server=redis://${redis_ip}:${redis_port} --local_server_name=${local_ip} segment_id=${target_ip} --block_size=$((1*1024*1024)) --operation=write --duration=30 --threads=16 --auto_discovery --use_vram=true --gpu_id=-1

And the result show that When allocating memory from eight H20 GPUs simultaneously (which is the purpose of this PR), the Mooncake Transfer Engine can fully utilize the bandwidth of eight 400Gbps RDMA NIC

_{🔄 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/675 **Author:** [@staryxchen](https://github.com/staryxchen) **Created:** 7/26/2025 **Status:** ✅ Merged **Merged:** 7/28/2025 **Merged by:** [@alogfans](https://github.com/alogfans) **Base:** `main` ← **Head:** `opt/bench` --- ### 📝 Commits (4) - [`c94f335`](https://github.com/kvcache-ai/Mooncake/commit/c94f3353eb5248819a6a337ef67418e62265beaf) feat(transfer_engine_bench): add multi-GPU support and improve memory allocation logic - [`3fac6ad`](https://github.com/kvcache-ai/Mooncake/commit/3fac6ad288076edcd51a65f772b167f075f2996f) replace thread array with vector - [`d77bc7f`](https://github.com/kvcache-ai/Mooncake/commit/d77bc7f28103e5a8aa7a96a3afeabb8845f5f59c) move DRAM log messages to initialization blocks - [`8c6a218`](https://github.com/kvcache-ai/Mooncake/commit/8c6a2185b0f81c0621707509d846ac531631f83d) fix typo ### 📊 Changes **3 files changed** (+82 additions, -31 deletions) <details> <summary>View changed files</summary> 📝 `mooncake-transfer-engine/example/transfer_engine_bench.cpp` (+78 -27) 📝 `mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp` (+2 -2) 📝 `mooncake-transfer-engine/tests/rdma_transport_test.cpp` (+2 -2) </details> ### 📄 Description # Summary This PR enhances the mooncake transfer engine benchmark to support multi-GPU configurations and improves memory allocation logic for better performance demonstration. # Change - Added support for using all available GPUs: When gpu_id=-1, the benchmark now automatically detects and utilizes all available GPUs - Dynamic buffer allocation: The number of buffers now scales with the number of GPUs when using VRAM - Enhanced logging: Added detailed logging for GPU memory allocation to aid in debugging and monitoring # Usage - Single GPU: --gpu_id=0 (existing behavior) - Multi-GPU: --gpu_id=-1 (**new feature - uses all available GPUs**) # Testing Result The changes maintain full backward compatibility with existing single-GPU configurations while enabling multi-GPU scenarios for enhanced performance testing. I have test the performance with the below command: - target: ``` ./transfer_engine_bench --mode=target --metadata_server=redis://${redis_ip}:${redis_port} --local_server_name=${local_ip} --use_vram=true --block_size=$((1*1024*1024)) --threads=16 --auto_discovery --gpu_id=-1 ``` - initiator: ``` ./transfer_engine_bench --metadata_server=redis://${redis_ip}:${redis_port} --local_server_name=${local_ip} segment_id=${target_ip} --block_size=$((1*1024*1024)) --operation=write --duration=30 --threads=16 --auto_discovery --use_vram=true --gpu_id=-1 ``` And the result show that **When allocating memory from eight H20 GPUs simultaneously (which is the purpose of this PR), the Mooncake Transfer Engine can fully utilize the bandwidth of eight 400Gbps RDMA NIC** <img width="672" height="31" alt="Clipboard_Screenshot_1753502699" src="https://github.com/user-attachments/assets/84c0968d-d19e-478b-96fb-061a6c093fa5" /> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-29 21:29:56 +00:00