[RFC]: RL Disaggregation via Rollout / Update Separation and Shared Mooncake Data Plane

### Changes proposed

## Summary

This RFC proposes a design for RL disaggregation in Mooncake from Issue #1883 **Milestone 1** by separating the rollout phase from the update phase and using Mooncake's data plane to move rollout artifacts, training samples, and versioned state between independently scaled worker pools.

In RL post-training workflows such as PPO, GRPO, and RLHF, rollout and update have different resource profiles:

* rollout is inference-heavy and KV-cache-heavy,
* update is training/optimizer-heavy and collective-heavy.
Co-locating rollout and update on the same GPU pool can lead to poor utilization because one phase may idle while the other phase is active. The proposed design introduces a clearer boundary between:

* **Rollout workers**, which generate samples,
* **Update workers**, which consume samples and perform training steps,
* **Mooncake Store / Transfer Engine**, which provides the shared data plane.
The first version should be design-first and prototype-oriented. It should define the data contracts and versioning semantics before attempting a full RL framework integration.

## Motivation

RL post-training pipelines often alternate between generation and training. These phases stress the system differently:

| Phase | Primary Work | Main Bottleneck |
| ------- | ------------------------- | ----------------------------------------------- |
| Rollout | sampling / generation | inference throughput, KV cache, serving latency |
| Update | gradient step / optimizer | training throughput, collectives, weight update |

A single colocated setup is simple, but it can waste resources when rollout and update do not need the same number or type of GPUs.

Mooncake already focuses on disaggregated serving and KV-cache-centric infrastructure. RL disaggregation is a natural extension: rollout and update should be able to scale independently while sharing a reliable data plane for samples, metadata, and versioned state.

The current RL sample flow demonstrates the idea of passing data through Mooncake Store, but it does not yet define a production-oriented rollout/update boundary.

## Goals

1. Define the boundary between rollout workers and update workers.
2. Define the minimum data contract for RL rollout artifacts.
3. Use Mooncake Store and Transfer Engine as the shared data plane.
4. Track policy/weight versions explicitly.
5. Allow rollout and update pools to scale independently.
6. Support a minimal prototype before full RL framework integration.
7. Keep KV reuse semantics safe by default.
8. Gather maintainer input on the first target RL framework.
## Non-Goals

1. This RFC does not build a full RL training framework.
2. This RFC does not implement every PPO, GRPO, or RLHF variant.
3. This RFC does not redesign optimizer or gradient collective logic.
4. This RFC does not assume KV cache can be reused across weight updates by default.
5. This RFC does not require immediate integration with every serving backend.
6. This RFC does not define a final public API without maintainer feedback.
## Key Design Principle: Versioned State

KV cache and rollout artifacts must be associated with a policy or weight version.

A KV cache produced under policy version `v1` should not be treated as valid for exact generation under policy version `v2` unless the system explicitly chooses an off-policy or stale-policy mode.

Default behavior should be conservative:

```text
KV reuse is valid only when the cached KV version matches the rollout worker's active policy version.
```

This avoids silently mixing model states across update boundaries.

## Proposed Architecture

```text
+--------------------+ +-------------------------+
| Rollout Workers | | Update Workers |
| | | |
| - sampling | | - consume trajectories |
| - generation | | - compute loss |
| - logprobs | | - optimizer step |
| - rewards metadata | | - publish new version |
+---------+----------+ +-----------+-------------+
| ^
| write rollout artifacts |
v |
+---------------------------------------------------------+
| Mooncake Shared Data Plane |
| |
| - Mooncake Store |
| - Transfer Engine |
| - versioned sample metadata |
| - optional KV references |
| - trajectory/logprob/reward records |
| - weight-version metadata |
+---------------------------------------------------------+
^ |
| read latest policy version |
| / fetch state |
v |
+---------------------------------------------------------+
| Weight / Policy Version Registry |
| |
| - current policy version |
| - previous versions |
| - staleness policy |
| - update step metadata |
+---------------------------------------------------------+
```

## Core Concepts

### RolloutSession

A `RolloutSession` represents a rollout-side unit of generation.

Possible fields:

```python
@dataclass
class RolloutSession:
session_id: str
episode_id: str
request_id: str
policy_version: str
prompt_ref: str
kv_ref: Optional[str]
output_ref: Optional[str]
logprob_ref: Optional[str]
reward_ref: Optional[str]
metadata: Dict[str, Any]
```

The exact implementation does not need to use this dataclass. The purpose is to define the conceptual boundary.

### UpdateSession

An `UpdateSession` represents an update-side training step or mini-batch consumption unit.

Possible fields:

```python
@dataclass
class UpdateSession:
update_id: str
input_policy_version: str
output_policy_version: Optional[str]
trajectory_refs: List[str]
reward_refs: List[str]
logprob_refs: List[str]
status: str
metadata: Dict[str, Any]
```

### PolicyVersion

A `PolicyVersion` identifies the model/weight version used to generate samples.

Possible fields:

```python
@dataclass
class PolicyVersion:
version_id: str
parent_version_id: Optional[str]
created_at_step: int
weight_ref: Optional[str]
metadata: Dict[str, Any]
```

## Data Plane Contract

The rollout-to-update boundary should include:

1. prompt or input reference,
2. generated output tokens,
3. logprobs,
4. reward or reward reference,
5. trajectory metadata,
6. policy/weight version,
7. optional KV reference,
8. status and lifecycle metadata.
Suggested logical record:

```json
{
"episode_id": "episode-123",
"request_id": "request-456",
"policy_version": "policy-v17",
"prompt_ref": "store://prompts/episode-123",
"output_ref": "store://outputs/episode-123",
"logprob_ref": "store://logprobs/episode-123",
"reward_ref": "store://rewards/episode-123",
"kv_ref": "store://kv/policy-v17/episode-123",
"status": "ready_for_update",
"metadata": {
"algorithm": "ppo",
"rollout_worker": "rollout-0",
"created_at_step": 17
}
}
```

## KV Reuse Semantics

### Safe Default

KV reuse is allowed only when:

```text
cached_kv.policy_version == rollout_worker.active_policy_version
```

If the policy version differs, the rollout worker should recompute or ignore the cached KV by default.

### Optional Staleness Mode

Some RL algorithms may tolerate stale samples or off-policy data. If maintainers want to support this later, the mode should be explicit.

Example:

```text
--allow-stale-rollout-samples
--max-policy-lag 1
```

This should apply to sample consumption semantics, not silently to exact KV reuse.

### Why This Matters

KV tensors are produced by a specific model state. Reusing KV across weight updates can be incorrect if the new model weights differ from the weights that generated the cache.

Therefore, this RFC treats versioning as part of the core design rather than an implementation detail.

## Proposed API Shape

This section is illustrative only.

### Rollout Side

```python
session = RolloutSession(
session_id="rollout-session-1",
episode_id="episode-123",
request_id="request-456",
policy_version=current_policy_version,
)

store.put(session.prompt_ref, prompt)
store.put(session.output_ref, generated_tokens)
store.put(session.logprob_ref, logprobs)
store.put(session.reward_ref, reward)

if kv_cache_is_materialized:
store.put(session.kv_ref, kv_cache_metadata)

store.put(
f"rollout_sessions/{session.session_id}",
session.to_json(),
)
```

### Update Side

```python
session = store.get("rollout_sessions/rollout-session-1")

if session.policy_version not in accepted_policy_versions:
handle_stale_sample(session)

trajectory = store.get(session.output_ref)
logprobs = store.get(session.logprob_ref)
reward = store.get(session.reward_ref)

loss = compute_loss(trajectory, logprobs, reward)
new_policy_version = optimizer_step(loss)

publish_policy_version(new_policy_version)
```

## Weight Version Flow

The update side should publish new policy versions after optimizer steps.

Possible flow:

```text
1. Rollout workers generate samples with policy-vN.
2. Rollout workers write samples and metadata to Mooncake Store.
3. Update workers consume samples for policy-vN.
4. Update workers perform optimizer step.
5. New weights are published as policy-vN+1.
6. Rollout workers observe policy-vN+1.
7. New rollout sessions use policy-vN+1.
8. KV generated under policy-vN is not reused for exact generation under policy-vN+1 by default.
```

## Scaling Model

Rollout and update worker pools should scale independently.

Example:

