vllm-omni/docs/design/ar_diffusion_pipeline_capability.md

AR-Diffusion pipeline capability

Status: experimental and non-normative. This note documents the current experimental/ar_diffusion contract. It follows the ownership and lifecycle direction in #5137, but it is not yet part of the stable module-design hierarchy.

ARDiffusionEngine is an opt-in diffusion backend for pipelines that keep autoregressive attention KV across requests. The engine only selects ARDiffusionModelRunner; the runner owns paged KV pools, session lookup, LRU eviction, and failure cleanup.

A pipeline opts in by implementing SupportsARDiffusionPipeline. Its immutable ARDiffusionKVCacheSpec reports worker-local KV geometry and policy:

• self-attention layers, TP-local KV heads, head size, and tokens per latent frame;
• latent frames committed by one model block, sliding window, and attention sink;
• logical KV branches mapped to worker-local storage slots;
• named fixed-length cross-attention KV, model-specific scratch geometry, and the pipeline's requested session-capacity upper bound.

The runner binds an ARDiffusionKVState only around one forward() call. The pipeline may use it during that context but must not retain it. The session survives subsequent requests with the same session_id. Request reset, explicit close, LRU eviction, and forward exceptions all release runner-owned KV and call the pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily per session and logical KV branch, and is released on the same path. A pipeline publishes one named cache transactionally by yielding exactly one (k, v) pair per layer:

state.populate_cross_attention("main", "text", project_text_kv_by_layer())

The cache becomes readable only after the iterable completes successfully. Logical branches remain independent even when they share a worker-local local_index; that index only maps lazily executed self-attention and scratch work to worker-local capacity slots. It does not define cross-attention cache ownership.

Requests select the session with extra_args["session_id"]; setting reset=True replaces it before the forward, while close_session=True releases it after a successful forward. The runner also exposes reset_session() and close_session() for an owning serving layer.

Current execution limits

The experimental runner supports request-mode execution with max_num_seqs=1. It rejects step execution and request-batch execution at load time, and its inherited batch/step entry points also fail explicitly. Supporting those paths requires a batch-aware state-binding contract rather than silently bypassing bind_ar_diffusion_state().

The current implementation retains one session per model instance. A request for a different session_id closes the resident session through the same LRU lifecycle used by explicit close and failure cleanup before allocating the new one. Self-attention and cross-attention memory are therefore sized for one resident session even if the pipeline declares a larger session_capacity. This prioritizes rollout throughput and prevents block-pool exhaustion before a count-based eviction can run. Multi-session residency can be added later with a memory-backed capacity calculation or allocation-pressure eviction.

Scratch storage is also capability-driven. frames_per_block reserves space for an uncommitted current video block, while max_scratch_tokens_per_branch declares the maximum model-specific tokens, such as action/state registers, that must coexist with it. AR_DIFFUSION_KV_SCRATCH_BLOCKS_PER_BRANCH may increase that derived minimum for deployment experiments but cannot reduce it.

DreamZero adapter

DreamZero implements the capability directly on DreamZeroPipeline. Its adapter reads DreamZero transformer geometry, describes the positive and negative CFG branches, binds the state with a context manager, and translates the existing boolean CFG calls to the explicit "positive" and "negative" KV branch names. Its optional ar_diffusion_warmup_requests() provider owns the robot observation schema and the one-frame/four-frame warmup convention; none of those details are present in the generic runner.

Code migrating from the former experimental KV state API should use the mechanical mapping below:

# Before: state.pos, state.neg, state.get_kv_caches(is_negative, ...)
positive = state.adapter("positive")
negative = state.adapter("negative")
contexts = state.get_kv_caches("negative", seq_len=frame_tokens, commit_current=True)

Direct experimental ARDiffusionKVCache construction also migrates mechanically: replace local_branches with explicit kv_branches, num_frame_per_block with frames_per_block, and the fixed text/image length arguments with a named cross_attention_lengths mapping plus session_capacity. Normal engine users do not construct this object directly.

Cross-attention writers migrate from incremental publication:

# Before
for layer_index, (key, value) in enumerate(project_text_kv_by_layer()):
    state.write_cross_attention_kv("text", layer_index, branch, key, value)
state.mark_cross_attention_populated(branch, "text")

# After
state.populate_cross_attention(branch, "text", project_text_kv_by_layer())

The iterable is consumed once and must yield exactly num_layers pairs. Failed or incomplete projection leaves the previous complete cache unchanged, or leaves the cache absent on its first population.

Single-KV-branch pipeline sketch

A future single-branch causal DMD pipeline can opt in without changes to the runner or KV core:

class WorldPipeline:
    def ar_diffusion_kv_cache_spec(self):
        return ARDiffusionKVCacheSpec(
            num_layers=self.transformer.num_layers,
            num_kv_heads=self.transformer.tp_local_kv_heads,
            head_size=self.transformer.head_size,
            tokens_per_frame=self.transformer.tokens_per_latent_frame,
            frames_per_block=3,
            window_frames=self.transformer.sliding_window_frames,
            sink_frames=self.transformer.sink_frames,
            kv_branches=(ARDiffusionKVBranchSpec("main", 0),),
            cross_attention=(ARDiffusionCrossAttentionKVSpec("text", self.text_length),),
            max_scratch_tokens_per_branch=self.transformer.max_register_tokens,
            session_capacity=1,
        )

    @contextmanager
    def bind_ar_diffusion_state(self, session_id, state):
        self.current_ar_state = state
        try:
            yield
        finally:
            self.current_ar_state = None

    def reset_ar_diffusion_session(self, session_id):
        self.model_states.pop(session_id, None)

    def close_ar_diffusion_session(self, session_id):
        self.model_states.pop(session_id, None)

The model writes/reads named cross-attention KV through the bound state and uses the "main" KV branch for paged self-attention. Pipelines without a warmup provider are loaded normally and skip AR rollout warmup.

Relationship to the multi-stage KV manager RFC

#5244 describes a future multi-stage KV transfer manager, connector integration, and cache-aware scheduling. This capability has a narrower scope: it owns model-local paged KV and session lifecycle inside one AR-Diffusion runner. It does not define cross-stage transfer, connector behavior, or scheduler-visible cache affinity.

When the experimental runtime migrates toward #5244, the implementation should reconcile its current full KVCacheManager use with the RFC's proposed BlockPool-oriented DiT manager. The pipeline capability should remain the model-facing geometry and state-binding boundary where possible, while the experimental ARDiffusionKVCache class itself should not be treated as a stable public API.