[GH-ISSUE #4990] [RFC]: Fix MemFabric SDMA memory allocation and NPU transport initialization for Mooncake Store #1265

New issue

Open

opened 2026-10-29 21:24:37 +00:00 by bsakaguchi · 1 comment

bsakaguchi commented

2026-10-29 21:24:37 +00:00

Owner

Originally created by @jamesjxliu on GitHub (Oct 27, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4990

Summary

This RFC describes the MemFabric integration fixes in https://github.com/kvcache-ai/Mooncake/pull/4834. The goal is to support distributed Mooncake Store caching and prefill/decode (PD) transfers on Ascend NPUs, including SDMA over HCCS and independent protocol selection for PD and Store.

Problems

When integrating SGLang HiCache with Mooncake Store and the TENT MemFabric transport, three integration gaps can prevent initialization, registration, or data transfer:

Host memory allocation for SDMA. The existing ACL Host allocation path does not necessarily provide the exportable/mappable Host memory required by MemFabric SDMA. Store segments and client transfer buffers need an allocation path compatible with the selected transport.
Initialization on a thread without an ACL device. Transport installation can run on a framework helper thread. Initializing MemFabric there assumes a current device that the thread may not have, or selects a device different from the one used for memory registration.
Overlapping native registrations after alignment. MemFabric aligns NPU memory registrations to 2 MiB pages. Separate logical buffers can therefore overlap after alignment, even when their original address ranges do not. Batching one registration call does not solve overlaps with earlier calls or logical remove/re-add operations.

For example, two adjacent 3 MiB buffers produce these aligned ranges:

Buffer A: [0x200000, 0x500000) -> [0x200000, 0x600000)
Buffer B: [0x500000, 0x800000) -> [0x400000, 0x800000)

Registering both aligned ranges independently can cause the native transport to reject an overlapping memory region.

Proposed changes

Allocate SDMA Store Host memory through MemFabric's public smem_mem_alloc / smem_mem_free APIs. Keep allocation ownership and library lifetime valid until outstanding buffers are released. Leave native memory/page management inside MemFabric.
Keep existing allocation paths for RDMA/URMA, including Host allocation callbacks and hugepage handling. Let the native SDMA allocator manage its page policy rather than passing it through Mooncake's mmap/hugetlb allocator.
Defer native MemFabric initialization until the first memory registration, on a device-bound caller thread. Record that device and bind subsequent registration, submission, and status-query calls to it.
Track aligned native registration coverage for each transport handle and register only uncovered ranges. Preserve successful registrations across partial failures and logical deregistration, consistent with SMEM's native MR lifetime.
Split native copies at both source and destination 2 MiB boundaries so each copy stays within its registered ranges, while preserving the original logical task's completion status and byte count.
Preserve the SDMA kLocalReadWrite registration bypass for initiator-local buffers.

Independent PD and Store protocols

MF_PROTOCOL selects the standalone Transfer Engine / PD protocol. MF_STORE_PROTOCOL overrides it for the Store Transfer Engine and Store allocation path. When the override is unset or empty, Store falls back to MF_PROTOCOL.

For example, with MemFabric enabled in the TENT configuration:

export MC_USE_TENT=1
export MC_TENT_CONF=/path/to/memfabric_tent.json
export MF_PROTOCOL=device_sdma
export MF_STORE_PROTOCOL=device_rdma

Using device_sdma for both variables selects SDMA for both PD and Store. A Store-only application can set MF_STORE_PROTOCOL without setting MF_PROTOCOL. The override does not rewrite the PD environment variable.

Validation

The PR extends existing MemFabric transport tests to cover allocation/free lifetime, deferred initialization, device ownership, protocol overrides, SDMA local buffers, aligned registration reuse, partial-failure retry, and boundary-safe READ/WRITE transfers.

End-to-end validation should include SGLang PD transfers, L3 cache population, clearing local caches, and L3 readback. Check per-rank cache counters and loaded native libraries in addition to HTTP responses. RDMA and SDMA must be validated on hardware supporting the corresponding path.

Originally created by @jamesjxliu on GitHub (Oct 27, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4990 ## Summary This RFC describes the MemFabric integration fixes in https://github.com/kvcache-ai/Mooncake/pull/4834. The goal is to support distributed Mooncake Store caching and prefill/decode (PD) transfers on Ascend NPUs, including SDMA over HCCS and independent protocol selection for PD and Store. ## Problems When integrating SGLang HiCache with Mooncake Store and the TENT MemFabric transport, three integration gaps can prevent initialization, registration, or data transfer: 1. **Host memory allocation for SDMA.** The existing ACL Host allocation path does not necessarily provide the exportable/mappable Host memory required by MemFabric SDMA. Store segments and client transfer buffers need an allocation path compatible with the selected transport. 2. **Initialization on a thread without an ACL device.** Transport installation can run on a framework helper thread. Initializing MemFabric there assumes a current device that the thread may not have, or selects a device different from the one used for memory registration. 3. **Overlapping native registrations after alignment.** MemFabric aligns NPU memory registrations to 2 MiB pages. Separate logical buffers can therefore overlap after alignment, even when their original address ranges do not. Batching one registration call does not solve overlaps with earlier calls or logical remove/re-add operations. For example, two adjacent 3 MiB buffers produce these aligned ranges: ```text Buffer A: [0x200000, 0x500000) -> [0x200000, 0x600000) Buffer B: [0x500000, 0x800000) -> [0x400000, 0x800000) ``` Registering both aligned ranges independently can cause the native transport to reject an overlapping memory region. ## Proposed changes - Allocate SDMA Store Host memory through MemFabric's public `smem_mem_alloc` / `smem_mem_free` APIs. Keep allocation ownership and library lifetime valid until outstanding buffers are released. Leave native memory/page management inside MemFabric. - Keep existing allocation paths for RDMA/URMA, including Host allocation callbacks and hugepage handling. Let the native SDMA allocator manage its page policy rather than passing it through Mooncake's mmap/hugetlb allocator. - Defer native MemFabric initialization until the first memory registration, on a device-bound caller thread. Record that device and bind subsequent registration, submission, and status-query calls to it. - Track aligned native registration coverage for each transport handle and register only uncovered ranges. Preserve successful registrations across partial failures and logical deregistration, consistent with SMEM's native MR lifetime. - Split native copies at both source and destination 2 MiB boundaries so each copy stays within its registered ranges, while preserving the original logical task's completion status and byte count. - Preserve the SDMA `kLocalReadWrite` registration bypass for initiator-local buffers. ## Independent PD and Store protocols `MF_PROTOCOL` selects the standalone Transfer Engine / PD protocol. `MF_STORE_PROTOCOL` overrides it for the Store Transfer Engine and Store allocation path. When the override is unset or empty, Store falls back to `MF_PROTOCOL`. For example, with MemFabric enabled in the TENT configuration: ```bash export MC_USE_TENT=1 export MC_TENT_CONF=/path/to/memfabric_tent.json export MF_PROTOCOL=device_sdma export MF_STORE_PROTOCOL=device_rdma ``` Using `device_sdma` for both variables selects SDMA for both PD and Store. A Store-only application can set `MF_STORE_PROTOCOL` without setting `MF_PROTOCOL`. The override does not rewrite the PD environment variable. ## Validation The PR extends existing MemFabric transport tests to cover allocation/free lifetime, deferred initialization, device ownership, protocol overrides, SDMA local buffers, aligned registration reuse, partial-failure retry, and boundary-safe READ/WRITE transfers. End-to-end validation should include SGLang PD transfers, L3 cache population, clearing local caches, and L3 readback. Check per-rank cache counters and loaded native libraries in addition to HTTP responses. RDMA and SDMA must be validated on hardware supporting the corresponding path.

bsakaguchi added the

RFC

label

2026-10-29 21:24:37 +00:00

bsakaguchi commented

2026-10-29 21:24:38 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Oct 27, 2026):

Thanks for opening this issue, @jamesjxliu!

Field	Value
Issue	#4990
GitHub user ID	`20507750`
Reporter	@jamesjxliu

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Oct 27, 2026): Thanks for opening this issue, @jamesjxliu! | Field | Value | |-------|-------| | **Issue** | #4990 | | **GitHub user ID** | `20507750` | | **Reporter** | @jamesjxliu | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi referenced this issue

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

[PR #1265] [MERGED] [store] add async api #2118