[RFC]: Disaggregated Diffusion Model Serving with Mooncake

### Changes proposed

# [RFC] Disaggregated Diffusion Model Serving with Mooncake

## Introduction

This RFC proposes first-class support in Mooncake for **disaggregated serving of diffusion models** (image and video generation), decoupling the pipeline into independently scalable stages — **Condition Encode (E)**, **Denoise (D, the DiT/U-Net transformer)**, and **VAE Decode (V)** — connected by the Mooncake Transfer Engine and Mooncake Store, and orchestrated through a **queue-based multi-stage scheduler**.

This directly targets two roadmap items:
- *(Disaggregated Diffusion Deployment)*: decentralized disaggregated deployment for video generation models, decoupling transformer and decoder stages.
- *(Queue-based Orchestration)*: queue-based scheduling for multi-stage diffusion serving requests.

## Motivation

LLM-style Prefill/Decode (PD) disaggregation exists because prefill is compute-bound and decode is memory-bound. Diffusion pipelines have an analogous — but structurally different — resource asymmetry:

| Stage | Workload character | Typical hardware fit |
|---|---|---|
| Condition encoders (T5/CLIP/LLM captioner, image/video encoders) | Small, bursty, compute-light; highly cacheable (prompts repeat) | Small GPUs / MIG slices / CPU |
| Denoise transformer (DiT) | Compute-bound, iterative (20–100 steps), long-running (seconds–minutes for video); scales with sequence/CFG/Ulysses parallelism | Large multi-GPU groups (TP/SP) |
| VAE decode | Memory-bandwidth- and VRAM-capacity-bound (activations at pixel resolution, especially video); embarrassingly parallel across temporal chunks | Fewer, memory-rich GPUs |

Co-locating these stages on one GPU group causes the same problems PD disaggregation solved for LLMs:

1. **VRAM contention**: for video models (e.g., Wan, HunyuanVideo, Qwen-series generators), VAE decode activation memory competes with DiT weights/activations, forcing tiled decode or smaller batch sizes.
2. **Poor utilization**: text encoders and the VAE are idle >90% of the time while the DiT denoises; the DiT is idle during decode. Independent scaling fixes this (e.g., N denoise groups sharing one decode pool).
3. **Interference and head-of-line blocking**: a 60-second video decode stalls the next request's denoise steps; step-latency jitter breaks SLOs for interactive image generation.
4. **No cross-instance reuse**: identical prompts/negative prompts re-run the text encoder everywhere; there is no shared cache for condition embeddings or intermediate latents.

**Prior art already validated the approach on Mooncake**: LightX2V ships encoder/transformer disaggregation over the Transfer Engine ([blog](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/)); SGLang EPD disaggregation decouples multimodal encoders using Mooncake RDMA; SGLang's diffusion disaggregation RFC (sgl-project/sglang#19512) decomposes the same three stages inside the engine; vLLM-Omni uses `MooncakeTransferEngineConnector`/`MooncakeStoreConnector` for its AR → Generation → Diffusion pipeline. This RFC generalizes these point integrations into a reusable Mooncake-level capability.

## Non-goals

- Implementing a diffusion inference engine (this stays in LightX2V / SGLang / vLLM-Omni / xDiT etc.); Mooncake provides transfer, storage, and orchestration primitives.
- Intra-stage parallelism (TP/SP/Ulysses inside the DiT) — owned by the serving framework.
- Model-quality techniques (distillation, step reduction).

## Background: what actually moves between stages

Unlike LLM PD (paged KV cache blocks), diffusion inter-stage payloads are **dense, contiguous tensors, transferred once per stage boundary**, not per token:

| Payload | Producer → Consumer | Size (typical) |
|---|---|---|
| Condition embeddings (text/image/video) | Encode → Denoise | 1–50 MB |
| Denoised latents | Denoise → Decode | image: 0.1–2 MB; video: 100 MB–several GB (e.g., 720p/121-frame latents) |
| Intermediate latents (per denoise step) | Denoise → Denoise (preemption/migration), or → preview decode | same order as final latents |
| Reusable features (CFG branch, TeaCache-style skipped-step features) | Denoise ↔ Store | tens of MB |

