Mooncake EPD: Multimodal and Agent-State Disaggregation

[Pinoeer-kingxi]

Mooncake EPD: Multimodal and Agent-State Disaggregation

Summary

This integration PR adds Encoder-Prefill-Decode disaggregation to Mooncake
and vLLM. It adds real Vision Hidden State transfer, layered paged-KV handoff,
multimodal hidden-state reuse, Agent KV state cloning, state-aware scheduling,
backpressure, and artifact-gated evaluation.

The implementation is intentionally submitted as one cohesive PR because the
data-plane fast paths, ownership protocol, vLLM hooks, scheduler, and strict
evaluation gates share end-to-end contracts. The commits and sections follow
existing Mooncake ownership boundaries so each subsystem can still be reviewed
independently within the same PR.

Motivation

Multimodal and Agent workloads expose reusable state that is not represented
by a conventional colocated request lifecycle:

• Vision Encoder output can be reused across requests and Agent steps.
• Prefill KV cache can be transferred to independently scaled Decode workers.
• Parallel Agent branches can share immutable KV pages and materialize only
  modified pages.
• Thinking, interactive, and hybrid tasks benefit from different Prefill,
  Decode, priority, and latency policies.

Mooncake provides the transport and shared-state foundation needed to make
these states first-class serving objects.

Task coverage

Foundation

• vLLM EPD prototype: Vision Hidden State moves from Encoder to Prefill,
  then paged KV cache moves from Prefill to Decode.
• Agent State Cloning: branches share page references with Copy-on-Write
  and explicit lifecycle management.
• Qwen-VL end-to-end demo: real-model launchers, colocated baseline,
  scale-out runner, metrics, and strict direct-path gates.

Advanced

• Worker-level transfer primitives: registered pointers, peer buffers,
  layer grouping, topology affinity, and transport telemetry.
• Agent PD scheduling: task-type pools, deadline/rho-aware admission,
  backpressure, and affinity-aware dispatch.
• Hidden State Prefix Caching: stable keys, event prefetch, FeatureHandles,
  vLLM hidden-state injection, and Vision Encoder skip.
• Upstream contribution package: public design, evaluation, tests, and a
  subsystem-oriented review structure within one integration PR.

Architecture

flowchart LR
    Client --> Proxy[EPD Proxy and Scheduler]
    Proxy --> Encoder[Vision Encoder]
    Encoder -->|FeatureHandle and peer buffer| Prefill[vLLM Prefill]
    Prefill -->|Grouped layered KV| Decode[vLLM Decode]
    Decode --> Client

    Proxy <--> Registry[Workflow Registry and WAL]
    Proxy <--> Directory[Owner-sharded KV Directory]
    Registry <--> Agent[Agent Clone and CoW]
    Encoder <--> MMCache[MM Hidden State Cache]

Loading

See DESIGN.md for interfaces,
lifecycles, data flow, and tradeoffs.

Main changes

Mooncake Transfer Engine

• Add registered-pointer batch transfer for persistent KV regions.
• Preserve registration ownership when memory is already registered.
• Improve intra-node CUDA stream dependency handling.
• Expose prepare, write, registration, and bandwidth telemetry.

Encoder-to-Prefill FeatureHandle

• Add Prefill-owned direct feature buffers with readiness, reference counting,
  bounded persistent cache, and explicit release.
• Add validated FeatureBundle descriptors for Qwen-VL hidden state, DeepStack
  intermediates, and grid_thw.
• Add vLLM-side epd-direct:// materialization and precomputed
  image_embeds injection.
• Add event-driven cache prefetch and concurrent request coalescing.

Prefill-to-Decode layered KV

• Connect grouped layer descriptors to the vLLM producer/consumer save/load
  path.
• Add direct peer-buffer transfer, topology affinity, receive progress, and
  strict transfer counters.
• Preserve handoff identity across prepare, commit, rollback, and release.

Agent state and scheduling

• Add page-reference clone and page-level Copy-on-Write semantics.
• Add owner-sharded directory state and workflow registry/WAL transitions.
• Add thinking, interactive, and hybrid task routing.
• Add queue, service-rate, deadline, rho, backpressure, and rejection metrics.

Benchmarks and documentation

• Add real Qwen-VL EPD, colocated baseline, 2P2D scale-out, Agent clone, and
  evaluation-matrix runners.
• Add fail-closed gates for real services, direct transfer, cache reuse,
  resource scope, output evidence, and lifecycle cleanup.
• Add public README.md, DESIGN.md, EVALUATION.md, submission material,
  and a compact evidence snapshot.

Compatibility and security

• Existing Mooncake and non-EPD vLLM startup remains unchanged.
• Workspace-scoped vLLM patches are enabled only when generated EPD workers set
  MOONCAKE_EPD_ENABLE_VLLM_PATCHES=1.
• Direct FeatureBuffer routes require a generated deployment token through
  X-Mooncake-EPD-Token.
