[PR #1179] [MERGED] [TE/Examples] Memory initialization and HIP cleanup fixes

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#2053

bsakaguchi commented

2026-10-29 21:32:25 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/Mooncake/pull/1179
Author: @amd-arozanov
Created: 12/8/2025
Status: ✅ Merged
Merged: 12/11/2025
Merged by: @stmatengss

Base: main ← Head: example_fixes

📝 Commits (5)

413e525 [TE/Examples] Initialize allocated memory
4250743 [TransferEngineValidator] Check data at least once
606ecdd [TE/Examples] Directly call freePinnedLocalMemory for HIP
6002f31 [TE/Examples] Directly call freePinnedLocalMemory for CUDA and MUSA too
1c428a8 [TE/Examples] Remove unnecessary copyrights

📊 Changes

4 files changed (+36 additions, -14 deletions)

View changed files

📝 mooncake-transfer-engine/example/transfer_engine_bench.cpp (+12 -4)
📝 mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp (+11 -4)
📝 mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp (+9 -0)
📝 mooncake-transfer-engine/example/transfer_engine_validator.cpp (+4 -6)

📄 Description

Description

Transfer engine examples changes:

Initialize allocated memory:
- GPU Driver may optimize transfers for non-initialized memory, while on real applications the memory is always initialized. So, need to initialize the memory in the benchmarks too.
Check data at least once in TransferEngineValidator:
- transfer_engine_validator checks data on random iterations. It may happen that it does not check data at all. So, added checking at least on the first iteration.
Directly call freePinnedLocalMemory for HIP:
- hipMemRetainAllocationHandle fails if memory was allocated using hipMalloc, which is possible inside NvlinkTransport::allocatePinnedLocalMemory. So, call just freePinnedLocalMemory, which will handle how to free.
- UPDATE: Directly call freePinnedLocalMemory for CUDA and MUSA too as suggested by gemini

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 nodes with AMG GPU

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/1179 **Author:** [@amd-arozanov](https://github.com/amd-arozanov) **Created:** 12/8/2025 **Status:** ✅ Merged **Merged:** 12/11/2025 **Merged by:** [@stmatengss](https://github.com/stmatengss) **Base:** `main` ← **Head:** `example_fixes` --- ### 📝 Commits (5) - [`413e525`](https://github.com/kvcache-ai/Mooncake/commit/413e525e7200139e85ab0e5905612194be1d6783) [TE/Examples] Initialize allocated memory - [`4250743`](https://github.com/kvcache-ai/Mooncake/commit/42507433020d044dc7850fea7715e68df3fbd009) [TransferEngineValidator] Check data at least once - [`606ecdd`](https://github.com/kvcache-ai/Mooncake/commit/606ecddd6b6de218ab5a137b6945a44fdea435f4) [TE/Examples] Directly call freePinnedLocalMemory for HIP - [`6002f31`](https://github.com/kvcache-ai/Mooncake/commit/6002f317e83a6c6437b7751d57009a47e17fa58a) [TE/Examples] Directly call freePinnedLocalMemory for CUDA and MUSA too - [`1c428a8`](https://github.com/kvcache-ai/Mooncake/commit/1c428a88d814d979cbf7c7b43ef13e713fe6960c) [TE/Examples] Remove unnecessary copyrights ### 📊 Changes **4 files changed** (+36 additions, -14 deletions) <details> <summary>View changed files</summary> 📝 `mooncake-transfer-engine/example/transfer_engine_bench.cpp` (+12 -4) 📝 `mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp` (+11 -4) 📝 `mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp` (+9 -0) 📝 `mooncake-transfer-engine/example/transfer_engine_validator.cpp` (+4 -6) </details> ### 📄 Description ## Description Transfer engine examples changes: 1. Initialize allocated memory: - GPU Driver may optimize transfers for non-initialized memory, while on real applications the memory is always initialized. So, need to initialize the memory in the benchmarks too. 2. Check data at least once in TransferEngineValidator: - transfer_engine_validator checks data on random iterations. It may happen that it does not check data at all. So, added checking at least on the first iteration. 3. Directly call freePinnedLocalMemory for HIP: - hipMemRetainAllocationHandle fails if memory was allocated using hipMalloc, which is possible inside NvlinkTransport::allocatePinnedLocalMemory. So, call just freePinnedLocalMemory, which will handle how to free. - **UPDATE**: Directly call freePinnedLocalMemory for CUDA and MUSA too as suggested by gemini ## Type of Change * Types - [x] Bug fix - [ ] 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 nodes with AMG GPU ## 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 #1179] [MERGED] [TE/Examples] Memory initialization and HIP cleanup fixes #2053