kvcache-ai/Mooncake

Fork 0

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

[PR #1389] [MERGED] [Store] Unified parallel tensor IO #2213

New issue

Closed

opened 2026-10-29 21:33:42 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:33:42 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/Mooncake/pull/1389
Author: @zxpdemonio
Created: 1/16/2026
Status: ✅ Merged
Merged: 4/27/2026
Merged by: @stmatengss

Base: main ← Head: cruz/flex_put_get

📝 Commits (10+)

2a674c2 [Store] add tensor object metadata and TP upsert APIs
d46f11f [Store] align unified tensor parallelism APIs
6c9d006 [Store] route tensor into reads through ranged gets
7efb1d1 [Store] fix TP ragged reconstruction edge cases
bc7f0d9 [Store] add exact-match multi-axis shard IO support
52e7f67 [Store] optimize raw-buffer shard writes and reconstruction reads
30101c6 [Store] add batch writer-partition unified APIs
b6aaa25 [Store] format writer-shard updates
a0c1dea [Store] tighten tensor write wrapper layering
a36a9a0 [Store] restore legacy tensor wrapper boundaries

📊 Changes

11 files changed (+7639 additions, -148 deletions)

View changed files

➕ docs/source/design/unified-parallel-tensor-io.md (+405 -0)
📝 docs/source/python-api-reference/mooncake-store.md (+228 -2)
📝 mooncake-integration/integration_utils.h (+199 -2)
📝 mooncake-integration/store/store_py.cpp (+357 -137)
➕ mooncake-integration/store/store_py_internal.h (+1483 -0)
➕ mooncake-integration/store/store_py_parallel_read.h (+1312 -0)
➕ mooncake-integration/store/store_py_parallel_write.h (+1443 -0)
📝 mooncake-store/include/real_client.h (+9 -0)
📝 mooncake-store/src/real_client.cpp (+22 -0)
📝 scripts/test_tensor_api.py (+1578 -3)
📝 scripts/test_upsert_api.py (+603 -4)

📄 Description

Motivation

This work is primarily driven by weight synchronization optimization for RL workloads. In RL systems, different components often consume the same model weights under different parallel layouts: some hold the full tensor, some only need a TP shard, and some run with mixed layouts such as DP+TP, EP+TP, or PP+TP. A storage path that only supports full tensors or plain TP shards causes unnecessary write amplification, unnecessary read amplification, and limited layout expressiveness. The goal of this update is to make the unified tensor IO layer suitable for RL-style weight sync by supporting richer layout identities while also optimizing both shard-oriented writes and shard/full reads.

What is supported now

The unified tensor IO path now supports not only heterogeneous TP, but also DP, EP, PP, and their combinations, including exact-match multi-axis shard identities such as DP+TP, PP+TP, and EP+TP. In other words, the system is no longer limited to “TP-only sharded tensors”; it can now persist and resolve multi-axis layout metadata through a single unified API surface.

Write-path optimization

On the write side, the key optimization for RL weight sync is to avoid having every writer store the full tensor when only a local shard needs to be persisted. This is implemented through writer_partition for single-key APIs and writer_partitions for batch APIs. In practice, this allows each writer to hold a full tensor locally but write only the shard it owns, directly reducing write amplification and network/storage overhead during synchronization. The _from paths were also optimized to use raw-buffer shard planning instead of decoding tensor-object buffers into Python tensors and then reshaping or re-sharding them, which removes unnecessary decode, tensor materialization, and extra memcpy/clone work from the hot path.

Read-path optimization

On the read side, the system now supports exact-match shard reads for multi-axis layouts, so consumers can fetch only the shard they need instead of reading a full tensor and slicing locally. This improves efficiency for TP-only consumers as well as mixed DP/TP, PP/TP, and EP/TP layouts. For cases that need the full tensor view, reconstruction is optimized by reusing the existing get_into_ranges(...) ranged-read substrate and avoiding an extra intermediate tensor conversion. As a result, both full reconstruction and _into reconstruction remain efficient while still working on top of the unified metadata model.

RL semantics

For RL workloads, an important distinction is that DP metadata and writer-side partial persistence are not the same thing. If the requirement is “each DP rank has the full tensor, but only writes its owned slice for synchronization efficiency,” that is supported through writer_partition / writer_partitions. If the requirement is “a DP parallelism spec alone should imply automatic partial persistence of a full tensor,” that is not the semantics of the DP metadata path. In the unified API, DP/EP/PP/TP describe layout identity, while writer_partition(s) provides the writer-side shard-only persistence optimization needed for RL weight sync.

Additional improvements

This update also closes an important API symmetry gap by adding batch support for writer_partitions across the unified put/upsert APIs, including _from variants. Batch equests now explicitly reject ambiguous mixed inputs such as parallelisms + writer_partitions. In addition, multi-axis shard metadata matching was fixed by canonicalizing the requested parallelism before comparison, which resolved incorrect mismatch behavior in cases such as PP+TP where axis order differed syntactically but not semantically.

Validation

The changes were validated with focused local tests covering batch writer-partition full reconstruction, batch writer-partition _into round trips, batch writer-partition upsert and upsert-from paths, multi-axis batch shard reads, multi-axis batch _into reads, and explicit invalid mixed-routing cases. Formatting was also rechecked and fixed with the repository formatter.

Support Matrix

Scenario	Supported	Read support	Write optimization	Notes
TP only	Yes	Shard read, full reconstruction, `_into`	Raw-buffer shard write fast path	Mature path
Heterogeneous TP	Yes	Shard read, full reconstruction, batch `_into`	Same as TP path	Supported through unified API
DP metadata layout	Yes	Exact-match shard identity support	Metadata/layout only	Does not by itself mean “write only one slice of a full tensor”
PP metadata layout	Yes	Exact-match shard identity support	Metadata/layout only	Supports stage-aware layout identity
EP metadata layout	Yes	Exact-match shard identity support	Metadata/layout only	Supports expert-aware layout identity
DP + TP	Yes	Exact-match shard read, batch read, `_into`, full reconstruction where applicable	TP shard persistence + unified metadata	Good fit for RL mixed layouts
PP + TP	Yes	Exact-match shard read, batch read, `_into`	TP shard persistence + unified metadata	Axis-order matching issue was fixed
EP + TP	Yes	Exact-match shard read	TP shard persistence + unified metadata	Supports expert-aware shard identity
Writer-side partial persistence from full tensor	Yes	Full reconstruction and shard access	`writer_partition` / `writer_partitions`	Main optimization for RL weight sync
Batch writer-side partial persistence	Yes	Full reconstruction and shard access	Batch `writer_partitions`	Newly added in this update
Mixed `parallelisms` + `writer_partitions` in one batch request	Rejected	N/A	N/A	Explicitly invalid to avoid ambiguous semantics

High-level architecture

flowchart TB                                                                                                                                                                       
    Caller["Python caller"]                                                                                                                                                        
                                                                                                                                                                                   
    API["Unified with_parallelism API<br/>put / get / upsert"]                                                                                                                     
                                                                                                                                                                                   
    Request["Request objects<br/>TensorParallelism / ReadTarget / writer_partition"]                                                                                               
                                                                                                                                                                                   
    Router["Internal route resolver"]                                                                                                                                              
                                                                                                                                                                                   
    Full["Full tensor path"]                                                                                                                                                       
                                                                                                                                                                                   
    TPCompat["TP compatibility path"]                                                                                                                                              
                                                                                                                                                                                   
    MultiAxis["Multi-axis parallelism path"]                                                                                                                                       
                                                                                                                                                                                   
    WriterPartition["writer_partition path"]                                                                                                                                       
                                                                                                                                                                                   
    Metadata["Tensor metadata<br/>global/local shape, dtype, layout axes"]                                                                                                         
                                                                                                                                                                                   
    Manifest["Shard manifest<br/>global shape, dtype, split_dim, shard_count"]                                                                                                     
                                                                                                                                                                                   
    Store["Mooncake object store"]                                                                                                                                                 
                                                                                                                                                                                   
    ReadPlanner["Read planner"]                                                                                                                                                    
                                                                                                                                                                                   
    RangeRead["Range-based reconstruction<br/>get_into_ranges(...)"]                                                                                                               
                                                                                                                                                                                   
    Output["torch.Tensor / user buffer"]                                                                                                                                           
                                                                                                                                                                                   
    Caller --> API                                                                                                                                                                 
    API --> Request                                                                                                                                                                
    Request --> Router                                                                                                                                                             
                                                                                                                                                                                   
    Router --> Full                                                                                                                                                                
    Router --> TPCompat                                                                                                                                                            
    Router --> MultiAxis                                                                                                                                                           
    Router --> WriterPartition                                                                                                                                                     
                                                                                                                                                                                   
    Full --> Metadata                                                                                                                                                              
    TPCompat --> Metadata                                                                                                                                                          
    MultiAxis --> Metadata                                                                                                                                                         
    WriterPartition --> Metadata                                                                                                                                                   
                                                                                                                                                                                   
    Metadata --> Store                                                                                                                                                             
    MultiAxis --> Manifest                                                                                                                                                         
    WriterPartition --> Manifest                                                                                                                                                   
    Manifest --> Store                                                                                                                                                             
                                                                                                                                                                                   
    Store --> ReadPlanner                                                                                                                                                          
    Manifest --> ReadPlanner                                                                                                                                                       
                                                                                                                                                                                   
    ReadPlanner --> Output                                                                                                                                                         
    ReadPlanner --> RangeRead                                                                                                                                                      
    RangeRead --> Output

Public API shape

This PR centers tensor IO around the following API families:

put_tensor_with_parallelism(...)                                                                                                                                                   
batch_put_tensor_with_parallelism(...)                                                                                                                                             
                                                                                                                                                                                   
get_tensor_with_parallelism(...)                                                                                                                                               
batch_get_tensor_with_parallelism(...)                                                                                                                                             
                                                                                                                                                                                   
get_tensor_with_parallelism_into(...)                                                                                                                                              
batch_get_tensor_with_parallelism_into(...)                                                                                                                                        
                                                                                                                                                                               
upsert_tensor_with_parallelism(...)                                                                                                                                                
batch_upsert_tensor_with_parallelism(...)

The core public request objects are:

TensorParallelism                
ParallelAxis                     
ReadTarget

TensorParallelism describes where a tensor object belongs in the parallel layout.

ReadTarget describes what the caller wants to materialize on read:

stored object
target shard
reconstructed full tensor

Multi-axis write flow

For a request such as:

put_tensor_with_parallelism(                                                                                                                                                       
    key,                                                                                                                                                                           
    full_tensor,                                                                                                                                                                   
    TensorParallelism([                                                                                                                                                            
        ParallelAxis(kind="dp", rank=0, size=2),                                                                                                                               
        ParallelAxis(kind="tp", rank=1, size=4, split_dim=0),                                                                                                                      
    ]),                                                                                                                                                                            
)

Mooncake internally performs:

1. Parse TensorParallelism                                                                                                                                                     
2. Find the TP axis              
3. Validate uniform sharding   
4. Compute the requested TP shard range       
5. Slice the requested shard from the full tensor                                                                                                                                  
6. Store only that shard                                                                                                                                                           
7. Store metadata and manifest for later reads

This lets callers express DP+TP, PP+TP, and EP+TP writes without pre-splitting tensors themselves.

Read-side semantics

Read behavior is explicit through ReadTarget.

ReadTarget mode	Caller intent	Internal behavior
`as_stored`	Return exactly what was stored	Direct object read
`shard`	Return a target shard	Direct shard read if available; otherwise reconstruct the requested shard
`full`	Return the full tensor	Reconstruct from shard metadata / manifest

This avoids guessing read behavior from parallelism alone.

For example, a TP axis can appear in multiple valid read requests:

ReadTarget(mode="as_stored")                
ReadTarget(mode="shard", parallelism=tp_parallelism)                                                                                                                               
ReadTarget(mode="full", parallelism=tp_parallelism)

These requests mean different things, so the mode is explicit.

Reconstruction path

flowchart TB                                                                                                                                                                       
    Read["ReadTarget: shard or full"]                                                                                                                                              
                                                                                                                                                                                   
    Resolve["Resolve exact stored object"]                                                                                                                                     
                                            
    Exact["Exact shard/full object found"]                                                                                                                                         
                                                                                                                                                                                   
    Reconstruct["Reconstruction needed"]                                                                                                                                           
                                                                                                                                                                                   
    Manifest["Load manifest / metadata"]                                                                                                                                       
                               
    Sources["Discover source shards"]         
                                 
    Validate["Validate dtype, shape, split_dim, shard_count"]
                                 
    Plan["Build TensorIntoPlan"]
                                            
    Ranges["get_into_ranges(...)"]                                                                                                                                                 
                                                                                                                                                                                   
    Header["Write tensor metadata header after successful transfer"]                                                                                                               
                                                                                                                                                                                   
    Output["Return tensor / user buffer"]                                                                                                                                      
                               
    Read --> Resolve                          
                                 
    Resolve --> Exact          
    Resolve --> Reconstruct      
                             
    Exact --> Output                        
                                                                                                                                                                                   
    Reconstruct --> Manifest                                                                                                                                                       
    Manifest --> Sources                                                                                                                                                           
    Sources --> Validate                                                                                                                                                           
    Validate --> Plan                                                                                                                                                          
    Plan --> Ranges                           
    Ranges --> Header            
    Header --> Output

Full and reconstructed-shard reads lower to internal range plans and reuse:

get_into_ranges(...)

The public API does not expose:

byte ranges
source offsets
destination offsets
reconstruction plans

Those remain internal planner/runtime details.

Metadata model

Each stored tensor object carries enough metadata to describe both the logical tensor and the stored object.

Metadata field	Purpose
`layout_kind`	Full tensor or shard tensor
`global_shape`	Logical full tensor shape
`local_shape`	Stored object shape
`dtype`	Tensor dtype
`layout_axes`	Parallel layout axes such as TP / DP / PP / EP
`split_dim`	Dimension split by TP, when applicable
`shard_rank`	Rank of this shard along the split axis
`shard_count`	Number of shards along the split axis

For shard sets, Mooncake also stores a manifest containing reconstruction-level information:

Manifest field	Purpose
`global_shape`	Shape of the reconstructed tensor
`dtype`	Tensor dtype
`split_dim`	Split dimension
`shard_count`	Number of uniform shards

How Different Parallelism Axes Are Handled

The unified with_parallelism API models parallelism as an explicit list of axes, but not every axis affects tensor storage in the same way.

The key distinction is:

layout-affecting axes: determine how tensor bytes are physically split
scope / ownership axes: identify which replica, stage, or expert owns the tensor object, but do not define a byte split by themselves

Axis semantics

Axis	Meaning	Affects tensor byte layout?	Handling
`TP`	Tensor-parallel shard	Yes	Defines `rank`, `size`, and `split_dim`; Mooncake can derive the shard range.
`DP`	Data-parallel replica	No	Treated as replica scope; does not imply tensor splitting.
`PP`	Pipeline stage	No	Treated as stage ownership; does not imply tensor splitting.
`EP`	Expert identity	No, by itself	Treated as expert scope; does not imply TP-style tensor splitting.
`DP+TP`	Replica-scoped tensor shard	TP does	DP scopes the object; TP defines the physical shard.
`PP+TP`	Stage-scoped tensor shard	TP does	PP scopes the object; TP defines the physical shard.
`EP+TP`	Expert-scoped tensor shard	TP does	EP scopes the object; TP defines the physical shard.

Why TP is handled differently

TP is the only axis in the current implementation that directly tells Mooncake how to split a tensor.

A TP axis carries:

rank                                                                                                                                                                               
size                                                                                                                                                                               
split_dim

From those fields, Mooncake can compute:

shard_extent = global_shape[split_dim] / tp_size                                                                                                                                   
shard_start  = tp_rank * shard_extent

So for TP-containing layouts, Mooncake can safely derive the requested shard from a full source tensor.

That is why layouts such as:

TensorParallelism([DP, TP])                                                                                                                                                        
TensorParallelism([PP, TP])                                                                                                                                                        
TensorParallelism([EP, TP])

can accept full-tensor input and internally store only the requested TP shard.

Why DP / PP / EP do not imply splitting

DP, PP, and EP identify ownership, not a generic tensor slicing rule.

For example:

TensorParallelism([DP(rank=0, size=2)])

means:

this tensor belongs to data-parallel replica 0

It does not tell Mooncake:

which dimension should be split?                                                                                                                                                   
how large is each shard?                                                                                                                                                           
which byte range belongs to this rank?

The same applies to PP and EP:

PP identifies a pipeline stage.
EP identifies an expert or expert-local scope.
Neither defines a universal tensor split by itself.

Therefore pure DP, PP, or EP layouts store the passed tensor object directly. Mooncake does not invent an implicit split axis.

Why multi-axis layouts are split only through TP

For multi-axis layouts, Mooncake separates scope from layout.

Example:

TensorParallelism([                                                                                                                                                                
    ParallelAxis(kind="dp", rank=0, size=2),                                                                                                                                       
    ParallelAxis(kind="tp", rank=1, size=4, split_dim=0),                                                                                                                          
])

This means:

DP says: this object belongs to replica 0                                                                                                                                          
TP says: store rank 1 of a 4-way split along dimension 0

So the handling is:

full tensor input                                                                                                                                                                  
    -> DP scopes the stored object key/metadata                                                                                                                                    
    -> TP determines the shard slice                                                                                                                                               
    -> Mooncake stores only the requested TP shard

This keeps the semantics explicit and avoids treating mixed parallelism as a vague special case.

Current write behavior by layout

Layout	Caller input	Stored object	Reason
`parallelism=None`	Full tensor	Full tensor	No parallel layout requested.
`[TP]`	TP shard tensor	TP shard	Preserves existing single-axis TP compatibility behavior.
`[DP]`	Tensor object	Tensor object	DP is replica scope, not a split rule.
`[PP]`	Tensor object	Tensor object	PP is stage scope, not a split rule.
`[EP]`	Tensor object	Tensor object	EP is expert scope, not a split rule.
`[DP, TP]`	Full tensor	Requested TP shard	DP scopes the object; TP defines the shard.
`[PP, TP]`	Full tensor	Requested TP shard	PP scopes the object; TP defines the shard.
`[EP, TP]`	Full tensor	Requested TP shard	EP scopes the object; TP defines the shard.

Summary

In short, this work extends unified tensor IO from a TP-centric design into a layout-aware IO layer that supports TP, DP, EP, PP, and their combinations, while also introducing the writer-side optimization that RL workloads actually need: each writer can persist only its owned shard instead of redundantly writing the full tensor. Combined with shard-direct reads, optimized full reconstruction, raw-buffer _from fast paths, and batch API symmetry, this makes the storage layer substantially better suited for RL weight synchronization and mixed-parallel deployment scenarios.

Type of Change

Types
- Bug fix
- New feature
  - Transfer Engine
  - Mooncake Store
  - Mooncake EP
  - Integration
  - P2P Store
  - Python Wheel
- Breaking change
- CI/CD
- Documentation update
- Other

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have formatted my own code using ./scripts/code_format.sh before submitting.
I have updated the documentation.
I have added tests to prove my changes are effective.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/Mooncake/pull/1389 **Author:** [@zxpdemonio](https://github.com/zxpdemonio) **Created:** 1/16/2026 **Status:** ✅ Merged **Merged:** 4/27/2026 **Merged by:** [@stmatengss](https://github.com/stmatengss) **Base:** `main` ← **Head:** `cruz/flex_put_get` --- ### 📝 Commits (10+) - [`2a674c2`](https://github.com/kvcache-ai/Mooncake/commit/2a674c22ca0592d653c3e0ba824cd09531877ec5) [Store] add tensor object metadata and TP upsert APIs - [`d46f11f`](https://github.com/kvcache-ai/Mooncake/commit/d46f11f8e78f408ebe71b19bacbed90573c803d5) [Store] align unified tensor parallelism APIs - [`6c9d006`](https://github.com/kvcache-ai/Mooncake/commit/6c9d0064e17ab9caf1b005f2c4f4c2883c63bf0a) [Store] route tensor into reads through ranged gets - [`7efb1d1`](https://github.com/kvcache-ai/Mooncake/commit/7efb1d1c0c5265b0c72dc1ea545037223631f893) [Store] fix TP ragged reconstruction edge cases - [`bc7f0d9`](https://github.com/kvcache-ai/Mooncake/commit/bc7f0d99890e05030187f6184037099ae6fe039f) [Store] add exact-match multi-axis shard IO support - [`52e7f67`](https://github.com/kvcache-ai/Mooncake/commit/52e7f67b833020de6a4473bdc0fd7fd677c13064) [Store] optimize raw-buffer shard writes and reconstruction reads - [`30101c6`](https://github.com/kvcache-ai/Mooncake/commit/30101c69a7a253ca88feca6c8559ea6af28518ef) [Store] add batch writer-partition unified APIs - [`b6aaa25`](https://github.com/kvcache-ai/Mooncake/commit/b6aaa2569560bc25cbfb5b9a8664e7ed98a598a4) [Store] format writer-shard updates - [`a0c1dea`](https://github.com/kvcache-ai/Mooncake/commit/a0c1dea681dcdb486169c175df0476f98d44aee5) [Store] tighten tensor write wrapper layering - [`a36a9a0`](https://github.com/kvcache-ai/Mooncake/commit/a36a9a0f56626aa5308cfe123e6196d94462d05d) [Store] restore legacy tensor wrapper boundaries ### 📊 Changes **11 files changed** (+7639 additions, -148 deletions) <details> <summary>View changed files</summary> ➕ `docs/source/design/unified-parallel-tensor-io.md` (+405 -0) 📝 `docs/source/python-api-reference/mooncake-store.md` (+228 -2) 📝 `mooncake-integration/integration_utils.h` (+199 -2) 📝 `mooncake-integration/store/store_py.cpp` (+357 -137) ➕ `mooncake-integration/store/store_py_internal.h` (+1483 -0) ➕ `mooncake-integration/store/store_py_parallel_read.h` (+1312 -0) ➕ `mooncake-integration/store/store_py_parallel_write.h` (+1443 -0) 📝 `mooncake-store/include/real_client.h` (+9 -0) 📝 `mooncake-store/src/real_client.cpp` (+22 -0) 📝 `scripts/test_tensor_api.py` (+1578 -3) 📝 `scripts/test_upsert_api.py` (+603 -4) </details> ### 📄 Description ## Motivation This work is primarily driven by **weight synchronization optimization for RL workloads**. In RL systems, different components often consume the same model weights under different parallel layouts: some hold the full tensor, some only need a TP shard, and some run with mixed layouts such as DP+TP, EP+TP, or PP+TP. A storage path that only supports full tensors or plain TP shards causes unnecessary write amplification, unnecessary read amplification, and limited layout expressiveness. The goal of this update is to make the unified tensor IO layer suitable for RL-style weight sync by supporting richer layout identities while also optimizing both shard-oriented writes and shard/full reads. ## What is supported now The unified tensor IO path now supports not only heterogeneous **TP**, but also **DP**, **EP**, **PP**, and their combinations, including exact-match multi-axis shard identities such as **DP+TP**, **PP+TP**, and **EP+TP**. In other words, the system is no longer limited to “TP-only sharded tensors”; it can now persist and resolve multi-axis layout metadata through a single unified API surface. ## Write-path optimization On the write side, the key optimization for RL weight sync is to avoid having every writer store the full tensor when only a local shard needs to be persisted. This is implemented through `writer_partition` for single-key APIs and `writer_partitions` for batch APIs. In practice, this allows each writer to hold a full tensor locally but write only the shard it owns, directly reducing write amplification and network/storage overhead during synchronization. The `_from` paths were also optimized to use **raw-buffer shard planning** instead of decoding tensor-object buffers into Python tensors and then reshaping or re-sharding them, which removes unnecessary decode, tensor materialization, and extra memcpy/clone work from the hot path. ## Read-path optimization On the read side, the system now supports **exact-match shard reads** for multi-axis layouts, so consumers can fetch only the shard they need instead of reading a full tensor and slicing locally. This improves efficiency for TP-only consumers as well as mixed DP/TP, PP/TP, and EP/TP layouts. For cases that need the full tensor view, reconstruction is optimized by reusing the existing `get_into_ranges(...)` ranged-read substrate and avoiding an extra intermediate tensor conversion. As a result, both full reconstruction and `_into` reconstruction remain efficient while still working on top of the unified metadata model. ## RL semantics For RL workloads, an important distinction is that **DP metadata** and **writer-side partial persistence** are not the same thing. If the requirement is “each DP rank has the full tensor, but only writes its owned slice for synchronization efficiency,” that is supported through `writer_partition` / `writer_partitions`. If the requirement is “a DP parallelism spec alone should imply automatic partial persistence of a full tensor,” that is **not** the semantics of the DP metadata path. In the unified API, DP/EP/PP/TP describe **layout identity**, while `writer_partition(s)` provides the **writer-side shard-only persistence optimization** needed for RL weight sync. ## Additional improvements This update also closes an important API symmetry gap by adding batch support for `writer_partitions` across the unified put/upsert APIs, including `_from` variants. Batch equests now explicitly reject ambiguous mixed inputs such as `parallelisms + writer_partitions`. In addition, multi-axis shard metadata matching was fixed by canonicalizing the requested parallelism before comparison, which resolved incorrect mismatch behavior in cases such as `PP+TP` where axis order differed syntactically but not semantically. ## Validation The changes were validated with focused local tests covering batch writer-partition full reconstruction, batch writer-partition `_into` round trips, batch writer-partition upsert and upsert-from paths, multi-axis batch shard reads, multi-axis batch `_into` reads, and explicit invalid mixed-routing cases. Formatting was also rechecked and fixed with the repository formatter. ## Support Matrix | Scenario | Supported | Read support | Write optimization | Notes | |---|---|---|---|---| | TP only | Yes | Shard read, full reconstruction, `_into` | Raw-buffer shard write fast path | Mature path | | Heterogeneous TP | Yes | Shard read, full reconstruction, batch `_into` | Same as TP path | Supported through unified API | | DP metadata layout | Yes | Exact-match shard identity support | Metadata/layout only | Does not by itself mean “write only one slice of a full tensor” | | PP metadata layout | Yes | Exact-match shard identity support | Metadata/layout only | Supports stage-aware layout identity | | EP metadata layout | Yes | Exact-match shard identity support | Metadata/layout only | Supports expert-aware layout identity | | DP + TP | Yes | Exact-match shard read, batch read, `_into`, full reconstruction where applicable | TP shard persistence + unified metadata | Good fit for RL mixed layouts | | PP + TP | Yes | Exact-match shard read, batch read, `_into` | TP shard persistence + unified metadata | Axis-order matching issue was fixed | | EP + TP | Yes | Exact-match shard read | TP shard persistence + unified metadata | Supports expert-aware shard identity | | Writer-side partial persistence from full tensor | Yes | Full reconstruction and shard access | `writer_partition` / `writer_partitions` | Main optimization for RL weight sync | | Batch writer-side partial persistence | Yes | Full reconstruction and shard access | Batch `writer_partitions` | Newly added in this update | | Mixed `parallelisms` + `writer_partitions` in one batch request | Rejected | N/A | N/A | Explicitly invalid to avoid ambiguous semantics | ## High-level architecture ```mermaid flowchart TB Caller["Python caller"] API["Unified with_parallelism API put / get / upsert"] Request["Request objects TensorParallelism / ReadTarget / writer_partition"] Router["Internal route resolver"] Full["Full tensor path"] TPCompat["TP compatibility path"] MultiAxis["Multi-axis parallelism path"] WriterPartition["writer_partition path"] Metadata["Tensor metadata global/local shape, dtype, layout axes"] Manifest["Shard manifest global shape, dtype, split_dim, shard_count"] Store["Mooncake object store"] ReadPlanner["Read planner"] RangeRead["Range-based reconstruction get_into_ranges(...)"] Output["torch.Tensor / user buffer"] Caller --> API API --> Request Request --> Router Router --> Full Router --> TPCompat Router --> MultiAxis Router --> WriterPartition Full --> Metadata TPCompat --> Metadata MultiAxis --> Metadata WriterPartition --> Metadata Metadata --> Store MultiAxis --> Manifest WriterPartition --> Manifest Manifest --> Store Store --> ReadPlanner Manifest --> ReadPlanner ReadPlanner --> Output ReadPlanner --> RangeRead RangeRead --> Output ``` ## Public API shape This PR centers tensor IO around the following API families: ```python put_tensor_with_parallelism(...) batch_put_tensor_with_parallelism(...) get_tensor_with_parallelism(...) batch_get_tensor_with_parallelism(...) get_tensor_with_parallelism_into(...) batch_get_tensor_with_parallelism_into(...) upsert_tensor_with_parallelism(...) batch_upsert_tensor_with_parallelism(...) ``` The core public request objects are: ```python TensorParallelism ParallelAxis ReadTarget ``` `TensorParallelism` describes where a tensor object belongs in the parallel layout. `ReadTarget` describes what the caller wants to materialize on read: - stored object - target shard - reconstructed full tensor ## Multi-axis write flow For a request such as: ```python put_tensor_with_parallelism( key, full_tensor, TensorParallelism([ ParallelAxis(kind="dp", rank=0, size=2), ParallelAxis(kind="tp", rank=1, size=4, split_dim=0), ]), ) ``` Mooncake internally performs: ```text 1. Parse TensorParallelism 2. Find the TP axis 3. Validate uniform sharding 4. Compute the requested TP shard range 5. Slice the requested shard from the full tensor 6. Store only that shard 7. Store metadata and manifest for later reads ``` This lets callers express `DP+TP`, `PP+TP`, and `EP+TP` writes without pre-splitting tensors themselves. ## Read-side semantics Read behavior is explicit through `ReadTarget`. | ReadTarget mode | Caller intent | Internal behavior | |---|---|---| | `as_stored` | Return exactly what was stored | Direct object read | | `shard` | Return a target shard | Direct shard read if available; otherwise reconstruct the requested shard | | `full` | Return the full tensor | Reconstruct from shard metadata / manifest | This avoids guessing read behavior from `parallelism` alone. For example, a TP axis can appear in multiple valid read requests: ```python ReadTarget(mode="as_stored") ReadTarget(mode="shard", parallelism=tp_parallelism) ReadTarget(mode="full", parallelism=tp_parallelism) ``` These requests mean different things, so the mode is explicit. ## Reconstruction path ```mermaid flowchart TB Read["ReadTarget: shard or full"] Resolve["Resolve exact stored object"] Exact["Exact shard/full object found"] Reconstruct["Reconstruction needed"] Manifest["Load manifest / metadata"] Sources["Discover source shards"] Validate["Validate dtype, shape, split_dim, shard_count"] Plan["Build TensorIntoPlan"] Ranges["get_into_ranges(...)"] Header["Write tensor metadata header after successful transfer"] Output["Return tensor / user buffer"] Read --> Resolve Resolve --> Exact Resolve --> Reconstruct Exact --> Output Reconstruct --> Manifest Manifest --> Sources Sources --> Validate Validate --> Plan Plan --> Ranges Ranges --> Header Header --> Output ``` Full and reconstructed-shard reads lower to internal range plans and reuse: ```python get_into_ranges(...) ``` The public API does not expose: - byte ranges - source offsets - destination offsets - reconstruction plans Those remain internal planner/runtime details. ## Metadata model Each stored tensor object carries enough metadata to describe both the logical tensor and the stored object. | Metadata field | Purpose | |---|---| | `layout_kind` | Full tensor or shard tensor | | `global_shape` | Logical full tensor shape | | `local_shape` | Stored object shape | | `dtype` | Tensor dtype | | `layout_axes` | Parallel layout axes such as TP / DP / PP / EP | | `split_dim` | Dimension split by TP, when applicable | | `shard_rank` | Rank of this shard along the split axis | | `shard_count` | Number of shards along the split axis | For shard sets, Mooncake also stores a manifest containing reconstruction-level information: | Manifest field | Purpose | |---|---| | `global_shape` | Shape of the reconstructed tensor | | `dtype` | Tensor dtype | | `split_dim` | Split dimension | | `shard_count` | Number of uniform shards | ## How Different Parallelism Axes Are Handled The unified `with_parallelism` API models parallelism as an explicit list of axes, but not every axis affects tensor storage in the same way. The key distinction is: - **layout-affecting axes**: determine how tensor bytes are physically split - **scope / ownership axes**: identify which replica, stage, or expert owns the tensor object, but do not define a byte split by themselves ### Axis semantics | Axis | Meaning | Affects tensor byte layout? | Handling | |---|---|---:|---| | `TP` | Tensor-parallel shard | Yes | Defines `rank`, `size`, and `split_dim`; Mooncake can derive the shard range. | | `DP` | Data-parallel replica | No | Treated as replica scope; does not imply tensor splitting. | | `PP` | Pipeline stage | No | Treated as stage ownership; does not imply tensor splitting. | | `EP` | Expert identity | No, by itself | Treated as expert scope; does not imply TP-style tensor splitting. | | `DP+TP` | Replica-scoped tensor shard | TP does | DP scopes the object; TP defines the physical shard. | | `PP+TP` | Stage-scoped tensor shard | TP does | PP scopes the object; TP defines the physical shard. | | `EP+TP` | Expert-scoped tensor shard | TP does | EP scopes the object; TP defines the physical shard. | ### Why TP is handled differently `TP` is the only axis in the current implementation that directly tells Mooncake how to split a tensor. A TP axis carries: ```text rank size split_dim ``` From those fields, Mooncake can compute: ```text shard_extent = global_shape[split_dim] / tp_size shard_start = tp_rank * shard_extent ``` So for TP-containing layouts, Mooncake can safely derive the requested shard from a full source tensor. That is why layouts such as: ```python TensorParallelism([DP, TP]) TensorParallelism([PP, TP]) TensorParallelism([EP, TP]) ``` can accept full-tensor input and internally store only the requested TP shard. ### Why DP / PP / EP do not imply splitting `DP`, `PP`, and `EP` identify ownership, not a generic tensor slicing rule. For example: ```python TensorParallelism([DP(rank=0, size=2)]) ``` means: ```text this tensor belongs to data-parallel replica 0 ``` It does not tell Mooncake: ```text which dimension should be split? how large is each shard? which byte range belongs to this rank? ``` The same applies to PP and EP: - `PP` identifies a pipeline stage. - `EP` identifies an expert or expert-local scope. - Neither defines a universal tensor split by itself. Therefore pure `DP`, `PP`, or `EP` layouts store the passed tensor object directly. Mooncake does not invent an implicit split axis. ### Why multi-axis layouts are split only through TP For multi-axis layouts, Mooncake separates scope from layout. Example: ```python TensorParallelism([ ParallelAxis(kind="dp", rank=0, size=2), ParallelAxis(kind="tp", rank=1, size=4, split_dim=0), ]) ``` This means: ```text DP says: this object belongs to replica 0 TP says: store rank 1 of a 4-way split along dimension 0 ``` So the handling is: ```text full tensor input -> DP scopes the stored object key/metadata -> TP determines the shard slice -> Mooncake stores only the requested TP shard ``` This keeps the semantics explicit and avoids treating `mixed` parallelism as a vague special case. ### Current write behavior by layout | Layout | Caller input | Stored object | Reason | |---|---|---|---| | `parallelism=None` | Full tensor | Full tensor | No parallel layout requested. | | `[TP]` | TP shard tensor | TP shard | Preserves existing single-axis TP compatibility behavior. | | `[DP]` | Tensor object | Tensor object | DP is replica scope, not a split rule. | | `[PP]` | Tensor object | Tensor object | PP is stage scope, not a split rule. | | `[EP]` | Tensor object | Tensor object | EP is expert scope, not a split rule. | | `[DP, TP]` | Full tensor | Requested TP shard | DP scopes the object; TP defines the shard. | | `[PP, TP]` | Full tensor | Requested TP shard | PP scopes the object; TP defines the shard. | | `[EP, TP]` | Full tensor | Requested TP shard | EP scopes the object; TP defines the shard. | ## Summary In short, this work extends unified tensor IO from a TP-centric design into a layout-aware IO layer that supports **TP, DP, EP, PP, and their combinations**, while also introducing the writer-side optimization that RL workloads actually need: **each writer can persist only its owned shard instead of redundantly writing the full tensor**. Combined with shard-direct reads, optimized full reconstruction, raw-buffer `_from` fast paths, and batch API symmetry, this makes the storage layer substantially better suited for RL weight synchronization and mixed-parallel deployment scenarios. ## Type of Change * Types - [ ] Bug fix - [x] New feature - [ ] Transfer Engine - [x] Mooncake Store - [ ] Mooncake EP - [x] Integration - [ ] P2P Store - [ ] Python Wheel - [ ] Breaking change - [ ] CI/CD - [ ] Documentation update - [ ] Other ## How Has This Been Tested?  ## Checklist - [x] I have performed a self-review of my own code. - [x] I have formatted my own code using `./scripts/code_format.sh` before submitting. - [x] I have updated the documentation. - [x] I have added tests to prove my changes are effective. --- 🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.

bsakaguchi

2026-10-29 21:33:42 +00:00