278 lines
11 KiB
Python
278 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
"""Native multimodal block hashing and diffusion-specific identity helpers."""
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
|
|
from vllm.utils.hashing import get_hash_fn_by_name
|
|
from vllm.v1.core.kv_cache_utils import BlockHash, generate_block_hash_extra_keys, hash_block_tokens
|
|
|
|
import vllm_omni.diffusion.diffusion_engine as engine_module
|
|
import vllm_omni.diffusion.diffusion_kv.kv_cache_utils as kv_utils
|
|
import vllm_omni.diffusion.diffusion_kv.request as request_module
|
|
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
|
from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
|
|
from vllm_omni.diffusion.diffusion_kv.kv_cache_utils import get_cache_namespace, hash_prefix_cache_value
|
|
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
|
|
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
|
|
|
|
|
def _request(*, num_rows=1, seed=7, adapter_id=None, scale=1.0):
|
|
request = OmniDiffusionRequest(
|
|
request_id="fake", prompt="fake model", sampling_params=OmniDiffusionSamplingParams(seed=seed)
|
|
)
|
|
request.sampling_params.lora_request = None if adapter_id is None else SimpleNamespace(lora_int_id=adapter_id)
|
|
request.sampling_params.lora_scale = scale
|
|
request.diffusion_kv_requests = tuple(
|
|
DiffusionKVRequest(f"fake/{row}", sequence_id=row, prefix_len=8, target_len=4, seq_len=12)
|
|
for row in range(num_rows)
|
|
)
|
|
return request
|
|
|
|
|
|
def _feature(*, value="image", start=4, end=8):
|
|
return MultiModalFeatureSpec(
|
|
data=None,
|
|
modality="image",
|
|
identifier=hash_prefix_cache_value((value, start, end)).hex(),
|
|
mm_position=PlaceholderRange(offset=start, length=end - start),
|
|
)
|
|
|
|
|
|
def _prepare(request, feature=None):
|
|
namespace = get_cache_namespace("fake-model-v1", request.sampling_params)
|
|
for row in request.diffusion_kv_requests:
|
|
row.cache_token_ids = tuple(range(row.prefix_len))
|
|
row.mm_features = [feature or _feature()]
|
|
row.cache_namespace = namespace
|
|
|
|
|
|
def _hashes(request, feature=None):
|
|
_prepare(request, feature)
|
|
row = request.diffusion_kv_requests[0]
|
|
row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
|
|
return row.block_hashes
|
|
|
|
|
|
def test_partial_hits_and_randomness_are_model_opt_in():
|
|
base = _hashes(_request())
|
|
assert base == _hashes(_request(seed=999))
|
|
changed = _hashes(_request(), _feature(value="other image"))
|
|
assert base[0] == changed[0]
|
|
assert base[1] != changed[1]
|
|
changed_geometry = _hashes(_request(), _feature(end=9))
|
|
assert base[0] == changed_geometry[0]
|
|
assert base[1] != changed_geometry[1]
|
|
|
|
|
|
@pytest.mark.parametrize(("adapter_id", "scale"), [(None, 1.0), (2, 0.5), (1, 1.0)])
|
|
def test_lora_is_common_request_wide_identity(adapter_id, scale):
|
|
base = _hashes(_request(adapter_id=1, scale=0.5))
|
|
changed = _hashes(_request(adapter_id=adapter_id, scale=scale))
|
|
assert all(a != b for a, b in zip(base, changed, strict=True))
|
|
assert _hashes(_request()) == _hashes(_request(scale=0.5))
|
|
|
|
|
|
def test_block_hashing_uses_upstream_extra_keys_directly(monkeypatch):
|
|
row = DiffusionKVRequest(
|
|
"native",
|
|
sequence_id=0,
|
|
prefix_len=14,
|
|
target_len=2,
|
|
seq_len=16,
|
|
cache_token_ids=range(16),
|
|
mm_features=[
|
|
_feature(start=2, end=7),
|
|
_feature(start=7, end=10),
|
|
],
|
|
)
|
|
upstream = Mock(wraps=generate_block_hash_extra_keys)
|
|
monkeypatch.setattr(request_module, "generate_block_hash_extra_keys", upstream)
|
|
hash_fn = get_hash_fn_by_name("sha256")
|
|
row.build_block_hashes(4, hash_fn)
|
|
assert [(call.args[1], call.args[2]) for call in upstream.call_args_list] == [(0, 4), (4, 8), (8, 12)]
|
|
# Independent expected native extras: a crossing item has a negative
|
|
# relative offset; a block can include the end of one item and the next.
|
|
first, second = (feature.identifier for feature in row.mm_features)
|
|
extras = [(("mm", first, 2),), (("mm", first, -2), ("mm", second, 3)), (("mm", second, -1),)]
|
|
parent = BlockHash(hash_fn(("diffusion-prefix-root-v1", row.cache_namespace)))
|
|
expected = []
|
|
for start, keys in zip(range(0, 12, 4), extras, strict=True):
|
|
parent = hash_block_tokens(hash_fn, parent, list(range(start, start + 4)), keys)
|
|
expected.append(parent)
|
|
assert row.block_hashes == expected
|
|
# Never hash the partial prefix tail or mutable target [14, 16), even
|
|
# though cache_token_ids deliberately contains the whole sequence.
|
|
assert len(row.block_hashes) == 3
|
|
|
|
|
|
def test_same_identifier_at_different_offsets_changes_block_hash():
|
|
first = _feature(start=4, end=6)
|
|
shifted = _feature(start=5, end=7)
|
|
shifted.identifier = first.identifier
|
|
a = _hashes(_request(), first)
|
|
b = _hashes(_request(), shifted)
|
|
assert a[0] == b[0]
|
|
assert a[1] != b[1]
|
|
|
|
|
|
def test_multimodal_identity_beyond_prefix_does_not_affect_hashes():
|
|
assert _hashes(_request(), _feature(value="a", start=8, end=12)) == _hashes(
|
|
_request(), _feature(value="b", start=8, end=12)
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"positions",
|
|
[
|
|
[(-1, 2)],
|
|
[(4, 4)],
|
|
[(4, 13)],
|
|
[(5, 7), (4, 5)],
|
|
[(3, 7), (4, 6)],
|
|
],
|
|
)
|
|
def test_invalid_or_overlapping_positions_fail_before_hashing(positions):
|
|
row = _request().diffusion_kv_requests[0]
|
|
row.cache_token_ids = tuple(range(8))
|
|
row.mm_features = [_feature(start=start, end=end) for start, end in positions]
|
|
fail = Mock(side_effect=AssertionError("invalid positions must fail before hashing"))
|
|
with pytest.raises(ValueError, match="sorted, non-overlapping"):
|
|
row.build_block_hashes(4, fail)
|
|
fail.assert_not_called()
|
|
assert row.block_hashes == []
|
|
|
|
|
|
def test_empty_multimodal_identifier_is_rejected():
|
|
row = _request().diffusion_kv_requests[0]
|
|
row.cache_token_ids = tuple(range(8))
|
|
feature = _feature()
|
|
feature.identifier = ""
|
|
row.mm_features = [feature]
|
|
with pytest.raises(ValueError, match="content identifiers"):
|
|
row.build_block_hashes(4, get_hash_fn_by_name("sha256"))
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16, torch.int64, torch.bool])
|
|
def test_tensor_identity_includes_dtype_shape_and_values(dtype):
|
|
tensor = torch.tensor([[0, 1], [1, 0]], dtype=dtype)
|
|
assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.clone())
|
|
assert hash_prefix_cache_value(tensor) == hash_prefix_cache_value(tensor.t().contiguous().t())
|
|
assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.reshape(4))
|
|
changed = tensor.clone()
|
|
changed[0, 0] = 1
|
|
assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(changed)
|
|
assert hash_prefix_cache_value(tensor) != hash_prefix_cache_value(tensor.to(torch.float64))
|
|
|
|
|
|
def test_numpy_scalar_shape_and_nested_dict_order():
|
|
assert hash_prefix_cache_value({"b": {"y": 2, "x": 1}, "a": None}) == hash_prefix_cache_value(
|
|
{"a": None, "b": {"x": 1, "y": 2}}
|
|
)
|
|
assert hash_prefix_cache_value(np.array(1)) != hash_prefix_cache_value(np.array([1]))
|
|
assert hash_prefix_cache_value(np.int64(1)) == hash_prefix_cache_value(1)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("first", "second"), [(1, True), (1, 1.0), ("1", b"1"), ([], {}), ([None], []), (["ab", "c"], ["a", "bc"])]
|
|
)
|
|
def test_identity_has_unambiguous_type_and_container_boundaries(first, second):
|
|
assert hash_prefix_cache_value(first) != hash_prefix_cache_value(second)
|
|
|
|
|
|
@pytest.mark.parametrize("value", [object(), {1: "value"}, np.array([object()]), np.longdouble(1)])
|
|
def test_identity_rejects_unsupported_values_without_pickle(value):
|
|
with pytest.raises(TypeError):
|
|
hash_prefix_cache_value(value)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("mode", "enabled", "expect_hook"),
|
|
[
|
|
(DiffusionKVCacheMode.PAGED_SCHEDULER, False, False),
|
|
(DiffusionKVCacheMode.PAGED_SCHEDULER, True, True),
|
|
(DiffusionKVCacheMode.DENSE_LEGACY, True, False),
|
|
],
|
|
)
|
|
def test_engine_gates_hook_loading_and_admission(monkeypatch, mode, enabled, expect_hook):
|
|
events = []
|
|
request = _request()
|
|
|
|
def preprocess(req):
|
|
events.append("layout")
|
|
req.prepared_layout = "fake layout"
|
|
return req
|
|
|
|
def prepare(req):
|
|
assert req.prepared_layout == "fake layout"
|
|
assert req.diffusion_kv_requests[0].cache_token_ids == ()
|
|
events.append("cache_inputs")
|
|
_prepare(req)
|
|
|
|
hook = Mock(side_effect=prepare if expect_hook else AssertionError("disabled hook"))
|
|
load_hook = Mock(return_value=hook)
|
|
if not expect_hook:
|
|
load_hook.side_effect = AssertionError("disabled hook factory")
|
|
monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", Mock(side_effect=AssertionError("disabled hash")))
|
|
monkeypatch.setattr(engine_module, "get_diffusion_pre_process_func", lambda _config: preprocess)
|
|
monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None)
|
|
monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", load_hook)
|
|
engine = DiffusionEngine.__new__(DiffusionEngine)
|
|
engine.od_config = SimpleNamespace(diffusion_kv_mode=mode, enable_prefix_caching=enabled)
|
|
engine._init_process_hooks(engine.od_config)
|
|
engine.prefix_cache_func = hook
|
|
assert engine._prepare_request_for_admission(request) is request
|
|
row = request.diffusion_kv_requests[0]
|
|
if expect_hook:
|
|
assert events == ["layout", "cache_inputs"]
|
|
assert row.cache_token_ids == tuple(range(8))
|
|
load_hook.assert_called_once_with(engine.od_config)
|
|
hook.assert_called_once_with(request)
|
|
else:
|
|
assert events == ["layout"]
|
|
assert row.cache_token_ids == ()
|
|
assert row.mm_features == []
|
|
load_hook.assert_not_called()
|
|
hook.assert_not_called()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"model_class_name,load_format",
|
|
[(None, "default"), ("UnregisteredPipeline", "default"), ("HunyuanImage3ForCausalMM", "diffusers")],
|
|
)
|
|
def test_engine_rejects_enabled_prefix_caching_without_hook(monkeypatch, model_class_name, load_format):
|
|
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
|
|
|
monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500)
|
|
engine = DiffusionEngine.__new__(DiffusionEngine)
|
|
engine.od_config = OmniDiffusionConfig.from_kwargs(
|
|
model_class_name=model_class_name,
|
|
diffusion_load_format=load_format,
|
|
diffusion_kv_mode="paged_scheduler",
|
|
diffusion_kv_max_rows_per_request=2,
|
|
enable_prefix_caching=True,
|
|
)
|
|
with pytest.raises(ValueError, match="requires a registered prefix-cache hook"):
|
|
engine._init_process_hooks(engine.od_config)
|
|
|
|
|
|
def test_engine_rejects_out_of_profile_request_before_hashing():
|
|
engine = DiffusionEngine.__new__(DiffusionEngine)
|
|
engine.od_config = SimpleNamespace(
|
|
diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, enable_prefix_caching=True
|
|
)
|
|
engine._diffusion_kv_profile_limits = (1, 8, 4)
|
|
engine.prefix_cache_func = Mock(side_effect=AssertionError("invalid request must fail before identity work"))
|
|
with pytest.raises(ValueError, match="exceeds the startup memory-profile envelope"):
|
|
engine._prepare_request_for_admission(_request())
|
|
engine.prefix_cache_func.assert_not_called()
|