Consequences for the design:
- Transfers are **large, batched, and latency-tolerant relative to a denoise step** — a perfect fit for `batch_transfer_sync_write` / `BatchTransfer` over multi-NIC RDMA; no need for per-layer streaming like KV cache.
- Video latents are large enough that **Denoise → Decode transfer should overlap with denoising** (chunk-wise: transfer temporal chunk *t* while chunk *t+1* is still being processed, or transfer as soon as the final step completes per chunk).
- Because payloads are immutable once produced, they map cleanly onto Mooncake Store's `Put`/`Get` object model with hash-based keys (prompt hash → embedding; request-id + step → latent).

## Proposed design

Three deliverables, in dependency order:

### 1. Inter-stage transfer connector (`DiffusionConnector`)

A thin, framework-facing connector (Python, under `mooncake-integration/`) mirroring the vLLM `MooncakeConnector` split:

- **Producer side** registers output tensor memory with the Transfer Engine (`register_memory`) at stage start-up — the actual inference buffers (encoder output, DiT latent output, VAE input), so all transfers are DMA-only with no CPU-side copies or serialization — and exposes them as segments.
- **Consumer side** pulls via `batch_transfer_sync_read` (or the producer pushes with `batch_transfer_sync_write`), using the existing `P2PHANDSHAKE` metadata path so no extra metadata service is required.
- **Two transport modes**, chosen per payload:
- *Direct P2P* (Transfer Engine): GPUDirect VRAM-to-VRAM, lowest latency, for the hot path Denoise → Decode when the consumer is known and has capacity at schedule time.
- *Store-relayed* (Mooncake Store `Put`/`Get`): the producer publishes into the Store's cluster-wide DRAM pool (with SSD spill) and frees its VRAM immediately — decouples producer/consumer lifetimes, which is required for queue-based orchestration (below), elastic pools, and any-to-any routing. This is the same relay pattern SGLang-Omni already uses for thinker/talker/vocoder stages.
- **Chunked video-latent transfer**: an API to transfer a latent tensor in temporal chunks with a completion callback per chunk, so VAE decode can start on chunk 0 while the tail is still in flight. Each chunk is addressed as an offset/length slice of the already-registered latent buffer. This is the main new mechanism vs. existing connectors (which transfer whole objects).
- **Request-level asynchrony**: `publish()` is asynchronous per request (not per batch), so one slow request in a producer batch never gates its batchmates' handoff.

### 2. Queue-based orchestration

LLM PD routing is point-to-point (a prefill instance is paired with a decode instance per request). Diffusion stages have wildly different service times (encode: ms; denoise: seconds–minutes; decode: seconds), so static pairing wastes capacity. We propose **per-stage queues** with Store-relayed handoff — deliberately **asymmetric**, because the stages are not symmetric:

- **Encode and Decode get true work queues.** Both are stateless single-shot operations; any pool worker can serve any item, so work-stealing across the pool is safe and maximizes utilization.
- **Denoise gets an admission queue only.** A denoise job is 20–50+ *sequential, stateful* steps (evolving latent, sampler state, sequence-parallel attention sharding). Once dequeued, the job is pinned to one denoise worker group until completion — the queue decides *when* and *where* a job starts, never per-step placement. Re-queuing individual steps would add a Store round-trip per step (≈50x the transfer volume of stage-boundary handoff) and break TP/SP shard locality. Mid-job movement happens only through the coarse-grained step-checkpoint mechanism (Section 3), reserved for preemption, elastic scale-down, and failure recovery.

Mechanics:

- Each stage pool consumes from its queue; the payload reference (Store key + shape/dtype metadata) travels in the queue message, not the data.
- The orchestrator (a small service, or a library embedded in an existing router such as Mooncake Conductor) implements:
- **Priority and SLO classes** — interactive image requests can jump ahead of batch video jobs at the decode queue.
- **Credit-based flow control** — downstream workers advertise free buffer slots (credits); the orchestrator dispatches only against available credits, so a slow decode pool throttles denoise admission *before* multi-GB latents pile up, rather than reacting to queue depth after the fact.
- **Affinity hints** — prefer a decode worker on the same rail/NUMA domain as the producing denoise group (reuse Transfer Engine topology awareness).
- **Lifecycle/eviction** — latents are single-consumer: delete-on-consume leases (aligned with Store's lease + eviction machinery) so the pool doesn't fill with dead latents; embeddings are multi-consumer and cached with normal LRU/soft-pin.
- **Failure semantics**: stages are idempotent given their inputs. If a decode worker dies, the latent is still in the Store and the work item is re-queued — this is a *correctness win over direct P2P pairing* and the main argument for store-relayed handoff as the default for video.
- **Centralized first, decentralized later.** Diffusion dispatch rates are requests/sec, not tokens/sec — each job runs seconds to minutes, so even a large cluster generates trivial control-plane traffic. We start with a single orchestrator and keep the routing policy pluggable so a decentralized control plane (e.g., etcd-based discovery with direct P2P dispatch) can replace it later without touching the data plane, matching the roadmap's "decentralized disaggregated deployment" wording as a later milestone rather than a prerequisite.

### 3. Cross-instance caching (reuse, not just transfer)

Mooncake Store as a **shared cache pool** for diffusion, analogous to prefix caching / the SGLang Encoder Global Cache Manager:

- **Condition-embedding cache**: key = hash(encoder-id, prompt, params). Negative prompts are extremely hot (often a handful of strings cluster-wide) — near-100% hit rate. This makes the Encode pool tiny.
- **CFG branch sharing**: with classifier-free guidance the unconditional branch depends only on the negative prompt; its features can be cached and shared across requests with the same negative prompt at matched steps (best-effort, framework opt-in).
- **Step-checkpoint cache**: persisting intermediate latents at step *k* enables (a) preemption/migration of long video jobs across denoise groups (elastic scaling, spot recovery), and (b) "edit and resume" workflows that re-denoise from a mid-trajectory checkpoint.

### Proposed API sketch (Python)

```python
from mooncake.diffusion import DiffusionConnector, StagePayload

conn = DiffusionConnector(role="denoise", # encode | denoise | decode
transfer_mode="auto", # p2p | store | auto
store_config=..., te_config=...)

# producer (denoise worker): publish latents chunk-by-chunk
handle = conn.publish(StagePayload(
request_id=req_id, kind="latent",
tensor=latents, chunk_dim=2, num_chunks=8, # temporal chunking
lease="delete_on_consume"))

# consumer (decode worker): pulled from queue, fetch chunks as they land
# `dst` is a pre-allocated, pre-registered device tensor — chunks are written
# directly into it via RDMA/NVLink, `chunk` below is a view, not a copy.
dst = torch.empty_like(expected_latent_shape, device="cuda:0")
for chunk in conn.consume(handle, dst=dst):
vae.decode_chunk(chunk)
```

## Architecture

```mermaid
flowchart TB
ORCH["Diffusion Orchestrator
(queue-based scheduler; per-stage work queues)"]

ORCH -->|enqueue| ENC["Encode pool
(T5 / CLIP / …)"]
ORCH -->|enqueue| DEN["Denoise pool
(DiT, TP/SP)"]
ORCH -->|enqueue| DEC["Decode pool
(VAE)"]

ENC -->|embeddings| TE["Transfer Engine
(RDMA / NVLink / TCP)"]
DEN -->|latents| TE
DEC --> TE

TE --- STORE[("Mooncake Store
(embedding / latent cache)")]

classDef pool fill:#2563eb,color:#fff,stroke:#1e40af;
classDef infra fill:#059669,color:#fff,stroke:#065f46;
classDef orch fill:#7c3aed,color:#fff,stroke:#5b21b6;
class ENC,DEN,DEC pool;
class TE,STORE infra;
class ORCH orch;
```

Chunked Denoise → Decode handoff, overlapping transfer and decode:

```mermaid
sequenceDiagram
participant Denoise as Denoise worker (DiT)
participant TE as Transfer Engine
participant Decode as Decode worker (VAE)

Note over Denoise: final step, chunk 0 ready
Denoise->>TE: publish(chunk 0)
Note over Denoise: continues denoising chunk 1..N
TE->>Decode: chunk 0 available
Decode->>Decode: decode chunk 0
Denoise->>TE: publish(chunk 1)
TE->>Decode: chunk 1 available
Decode->>Decode: decode chunk 1
Note over Denoise,Decode: ... overlapped for remaining chunks ...
Denoise->>TE: publish(chunk N)
TE->>Decode: chunk N available
Decode->>Decode: decode chunk N
Decode-->>Denoise: lease released (delete-on-consume)
```

## Implementation plan

| Phase | Scope | Depends on |
|---|---|---|
| 1 | `DiffusionConnector` with whole-tensor P2P + Store-relay modes; reference integration in **LightX2V** (upstream its ad-hoc transfer onto the connector) | existing TE / Store APIs |
| 2 | Chunked latent transfer + decode overlap; delete-on-consume lease mode in Store | Phase 1 |
| 3 | Queue-based orchestrator (standalone service + library mode); credit-based flow control, priority classes | Phase 1 |
| 4 | Embedding / CFG / step-checkpoint caching; second integration (vLLM-Omni Generation→Diffusion stage or SGLang diffusion runtime) | Phases 1–3 |

Benchmarks to publish with Phase 2/3: (a) end-to-end video-gen throughput and P99 TTFF (time-to-first-frame) for colocated vs. disaggregated at equal GPU count; (b) Denoise→Decode transfer overlap efficiency; (c) decode-pool sharing ratio (N denoise groups : M decode workers) at target SLO; (d) connector overhead vs. a raw `batch_transfer_sync_write` baseline, confirming the connector adds metadata/scheduling cost only, not data copies.

## Alternatives considered

- **Framework-internal solutions only** (keep this in LightX2V / vLLM-Omni / SGLang): each framework re-implements chunked transfer, staging buffers, leases, credits, and queueing; no shared cache or staging pool across engines. The SGLang/vLLM connector history shows the primitives belong in Mooncake.
- **Direct P2P only, no Store relay**: lower latency but couples stage lifetimes — the producer must hold multi-GB outputs in scarce VRAM while the downstream pool is busy — and breaks elastic scaling and retry-on-failure for multi-minute video jobs.
- **Per-node pinned-host staging buffers** (the sglang#19512 approach): solves the VRAM-residency problem but statically partitions staging capacity per node (one node's idle staging RAM can't absorb another's burst), always pays two PCIe hops even when the consumer is ready, and requires a bespoke allocator; the Store's pooled DRAM tier provides the same decoupling cluster-wide with existing allocator/lease/eviction machinery.
- **Generic message queue (Redis/Kafka) carrying tensors**: cannot approach RDMA bandwidth for multi-GB latents; the queue should carry references, data goes over the Transfer Engine — which is exactly the proposed split.

## Open questions

1. Should the orchestrator live in **Mooncake Conductor** (extending its routing role to multi-stage pipelines) or as a separate lightweight component under `mooncake-integration/`?
2. Chunked-transfer API: expose at the Transfer Engine level (generalizes to Omni stage transfer, roadmap item *Worker-Level Inter-Stage Transfer*) or keep it in the diffusion connector initially?
3. Do we need a first-class **tensor metadata schema** (dtype/shape/layout/device) in Store objects, or continue with side-band metadata in the queue message as vLLM-Omni does today?
4. Step-checkpoint cache keys: standardize a scheduler-state hash (sampler, step index, seed, guidance scale) so checkpoints are portable across denoise workers?

## Related work

- Mooncake roadmap Milestone 10 — Diffusion Collaboration items #1883
- [RFC]: Mooncake and vLLM-Omni Collaboration Roadmap (vllm-project/vllm-omni#2904)
- [RFC]: Qwen3-Omni Stage Transfer via Mooncake Transfer Engine (vllm-project/vllm-omni#3635)
- Mooncake #2098 — [RFC] Agent-Aware KV Cache Support in Mooncake (Phase 1)
- LightX2V disaggregated deployment on Mooncake (ModelTC/LightX2V#893)
- SGLang EPD Disaggregation with Mooncake transfer backend; SGLang Encoder Global Cache Manager (sgl-project/sglang#16137)
- [RFC] SGLang Diffusion Disaggregation ([sgl-project/sglang#19512](https://github.com/sgl-project/sglang/issues/19512)) — engine-internal design with the same stage decomposition; this RFC adopts its credit-based dispatch and request-level asynchrony

### Before submitting a new issue...

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

No implementation file or test is named; start by reviewing the existing integrations under mooncake-integration and the Transfer Engine and Store APIs used by current connectors. The proposed DiffusionConnector, queue orchestration, chunked transfer, and caching need an agreed milestone and validation plan with maintainers before implementation can begin.

Written by the indexing model from the issue text.