[Diffusion] Add paged KV cache support for HunyuanImage3 DiT on NPU and GPU (#6563)
Signed-off-by: Acerak01-fy <wfy2003324@163.com> Signed-off-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com> Signed-off-by: zwhzzz0821 <2831474076@qq.com> Co-authored-by: zwhzzz0821 <zwhzzz0821@users.noreply.github.com> Co-authored-by: zwhzzz0821 <2831474076@qq.com>
This commit is contained in:
28 files changed
+1789
-205
No files matched your search
@@ -11,6 +11,7 @@ Tests cover:
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -466,6 +467,120 @@ class TestCurrentDiffusionConfig:
|
||||
|
||||
|
||||
class TestAttentionInitUsesCurrentDiffusionConfig:
|
||||
@pytest.mark.parametrize("mindiesd_available", [False, True])
|
||||
def test_paged_npu_memory_profile_uses_sdpa_only_without_mindiesd(self, monkeypatch, mindiesd_available):
|
||||
attention = Attention.__new__(Attention)
|
||||
attention._scheduler_paged_kv = True
|
||||
attention.paged_kv_cache_role = "primary"
|
||||
attention.backend_pref = "FLASH_ATTN"
|
||||
attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
|
||||
flash_output = torch.ones(1)
|
||||
sdpa_output = torch.zeros(1)
|
||||
dense_flash = Mock(return_value=flash_output)
|
||||
sdpa = Mock(return_value=sdpa_output)
|
||||
attention.attention = SimpleNamespace(forward=dense_flash)
|
||||
attention.sdpa_fallback = SimpleNamespace(forward=sdpa)
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
layer_mod,
|
||||
"get_forward_context",
|
||||
lambda: SimpleNamespace(in_diffusion_kv_memory_profile=True),
|
||||
)
|
||||
monkeypatch.setattr(layer_mod.current_omni_platform, "is_npu", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
layer_mod.current_omni_platform,
|
||||
"supports_diffusion_dense_flash_attention",
|
||||
lambda: mindiesd_available,
|
||||
)
|
||||
|
||||
result = attention._run_local_attention(
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
AttentionMetadata(),
|
||||
)
|
||||
|
||||
if mindiesd_available:
|
||||
assert result is flash_output
|
||||
dense_flash.assert_called_once()
|
||||
sdpa.assert_not_called()
|
||||
else:
|
||||
assert result is sdpa_output
|
||||
dense_flash.assert_not_called()
|
||||
sdpa.assert_called_once()
|
||||
|
||||
def test_dense_flash_import_error_is_not_reclassified(self, monkeypatch):
|
||||
attention = Attention.__new__(Attention)
|
||||
attention._scheduler_paged_kv = False
|
||||
attention.paged_kv_cache_role = "primary"
|
||||
attention.backend_pref = "FLASH_ATTN"
|
||||
attention.attn_backend = SimpleNamespace(supports_piecewise_spans=True, get_name=lambda: "FLASH_ATTN")
|
||||
attention.attention = SimpleNamespace(
|
||||
forward=Mock(side_effect=ModuleNotFoundError("No module named 'mindiesd'", name="mindiesd"))
|
||||
)
|
||||
attention.sdpa_fallback = SimpleNamespace(forward=Mock(side_effect=AssertionError("must not fall back")))
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
|
||||
|
||||
with pytest.raises(ModuleNotFoundError, match="mindiesd"):
|
||||
attention._run_local_attention(
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
torch.zeros(1, dtype=torch.bfloat16),
|
||||
AttentionMetadata(),
|
||||
)
|
||||
|
||||
def test_paged_marker_does_not_change_dense_backend(self, monkeypatch):
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
|
||||
def forward(self, query, key, value, attn_metadata=None):
|
||||
return query
|
||||
|
||||
class _DenseBackend:
|
||||
supports_paged_kv = False
|
||||
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "TORCH_SDPA"
|
||||
|
||||
@staticmethod
|
||||
def get_impl_cls():
|
||||
return _FakeAttentionImpl
|
||||
|
||||
calls = []
|
||||
|
||||
def _fake_get_attn_backend_for_role(*args, **kwargs):
|
||||
calls.append((args, kwargs))
|
||||
return _DenseBackend, None
|
||||
|
||||
monkeypatch.setattr(layer_mod, "get_attn_backend_for_role", _fake_get_attn_backend_for_role)
|
||||
monkeypatch.setattr(layer_mod.SDPABackend, "get_impl_cls", staticmethod(lambda: _FakeAttentionImpl))
|
||||
monkeypatch.setattr(layer_mod, "build_parallel_attention_strategy", lambda **kwargs: object())
|
||||
monkeypatch.setattr(layer_mod, "is_forward_context_available", lambda: False)
|
||||
|
||||
od_config = SimpleNamespace(
|
||||
diffusion_attention_config=AttentionConfig(),
|
||||
diffusion_kv_mode=layer_mod.DiffusionKVCacheMode.DENSE_LEGACY,
|
||||
parallel_config=SimpleNamespace(ring_degree=1),
|
||||
diffusion_kv_cache_dtype=None,
|
||||
diffusion_kv_cache_skip_step_indices=None,
|
||||
diffusion_kv_cache_skip_layer_indices=None,
|
||||
)
|
||||
|
||||
with set_current_diffusion_config(od_config):
|
||||
attention = Attention(
|
||||
num_heads=4,
|
||||
head_size=64,
|
||||
causal=False,
|
||||
softmax_scale=1.0,
|
||||
paged_kv_cache_role="primary",
|
||||
)
|
||||
|
||||
assert len(calls) == 1
|
||||
assert attention.attn_backend is _DenseBackend
|
||||
assert attention.backend_pref is None
|
||||
|
||||
def test_attention_init_uses_current_diffusion_config_without_forward_context(self, monkeypatch):
|
||||
class _FakeAttentionImpl:
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""End-to-end test for ``piecewise_attn`` (CPU).
|
||||
|
||||
Verify that running attention in segments (causal outside full-attn spans,
|
||||
@@ -98,9 +98,12 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
key = query + 100
|
||||
value = query + 200
|
||||
calls = []
|
||||
calls: list[object] = []
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, segment_batch):
|
||||
output_buffer = torch.empty_like(query)
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, segment_batch, segment_output_buffer):
|
||||
assert segment_output_buffer is None
|
||||
calls.append(
|
||||
(
|
||||
[segment.mode for segment in segment_batch.row_segments],
|
||||
@@ -119,9 +122,13 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
output_buffer=output_buffer,
|
||||
)
|
||||
|
||||
assert output is output_buffer
|
||||
torch.testing.assert_close(output, query)
|
||||
assert all(segment.query_range is None for segment in plan.segments)
|
||||
assert not plan.segments_cover_query_contiguously
|
||||
assert calls == [
|
||||
(["causal", "causal"], [2, 4], [0, 1, 6, 7]),
|
||||
(["full", "full"], [5, 7], [2, 3, 4, 8, 9, 10]),
|
||||
@@ -129,6 +136,161 @@ def test_paged_piecewise_runner_uses_packed_indices_and_kv_endpoints():
|
||||
]
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_preserves_homogeneous_batch_without_scatter(monkeypatch: pytest.MonkeyPatch):
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)], [(2, 5)]],
|
||||
query_offsets=[0, 0],
|
||||
query_lens=[6, 6],
|
||||
seq_lens=[6, 6],
|
||||
device=DEVICE,
|
||||
)
|
||||
assert plan.homogeneous_batch_shape == (2, 6)
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
calls: list[torch.Tensor] = []
|
||||
|
||||
def fail_index_copy(*_args, **_kwargs):
|
||||
raise AssertionError("homogeneous piecewise path must not use index_copy_")
|
||||
|
||||
monkeypatch.setattr(torch.Tensor, "index_copy_", fail_index_copy)
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, _metadata, segment_output_buffer):
|
||||
assert segment_key is None
|
||||
assert segment_value is None
|
||||
assert segment_output_buffer is None
|
||||
calls.append(segment_query.clone())
|
||||
return segment_query + 1
|
||||
|
||||
output_buffer = torch.empty_like(query)
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
None,
|
||||
None,
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
output_buffer=output_buffer,
|
||||
use_homogeneous_batch=True,
|
||||
)
|
||||
|
||||
assert output is output_buffer
|
||||
torch.testing.assert_close(output, query + 1)
|
||||
assert [call.shape[0] for call in calls] == [4, 6, 2]
|
||||
torch.testing.assert_close(calls[0], query[[0, 1, 6, 7]])
|
||||
torch.testing.assert_close(calls[1], query[[2, 3, 4, 8, 9, 10]])
|
||||
torch.testing.assert_close(calls[2], query[[5, 11]])
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_keeps_indexed_fallback_for_heterogeneous_layout():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)], [(1, 4)]],
|
||||
query_offsets=[0, 0],
|
||||
query_lens=[6, 6],
|
||||
seq_lens=[6, 6],
|
||||
device=DEVICE,
|
||||
)
|
||||
assert plan.homogeneous_batch_shape is None
|
||||
query = torch.arange(12 * 2 * 4, dtype=torch.float32).reshape(12, 2, 4)
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
None,
|
||||
None,
|
||||
plan,
|
||||
plan.segments,
|
||||
lambda segment_query, *_args: segment_query,
|
||||
use_homogeneous_batch=True,
|
||||
)
|
||||
torch.testing.assert_close(output, query)
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_slices_contiguous_segments():
|
||||
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)
|
||||
key = query + 100
|
||||
value = query + 200
|
||||
expected_ranges = [(0, 2), (2, 5), (5, 6)]
|
||||
calls: list[tuple[int, int]] = []
|
||||
|
||||
def run_segment(segment_query, segment_key, segment_value, _segment_batch, segment_output_buffer):
|
||||
assert segment_output_buffer is None
|
||||
expected_range = expected_ranges[len(calls)]
|
||||
start, end = expected_range
|
||||
calls.append(expected_range)
|
||||
assert segment_query.untyped_storage().data_ptr() == query.untyped_storage().data_ptr()
|
||||
assert segment_key.untyped_storage().data_ptr() == key.untyped_storage().data_ptr()
|
||||
assert segment_value.untyped_storage().data_ptr() == value.untyped_storage().data_ptr()
|
||||
torch.testing.assert_close(segment_query, query[start:end])
|
||||
return segment_query
|
||||
|
||||
output = run_paged_piecewise_plan(query, key, value, plan, plan.segments, run_segment)
|
||||
|
||||
assert calls == expected_ranges
|
||||
assert [segment.query_range for segment in plan.segments] == expected_ranges
|
||||
assert plan.segments_cover_query_contiguously
|
||||
torch.testing.assert_close(output, query)
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_writes_contiguous_segments_into_output_buffer():
|
||||
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_buffer = torch.empty(6, 2, 3)
|
||||
targets = []
|
||||
|
||||
def run_segment(segment_query, _key, _value, _metadata, segment_output_buffer):
|
||||
assert segment_output_buffer is not None
|
||||
assert segment_output_buffer.untyped_storage().data_ptr() == output_buffer.untyped_storage().data_ptr()
|
||||
targets.append(segment_output_buffer)
|
||||
return segment_output_buffer.copy_(segment_query[..., :3])
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
query,
|
||||
query,
|
||||
plan,
|
||||
plan.segments,
|
||||
run_segment,
|
||||
output_buffer=output_buffer,
|
||||
)
|
||||
|
||||
assert output is output_buffer
|
||||
assert [target.shape[0] for target in targets] == [2, 3, 1]
|
||||
torch.testing.assert_close(output, query[..., :3])
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_returns_single_full_segment_directly():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[]],
|
||||
query_offsets=[0],
|
||||
query_lens=[6],
|
||||
seq_lens=[6],
|
||||
device=DEVICE,
|
||||
)
|
||||
query = torch.randn(6, 2, 4)
|
||||
runner_output = query + 1
|
||||
|
||||
output = run_paged_piecewise_plan(
|
||||
query,
|
||||
query,
|
||||
query,
|
||||
plan,
|
||||
plan.segments,
|
||||
lambda *_args: runner_output,
|
||||
)
|
||||
|
||||
assert output is runner_output
|
||||
|
||||
|
||||
def test_paged_piecewise_runner_uses_runner_output_shape():
|
||||
plan = build_paged_piecewise_plan(
|
||||
[[(2, 5)]],
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -485,8 +485,16 @@ def test_valid_sequence_and_context_install_into_native_rows(monkeypatch: pytest
|
||||
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)
|
||||
assert (sequence_binding.row_index, sequence_binding.max_seq_len, sequence_binding.block_ids) == (
|
||||
0,
|
||||
9,
|
||||
((1, 2, 3), (4, 5)),
|
||||
)
|
||||
assert (context_binding.row_index, context_binding.max_seq_len, context_binding.block_ids) == (
|
||||
1,
|
||||
5,
|
||||
((6, 7), (8,)),
|
||||
)
|
||||
|
||||
|
||||
def test_block_table_mutations_invalidate_prepared_attention_batches(
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -36,7 +36,7 @@ def test_non_mapping_omni_kv_config_is_rejected() -> None:
|
||||
OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler")
|
||||
|
||||
|
||||
def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None:
|
||||
def test_paged_scheduler_rejects_dense_connector_kv_receive() -> None:
|
||||
with pytest.raises(ValueError, match="does not support imported AR KV"):
|
||||
OmniDiffusionConfig.from_kwargs(
|
||||
diffusion_kv_mode="paged_scheduler",
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
@@ -290,6 +290,7 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
|
||||
dtype=torch.bfloat16,
|
||||
non_causal=True,
|
||||
)
|
||||
model_limit = vllm_config.model_config.max_model_len
|
||||
|
||||
initialization.build_native_kv_cache_configs(
|
||||
vllm_config,
|
||||
@@ -298,4 +299,8 @@ def test_minus_one_max_model_len_is_auto_fitted_by_native_cache_sizing() -> None
|
||||
)
|
||||
|
||||
assert vllm_config.model_config.original_max_model_len == -1
|
||||
assert vllm_config.model_config.max_model_len == 256
|
||||
# Sixteen physical pages provide ``16 * block_size`` tokens. The native
|
||||
# backend chooses the block geometry, so keep this assertion backend
|
||||
# independent instead of assuming the generic 16-token block size.
|
||||
expected_max_model_len = min(model_limit, spec.block_size * 16)
|
||||
assert vllm_config.model_config.max_model_len == expected_max_model_len
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -20,12 +20,16 @@ from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.diffusion_kv import paged_attention_adapter as adapter_module
|
||||
from vllm_omni.diffusion.diffusion_kv.paged_attention_adapter import (
|
||||
DiffusionPagedAttentionAdapter,
|
||||
DiffusionPagedAttentionMetadata,
|
||||
DiffusionPagedAttentionRow,
|
||||
DiffusionPagedAttentionRowBinding,
|
||||
DiffusionPagedAttentionRuntime,
|
||||
)
|
||||
from vllm_omni.diffusion.forward_context import (
|
||||
get_forward_context,
|
||||
override_paged_kv_adapter,
|
||||
set_forward_context,
|
||||
set_forward_context_denoise_step_idx,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
@@ -122,8 +126,8 @@ class _FakeNativeImpl:
|
||||
def do_kv_cache_update(
|
||||
self,
|
||||
_layer,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
key: torch.Tensor | None,
|
||||
value: torch.Tensor | None,
|
||||
_kv_cache,
|
||||
slot_mapping: torch.Tensor,
|
||||
) -> None:
|
||||
@@ -155,6 +159,11 @@ def _make_adapter(
|
||||
capacity: int = 16,
|
||||
) -> tuple[DiffusionPagedAttentionAdapter, _FakeBlockTables, _FakeLayer, list[tuple]]:
|
||||
events: list[tuple] = []
|
||||
monkeypatch.setattr(
|
||||
adapter_module.current_omni_platform,
|
||||
"requires_diffusion_paged_kv_prewrite",
|
||||
lambda: False,
|
||||
)
|
||||
block_tables = _FakeBlockTables()
|
||||
block_tables.max_num_batched_tokens = capacity
|
||||
block_tables.num_blocks.np.fill(math.ceil(capacity / block_tables.block_sizes[0]))
|
||||
@@ -301,6 +310,47 @@ def test_paged_adapter_accepts_native_gqa_shapes(monkeypatch: pytest.MonkeyPatch
|
||||
assert context.value_write.shape == (3, 8, 4)
|
||||
|
||||
|
||||
def test_paged_adapter_preserves_strided_kv_for_platform_writer(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
|
||||
)
|
||||
packed_qkv = torch.randn(1, 5, 3 * 2 * 4)
|
||||
query, key, value = packed_qkv.split(2 * 4, dim=-1)
|
||||
query = query.reshape(1, 5, 2, 4)
|
||||
key = key.reshape(1, 5, 2, 4)
|
||||
value = value.reshape(1, 5, 2, 4)
|
||||
assert not key.is_contiguous()
|
||||
assert not value.is_contiguous()
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
||||
|
||||
assert context.query.is_contiguous()
|
||||
assert not context.key_write.is_contiguous()
|
||||
assert not context.value_write.is_contiguous()
|
||||
assert context.key_write.data_ptr() == key.data_ptr()
|
||||
assert context.value_write.data_ptr() == value.data_ptr()
|
||||
|
||||
|
||||
def test_paged_adapter_packs_noncontiguous_kv_head_layout(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=5, seq_len=5)]
|
||||
)
|
||||
query = torch.randn(1, 5, 2, 4)
|
||||
key = torch.randn(1, 5, 4, 2).transpose(-1, -2)
|
||||
value = torch.randn(1, 5, 4, 2).transpose(-1, -2)
|
||||
assert key.shape == (1, 5, 2, 4)
|
||||
assert key.stride(-1) != 1
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", query, key, value)
|
||||
|
||||
assert context.key_write.is_contiguous()
|
||||
assert context.value_write.is_contiguous()
|
||||
|
||||
|
||||
def test_omni_paged_backend_defers_backend_owned_cache_update(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
||||
layer.attn_backend.forward_includes_kv_cache_update = True
|
||||
@@ -350,10 +400,48 @@ def test_omni_piecewise_backend_defers_backend_owned_cache_update(monkeypatch: p
|
||||
assert [metadata.query_start_loc_cpu.tolist() for metadata in segment_metadata] == [[0, 1], [0, 2], [0, 1]]
|
||||
assert [metadata.positions.tolist() for metadata in segment_metadata] == [[0], [1, 2], [3]]
|
||||
assert [metadata.slot_mappings.tolist() for metadata in segment_metadata] == [[[0]], [[1, 2]], [[3]]]
|
||||
for event in events[1:]:
|
||||
build_kwargs = event[1]
|
||||
assert build_kwargs["positions"].untyped_storage().data_ptr() == batch.positions.untyped_storage().data_ptr()
|
||||
assert (
|
||||
build_kwargs["slot_mappings"].untyped_storage().data_ptr()
|
||||
== batch.slot_mappings.untyped_storage().data_ptr()
|
||||
)
|
||||
|
||||
|
||||
def test_platform_prewrite_updates_once_then_piecewise_reads_cache(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, _ = _make_adapter(monkeypatch)
|
||||
layer.attn_backend.forward_includes_kv_cache_update = True
|
||||
monkeypatch.setattr(
|
||||
adapter_module.current_omni_platform,
|
||||
"requires_diffusion_paged_kv_prewrite",
|
||||
lambda: True,
|
||||
)
|
||||
adapter.resolve_row = lambda *_args: DiffusionPagedAttentionRowBinding(
|
||||
row_index=2,
|
||||
max_seq_len=16,
|
||||
block_ids=((3, 7, 9, 11),),
|
||||
)
|
||||
row = DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)
|
||||
batch = adapter.prepare_batch([row])
|
||||
qkv = torch.randn(1, 4, 2, 4)
|
||||
metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv, omni_attn_metadata=metadata)
|
||||
output = _run_omni_paged_backend(context)
|
||||
|
||||
assert torch.equal(output, qkv)
|
||||
assert layer.native_events == ["update", "forward", "forward", "forward"]
|
||||
assert len(layer.updates) == 1
|
||||
assert all(call[1] is None and call[2] is None for call in layer.calls)
|
||||
|
||||
|
||||
def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, events = _make_adapter(monkeypatch)
|
||||
# Exercise the Ascend-only homogeneous batch path. CUDA keeps its native
|
||||
# output-buffer contract and is covered by the GPU adapter tests.
|
||||
monkeypatch.setattr(adapter_module.current_omni_platform, "is_npu", lambda: True)
|
||||
batch = adapter.prepare_batch(
|
||||
[
|
||||
DiffusionPagedAttentionRow(
|
||||
@@ -450,6 +538,48 @@ def test_omni_paged_backend_runs_hunyuan_piecewise_segments(monkeypatch: pytest.
|
||||
]
|
||||
|
||||
|
||||
def test_piecewise_metadata_snapshots_reused_native_scheduler_buffer(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
shared_scheduler_metadata = torch.zeros(1, dtype=torch.int32)
|
||||
build_id = 0
|
||||
|
||||
def build_metadata(**kwargs):
|
||||
nonlocal build_id
|
||||
shared_scheduler_metadata.fill_(build_id)
|
||||
metadata = SimpleNamespace(
|
||||
build_id=build_id,
|
||||
causal=kwargs["causal"],
|
||||
seq_lens=kwargs["seq_lens"].clone(),
|
||||
query_start_loc_cpu=kwargs["query_start_loc_cpu"].clone(),
|
||||
positions=kwargs["positions"].clone(),
|
||||
slot_mappings=kwargs["slot_mappings"].clone(),
|
||||
scheduler_metadata=shared_scheduler_metadata,
|
||||
prefix_scheduler_metadata=None,
|
||||
)
|
||||
build_id += 1
|
||||
return {"layer-0": metadata}
|
||||
|
||||
monkeypatch.setattr(adapter_module.current_omni_platform, "build_diffusion_kv_attn_metadata", build_metadata)
|
||||
batch = adapter.prepare_batch(
|
||||
[DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=4, seq_len=4)]
|
||||
)
|
||||
qkv = torch.randn(1, 4, 2, 4)
|
||||
metadata = SimpleNamespace(full_attn_spans=[[(1, 3)]], extra={})
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context(
|
||||
"layer-0",
|
||||
qkv,
|
||||
qkv,
|
||||
qkv,
|
||||
omni_attn_metadata=metadata,
|
||||
)
|
||||
|
||||
segment_schedules = [item.scheduler_metadata for item in context.piecewise_native_metadata]
|
||||
assert [item.item() for item in segment_schedules] == [1, 2, 3]
|
||||
assert len({item.data_ptr() for item in segment_schedules}) == 3
|
||||
|
||||
|
||||
def test_omni_paged_backend_treats_empty_full_spans_as_causal(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, layer, events = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
@@ -851,6 +981,28 @@ def test_forward_rejects_batched_layout_that_does_not_match_rows(monkeypatch: py
|
||||
adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
||||
|
||||
|
||||
def test_forward_packs_and_restores_right_padded_heterogeneous_rows(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
[
|
||||
DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=1, seq_len=1),
|
||||
DiffusionPagedAttentionRow(request_id="req-1", sequence_id=1, query_len=3, seq_len=3),
|
||||
]
|
||||
)
|
||||
qkv = torch.arange(2 * 3 * 2 * 4, dtype=torch.float32).reshape(2, 3, 2, 4)
|
||||
|
||||
with adapter.activate(batch):
|
||||
context = adapter.prepare_layer_context("layer-0", qkv, qkv, qkv)
|
||||
|
||||
assert context.output_scatter_indices.tolist() == [0, 3, 4, 5]
|
||||
torch.testing.assert_close(context.query, qkv.reshape(-1, 2, 4)[[0, 3, 4, 5]])
|
||||
restored = context.restore_output(context.query)
|
||||
assert restored.shape == qkv.shape
|
||||
torch.testing.assert_close(restored[0, 0], qkv[0, 0])
|
||||
assert torch.count_nonzero(restored[0, 1:]) == 0
|
||||
torch.testing.assert_close(restored[1], qkv[1])
|
||||
|
||||
|
||||
def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
batch = adapter.prepare_batch(
|
||||
@@ -878,6 +1030,41 @@ def test_forward_rejects_padded_full_kv(monkeypatch: pytest.MonkeyPatch) -> None
|
||||
adapter.prepare_layer_context("layer-0", query, full_kv, full_kv)
|
||||
|
||||
|
||||
def test_runner_runtime_switches_from_prefill_to_reusable_denoise_metadata(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch)
|
||||
metadata = DiffusionPagedAttentionMetadata(
|
||||
prefill_rows=(DiffusionPagedAttentionRow(request_id="req-0", sequence_id=0, query_len=6, seq_len=6),),
|
||||
denoise_rows=(
|
||||
DiffusionPagedAttentionRow(
|
||||
request_id="req-0",
|
||||
sequence_id=0,
|
||||
query_len=2,
|
||||
seq_len=5,
|
||||
kv_start_pos=3,
|
||||
),
|
||||
),
|
||||
)
|
||||
runtime = DiffusionPagedAttentionRuntime(adapter, metadata)
|
||||
|
||||
with set_forward_context(paged_kv_runtime=runtime), runtime.activate():
|
||||
set_forward_context_denoise_step_idx(0)
|
||||
assert Attention._active_paged_kv_adapter() is adapter
|
||||
assert adapter._active_batch.rows == metadata.prefill_rows
|
||||
|
||||
set_forward_context_denoise_step_idx(1)
|
||||
assert Attention._active_paged_kv_adapter() is adapter
|
||||
denoise_batch = adapter._active_batch
|
||||
assert denoise_batch.rows == metadata.denoise_rows
|
||||
|
||||
set_forward_context_denoise_step_idx(2)
|
||||
assert Attention._active_paged_kv_adapter() is adapter
|
||||
assert adapter._active_batch is denoise_batch
|
||||
|
||||
assert adapter._active_batch is None
|
||||
|
||||
|
||||
def test_causal_batch_rejects_non_suffix_write(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
adapter, _, _, _ = _make_adapter(monkeypatch, non_causal=False)
|
||||
row = DiffusionPagedAttentionRow(
|
||||
@@ -992,7 +1179,8 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
selected_backends = []
|
||||
impl_cls = Mock(return_value=SimpleNamespace(forward=Mock()))
|
||||
specialized_backends = []
|
||||
impl_cls = Mock(return_value=SimpleNamespace(forward=Mock(), do_kv_cache_update=Mock()))
|
||||
native_backend = SimpleNamespace(
|
||||
get_name=lambda: "ASCEND",
|
||||
indexes_kv_by_block_stride=lambda: True,
|
||||
@@ -1003,16 +1191,28 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
original_backend_per_kind = {"full": object()}
|
||||
config = SimpleNamespace(
|
||||
attention_config=SimpleNamespace(backend=None, backend_per_kind=original_backend_per_kind),
|
||||
parallel_config=SimpleNamespace(prefill_context_parallel_size=2),
|
||||
model_config=SimpleNamespace(dtype=torch.float16),
|
||||
cache_config=SimpleNamespace(cache_dtype="auto"),
|
||||
)
|
||||
|
||||
def select_backend(**_kwargs):
|
||||
selected_backends.append((config.attention_config.backend, config.attention_config.backend_per_kind))
|
||||
selected_backends.append(
|
||||
(
|
||||
config.attention_config.backend,
|
||||
config.attention_config.backend_per_kind,
|
||||
config.parallel_config.prefill_context_parallel_size,
|
||||
)
|
||||
)
|
||||
return native_backend
|
||||
|
||||
monkeypatch.setattr(adapter_module, "get_attn_backend", select_backend)
|
||||
monkeypatch.setattr(adapter_module, "set_current_vllm_config", lambda _config: nullcontext())
|
||||
monkeypatch.setattr(
|
||||
adapter_module.current_omni_platform,
|
||||
"get_diffusion_paged_kv_attn_backend",
|
||||
lambda backend, *, ulysses_degree: specialized_backends.append((backend, ulysses_degree)) or backend,
|
||||
)
|
||||
layer = SimpleNamespace(num_heads=8, softmax_scale=0.125)
|
||||
spec = FullAttentionSpec(
|
||||
block_size=16,
|
||||
@@ -1031,9 +1231,11 @@ def test_layer_adapter_accepts_platform_native_backend_and_uses_rank_local_heads
|
||||
ulysses_degree=2,
|
||||
)
|
||||
|
||||
assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {})]
|
||||
assert selected_backends == [(adapter_module.AttentionBackendEnum.FLASH_ATTN, {}, 2)]
|
||||
assert specialized_backends == [(native_backend, 2)]
|
||||
assert config.attention_config.backend is None
|
||||
assert config.attention_config.backend_per_kind is original_backend_per_kind
|
||||
assert config.parallel_config.prefill_context_parallel_size == 2
|
||||
assert native_layer.num_heads == 4
|
||||
assert native_layer.num_kv_heads == 2
|
||||
assert native_layer.spec.num_kv_heads == 2
|
||||
@@ -1094,6 +1296,51 @@ def test_omni_attention_wraps_paged_kernel_with_sp_hooks() -> None:
|
||||
assert torch.equal(output, torch.full_like(original, 6))
|
||||
|
||||
|
||||
def test_omni_attention_strips_paged_ulysses_padding_around_kernel() -> None:
|
||||
events: list[tuple[str, tuple[int, ...]]] = []
|
||||
|
||||
class Strategy:
|
||||
name = "ulysses"
|
||||
|
||||
def pre_attention(self, query, key, value, metadata):
|
||||
return query, key, value, metadata, SimpleNamespace(joint_len=0, joint_strategy="front")
|
||||
|
||||
def post_attention(self, output, _ctx):
|
||||
events.append(("post", tuple(output.shape)))
|
||||
return output
|
||||
|
||||
class Adapter:
|
||||
def prepare_layer_context(self, _layer_name, query, key, value, *, omni_attn_metadata):
|
||||
assert omni_attn_metadata is None
|
||||
events.append(("kernel_input", tuple(query.shape)))
|
||||
assert query[:, -1].tolist() == [[[3.0]]]
|
||||
assert torch.equal(query, key)
|
||||
assert torch.equal(query, value)
|
||||
return SimpleNamespace(query=query)
|
||||
|
||||
class Backend:
|
||||
def forward_paged(self, context):
|
||||
return context.query
|
||||
|
||||
layer = Attention.__new__(Attention)
|
||||
nn.Module.__init__(layer)
|
||||
layer.prefix = "layer-0"
|
||||
layer.paged_kv_cache_role = "primary"
|
||||
layer.attn_backend = SimpleNamespace(supports_paged_kv=True, get_name=lambda: "FLASH_ATTN")
|
||||
layer.attention = Backend()
|
||||
layer.use_ring = False
|
||||
layer._no_parallel_strategy = object()
|
||||
layer._get_active_parallel_strategy = lambda: Strategy()
|
||||
tensor = torch.arange(5, dtype=torch.float32).reshape(1, 5, 1, 1)
|
||||
|
||||
with set_forward_context(), override_paged_kv_adapter(Adapter()):
|
||||
get_forward_context().sp_padding_size = 1
|
||||
output = layer._forward_impl(tensor, tensor, tensor)
|
||||
|
||||
assert events == [("kernel_input", (1, 4, 1, 1)), ("post", (1, 5, 1, 1))]
|
||||
assert output.flatten().tolist() == [0.0, 1.0, 2.0, 3.0, 0.0]
|
||||
|
||||
|
||||
def test_omni_attention_rejects_paged_request_for_backend_without_paged_support() -> None:
|
||||
class Adapter:
|
||||
def prepare_layer_context(self, *args, **kwargs):
|
||||
@@ -1145,3 +1392,15 @@ def test_omni_attention_keeps_dense_kernel_without_active_adapter() -> None:
|
||||
|
||||
assert events == ["pre", "dense", "post"]
|
||||
assert output is qkv
|
||||
|
||||
|
||||
def test_omni_attention_rejects_scheduler_paged_forward_without_active_adapter() -> None:
|
||||
layer = Attention.__new__(Attention)
|
||||
nn.Module.__init__(layer)
|
||||
layer._scheduler_paged_kv = True
|
||||
layer.paged_kv_cache_role = "primary"
|
||||
layer._get_active_parallel_strategy = lambda: object()
|
||||
qkv = torch.zeros(1, 2, 2, 4)
|
||||
|
||||
with set_forward_context(), pytest.raises(RuntimeError, match="without an active Worker adapter"):
|
||||
layer._forward_impl(qkv, qkv, qkv)
|
||||
@@ -154,8 +154,17 @@ def test_adapter_executes_native_paged_attention_on_non_contiguous_blocks() -> N
|
||||
|
||||
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)
|
||||
packed_qkv = torch.randn(
|
||||
19,
|
||||
3 * _NUM_HEADS * _HEAD_SIZE,
|
||||
dtype=torch.float16,
|
||||
device=device,
|
||||
)
|
||||
_, key_flat, value_flat = packed_qkv.split(_NUM_HEADS * _HEAD_SIZE, dim=-1)
|
||||
key = key_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
|
||||
value = value_flat.reshape(19, _NUM_HEADS, _HEAD_SIZE)
|
||||
assert not key.is_contiguous()
|
||||
assert not value.is_contiguous()
|
||||
omni_backend = FlashAttentionImpl(
|
||||
num_heads=_NUM_HEADS,
|
||||
head_size=_HEAD_SIZE,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
@@ -42,9 +42,10 @@ class _RunnerKVBackend:
|
||||
self.kv_cache_config = config
|
||||
|
||||
|
||||
def _attention(*, enabled: bool) -> Attention:
|
||||
def _attention(*, enabled: bool, prefix: str = "image_attention") -> Attention:
|
||||
attention = Attention.__new__(Attention)
|
||||
nn.Module.__init__(attention)
|
||||
attention.prefix = prefix
|
||||
attention.paged_kv_cache_role = "primary" if enabled else None
|
||||
attention.paged_kv_cache_dtype = torch.bfloat16
|
||||
attention.num_kv_heads = 2
|
||||
@@ -82,6 +83,23 @@ def test_runner_discovers_native_spec_from_loaded_attention() -> None:
|
||||
assert spec.non_causal is True
|
||||
|
||||
|
||||
def test_runner_uses_attention_prefix_as_canonical_layer_name() -> None:
|
||||
runner = _runner(_attention(enabled=True, prefix="layers.0.self_attn.image_attn.attn"))
|
||||
|
||||
specs = runner.get_kv_cache_spec()
|
||||
|
||||
assert set(specs) == {"layers.0.self_attn.image_attn.attn"}
|
||||
|
||||
|
||||
def test_runner_rejects_duplicate_attention_prefixes() -> None:
|
||||
runner = _runner(_attention(enabled=True, prefix="shared"))
|
||||
second_attention = _attention(enabled=True, prefix="shared")
|
||||
runner.pipeline.second_image_attention = second_attention
|
||||
|
||||
with pytest.raises(RuntimeError, match="Duplicate canonical paged Diffusion Attention prefix 'shared'"):
|
||||
runner.get_kv_cache_spec()
|
||||
|
||||
|
||||
def test_runner_rejects_paged_mode_without_cache_enabled_attention() -> None:
|
||||
runner = _runner(_attention(enabled=False))
|
||||
|
||||
|
||||
@@ -42,7 +42,10 @@ def make_metadata(request_id: str = "req-0") -> DiffusionKVMetadata:
|
||||
|
||||
def make_runner(mode: DiffusionKVCacheMode) -> DiffusionModelRunner:
|
||||
runner = object.__new__(DiffusionModelRunner)
|
||||
runner.od_config = SimpleNamespace(diffusion_kv_mode=mode)
|
||||
runner.od_config = SimpleNamespace(
|
||||
diffusion_kv_mode=mode,
|
||||
parallel_config=SimpleNamespace(cfg_parallel_size=1),
|
||||
)
|
||||
return runner
|
||||
|
||||
|
||||
@@ -85,6 +88,51 @@ def test_new_request_data_carries_scheduler_allocation_atomically() -> None:
|
||||
assert new_req.diffusion_kv_metadata is metadata
|
||||
|
||||
|
||||
def test_runner_builds_prefill_and_denoise_rows_from_scheduler_metadata() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
metadata = DiffusionKVMetadata(
|
||||
request_id="req-0",
|
||||
allocation_generation=1,
|
||||
sequences=(
|
||||
DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
|
||||
DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
|
||||
),
|
||||
)
|
||||
|
||||
attn_metadata = runner._build_paged_attention_metadata([metadata])
|
||||
|
||||
assert [
|
||||
(row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
|
||||
for row in attn_metadata.prefill_rows
|
||||
] == [("req-0", 0, 0, 8, 8), ("req-0", 1, 0, 9, 9)]
|
||||
assert [
|
||||
(row.request_id, row.sequence_id, row.kv_start_pos, row.query_len, row.seq_len)
|
||||
for row in attn_metadata.denoise_rows
|
||||
] == [("req-0", 0, 4, 2, 6), ("req-0", 1, 5, 2, 7)]
|
||||
|
||||
|
||||
def test_runner_selects_only_local_cfg_parallel_row(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner.od_config.parallel_config.cfg_parallel_size = 2
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.worker.diffusion_model_runner.get_classifier_free_guidance_rank",
|
||||
lambda: 1,
|
||||
)
|
||||
metadata = DiffusionKVMetadata(
|
||||
request_id="req-0",
|
||||
allocation_generation=1,
|
||||
sequences=(
|
||||
DiffusionKVSequenceMetadata(0, 4, 2, 8, ([1, 2],)),
|
||||
DiffusionKVSequenceMetadata(1, 5, 2, 9, ([3, 4],)),
|
||||
),
|
||||
)
|
||||
|
||||
attn_metadata = runner._build_paged_attention_metadata([metadata])
|
||||
|
||||
assert [row.sequence_id for row in attn_metadata.prefill_rows] == [1]
|
||||
assert [row.sequence_id for row in attn_metadata.denoise_rows] == [1]
|
||||
|
||||
|
||||
def test_paged_request_without_metadata_fails_before_request_forward() -> None:
|
||||
runner = make_runner(DiffusionKVCacheMode.PAGED_SCHEDULER)
|
||||
runner._execute_request_list = Mock()
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -270,6 +270,17 @@ def test_grouped_denoise_allows_sdpa_attention_backend():
|
||||
pipeline._ensure_grouped_attention_backend_supported(2)
|
||||
|
||||
|
||||
def test_scheduler_paged_step_execution_is_rejected():
|
||||
pipeline = _pipeline()
|
||||
pipeline.od_config.diffusion_kv_mode = hy3_module.DiffusionKVCacheMode.PAGED_SCHEDULER
|
||||
state = _state("paged-step", 0)
|
||||
|
||||
with pytest.raises(ValueError, match="request-level execution only"):
|
||||
pipeline.prepare_encode(state)
|
||||
with pytest.raises(ValueError, match="request-level execution only"):
|
||||
pipeline.denoise_step(InputBatch.make_batch([state]))
|
||||
|
||||
|
||||
def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches():
|
||||
pipeline = _pipeline()
|
||||
pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {})
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for ImageKVCacheManager.
|
||||
|
||||
Covers: cache → reuse flow, AR KV injection, CFG (sequential & parallel), SP, cross-request isolation.
|
||||
@@ -37,12 +37,21 @@ class MockAttention(nn.Module):
|
||||
super().__init__()
|
||||
self.paged_kv_active = False
|
||||
self.calls = []
|
||||
self.paged_calls = []
|
||||
self.paged_metadata = []
|
||||
|
||||
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))
|
||||
from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
|
||||
|
||||
if is_forward_context_available() and get_forward_context().paged_kv_adapter is not None:
|
||||
self.paged_calls.append((query, key, value))
|
||||
self.paged_metadata.append(attn_metadata)
|
||||
if attn_metadata is not None and attn_metadata.joint_query is not None:
|
||||
return torch.cat((attn_metadata.joint_query, query), dim=1)
|
||||
return query
|
||||
|
||||
|
||||
@@ -105,6 +114,7 @@ def _call_mgr(
|
||||
shard_image_size=None,
|
||||
gen_timestep_scatter_index=None,
|
||||
position_ids=None,
|
||||
full_attn_spans=None,
|
||||
):
|
||||
query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
|
||||
attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
|
||||
@@ -121,6 +131,7 @@ def _call_mgr(
|
||||
shard_image_size=shard_image_size,
|
||||
gen_timestep_scatter_index=gen_timestep_scatter_index,
|
||||
position_ids=position_ids,
|
||||
full_attn_spans=full_attn_spans,
|
||||
)
|
||||
|
||||
|
||||
@@ -132,6 +143,83 @@ def test_cache_manager_registers_attention_without_adding_dense_state() -> None:
|
||||
assert mgr.state_dict() == {}
|
||||
|
||||
|
||||
def test_scheduler_paged_kv_runs_piecewise_for_first_and_later_steps() -> None:
|
||||
from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
|
||||
|
||||
mgr = _make_cache_mgr()
|
||||
mgr.attn.paged_kv_active = True
|
||||
bs = 2
|
||||
prefix_lens = torch.tensor([[2], [4]], dtype=torch.long)
|
||||
first_q_len = 12
|
||||
first_k, first_v = _make_known_kv(bs * first_q_len, base=1.0)
|
||||
full_attn_spans = [[(2, 12)], [(2, 12)]]
|
||||
|
||||
with set_forward_context(), override_paged_kv_adapter(object()):
|
||||
_call_mgr(
|
||||
mgr,
|
||||
bs,
|
||||
q_len=first_q_len,
|
||||
seq_len=first_q_len,
|
||||
key_flat=first_k,
|
||||
value_flat=first_v,
|
||||
first_step=True,
|
||||
gen_timestep_scatter_index=prefix_lens,
|
||||
full_attn_spans=full_attn_spans,
|
||||
)
|
||||
|
||||
assert mgr.image_kv_cache_map is None
|
||||
assert mgr.image_kv_cache_lens is None
|
||||
assert len(mgr.attn.paged_calls) == 1
|
||||
assert mgr.attn.paged_calls[0][0].shape == (bs, first_q_len, NUM_HEADS, HEAD_DIM)
|
||||
assert mgr.attn.paged_calls[0][1].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
assert mgr.attn.paged_calls[0][2].shape == (bs, first_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
assert mgr.attn.paged_metadata[0].full_attn_spans == full_attn_spans
|
||||
|
||||
current_q_len = IMAGE_TOKEN_LEN
|
||||
current_k, current_v = _make_known_kv(bs * current_q_len, base=50.0)
|
||||
_call_mgr(
|
||||
mgr,
|
||||
bs,
|
||||
q_len=current_q_len,
|
||||
seq_len=prefix_lens[-1].item() + current_q_len,
|
||||
key_flat=current_k,
|
||||
value_flat=current_v,
|
||||
position_ids=torch.arange(4, 4 + current_q_len).repeat(bs, 1),
|
||||
full_attn_spans=full_attn_spans,
|
||||
)
|
||||
|
||||
assert len(mgr.attn.paged_calls) == 2
|
||||
assert mgr.attn.paged_calls[1][0].shape == (bs, current_q_len, NUM_HEADS, HEAD_DIM)
|
||||
assert mgr.attn.paged_calls[1][1].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
assert mgr.attn.paged_calls[1][2].shape == (bs, current_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
assert mgr.attn.paged_metadata[1].full_attn_spans == full_attn_spans
|
||||
|
||||
|
||||
def test_cache_manager_rejects_dense_imported_prefix_for_paged_attention() -> None:
|
||||
from vllm_omni.diffusion.forward_context import override_paged_kv_adapter, set_forward_context
|
||||
|
||||
mgr = _make_cache_mgr()
|
||||
mgr.attn.paged_kv_active = True
|
||||
prefix_len = 4
|
||||
prefix_key, prefix_value = _make_known_kv(prefix_len, base=100.0)
|
||||
mgr._injected_ar_kv = [(prefix_key, prefix_value)]
|
||||
current_key, current_value = _make_known_kv(3, base=1.0)
|
||||
|
||||
with set_forward_context(), override_paged_kv_adapter(object()):
|
||||
with pytest.raises(NotImplementedError, match="does not support imported AR KV"):
|
||||
_call_mgr(
|
||||
mgr,
|
||||
bs=1,
|
||||
q_len=3,
|
||||
seq_len=7,
|
||||
key_flat=current_key,
|
||||
value_flat=current_value,
|
||||
first_step=True,
|
||||
gen_timestep_scatter_index=torch.tensor([[0]]),
|
||||
full_attn_spans=[[(4, 7)]],
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("paged_kv_active", "expected_kv_heads"),
|
||||
[(False, NUM_HEADS), (True, NUM_KV_HEADS)],
|
||||
@@ -155,6 +243,7 @@ def test_gqa_kv_stays_compressed_for_paged_attention(
|
||||
first_step=True,
|
||||
num_image_tokens=0,
|
||||
gen_timestep_scatter_index=_gen_timestep_index(1, 0),
|
||||
full_attn_spans=[[(0, q_len)]] if paged_kv_active else None,
|
||||
)
|
||||
|
||||
_, key_input, value_input = mgr.attn.calls[-1]
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
@@ -533,7 +533,14 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
|
||||
return original_execute(*args, **kwargs)
|
||||
|
||||
runner._execute_request_list = execute
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
forward_context_calls = []
|
||||
|
||||
@contextmanager
|
||||
def record_forward_context(*args, **kwargs):
|
||||
forward_context_calls.append((args, kwargs))
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
|
||||
reset_peak_memory_stats = Mock()
|
||||
monkeypatch.setattr(
|
||||
model_runner_module.current_omni_platform,
|
||||
@@ -549,6 +556,8 @@ def test_profile_run_executes_forward_without_scheduler_kv_validation(monkeypatc
|
||||
|
||||
assert runner.pipeline.forward_calls == 1
|
||||
assert record_names == ["pipeline_memory_profile"]
|
||||
assert len(forward_context_calls) == 1
|
||||
assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
|
||||
runner._validate_diffusion_kv_metadata.assert_not_called()
|
||||
reset_peak_memory_stats.assert_not_called()
|
||||
synchronize.assert_called_once_with()
|
||||
@@ -575,7 +584,14 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
|
||||
|
||||
runner.pipeline.denoise_step = denoise_step
|
||||
runner._validate_diffusion_kv_metadata = Mock(side_effect=AssertionError("profile must bypass admission"))
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
forward_context_calls = []
|
||||
|
||||
@contextmanager
|
||||
def record_forward_context(*args, **kwargs):
|
||||
forward_context_calls.append((args, kwargs))
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", record_forward_context)
|
||||
reset_peak_memory_stats = Mock()
|
||||
monkeypatch.setattr(
|
||||
model_runner_module.current_omni_platform,
|
||||
@@ -592,6 +608,8 @@ def test_profile_run_executes_maximum_step_batch_without_resetting_peak(monkeypa
|
||||
assert observed_batch_rows == [(2, 2)]
|
||||
assert runner.state_cache == {}
|
||||
assert runner.input_batch is None
|
||||
assert len(forward_context_calls) == 1
|
||||
assert forward_context_calls[0][1]["in_diffusion_kv_memory_profile"] is True
|
||||
runner._validate_diffusion_kv_metadata.assert_not_called()
|
||||
reset_peak_memory_stats.assert_not_called()
|
||||
synchronize.assert_called_once_with()
|
||||
|
||||
Reference in new issue
Block a user