# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from types import SimpleNamespace from unittest.mock import Mock import pytest import torch import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers from transformers.utils.generic import ModelOutput from vllm.lora.request import LoRARequest from vllm.utils.hashing import get_hash_fn_by_name from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec 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.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module import vllm_omni.diffusion.models.hunyuan_image3.request_layout as layout_module from tests.helpers.kv_layout import build_kv_cache_tensor from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.diffusion_kv.manager import DiffusionKVCacheManager from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ( ImageInfo, JointImageInfo, ) from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import ( HunyuanImage3Pipeline, get_hunyuan_image_3_pre_process_func, ) from vllm_omni.diffusion.models.hunyuan_image3.request_layout import ( HunyuanPreparedLayout, build_hunyuan_diffusion_kv_requests, extract_hunyuan_prompt_inputs, hunyuan_num_image_tokens, hunyuan_num_special_tokens, normalize_hunyuan_cot_text, prepare_hunyuan_layout, prepare_hunyuan_prefix_cache, ) 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 test_dense_legacy_preprocess_does_not_initialize_layout_tokenizer(monkeypatch) -> None: hf_config = SimpleNamespace(vae_downsample_factor=(8, 8), patch_size=2) image_processor = SimpleNamespace(vision_encoder_processor=SimpleNamespace(patch_size=1)) monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config) monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor) monkeypatch.setattr( pipeline_module, "TokenizerWrapper", lambda _model: pytest.fail("dense_legacy must not initialize the Diffusion KV tokenizer"), ) get_hunyuan_image_3_pre_process_func( SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.DENSE_LEGACY) ) class _FakeTokenizerWrapper: def __init__(self, *, prefix_lens: list[int]) -> None: self.prefix_lens = prefix_lens self.calls: list[dict] = [] self.pad_token_id = 0 self.eos_token_id = 2 self.boi_token_id = 3 self.end_recaption_token_id = 4 self.end_answer_token_id = 5 self.special_token_map = {f"": 100 + index for index in range(33)} def apply_chat_template(self, **kwargs): self.calls.append(kwargs) rows = len(self.prefix_lens) has_cond_image = kwargs.get("batch_cond_image_info") is not None seq_lens = [prefix_len + 20 for prefix_len in self.prefix_lens] padded_seq_len = max(seq_lens) positions = torch.arange(padded_seq_len) # Like the real tokenizer, joint spans enclose VAE + separator + ViT, # while all_image_slices contains the individual RoPE spans. Leave one # full text block before the image so partial reuse can be tested. joint_image_slices = [[slice(5, 10)] if has_cond_image else [] for _ in range(rows)] cond_vae_image_slices = [[slice(5, 7)] if has_cond_image else [] for _ in range(rows)] cond_vit_image_slices = [[slice(8, 10)] if has_cond_image else [] for _ in range(rows)] gen_image_slices = [[slice(length + 1, length + 17)] for length in self.prefix_lens] all_image_slices = [ vae + vit + generated for vae, vit, generated in zip(cond_vae_image_slices, cond_vit_image_slices, gen_image_slices) ] cond_vae_image_mask = ( torch.stack([(positions >= 5) & (positions < 7) for _ in range(rows)]) if has_cond_image else None ) cond_vit_image_mask = ( torch.stack([(positions >= 8) & (positions < 10) for _ in range(rows)]) if has_cond_image else None ) sections = [] for _ in range(rows): row_sections: list[dict[str, object]] = [] if has_cond_image: row_sections.append( { "type": "joint_image", "token_height": [1, 1], "token_width": [2, 2], } ) row_sections.append({"type": "gen_image", "token_height": 4, "token_width": 4}) sections.append(row_sections) return { "output": TokenizerEncodeOutput( tokens=torch.arange(rows * padded_seq_len, dtype=torch.long).reshape(rows, padded_seq_len), gen_timestep_scatter_index=torch.tensor(self.prefix_lens, dtype=torch.long).reshape(rows, 1), cond_timestep_scatter_index=(torch.zeros(rows, 1, dtype=torch.long) if has_cond_image else None), all_image_slices=all_image_slices, gen_image_mask=torch.stack( [(positions >= length + 1) & (positions < length + 17) for length in self.prefix_lens] ), cond_vae_image_mask=cond_vae_image_mask, cond_vit_image_mask=cond_vit_image_mask, cond_vae_image_slices=cond_vae_image_slices, cond_vit_image_slices=cond_vit_image_slices, joint_image_slices=joint_image_slices, gen_image_slices=gen_image_slices, real_pos=torch.tensor(seq_lens, dtype=torch.long).reshape(rows, 1), ), "sections": sections, } class _FakeImageProcessor: def __init__(self, image_token_length: int = 16) -> None: self.image_token_length = image_token_length self.image_sizes: list[tuple[int, int]] = [] self.vision_encoder_processor = SimpleNamespace(patch_size=1) def build_image_info(self, image_size, **kwargs): self.image_sizes.append(image_size) return ImageInfo( image_type="gen_image", image_width=image_size[1], image_height=image_size[0], token_width=4, token_height=4, image_token_length=self.image_token_length, base_size=1024, ratio_index=0, **kwargs, ) def _components(prefix_lens: list[int]): return _FakeTokenizerWrapper(prefix_lens=prefix_lens), _FakeImageProcessor() def _prepare( request: OmniDiffusionRequest, tokenizer: _FakeTokenizerWrapper, image_processor: _FakeImageProcessor, ) -> HunyuanPreparedLayout: prepared_layout = prepare_hunyuan_layout( request, tokenizer_wrapper=tokenizer, image_processor=image_processor, generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False), image_base_size=1024, ) request.prepared_layout = prepared_layout return prepared_layout def _request( *, guidance_scale: float, prompt="draw a cat", request_id: str = "req", seed: int | None = None, ) -> OmniDiffusionRequest: return OmniDiffusionRequest( prompt=prompt, sampling_params=OmniDiffusionSamplingParams( height=768, width=1024, guidance_scale=guidance_scale, num_inference_steps=4, seed=seed, ), request_id=request_id, ) def test_distilled_layout_uses_embedded_cfg_and_meanflow_tokens() -> None: tokenizer, image_processor = _components([12]) request = _request(guidance_scale=2.5) prepared_layout = prepare_hunyuan_layout( request, tokenizer_wrapper=tokenizer, image_processor=image_processor, generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False), image_base_size=1024, cfg_distilled=True, use_meanflow=True, ) image_info = prepared_layout.generated_image_info assert image_info.add_guidance_token assert image_info.add_timestep_r_token assert tokenizer.calls[0]["cfg_factor"] == 1 assert hunyuan_num_special_tokens(image_info) == 3 assert hunyuan_num_image_tokens(image_info) == 19 assert len(build_hunyuan_diffusion_kv_requests(request, prepared_layout)) == 1 def test_builds_kv_request_lengths_without_model_execution() -> None: tokenizer, image_processor = _components([12]) request = _request(guidance_scale=1.0) prepared_layout = _prepare(request, tokenizer, image_processor) kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout) assert isinstance(request.prepared_layout, HunyuanPreparedLayout) assert all(isinstance(item, DiffusionKVRequest) for item in kv_requests) assert [(item.prefix_len, item.target_len, item.seq_len) for item in kv_requests] == [(12, 17, 32)] assert image_processor.image_sizes == [(768, 1024)] assert tokenizer.calls[0]["sequence_template"] == "instruct" assert tokenizer.calls[0]["cfg_factor"] == 1 assert kv_requests[0].block_hashes == [] assert kv_requests[0].cache_token_ids == () assert kv_requests[0].mm_features == [] assert kv_requests[0].kv_contexts == () assert kv_requests[0].skip_reading_prefix_cache is True def test_builds_one_kv_request_per_cfg_row() -> None: tokenizer, image_processor = _components([12, 14]) request = _request(guidance_scale=5.0) prepared_layout = _prepare(request, tokenizer, image_processor) kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout) assert [item.sequence_id for item in kv_requests] == [0, 1] assert [item.prefix_len for item in kv_requests] == [12, 14] assert [item.seq_len for item in kv_requests] == [32, 34] assert tokenizer.calls[0]["cfg_factor"] == 2 def _reference_image() -> JointImageInfo: return JointImageInfo( vae_image_info=ImageInfo( image_type="vae", image_tensor=torch.arange(12, dtype=torch.float32).reshape(1, 3, 2, 2), token_width=8, token_height=8, image_token_length=64, ), vision_image_info=ImageInfo( image_type="siglip2", image_tensor=torch.arange(8, dtype=torch.float32).reshape(1, 2, 2, 2), token_width=4, token_height=4, image_token_length=16, ), vision_encoder_kwargs={ "spatial_shapes": torch.tensor([4, 4]), "pixel_attention_mask": torch.ones(4, dtype=torch.bool), }, ) def _prepare_cached_requests(request, tokenizer, image_processor): layout = _prepare(request, tokenizer, image_processor) request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout) prepare_hunyuan_prefix_cache(request) return request.diffusion_kv_requests def test_reference_prefix_identity_includes_image_content_and_vae_seed() -> None: tokenizer, image_processor = _components([20]) def build(request_id: str, image: JointImageInfo, seed: int): request = _request( guidance_scale=1.0, request_id=request_id, seed=seed, prompt={ "prompt": "edit this image", "additional_information": {"batch_cond_image_info": [image]}, }, ) return _prepare_cached_requests(request, tokenizer, image_processor)[0] first = build("first", _reference_image(), 7) same = build("same", _reference_image(), 7) changed_seed = build("changed-seed", _reference_image(), 8) changed_image_value = _reference_image() changed_image_value.vae_image_info.image_tensor[0, 0, 0, 0] = 99 changed_image = build("changed-image", changed_image_value, 7) assert first.mm_features == same.mm_features assert [(feature.mm_position.offset, feature.mm_position.length) for feature in first.mm_features] == [ (5, 2), (8, 2), ] for changed in (changed_seed, changed_image): # Both subspans must carry the identity, not just the first one. assert all(a.identifier != b.identifier for a, b in zip(first.mm_features, changed.mm_features, strict=True)) def test_reference_prefix_rejects_missing_inputs_with_nested_spans() -> None: tokenizer, image_processor = _components([20]) request = _request( guidance_scale=1.0, prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [_reference_image()]}}, ) layout = _prepare(request, tokenizer, image_processor) request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout) request.prompt["additional_information"].pop("batch_cond_image_info") with pytest.raises(ValueError, match="missing its canonical inputs"): prepare_hunyuan_prefix_cache(request) def test_reference_content_is_hashed_once_across_cfg_rows(monkeypatch): tokenizer, image_processor = _components([20, 20]) image = _reference_image() hasher = Mock(wraps=layout_module.hash_prefix_cache_value) monkeypatch.setattr(layout_module, "hash_prefix_cache_value", hasher) def build(request_id): request = _request( guidance_scale=5.0, request_id=request_id, seed=7, prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [image]}}, ) return _prepare_cached_requests(request, tokenizer, image_processor) first = build("first") assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 1 assert len(first[0].mm_features) == len(first[1].mm_features) == 2 assert first[0].mm_features == first[1].mm_features image.vae_image_info.image_tensor[0, 0, 0, 0] += 1 changed = build("changed") assert sum(call.args[0][0] == "hunyuan-reference-v1" for call in hasher.call_args_list) == 2 assert all(a.identifier != b.identifier for a, b in zip(first[0].mm_features, changed[0].mm_features, strict=True)) def test_failed_cfg_cache_preparation_does_not_install_partial_inputs(): tokenizer, image_processor = _components([20, 20]) request = _request( guidance_scale=5.0, prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}}, ) layout = _prepare(request, tokenizer, image_processor) request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout) # First row prepares successfully; the second exposes inconsistent inputs. layout.rope_image_info[0] = [] request.prompt["additional_information"].pop("batch_cond_image_info") with pytest.raises(ValueError, match="missing its canonical inputs"): prepare_hunyuan_prefix_cache(request) assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests) @pytest.mark.parametrize("generator_list", [False, True]) def test_reference_identity_uses_explicit_generator_state_before_seed(generator_list): tokenizer, image_processor = _components([20]) def build(request_id, seed, advance): request = _request( guidance_scale=1.0, request_id=request_id, seed=seed, prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}}, ) generator = torch.Generator(device="cpu").manual_seed(123) if advance: torch.rand(1, generator=generator) request.sampling_params.generator = [generator] if generator_list else generator before = generator.get_state().clone() row = _prepare_cached_requests(request, tokenizer, image_processor)[0] torch.testing.assert_close(before, generator.get_state()) row.build_block_hashes(4, get_hash_fn_by_name("sha256")) return row.block_hashes first = build("first", 7, False) assert first == build("same-generator-different-unused-seed", 8, False) changed = build("advanced-generator-same-seed", 7, True) assert first[:1] == changed[:1] assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True)) def test_unseeded_reference_identity_is_request_local(): tokenizer, image_processor = _components([20]) def build(request_id): request = _request( guidance_scale=1.0, request_id=request_id, prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}}, ) # Real requests normally receive an automatic seed in __post_init__. request.sampling_params.seed = None row = _prepare_cached_requests(request, tokenizer, image_processor)[0] row.build_block_hashes(4, get_hash_fn_by_name("sha256")) return row.block_hashes first = build("first") assert first == build("first") changed = build("second") assert first[:1] == changed[:1] assert all(a != b for a, b in zip(first[1:], changed[1:], strict=True)) @pytest.mark.parametrize("guidance_scale", [1.0, 5.0]) def test_text_prefix_hashes_include_lora_identity_and_scale(guidance_scale: float) -> None: tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0))) def build(request_id: str, adapter_id: int | None, scale: float): request = _request(guidance_scale=guidance_scale, request_id=request_id) if adapter_id is not None: request.sampling_params.lora_request = LoRARequest( lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}" ) request.sampling_params.lora_scale = scale kv_requests = _prepare_cached_requests(request, tokenizer, image_processor) for kv_request in kv_requests: kv_request.build_block_hashes(4, get_hash_fn_by_name("sha256")) return kv_requests base = build("base", None, 1.0) same_base = build("same-base", None, 0.5) # Scale has no effect without an adapter. first = build("first", 1, 0.5) same = build("same", 1, 0.5) changed_id = build("changed-id", 2, 0.5) changed_scale = build("changed-scale", 1, 1.0) for row in range(len(first)): assert base[row].block_hashes == same_base[row].block_hashes assert first[row].block_hashes == same[row].block_hashes assert len(first[row].block_hashes) == 5 for changed in (base, changed_id, changed_scale): assert first[row].cache_token_ids == changed[row].cache_token_ids # LoRA changes even the text prefix, so every block must differ. assert all(a != b for a, b in zip(first[row].block_hashes, changed[row].block_hashes, strict=True)) @pytest.fixture def prefix_cache_manager(request): native_kv_managers.register_all_kvcache_specs(None) spec = FullAttentionSpec(block_size=4, num_kv_heads=2, head_size=8, dtype=torch.bfloat16) num_blocks = 64 config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=[build_kv_cache_tensor(spec, num_blocks, ["layer0"])], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec)], ) manager = DiffusionKVCacheManager( config, max_model_len=64, scheduler_block_size=4, hash_block_size=4, enable_prefix_caching=getattr(request, "param", True), ) yield manager manager.close() @pytest.mark.parametrize("prefix_cache_manager", [False], indirect=True) @pytest.mark.parametrize("guidance_scale", [1.0, 5.0]) def test_paged_cache_off_prepares_and_allocates_without_identity_work( monkeypatch, prefix_cache_manager, guidance_scale ): tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0))) config = SimpleNamespace( model="fake", model_class_name="HunyuanImage3ForCausalMM", diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER, enable_prefix_caching=False, ) monkeypatch.setattr( pipeline_module, "get_config", lambda *_args, **_kwargs: SimpleNamespace( vae_downsample_factor=(8, 8), patch_size=2, image_base_size=1024, ), ) monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor) monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer) monkeypatch.setattr(pipeline_module.GenerationConfig, "from_pretrained", lambda _model: SimpleNamespace()) monkeypatch.setattr(engine_module, "get_diffusion_post_process_func", lambda _config: None) fail = Mock(side_effect=AssertionError("disabled prefix cache must not prepare identity")) monkeypatch.setattr(engine_module, "get_diffusion_prefix_cache_func", fail) monkeypatch.setattr(layout_module, "prepare_hunyuan_prefix_cache", fail) monkeypatch.setattr(layout_module, "hash_prefix_cache_value", fail) monkeypatch.setattr(kv_utils, "hash_prefix_cache_value", fail) monkeypatch.setattr(kv_utils.MultiModalHasher, "hash_kwargs", fail) monkeypatch.setattr(torch.Tensor, "tolist", fail) monkeypatch.setattr(prefix_cache_manager, "_hash_function", fail) engine = DiffusionEngine.__new__(DiffusionEngine) engine.od_config = config engine._init_process_hooks(config) # Guard admission too, even if someone injects a hook after initialization. engine.prefix_cache_func = fail request = _request( guidance_scale=guidance_scale, prompt={"prompt": "edit", "additional_information": {"batch_cond_image_info": [_reference_image()]}}, ) request.sampling_params.seed = None request.sampling_params.generator = SimpleNamespace(get_state=fail) engine._prepare_request_for_admission(request) assert isinstance(request.prepared_layout, HunyuanPreparedLayout) assert all(row.cache_token_ids == () and row.mm_features == [] for row in request.diffusion_kv_requests) metadata = prefix_cache_manager.reserve_request(request.request_id, request.diffusion_kv_requests) assert metadata is not None assert all(row.cached_prefix_len == 0 for row in metadata.sequences) prefix_cache_manager.publish_request(request.request_id) prefix_cache_manager.free_request(request.request_id) fail.assert_not_called() @pytest.mark.parametrize("guidance_scale", [1.0, 5.0]) @pytest.mark.parametrize( ("change", "expected_hit"), [ ("same", 20), ("prompt_suffix", 12), ("target", 20), ("vae", 4), ("vit", 4), ("vision_kwargs", 4), ("rope", 4), ("seed", 4), ("lora_id", 0), ("lora_scale", 0), ("base", 0), ], ) def test_prefix_lookup_respects_reference_and_lora_identity( prefix_cache_manager: DiffusionKVCacheManager, guidance_scale: float, change: str, expected_hit: int ) -> None: tokenizer, image_processor = _components([20] * (1 + int(guidance_scale > 1.0))) def build(request_id: str, change: str): image = _reference_image() if change == "vae": image.vae_image_info.image_tensor[0, 0, 0, 0] = 99 elif change == "vit": image.vision_image_info.image_tensor[0, 0, 0, 0] = 99 elif change == "vision_kwargs": image.vision_encoder_kwargs["pixel_attention_mask"][0] = False request = _request( guidance_scale=guidance_scale, request_id=request_id, seed=8 if change == "seed" else 7, prompt={"prompt": "edit this image", "additional_information": {"batch_cond_image_info": [image]}}, ) if change != "base": adapter_id = 2 if change == "lora_id" else 1 request.sampling_params.lora_request = LoRARequest( lora_name=f"adapter-{adapter_id}", lora_int_id=adapter_id, lora_path=f"/adapters/{adapter_id}" ) request.sampling_params.lora_scale = 1.0 if change == "lora_scale" else 0.5 layout = _prepare(request, tokenizer, image_processor) if change == "rope": for spans in layout.rope_image_info: spans[0] = (slice(5, 7), (2, 1)) elif change == "prompt_suffix": layout.tokenizer_output.tokens[:, 12] += 100 elif change == "target": layout.tokenizer_output.tokens[:, 20:] += 100 request.diffusion_kv_requests = build_hunyuan_diffusion_kv_requests(request, layout) prepare_hunyuan_prefix_cache(request) return request.diffusion_kv_requests manager = prefix_cache_manager cold = build("cold", "same") cold_metadata = manager.reserve_request("cold", cold) assert cold_metadata is not None assert all(sequence.cached_prefix_len == 0 for sequence in cold_metadata.sequences) # CPU-only: simulate successful materialization to test native cache lookup. manager.publish_request("cold") manager.free_request("cold") warm = build("warm", change) warm_metadata = manager.reserve_request("warm", warm) assert warm_metadata is not None for cold_row, warm_row, cold_sequence, warm_sequence in zip( cold, warm, cold_metadata.sequences, warm_metadata.sequences, strict=True ): if change == "prompt_suffix": assert cold_row.cache_token_ids[:12] == warm_row.cache_token_ids[:12] assert cold_row.cache_token_ids[12] != warm_row.cache_token_ids[12] else: assert cold_row.cache_token_ids == warm_row.cache_token_ids assert warm_sequence.cached_prefix_len == expected_hit hit_blocks = expected_hit // 4 assert cold_sequence.block_ids[0][:hit_blocks] == warm_sequence.block_ids[0][:hit_blocks] assert cold_row.block_hashes[:hit_blocks] == warm_row.block_hashes[:hit_blocks] assert all( a != b for a, b in zip(cold_row.block_hashes[hit_blocks:], warm_row.block_hashes[hit_blocks:], strict=True) ) def test_passes_preprocessed_reference_image_geometry_to_tokenizer() -> None: tokenizer, image_processor = _components([20]) joint_image = _reference_image() prompt = { "prompt": "edit this image", "additional_information": {"batch_cond_image_info": [joint_image]}, } request = _request(guidance_scale=1.0, prompt=prompt) prepared_layout = _prepare(request, tokenizer, image_processor) kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout) assert tokenizer.calls[0]["batch_cond_image_info"] == [[joint_image]] assert kv_requests[0].prefix_len == 20 assert kv_requests[0].target_len == 17 assert kv_requests[0].seq_len == 40 assert kv_requests[0].kv_contexts == () def _assert_nested_equal(actual, expected) -> None: if isinstance(expected, torch.Tensor): torch.testing.assert_close(actual, expected) elif isinstance(expected, dict): assert actual.keys() == expected.keys() for key in expected: _assert_nested_equal(actual[key], expected[key]) elif isinstance(expected, list | tuple): assert isinstance(actual, type(expected)) assert len(actual) == len(expected) for actual_item, expected_item in zip(actual, expected): _assert_nested_equal(actual_item, expected_item) else: assert actual == expected @pytest.mark.parametrize( ("prefix_lens", "guidance_scale", "with_reference_image"), [ pytest.param([12], 1.0, False, id="cfg-1"), pytest.param([12, 14], 5.0, False, id="cfg-2"), pytest.param([20], 1.0, True, id="cfg-1-reference-image"), pytest.param([20, 22], 5.0, True, id="cfg-2-reference-image"), ], ) def test_prepared_model_inputs_match_local_tokenization( monkeypatch, prefix_lens: list[int], guidance_scale: float, with_reference_image: bool, ) -> None: tokenizer, image_processor = _components(prefix_lens) prompt: str | dict = "draw a cat" if with_reference_image: prompt = { "prompt": "edit this image", "additional_information": {"batch_cond_image_info": [_reference_image()]}, } request = _request(guidance_scale=guidance_scale, prompt=prompt) prepared_layout = _prepare(request, tokenizer, image_processor) prompts, cot_texts, system_prompt, batch_cond_image_info, bot_task = extract_hunyuan_prompt_inputs( [request.prompt], request.sampling_params.extra_args or {}, request_id=request.request_id, allow_cond_image=True, ) cot_text = ( [normalize_hunyuan_cot_text(text) for text in cot_texts] if any(text is not None for text in cot_texts) else None ) pipeline = object.__new__(HunyuanImage3Pipeline) pipeline._tkwrapper = tokenizer pipeline.image_processor = image_processor pipeline.generation_config = SimpleNamespace(sequence_template="instruct", drop_think=False) pipeline.config = SimpleNamespace( image_base_size=1024, attention_head_dim=2, rope_theta=10000.0, ) monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda _self: torch.device("cpu"))) def fake_encode_cond_image(batch_cond_image_info_list, cfg_factor=1, generator=None): del generator rows = len(batch_cond_image_info_list) * cfg_factor return ( torch.arange(rows * 4, dtype=torch.float32).reshape(rows, 1, 2, 2), torch.arange(rows, dtype=torch.float32), torch.arange(rows * 3, dtype=torch.float32).reshape(rows, 3), ) pipeline._encode_cond_image = fake_encode_cond_image def fake_build_batch_2d_rope(*, image_infos, seq_len, **_kwargs): geometry_sum = sum( int(image_slice.start or 0) + int(image_slice.stop or 0) + height + width for row in image_infos for image_slice, (height, width) in row ) values = torch.tensor([len(image_infos), seq_len, geometry_sum], dtype=torch.float32) return values, values + 1 monkeypatch.setattr(pipeline_module, "build_batch_2d_rope", fake_build_batch_2d_rope) common_kwargs = dict( prompt=prompts, cot_text=cot_text, system_prompt=system_prompt, mode="gen_image", guidance_scale=guidance_scale, image_size=(request.sampling_params.height, request.sampling_params.width), generator=[torch.Generator().manual_seed(0)], batch_cond_image_info=batch_cond_image_info, bot_task=bot_task, ) local_inputs = pipeline.prepare_model_inputs(**common_kwargs) prepared_inputs = pipeline.prepare_model_inputs(**common_kwargs, prepared_layout=prepared_layout) if with_reference_image: # Preparation runs without a forward context. Coverage must come from # the scheduler metadata copied onto this request, including every CFG row. reused = pipeline.prepare_model_inputs( # The image span is [5, 10); native KV must cover its end, not # merely reach the first image token. **common_kwargs, prepared_layout=prepared_layout, kv_computed_tokens=(10,) * len(prefix_lens), ) assert reused["cond_vae_images"] is None and reused["cond_vit_images"] is None partial = pipeline.prepare_model_inputs( **common_kwargs, prepared_layout=prepared_layout, kv_computed_tokens=(10,) * (len(prefix_lens) - 1) + (9,), ) assert partial["cond_vae_images"] is not None and partial["cond_vit_images"] is not None comparable_fields = ( "input_ids", "position_ids", "custom_pos_emb", "image_mask", "gen_timestep_scatter_index", "cond_vae_images", "cond_timestep", "cond_vae_image_mask", "cond_vit_images", "cond_vit_image_mask", "vit_kwargs", "cond_timestep_scatter_index", ) for field in comparable_fields: _assert_nested_equal(prepared_inputs[field], local_inputs[field]) tokenizer_fields = ( "tokens", "gen_image_mask", "cond_vae_image_mask", "cond_vit_image_mask", "gen_timestep_scatter_index", "cond_timestep_scatter_index", "real_pos", "all_image_slices", "joint_image_slices", "gen_image_slices", ) for field in tokenizer_fields: _assert_nested_equal( getattr(prepared_inputs["tokenizer_output"], field), getattr(local_inputs["tokenizer_output"], field), ) local_mask = pipeline._prepare_attention_mask_for_generation( local_inputs["input_ids"], pipeline.generation_config, local_inputs, ) prepared_mask = pipeline._prepare_attention_mask_for_generation( prepared_inputs["input_ids"], pipeline.generation_config, prepared_inputs, ) torch.testing.assert_close(prepared_mask, local_mask) assert prepared_inputs["full_attn_spans"] == local_inputs["full_attn_spans"] image_info = prepared_layout.generated_image_info generation_token_counts = { "num_image_tokens": hunyuan_num_image_tokens(image_info), "num_special_tokens": hunyuan_num_special_tokens(image_info), } local_inputs.update(attention_mask=local_mask, **generation_token_counts) prepared_inputs.update(attention_mask=prepared_mask, **generation_token_counts) local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs) prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs) for field in ("position_ids", "attention_mask", "gen_timestep_scatter_index", "full_attn_spans"): _assert_nested_equal(prepared_step_inputs[field], local_step_inputs[field]) def test_rejects_tokenizer_cfg_row_mismatch() -> None: tokenizer, image_processor = _components([12]) request = _request(guidance_scale=5.0) prepared_layout = _prepare(request, tokenizer, image_processor) with pytest.raises(ValueError, match="sequence count does not match"): build_hunyuan_diffusion_kv_requests(request, prepared_layout) def test_sender_endpoint_does_not_change_local_kv_request() -> None: tokenizer, image_processor = _components([12]) request = _request(guidance_scale=1.0) request.kv_sender_info = {"host": "127.0.0.1", "zmq_port": 5000} prepared_layout = _prepare(request, tokenizer, image_processor) kv_requests = build_hunyuan_diffusion_kv_requests(request, prepared_layout) assert kv_requests[0].seq_len == 32 assert kv_requests[0].num_computed_tokens == 0 def test_paged_preprocess_attaches_layout_and_scheduler_kv_requests(monkeypatch) -> None: tokenizer, image_processor = _components([12]) hf_config = SimpleNamespace( vae_downsample_factor=(8, 8), patch_size=2, image_base_size=1024, ) monkeypatch.setattr(pipeline_module, "get_config", lambda *_args, **_kwargs: hf_config) monkeypatch.setattr(pipeline_module, "HunyuanImage3ImageProcessor", lambda _config: image_processor) monkeypatch.setattr(pipeline_module, "TokenizerWrapper", lambda _model: tokenizer) monkeypatch.setattr( pipeline_module.GenerationConfig, "from_pretrained", lambda _model: SimpleNamespace(sequence_template="instruct", drop_think=False), ) preprocess = get_hunyuan_image_3_pre_process_func( SimpleNamespace(model="model", diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER) ) request = _request(guidance_scale=1.0) prepared_request = preprocess(request) assert prepared_request is request assert isinstance(request.prepared_layout, HunyuanPreparedLayout) assert request.diffusion_kv_requests is not None assert [(item.prefix_len, item.target_len, item.seq_len) for item in request.diffusion_kv_requests] == [ (12, 17, 32) ] assert len(tokenizer.calls) == 1 @pytest.mark.parametrize("transfer_tokens, expected", [(9, [7, 5]), (4, [4, 4])]) def test_cfg_rows_preserve_reusable_token_ids_and_native_transfer_boundary(transfer_tokens, expected): from vllm_omni.diffusion.diffusion_kv.kv_connector import prepare_kv_requests tokenizer, image_processor = _components([12, 14]) request = _request(guidance_scale=5.0) # Prompt token IDs are only materialized for the native AR -> DiT # transfer path; local prefix caching does not need them. request.kv_transfer_params = {"num_transfer_tokens": 0} layout = _prepare(request, tokenizer, image_processor) layout.tokenizer_output.think_recaption_end_pos = [[7], [5]] layout.tokenizer_output.uncond_cfg_start_pos = [[None], [5]] rows = build_hunyuan_diffusion_kv_requests(request, layout) assert [row.prompt_token_ids for row in rows] == [ layout.tokenizer_output.tokens[0, :7].tolist(), layout.tokenizer_output.tokens[1, :5].tolist(), ] prepare_kv_requests(rows, {"num_transfer_tokens": transfer_tokens}) assert [row.num_prompt_tokens for row in rows] == expected assert all(row.kv_transfer_params["num_transfer_tokens"] == transfer_tokens for row in rows) layout.tokenizer_output.uncond_cfg_start_pos = [[5]] with pytest.raises(ValueError, match="one boundary per CFG row"): build_hunyuan_diffusion_kv_requests(request, layout) @pytest.mark.parametrize("computed, expected", [((8, 8), True), ((8, 5), False), ((0, 0), False), ((), False)]) def test_native_image_reuse_requires_coverage_in_every_cfg_row(computed, expected): from vllm_omni.diffusion.models.hunyuan_image3.request_layout import native_kv_covers_cond_images output = TokenizerEncodeOutput(joint_image_slices=[[slice(2, 8)], [slice(2, 6)]]) assert native_kv_covers_cond_images(output, computed) is expected