kvcache-ai / kvcache-ai/Mooncake
[RFC]: EGM-backed DRAM Store tier over Multi-Node NVLink
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
### Changes proposed
## Summary
This RFC proposes an **EGM-backed DRAM tier for Mooncake Store** within one
Multi-Node NVLink (MNNVL) scale-up domain.
Extended GPU Memory
([EGM](https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/extended-gpu-memory.html))
allows GPUs to access CPU-attached memory through the NVLink/NVSwitch fabric.
This proposal uses the CUDA VMM path for EGM: Store providers allocate CPU
DRAM with `CU_MEM_LOCATION_TYPE_HOST_NUMA` and export the allocations with
`CU_MEM_HANDLE_TYPE_FABRIC`.
The allocations are mounted as existing `nvlink` Store segments. Consumers
reuse the current Fabric import and mapping path for transfers between GPU HBM
and remote EGM DRAM:
```text
Put: GPU HBM -> CUDA copy over NVLink/NVSwitch -> remote EGM DRAM
Get: remote EGM DRAM -> CUDA copy over NVLink/NVSwitch -> GPU HBM
```
No application-visible CPU staging buffer is required.
## Motivation
GB200/NVL72-class systems provide CPU-attached DRAM that GPUs can access
through Extended GPU Memory (EGM). Mooncake already supports Fabric-handle
publication, remote mapping, and CUDA copies in its `nvlink` transport, but
Store does not currently have a focused way to allocate and publish
`HOST_NUMA` DRAM as Store capacity.
Two workload directions motivate this proposal.
### Higher-bandwidth access to Store DRAM
A Store deployment can access remote DRAM through an external RDMA path.
On GB200/NVL72 systems, EGM provides another path in which GPU traffic to
CPU-attached DRAM is routed through the NVLink/NVSwitch scale-up fabric.
The NVLink fabric offers a substantially higher aggregate interconnect
bandwidth ceiling than a path constrained by individual RNICs. Reusing that
fabric for HBM-to-DRAM transfers may therefore provide higher effective
bandwidth than RDMA-based DRAM access, particularly when capacity is
distributed across multiple GPU-local host NUMA nodes.
This is a performance motivation rather than a guaranteed speedup. Effective
bandwidth depends on host-memory locality, memory-controller bandwidth, CUDA
copy behavior, and contention with other NVLink traffic. RDMA and EGM
performance should be measured on the final implementation commits in the same
environment.
### Future sparse-attention runtimes
Sparse-attention architectures such as
[DeepSeek Sparse Attention (DSA)](https://github.com/deepseek-ai/DeepSeek-V3.2-Exp)
select a subset of historical tokens for each attention operation. A future
serving runtime could keep less frequently accessed KV data in EGM DRAM, fetch
or access only the selected KV entries, and reserve more HBM capacity for
active KV data and model execution.
CUDA EGM mappings can be made accessible to GPU threads and can therefore
support GPU load/store access to mapped system memory. In principle, this
allows a sparse-attention kernel or runtime to access selected KV entries over
the MNNVL fabric without first restoring the complete KV cache to HBM.
This RFC does not implement that model-execution path. It does not expose a
remote-pointer API, define sparse-KV placement, or integrate EGM load/store
access with DSA, vLLM, or attention kernels. The proposed Store data path
remains ordinary Put/Get operations implemented with CUDA copies.
The immediate proposal is limited to exposing provider-owned EGM DRAM as
existing `nvlink` Store segments, while preserving the current Consumer
transfer interface.
## Goals
- Allocate provider Store capacity from CPU DRAM using CUDA VMM `HOST_NUMA`.
- Export the allocations with `CU_MEM_HANDLE_TYPE_FABRIC`.
- Support HBM-to-EGM Put and EGM-to-HBM Get within one MNNVL domain.
- Use GPU-local host NUMA nodes by default.
- Reuse the existing `nvlink` protocol, Fabric-handle metadata, Consumer
import/map/cache path, and Store segment model.
- Keep the feature disabled by default.
- Ensure setup failure and teardown preserve enough ownership for safe cleanup
retry.
## Non-goals
This proposal does not include:
- unified HBM/DRAM allocation, tier placement, promotion/demotion, eviction
policy, or NUMA-aware object placement in Master;
- sparse-KV placement, DSA/vLLM model-path integration, or a public GPU
load/store interface for remote EGM mappings;
- a new Store protocol, transport protocol, or segment wire schema;
- cross-MNNVL-domain access, RDMA registration, or RDMA fallback;
- vLLM model-path or Transfer Engine NEXT (TENT) integration;
- online pool expansion, generic lifecycle refactoring, new metrics, health
codes, or a standalone hardware framework.
These areas should be handled by separate RFCs or PRs if they are pursued.
## Design
### Provider and Consumer roles
The Provider owns the EGM allocations and publishes them through the existing
Store and Transfer Engine metadata paths:
```text
Provider Store
├─ allocates HOST_NUMA EGM DRAM
├─ registers each exact range with NvlinkTransport
│ └─ publishes the existing BufferDesc Fabric handle
└─ mounts each range as an existing nvlink Store segment
└─ publishes the segment through Master
```
The Consumer does not allocate EGM capacity and does not enable the EGM Store
configuration:
```text
Consumer Store
├─ resolves the object and segment through Master
└─ imports and maps the Fabric handle through NvlinkTransport
└─ performs CUDA copies between local HBM and mapped remote EGM
```
The Provider process is involved in allocation and publication, but it is not
an application-level staging hop in the transfer data path.
Consumers continue to use ordinary Put/Get operations and do not need to know
whether a remote `nvlink` segment is backed by HBM or host DRAM.
### NUMA selection and capacity planning
`egm_numa_nodes=auto` is the default. It queries
`CU_DEVICE_ATTRIBUTE_HOST_NUMA_ID` for every visible CUDA GPU, then validates,
sorts, and deduplicates the returned host NUMA IDs. Auto discovery does not
depend on GPU PCI sysfs paths or RDMA NIC topology.
An explicit comma-separated NUMA list acts as an operator placement override.
Each selected node must support the CUDA `HOST_NUMA` allocation path; setup
fails if its granularity cannot be queried or its allocation cannot be created.
For the selected nodes, Store:
1. queries the CUDA allocation granularity for every node;
2. computes a common alignment with the Cachelib slab alignment;
3. rounds total `global_segment_size` down to that alignment;
4. distributes aligned capacity units as evenly as possible across the sorted
nodes;
5. rounds the maximum chunk size down to an aligned value no larger than
`max_mr_size`;
6. splits each node's capacity into aligned chunks.
Each chunk corresponds to one CUDA VMM allocation, one Transfer Engine
registered range, and one existing `nvlink` Store segment. An allocation never
crosses a NUMA-node boundary.
### Setup and cleanup
The Store client owns the EGM pool for its lifetime.
Setup first validates the configuration and calculates the complete capacity
plan. It then creates all allocations before registering and mounting the
segments. Setup succeeds only after every segment is published.
Failure and normal close unwind completed work in reverse order:
1. unmount Store segments;
2. unregister their Transfer Engine ranges;
3. release the CUDA VMM allocations.
If a cleanup step fails, ownership of the remaining resource is retained so a
later `close()` can retry. Cleanup that already succeeded is not repeated.
The Python wrapper destroys the native client only after teardown completes
successfully.
### Transfer Engine prerequisite
Transfer Engine needs a focused provider-side primitive for Mooncake-owned
`HOST_NUMA` allocations:
- discover GPU-local host NUMA IDs and query allocation granularity;
- create and map Fabric-capable `HOST_NUMA` allocations;
- register a Mooncake-owned exact range using the bounds known at allocation
time;
- retain the CUDA allocation handle while the range is registered;
- unregister the range before releasing the allocation.
Registration publishes the existing `BufferDesc` Fabric-handle representation.
The current Consumer Fabric import/map/cache path remains unchanged, as do
non-EGM allocation, registration, and transfer-completion paths.
## Configuration
The proposal adds two Python `ConfigDict` keys:
```text
enable_egm_store_pool = false | true
egm_numa_nodes = auto |
```
The EGM Store tier is enabled only when all of the following are true:
```text
enable_egm_store_pool = true
protocol = nvlink
global_segment_size > 0
local_buffer_size = 0
```
`global_segment_size` is the total EGM capacity across the selected nodes.
Consumers do not enable `enable_egm_store_pool`.
When the feature is disabled, existing Store configuration and allocation
behavior remain unchanged.
## Validation
### Automated validation
Focused tests should cover:
- default-off behavior and invalid configuration combinations;
- automatic and explicit NUMA selection;
- common alignment, total-capacity rounding, node distribution, and
`max_mr_size` chunking;
- allocation, registration, and mount failures;
- reverse-order rollback and teardown;
- partial cleanup followed by retry;
- unchanged non-EGM behavior in CUDA-enabled and CUDA-disabled builds.
The implementation must also pass formatting, CUDA-enabled, CUDA-disabled, and
the existing Store and Transfer Engine CI matrix without broadening unrelated
unit-test jobs.
### GB200/NVL72 validation
The final implementation commits must be validated on GB200/NVL72-class
hardware:
- when automatic discovery returns multiple GPU-local NUMA nodes, capacity is
distributed across those nodes;
- HBM-to-remote-EGM-to-HBM data is byte-correct;
- first-time Fabric import and reuse of the cached mapping both work;
- multiple GPUs can transfer concurrently;
- successful teardown leaves no published EGM capacity in Master or Transfer
Engine metadata;
- disabling the feature preserves existing Store and NVLink behavior.
Performance results should be reported for the final validated commits, but
this RFC does not define a fixed bandwidth threshold. Results from older
development branches are historical evidence rather than acceptance evidence.
## Implementation roadmap
1. [#2966](https://github.com/kvcache-ai/Mooncake/pull/2966) provides the
minimal, default-off Transfer Engine `HOST_NUMA` prerequisite.
2. After #2966 merges, the Store-only implementation will be rebuilt from the
then-current `main` and published as a separate PR. This keeps the Store
review limited to its own incremental diff.
3. Automated validation and GB200/NVL72 validation will be rerun on the final
commits before the corresponding PRs are marked ready.
4. vLLM integration, RDMA fallback, observability, generic lifecycle changes,
Provider-restart mapping identity
([#2832](https://github.com/kvcache-ai/Mooncake/issues/2832)), and TP/EP
interference evaluation remain independent follow-up work.
The closed #2934/#2935 stack is historical and is not the implementation under
review.
### Before submitting a new issue...
- [x] Make sure you already searched for relevant issues and read the
documentation.
Contributor guide
Research direction
Start by reading prerequisite PR #2966 and the existing nvlink Fabric-handle publication and import/map paths described in the RFC. Then trace Store setup, segment mounting, and teardown around the proposed HOST_NUMA configuration. Done means focused tests cover planning, rollback, retryable cleanup, and unchanged non-EGM behavior, followed by GB200/NVL72 validation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- backend, distributed-systems, infrastructure
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100