[PR #1353] [MERGED] [TE/HIP] Add stream and event pools for async transfer operations

No labels

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/Mooncake#2185

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/Mooncake/pull/1353
Author: @amd-arozanov
Created: 1/9/2026
Status: ✅ Merged
Merged: 1/12/2026
Merged by: @alogfans

Base: main ← Head: hip_multi_stream

📝 Commits (1)

7d34677 [TE/HIP] Add stream and event pools for async transfer operations

📊 Changes

6 files changed (+491 additions, -39 deletions)

View changed files

➕ mooncake-transfer-engine/include/transport/hip_transport/event_pool.h (+62 -0)
📝 mooncake-transfer-engine/include/transport/hip_transport/hip_transport.h (+19 -3)
➕ mooncake-transfer-engine/include/transport/hip_transport/stream_pool.h (+54 -0)
➕ mooncake-transfer-engine/src/transport/hip_transport/event_pool.cpp (+113 -0)
📝 mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp (+151 -36)
➕ mooncake-transfer-engine/src/transport/hip_transport/stream_pool.cpp (+92 -0)

📄 Description

Description

(AI assist)

This PR improves hip transport performance by using streams/events and enabling asynchronous data transfers across multiple streams. Each transfer in a batch executes on a separate stream, allowing multiple transfers to run concurrently on the GPU.

Changes:

Add StreamPool class to manage reusable HIP streams per device
Add EventPool class to manage reusable HIP events per device
Refactor HipTransport to use async transfers with stream/event pools
Add support for configurable pool sizes via environment variables:
- MC_HIP_NUM_STREAMS (default: 64)
- MC_HIP_NUM_EVENTS (default: 64)

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?

Ran transfer_engine_bench and transfer_engine_validator on the nodes with AMG GPU. Performance was checked by running transfer_engine_bench with different values of -batch_size parameter.

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.

_{🔄 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/1353 **Author:** [@amd-arozanov](https://github.com/amd-arozanov) **Created:** 1/9/2026 **Status:** ✅ Merged **Merged:** 1/12/2026 **Merged by:** [@alogfans](https://github.com/alogfans) **Base:** `main` ← **Head:** `hip_multi_stream` --- ### 📝 Commits (1) - [`7d34677`](https://github.com/kvcache-ai/Mooncake/commit/7d34677804af06503ce04e0d6c8a629a3d24c431) [TE/HIP] Add stream and event pools for async transfer operations ### 📊 Changes **6 files changed** (+491 additions, -39 deletions) <details> <summary>View changed files</summary> ➕ `mooncake-transfer-engine/include/transport/hip_transport/event_pool.h` (+62 -0) 📝 `mooncake-transfer-engine/include/transport/hip_transport/hip_transport.h` (+19 -3) ➕ `mooncake-transfer-engine/include/transport/hip_transport/stream_pool.h` (+54 -0) ➕ `mooncake-transfer-engine/src/transport/hip_transport/event_pool.cpp` (+113 -0) 📝 `mooncake-transfer-engine/src/transport/hip_transport/hip_transport.cpp` (+151 -36) ➕ `mooncake-transfer-engine/src/transport/hip_transport/stream_pool.cpp` (+92 -0) </details> ### 📄 Description ## Description (AI assist) This PR improves hip transport performance by using streams/events and enabling asynchronous data transfers across multiple streams. Each transfer in a batch executes on a separate stream, allowing multiple transfers to run concurrently on the GPU. ### Changes: - Add StreamPool class to manage reusable HIP streams per device - Add EventPool class to manage reusable HIP events per device - Refactor HipTransport to use async transfers with stream/event pools - Add support for configurable pool sizes via environment variables: * MC_HIP_NUM_STREAMS (default: 64) * MC_HIP_NUM_EVENTS (default: 64) ## 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? Ran transfer_engine_bench and transfer_engine_validator on the nodes with AMG GPU. Performance was checked by running transfer_engine_bench with different values of -batch_size parameter. ## Checklist - [x] I have performed a self-review of my own code. - [x] I have updated the documentation. - [x] 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

closed this issue
added the
pull-request
label

Rows
Columns

[PR #1353] [MERGED] [TE/HIP] Add stream and event pools for async transfer operations #2185