Skip to content

[TransferEngine] Share one host KV segment across a TP group - #3285

Open
ascend-direct-dev wants to merge 1 commit into
kvcache-ai:mainfrom
ascend-direct-dev:feat/shared-host-segment
Open

[TransferEngine] Share one host KV segment across a TP group#3285
ascend-direct-dev wants to merge 1 commit into
kvcache-ai:mainfrom
ascend-direct-dev:feat/shared-host-segment

Conversation

@ascend-direct-dev

@ascend-direct-dev ascend-direct-dev commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

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:

  1. SharedSegment::Create reserves address space (owner also allocates and exports) and returns a fixed-size blob.
  2. The caller all_gathers the blob over its own group.
  3. 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 size rather than size * world_size.

Backends. Compiled in per platform and both optional:

Platform Mechanism Handle
Ascend ACL VMM via adxl::AdxlEngine aclrtMemFabricHandle
NVIDIA CUDA VMM (cuMemCreate/cuMemMap) CUmemFabricHandle, CUDA 12.4+

SharedSegment::Supported() reports whether the running system can actually share memory. Builds without either backend return kNotImplemented rather than failing to link.

Python API

from mooncake.shared_segment import create_shared_segment

seg = create_shared_segment(
    "vllm_sparse_kv",
    blocks={
        "k": dict(count=num_layers, shape=k_shape, dtype=torch.bfloat16),
        "v": dict(count=num_layers, shape=v_shape, dtype=torch.bfloat16),
    },
    world_size=tp_size,
    rank_id=tp_rank,
    tp_group=tp_group,
)
k_caches_cpu = seg.tensors("k")  # zero-copy views onto the shared pages

Test plan

  • shared_segment_test: fingerprint, blob encode/decode, and the full two-phase protocol against a fake backend.
  • Ascend build and pybind smoke test on a CANN toolchain.
  • CUDA backend compiles clean under -Wall -Wextra against the CUDA driver headers.
  • Multi-rank run on real Ascend hardware, paired with HiXL https://gitcode.com/cann/hixl/merge_requests/851

@ascend-direct-dev

Copy link
Copy Markdown
Collaborator Author

Comment thread mooncake-transfer-engine/src/shared_segment/shared_segment_internal.h Outdated
@alogfans

Copy link
Copy Markdown
Collaborator

Feel free to mark it ready for review after performing coding style changes.

@ascend-direct-dev
ascend-direct-dev marked this pull request as ready for review August 11, 2026 07:49
@ascend-direct-dev
ascend-direct-dev force-pushed the feat/shared-host-segment branch 3 times, most recently from a51e7e9 to ec87b07 Compare August 11, 2026 11:59
Comment on lines 206 to 208
/**
* @brief 释放MallocMem申请的内存内存
* @param [in] ptr 释放的虚拟内存ptr

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Translate as well.

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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Once a rank raises an error, the whole cluster cannot continue or shutdown gracefully.

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should precesion (bf16/fp16) and/or layout be considered?

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)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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
ascend-direct-dev force-pushed the feat/shared-host-segment branch from ec87b07 to 5047594 Compare August 12, 2026 03:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants