184 lines
8.6 KiB
Python
184 lines
8.6 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""DreamZero-facing AR-Diffusion KV bridge backed by paged self-attention."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import (
|
|
ARDiffusionPagedForwardContext,
|
|
ARDiffusionPagedLayerContext,
|
|
)
|
|
|
|
_log = init_logger(__name__)
|
|
|
|
|
|
class ARDiffusionKVState:
|
|
"""Per-request bridge between a DreamZero rollout and the AR-Diffusion KV pool.
|
|
|
|
The self-attention path is paged-only: ``get_kv_caches`` returns lightweight
|
|
per-layer contexts, and each attention layer writes current K/V into the
|
|
pool before calling the FlashAttention block-table kernel. No resident
|
|
window is gathered into a contiguous tensor.
|
|
"""
|
|
|
|
def __init__(self, kv_cache, pos_adapter, neg_adapter, num_layers: int) -> None:
|
|
"""One state per DreamZero session, holding BOTH CFG branches' adapters.
|
|
|
|
``pos``/``neg`` are the classifier-free-guidance branches: ``pos`` is
|
|
the conditional pass (real prompt), ``neg`` the unconditional pass
|
|
(negative prompt). Each owns an independent request adapter — and
|
|
therefore independent pool blocks — because the two branches produce
|
|
different K/V for the same frames. Every ``is_negative`` flag in this
|
|
module selects between them (``False`` -> ``pos``, ``True`` -> ``neg``).
|
|
"""
|
|
self.kv_cache = kv_cache
|
|
self.pos = pos_adapter
|
|
self.neg = neg_adapter
|
|
self.num_layers = num_layers
|
|
self._committed: dict[bool, int] = {False: 0, True: 0}
|
|
self._paged_pending: dict[bool, ARDiffusionPagedForwardContext | None] = {False: None, True: None}
|
|
# Cross-attn KV is populated once after text encoding (eager), then
|
|
# read every denoising step. Tracked per half: the text K/V depend only
|
|
# on the session prompt (survive window-boundary resets), while the I2V
|
|
# image-token K/V come from the current observation's CLIP features and
|
|
# must be re-projected on every window restart.
|
|
self._cross_text_populated: dict[bool, bool] = {False: False, True: False}
|
|
self._cross_img_populated: dict[bool, bool] = {False: False, True: False}
|
|
|
|
def _adapter(self, is_negative: bool):
|
|
return self.neg if is_negative else self.pos
|
|
|
|
def get_kv_caches(
|
|
self,
|
|
is_negative: bool,
|
|
seq_len: int | None = None,
|
|
commit_current: bool = False,
|
|
) -> list[ARDiffusionPagedLayerContext]:
|
|
if seq_len is None:
|
|
raise ValueError("AR-Diffusion paged self-attention requires seq_len in get_kv_caches()")
|
|
return self.prepare_paged_context(is_negative, seq_len, commit_current)
|
|
|
|
def prepare_paged_context(
|
|
self,
|
|
is_negative: bool,
|
|
seq_len: int,
|
|
commit_current: bool,
|
|
) -> list[ARDiffusionPagedLayerContext]:
|
|
"""Return per-layer paged attention contexts for one branch forward.
|
|
|
|
Allocation is lazy: constructing both CFG branches' kwargs must not
|
|
allocate manager blocks for a branch that a CFG-parallel rank will not run.
|
|
The first self-attn layer that consumes the context allocates the current
|
|
video slots, then all layers share that branch-level state.
|
|
"""
|
|
cs = self.kv_cache.spec.chunk_size
|
|
if int(seq_len) % cs != 0:
|
|
raise AssertionError(
|
|
f"AR-Diffusion expects frame-aligned seq_len (multiple of chunk_size={cs}), got {seq_len}"
|
|
)
|
|
|
|
pending = self._paged_pending.get(is_negative)
|
|
if pending is not None and pending.commit_current and pending._allocated_video and not pending._committed:
|
|
raise RuntimeError("AR-Diffusion paged context replaced before its managed current chunk was committed")
|
|
|
|
adapter = self._adapter(is_negative)
|
|
branch = "neg" if is_negative else "pos"
|
|
forward_ctx = ARDiffusionPagedForwardContext(
|
|
kv_cache=self.kv_cache,
|
|
adapter=adapter,
|
|
is_negative=is_negative,
|
|
history_block_ids=self.kv_cache.window_block_ids(adapter),
|
|
seq_len=int(seq_len),
|
|
commit_current=bool(commit_current),
|
|
max_video_tokens=int(self.kv_cache.spec.sliding_window),
|
|
)
|
|
self._paged_pending[is_negative] = forward_ctx
|
|
_log.debug(
|
|
"AR-Diffusion GET [%s] source=paged-attn layers=%d history_blocks=%d seq_len=%d commit_current=%s",
|
|
branch,
|
|
self.num_layers,
|
|
len(forward_ctx.history_block_ids),
|
|
int(seq_len),
|
|
bool(commit_current),
|
|
)
|
|
return [ARDiffusionPagedLayerContext(layer_idx=i, forward_ctx=forward_ctx) for i in range(self.num_layers)]
|
|
|
|
def commit_paged_context(self, is_negative: bool) -> None:
|
|
"""Commit the managed current video blocks after a successful forward."""
|
|
ctx = self._paged_pending.get(is_negative)
|
|
if ctx is None:
|
|
return
|
|
branch = "neg" if is_negative else "pos"
|
|
if ctx.commit_current and ctx._allocated_video:
|
|
n_chunks = ctx.seq_len // self.kv_cache.spec.chunk_size
|
|
for _ in range(n_chunks):
|
|
ctx.adapter.on_chunk_committed()
|
|
self._committed[is_negative] += ctx.seq_len
|
|
_log.debug(
|
|
"AR-Diffusion COMMIT [%s] paged-attn new_tokens=%d chunks=%d resident=%d/%d",
|
|
branch,
|
|
ctx.seq_len,
|
|
n_chunks,
|
|
len(self.kv_cache.window_block_ids(ctx.adapter)),
|
|
self.kv_cache.spec.window_chunks,
|
|
)
|
|
ctx.mark_committed()
|
|
self._paged_pending[is_negative] = None
|
|
|
|
def get_cross_kv_caches(self, is_negative: bool) -> list[dict]:
|
|
"""Return pool-backed cross-attn cache dicts.
|
|
|
|
Under AR-Diffusion the cross-attn pool is always populated by ``_kv_populate_cross``
|
|
before the first read, so this must never be reached unpopulated — if it
|
|
were, the model would lazily project and write cross KV itself, violating
|
|
engine ownership. Fail loud rather than fall back to the model.
|
|
"""
|
|
img_ok = not self.kv_cache._cross_k_img or self._cross_img_populated[is_negative]
|
|
if not (self._cross_text_populated[is_negative] and img_ok and self.kv_cache.cross_attn_length > 0):
|
|
raise RuntimeError(
|
|
f"AR-Diffusion cross-attn read before _kv_populate_cross (neg={is_negative}, "
|
|
f"cross_attn_length={self.kv_cache.cross_attn_length}, "
|
|
f"text={self._cross_text_populated[is_negative]} "
|
|
f"img={self._cross_img_populated[is_negative]}) — the engine must own all cross KV"
|
|
)
|
|
return [self.kv_cache.read_cross_kv(i, is_negative) for i in range(self.num_layers)]
|
|
|
|
def close(self) -> None:
|
|
"""Final teardown: free both branches' resident pool blocks."""
|
|
self.kv_cache.end_request(self.pos)
|
|
self.kv_cache.end_request(self.neg)
|
|
self._paged_pending = {False: None, True: None}
|
|
|
|
def reset(self, *, keep_cross_text: bool = False) -> None:
|
|
"""Drop this session's KV — mirrors the model-local ``state.reset()``.
|
|
|
|
DreamZero resets at the attention-window boundary (``should_reset``); the
|
|
model starts a fresh sliding window, so AR-Diffusion must free the resident pool
|
|
blocks and start a fresh adapter for each branch. The ``ARDiffusionKVState``
|
|
object (and thus ``pipeline._ar_diffusion_kv_state``) is preserved across the reset
|
|
so the runner's session mapping stays valid — only the pool-backed state
|
|
is recycled.
|
|
|
|
Args:
|
|
keep_cross_text: On window ("inference") resets the session prompt is
|
|
unchanged, so the text cross-attn K/V in the static pool buffers are
|
|
still valid — keep them and only force the image half to repopulate.
|
|
Session resets (prompt change / explicit reset) pass ``False``.
|
|
"""
|
|
pos_id, neg_id = self.pos.request_id, self.neg.request_id
|
|
self.close()
|
|
self.pos = self.kv_cache.begin_request(pos_id)
|
|
self.neg = self.kv_cache.begin_request(neg_id)
|
|
self._committed = {False: 0, True: 0}
|
|
self._paged_pending = {False: None, True: None}
|
|
if not keep_cross_text:
|
|
self._cross_text_populated = {False: False, True: False}
|
|
self._cross_img_populated = {False: False, True: False}
|
|
_log.info(
|
|
"AR-Diffusion RESET [%s/%s] session KV cleared (window boundary; cross text %s)",
|
|
pos_id,
|
|
neg_id,
|
|
"kept" if keep_cross_text else "cleared",
|
|
)
|