223 lines
8.2 KiB
Python
223 lines
8.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for the ARDiffusionKVState paged-attention pipeline bridge."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache import (
|
|
ARDiffusionKVCache,
|
|
ARDiffusionKVConfig,
|
|
ARDiffusionPagedLayerContext,
|
|
)
|
|
from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState
|
|
|
|
BLOCK = 16
|
|
N_HEADS = 4
|
|
HEAD_DIM = 64
|
|
|
|
|
|
def make_state(num_layers=1, window_chunks=4, cross_attn_length=0):
|
|
cfg = ARDiffusionKVConfig(enable=True, chunk_size=BLOCK, window_chunks=window_chunks)
|
|
kv = ARDiffusionKVCache(
|
|
cfg,
|
|
num_layers=num_layers,
|
|
num_kv_heads=N_HEADS,
|
|
head_size=HEAD_DIM,
|
|
dtype=torch.float32,
|
|
block_size=BLOCK,
|
|
max_model_len=4096,
|
|
available_bytes=1 << 24,
|
|
cross_attn_length=cross_attn_length,
|
|
device=torch.device("cpu"),
|
|
)
|
|
pos = kv.begin_request("r-pos")
|
|
neg = kv.begin_request("r-neg")
|
|
return kv, ARDiffusionKVState(kv, pos, neg, num_layers=num_layers)
|
|
|
|
|
|
def _prepare_and_commit(st: ARDiffusionKVState, is_negative: bool, n_chunks: int) -> None:
|
|
ctx = st.get_kv_caches(is_negative, seq_len=n_chunks * BLOCK, commit_current=True)[0].forward_ctx
|
|
ctx.ensure_video_slots(torch.device("cpu"))
|
|
st.commit_paged_context(is_negative)
|
|
|
|
|
|
def test_get_returns_paged_layer_contexts_when_nothing_committed():
|
|
_, st = make_state(num_layers=2)
|
|
for neg in (False, True):
|
|
contexts = st.get_kv_caches(neg, seq_len=BLOCK, commit_current=False)
|
|
assert len(contexts) == 2
|
|
assert all(isinstance(ctx, ARDiffusionPagedLayerContext) for ctx in contexts)
|
|
assert contexts[0].history_block_ids == []
|
|
|
|
|
|
def test_managed_paged_context_commits_only_after_forward():
|
|
kv, st = make_state(num_layers=1)
|
|
contexts = st.get_kv_caches(False, seq_len=2 * BLOCK, commit_current=True)
|
|
ctx = contexts[0].forward_ctx
|
|
|
|
assert st.pos.completed_chunks == 0
|
|
assert kv.window_block_ids(st.pos) == []
|
|
ctx.ensure_video_slots(torch.device("cpu"))
|
|
assert st.pos.completed_chunks == 0
|
|
assert len(ctx.current_video_block_ids) == 2
|
|
|
|
st.commit_paged_context(False)
|
|
|
|
assert st.pos.completed_chunks == 2
|
|
assert st._committed[False] == 2 * BLOCK
|
|
assert len(kv.window_block_ids(st.pos)) == 2
|
|
|
|
|
|
def test_scratch_paged_context_does_not_commit():
|
|
kv, st = make_state(num_layers=1)
|
|
contexts = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)
|
|
ctx = contexts[0].forward_ctx
|
|
ctx.ensure_video_slots(torch.device("cpu"))
|
|
|
|
assert ctx.current_video_block_ids == kv.scratch_block_ids(False, 0, 1)
|
|
st.commit_paged_context(False)
|
|
|
|
assert st.pos.completed_chunks == 0
|
|
assert st._committed[False] == 0
|
|
assert kv.window_block_ids(st.pos) == []
|
|
|
|
|
|
def test_eviction_bounds_resident_window():
|
|
kv, st = make_state(num_layers=1, window_chunks=3)
|
|
_prepare_and_commit(st, False, 3)
|
|
assert len(kv.window_block_ids(st.pos)) == 3
|
|
|
|
_prepare_and_commit(st, False, 2)
|
|
ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=False)[0].forward_ctx
|
|
visible_blocks, video_len = ctx.video_block_table(torch.device("cpu"))
|
|
assert len(visible_blocks) == 3
|
|
assert video_len == 3 * BLOCK
|
|
|
|
|
|
def test_branches_are_independent():
|
|
kv, st = make_state()
|
|
_prepare_and_commit(st, False, 1)
|
|
assert len(kv.window_block_ids(st.pos)) == 1
|
|
assert kv.window_block_ids(st.neg) == []
|
|
|
|
|
|
def test_reset_clears_session_window():
|
|
kv, st = make_state(num_layers=1, window_chunks=4)
|
|
_prepare_and_commit(st, False, 2)
|
|
_prepare_and_commit(st, True, 1)
|
|
assert st._committed[False] == 2 * BLOCK and st._committed[True] == BLOCK
|
|
free_before = kv.manager.block_pool.get_num_free_blocks()
|
|
|
|
st.reset()
|
|
|
|
assert st._committed == {False: 0, True: 0}
|
|
assert st._paged_pending == {False: None, True: None}
|
|
assert kv.window_block_ids(st.pos) == []
|
|
assert kv.window_block_ids(st.neg) == []
|
|
assert kv.manager.block_pool.get_num_free_blocks() > free_before
|
|
_prepare_and_commit(st, False, 1)
|
|
assert len(kv.window_block_ids(st.pos)) == 1
|
|
|
|
|
|
def test_get_cross_kv_caches_raises_before_populate():
|
|
"""Engine ownership guard: a cross-attn read before _kv_populate_cross must fail loud."""
|
|
_, st = make_state() # cross_attn_length == 0, nothing populated
|
|
for neg in (False, True):
|
|
with pytest.raises(RuntimeError, match="cross-attn read before"):
|
|
st.get_cross_kv_caches(neg)
|
|
# Guards the full AND: the populated flag alone is not enough without a pool.
|
|
st._cross_text_populated[False] = True
|
|
with pytest.raises(RuntimeError, match="cross-attn read before"):
|
|
st.get_cross_kv_caches(False)
|
|
|
|
|
|
def test_get_cross_kv_caches_returns_pool_dicts_when_populated():
|
|
L = 8
|
|
kv, st = make_state(num_layers=2, cross_attn_length=L)
|
|
written = []
|
|
for i in range(2):
|
|
k = torch.randn(1, L, N_HEADS, HEAD_DIM)
|
|
v = torch.randn(1, L, N_HEADS, HEAD_DIM)
|
|
kv.write_cross_kv(i, False, k, v)
|
|
written.append((k, v))
|
|
st._cross_text_populated[False] = True
|
|
|
|
out = st.get_cross_kv_caches(False)
|
|
assert len(out) == 2
|
|
for i, (k, v) in enumerate(written):
|
|
assert out[i]["is_init"] is True
|
|
assert out[i]["k"].shape == (1, L, N_HEADS, HEAD_DIM)
|
|
assert torch.equal(out[i]["k"], k)
|
|
assert torch.equal(out[i]["v"], v)
|
|
|
|
|
|
def test_kv_create_is_noop_engine_owns_allocation():
|
|
from unittest.mock import MagicMock
|
|
|
|
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import DreamZeroPipeline
|
|
|
|
p = DreamZeroPipeline.__new__(DreamZeroPipeline)
|
|
p._ar_diffusion_kv_state = object()
|
|
state = MagicMock()
|
|
p._kv_create(state, 1, "float32", "cpu", 24, 4, 64)
|
|
state.create_kv_caches.assert_not_called()
|
|
|
|
|
|
def test_get_kv_cache_requires_frame_aligned_seqlen():
|
|
_, st = make_state()
|
|
with pytest.raises(AssertionError, match="frame-aligned"):
|
|
st.get_kv_caches(False, seq_len=BLOCK + 1, commit_current=True)
|
|
|
|
|
|
def test_prepare_one_branch_allocates_nothing_for_the_other():
|
|
"""CFG-parallel laziness: a rank prepares only the branch it runs."""
|
|
kv, st = make_state(num_layers=1, window_chunks=4)
|
|
pos_ctx = st.get_kv_caches(False, seq_len=BLOCK, commit_current=True)[0].forward_ctx
|
|
neg_ctx = st.get_kv_caches(True, seq_len=BLOCK, commit_current=True)[0].forward_ctx
|
|
|
|
pos_ctx.prepare(device=torch.device("cpu"), action_len=0, query_len=BLOCK)
|
|
|
|
assert pos_ctx._allocated_video
|
|
assert not neg_ctx._allocated_video
|
|
assert kv.window_block_ids(st.neg) == []
|
|
|
|
|
|
def test_failed_forward_tears_down_session(monkeypatch):
|
|
"""Review follow-up (hsliuustc0106): a forward that dies partway must not
|
|
leave allocated-but-uncommitted KV behind — the session is torn down (pool
|
|
blocks freed, model-local state dropped) and the error propagates."""
|
|
from types import SimpleNamespace
|
|
|
|
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
|
from vllm_omni.experimental.ar_diffusion.runner import ARDiffusionModelRunner
|
|
|
|
kv, st = make_state(num_layers=1, window_chunks=4)
|
|
_prepare_and_commit(st, False, 2) # session owns pool blocks
|
|
free_before_failure = kv.manager.block_pool.get_num_free_blocks()
|
|
|
|
runner = object.__new__(ARDiffusionModelRunner)
|
|
runner.kv_cache = kv
|
|
runner._ar_diffusion_states = __import__("collections").OrderedDict({"s1": st})
|
|
runner._max_ar_diffusion_states = 4
|
|
runner.pipeline = SimpleNamespace(_states={"s1": object()}, _ar_diffusion_kv_state=None)
|
|
runner.device = None
|
|
runner._perf_e2e_times = []
|
|
|
|
def boom(self, req):
|
|
raise RuntimeError("layer 17 exploded")
|
|
|
|
monkeypatch.setattr(DiffusionModelRunner, "execute_model", boom)
|
|
req = SimpleNamespace(
|
|
request_id="r0",
|
|
sampling_params=SimpleNamespace(extra_args={"session_id": "s1"}),
|
|
)
|
|
|
|
with pytest.raises(RuntimeError, match="layer 17 exploded"):
|
|
ARDiffusionModelRunner.execute_model(runner, req)
|
|
|
|
assert "s1" not in runner._ar_diffusion_states
|
|
assert "s1" not in runner.pipeline._states
|
|
assert runner.pipeline._ar_diffusion_kv_state is None
|
|
# close() freed the session's resident blocks.
|
|
assert kv.manager.block_pool.get_num_free_blocks() > free_before_failure
|