kvcache-ai/Mooncake

Fork 0

mirror of https://github.com/kvcache-ai/Mooncake.git synced 2026-11-05 01:17:55 +00:00

[GH-ISSUE #4852] [RFC]: Add fixed-capacity elastic peer management to the existing Mooncake EP Buffer #1234

New issue

Closed

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

bsakaguchi commented

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

Owner

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

Originally assigned to: @UNIDY2002 on GitHub.

Proposal

Add an opt-in elastic path to the existing mooncake.mooncake_ep_buffer.Buffer. Callers reserve a fixed capacity, connect or disconnect IPC/RDMA peers individually by stable global-rank slot, and use a mask to control which peers participate in communication.

The existing Buffer(group, ...) constructor remains available. The new path manages buffer capacity, dispatch rank range, and the active peer set separately.

Motivation

A fixed peer set can be connected when a Buffer is created. Elastic expansion and replacement also need separate control over:

The maximum number of ranks the buffer can hold.
The global-rank range covered by a dispatch or combine operation.
Which slots are connected, masked, or active.

Treating these as one world size makes sparse membership and replacement difficult to express. Renumbering ranks to close a gap would also change kernel layouts and the caller's global-rank mapping.

1. Slots identify logical ranks

With a capacity of 8, the buffer reserves slots 0 through 7. Each slot index identifies a logical rank. The caller maps ranks to physical GPUs and experts.

capacity          = 8
slot                0  1  2  3  4  5  6  7
initial active      1  1  1  1  0  0  0  0
after expansion     1  1  1  1  1  1  1  1
after removing 2    1  1  0  1  1  1  1  1

After slot 2 is removed, its index remains reserved. A replacement worker can reuse slot 2 with a new IPC handle and RDMA metadata. Other ranks keep their original indices.

Capacity, rank bound, and active count are independent. If only slots 0 and 3 participate, the active count is 2, but a global-rank-indexed layout must still cover through slot 3. Allocation reserves resources for the maximum rank count; a larger capacity requires a new Buffer.

QP row stride remains based on the fixed capacity, while dispatch rank_bound may be smaller than capacity. Sparse active ranks therefore keep stable global-rank slots.

2. Each slot stores its current peer connection

The caller exchanges peer metadata through TCPStore. Mooncake's device transport continues to carry token payloads. Connections use IPC/P2P or available RDMA resources according to the local topology.

disconnect_ranks() clears the connection resources for a slot. A replacement worker supplies new IPC and RDMA metadata when reconnecting.

3. Prepare a connection before activating it

Peer connection has separate prepare and activate steps:

disconnected, masked
        |
        | connect_ranks([peer], activate=False)
        v
staged, still masked; slot is reserved by the connection
        |
        | update_mask_buffer(peer, mask=False)
        v
wait for this process's earlier GPU work, then publish the new mapping
        |
        v
active; eligible for dispatch/combine

connect_ranks(..., activate=False) connects a peer and keeps it masked. The mask controls whether the peer participates; disconnect_ranks() releases its connection resources.

For replacement, the side holding the old connection first calls disconnect_ranks([slot]). Both sides then exchange the new peer metadata through TCPStore and reconnect. The caller serializes connect_ranks() and disconnect_ranks() operations on the same buffer.

Preparation sets up the new peer mapping and QP while preserving current connections. In the 3-rank communication test, preparation overlapped with GPU communication among existing peers.

When update_mask_buffer(..., mask=False) publishes an uncommitted peer, commit_staged_peers() calls cudaDeviceSynchronize() to wait for GPU work already submitted on this device, then publishes the mapping. Disconnect also waits for this device's submitted GPU work before clearing the connection, and synchronizes the communication stream where needed.

4. API example

buffer = Buffer.create_elastic(
    rank=global_rank,
    tcp_store_group=store,
    timeout_ms=30000,
    metadata_timeout_ms=300000,
)
buffer.update_memory_buffers(
    num_ranks=capacity,
    num_experts_per_rank=experts_per_rank,
    num_rdma_bytes=buffer_bytes,
)

buffer.connect_ranks(peer_ranks, activate=False)
buffer.update_mask_buffer(peer_rank, mask=False)

# During replacement, disconnect the old peer before both sides exchange new metadata.
buffer.disconnect_ranks([peer_rank])
buffer.connect_ranks([peer_rank])

create_elastic() creates the Python wrapper. update_memory_buffers() allocates the native buffer and fixes its capacity.

query_mask_buffer() writes the mask to a caller-provided contiguous CUDA int32 tensor. The public convention is 1 = masked; the API converts from the internal active representation. The caller manages worker startup, expert reassignment, and request retries.

Dispatch can optionally receive cumulative_local_expert_recv_stats, which accumulates the number of tokens received by each local expert. Send, receive, and wait paths check the active mask and skip masked peers. The native layer also provides an optional barrier.

5. Metadata and GPU communication timeouts

metadata_timeout_ms defaults to 300,000 ms and controls only TCPStore peer-metadata waits. Native GPU communication uses the separate timeout_ms parameter. The caller can adjust the metadata timeout for its deployment's exchange time.

After a GPU receive wait times out, the local active mask marks the corresponding peer inactive. Communication kernels use the mask to skip that peer.

6. Scope

Enable the elastic peer lifecycle with Buffer.create_elastic(); keep the existing Buffer(group, ...) constructor available.
Allocate a fixed capacity and manage elastic membership with logical-rank slots, an active mask, and per-peer connections.
The validated path uses zero_copy=False, BF16, and standard FP8.
The caller coordinates workers, experts, and request retries.

7. Validation

Typos, Ruff, pre-commit, C++ changed-line clang-format, and git diff --check passed. A fresh native build passed with CUDA 13.0.88 and Python 3.12.15. CPU checks passed: test_review_scope_cpu.py 10/10 and test_launcher_cpu.py 14/14.

All 7 focused plans passed on two H20 nodes with 8 GPUs each. They covered single-rank execution, successive growth, 8-to-16 ranks, clean removal, fault rejoin, and a final sparse layout with active ranks 0 and 3. The tests used BF16, standard FP8, and zero_copy=False; covered event, hook, stats, and out paths; and used native P2P within a node and IBGDA across nodes.

The 3-rank connection contract verified preparation alongside existing GPU work, disconnect/reconnect, 8 barrier rounds, and mask changes after a missing-peer timeout.

Originally created by @niwehao on GitHub (Oct 22, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4852 Originally assigned to: @UNIDY2002 on GitHub. ## Proposal Add an opt-in elastic path to the existing `mooncake.mooncake_ep_buffer.Buffer`. Callers reserve a fixed capacity, connect or disconnect IPC/RDMA peers individually by stable global-rank slot, and use a mask to control which peers participate in communication. The existing `Buffer(group, ...)` constructor remains available. The new path manages buffer capacity, dispatch rank range, and the active peer set separately. ## Motivation A fixed peer set can be connected when a `Buffer` is created. Elastic expansion and replacement also need separate control over: 1. The maximum number of ranks the buffer can hold. 2. The global-rank range covered by a dispatch or combine operation. 3. Which slots are connected, masked, or active. Treating these as one world size makes sparse membership and replacement difficult to express. Renumbering ranks to close a gap would also change kernel layouts and the caller's global-rank mapping. ## 1. Slots identify logical ranks With a capacity of 8, the buffer reserves slots 0 through 7. Each slot index identifies a logical rank. The caller maps ranks to physical GPUs and experts. ```text capacity = 8 slot 0 1 2 3 4 5 6 7 initial active 1 1 1 1 0 0 0 0 after expansion 1 1 1 1 1 1 1 1 after removing 2 1 1 0 1 1 1 1 1 ``` After slot 2 is removed, its index remains reserved. A replacement worker can reuse slot 2 with a new IPC handle and RDMA metadata. Other ranks keep their original indices. Capacity, rank bound, and active count are independent. If only slots 0 and 3 participate, the active count is 2, but a global-rank-indexed layout must still cover through slot 3. Allocation reserves resources for the maximum rank count; a larger capacity requires a new `Buffer`. QP row stride remains based on the fixed capacity, while dispatch `rank_bound` may be smaller than capacity. Sparse active ranks therefore keep stable global-rank slots. ## 2. Each slot stores its current peer connection The caller exchanges peer metadata through TCPStore. Mooncake's device transport continues to carry token payloads. Connections use IPC/P2P or available RDMA resources according to the local topology. `disconnect_ranks()` clears the connection resources for a slot. A replacement worker supplies new IPC and RDMA metadata when reconnecting. ## 3. Prepare a connection before activating it Peer connection has separate prepare and activate steps: ```text disconnected, masked | | connect_ranks([peer], activate=False) v staged, still masked; slot is reserved by the connection | | update_mask_buffer(peer, mask=False) v wait for this process's earlier GPU work, then publish the new mapping | v active; eligible for dispatch/combine ``` `connect_ranks(..., activate=False)` connects a peer and keeps it masked. The mask controls whether the peer participates; `disconnect_ranks()` releases its connection resources. For replacement, the side holding the old connection first calls `disconnect_ranks([slot])`. Both sides then exchange the new peer metadata through TCPStore and reconnect. The caller serializes `connect_ranks()` and `disconnect_ranks()` operations on the same buffer. Preparation sets up the new peer mapping and QP while preserving current connections. In the 3-rank communication test, preparation overlapped with GPU communication among existing peers. When `update_mask_buffer(..., mask=False)` publishes an uncommitted peer, `commit_staged_peers()` calls `cudaDeviceSynchronize()` to wait for GPU work already submitted on this device, then publishes the mapping. Disconnect also waits for this device's submitted GPU work before clearing the connection, and synchronizes the communication stream where needed. ## 4. API example ```python buffer = Buffer.create_elastic( rank=global_rank, tcp_store_group=store, timeout_ms=30000, metadata_timeout_ms=300000, ) buffer.update_memory_buffers( num_ranks=capacity, num_experts_per_rank=experts_per_rank, num_rdma_bytes=buffer_bytes, ) buffer.connect_ranks(peer_ranks, activate=False) buffer.update_mask_buffer(peer_rank, mask=False) # During replacement, disconnect the old peer before both sides exchange new metadata. buffer.disconnect_ranks([peer_rank]) buffer.connect_ranks([peer_rank]) ``` `create_elastic()` creates the Python wrapper. `update_memory_buffers()` allocates the native buffer and fixes its capacity. `query_mask_buffer()` writes the mask to a caller-provided contiguous CUDA `int32` tensor. The public convention is `1 = masked`; the API converts from the internal active representation. The caller manages worker startup, expert reassignment, and request retries. Dispatch can optionally receive `cumulative_local_expert_recv_stats`, which accumulates the number of tokens received by each local expert. Send, receive, and wait paths check the active mask and skip masked peers. The native layer also provides an optional barrier. ## 5. Metadata and GPU communication timeouts `metadata_timeout_ms` defaults to 300,000 ms and controls only TCPStore peer-metadata waits. Native GPU communication uses the separate `timeout_ms` parameter. The caller can adjust the metadata timeout for its deployment's exchange time. After a GPU receive wait times out, the local active mask marks the corresponding peer inactive. Communication kernels use the mask to skip that peer. ## 6. Scope - Enable the elastic peer lifecycle with `Buffer.create_elastic()`; keep the existing `Buffer(group, ...)` constructor available. - Allocate a fixed capacity and manage elastic membership with logical-rank slots, an active mask, and per-peer connections. - The validated path uses `zero_copy=False`, BF16, and standard FP8. - The caller coordinates workers, experts, and request retries. ## 7. Validation Typos, Ruff, pre-commit, C++ changed-line clang-format, and `git diff --check` passed. A fresh native build passed with CUDA 13.0.88 and Python 3.12.15. CPU checks passed: `test_review_scope_cpu.py` 10/10 and `test_launcher_cpu.py` 14/14. All 7 focused plans passed on two H20 nodes with 8 GPUs each. They covered single-rank execution, successive growth, 8-to-16 ranks, clean removal, fault rejoin, and a final sparse layout with active ranks 0 and 3. The tests used BF16, standard FP8, and `zero_copy=False`; covered event, hook, stats, and `out` paths; and used native P2P within a node and IBGDA across nodes. The 3-rank connection contract verified preparation alongside existing GPU work, disconnect/reconnect, 8 barrier rounds, and mask changes after a missing-peer timeout.

bsakaguchi closed this issue

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

bsakaguchi commented

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

Author

Owner

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

Thanks for opening this issue, @niwehao!

Field	Value
Issue	#4852
GitHub user ID	`125028662`
Reporter	@niwehao

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 22, 2026): Thanks for opening this issue, @niwehao! | Field | Value | |-------|-------| | **Issue** | #4852 | | **GitHub user ID** | `125028662` | | **Reporter** | @niwehao | 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:43 +00:00

[PR #1234] [MERGED] [Feature] Support H20 intraNode nvlink by introducing fallback mechanism to leverage cudaIPC #2092