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](https://github.com/vllm-project/vllm-omni/issues/5137),
> but it is not yet part of the stable module-design hierarchy.
>
> The end-to-end request, identity, failure, and memory contracts are documented
> in [Realtime AR-Diffusion sessions](feature/realtime_ar_diffusion.md).

`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 runner invocation
(`execute_model` or `execute_stepwise`). The pipeline may use it during that
context but must not retain it. Scratch and committed pages stay on the
runner-owned session object, so unbinding does not drop an in-flight stepwise
chunk. Tick-mode sessions survive subsequent requests with the same
`session_id`. On the stepwise path the session id is the request id and is
released when that request finishes or fails. 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:

```python
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 `max_num_seqs=1`. Request-batch execution is
still rejected. Step execution is allowed only when the loaded pipeline
implements `SupportsStepExecution`; other AR-Diffusion pipelines still fail at
load if `step_execution=True`. Capable pipelines bind KV around each
`execute_stepwise` invocation instead of inheriting the generic runner entry
point unwrapped.

The implementation may retain multiple sessions per model instance. It selects
the largest feasible capacity up to the pipeline's `session_capacity`, based on
self-attention pages, cross-attention allocations, scratch, and declared
model-owned state. One session is admitted whenever it fits actual free device
memory; `gpu_memory_fraction` controls additional residency. Requests beyond
that capacity evict the least-recently-used session through the same lifecycle
hook used by explicit close and failure cleanup.

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:

```python
# 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:

```python
# 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:

```python
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](https://github.com/vllm-project/vllm-omni/issues/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.