Fix NVLink IPC offset corruption for sub-allocated GPU tensors by ishandhanani · Pull Request #1622 · kvcache-ai/Mooncake

ishandhanani · 2026-03-06T06:26:10Z

Problem

registerLocalMemory() in IntraNodeNvlinkTransport stores the caller-provided pointer as the buffer base address. When frameworks like PyTorch sub-allocate tensors within larger cudaMalloc segments (via caching allocators), this pointer is not the cudaMalloc base. Since cudaIpcGetMemHandle always returns a handle for the entire cudaMalloc allocation, relocateSharedMemoryAddress() computes the wrong offset on the target side -- the sub-allocation offset within the segment is lost.

This corrupts all small tensors (<1MB) that share a cudaMalloc segment. Large tensors (>1MB) are unaffected because PyTorch gives them their own segment.

Fix

Use cuMemGetAddressRange() to resolve the true cudaMalloc base address and allocation size before registration. Register at segment granularity and skip duplicates when multiple tensors share the same segment.

Testing

Reproduced with a two-process test (seed on GPU 0, target on GPU 1) transferring tensors of varying sizes (64B, 16KB, 128KB, 512KB, 4MB). Before fix: 4/6 tensors corrupted. After fix: 6/6 correct.

Also validated end-to-end with SGLang remote instance weight loading (Qwen3-235B-A22B-FP8, TP=4, NVLink transport).

Note

Claude code helped me find this bug while working on sharding weights for transfer in sgl-project/sglang#19983. But I reviewed the fix and have tried it and it works

…nsport When framework caching allocators (PyTorch, etc.) sub-allocate tensors within larger cudaMalloc segments, cudaIpcGetMemHandle returns a handle for the entire segment, not the sub-allocation. The existing code stored the sub-allocation address as the buffer base, causing relocateShared- MemoryAddress to compute an incorrect offset on the target side. Small tensors (<1MB) sharing a segment would read from the wrong location. Fix: use cuMemGetAddressRange() to resolve the true cudaMalloc base address before registration. Register at segment granularity and skip duplicate registrations when multiple tensors share the same segment.

gemini-code-assist · 2026-03-06T06:27:15Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical issue in the NVLink IPC transport mechanism where sub-allocated GPU tensors, particularly those smaller than 1MB, were experiencing data corruption due to incorrect memory registration. The core problem stemmed from registering sub-pointers instead of the underlying cudaMalloc base addresses. The solution involves accurately identifying and registering the true cudaMalloc segments, thereby ensuring proper offset calculation and preventing data integrity issues during inter-process communication.

Highlights

Problem Identification: Identified and addressed a bug where registerLocalMemory in IntraNodeNvlinkTransport incorrectly used sub-allocated tensor pointers instead of the true cudaMalloc base addresses, leading to IPC offset corruption for small tensors (<1MB) when frameworks like PyTorch sub-allocate memory.
Resolution Strategy: Implemented a fix to use cuMemGetAddressRange() to resolve the actual cudaMalloc base address and allocation size. Memory is now registered at the segment granularity, and duplicate registrations for the same cudaMalloc segment are skipped.
Impact and Validation: The fix prevents corruption of sub-allocated tensors, ensuring correct data transfer. Validation included a two-process test showing 6/6 tensors correct after the fix (from 4/6 corrupted) and end-to-end validation with SGLang remote instance weight loading.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

mooncake-transfer-engine/include/transport/intranode_nvlink_transport/intranode_nvlink_transport.h
- Included <unordered_set> for efficient tracking of registered memory blocks.
- Added registered_base_addrs_ as a private member to store uint64_t representations of CUDA memory base pointers.
mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp
- Included <cuda.h> to access cuMemGetAddressRange.
- In registerLocalMemory, implemented cuMemGetAddressRange to determine the actual cudaMalloc base address and allocation size for a given pointer.
- Added a check using registered_base_addrs_ to prevent re-registering the same cudaMalloc segment, returning 0 if already registered.
- Changed cudaIpcGetMemHandle to operate on the cudaMalloc base address instead of the sub-allocated pointer.
- Updated the BufferDesc to store the cudaMalloc base address and its full allocation size.
- In unregisterLocalMemory, added logic to remove the addr from registered_base_addrs_ within a mutex lock.

Activity

No specific activity (comments, reviews, or progress updates) has been recorded for this pull request yet.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request effectively addresses the NVLink IPC offset corruption by resolving and registering the true base address of CUDA memory segments, with the registerLocalMemory implementation appearing sound. However, a critical vulnerability exists in the unregistration logic: unregisterLocalMemory fails to resolve the base address of sub-allocated tensors, leading to resource leaks and preventing correct re-registration. Additionally, a race condition was identified in the unregistration flow, which could lead to an inconsistent state if metadata registration fails. These issues require immediate attention to ensure the robustness and correctness of the transport engine.

- unregisterLocalMemory: resolve base address via cuMemGetAddressRange before erasing from tracking set and metadata (matches register path) - registerLocalMemory: insert into registered_base_addrs_ only after addLocalMemoryBuffer succeeds to avoid inconsistent state on failure

ShangmingCai · 2026-03-06T06:49:12Z

+#include <cuda.h>
 #include "cuda_alike.h"


@alogfans can you check this?

It is double defined.

codecov-commenter · 2026-03-06T07:15:46Z

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

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

alogfans

relocateSharedMemoryAddress should be changed accordingly.

stmatengss · 2026-03-11T13:44:30Z

Intro-node NVLink feature is contributed by @TTThanos . PTAL. Thanks

stmatengss · 2026-03-11T13:45:45Z

