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

[ascend-direct-dev]

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

[codecov-commenter]

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

Codecov Report

❌ Patch coverage is 89.93289% with 45 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...-integration/transfer_engine/shared_segment_py.cpp	0.00%	34 Missing ⚠️
...cake-transfer-engine/tests/shared_segment_test.cpp	97.97%	5 Missing ⚠️
...nsfer-engine/src/shared_segment/shared_segment.cpp	98.06%	3 Missing ⚠️
...ooncake-transfer-engine/src/common/base/status.cpp	0.00%	2 Missing ⚠️
...integration/transfer_engine/transfer_engine_py.cpp	0.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[ascend-direct-dev]

Superseded by #3285 (same branch tip 221e6ec9; this PR could not be reopened after a force-push).