[TransferEngine] Share one host KV segment across a TP group - #3285
Open
ascend-direct-dev wants to merge 1 commit into
Open
[TransferEngine] Share one host KV segment across a TP group#3285ascend-direct-dev wants to merge 1 commit into
ascend-direct-dev wants to merge 1 commit into
Conversation
ascend-direct-dev
requested review from
ShangmingCai,
alogfans,
chestnut-Q,
doujiang24 and
staryxchen
as code owners
August 4, 2026 12:39
4 tasks
ascend-direct-dev
force-pushed
the
feat/shared-host-segment
branch
from
August 4, 2026 12:40
72f02d2 to
221e6ec
Compare
ascend-direct-dev
marked this pull request as draft
August 4, 2026 12:46
ascend-direct-dev
force-pushed
the
feat/shared-host-segment
branch
from
August 4, 2026 14:14
221e6ec to
8bcb8a9
Compare
ascend-direct-dev
force-pushed
the
feat/shared-host-segment
branch
4 times, most recently
from
August 7, 2026 07:55
fe5909c to
1553c6d
Compare
ascend-direct-dev
force-pushed
the
feat/shared-host-segment
branch
3 times, most recently
from
August 10, 2026 06:51
7d7d9c6 to
624f202
Compare
Collaborator
Author
ascend-direct-dev
force-pushed
the
feat/shared-host-segment
branch
from
August 11, 2026 02:05
51eff91 to
d553e85
Compare
alogfans
reviewed
Aug 11, 2026
Collaborator
|
Feel free to mark it ready for review after performing coding style changes. |
ascend-direct-dev
marked this pull request as ready for review
August 11, 2026 07:49
ascend-direct-dev
force-pushed
the
feat/shared-host-segment
branch
3 times, most recently
from
August 11, 2026 11:59
a51e7e9 to
ec87b07
Compare
alogfans
reviewed
Aug 12, 2026
Comment on lines
206
to
208
| /** | ||
| * @brief 释放MallocMem申请的内存内存 | ||
| * @param [in] ptr 释放的虚拟内存ptr |
alogfans
reviewed
Aug 12, 2026
Comment on lines
+115
to
+127
| except RuntimeError as exc: | ||
| raise SharedSegmentError(str(exc)) from exc | ||
|
|
||
| try: | ||
| if world_size == 1: | ||
| blobs = [blob] | ||
| elif tp_group is None: | ||
| raise SharedSegmentError("tp_group is required when world_size > 1") | ||
| else: | ||
| blobs = _all_gather_blob(blob, world_size, tp_group) | ||
| segment.complete(blobs) | ||
| except RuntimeError as exc: | ||
| raise SharedSegmentError(str(exc)) from exc |
Collaborator
There was a problem hiding this comment.
Once a rank raises an error, the whole cluster cannot continue or shutdown gracefully.
alogfans
reviewed
Aug 12, 2026
Comment on lines
+93
to
+98
| names = sorted(blocks) | ||
| specs = { | ||
| block_name: _parse_block(block_name, blocks[block_name]) | ||
| for block_name in names | ||
| } | ||
| offsets, stride, total = _build_layout(specs, names) |
Collaborator
There was a problem hiding this comment.
Should precesion (bf16/fp16) and/or layout be considered?
alogfans
reviewed
Aug 12, 2026
Comment on lines
+14
to
+18
| # Private headers live next to the .cpp sources. CUDA builds compile from the | ||
| # source tree so quoted includes work; USE_HIP hipifies into the build tree and | ||
| # needs this path explicitly. | ||
| target_include_directories(transfer_engine | ||
| PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/shared_segment) |
Collaborator
There was a problem hiding this comment.
This requires CUDA 12.4 or above.
Create a single host-side shared segment for a TP group so ranks can map the same H memory. mmap defaults on; host_register is optional and gates HostRegister for ROCE D2rH via device VA. Co-authored-by: Cursor <cursoragent@cursor.com>
ascend-direct-dev
force-pushed
the
feat/shared-host-segment
branch
from
August 12, 2026 03:44
ec87b07 to
5047594
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
MLA keeps the same KV cache in every rank of a tensor-parallel group. When those ranks offload KV to host memory during decode, each one allocates its own copy, so host usage grows with the TP size while every copy holds identical bytes.
This PR adds a shared segment: the owner rank allocates the pages once through the platform's virtual memory manager, and every other rank maps the very same physical memory. Host usage for offloaded MLA KV drops to 1/TP.
Design
C++ shares one contiguous host span. Tensor layout (which offset holds which layer's K/V) is computed in Python and folded into the segment name so peers that disagree still fail the fingerprint check.
Two-phase creation. The handle exchange belongs to the caller's process group:
SharedSegment::Createreserves address space (owner also allocates and exports) and returns a fixed-size blob.all_gathers the blob over its own group.segment->Complete(blobs)maps the owner's pages and verifies that every rank declared the same thing.Doing it this way avoids calling back into Python collectives from C++ under the GIL. Non-owner blobs are the same length as the owner's, so a plain fixed-size all-gather carries them.
Addresses are not forced to match. Ranks agree on the byte layout, not on virtual addresses, so each rank reserves
sizerather thansize * world_size.Backends. Compiled in per platform and both optional:
adxl::AdxlEngineaclrtMemFabricHandlecuMemCreate/cuMemMap)CUmemFabricHandle, CUDA 12.4+SharedSegment::Supported()reports whether the running system can actually share memory. Builds without either backend returnkNotImplementedrather than failing to link.Python API
Test plan
shared_segment_test: fingerprint, blob encode/decode, and the full two-phase protocol against a fake backend.-Wall -Wextraagainst the CUDA driver headers.