relocateSharedMemoryAddress should be changed accordingly.

True. @ishandhanani, could you also modify relocateSharedMemoryAddress? Alternatively, we can implement it in the next PR.

stmatengss · 2026-03-12T02:12:29Z

+    // Resolve the true cudaMalloc base address. Framework caching allocators
+    // (PyTorch, etc.) sub-allocate tensors within larger cudaMalloc segments.
+    // cudaIpcGetMemHandle always returns a handle for the entire segment, so
+    // we must register at segment granularity for correct IPC relocation.
+    CUdeviceptr base_ptr = 0;
+    size_t alloc_size = 0;
+    CUresult cu_err =
+        cuMemGetAddressRange(&base_ptr, &alloc_size, (CUdeviceptr)addr);
+    if (cu_err != CUDA_SUCCESS) {
+        LOG(ERROR) << "IntraNodeNvlinkTransport: cuMemGetAddressRange failed "
+                   << "for addr " << addr << " (error " << cu_err << ")";
+        return -1;
+    }
+
+    // Skip if this cudaMalloc block is already registered
+    if (registered_base_addrs_.count((uint64_t)base_ptr)) {
+        return 0;
+    }
+


Learn a lot from these codes.

…ort/intranode_nvlink_transport.cpp

I will update relocate API.

TTThanos · 2026-03-24T03:26:56Z

Intro-node NVLink feature is contributed by @TTThanos . PTAL. Thanks

Sorry for the late response — I saw this message a bit too late. Thanks for the @alogfans 's contribution! I noticed the issue has already been addressed !

…he-ai#1622) * Fix IPC offset corruption for sub-allocated GPU tensors in NVLink transport When framework caching allocators (PyTorch, etc.) sub-allocate tensors within larger cudaMalloc segments, cudaIpcGetMemHandle returns a handle for the entire segment, not the sub-allocation. The existing code stored the sub-allocation address as the buffer base, causing relocateShared- MemoryAddress to compute an incorrect offset on the target side. Small tensors (<1MB) sharing a segment would read from the wrong location. Fix: use cuMemGetAddressRange() to resolve the true cudaMalloc base address before registration. Register at segment granularity and skip duplicate registrations when multiple tensors share the same segment. * Address review: fix unregister and insert-before-confirm - unregisterLocalMemory: resolve base address via cuMemGetAddressRange before erasing from tracking set and metadata (matches register path) - registerLocalMemory: insert into registered_base_addrs_ only after addLocalMemoryBuffer succeeds to avoid inconsistent state on failure * Update mooncake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp --------- Co-authored-by: Ishan Dhanani <ishan@dhanani.dev> Co-authored-by: Teng Ma <teng-ma@linux.alibaba.com>

ishandhanani requested review from alogfans, chestnut-Q and doujiang24 as code owners March 6, 2026 06:26

github-actions Bot added run-ci Transfer Engine labels Mar 6, 2026

gemini-code-assist Bot reviewed Mar 6, 2026

View reviewed changes

Comment thread ...cake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp Outdated

Comment thread ...cake-transfer-engine/src/transport/intranode_nvlink_transport/intranode_nvlink_transport.cpp Outdated

ShangmingCai reviewed Mar 6, 2026

View reviewed changes

alogfans previously requested changes Mar 9, 2026

View reviewed changes

stmatengss added the high priority label Mar 10, 2026

stmatengss reviewed Mar 12, 2026

View reviewed changes

Update mooncake-transfer-engine/src/transport/intranode_nvlink_transp…

46c02cd

…ort/intranode_nvlink_transport.cpp

stmatengss approved these changes Mar 12, 2026

View reviewed changes

stmatengss requested a review from alogfans March 12, 2026 03:17

alogfans approved these changes Mar 12, 2026

View reviewed changes

alogfans merged commit 68c1da7 into kvcache-ai:main Mar 12, 2026
17 checks passed

shuoerw mentioned this pull request Apr 7, 2026

[Bug][TENT]: SGLang P/D intra-node communication fails due to wrong IPC address #1829

Closed

he-yufeng mentioned this pull request Apr 7, 2026

[TENT] Fix NVLink IPC address for sub-allocated GPU tensors #1831

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix NVLink IPC offset corruption for sub-allocated GPU tensors#1622

Fix NVLink IPC offset corruption for sub-allocated GPU tensors#1622
alogfans merged 3 commits intokvcache-ai:mainfrom
ishandhanani:idhanani/fix-nvlink-ipc-offset

ishandhanani commented Mar 6, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Mar 6, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

ShangmingCai Mar 6, 2026

Uh oh!

stmatengss Mar 11, 2026

Uh oh!

codecov-commenter commented Mar 6, 2026

Uh oh!

alogfans left a comment

Uh oh!

stmatengss commented Mar 11, 2026

Uh oh!

stmatengss commented Mar 11, 2026

Uh oh!

Uh oh!

stmatengss Mar 12, 2026

Uh oh!

Uh oh!

TTThanos commented Mar 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

ishandhanani commented Mar 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

Testing

Note

Uh oh!

gemini-code-assist Bot commented Mar 6, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

ShangmingCai Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

stmatengss Mar 11, 2026

Choose a reason for hiding this comment

Uh oh!

codecov-commenter commented Mar 6, 2026

Codecov Report

Uh oh!

alogfans left a comment

Choose a reason for hiding this comment

Uh oh!

stmatengss commented Mar 11, 2026

Uh oh!

stmatengss commented Mar 11, 2026

Uh oh!

Uh oh!

stmatengss Mar 12, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

TTThanos commented Mar 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

ishandhanani commented Mar 6, 2026 •

edited

Loading