• Strict serving confines file-backed FeatureHandles to configured store roots.
• Public interfaces retain existing defaults; EPD behavior is activated by
  explicit launcher and connector configuration.
• TCP and same-host nvlink_intra are covered by public measurements. SHM and
  RDMA remain selectable through the transport policy for compatible hosts.

Validation

From the Mooncake repository root:

export PYTHONPATH=$PWD

PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
MOONCAKE_EPD_ENABLE_VLLM_PATCHES=0 \
python -m pytest -q mooncake_epd/tests

PYTHONPATH=$PWD/mooncake-wheel \
PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
python -m pytest -q mooncake-wheel/tests/test_mooncake_store_service_api.py

Verification	Result
EPD unit and integration suite	371 passed, 1 skipped
Mooncake Store service API	72 passed
Runtime/security boundary subset	52 passed
Compile, JSON, whitespace checks	Pass
Strict real-EPD artifact gate	Pass

Benchmark summary

Real multimodal EPD

Qwen3-VL-8B-Instruct, 1 Encoder + 1 Prefill + 2 Decode GPUs, 8 warmups,
16 measured requests, concurrency 8:

Metric	Result
Request throughput	8.143 RPS
Output throughput	232.08 tok/s
Mean / P95 TTFT	288.82 / 344.01 ms
HTTP and EPD route success	16/16
P-to-D peer-buffer traffic	46 batches / 301,989,888 bytes
Fallback / send / receive failures	0 / 0 / 0

Scale-out capacity

The recorded comparison uses four GPUs for 2P2D EPD and one GPU for the
colocated baseline. It demonstrates scale-out capacity, not equal-resource
efficiency.

Metric	1-GPU colocated	4-GPU 2P2D EPD	Change
Request throughput	2.580 RPS	3.685 RPS	+42.8%
Output throughput	74.83 tok/s	108.24 tok/s	+44.6%
Mean latency	2917.69 ms	2143.09 ms	-26.5%
Mean TTFT	1271.71 ms	1413.93 ms	+11.2%
P95 TTFT	2006.39 ms	1481.26 ms	-26.2%

Agent State Cloning

• 10.03x mean clone speedup over deep copy.
• 0 bytes copied for page-reference branch creation versus 67,108,864 bytes
  for deep copy.
• 0 retained states and 0 orphan directory blocks after release.

Full methodology and interpretation are in
EVALUATION.md. The compact evidence
record is
artifacts/2026-07-10/benchmark_summary.json.

Review guide

This is the primary integration PR. Review can proceed by subsystem without
splitting the end-to-end contract:

1. Transfer Engine / intra-node transport
   • Registered-pointer fast path, synchronization, registration ownership,
     and telemetry.
2. Encoder-to-Prefill FeatureHandle
   • Direct feature buffers, descriptors, hidden-state injection, and cache.
3. Serving and Agent control plane
   • Layered KV connector, scheduler, workflow state, directory, handoff, and
     backpressure.
4. Benchmarks and documentation
   • Dataset builder, runners, gates, tests, public reports, and evidence.

Small follow-up PRs may extract narrowly scoped transport or framework adapters
after the integrated interfaces are agreed, without changing the public EPD
state and lifecycle semantics established here.

Checklist

• Integrated into the Mooncake repository layout.
• Preserves default behavior unless EPD is explicitly enabled.
• Includes error handling, authorization, lifecycle, and cleanup logic.
• Includes unit, integration, Store API, and real-serving validation.
• Includes baseline and optimized results with resource scope.
• Includes public run instructions using repository-relative paths.
• Includes one cohesive integration PR with subsystem-oriented commits.

[gemini-code-assist]

Code Review

This pull request introduces a sophisticated Encoder-Prefill-Decode (EPD) disaggregation framework for multimodal LLMs and Agent workflows. Key features include an event-driven multimodal store with asynchronous prefetch, a three-tier KV reuse pipeline, and a distributed KV directory for cross-shard state management. My review identified several critical issues, including a memory leak in the MMStore handle management, a race condition in the page manager's lease mechanism, and potential runtime crashes due to missing defensive checks in state advancement and configuration loading. I have provided actionable suggestions to address these stability and performance concerns.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Storing all historical MMStoreHandle objects in self._handles without any eviction or cleanup mechanism introduces a severe memory leak on long-running servers. Since the caller receives the MMStoreHandle directly and can wait on it asynchronously, self._handles is not needed for coordination. We can completely eliminate self._handles by storing the MMStoreHandle objects directly in self._waiters_by_event instead of string UUIDs, allowing them to be garbage collected naturally once resolved.