```text
rollout_replicas: 16
update_replicas: 4
```

The exact scheduler is outside the scope of this RFC. The RFC only requires that Mooncake's data plane can support the separation cleanly.

## Implementation Plan

### Phase 1 — Design and Interfaces

* Define rollout/update data contract.
* Define policy version metadata.
* Define safe KV reuse rules.
* Add a design document under `mooncake-rl/` or `docs/`.
### Phase 2 — Minimal Prototype

* Extend the existing RL sample flow into a more realistic producer/consumer example.
* Add separate rollout and update processes.
* Use Mooncake Store for rollout artifacts.
* Track policy version in metadata.
* Demonstrate stale-version rejection or recomputation behavior.
### Phase 3 — Framework Selection

Ask maintainers which RL framework should be the first integration target.

Candidates may include:

* slime,
* verl,
* OpenRLHF,
* an existing Mooncake/SGLang-adjacent integration,
* another maintainer-preferred framework.
### Phase 4 — Reference Integration

After framework selection:

* implement minimal integration,
* run PPO-style throughput comparison,
* measure colocated vs disaggregated flow,
* document limitations.
### Phase 5 — Optimization

Potential later optimizations:

* batched sample movement,
* async prefetch,
* bulk transfer of weight versions,
* policy-version-aware cache retention,
* improved backpressure between rollout and update pools.
## Benchmark Plan

A useful benchmark should compare:

1. colocated rollout/update baseline,
2. disaggregated rollout/update with Mooncake data plane,
3. different rollout:update worker ratios,
4. different sample sizes,
5. different policy update frequencies.
Possible metrics:

* samples/sec,
* tokens/sec during rollout,
* update steps/sec,
* rollout GPU utilization,
* update GPU utilization,
* end-to-end RL step time,
* data-plane transfer latency,
* stale sample rate,
* recompute rate due to policy mismatch.
## CI Strategy

Initial CI should only validate lightweight behavior:

1. data contract serialization,
2. policy-version comparison logic,
3. mocked producer/consumer flow,
4. no dependency on a full RL framework.
Full RL benchmarks should remain manual or nightly until a stable test environment exists.

## Alternatives Considered

### Status Quo: Colocated Rollout and Update

This is simpler but can waste resources because rollout and update have different hardware and scheduling characteristics.

### Plain File/Object Store Without Mooncake

A generic store can move samples but does not exercise Mooncake's intended disaggregated data-plane role and does not integrate naturally with KV/cache-aware infrastructure.

### Immediate Full Framework Integration

Rejected for first step.

A full framework integration before agreeing on the data contract may lead to the wrong abstraction. The safer path is to define the boundary first, then integrate.

### Reuse KV Across Weight Updates by Default

Rejected.

KV cache is produced by a specific policy/weight version. Reusing it across weight updates can be incorrect unless a specific stale-policy mode is explicitly designed and accepted.

## Risks

1. RL frameworks differ significantly in data formats and lifecycle.
2. KV reuse semantics are easy to get wrong.
3. Weight versioning may require integration with serving/runtime components.
4. A prototype may not show benefits unless the benchmark workload is realistic.
5. Too much framework-specific logic could make the Mooncake abstraction less reusable.
## Open Questions

1. Which RL framework should be the first reference integration target?
2. Should Mooncake define generic rollout/update records, or adapt to one framework's schema first?
3. Where should policy-version metadata live?
4. Should stale rollout samples be supported initially, or rejected by default?
5. Should KV references be part of the v1 data contract, or deferred until the sample/logprob/reward path is stable?
6. What is the minimum benchmark that maintainers would consider useful for M1?
7. Should the first prototype be single-node multi-process, multi-node, or both?
## Expected Outcome

This RFC should produce a clear rollout/update boundary for RL disaggregation and a minimal prototype that uses Mooncake as the shared data plane. The design should be conservative about KV reuse and explicit about policy versioning.

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Start by reading the existing RL sample flow and inspecting the proposed `mooncake-rl/` or `docs/` locations. Trace how Mooncake Store and the Transfer Engine currently carry rollout data, then define the rollout/update contracts, policy-version semantics, and minimal prototype scope; done requires maintainer agreement on the first framework target.

Written by the indexing model from the issue text.