[Diffusion] Add native paged KV backend for diffusion workers (#6102)
Signed-off-by: Acerak01-fy <wfy2003324@163.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
Wang fuyinandHongsheng Liu authored and GitHub committed
1 parent
c2daa471c3
commit
72ee535fcd
34 files changed
+4595
-71
No files matched your search
@@ -23,8 +23,10 @@ import vllm_omni.diffusion.attention.backends.utils.piecewise_attn as piecewise_
|
||||
from vllm_omni.diffusion.attention.backends.abstract import QueryRange
|
||||
from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
|
||||
Segment,
|
||||
build_paged_piecewise_plan,
|
||||
build_segments,
|
||||
piecewise_attn,
|
||||
run_paged_piecewise_plan,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -84,6 +86,72 @@ def test_build_segments_full_span_retains_global_kv_end():
|
||||
assert segments == [Segment(q_start=4, q_end=6, kv_end=7, mode="full")]
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
full_attn_spans = [[(2, 5)], [(4, 7)]]
|
||||
plan = build_paged_piecewise_plan(
|
||||
full_attn_spans,
|
||||
query_offsets=[0, 2],
|
||||
query_lens=[6, 6],
|
||||
seq_lens=[6, 8],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
key = query + 100
|
||||
value = query + 200
|
||||
calls = []
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, segment_batch):
|
||||
calls.append(
|
||||
(
|
||||
[segment.mode for segment in segment_batch.row_segments],
|
||||
[segment.kv_end for segment in segment_batch.row_segments],
|
||||
segment_batch.query_indices.tolist(),
|
||||
)
|
||||
)
|
||||
torch.testing.assert_close(segment_key, segment_query + 100)
|
||||
torch.testing.assert_close(segment_value, segment_query + 200)
|
||||
return segment_query
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(output, query)
|
||||
assert calls == [
|
||||
(["causal", "causal"], [2, 4], [0, 1, 6, 7]),
|
||||
(["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
|
||||
(["causal", "causal"], [6, 8], [5, 11]),
|
||||
]
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_uses_runner_output_shape():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)]],
|
||||
query_offsets=[0],
|
||||
query_lens=[6],
|
||||
seq_lens=[6],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.arange(6 * 2 * 4, dtype=torch.float32).reshape(6, 2, 4)
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
query,
|
||||
query,
|
||||
plan,
|
||||
plan.segments,
|
||||
lambda segment_query, *_args: segment_query[..., :3],
|
||||
)
|
||||
|
||||
assert output.shape == (6, 2, 3)
|
||||
torch.testing.assert_close(output, query[..., :3])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("global_spans", SPAN_CASES)
|
||||
@pytest.mark.parametrize("q_range", Q_RANGE_CASES)
|
||||
@pytest.mark.parametrize("batch_size", BATCH_CASES)
|
||||
|
||||
@@ -0,0 +1,772 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import replace
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec
|
||||
from vllm.v1.worker import block_table as native_block_table
|
||||
|
||||
from vllm_omni.diffusion.diffusion_kv import model_runner_backend as model_runner_backend_module
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.diffusion_kv.metadata import (
|
||||
DiffusionKVContextMetadata,
|
||||
DiffusionKVMetadata,
|
||||
DiffusionKVSequenceMetadata,
|
||||
)
|
||||
from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
|
||||
from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _registration_backend(parallel_config: object):
|
||||
model_config = SimpleNamespace(set_attention_geometry=Mock())
|
||||
vllm_config = SimpleNamespace(
|
||||
attention_config=SimpleNamespace(use_non_causal=False),
|
||||
compilation_config=SimpleNamespace(static_forward_context={}),
|
||||
model_config=model_config,
|
||||
)
|
||||
backend = DiffusionKVModelRunnerBackend(
|
||||
vllm_config=vllm_config,
|
||||
od_config=SimpleNamespace(parallel_config=parallel_config),
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
return backend, vllm_config, model_config
|
||||
|
||||
|
||||
def _parallel_config(**overrides):
|
||||
values = dict(ulysses_degree=1, ring_degree=1, allgather_degree=1, ulysses_mode="strict")
|
||||
values.update(overrides)
|
||||
return SimpleNamespace(**values)
|
||||
|
||||
|
||||
def test_cache_layer_registration_installs_adapters_with_ulysses_local_geometry(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
backend, vllm_config, model_config = _registration_backend(_parallel_config(ulysses_degree=2))
|
||||
layer = SimpleNamespace(
|
||||
num_heads=8,
|
||||
softmax_scale=0.125,
|
||||
skip_sequence_parallel=False,
|
||||
attn_backend=SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN"),
|
||||
)
|
||||
unrelated = object()
|
||||
vllm_config.compilation_config.static_forward_context.update({"layer-0": layer, "unrelated": unrelated})
|
||||
spec = FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=4,
|
||||
head_size=8,
|
||||
dtype=torch.float16,
|
||||
non_causal=True,
|
||||
)
|
||||
constructor_calls = []
|
||||
|
||||
class FakeLayerAdapter:
|
||||
def __init__(self, *, layer_name, layer, spec, ulysses_degree, **_kwargs) -> None:
|
||||
constructor_calls.append((layer_name, layer, spec, ulysses_degree))
|
||||
self.num_heads = layer.num_heads // ulysses_degree
|
||||
self.num_kv_heads = spec.num_kv_heads // ulysses_degree
|
||||
self.head_size = spec.head_size
|
||||
self.spec = replace(spec, num_kv_heads=self.num_kv_heads)
|
||||
|
||||
monkeypatch.setattr(model_runner_backend_module, "DiffusionPagedAttentionLayerAdapter", FakeLayerAdapter)
|
||||
monkeypatch.setattr(model_runner_backend_module, "set_current_vllm_config", lambda _config: nullcontext())
|
||||
|
||||
specs = backend.register_kv_cache_layers({"layer-0": (layer, spec)})
|
||||
|
||||
installed = vllm_config.compilation_config.static_forward_context["layer-0"]
|
||||
assert isinstance(installed, FakeLayerAdapter)
|
||||
assert vllm_config.compilation_config.static_forward_context["unrelated"] is unrelated
|
||||
assert specs["layer-0"].num_kv_heads == 2
|
||||
assert constructor_calls == [("layer-0", layer, spec, 2)]
|
||||
model_config.set_attention_geometry.assert_called_once_with(
|
||||
num_heads=4,
|
||||
num_kv_heads=2,
|
||||
head_size=8,
|
||||
)
|
||||
assert vllm_config.attention_config.use_non_causal is True
|
||||
|
||||
|
||||
def test_cache_layer_registration_rejects_backend_without_paged_support() -> None:
|
||||
backend, _, _ = _registration_backend(_parallel_config())
|
||||
layer = SimpleNamespace(
|
||||
attn_backend=SimpleNamespace(supports_paged_kv=False, get_name=lambda: "SDPA"),
|
||||
)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=16,
|
||||
num_kv_heads=4,
|
||||
head_size=8,
|
||||
dtype=torch.float16,
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="layer-0.*paged support.*SDPA"):
|
||||
backend.register_kv_cache_layers({"layer-0": (layer, spec)})
|
||||
|
||||
|
||||
def test_cache_layer_registration_resolves_platform_paged_hooks(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
backend, _, _ = _registration_backend(_parallel_config())
|
||||
|
||||
def unsupported_block_tables():
|
||||
raise NotImplementedError("platform has no paged BlockTables")
|
||||
|
||||
monkeypatch.setattr(
|
||||
model_runner_backend_module.current_omni_platform,
|
||||
"get_diffusion_kv_block_tables_cls",
|
||||
unsupported_block_tables,
|
||||
)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="no paged BlockTables"):
|
||||
backend.register_kv_cache_layers({"layer-0": (object(), Mock(spec=FullAttentionSpec))})
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("parallel_config", "message"),
|
||||
[
|
||||
(_parallel_config(ring_degree=2), "Ring"),
|
||||
(_parallel_config(allgather_degree=2), "AllGather-KV"),
|
||||
(_parallel_config(ulysses_degree=2, ulysses_mode="advanced_uaa"), "strict Ulysses"),
|
||||
],
|
||||
)
|
||||
def test_cache_layer_registration_rejects_unsupported_sp_modes(parallel_config, message: str) -> None:
|
||||
backend, _, _ = _registration_backend(parallel_config)
|
||||
|
||||
with pytest.raises(NotImplementedError, match=message):
|
||||
backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=False))
|
||||
|
||||
|
||||
def test_cache_layer_registration_ignores_sp_for_opted_out_layer() -> None:
|
||||
backend, _, _ = _registration_backend(
|
||||
_parallel_config(ulysses_degree=2, ring_degree=2, ulysses_mode="advanced_uaa")
|
||||
)
|
||||
|
||||
assert backend._get_paged_attention_ulysses_degree(SimpleNamespace(skip_sequence_parallel=True)) == 1
|
||||
|
||||
|
||||
class _FakeSpec:
|
||||
def __init__(self, block_size: int) -> None:
|
||||
self.block_size = block_size
|
||||
self.capacity_calls: list[tuple[object, int]] = []
|
||||
|
||||
def max_num_blocks_per_req(self, vllm_config: object, max_len: int) -> int:
|
||||
self.capacity_calls.append((vllm_config, max_len))
|
||||
return math.ceil(max_len / self.block_size)
|
||||
|
||||
|
||||
class _FakeStagedTable:
|
||||
def __init__(self) -> None:
|
||||
self.staged: list[tuple[int, tuple[int, ...]]] = []
|
||||
self.clear_calls = 0
|
||||
|
||||
def clear_staged_writes(self) -> None:
|
||||
self.staged.clear()
|
||||
self.clear_calls += 1
|
||||
|
||||
|
||||
class _FakeNumBlocks:
|
||||
def __init__(self, num_groups: int, max_num_reqs: int) -> None:
|
||||
self.np = np.zeros((num_groups, max_num_reqs), dtype=np.int32)
|
||||
self.copy_calls = 0
|
||||
|
||||
def copy_to_uva(self) -> None:
|
||||
self.copy_calls += 1
|
||||
|
||||
|
||||
class _FakeBlockTables:
|
||||
instances: list[_FakeBlockTables] = []
|
||||
|
||||
def __init__(self, **kwargs: object) -> None:
|
||||
self.kwargs = kwargs
|
||||
self.block_sizes = list(kwargs["block_sizes"])
|
||||
self.max_num_reqs = int(kwargs["max_num_reqs"])
|
||||
self.max_num_batched_tokens = int(kwargs["max_num_batched_tokens"])
|
||||
self.cp_size = int(kwargs["cp_size"])
|
||||
self.cp_rank = int(kwargs["cp_rank"])
|
||||
self.cp_interleave = int(kwargs["cp_interleave"])
|
||||
self.blocks_per_kv_block = [1] * len(self.block_sizes)
|
||||
self.block_tables = [_FakeStagedTable() for _ in self.block_sizes]
|
||||
self.num_blocks = _FakeNumBlocks(len(self.block_sizes), self.max_num_reqs)
|
||||
self.append_calls: list[tuple[int, tuple[list[int], ...], bool]] = []
|
||||
self.apply_calls = 0
|
||||
self.layout_refreshes = 0
|
||||
self.fail_on_append_call: int | None = None
|
||||
self.fail_apply = False
|
||||
self.committed: list[dict[int, tuple[int, ...]]] = [dict() for _ in self.block_sizes]
|
||||
self.__class__.instances.append(self)
|
||||
|
||||
def append_block_ids(
|
||||
self,
|
||||
req_index: int,
|
||||
new_block_ids: tuple[list[int], ...],
|
||||
overwrite: bool,
|
||||
) -> None:
|
||||
call_number = len(self.append_calls) + 1
|
||||
if self.fail_on_append_call == call_number:
|
||||
raise RuntimeError("fake append failure")
|
||||
copied_ids = tuple(list(group_ids) for group_ids in new_block_ids)
|
||||
self.append_calls.append((req_index, copied_ids, overwrite))
|
||||
for group_index, group_ids in enumerate(copied_ids):
|
||||
self.num_blocks.np[group_index, req_index] = len(group_ids)
|
||||
if group_ids:
|
||||
self.block_tables[group_index].staged.append((req_index, tuple(group_ids)))
|
||||
|
||||
def apply_staged_writes(self) -> None:
|
||||
self.apply_calls += 1
|
||||
if self.fail_apply:
|
||||
raise RuntimeError("fake apply failure")
|
||||
for group_index, group_table in enumerate(self.block_tables):
|
||||
for row, group_ids in group_table.staged:
|
||||
self.committed[group_index][row] = group_ids
|
||||
group_table.staged.clear()
|
||||
self.num_blocks.copy_to_uva()
|
||||
|
||||
def init_block_table_layout_tensors(self) -> None:
|
||||
self.layout_refreshes += 1
|
||||
|
||||
|
||||
def _metadata(
|
||||
request_id: str = "req-0",
|
||||
generation: int = 1,
|
||||
*,
|
||||
seq_len: int = 9,
|
||||
sequence_block_ids: tuple[list[int], ...] = ([1, 2, 3], [4, 5]),
|
||||
contexts: tuple[DiffusionKVContextMetadata, ...] | None = None,
|
||||
) -> DiffusionKVMetadata:
|
||||
if contexts is None:
|
||||
contexts = (
|
||||
DiffusionKVContextMetadata(
|
||||
context_id="text",
|
||||
cache_role="cross_attention",
|
||||
num_tokens=5,
|
||||
block_ids=([6, 7], [8]),
|
||||
),
|
||||
)
|
||||
return DiffusionKVMetadata(
|
||||
request_id=request_id,
|
||||
allocation_generation=generation,
|
||||
sequences=(
|
||||
DiffusionKVSequenceMetadata(
|
||||
sequence_id=0,
|
||||
prefix_len=4,
|
||||
target_len=5,
|
||||
seq_len=seq_len,
|
||||
block_ids=sequence_block_ids,
|
||||
context_ids=tuple(context.context_id for context in contexts),
|
||||
),
|
||||
),
|
||||
contexts=contexts,
|
||||
)
|
||||
|
||||
|
||||
def _install_native_modules(monkeypatch: pytest.MonkeyPatch, events: list[tuple]) -> None:
|
||||
_FakeBlockTables.instances.clear()
|
||||
|
||||
def get_block_table_width(max_num_blocks, block_size):
|
||||
events.append(("width", max_num_blocks, block_size))
|
||||
return max_num_blocks + 1
|
||||
|
||||
def init_backend(kv_cache_config, vllm_config, device):
|
||||
events.append(("backend", kv_cache_config, vllm_config, device))
|
||||
group = SimpleNamespace(get_metadata_builder=lambda _index: SimpleNamespace(reorder_batch_threshold=None))
|
||||
return [[group], [group]], "cg-support", [4, 8]
|
||||
|
||||
def init_cache(
|
||||
runner_kv_caches,
|
||||
forward_context,
|
||||
kv_cache_config,
|
||||
attn_groups,
|
||||
device,
|
||||
cache_dtype,
|
||||
kernel_block_sizes,
|
||||
vllm_config,
|
||||
):
|
||||
events.append(("cache", kv_cache_config, attn_groups, kernel_block_sizes))
|
||||
runner_kv_caches.extend(["cache-0", "cache-1"])
|
||||
return {"layer-0": "tensor-0", "layer-1": "tensor-1"}
|
||||
|
||||
monkeypatch.setattr(model_runner_backend_module, "init_attn_backend", init_backend)
|
||||
monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", init_cache)
|
||||
monkeypatch.setattr(
|
||||
model_runner_backend_module.current_omni_platform,
|
||||
"get_diffusion_kv_block_tables_cls",
|
||||
lambda: _FakeBlockTables,
|
||||
)
|
||||
monkeypatch.setattr(native_block_table, "get_block_table_width", get_block_table_width, raising=False)
|
||||
|
||||
|
||||
def _runner(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
max_rows_per_request: int = 3,
|
||||
max_num_seqs: int = 2,
|
||||
max_len: int = 16,
|
||||
max_num_batched_tokens: int = 32,
|
||||
cp_size: int = 1,
|
||||
cp_rank: int = 0,
|
||||
cp_interleave: int = 1,
|
||||
) -> tuple[DiffusionKVModelRunnerBackend, _FakeBlockTables, list[tuple]]:
|
||||
events: list[tuple] = []
|
||||
_install_native_modules(monkeypatch, events)
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
diffusion_kv_max_rows_per_request=max_rows_per_request,
|
||||
max_num_seqs=max_num_seqs,
|
||||
)
|
||||
vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(max_model_len=max_len),
|
||||
scheduler_config=SimpleNamespace(
|
||||
max_num_batched_tokens=max_num_batched_tokens,
|
||||
max_num_seqs=max_num_seqs,
|
||||
),
|
||||
parallel_config=SimpleNamespace(
|
||||
decode_context_parallel_size=cp_size,
|
||||
cp_kv_cache_interleave_size=cp_interleave,
|
||||
),
|
||||
compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
|
||||
cache_config=SimpleNamespace(cache_dtype="auto"),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
model_runner_backend_module,
|
||||
"get_dcp_group",
|
||||
lambda: SimpleNamespace(rank_in_group=cp_rank),
|
||||
)
|
||||
backend = DiffusionKVModelRunnerBackend(
|
||||
vllm_config=vllm_config,
|
||||
od_config=od_config,
|
||||
device="cuda:1",
|
||||
)
|
||||
backend._kv_cache_layer_adapters = {
|
||||
"layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
|
||||
"layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None),
|
||||
}
|
||||
|
||||
config = SimpleNamespace(
|
||||
num_blocks=32,
|
||||
kv_cache_groups=[
|
||||
SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
|
||||
SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
|
||||
],
|
||||
)
|
||||
backend.initialize_kv_cache(config)
|
||||
return backend, _FakeBlockTables.instances[-1], events
|
||||
|
||||
|
||||
def test_initialize_builds_native_block_tables_from_rank_local_config(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, events = _runner(monkeypatch)
|
||||
|
||||
assert events[:2] == [("width", 4, 4), ("width", 2, 8)]
|
||||
assert [event[0] for event in events[2:]] == ["backend", "cache"]
|
||||
assert block_tables.kwargs == {
|
||||
"block_sizes": [4, 8],
|
||||
"max_num_reqs": 6,
|
||||
"max_num_batched_tokens": 32,
|
||||
"max_num_blocks_per_group": [5, 3],
|
||||
"device": "cuda:1",
|
||||
"kernel_block_sizes": [4, 8],
|
||||
"cp_size": 1,
|
||||
"cp_rank": 0,
|
||||
"cp_interleave": 1,
|
||||
}
|
||||
assert runner.block_tables is block_tables
|
||||
assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
|
||||
assert runner.kv_caches == ["cache-0", "cache-1"]
|
||||
assert runner.vllm_config.scheduler_config.max_num_seqs == 6
|
||||
|
||||
|
||||
def test_initialize_propagates_native_context_parallel_layout(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
_, block_tables, _ = _runner(
|
||||
monkeypatch,
|
||||
cp_size=2,
|
||||
cp_rank=1,
|
||||
cp_interleave=4,
|
||||
)
|
||||
|
||||
assert block_tables.kwargs["cp_size"] == 2
|
||||
assert block_tables.kwargs["cp_rank"] == 1
|
||||
assert block_tables.kwargs["cp_interleave"] == 4
|
||||
|
||||
|
||||
def test_initialize_failure_does_not_publish_partial_cache_state(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
events: list[tuple] = []
|
||||
_install_native_modules(monkeypatch, events)
|
||||
successful_init_cache = model_runner_backend_module.init_kv_cache
|
||||
vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(max_model_len=16),
|
||||
scheduler_config=SimpleNamespace(max_num_batched_tokens=32, max_num_seqs=1),
|
||||
parallel_config=SimpleNamespace(
|
||||
decode_context_parallel_size=1,
|
||||
cp_kv_cache_interleave_size=1,
|
||||
),
|
||||
compilation_config=SimpleNamespace(static_forward_context={"layer-0": object(), "layer-1": object()}),
|
||||
cache_config=SimpleNamespace(cache_dtype="auto"),
|
||||
)
|
||||
backend = DiffusionKVModelRunnerBackend(
|
||||
vllm_config=vllm_config,
|
||||
od_config=SimpleNamespace(
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER,
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
),
|
||||
device="cuda:1",
|
||||
)
|
||||
backend._kv_cache_layer_adapters = {
|
||||
"layer-0": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-0"),
|
||||
"layer-1": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache="placeholder-1"),
|
||||
}
|
||||
config = SimpleNamespace(
|
||||
num_blocks=32,
|
||||
kv_cache_groups=[
|
||||
SimpleNamespace(layer_names=["layer-0"], kv_cache_spec=_FakeSpec(4)),
|
||||
SimpleNamespace(layer_names=["layer-1"], kv_cache_spec=_FakeSpec(8)),
|
||||
],
|
||||
)
|
||||
|
||||
def fail_init_cache(runner_kv_caches, *_args, **_kwargs):
|
||||
runner_kv_caches.append("partial-cache")
|
||||
backend._kv_cache_layer_adapters["layer-0"].kv_cache = "partial-binding"
|
||||
raise RuntimeError("cache binding failed")
|
||||
|
||||
monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", fail_init_cache)
|
||||
with pytest.raises(RuntimeError, match="cache binding failed"):
|
||||
backend.initialize_kv_cache(config)
|
||||
|
||||
assert backend.kv_cache_config is None
|
||||
assert backend.kv_caches == []
|
||||
assert backend.block_tables is None
|
||||
assert backend.paged_attention_adapter is None
|
||||
assert backend._kv_cache_layer_adapters["layer-0"].kv_cache == "placeholder-0"
|
||||
assert backend.vllm_config.scheduler_config.max_num_seqs == 1
|
||||
|
||||
monkeypatch.setattr(model_runner_backend_module, "init_kv_cache", successful_init_cache)
|
||||
backend.initialize_kv_cache(config)
|
||||
assert backend.kv_cache_config is not None
|
||||
assert backend.kv_caches == ["cache-0", "cache-1"]
|
||||
assert backend.vllm_config.scheduler_config.max_num_seqs == 2
|
||||
|
||||
|
||||
def test_initialize_rejects_config_for_different_registered_layers() -> None:
|
||||
backend = DiffusionKVModelRunnerBackend(
|
||||
vllm_config=SimpleNamespace(),
|
||||
od_config=SimpleNamespace(),
|
||||
device="cuda:1",
|
||||
)
|
||||
backend._kv_cache_layer_adapters = {
|
||||
"expected-layer": SimpleNamespace(spec=SimpleNamespace(non_causal=False), kv_cache=None)
|
||||
}
|
||||
config = SimpleNamespace(
|
||||
kv_cache_groups=[SimpleNamespace(layer_names=["other-layer"])],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="layer mismatch"):
|
||||
backend.initialize_kv_cache(config)
|
||||
|
||||
|
||||
def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
|
||||
assert runner.install_diffusion_kv_metadata(_metadata()) is True
|
||||
|
||||
assert block_tables.append_calls == [
|
||||
(0, ([1, 2, 3], [4, 5]), True),
|
||||
(1, ([6, 7], [8]), True),
|
||||
]
|
||||
assert block_tables.apply_calls == 1
|
||||
assert runner.get_diffusion_kv_row("req-0", 0) == 0
|
||||
assert runner.get_diffusion_kv_row("req-0", 0, "text") == 1
|
||||
sequence_binding = runner._resolve_paged_attention_row("req-0", 0, None)
|
||||
context_binding = runner._resolve_paged_attention_row("req-0", None, "text")
|
||||
assert (sequence_binding.row_index, sequence_binding.max_seq_len) == (0, 9)
|
||||
assert (context_binding.row_index, context_binding.max_seq_len) == (1, 5)
|
||||
|
||||
|
||||
def test_block_table_mutations_invalidate_prepared_attention_batches(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
runner, _, _ = _runner(monkeypatch)
|
||||
assert runner.paged_attention_adapter is not None
|
||||
runner.paged_attention_adapter.invalidate_prepared_batches = Mock()
|
||||
|
||||
runner.install_diffusion_kv_metadata(_metadata())
|
||||
runner.paged_attention_adapter.invalidate_prepared_batches.assert_called_once_with()
|
||||
|
||||
runner.remove_diffusion_kv_requests(["req-0"])
|
||||
assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 2
|
||||
|
||||
runner.refresh_block_table_layout()
|
||||
assert runner.paged_attention_adapter.invalidate_prepared_batches.call_count == 3
|
||||
|
||||
|
||||
def test_cache_layer_registration_is_frozen_after_physical_initialization(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
runner, _, _ = _runner(monkeypatch)
|
||||
|
||||
with pytest.raises(RuntimeError, match="cannot be changed"):
|
||||
runner.register_kv_cache_layers({})
|
||||
|
||||
|
||||
def test_request_scoped_context_uses_one_row_across_sequences(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
first_sequence = _metadata().sequences[0]
|
||||
metadata = replace(
|
||||
_metadata(),
|
||||
sequences=(
|
||||
first_sequence,
|
||||
replace(
|
||||
first_sequence,
|
||||
sequence_id=1,
|
||||
block_ids=([9, 10, 11], [12, 13]),
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
assert runner.install_diffusion_kv_metadata(metadata) is True
|
||||
|
||||
assert block_tables.append_calls == [
|
||||
(0, ([1, 2, 3], [4, 5]), True),
|
||||
(1, ([9, 10, 11], [12, 13]), True),
|
||||
(2, ([6, 7], [8]), True),
|
||||
]
|
||||
assert runner.get_diffusion_kv_row("req-0", 0, "text") == 2
|
||||
assert runner.get_diffusion_kv_row("req-0", 1, "text") == 2
|
||||
|
||||
|
||||
def test_row_block_count_uses_native_spec_semantics(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
spec = runner.kv_cache_config.kv_cache_groups[0].kv_cache_spec
|
||||
spec.max_num_blocks_per_req = lambda _config, token_len: math.ceil(token_len / 8)
|
||||
|
||||
metadata = _metadata(
|
||||
sequence_block_ids=([1, 2], [4, 5]),
|
||||
contexts=(),
|
||||
)
|
||||
|
||||
assert runner.install_diffusion_kv_metadata(metadata) is True
|
||||
assert block_tables.append_calls == [(0, ([1, 2], [4, 5]), True)]
|
||||
|
||||
|
||||
def test_repeated_native_block_ids_are_allowed(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
metadata = _metadata(
|
||||
sequence_block_ids=([1, 1, 3], [4, 4]),
|
||||
contexts=(),
|
||||
)
|
||||
|
||||
assert runner.install_diffusion_kv_metadata(metadata) is True
|
||||
assert block_tables.append_calls == [(0, ([1, 1, 3], [4, 4]), True)]
|
||||
|
||||
|
||||
def test_generation_is_idempotent_stale_safe_and_conflict_safe(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
metadata = _metadata()
|
||||
runner.install_diffusion_kv_metadata(metadata)
|
||||
|
||||
assert runner.install_diffusion_kv_metadata(metadata) is False
|
||||
assert runner.install_diffusion_kv_metadata(replace(metadata, allocation_generation=2)) is False
|
||||
assert block_tables.apply_calls == 1
|
||||
|
||||
with pytest.raises(ValueError, match="Stale Diffusion KV allocation generation"):
|
||||
runner.install_diffusion_kv_metadata(metadata)
|
||||
with pytest.raises(ValueError, match="Conflicting Diffusion KV allocation snapshot"):
|
||||
runner.install_diffusion_kv_metadata(_metadata(generation=3, sequence_block_ids=([9, 2, 3], [4, 5])))
|
||||
assert block_tables.apply_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("metadata", "message"),
|
||||
[
|
||||
(_metadata(sequence_block_ids=([1, 2, 3],)), "block groups"),
|
||||
(_metadata(sequence_block_ids=([1, 2], [4, 5])), "expected 3"),
|
||||
(_metadata(sequence_block_ids=([1, 2, 32], [4, 5])), "outside"),
|
||||
(
|
||||
_metadata(
|
||||
seq_len=17,
|
||||
sequence_block_ids=([1, 2, 3, 4, 5], [6, 7, 8]),
|
||||
contexts=(),
|
||||
),
|
||||
"exceeds row capacity",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_invalid_group_count_block_count_range_and_row_capacity_are_atomic(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
metadata: DiffusionKVMetadata,
|
||||
message: str,
|
||||
) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
runner.install_diffusion_kv_metadata(metadata)
|
||||
|
||||
assert block_tables.append_calls == []
|
||||
assert block_tables.apply_calls == 0
|
||||
assert runner._diffusion_kv_identity_to_row == {}
|
||||
assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
|
||||
|
||||
|
||||
def test_invalid_later_context_does_not_stage_earlier_valid_row(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
bad_context = DiffusionKVContextMetadata(
|
||||
context_id="text",
|
||||
cache_role="cross_attention",
|
||||
num_tokens=5,
|
||||
block_ids=([6, 32], [8]),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="outside"):
|
||||
runner.install_diffusion_kv_metadata(_metadata(contexts=(bad_context,)))
|
||||
|
||||
assert block_tables.append_calls == []
|
||||
assert all(not group_table.staged for group_table in block_tables.block_tables)
|
||||
|
||||
|
||||
def test_native_append_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
block_tables.fail_on_append_call = 2
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake append failure"):
|
||||
runner.install_diffusion_kv_metadata(_metadata())
|
||||
|
||||
assert np.count_nonzero(block_tables.num_blocks.np) == 0
|
||||
assert all(not group_table.staged for group_table in block_tables.block_tables)
|
||||
assert runner._diffusion_kv_identity_to_row == {}
|
||||
assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
|
||||
|
||||
|
||||
def test_native_apply_failure_rolls_back_staged_state_and_registry(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
block_tables.fail_apply = True
|
||||
|
||||
with pytest.raises(RuntimeError, match="fake apply failure"):
|
||||
runner.install_diffusion_kv_metadata(_metadata())
|
||||
|
||||
assert block_tables.apply_calls == 1
|
||||
assert np.count_nonzero(block_tables.num_blocks.np) == 0
|
||||
assert all(not group_table.staged for group_table in block_tables.block_tables)
|
||||
assert all(group_table.clear_calls == 1 for group_table in block_tables.block_tables)
|
||||
assert runner._diffusion_kv_identity_to_row == {}
|
||||
assert runner._diffusion_kv_free_rows == [5, 4, 3, 2, 1, 0]
|
||||
|
||||
|
||||
def test_adapter_and_global_row_capacity_are_enforced(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
|
||||
second_context = DiffusionKVContextMetadata(
|
||||
context_id="image",
|
||||
cache_role="cross_attention",
|
||||
num_tokens=1,
|
||||
block_ids=([9], [10]),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="adapter limit is 2"):
|
||||
runner.install_diffusion_kv_metadata(_metadata(contexts=(_metadata().contexts[0], second_context)))
|
||||
assert block_tables.append_calls == []
|
||||
|
||||
runner.install_diffusion_kv_metadata(_metadata())
|
||||
append_count = len(block_tables.append_calls)
|
||||
with pytest.raises(ValueError, match="only 0 rows are free"):
|
||||
runner.install_diffusion_kv_metadata(_metadata("req-1"))
|
||||
assert len(block_tables.append_calls) == append_count
|
||||
|
||||
|
||||
def test_empty_and_duplicate_sequence_metadata_are_rejected(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
empty = DiffusionKVMetadata(request_id="req-0", allocation_generation=0, sequences=())
|
||||
duplicate_sequence = _metadata().sequences[0]
|
||||
duplicate = DiffusionKVMetadata(
|
||||
request_id="req-0",
|
||||
allocation_generation=0,
|
||||
sequences=(duplicate_sequence, duplicate_sequence),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="at least one sequence"):
|
||||
runner.install_diffusion_kv_metadata(empty)
|
||||
with pytest.raises(ValueError, match="Duplicate diffusion KV row identity"):
|
||||
runner.install_diffusion_kv_metadata(duplicate)
|
||||
assert block_tables.append_calls == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("metadata", "message"),
|
||||
[
|
||||
(_metadata(generation=-1), "non-negative integer"),
|
||||
(_metadata(request_id=""), "non-empty string"),
|
||||
(_metadata(sequence_block_ids=([1, 2, True], [4, 5])), "non-integer block ID"),
|
||||
],
|
||||
)
|
||||
def test_request_generation_and_block_id_types_are_validated(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
metadata: DiffusionKVMetadata,
|
||||
message: str,
|
||||
) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
|
||||
with pytest.raises(ValueError, match=message):
|
||||
runner.install_diffusion_kv_metadata(metadata)
|
||||
|
||||
assert block_tables.append_calls == []
|
||||
assert runner._diffusion_kv_request_states == {}
|
||||
|
||||
|
||||
def test_cleanup_is_idempotent_and_reuses_native_rows(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch, max_rows_per_request=2, max_num_seqs=1)
|
||||
runner.install_diffusion_kv_metadata(_metadata())
|
||||
|
||||
assert runner.remove_diffusion_kv_requests(["missing"]) == 0
|
||||
assert runner.remove_diffusion_kv_requests(["req-0", "req-0"]) == 2
|
||||
assert block_tables.append_calls[-2:] == [(0, ([], []), True), (1, ([], []), True)]
|
||||
assert block_tables.num_blocks.np[:, :2].tolist() == [[0, 0], [0, 0]]
|
||||
assert runner.remove_diffusion_kv_requests(["req-0"]) == 0
|
||||
with pytest.raises(KeyError, match="No native diffusion KV row"):
|
||||
runner.get_diffusion_kv_row("req-0", 0)
|
||||
|
||||
runner.install_diffusion_kv_metadata(_metadata("req-1"))
|
||||
assert runner.get_diffusion_kv_row("req-1", 0) == 0
|
||||
assert runner.get_diffusion_kv_row("req-1", 0, "text") == 1
|
||||
|
||||
|
||||
def test_cached_step_does_not_reinstall_metadata() -> None:
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner.pipeline = object()
|
||||
runner.install_diffusion_kv_metadata = Mock()
|
||||
runner._supports_step_mode = Mock(return_value=False)
|
||||
scheduler_output = DiffusionSchedulerOutput(
|
||||
step_id=1,
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData(request_ids=["req-0"]),
|
||||
finished_req_ids=set(),
|
||||
num_running_reqs=1,
|
||||
num_waiting_reqs=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="does not support step execution"):
|
||||
runner.execute_stepwise(scheduler_output)
|
||||
runner.install_diffusion_kv_metadata.assert_not_called()
|
||||
|
||||
|
||||
def test_dense_cleanup_and_wake_are_noops_but_paged_wake_refreshes(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner, block_tables, _ = _runner(monkeypatch)
|
||||
runner.refresh_block_table_layout()
|
||||
assert block_tables.layout_refreshes == 1
|
||||
|
||||
dense_backend = DiffusionKVModelRunnerBackend(
|
||||
vllm_config=SimpleNamespace(),
|
||||
od_config=SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY),
|
||||
device="cuda:1",
|
||||
)
|
||||
dense_backend.block_tables = Mock()
|
||||
assert dense_backend.remove_diffusion_kv_requests(["req-0"]) == 0
|
||||
dense_backend.refresh_block_table_layout()
|
||||
dense_backend.block_tables.init_block_table_layout_tensors.assert_not_called()
|
||||
with pytest.raises(ValueError, match="Dense diffusion execution"):
|
||||
dense_backend.install_diffusion_kv_metadata(_metadata())
|
||||
@@ -3,6 +3,7 @@
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.config.omni_config import _DiffusionConfigProjection
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
|
||||
@@ -39,6 +40,8 @@ def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
|
||||
with pytest.raises(ValueError, match="does not support imported AR KV"):
|
||||
OmniDiffusionConfig.from_kwargs(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=1,
|
||||
max_num_batched_tokens=1,
|
||||
omni_kv_config={"need_recv_cache": True},
|
||||
)
|
||||
|
||||
@@ -47,6 +50,33 @@ def test_paged_scheduler_does_not_depend_on_model_registry() -> None:
|
||||
config = OmniDiffusionConfig.from_kwargs(
|
||||
model_class_name="FutureDiffusionModel",
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=1,
|
||||
max_num_batched_tokens=1,
|
||||
)
|
||||
|
||||
assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config_cls", [OmniDiffusionConfig, _DiffusionConfigProjection])
|
||||
def test_paged_scheduler_mode_is_platform_agnostic(config_cls) -> None:
|
||||
config = config_cls.from_kwargs(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=1,
|
||||
)
|
||||
|
||||
assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||
|
||||
|
||||
def test_paged_scheduler_requires_a_worker_row_limit() -> None:
|
||||
with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
|
||||
OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_scheduler")
|
||||
with pytest.raises(ValueError, match="requires diffusion_kv_max_rows_per_request"):
|
||||
_DiffusionConfigProjection.from_kwargs(diffusion_kv_mode="paged_scheduler")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("invalid_limit", [0, -1, True, 1.5])
|
||||
def test_worker_row_limit_must_be_a_positive_integer(invalid_limit: object) -> None:
|
||||
with pytest.raises((TypeError, ValueError)):
|
||||
OmniDiffusionConfig.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)
|
||||
with pytest.raises((TypeError, ValueError)):
|
||||
_DiffusionConfigProjection.from_kwargs(diffusion_kv_max_rows_per_request=invalid_limit)
|
||||
@@ -22,6 +22,7 @@ class _Executor:
|
||||
self.specs = [{"layer0": object()}, {"layer0": object()}]
|
||||
self.memory = [1024, 1024]
|
||||
self.configs = None
|
||||
self.resolved_max_model_len = None
|
||||
self.profile_requests = None
|
||||
|
||||
def get_kv_cache_specs(self):
|
||||
@@ -31,8 +32,9 @@ class _Executor:
|
||||
self.profile_requests = profile_requests
|
||||
return self.memory
|
||||
|
||||
def set_kv_cache_configs(self, configs) -> None:
|
||||
def set_kv_cache_configs(self, configs, resolved_max_model_len) -> None:
|
||||
self.configs = configs
|
||||
self.resolved_max_model_len = resolved_max_model_len
|
||||
|
||||
|
||||
def _od_config(**overrides):
|
||||
@@ -55,6 +57,7 @@ def _od_config(**overrides):
|
||||
),
|
||||
gpu_memory_utilization=0.9,
|
||||
kv_cache_memory_bytes=None,
|
||||
max_num_seqs=1,
|
||||
max_num_batched_tokens=64,
|
||||
num_gpus=1,
|
||||
)
|
||||
@@ -64,9 +67,9 @@ def _od_config(**overrides):
|
||||
|
||||
def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
|
||||
executor = _Executor()
|
||||
od_config = _od_config(num_gpus=2)
|
||||
od_config = _od_config(num_gpus=2, max_model_len=-1)
|
||||
vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(max_model_len=4096),
|
||||
model_config=SimpleNamespace(max_model_len=256),
|
||||
max_in_flight_tokens=128,
|
||||
)
|
||||
worker_configs = [object(), object()]
|
||||
@@ -91,6 +94,7 @@ def test_control_plane_builds_worker_and_scheduler_configs(monkeypatch) -> None:
|
||||
build.assert_called_once_with(vllm_config, executor.specs, executor.memory)
|
||||
assert result == (scheduler_kv_cache_config, 16, 16, vllm_config)
|
||||
assert executor.configs == worker_configs
|
||||
assert executor.resolved_max_model_len == vllm_config.model_config.max_model_len
|
||||
assert executor.profile_requests is profile_requests
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from importlib import import_module, util
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
|
||||
|
||||
|
||||
def test_native_block_tables_slot_mapping_handles_non_aligned_tail() -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA is required for the native BlockTables slot-mapping smoke test")
|
||||
module_name = "vllm.v1.worker.gpu.block_table"
|
||||
try:
|
||||
module_spec = util.find_spec(module_name)
|
||||
except ModuleNotFoundError as exc:
|
||||
if exc.name is None or not module_name.startswith(f"{exc.name}."):
|
||||
raise
|
||||
pytest.skip("installed vLLM does not provide GPU BlockTables")
|
||||
if module_spec is None:
|
||||
pytest.skip("installed vLLM does not provide GPU BlockTables")
|
||||
block_table_module = import_module(module_name)
|
||||
block_tables_cls = getattr(block_table_module, "BlockTables", None)
|
||||
if block_tables_cls is None or not hasattr(block_tables_cls, "compute_slot_mappings"):
|
||||
pytest.skip("installed vLLM does not provide the required BlockTables API")
|
||||
|
||||
device = torch.device("cuda", torch.accelerator.current_device_index())
|
||||
block_tables = block_tables_cls(
|
||||
block_sizes=[4],
|
||||
max_num_reqs=1,
|
||||
max_num_batched_tokens=16,
|
||||
max_num_blocks_per_group=[4],
|
||||
device=device,
|
||||
kernel_block_sizes=[4],
|
||||
)
|
||||
block_tables.append_block_ids(0, ([1, 2, 3],), overwrite=True)
|
||||
block_tables.apply_staged_writes()
|
||||
idx_mapping = torch.tensor([0], dtype=torch.int32, device=device)
|
||||
query_start_loc = torch.tensor([0, 9], dtype=torch.int32, device=device)
|
||||
positions = torch.arange(9, dtype=torch.int64, device=device)
|
||||
|
||||
slot_mappings = block_tables.compute_slot_mappings(
|
||||
idx_mapping,
|
||||
query_start_loc,
|
||||
positions,
|
||||
num_tokens_padded=9,
|
||||
)
|
||||
torch.accelerator.synchronize(device)
|
||||
|
||||
assert slot_mappings.cpu().tolist() == [[4, 5, 6, 7, 8, 9, 10, 11, 12]]
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,251 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
from vllm.config import AttentionConfig, CacheConfig, VllmConfig, set_current_vllm_config
|
||||
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheTensor
|
||||
from vllm.v1.worker.gpu.attn_utils import init_attn_backend, init_kv_cache
|
||||
from vllm.v1.worker.gpu.block_table import BlockTables
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
|
||||
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
|
||||
DiffusionPagedAttentionAdapter,
|
||||
DiffusionPagedAttentionLayerAdapter,
|
||||
DiffusionPagedAttentionRow,
|
||||
DiffusionPagedAttentionRowBinding,
|
||||
)
|
||||
from vllm_omni.diffusion.vllm_config import _DiffusionVllmModelConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
|
||||
|
||||
_LAYER_NAME = "model.layers.0.attn"
|
||||
_NUM_HEADS = 2
|
||||
_HEAD_SIZE = 64
|
||||
_BLOCK_SIZE = 16
|
||||
|
||||
|
||||
class _SmokeDiffusionAttention(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.num_heads = _NUM_HEADS
|
||||
self.num_kv_heads = _NUM_HEADS
|
||||
self.head_size = _HEAD_SIZE
|
||||
self.softmax_scale = _HEAD_SIZE**-0.5
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for native paged attention")
|
||||
def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> None:
|
||||
device = torch.device("cuda", torch.accelerator.current_device_index())
|
||||
vllm_config = VllmConfig(
|
||||
cache_config=CacheConfig(
|
||||
block_size=_BLOCK_SIZE,
|
||||
cache_dtype="float16",
|
||||
enable_prefix_caching=False,
|
||||
),
|
||||
attention_config=AttentionConfig(use_non_causal=True),
|
||||
)
|
||||
model_config = _DiffusionVllmModelConfig(
|
||||
model=None,
|
||||
dtype=torch.float16,
|
||||
max_model_len=32,
|
||||
original_max_model_len=32,
|
||||
hf_config=SimpleNamespace(model_type="diffusion_paged_attention_smoke"),
|
||||
)
|
||||
vllm_config.model_config = model_config
|
||||
model_config.set_attention_geometry(
|
||||
num_heads=_NUM_HEADS,
|
||||
num_kv_heads=_NUM_HEADS,
|
||||
head_size=_HEAD_SIZE,
|
||||
)
|
||||
vllm_config.compilation_config.static_forward_context.clear()
|
||||
|
||||
diffusion_layer = _SmokeDiffusionAttention().to(device)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=_BLOCK_SIZE,
|
||||
num_kv_heads=_NUM_HEADS,
|
||||
head_size=_HEAD_SIZE,
|
||||
dtype=torch.float16,
|
||||
non_causal=True,
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
native_layer = DiffusionPagedAttentionLayerAdapter(
|
||||
layer_name=_LAYER_NAME,
|
||||
layer=diffusion_layer,
|
||||
spec=spec,
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
)
|
||||
vllm_config.compilation_config.static_forward_context[_LAYER_NAME] = native_layer
|
||||
|
||||
canonical_spec = native_layer.spec
|
||||
kv_cache_config = KVCacheConfig(
|
||||
num_blocks=4,
|
||||
kv_cache_tensors=[
|
||||
KVCacheTensor(
|
||||
size=canonical_spec.page_size_bytes * 4,
|
||||
shared_by=[_LAYER_NAME],
|
||||
)
|
||||
],
|
||||
kv_cache_groups=[
|
||||
KVCacheGroupSpec(
|
||||
layer_names=[_LAYER_NAME],
|
||||
kv_cache_spec=canonical_spec,
|
||||
)
|
||||
],
|
||||
)
|
||||
with set_current_vllm_config(vllm_config):
|
||||
attn_groups, _, kernel_block_sizes = init_attn_backend(
|
||||
kv_cache_config,
|
||||
vllm_config,
|
||||
device,
|
||||
)
|
||||
runner_kv_caches: list[torch.Tensor | list[torch.Tensor]] = []
|
||||
init_kv_cache(
|
||||
runner_kv_caches,
|
||||
vllm_config.compilation_config.static_forward_context,
|
||||
kv_cache_config,
|
||||
attn_groups,
|
||||
device,
|
||||
vllm_config.cache_config.cache_dtype,
|
||||
kernel_block_sizes,
|
||||
vllm_config,
|
||||
)
|
||||
block_tables = BlockTables(
|
||||
block_sizes=[_BLOCK_SIZE],
|
||||
max_num_reqs=1,
|
||||
max_num_batched_tokens=32,
|
||||
max_num_blocks_per_group=[2],
|
||||
device=device,
|
||||
kernel_block_sizes=kernel_block_sizes,
|
||||
)
|
||||
|
||||
block_tables.append_block_ids(0, ([3, 1],), overwrite=True)
|
||||
block_tables.apply_staged_writes()
|
||||
adapter = DiffusionPagedAttentionAdapter(
|
||||
vllm_config=vllm_config,
|
||||
device=device,
|
||||
kv_cache_config=kv_cache_config,
|
||||
block_tables=block_tables,
|
||||
attn_groups=attn_groups,
|
||||
layers={_LAYER_NAME: native_layer},
|
||||
resolve_row=lambda _request_id, _sequence_id, _context_id: DiffusionPagedAttentionRowBinding(
|
||||
row_index=0,
|
||||
max_seq_len=19,
|
||||
),
|
||||
)
|
||||
prefix_batch = adapter.prepare_batch(
|
||||
[
|
||||
DiffusionPagedAttentionRow(
|
||||
request_id="req-0",
|
||||
sequence_id=0,
|
||||
query_len=17,
|
||||
seq_len=17,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
torch.manual_seed(1)
|
||||
query = torch.randn(19, _NUM_HEADS, _HEAD_SIZE, dtype=torch.float16, device=device)
|
||||
key = torch.randn_like(query)
|
||||
value = torch.randn_like(query)
|
||||
omni_backend = FlashAttentionImpl(
|
||||
num_heads=_NUM_HEADS,
|
||||
head_size=_HEAD_SIZE,
|
||||
softmax_scale=diffusion_layer.softmax_scale,
|
||||
num_kv_heads=_NUM_HEADS,
|
||||
)
|
||||
with adapter.activate(prefix_batch):
|
||||
prefix_context = adapter.prepare_layer_context(_LAYER_NAME, query[:17], key[:17], value[:17])
|
||||
prefix_output = omni_backend.forward_paged(prefix_context)
|
||||
prefix_slot_mappings = prefix_batch.slot_mappings.clone()
|
||||
|
||||
suffix_batch = adapter.prepare_batch(
|
||||
[
|
||||
DiffusionPagedAttentionRow(
|
||||
request_id="req-0",
|
||||
sequence_id=0,
|
||||
query_len=2,
|
||||
seq_len=19,
|
||||
kv_start_pos=17,
|
||||
)
|
||||
]
|
||||
)
|
||||
with adapter.activate(suffix_batch):
|
||||
suffix_context = adapter.prepare_layer_context(_LAYER_NAME, query[17:], key[17:], value[17:])
|
||||
suffix_output = omni_backend.forward_paged(suffix_context)
|
||||
suffix_slot_mappings = suffix_batch.slot_mappings.clone()
|
||||
|
||||
piecewise_batch = adapter.prepare_batch(
|
||||
[
|
||||
DiffusionPagedAttentionRow(
|
||||
request_id="req-0",
|
||||
sequence_id=0,
|
||||
query_len=19,
|
||||
seq_len=19,
|
||||
)
|
||||
]
|
||||
)
|
||||
mixed_mask = torch.ones(19, 19, dtype=torch.bool, device=device).tril()
|
||||
mixed_mask[5:10, 5:10] = True
|
||||
mixed_metadata = AttentionMetadata(
|
||||
attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
|
||||
full_attn_spans=[[(5, 10)]],
|
||||
)
|
||||
with adapter.activate(piecewise_batch):
|
||||
piecewise_context = adapter.prepare_layer_context(
|
||||
_LAYER_NAME,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
omni_attn_metadata=mixed_metadata,
|
||||
)
|
||||
piecewise_output = omni_backend.forward_paged(piecewise_context)
|
||||
|
||||
prefix_reference = (
|
||||
F.scaled_dot_product_attention(
|
||||
query[:17].transpose(0, 1).unsqueeze(0),
|
||||
key[:17].transpose(0, 1).unsqueeze(0),
|
||||
value[:17].transpose(0, 1).unsqueeze(0),
|
||||
is_causal=False,
|
||||
scale=diffusion_layer.softmax_scale,
|
||||
)
|
||||
.squeeze(0)
|
||||
.transpose(0, 1)
|
||||
)
|
||||
suffix_reference = (
|
||||
F.scaled_dot_product_attention(
|
||||
query[17:].transpose(0, 1).unsqueeze(0),
|
||||
key.transpose(0, 1).unsqueeze(0),
|
||||
value.transpose(0, 1).unsqueeze(0),
|
||||
is_causal=False,
|
||||
scale=diffusion_layer.softmax_scale,
|
||||
)
|
||||
.squeeze(0)
|
||||
.transpose(0, 1)
|
||||
)
|
||||
piecewise_reference = (
|
||||
F.scaled_dot_product_attention(
|
||||
query.transpose(0, 1).unsqueeze(0),
|
||||
key.transpose(0, 1).unsqueeze(0),
|
||||
value.transpose(0, 1).unsqueeze(0),
|
||||
attn_mask=mixed_mask.unsqueeze(0).unsqueeze(0),
|
||||
scale=diffusion_layer.softmax_scale,
|
||||
)
|
||||
.squeeze(0)
|
||||
.transpose(0, 1)
|
||||
)
|
||||
torch.accelerator.synchronize(device)
|
||||
|
||||
assert prefix_slot_mappings.cpu().tolist() == [list(range(3 * _BLOCK_SIZE, 4 * _BLOCK_SIZE)) + [_BLOCK_SIZE]]
|
||||
assert suffix_slot_mappings.cpu().tolist() == [[_BLOCK_SIZE + 1, _BLOCK_SIZE + 2]]
|
||||
torch.testing.assert_close(prefix_output, prefix_reference, rtol=2e-2, atol=2e-2)
|
||||
torch.testing.assert_close(suffix_output, suffix_reference, rtol=2e-2, atol=2e-2)
|
||||
torch.testing.assert_close(piecewise_output, piecewise_reference, rtol=2e-2, atol=2e-2)
|
||||
@@ -25,6 +25,23 @@ class _Backend:
|
||||
return True
|
||||
|
||||
|
||||
class _RunnerKVBackend:
|
||||
def __init__(self) -> None:
|
||||
self.kv_cache_config = None
|
||||
|
||||
def register_kv_cache_layers(self, layers):
|
||||
return {layer_name: spec for layer_name, (_layer, spec) in layers.items()}
|
||||
|
||||
def initialize_kv_cache(self, config) -> None:
|
||||
configured_layers = {layer_name for group in config.kv_cache_groups for layer_name in group.layer_names}
|
||||
if configured_layers != {"image_attention"}:
|
||||
raise ValueError(
|
||||
"Rank-local Diffusion KVCacheConfig layer mismatch: "
|
||||
f"expected=['image_attention'], configured={sorted(configured_layers)}"
|
||||
)
|
||||
self.kv_cache_config = config
|
||||
|
||||
|
||||
def _attention(*, enabled: bool) -> Attention:
|
||||
attention = Attention.__new__(Attention)
|
||||
nn.Module.__init__(attention)
|
||||
@@ -47,6 +64,7 @@ def _runner(attention: Attention) -> DiffusionModelRunner:
|
||||
pipeline = nn.Module()
|
||||
pipeline.image_attention = attention
|
||||
runner.pipeline = pipeline
|
||||
runner.diffusion_kv_backend = _RunnerKVBackend()
|
||||
return runner
|
||||
|
||||
|
||||
@@ -102,14 +120,21 @@ def test_worker_selects_its_rank_local_config() -> None:
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.rank = 1
|
||||
worker.od_config = SimpleNamespace(num_gpus=2)
|
||||
worker.vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(max_model_len=None),
|
||||
cache_config=SimpleNamespace(num_gpu_blocks=None),
|
||||
)
|
||||
worker._maybe_get_memory_pool_context = lambda _tag: nullcontext()
|
||||
worker.model_runner = SimpleNamespace(set_kv_cache_config=lambda config: setattr(worker, "installed", config))
|
||||
configs = [object(), object()]
|
||||
configs = [SimpleNamespace(num_blocks=4), SimpleNamespace(num_blocks=8)]
|
||||
|
||||
worker.set_kv_cache_configs(configs)
|
||||
worker.set_kv_cache_configs(configs, 64)
|
||||
|
||||
assert worker.installed is configs[1]
|
||||
assert worker.vllm_config.model_config.max_model_len == 64
|
||||
assert worker.vllm_config.cache_config.num_gpu_blocks == 8
|
||||
with pytest.raises(ValueError, match="rank count mismatch"):
|
||||
worker.set_kv_cache_configs(configs[:1])
|
||||
worker.set_kv_cache_configs(configs[:1], 64)
|
||||
|
||||
|
||||
def test_worker_honors_explicit_kv_memory_budget(monkeypatch) -> None:
|
||||
|
||||
@@ -19,8 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_TRANSFORMER_MODULE = "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer"
|
||||
|
||||
NUM_HEADS = 4
|
||||
NUM_KV_HEADS = 2
|
||||
NUM_HEADS = 32
|
||||
NUM_KV_HEADS = 8
|
||||
HEAD_DIM = 16
|
||||
IMAGE_TOKEN_LEN = 8
|
||||
SUFFIX_TOKEN_LEN = 3
|
||||
@@ -35,8 +35,14 @@ SCALING = 1.0 / math.sqrt(HEAD_DIM)
|
||||
class MockAttention(nn.Module):
|
||||
def __init__(self, num_heads, head_size, causal=False, softmax_scale=None, num_kv_heads=None, **kwargs):
|
||||
super().__init__()
|
||||
self.paged_kv_active = False
|
||||
self.calls = []
|
||||
|
||||
def is_paged_kv_active(self):
|
||||
return self.paged_kv_active
|
||||
|
||||
def forward(self, query, key, value, attn_metadata=None, **kwargs):
|
||||
self.calls.append((query, key, value))
|
||||
return query
|
||||
|
||||
|
||||
@@ -126,6 +132,37 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
|
||||
assert mgr.state_dict() == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("paged_kv_active", "expected_kv_heads"),
|
||||
[(False, NUM_HEADS), (True, NUM_KV_HEADS)],
|
||||
)
|
||||
def test_gqa_kv_stays_compressed_for_paged_attention(
|
||||
paged_kv_active: bool,
|
||||
expected_kv_heads: int,
|
||||
) -> None:
|
||||
mgr = _make_cache_mgr()
|
||||
mgr.attn.paged_kv_active = paged_kv_active
|
||||
q_len = 3
|
||||
key, value = _make_known_kv(q_len)
|
||||
|
||||
_call_mgr(
|
||||
mgr,
|
||||
bs=1,
|
||||
q_len=q_len,
|
||||
seq_len=q_len,
|
||||
key_flat=key,
|
||||
value_flat=value,
|
||||
first_step=True,
|
||||
num_image_tokens=0,
|
||||
gen_timestep_scatter_index=_gen_timestep_index(1, 0),
|
||||
)
|
||||
|
||||
_, key_input, value_input = mgr.attn.calls[-1]
|
||||
expected_shape = (1, q_len, expected_kv_heads, HEAD_DIM)
|
||||
assert key_input.shape == expected_shape
|
||||
assert value_input.shape == expected_shape
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Test 1: No AR KV — basic cache → reuse
|
||||
# ============================================================
|
||||
|
||||
@@ -149,10 +149,18 @@ class TestCreateDefaultDiffusion:
|
||||
assert ea["enforce_eager"] is True
|
||||
assert ea["lora_path"] == "/tmp/lora"
|
||||
|
||||
def test_diffusion_kv_mode_roundtrip(self):
|
||||
od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler")
|
||||
def test_diffusion_kv_mode_roundtrip(self, monkeypatch):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
||||
od = _roundtrip_diffusion_config(
|
||||
model="x",
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
diffusion_kv_max_rows_per_request=2,
|
||||
)
|
||||
|
||||
assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||
assert od.diffusion_kv_max_rows_per_request == 2
|
||||
|
||||
def test_diffusion_kv_sizing_fields_roundtrip(self, monkeypatch):
|
||||
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda self: 29500)
|
||||
|
||||
@@ -12,6 +12,7 @@ import pytest
|
||||
from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -88,6 +89,50 @@ def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None:
|
||||
assert engine.scheduler.get_request_state(request_id) is None
|
||||
|
||||
|
||||
def _make_kv_cleanup_engine(mode: DiffusionKVCacheMode) -> tuple[DiffusionEngine, Mock]:
|
||||
engine = DiffusionEngine.__new__(DiffusionEngine)
|
||||
engine.od_config = SimpleNamespace(diffusion_kv_mode=mode)
|
||||
cleanup = Mock()
|
||||
engine.executor = SimpleNamespace(remove_diffusion_kv_requests=cleanup)
|
||||
return engine, cleanup
|
||||
|
||||
|
||||
def test_paged_terminal_output_clears_worker_rows() -> None:
|
||||
engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
engine._finalize_finished_request = lambda request_id, *_args: request_id
|
||||
engine._put_output = lambda *_args: None
|
||||
|
||||
engine._emit_finished_outputs({"req-0", "req-1"})
|
||||
|
||||
cleanup.assert_called_once()
|
||||
assert set(cleanup.call_args.args[0]) == {"req-0", "req-1"}
|
||||
|
||||
|
||||
def test_paged_abort_clears_worker_rows_before_scheduler_free() -> None:
|
||||
engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
engine.scheduler = SimpleNamespace(
|
||||
get_request_state=lambda _request_id: None,
|
||||
finish_requests=Mock(side_effect=AssertionError("unknown request must not be logically freed")),
|
||||
)
|
||||
|
||||
engine._abort_requests(["req-idle", "req-idle"])
|
||||
|
||||
cleanup.assert_called_once_with(["req-idle"])
|
||||
engine.scheduler.finish_requests.assert_not_called()
|
||||
|
||||
|
||||
def test_dense_terminal_and_abort_paths_skip_worker_row_cleanup() -> None:
|
||||
engine, cleanup = _make_kv_cleanup_engine(DiffusionKVCacheMode.DENSE_LEGACY)
|
||||
engine._finalize_finished_request = lambda request_id, *_args: request_id
|
||||
engine._put_output = lambda *_args: None
|
||||
engine.scheduler = SimpleNamespace(get_request_state=lambda _request_id: None)
|
||||
|
||||
engine._emit_finished_outputs({"req-0"})
|
||||
engine._abort_requests(["req-idle"])
|
||||
|
||||
cleanup.assert_not_called()
|
||||
|
||||
|
||||
def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
od_config = SimpleNamespace(
|
||||
custom_pipeline_args=None,
|
||||
|
||||
@@ -11,6 +11,7 @@ import torch
|
||||
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.diffusion_kv.model_runner_backend import DiffusionKVModelRunnerBackend
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
|
||||
|
||||
@@ -193,6 +194,11 @@ def _make_runner(cache_backend, cache_backend_name: str, enable_cache_dit_summar
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
streaming_output=False,
|
||||
)
|
||||
runner.diffusion_kv_backend = DiffusionKVModelRunnerBackend(
|
||||
vllm_config=runner.vllm_config,
|
||||
od_config=runner.od_config,
|
||||
device=runner.device,
|
||||
)
|
||||
runner.kv_transfer_manager = SimpleNamespace(
|
||||
receive_kv_cache=lambda req, target_device=None: None,
|
||||
receive_multi_kv_cache=lambda req, cfg_kv_collect_func=None, target_device=None: None,
|
||||
|
||||
@@ -40,6 +40,12 @@ class TestPlatformDiffusionHooks:
|
||||
result = OmniPlatform.get_diffusion_model_runner_cls()
|
||||
assert result == "vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner"
|
||||
|
||||
def test_get_diffusion_kv_block_tables_cls_default(self):
|
||||
"""Test the default paged-KV BlockTables implementation."""
|
||||
from vllm.v1.worker.gpu.block_table import BlockTables
|
||||
|
||||
assert OmniPlatform.get_diffusion_kv_block_tables_cls() is BlockTables
|
||||
|
||||
def test_oot_enum_exists(self):
|
||||
"""Test OOT is a valid platform enum value."""
|
||||
assert OmniPlatformEnum.OOT.value == "oot"
|
||||
|
||||
@@ -17,6 +17,8 @@ import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
||||
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
||||
from vllm_omni.diffusion.diffusion_kv.metadata import DiffusionKVMetadata
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
|
||||
from vllm_omni.diffusion.distributed.parallel_state import (
|
||||
@@ -306,11 +308,14 @@ def _make_vllm_config():
|
||||
)
|
||||
|
||||
|
||||
def _make_runner():
|
||||
def _make_runner(
|
||||
diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
):
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.vllm_config = _make_vllm_config()
|
||||
runner.od_config = SimpleNamespace(
|
||||
cache_backend=None,
|
||||
diffusion_kv_mode=diffusion_kv_mode,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
streaming_output=False,
|
||||
)
|
||||
@@ -330,6 +335,7 @@ def _make_distributed_runner(mode: str, device: torch.device):
|
||||
runner.vllm_config = _make_vllm_config()
|
||||
runner.od_config = SimpleNamespace(
|
||||
cache_backend=None,
|
||||
diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
streaming_output=False,
|
||||
)
|
||||
@@ -502,6 +508,81 @@ def test_step_profiler_reports_denoise_step_as_diffuse():
|
||||
class TestRunner:
|
||||
"""DiffusionModelRunner.execute_stepwise"""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("diffusion_kv_mode", "should_cleanup"),
|
||||
[
|
||||
(DiffusionKVCacheMode.DENSE_LEGACY, False),
|
||||
(DiffusionKVCacheMode.PAGED_SCHEDULER, True),
|
||||
],
|
||||
)
|
||||
def test_scheduler_finished_cleanup_only_runs_for_paged_kv(
|
||||
self,
|
||||
mocker: MockerFixture,
|
||||
diffusion_kv_mode: DiffusionKVCacheMode,
|
||||
should_cleanup: bool,
|
||||
):
|
||||
runner = _make_runner(diffusion_kv_mode)
|
||||
cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
|
||||
expected_output = object()
|
||||
mocker.patch.object(runner, "_execute_stepwise_core", return_value=expected_output)
|
||||
scheduler_output = DiffusionSchedulerOutput(
|
||||
step_id=1,
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
finished_req_ids={"finished-request"},
|
||||
num_running_reqs=0,
|
||||
num_waiting_reqs=0,
|
||||
)
|
||||
|
||||
output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
||||
|
||||
assert output is expected_output
|
||||
if should_cleanup:
|
||||
cleanup.assert_called_once_with(["finished-request"])
|
||||
else:
|
||||
cleanup.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("diffusion_kv_mode", "should_cleanup"),
|
||||
[
|
||||
(DiffusionKVCacheMode.DENSE_LEGACY, False),
|
||||
(DiffusionKVCacheMode.PAGED_SCHEDULER, True),
|
||||
],
|
||||
)
|
||||
def test_terminal_cleanup_only_runs_for_paged_kv(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
mocker: MockerFixture,
|
||||
diffusion_kv_mode: DiffusionKVCacheMode,
|
||||
should_cleanup: bool,
|
||||
):
|
||||
runner = _make_runner(diffusion_kv_mode)
|
||||
cleanup = mocker.patch.object(runner, "remove_diffusion_kv_requests")
|
||||
install = mocker.patch.object(runner, "install_diffusion_kv_metadata")
|
||||
req = _make_step_request()
|
||||
scheduler_output = _make_scheduler_output(req, step_id=0)
|
||||
if should_cleanup:
|
||||
scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata = DiffusionKVMetadata(
|
||||
request_id=req.request_id,
|
||||
allocation_generation=1,
|
||||
sequences=(),
|
||||
)
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
first_output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
|
||||
assert first_output.get_request_output(req.request_id).finished is False
|
||||
cleanup.assert_not_called()
|
||||
|
||||
output = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
|
||||
|
||||
assert output.get_request_output(req.request_id).finished is True
|
||||
if should_cleanup:
|
||||
install.assert_called_once_with(scheduler_output.scheduled_new_reqs[0].diffusion_kv_metadata)
|
||||
cleanup.assert_called_once_with([req.request_id])
|
||||
else:
|
||||
install.assert_not_called()
|
||||
cleanup.assert_not_called()
|
||||
|
||||
def test_completes_request_and_clears_state(self, monkeypatch):
|
||||
runner = _make_runner()
|
||||
req = _make_step_request()
|
||||
|
||||
Reference in new issue
Block a user