Suggested change

	self.queue_put_timeout_s = max(0.0, float(queue_put_timeout_s))
	self.inline_fallback_on_queue_full = (
	self._waiters_by_event: Dict[str, List[MMStoreHandle]] = {}

[gemini-code-assist]

In decref, when new_refcount <= 0, the physical page is deleted from self._pages immediately, even if lease_count > 0. This breaks the lease mechanism: leased pages are meant to be kept alive during asynchronous offloading or handoff. If the physical tensor is deleted from self._pages while leased, any concurrent read/write or restore operation on that page will fail with a KeyError. The physical page should only be removed from self._pages when the directory record is fully dropped (i.e., both refcount and lease_count are 0, and physical_released is True).

Suggested change

	if new_refcount <= 0:
	del self._pages[physical_id]
	self.kv_directory.release_physical(gid)
	return True
	if new_refcount <= 0:
	self.kv_directory.release_physical(gid)
	if self.kv_directory.get_record(gid) is None:
	self._pages.pop(physical_id, None)
	return True
	return False

[gemini-code-assist]

In handle_put_bytes, request.read() returns a bytes object (or b'' if empty), never None. Therefore, the check value is None will always evaluate to False, failing to catch empty request bodies. Using if not value: is more robust as it handles both None and empty bytes (b'').

Suggested change

	value = await request.read()
	if value is None:
	value = await request.read()
	if not value:

[gemini-code-assist]

The advance method relies heavily on self.sl (the StateLayer instance) to register and manage states. However, state_layer is optional in __init__ and defaults to None. If self.sl is None, calling advance will result in an unhandled AttributeError at line 260. Adding a defensive guard at the beginning of advance to raise a clear ValueError when self.sl is None will prevent unexpected crashes.

Suggested change

	def advance(
	self,
	new_token_ids: Sequence[int],
	new_pixel_values: Optional[torch.Tensor] = None,
	new_image_grid_thw: Optional[torch.Tensor] = None,
	new_image_hashes: Optional[Sequence[str]] = None,
	divergence_threshold: float = 0.6,
	) -> WorkflowStep:
	def advance(
	self,
	new_token_ids: Sequence[int],
	new_pixel_values: Optional[torch.Tensor] = None,
	new_image_grid_thw: Optional[torch.Tensor] = None,
	new_image_hashes: Optional[Sequence[str]] = None,
	divergence_threshold: float = 0.6,
	) -> WorkflowStep:
	"""Run the next step, reusing KV from the predecessor where possible."""
	if self.sl is None:
	raise ValueError("state_layer must be provided to advance the workflow")

[gemini-code-assist]

The tier-3 fallback similarity evaluation (_apply_tier3_fallback) is an optimization rather than a strict correctness dependency. If any unexpected error occurs during the tensor slicing or similarity computation (e.g., shape mismatches or out-of-bounds indices), it is safer to wrap this call in a try...except block to log a warning and fall back gracefully to the unoptimized relay_refs instead of crashing the entire serving request.

Suggested change

	relay_refs = self._apply_tier3_fallback(
	relay_refs=relay_refs,
	prev_refs=prev_refs,
	tier3_candidates=tier3_candidates,
	stats=stats,
	)
	# -- Tier 3: attention-similarity fallback on recomputed reusable pages ---
	try:
	relay_refs = self._apply_tier3_fallback(
	relay_refs=relay_refs,
	prev_refs=prev_refs,
	tier3_candidates=tier3_candidates,
	stats=stats,
	)
	except Exception as exc:
	# Fallback gracefully to unoptimized refs on any unexpected error
	pass

[gemini-code-assist]

In _python_store, if setup(config) raises a TypeError (indicating an older version of the SDK), the fallback positional setup is called. However, if the JSON config file loaded from config_path is missing any of the expected keys (like local_hostname or protocol), accessing them directly via config["key"] will raise a KeyError and crash. Using config.get("key", default) is much safer and prevents unexpected crashes.

Suggested change

	rc = setup(config)
	except TypeError:
	rc = setup(
	config["local_hostname"],
	config["metadata_server"],
	int(config["global_segment_size"]),
	int(config["local_buffer_size"]),
	config["protocol"],
	config.get("device_name", ""),
	config["master_server_address"],
	)
	try:
	rc = setup(config)
	except TypeError:
	rc = setup(
	config.get("local_hostname", "127.0.0.1"),
	config.get("metadata_server", ""),
	int(config.get("global_segment_size", 16 * 1024 * 1024)),
	int(config.get("local_buffer_size", 16 * 1024 * 1024)),
	config.get("protocol", "tcp"),
	config.get("device_name", ""),
	config.get("master_server_address", ""),
	)

[gemini-code-assist]

The greedy longest-common-substring matching algorithm in split_segments has a worst-case time complexity of $O(N \times M)$ where $N$ is the length of new_tokens and $M$ is the length of prev_tokens. On long sequences with highly repetitive tokens (e.g., common punctuation, padding, or frequent words), prev_index.get(int(new_tokens[i])) can return a very large list of indices, causing the nested while loop to execute frequently and severely degrade performance. Consider optimizing this matching logic or limiting the search depth to avoid quadratic complexity on long sequences.