From 6ae7ff78463683fc3d0fb3eb49f6a7bea2aaf4eb Mon Sep 17 00:00:00 2001 From: Yukim1 <121286183+zwhzzz0821@users.noreply.github.com> Date: Mon, 10 Aug 2026 08:48:56 +0800 Subject: [PATCH] [Diffusion] Prepare HunyuanImage3 for Scheduler-managed paged KV cache (#5541) Signed-off-by: zwhzzz0821 Co-authored-by: zwhzzz0821 Co-authored-by: Hongsheng Liu --- tests/config/test_omni_config.py | 31 ++ tests/diffusion/diffusion_kv/test_config.py | 52 ++ tests/diffusion/diffusion_kv/test_request.py | 200 ++++++++ .../test_diffusion_kv_request.py | 447 ++++++++++++++++++ .../test_hunyuan_image3_step_execution.py | 91 +++- .../test_diffusion_config_propagation.py | 6 + tests/diffusion/test_diffusion_engine.py | 39 ++ .../test_diffusion_engine_rpc_routing.py | 1 + .../diffusion/test_diffusion_model_runner.py | 19 + tests/diffusion/test_diffusion_scheduler.py | 124 ++++- .../test_multiproc_engine_concurrency.py | 1 + vllm_omni/config/omni_config.py | 4 + vllm_omni/diffusion/data.py | 23 + vllm_omni/diffusion/diffusion_engine.py | 30 +- vllm_omni/diffusion/diffusion_kv/__init__.py | 2 + vllm_omni/diffusion/diffusion_kv/config.py | 39 ++ vllm_omni/diffusion/diffusion_kv/request.py | 116 +++++ .../hunyuan_image3/pipeline_hunyuan_image3.py | 341 ++++--------- .../models/hunyuan_image3/request_layout.py | 377 +++++++++++++++ vllm_omni/diffusion/request.py | 15 +- vllm_omni/diffusion/sched/base_scheduler.py | 37 ++ vllm_omni/diffusion/sched/interface.py | 6 +- .../worker/diffusion_model_runner.py | 1 + vllm_omni/diffusion/worker/utils.py | 3 + 24 files changed, 1728 insertions(+), 277 deletions(-) create mode 100644 tests/diffusion/diffusion_kv/test_config.py create mode 100644 tests/diffusion/diffusion_kv/test_request.py create mode 100644 tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py create mode 100644 vllm_omni/diffusion/diffusion_kv/__init__.py create mode 100644 vllm_omni/diffusion/diffusion_kv/config.py create mode 100644 vllm_omni/diffusion/diffusion_kv/request.py create mode 100644 vllm_omni/diffusion/models/hunyuan_image3/request_layout.py diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py index c0a3f20e8..3d2caeabf 100644 --- a/tests/config/test_omni_config.py +++ b/tests/config/test_omni_config.py @@ -39,6 +39,7 @@ from vllm_omni.config.stage_config import ( load_deploy_config, merge_pipeline_deploy, ) +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.engine.stage_init_utils import build_legacy_engine_args_dict pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -966,20 +967,24 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch): cfg = omni_config_module._DiffusionConfigProjection.from_kwargs( diffusion_attention_backend="flash_attn", + fa_deterministic=True, kv_cache_dtype="fp8", kv_cache_skip_steps="0-1", kv_cache_skip_layers=[2], static_lora_scale=0.25, + diffusion_kv_mode="paged_scheduler", diffusers_load_kwargs=None, diffusers_call_kwargs=None, ) assert cfg.diffusion_attention_config.default.backend == "flash_attn" + assert cfg.fa_deterministic is True assert cfg.diffusion_kv_cache_dtype == "fp8" assert cfg.diffusion_kv_cache_skip_step_indices == {0, 1} assert cfg.diffusion_kv_cache_skip_layer_indices == {2} assert cfg.lora_scale == 0.25 assert cfg.cache_backend == "tea_cache" + assert cfg.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER assert cfg.diffusers_load_kwargs == {} assert cfg.diffusers_call_kwargs == {} @@ -994,6 +999,8 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar "stages:", " - stage_id: 0", " diffusion_attention_backend: flash_attn", + " fa_deterministic: true", + " diffusion_kv_mode: paged_scheduler", ] ) ) @@ -1005,6 +1012,29 @@ def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_ar assert isinstance(stage, VllmOmniDiffusionStageConfig) assert stage.diffusion_config.diffusion_attention_config.default.backend == "flash_attn" + assert stage.diffusion_config.fa_deterministic is True + assert stage.diffusion_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + + +def test_from_pipeline_config_rejects_reserved_diffusion_kv_mode(tmp_path): + deploy_path = tmp_path / "dreamzero_reserved_diffusion_kv.yaml" + deploy_path.write_text( + "\n".join( + [ + "pipeline: dreamzero", + "async_chunk: false", + "stages:", + " - stage_id: 0", + " diffusion_kv_mode: paged_worker_local", + ] + ) + ) + + with pytest.raises(ValueError, match="reserved but not implemented"): + _from_pipeline_key( + "dreamzero", + deploy_config_path=str(deploy_path), + ) def test_diffusion_config_field_classification_covers_current_fields(): @@ -1025,6 +1055,7 @@ def test_diffusion_config_field_classification_covers_current_fields(): "enable_session_state_manager", "prompt_embed_cache_size", "diffusion_kv_cache_dtype", + "diffusion_kv_mode", } <= omni_config_module._DIFFUSION_ONLY_CONFIG_FIELDS assert { "revision", diff --git a/tests/diffusion/diffusion_kv/test_config.py b/tests/diffusion/diffusion_kv/test_config.py new file mode 100644 index 000000000..f03120784 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_config.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + + +@pytest.fixture(autouse=True) +def _fixed_master_port(monkeypatch) -> None: + monkeypatch.setattr(OmniDiffusionConfig, "_resolve_master_port", lambda _self: 29500) + + +def test_dense_legacy_is_default() -> None: + config = OmniDiffusionConfig.from_kwargs() + + assert config.diffusion_kv_mode is DiffusionKVCacheMode.DENSE_LEGACY + + +def test_paged_worker_local_is_rejected_until_implemented() -> None: + with pytest.raises(ValueError, match="reserved but not implemented"): + OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="paged_worker_local") + + +def test_unknown_cache_mode_is_rejected() -> None: + with pytest.raises(ValueError, match="Unsupported Diffusion KV diffusion_kv_mode"): + OmniDiffusionConfig.from_kwargs(diffusion_kv_mode="unknown") + + +def test_non_mapping_omni_kv_config_is_rejected() -> None: + with pytest.raises(TypeError, match="omni_kv_config must be a mapping"): + OmniDiffusionConfig.from_kwargs(omni_kv_config="paged_scheduler") + + +def test_paged_scheduler_rejects_dense_legacy_kv_receive() -> None: + with pytest.raises(ValueError, match="does not support imported AR KV"): + OmniDiffusionConfig.from_kwargs( + diffusion_kv_mode="paged_scheduler", + omni_kv_config={"need_recv_cache": True}, + ) + + +def test_paged_scheduler_does_not_depend_on_model_registry() -> None: + config = OmniDiffusionConfig.from_kwargs( + model_class_name="FutureDiffusionModel", + diffusion_kv_mode="paged_scheduler", + ) + + assert config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER diff --git a/tests/diffusion/diffusion_kv/test_request.py b/tests/diffusion/diffusion_kv/test_request.py new file mode 100644 index 000000000..43734ca31 --- /dev/null +++ b/tests/diffusion/diffusion_kv/test_request.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import vllm.v1.core.single_type_kv_cache_manager as native_kv_managers +from vllm.v1.core.kv_cache_manager import KVCacheManager +from vllm.v1.core.kv_cache_utils import BlockHash +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, +) +from vllm.v1.request import RequestStatus + +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] + +BLOCK_SIZE = 4 + + +def _request(**overrides) -> DiffusionKVRequest: + values = dict( + request_id="req/diffusion-kv/0", + sequence_id=0, + prefix_len=4, + target_len=4, + seq_len=8, + ) + values.update(overrides) + return DiffusionKVRequest(**values) + + +def test_request_exposes_native_request_and_diffusion_semantics() -> None: + request = _request() + + assert request.sequence_id == 0 + assert request.prefix_len == 4 + assert request.target_len == 4 + assert request.num_tokens == request.seq_len == 8 + assert request.num_prompt_tokens == 4 + assert request.num_computed_tokens == 0 + assert request.block_hashes == [] + assert request.kv_contexts == () + assert request.skip_reading_prefix_cache is True + assert request.shared_prefix_boundary == 0 + assert request.status is RequestStatus.WAITING + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"sequence_id": -1}, "sequence_id"), + ({"prefix_len": -1}, "prefix_len"), + ({"target_len": 0}, "target_len"), + ({"seq_len": 0}, "seq_len"), + ({"prefix_len": 5, "target_len": 4, "seq_len": 8}, "must not exceed"), + ], +) +def test_request_rejects_invalid_lengths(overrides, message: str) -> None: + with pytest.raises(ValueError, match=message): + _request(**overrides) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"context_id": ""}, "context_id"), + ({"cache_role": ""}, "cache_role"), + ({"num_tokens": 0}, "num_tokens"), + ], +) +def test_context_rejects_invalid_identity_or_length(overrides, message: str) -> None: + values = dict(context_id="text", cache_role="cross.text", num_tokens=8) + values.update(overrides) + + with pytest.raises(ValueError, match=message): + DiffusionKVContext(**values) + + +def test_request_owns_independent_kv_contexts() -> None: + text = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8) + image = DiffusionKVContext(context_id="image", cache_role="cross.image", num_tokens=4) + + request = _request(kv_contexts=(text, image)) + + assert request.kv_contexts == (text, image) + + +def test_request_rejects_duplicate_context_ids() -> None: + contexts = ( + DiffusionKVContext(context_id="condition", cache_role="cross.text", num_tokens=8), + DiffusionKVContext(context_id="condition", cache_role="cross.image", num_tokens=4), + ) + + with pytest.raises(ValueError, match="unique context_id"): + _request(kv_contexts=contexts) + + +def test_request_rejects_non_context_values() -> None: + with pytest.raises(TypeError, match="DiffusionKVContext"): + _request(kv_contexts=(object(),)) + + +def _manager(*, enable_caching: bool) -> KVCacheManager: + # vLLM 0.26 lazily registers its built-in KV specs. Another test may have + # already registered an out-of-tree spec, making the registry non-empty + # before FullAttentionSpec is registered. Initialize the native mappings + # explicitly; vLLM 0.25 has no registration hook and needs no setup. + register_all = getattr(native_kv_managers, "register_all_kvcache_specs", None) + if register_all is not None: + register_all(None) + + spec = FullAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=2, + head_size=8, + dtype=torch.bfloat16, + ) + group = KVCacheGroupSpec(layer_names=["layer0"], kv_cache_spec=spec) + config = KVCacheConfig( + num_blocks=8, + kv_cache_tensors=[KVCacheTensor(size=spec.page_size_bytes * 8, shared_by=["layer0"])], + kv_cache_groups=[group], + ) + return KVCacheManager( + config, + max_model_len=64, + scheduler_block_size=BLOCK_SIZE, + hash_block_size=BLOCK_SIZE, + enable_caching=enable_caching, + ) + + +def test_empty_hash_request_conforms_to_native_request_local_allocation() -> None: + manager = _manager(enable_caching=True) + free_before = manager.block_pool.get_num_free_blocks() + request = _request() + + # vLLM 0.26 adds a shared-prefix-boundary return value. + computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request) + assert num_computed_tokens == 0 + blocks = manager.allocate_slots( + request, + num_new_tokens=request.num_tokens, + new_computed_blocks=computed_blocks, + delay_cache_blocks=True, + full_sequence_must_fit=True, + ) + + assert blocks is not None + assert len(manager.get_block_ids(request.request_id)[0]) == 2 + # Empty hashes intentionally mean request-local residency only. Calling + # native cache_blocks() would require a canonical hash for the full prefix + # block and is therefore outside this request's contract. + manager.free(request) + assert manager.block_pool.get_num_free_blocks() == free_before + + +def test_hashed_request_conforms_to_native_delayed_prefix_commit() -> None: + manager = _manager(enable_caching=True) + prefix_hash = BlockHash(b"prefix".ljust(32, b"\0")) + target_hash = BlockHash(b"target".ljust(32, b"\0")) + request = _request(block_hashes=[prefix_hash, target_hash]) + + computed_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(request) + assert num_computed_tokens == 0 + blocks = manager.allocate_slots( + request, + num_new_tokens=request.num_tokens, + new_computed_blocks=computed_blocks, + delay_cache_blocks=True, + full_sequence_must_fit=True, + ) + + assert blocks is not None + manager.cache_blocks(request, request.prefix_len) + manager.free(request) + + cached_request = _request( + request_id="req/diffusion-kv/cached", + block_hashes=[prefix_hash, target_hash], + ) + cached_blocks, num_computed_tokens, *_ = manager.get_computed_blocks(cached_request) + + # Only the stable prefix was published. The dynamic target remains a miss. + assert num_computed_tokens == cached_request.prefix_len + assert len(cached_blocks.blocks[0]) == 1 + cached_allocation = manager.allocate_slots( + cached_request, + num_new_tokens=cached_request.num_tokens - num_computed_tokens, + num_new_computed_tokens=num_computed_tokens, + new_computed_blocks=cached_blocks, + delay_cache_blocks=True, + full_sequence_must_fit=True, + ) + assert cached_allocation is not None + manager.free(cached_request) diff --git a/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py b/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py new file mode 100644 index 000000000..3876762fa --- /dev/null +++ b/tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py @@ -0,0 +1,447 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from transformers.utils.generic import ModelOutput + +import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as pipeline_module +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +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, + normalize_hunyuan_cot_text, + prepare_hunyuan_layout, +) +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) + joint_image_slices = [[slice(1, 3), slice(3, 5)] 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 = [joint + generated for joint, generated in zip(joint_image_slices, gen_image_slices)] + cond_vae_image_mask = ( + torch.stack([(positions >= 1) & (positions < 3) for _ in range(rows)]) if has_cond_image else None + ) + cond_vit_image_mask = ( + torch.stack([(positions >= 3) & (positions < 5) for _ in range(rows)]) if has_cond_image else None + ) + sections = [] + for _ in range(rows): + row_sections = [] + 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, + 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): + 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, + ) + + +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", +) -> OmniDiffusionRequest: + return OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams( + height=768, + width=1024, + guidance_scale=guidance_scale, + num_inference_steps=4, + ), + request_id=request_id, + ) + + +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].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", + token_width=8, + token_height=8, + image_token_length=64, + ), + vision_image_info=ImageInfo( + image_type="siglip2", + 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 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) + + 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"] + + num_image_tokens = hunyuan_num_image_tokens(prepared_layout.generated_image_info) + local_inputs.update(attention_mask=local_mask, num_image_tokens=num_image_tokens) + prepared_inputs.update(attention_mask=prepared_mask, num_image_tokens=num_image_tokens) + 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 diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py index d9d310ee4..dfb383e87 100644 --- a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py @@ -7,7 +7,10 @@ import pytest import torch import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as hy3_module +import vllm_omni.diffusion.models.hunyuan_image3.request_layout as hy3_layout_module from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec +from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_tokenizer import TokenizerEncodeOutput +from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ImageInfo from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import ( _STEP_AR_KV, _STEP_CFG_FACTOR, @@ -18,6 +21,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import ( _STEP_PROMPT_KV, HunyuanImage3Pipeline, ) +from vllm_omni.diffusion.models.hunyuan_image3.request_layout import HunyuanPreparedLayout from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.diffusion.worker.utils import StepRequestState @@ -76,6 +80,72 @@ def _sampling_params(**extra_args): ) +def _prepared_layout() -> HunyuanPreparedLayout: + return HunyuanPreparedLayout( + tokenizer_output=TokenizerEncodeOutput( + tokens=torch.arange(21).reshape(1, 21), + gen_image_mask=torch.ones(1, 21, dtype=torch.bool), + gen_timestep_scatter_index=torch.tensor([[4]]), + all_image_slices=[[slice(5, 21)]], + joint_image_slices=[[]], + gen_image_slices=[[slice(5, 21)]], + real_pos=torch.tensor([[21]]), + ), + rope_image_info=[[(slice(5, 21), (4, 4))]], + generated_image_info=ImageInfo( + image_type="gen_image", + image_width=512, + image_height=512, + token_width=4, + token_height=4, + image_token_length=16, + ), + ) + + +def test_prepare_model_inputs_reuses_prepared_layout(monkeypatch): + pipeline = _pipeline() + monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu"))) + rope_kwargs = {} + + def fake_build_batch_2d_rope(**kwargs): + rope_kwargs.update(kwargs) + return torch.zeros(1), torch.zeros(1) + + monkeypatch.setattr(hy3_module, "build_batch_2d_rope", fake_build_batch_2d_rope) + + def fail_apply_chat_template(**_kwargs): + pytest.fail("prepared Hunyuan execution must not apply the chat template again") + + pipeline._tkwrapper = SimpleNamespace( + apply_chat_template=fail_apply_chat_template, + eos_token_id=2, + boi_token_id=3, + end_recaption_token_id=4, + end_answer_token_id=5, + ) + pipeline.config = SimpleNamespace( + attention_head_dim=2, + rope_theta=10000.0, + ) + prepared = _prepared_layout() + + model_inputs = pipeline.prepare_model_inputs( + prompt="prompt", + mode="gen_image", + guidance_scale=1.0, + image_size=(512, 512), + device=torch.device("cpu"), + generator=[torch.Generator().manual_seed(0)], + prepared_layout=prepared, + ) + + assert model_inputs["tokenizer_output"] is prepared.tokenizer_output + assert model_inputs["batch_gen_image_info"] == [prepared.generated_image_info] + assert rope_kwargs["image_infos"] is prepared.rope_image_info + torch.testing.assert_close(model_inputs["input_ids"], prepared.tokenizer_output.tokens) + + def test_hunyuan_step_group_key_ignores_step_index_for_later_steps(): pipeline = _pipeline() states = [_state("req-0", 1), _state("req-1", 3)] @@ -131,12 +201,13 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics( captured.update(kwargs) raise RuntimeError("stop after prepare_model_inputs") - monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt) + monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt) pipeline.prepare_model_inputs = fake_prepare_model_inputs state = StepRequestState( request_id="req-bot-task", sampling=sampling, prompt=prompt_item, + prepared_layout=_prepared_layout(), ) with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"): @@ -144,6 +215,7 @@ def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics( assert captured["bot_task"] == expected_model_bot_task assert captured["system_prompt_bot_task"] == expected_system_bot_task + assert captured["prepared_layout"] is state.prepared_layout def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch): @@ -159,23 +231,22 @@ def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch): captured.update(kwargs) raise RuntimeError("stop after prepare_model_inputs") - monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt) + monkeypatch.setattr(hy3_layout_module, "get_system_prompt", fake_get_system_prompt) pipeline.prepare_model_inputs = fake_prepare_model_inputs - req = DiffusionRequestBatch( - requests=[ - SimpleNamespace( - request_id="req-forward-bot-task", - sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"), - prompt={"prompt": "prompt", "bot_task": "vanilla"}, - ) - ] + request = SimpleNamespace( + request_id="req-forward-bot-task", + sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"), + prompt={"prompt": "prompt", "bot_task": "vanilla"}, + prepared_layout=_prepared_layout(), ) + req = DiffusionRequestBatch(requests=[request]) with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"): pipeline.forward(req) assert captured["bot_task"] == "think" assert captured["system_prompt_bot_task"] == "think" + assert captured["prepared_layout"] is request.prepared_layout def test_grouped_denoise_rejects_non_sdpa_attention_backend(): diff --git a/tests/diffusion/test_diffusion_config_propagation.py b/tests/diffusion/test_diffusion_config_propagation.py index 91d366723..e25d6497c 100644 --- a/tests/diffusion/test_diffusion_config_propagation.py +++ b/tests/diffusion/test_diffusion_config_propagation.py @@ -15,6 +15,7 @@ from vllm_omni.diffusion.data import ( DiffusionParallelConfig, OmniDiffusionConfig, ) +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.model_metadata import ( HUNYUAN_IMAGE3_MAX_INPUT_IMAGES, QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES, @@ -124,6 +125,11 @@ class TestCreateDefaultDiffusion: assert ea["enforce_eager"] is True assert ea["lora_path"] == "/tmp/lora" + def test_diffusion_kv_mode_roundtrip(self): + od = _roundtrip_diffusion_config(model="x", diffusion_kv_mode="paged_scheduler") + + assert od.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + def test_qwen_image_edit_plus_sets_generic_multimodal_limit(): od_config = OmniDiffusionConfig(model="Qwen/Qwen-Image-Edit-2511", model_class_name="QwenImageEditPlusPipeline") diff --git a/tests/diffusion/test_diffusion_engine.py b/tests/diffusion/test_diffusion_engine.py index 3beb699a6..c0f79aff6 100644 --- a/tests/diffusion/test_diffusion_engine.py +++ b/tests/diffusion/test_diffusion_engine.py @@ -108,6 +108,17 @@ class MockScheduler: pass +def _make_admission_engine(pre_process_func) -> DiffusionEngine: + engine = object.__new__(DiffusionEngine) + engine.scheduler = MockScheduler() + engine.pre_process_func = pre_process_func + engine._out_streams = {} + engine._rpc_lock = threading.RLock() + engine._cv = threading.Condition(engine._rpc_lock) + engine._closed = False + return engine + + class _BatchCapablePipeline: supports_request_batch = True @@ -683,6 +694,34 @@ async def _consume_final_output(generator): return final_output +@pytest.mark.cpu +@pytest.mark.parametrize("entrypoint", ["add_request", "async_add_req_and_stream_response"]) +def test_engine_admission_preprocesses_request_once(entrypoint: str) -> None: + raw_request = OmniDiffusionRequest( + prompt="raw", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_id="admission", + ) + prepared_request = OmniDiffusionRequest( + prompt="prepared", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + request_id="admission", + prepared_layout=object(), + ) + preprocess_calls = [] + + def preprocess(request): + preprocess_calls.append(request) + return prepared_request + + engine = _make_admission_engine(preprocess) + + getattr(engine, entrypoint)(raw_request) + + assert preprocess_calls == [raw_request] + assert engine.scheduler._waiting_queue == [prepared_request] + + @pytest.mark.cpu @pytest.mark.asyncio async def test_async_add_req_and_stream_response(): diff --git a/tests/diffusion/test_diffusion_engine_rpc_routing.py b/tests/diffusion/test_diffusion_engine_rpc_routing.py index dff5ac7bc..5e4e9189a 100644 --- a/tests/diffusion/test_diffusion_engine_rpc_routing.py +++ b/tests/diffusion/test_diffusion_engine_rpc_routing.py @@ -127,6 +127,7 @@ def _make_request(tag: str): request_id=tag, prompt=f"prompt_{tag}", sampling_params=SimpleNamespace(**sampling_params), + diffusion_kv_requests=None, ) diff --git a/tests/diffusion/test_diffusion_model_runner.py b/tests/diffusion/test_diffusion_model_runner.py index e57275485..a1f8440e3 100644 --- a/tests/diffusion/test_diffusion_model_runner.py +++ b/tests/diffusion/test_diffusion_model_runner.py @@ -218,6 +218,25 @@ def _make_compile_runner( return runner +@pytest.mark.core_model +@pytest.mark.cpu +def test_update_states_carries_prepared_layout() -> None: + runner = _make_runner(cache_backend=None, cache_backend_name=None) + request = _make_request() + prepared_layout = object() + request.prepared_layout = prepared_layout + scheduler_output = SimpleNamespace( + finished_req_ids=set(), + scheduled_new_reqs=[SimpleNamespace(request_id=request.request_id, req=request)], + scheduled_cached_reqs=SimpleNamespace(request_ids=[]), + ) + + states, new_request_ids = DiffusionModelRunner._update_states(runner, scheduler_output) + + assert new_request_ids == [request.request_id] + assert states[0].prepared_layout is prepared_layout + + @pytest.mark.core_model @pytest.mark.cpu @pytest.mark.parametrize("use_hsdp", [False, True]) diff --git a/tests/diffusion/test_diffusion_scheduler.py b/tests/diffusion/test_diffusion_scheduler.py index a20d7e5a9..ec0bef264 100644 --- a/tests/diffusion/test_diffusion_scheduler.py +++ b/tests/diffusion/test_diffusion_scheduler.py @@ -12,6 +12,8 @@ from pytest_mock import MockerFixture from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVContext, DiffusionKVRequest from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched import ( BaseScheduler, @@ -259,6 +261,90 @@ class TestRequestScheduler: assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED assert self.scheduler.has_requests() is False + def test_diffusion_kv_request_moves_to_scheduler_state(self) -> None: + self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + request = _make_request("diffusion-kv") + prepared_layout = object() + request.prepared_layout = prepared_layout + text_context = DiffusionKVContext(context_id="text", cache_role="cross.text", num_tokens=8) + kv_request = DiffusionKVRequest( + "diffusion-kv/diffusion-kv/0", + sequence_id=0, + prefix_len=4, + target_len=8, + seq_len=16, + kv_contexts=(text_context,), + ) + request.diffusion_kv_requests = (kv_request,) + + self.scheduler.add_request(request) + state = self.scheduler.get_request_state(request.request_id) + scheduler_output = self.scheduler.schedule() + + assert state is not None + assert state.diffusion_kv_requests == (kv_request,) + assert state.diffusion_kv_requests[0].kv_contexts == (text_context,) + assert request.diffusion_kv_requests is None + assert scheduler_output.scheduled_new_reqs[0].req.prepared_layout is prepared_layout + + def test_paged_scheduler_rejects_missing_kv_request(self) -> None: + self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + request = _make_request("diffusion-kv") + + with pytest.raises(ValueError, match="did not produce DiffusionKVRequest"): + self.scheduler.add_request(request) + + @pytest.mark.parametrize( + ("owner_name", "field_name"), + [ + ("request", "past_key_values"), + ("sampling_params", "past_key_values"), + ("sampling_params", "cfg_text_past_key_values"), + ("sampling_params", "cfg_img_past_key_values"), + ("sampling_params", "cfg_branch_past_key_values"), + ], + ) + def test_diffusion_kv_request_rejects_legacy_dense_kv(self, owner_name: str, field_name: str) -> None: + self.scheduler.initialize(SimpleNamespace(diffusion_kv_mode=DiffusionKVCacheMode.PAGED_SCHEDULER)) + request = _make_request("diffusion-kv") + owner = request if owner_name == "request" else request.sampling_params + setattr(owner, field_name, object()) + request.diffusion_kv_requests = ( + DiffusionKVRequest( + "diffusion-kv/diffusion-kv/0", + sequence_id=0, + prefix_len=4, + target_len=8, + seq_len=16, + ), + ) + + with pytest.raises(ValueError, match=rf"{owner_name}\.{field_name}"): + self.scheduler.add_request(request) + + def test_dense_scheduler_rejects_scheduler_only_kv_request(self) -> None: + request = _make_request("dense-kv-request") + request.diffusion_kv_requests = ( + DiffusionKVRequest( + "dense-kv-request/diffusion-kv/0", + sequence_id=0, + prefix_len=4, + target_len=8, + seq_len=16, + ), + ) + + with pytest.raises(ValueError, match="dense_legacy request unexpectedly contains"): + self.scheduler.add_request(request) + + def test_dense_scheduler_accepts_legacy_dense_kv(self) -> None: + request = _make_request("dense-kv") + request.sampling_params.past_key_values = object() + + request_id = self.scheduler.add_request(request) + + assert request_id == request.request_id + def test_error_output_marks_finished_error(self) -> None: req_id = self.scheduler.add_request(_make_request("err")) @@ -562,12 +648,20 @@ class TestDiffusionEngine: engine.abort_queue = queue.Queue() request = _make_request("engine") + prepared_layout = object() + + def preprocess(req): + req.prepared_layout = prepared_layout + return req + + engine.pre_process_func = mocker.Mock(side_effect=preprocess) runner_output = _make_request_output("engine") engine.execute_fn = mocker.Mock(return_value=runner_output) output = engine.add_req_and_wait_for_response(request) assert output is runner_output.result + engine.pre_process_func.assert_called_once_with(request) engine.execute_fn.assert_called_once() def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> None: @@ -631,17 +725,29 @@ class TestDiffusionEngine: async def test_step_streaming_raises_aborted_error(self, mocker: MockerFixture) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) engine._check_and_start_background_loop = mocker.AsyncMock() - engine.pre_process_func = None + request = _make_request("req-abort") + prepared_layout = object() + + def preprocess(req): + req.prepared_layout = prepared_layout + return req + + engine.pre_process_func = mocker.Mock(side_effect=preprocess) async def _stream(_request): yield DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.") - engine.async_add_req_and_stream_response = mocker.Mock(return_value=_stream(None)) + engine._add_prepared_request = mocker.Mock(return_value=request.request_id) + engine.get_output_stream = mocker.Mock(return_value=_stream(None)) with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"): - async for _ in engine.step_streaming(_make_request("req-abort")): + async for _ in engine.step_streaming(request): pass + engine.pre_process_func.assert_called_once_with(request) + engine._add_prepared_request.assert_called_once_with(request) + engine.get_output_stream.assert_called_once_with(request.request_id) + def test_abort_queue_marks_request_finished_aborted(self) -> None: engine = DiffusionEngine.__new__(DiffusionEngine) engine._rpc_lock = threading.RLock() @@ -774,6 +880,18 @@ class TestDiffusionEngine: with pytest.raises(RuntimeError, match="Dummy run failed: boom"): engine._dummy_run() + def test_dummy_run_delegates_preprocessing_to_sync_admission(self, mocker: MockerFixture) -> None: + engine = DiffusionEngine.__new__(DiffusionEngine) + engine.od_config = SimpleNamespace(model_class_name="mock_model", diffusion_load_format="default") + engine.pre_process_func = mocker.Mock() + engine.add_req_and_wait_for_response = mocker.Mock(return_value=DiffusionOutput(output=None)) + + engine._dummy_run() + + engine.pre_process_func.assert_not_called() + admitted_request = engine.add_req_and_wait_for_response.call_args.args[0] + assert admitted_request.request_id == "dummy_req_id" + class TestStepScheduler: def setup_method(self) -> None: diff --git a/tests/diffusion/test_multiproc_engine_concurrency.py b/tests/diffusion/test_multiproc_engine_concurrency.py index 9ceca0949..aa40d0770 100644 --- a/tests/diffusion/test_multiproc_engine_concurrency.py +++ b/tests/diffusion/test_multiproc_engine_concurrency.py @@ -48,6 +48,7 @@ def _mock_request(tag: str): request_id=tag, prompt=f"prompt_{tag}", sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1), + diffusion_kv_requests=None, ) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 641997468..8997b12e1 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -37,6 +37,7 @@ from vllm_omni.config.stage_config import ( build_stage_runtime_overrides, load_deploy_config, ) +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode _EXECUTION_TYPE_TO_STAGE_WORKER: dict[StageExecutionType, tuple[StageType, str | None]] = { StageExecutionType.LLM_AR: (StageType.LLM, "ar"), @@ -549,6 +550,7 @@ class _DiffusionConfigProjection: cache_backend: str = "none" cache_config: Any = field(default_factory=dict) enable_cache_dit_summary: bool = False + diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY enable_prompt_embed_cache: bool = False prompt_embed_cache_size: int = Field(default=32, ge=1) enable_session_state_manager: bool = False @@ -623,6 +625,7 @@ class _DiffusionConfigProjection: build_attention_config, parse_kv_cache_skip_selector, ) + from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode from vllm_omni.quantization import build_quant_config if self.tf_model_config is None: @@ -684,6 +687,7 @@ class _DiffusionConfigProjection: f"got {type(self.diffusion_attention_config)!r}" ) + self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) self.diffusion_kv_cache_skip_step_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_steps) self.diffusion_kv_cache_skip_layer_indices = parse_kv_cache_skip_selector(self.diffusion_kv_cache_skip_layers) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 9ffb484a1..6874d9efe 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -23,6 +23,10 @@ from vllm.model_executor.layers.quantization.base_config import ( ) from vllm.transformers_utils.repo_utils import get_model_path +from vllm_omni.diffusion.diffusion_kv.config import ( + DiffusionKVCacheMode, + parse_diffusion_kv_cache_mode, +) from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata from vllm_omni.diffusion.utils.network_utils import is_port_available from vllm_omni.errors import client_error_metadata @@ -658,6 +662,9 @@ class OmniDiffusionConfig: # the string form only). engine_backend: str | type = "default" + # Local Diffusion KV ownership and cache-layout mode. + diffusion_kv_mode: DiffusionKVCacheMode = DiffusionKVCacheMode.DENSE_LEGACY + # Optional override for the diffusion model runner class (import path). # Precedence in the worker: this override > the runner declared by the # selected engine class (``default_diffusion_model_runner_cls``) > the @@ -923,6 +930,22 @@ class OmniDiffusionConfig: ) if not isinstance(self.diffusion_compile_dynamic, bool): raise TypeError(f"diffusion_compile_dynamic must be a bool, got {type(self.diffusion_compile_dynamic)!r}") + self.diffusion_kv_mode = parse_diffusion_kv_cache_mode(self.diffusion_kv_mode) + + if self.omni_kv_config is None: + self.omni_kv_config = {} + elif isinstance(self.omni_kv_config, Mapping): + self.omni_kv_config = dict(self.omni_kv_config) + else: + raise TypeError("omni_kv_config must be a mapping") + if self.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER and self.omni_kv_config.get( + "need_recv_cache", False + ): + raise ValueError( + "paged_scheduler Diffusion KV does not support imported AR KV in Phase 1; " + "disable need_recv_cache until connector-aware admission is implemented" + ) + self.master_port = self._resolve_master_port() self.request_batch_max_wait_ms = float(self.request_batch_max_wait_ms or 0.0) if self.request_batch_max_wait_ms < 0: diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 6bbf6cfae..e44d935fb 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -310,14 +310,16 @@ class DiffusionEngine: diffusion_engine_start_time = time.perf_counter() preprocess_time = 0.0 - if self.pre_process_func is not None: - preprocess_start_time = time.perf_counter() - request = self.pre_process_func(request) + has_preprocessor = getattr(self, "pre_process_func", None) is not None + preprocess_start_time = time.perf_counter() if has_preprocessor else None + request = self._prepare_request_for_admission(request) + if preprocess_start_time is not None: preprocess_time = time.perf_counter() - preprocess_start_time logger.debug("Pre-processing completed in %.4f seconds", preprocess_time) exec_start_time = time.perf_counter() - generator = self.async_add_req_and_stream_response(request) + request_id = self._add_prepared_request(request) + generator = self.get_output_stream(request_id) async for output in generator: exec_total_time = time.perf_counter() - exec_start_time # Async mode: wait for background D2H/SHM to complete. @@ -689,7 +691,17 @@ class DiffusionEngine: engine.run_startup_warmup() return engine - def add_request(self, request: OmniDiffusionRequest) -> str: + def _prepare_request_for_admission(self, request: OmniDiffusionRequest) -> OmniDiffusionRequest: + """Run model-owned preprocessing once, before entering Engine locks.""" + + pre_process_func = getattr(self, "pre_process_func", None) + if pre_process_func is None: + return request + return pre_process_func(request) + + def _add_prepared_request(self, request: OmniDiffusionRequest) -> str: + """Admit a request whose model-owned preprocessing is complete.""" + with self._cv: if self._closed: raise RuntimeError("DiffusionEngine is closed.") @@ -700,6 +712,10 @@ class DiffusionEngine: return request_id + def add_request(self, request: OmniDiffusionRequest) -> str: + request = self._prepare_request_for_admission(request) + return self._add_prepared_request(request) + async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]: with self._cv: queue = self._out_streams.get(request_id) @@ -742,6 +758,7 @@ class DiffusionEngine: return final_output def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput: + request = self._prepare_request_for_admission(request) with self._rpc_lock: if self._closed: raise RuntimeError("DiffusionEngine is closed.") @@ -886,8 +903,7 @@ class DiffusionEngine: ), ) logger.info("dummy run to warm up the model") - request = self.pre_process_func(req) if self.pre_process_func is not None else req - output = self.add_req_and_wait_for_response(request) + output = self.add_req_and_wait_for_response(req) if output.error: raise RuntimeError(f"Dummy run failed: {output.error}") diff --git a/vllm_omni/diffusion/diffusion_kv/__init__.py b/vllm_omni/diffusion/diffusion_kv/__init__.py new file mode 100644 index 000000000..208f01a7c --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm_omni/diffusion/diffusion_kv/config.py b/vllm_omni/diffusion/diffusion_kv/config.py new file mode 100644 index 000000000..cd69d8988 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/config.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from enum import Enum + + +class DiffusionKVCacheMode(str, Enum): + """Migration mode for diffusion KV ownership.""" + + DENSE_LEGACY = "dense_legacy" + # Reserved for the experimental Worker-owned paging design. Production + # configuration rejects it until that engine is migrated to the common + # Scheduler-owned path. + PAGED_WORKER_LOCAL = "paged_worker_local" + PAGED_SCHEDULER = "paged_scheduler" + + +def parse_diffusion_kv_cache_mode(value: object) -> DiffusionKVCacheMode: + """Parse a selectable Diffusion KV ownership mode. + + ``PAGED_WORKER_LOCAL`` remains in the enum to document the migration + topology, but it is not a valid production selection until that runtime is + wired. Keep this validation shared by structured and runtime configs so an + unsupported mode fails before managed-process startup. + """ + + try: + mode = DiffusionKVCacheMode(value) + except (TypeError, ValueError) as exc: + supported = ", ".join(mode.value for mode in DiffusionKVCacheMode) + raise ValueError(f"Unsupported Diffusion KV diffusion_kv_mode {value!r}; expected one of: {supported}") from exc + if mode is DiffusionKVCacheMode.PAGED_WORKER_LOCAL: + raise ValueError( + "Diffusion KV diffusion_kv_mode 'paged_worker_local' is reserved but not implemented; " + "use 'dense_legacy' or 'paged_scheduler'" + ) + return mode diff --git a/vllm_omni/diffusion/diffusion_kv/request.py b/vllm_omni/diffusion/diffusion_kv/request.py new file mode 100644 index 000000000..4db9701a1 --- /dev/null +++ b/vllm_omni/diffusion/diffusion_kv/request.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +from vllm.v1.core.kv_cache_utils import BlockHash +from vllm.v1.request import RequestStatus + + +@dataclass(frozen=True) +class DiffusionKVContext: + """An independently managed K/V context outside the primary sequence. + + ``context_id`` identifies the logical context within one execution + sequence. ``cache_role`` binds it to a logical attention cache role exposed + by the Worker. Physical cache geometry remains native ``KVCacheSpec`` / + ``KVCacheConfig`` state and is deliberately absent here. + """ + + context_id: str + cache_role: str + num_tokens: int + block_hashes: tuple[BlockHash, ...] = () + + def __post_init__(self) -> None: + if not self.context_id: + raise ValueError("context_id must be non-empty") + if not self.cache_role: + raise ValueError("cache_role must be non-empty") + if self.num_tokens <= 0: + raise ValueError(f"num_tokens must be positive, got {self.num_tokens}") + object.__setattr__(self, "block_hashes", tuple(self.block_hashes)) + + +class DiffusionKVRequest: + """Scheduler-owned KV state for one diffusion execution sequence. + + The primary sequence follows an ordered ``[prefix | target]`` policy, such + as one Hunyuan CFG row. ``prefix_len`` is the contiguous reusable prefix; + ``target_len`` is overwritten by every denoise step; ``num_tokens`` is the + complete first-step allocation boundary. + + Independent cross/joint-attention K/V does not belong to that token axis + and is described by ``kv_contexts``. This object also exposes the minimal + mutable Request surface consumed by native ``KVCacheManager``. + + An empty ``block_hashes`` sequence means the prefix has no canonical cache + identity yet. Such a request may use native request-local page allocation, + but consumers must not publish it through ``KVCacheManager.cache_blocks``. + Prefix publication becomes valid only after model preprocessing supplies + one canonical hash for every cacheable full block. + """ + + def __init__( + self, + request_id: str, + *, + sequence_id: int, + prefix_len: int, + target_len: int, + seq_len: int, + block_hashes: Sequence[BlockHash] = (), + kv_contexts: Sequence[DiffusionKVContext] = (), + ) -> None: + if not request_id: + raise ValueError("request_id must be non-empty") + if sequence_id < 0: + raise ValueError(f"sequence_id must be non-negative, got {sequence_id}") + if prefix_len < 0: + raise ValueError(f"prefix_len must be non-negative, got {prefix_len}") + if target_len <= 0: + raise ValueError(f"target_len must be positive, got {target_len}") + if seq_len <= 0: + raise ValueError(f"seq_len must be positive, got {seq_len}") + if prefix_len + target_len > seq_len: + raise ValueError( + "prefix_len + target_len must not exceed seq_len: " + f"prefix_len={prefix_len}, target_len={target_len}, seq_len={seq_len}" + ) + + contexts = tuple(kv_contexts) + if any(not isinstance(context, DiffusionKVContext) for context in contexts): + raise TypeError("kv_contexts must contain only DiffusionKVContext values") + context_ids = [context.context_id for context in contexts] + if len(context_ids) != len(set(context_ids)): + raise ValueError(f"kv_contexts must have unique context_id values, got {context_ids}") + + # Diffusion semantics consumed by the Scheduler-side facade. + self.sequence_id = sequence_id + self.prefix_len = prefix_len + self.target_len = target_len + self.kv_contexts = contexts + + # Native vLLM Request surface. Keep this list intentionally small and + # cover it with real-KVCacheManager conformance tests. + self.request_id = request_id + self.num_tokens = seq_len + self.num_prompt_tokens = prefix_len + self.num_computed_tokens = 0 + self.block_hashes = list(block_hashes) + self.skip_reading_prefix_cache = not self.block_hashes + self.status = RequestStatus.WAITING + self.num_preemptions = 0 + self.num_in_flight_tokens = 0 + # vLLM 0.26 uses this optional boundary when publishing cached blocks; + # zero means that no sparse-retention boundary is active. + self.shared_prefix_boundary = 0 + + @property + def seq_len(self) -> int: + """Complete first-step sequence length and allocation boundary.""" + + return self.num_tokens diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 36e69302e..4b1da8911 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -21,6 +21,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.transformers_utils.config import get_config from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader @@ -33,6 +34,7 @@ from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer +from . import request_layout as request_layout_utils from .hunyuan_image3_tokenizer import TokenizerEncodeOutput, TokenizerWrapper from .hunyuan_image3_transformer import ( CausalMMOutputWithPast, @@ -50,7 +52,7 @@ from .hunyuan_image3_transformer import ( real_batched_index_select, retrieve_timesteps, ) -from .system_prompt import get_system_prompt +from .request_layout import HunyuanPreparedLayout if TYPE_CHECKING: from vllm_omni.diffusion.worker.input_batch import InputBatch @@ -146,12 +148,6 @@ def _resize_and_crop_center(image: PILImage.Image, target_width: int, target_hei return resized.crop((crop_left, crop_top, crop_left + tw, crop_top + th)) -def _to_python_scalar(value: Any) -> Any: - if isinstance(value, np.generic): - return value.item() - return value - - def _shift_full_attn_spans( spans: list[tuple[int, int]], prefix_len: int | None, @@ -165,80 +161,12 @@ def _shift_full_attn_spans( return [(start + shift, end + shift) if start >= prefix_len else (start, end) for start, end in spans] -def _image_info_to_payload(image_info: ImageInfo) -> dict[str, Any]: - return { - "image_type": image_info.image_type, - "image_tensor": image_info.image_tensor, - "image_width": _to_python_scalar(image_info.image_width), - "image_height": _to_python_scalar(image_info.image_height), - "token_width": _to_python_scalar(image_info.token_width), - "token_height": _to_python_scalar(image_info.token_height), - "image_token_length": _to_python_scalar(image_info.image_token_length), - "base_size": _to_python_scalar(image_info.base_size), - "ratio_index": _to_python_scalar(image_info.ratio_index), - "add_timestep_token": image_info.add_timestep_token, - "add_guidance_token": image_info.add_guidance_token, - "use_front_boi_token": image_info.use_front_boi_token, - "add_image_shape_token": image_info.add_image_shape_token, - } - - -def _to_tensor_if_needed(value: Any) -> Any: - if isinstance(value, np.generic): - return value.item() - if isinstance(value, list): - return torch.tensor(value) - return value - - -def _image_info_from_payload(payload: dict[str, Any]) -> ImageInfo: - return ImageInfo( - image_type=payload.get("image_type"), - image_tensor=_to_tensor_if_needed(payload.get("image_tensor")), - image_width=payload.get("image_width"), - image_height=payload.get("image_height"), - token_width=payload.get("token_width"), - token_height=payload.get("token_height"), - image_token_length=payload.get("image_token_length"), - base_size=payload.get("base_size"), - ratio_index=payload.get("ratio_index"), - add_timestep_token=payload.get("add_timestep_token", True), - add_guidance_token=payload.get("add_guidance_token", False), - use_front_boi_token=payload.get("use_front_boi_token", True), - add_image_shape_token=payload.get("add_image_shape_token", True), - ) - - -def _joint_image_info_to_payload(joint_image_info: JointImageInfo) -> dict[str, Any]: - return { - "type": "joint_image_info", - "vae_image_info": _image_info_to_payload(joint_image_info.vae_image_info), - "vision_image_info": _image_info_to_payload(joint_image_info.vision_image_info), - "vision_encoder_kwargs": joint_image_info.vision_encoder_kwargs, - } - - -def _joint_image_info_from_payload(payload: Any) -> JointImageInfo: - if isinstance(payload, JointImageInfo): - return payload - if not isinstance(payload, dict): - raise TypeError(f"Expected dict or JointImageInfo for conditional image payload, got {type(payload)}.") - - vae_image_info = _image_info_from_payload(payload["vae_image_info"]) - vision_image_info = _image_info_from_payload(payload["vision_image_info"]) - vision_encoder_kwargs = payload.get("vision_encoder_kwargs") or {} - if isinstance(vision_encoder_kwargs, dict): - vision_encoder_kwargs = {key: _to_tensor_if_needed(value) for key, value in vision_encoder_kwargs.items()} - return JointImageInfo( - vae_image_info=vae_image_info, - vision_image_info=vision_image_info, - vision_encoder_kwargs=vision_encoder_kwargs, - ) - - def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig): hf_config = get_config(od_config.model, trust_remote_code=True) image_processor = HunyuanImage3ImageProcessor(hf_config) + prepare_diffusion_kv_layout = od_config.diffusion_kv_mode is DiffusionKVCacheMode.PAGED_SCHEDULER + tokenizer_wrapper = TokenizerWrapper(od_config.model) if prepare_diffusion_kv_layout else None + generation_config = GenerationConfig.from_pretrained(od_config.model) if prepare_diffusion_kv_layout else None vae_h_factor = hf_config.vae_downsample_factor[0] * hf_config.patch_size vae_w_factor = hf_config.vae_downsample_factor[1] * hf_config.patch_size vit_patch_size = getattr(image_processor.vision_encoder_processor, "patch_size", 1) @@ -291,7 +219,7 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig): image_token_length=int(vit_tensor.shape[1]), ) - return _joint_image_info_to_payload( + return request_layout_utils.joint_image_info_to_payload( JointImageInfo( vae_image_info=vae_info, vision_image_info=vit_info, @@ -329,6 +257,20 @@ def get_hunyuan_image_3_pre_process_func(od_config: OmniDiffusionConfig): request.sampling_params.height = int(bridge_h or first_image_h) request.prompt = prompt + if prepare_diffusion_kv_layout: + assert tokenizer_wrapper is not None and generation_config is not None + prepared_layout = request_layout_utils.prepare_hunyuan_layout( + request, + tokenizer_wrapper=tokenizer_wrapper, + image_processor=image_processor, + generation_config=generation_config, + image_base_size=hf_config.image_base_size, + ) + request.prepared_layout = prepared_layout + request.diffusion_kv_requests = request_layout_utils.build_hunyuan_diffusion_kv_requests( + request, + prepared_layout, + ) return request @@ -543,94 +485,6 @@ class HunyuanImage3Pipeline( if getattr(self.od_config, "diffusion_kv_cache_skip_step_indices", None): raise ValueError("HunyuanImage3 step execution does not support diffusion KV cache step skips yet.") - @staticmethod - def _normalize_single_stage_bot_task(bot_task: Any) -> str: - if isinstance(bot_task, str) and bot_task.lower() == "none": - bot_task = None - tokenizer_bot_task = bot_task - if tokenizer_bot_task == "think_recaption": - tokenizer_bot_task = "think" - elif tokenizer_bot_task == "vanilla": - tokenizer_bot_task = "image" - supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"} - if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks: - raise ValueError( - f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. " - f"Supported values are: {sorted(supported_bot_tasks)}." - ) - return tokenizer_bot_task or "auto" - - def _extract_prompt_inputs( - self, - prompts: list[Any], - extra_args: dict[str, Any], - *, - request_id: str, - allow_cond_image: bool, - ) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]: - is_dummy_warmup = OmniDiffusionRequest.is_dummy_run_request_id(request_id) - bot_task = extra_args.get("bot_task") - use_system_prompt = extra_args.get("use_system_prompt") - system_prompt = extra_args.get("system_prompt") - - first_prompt = prompts[0] if prompts else None - if isinstance(first_prompt, dict): - if bot_task is None: - bot_task = first_prompt.get("bot_task") - if use_system_prompt is None: - use_system_prompt = first_prompt.get("use_system_prompt") - if system_prompt is None: - system_prompt = first_prompt.get("system_prompt") - tokenizer_bot_task = self._normalize_single_stage_bot_task(bot_task) - if use_system_prompt is not None: - system_prompt_bot_task = "image" if tokenizer_bot_task == "auto" else tokenizer_bot_task - system_prompt = get_system_prompt(use_system_prompt, system_prompt_bot_task, system_prompt) - system_prompt = system_prompt.strip() if system_prompt is not None else "" - - prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] - cot_text_list = [ - (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in prompts - ] - - batch_cond_image_info: list[list[JointImageInfo]] | None = None - if any(not isinstance(p, str) for p in prompts): - batch_cond_image_info = [] - for prompt_item in prompts: - if isinstance(prompt_item, str): - batch_cond_image_info.append([]) - continue - additional_info = prompt_item.get("additional_information") or {} - cond_infos = additional_info.get("batch_cond_image_info", []) - if isinstance(cond_infos, JointImageInfo | dict): - cond_infos = [cond_infos] - if cond_infos is None: - cond_infos = [] - batch_cond_image_info.append([_joint_image_info_from_payload(item) for item in cond_infos]) - - has_cond_image = [len(cond_infos) > 0 for cond_infos in batch_cond_image_info] - if any(has_cond_image) and (not allow_cond_image) and not is_dummy_warmup: - raise ValueError("HunyuanImage3 step execution does not support image editing requests yet.") - if allow_cond_image and any(has_cond_image) and not all(has_cond_image): - raise ValueError( - "When batching Hunyuan image editing requests, every prompt must include input image(s)." - ) - if not allow_cond_image or not any(has_cond_image): - batch_cond_image_info = None - - return prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task - - def _extract_step_prompt_inputs( - self, - state: "StepRequestState", - ) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]: - sampling = state.sampling - return self._extract_prompt_inputs( - [state.prompt] if state.prompt is not None else [], - getattr(sampling, "extra_args", {}) or {}, - request_id=state.request_id, - allow_cond_image=False, - ) - def _snapshot_injected_ar_kv(self) -> list[list[tuple[torch.Tensor, torch.Tensor]] | None] | None: snapshot: list[list[tuple[torch.Tensor, torch.Tensor]] | None] = [] found = False @@ -1081,27 +935,6 @@ class HunyuanImage3Pipeline( pred = self.final_layer(image_output, timestep_emb, token_h, token_w) return pred - @staticmethod - def build_batch_rope_image_info(output, sections): - rope_image_info = [] - for image_slices, sections_i in zip(output.all_image_slices, sections): - image_shapes = [] - for section in sections_i: - if "image" in section["type"]: - if isinstance(section["token_height"], list): - assert len(section["token_height"]) == len(section["token_width"]), ( - f"token_height and token_width should have the same length, " - f"but got {len(section['token_height'])} and {len(section['token_width'])}" - ) - image_shapes.extend(list(zip(section["token_height"], section["token_width"]))) - else: - image_shapes.append((section["token_height"], section["token_width"])) - assert len(image_slices) == len(image_shapes), ( - f"Size miss matching: Image slices({len(image_slices)}) != image shapes({len(image_shapes)})" - ) - rope_image_info.append(list(zip(image_slices, image_shapes))) - return rope_image_info - def vae_encode(self, image, cfg_factor=1, generator=None): config = self.vae.config @@ -1210,26 +1043,6 @@ class HunyuanImage3Pipeline( f"Each message should be a list of dicts or a dict, but got {type(message)}." ) - @staticmethod - def _normalize_cot_text(cot: str | None) -> str | None: - """ - Ensure cot_text starts with the appropriate opening tag. - - AR-generated text may omit the leading or tag - (since it was used as a generation trigger). This normalizes the text - so downstream parsing in get_cot_sections works correctly. - """ - if not cot: - return cot - - if "" in cot and not cot.startswith(""): - cot = "" + cot - return cot - if "" in cot and not cot.startswith(""): - cot = "" + cot - return cot - return cot - def prepare_model_inputs( self, prompt=None, @@ -1242,6 +1055,7 @@ class HunyuanImage3Pipeline( message_list=None, device=None, max_new_tokens=None, + prepared_layout: HunyuanPreparedLayout | None = None, **kwargs, ): # 1. Sanity check @@ -1297,7 +1111,16 @@ class HunyuanImage3Pipeline( ) if mode == "gen_image": - batch_gen_image_info = [self.image_processor.build_image_info(image_size) for _ in range(batch_size)] + if prepared_layout is not None: + # TODO: Support batch_size > 1 once prepared layouts can + # preserve the public-request-to-CFG-branch mapping. + if batch_size != 1: + raise ValueError("Hunyuan prepared layout currently supports one public request at a time.") + batch_gen_image_info = [prepared_layout.generated_image_info] + else: + batch_gen_image_info = [ + self.image_processor.build_image_info(image_size) for _ in range(batch_size) + ] if batch_cond_image_info is not None: assert isinstance(batch_cond_image_info, list) and len(batch_cond_image_info) == batch_size, ( @@ -1314,30 +1137,41 @@ class HunyuanImage3Pipeline( # 3. apply chat template cfg_factor = {"gen_text": 1, "gen_image": 1 + int(guidance_scale > 1.0)} bot_task = kwargs.pop("bot_task", "auto") - # If `drop_think` enabled, always drop parts in the context. - drop_think = kwargs.get("drop_think", self.generation_config.drop_think) - # Pull sequence_template from the model's generation_config so the DiT - # text prefix matches how the model was trained (Instruct for the - # HunyuanImage-3.0-Instruct checkpoint). Falling back to "pretrain" - # only if the config does not specify it. - sequence_template = getattr(self.generation_config, "sequence_template", "pretrain") - # Apply batched prompt or batched message_list to build input sequence with associated info. - out = self._tkwrapper.apply_chat_template( - batch_prompt=batch_prompt, - batch_message_list=batch_message_list, - mode=mode, - batch_gen_image_info=batch_gen_image_info, - batch_cond_image_info=batch_cond_image_info, - batch_system_prompt=batch_system_prompt, - batch_cot_text=batch_cot_text, - max_length=kwargs.get("max_length"), - bot_task=bot_task, - image_base_size=self.config.image_base_size, - sequence_template=sequence_template, - cfg_factor=cfg_factor[mode], - drop_think=drop_think, - ) - output, sections = out["output"], out["sections"] + if prepared_layout is None: + # Pull template options from generation_config only when the + # Worker must tokenize locally. + drop_think = kwargs.get("drop_think", self.generation_config.drop_think) + sequence_template = getattr(self.generation_config, "sequence_template", "pretrain") + # Apply batched prompt or batched message_list to build input + # sequence with associated info. + out = self._tkwrapper.apply_chat_template( + batch_prompt=batch_prompt, + batch_message_list=batch_message_list, + mode=mode, + batch_gen_image_info=batch_gen_image_info, + batch_cond_image_info=batch_cond_image_info, + batch_system_prompt=batch_system_prompt, + batch_cot_text=batch_cot_text, + max_length=kwargs.get("max_length"), + bot_task=bot_task, + image_base_size=self.config.image_base_size, + sequence_template=sequence_template, + cfg_factor=cfg_factor[mode], + drop_think=drop_think, + ) + output = out["output"] + rope_image_info = request_layout_utils.build_hunyuan_batch_rope_image_info(output, out["sections"]) + else: + if mode != "gen_image" or batch_message_list is not None: + raise ValueError("Hunyuan prepared layout only supports prompt-based gen_image execution.") + output = prepared_layout.tokenizer_output + rope_image_info = prepared_layout.rope_image_info + expected_rows = batch_size * cfg_factor[mode] + if prepared_layout.num_branches != expected_rows: + raise ValueError( + "Hunyuan prepared branches do not match execution CFG layout: " + f"expected_rows={expected_rows}, prepared_branches={prepared_layout.num_branches}" + ) # 4. Encode conditional images # Skip encoding if AR KV reuse is enabled @@ -1362,7 +1196,6 @@ class HunyuanImage3Pipeline( vit_kwargs = None # 5. Build position embeddings - rope_image_info = self.build_batch_rope_image_info(output, sections) if mode == "gen_text": seq_len = self.generation_config.max_length else: @@ -1616,12 +1449,7 @@ class HunyuanImage3Pipeline( raise ValueError("`batch_gen_image_info` should be provided when `mode` is `gen_image`.") image_info: ImageInfo = batch_gen_image_info[0] - num_image_tokens = ( - image_info.image_token_length - + (1 if image_info.add_timestep_token else 0) - + (1 if image_info.add_guidance_token else 0) - ) - kwargs["num_image_tokens"] = num_image_tokens + kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info) # 50 and 5.0 hard code results = self.pipeline( batch_size=len(batch_gen_image_info), @@ -1888,9 +1716,14 @@ class HunyuanImage3Pipeline( system_prompt, batch_cond_image_info, tokenizer_bot_task, - ) = self._extract_step_prompt_inputs(state) + ) = request_layout_utils.extract_hunyuan_prompt_inputs( + [state.prompt] if state.prompt is not None else [], + getattr(sampling, "extra_args", {}) or {}, + request_id=state.request_id, + allow_cond_image=False, + ) cot_text = ( - [self._normalize_cot_text(text) for text in cot_text_list] + [request_layout_utils.normalize_hunyuan_cot_text(text) for text in cot_text_list] if any(text is not None for text in cot_text_list) else None ) @@ -1899,7 +1732,7 @@ class HunyuanImage3Pipeline( width = sampling.width or 1024 image_size = (height, width) num_inference_steps = sampling.num_inference_steps or 50 - guidance_scale = sampling.guidance_scale if sampling.guidance_scale_provided else 5.0 + guidance_scale = request_layout_utils.resolve_hunyuan_guidance_scale(sampling) if guidance_scale <= 1.0: logger.info("HunyuanImage3.0 step execution runs without classifier-free guidance.") pipe._guidance_scale = guidance_scale @@ -1916,6 +1749,7 @@ class HunyuanImage3Pipeline( guidance_scale=guidance_scale, batch_cond_image_info=batch_cond_image_info, bot_task=tokenizer_bot_task, + prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(state), ) model_kwargs.update(self._extract_ar_kv_from_sampling(sampling)) model_kwargs["use_cache"] = False @@ -1923,14 +1757,9 @@ class HunyuanImage3Pipeline( input_ids = model_kwargs.pop("input_ids") batch_gen_image_info: list[ImageInfo] = model_kwargs["batch_gen_image_info"] image_info = batch_gen_image_info[0] - target_height = int(_to_python_scalar(image_info.image_height)) - target_width = int(_to_python_scalar(image_info.image_width)) - num_image_tokens = ( - image_info.image_token_length - + (1 if image_info.add_timestep_token else 0) - + (1 if image_info.add_guidance_token else 0) - ) - model_kwargs["num_image_tokens"] = num_image_tokens + target_height = int(image_info.image_height) + target_width = int(image_info.image_width) + model_kwargs["num_image_tokens"] = request_layout_utils.hunyuan_num_image_tokens(image_info) timesteps, _ = retrieve_timesteps( self.scheduler, @@ -2294,7 +2123,7 @@ class HunyuanImage3Pipeline( system_prompt, batch_cond_image_info, tokenizer_bot_task, - ) = self._extract_prompt_inputs( + ) = request_layout_utils.extract_hunyuan_prompt_inputs( req.prompts, extra_args, request_id=req.request_id, @@ -2302,15 +2131,16 @@ class HunyuanImage3Pipeline( ) prompt = prompt_from_req or prompt cot_text = ( - [self._normalize_cot_text(t) for t in cot_text_list] if any(t is not None for t in cot_text_list) else None + [request_layout_utils.normalize_hunyuan_cot_text(t) for t in cot_text_list] + if any(t is not None for t in cot_text_list) + else None ) generator = req.sampling_params.generator or generator height = req.sampling_params.height or height width = req.sampling_params.width or width num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps - if req.sampling_params.guidance_scale_provided: - guidance_scale = req.sampling_params.guidance_scale + guidance_scale = request_layout_utils.resolve_hunyuan_guidance_scale(req.sampling_params, guidance_scale) if guidance_scale <= 1.0: logger.info("HunyuanImage3.0 runs without classifier-free guidance when guidance_scale <= 1.0.") image_size = (height, width) @@ -2329,6 +2159,7 @@ class HunyuanImage3Pipeline( guidance_scale=guidance_scale, batch_cond_image_info=batch_cond_image_info, bot_task=tokenizer_bot_task, + prepared_layout=request_layout_utils.get_hunyuan_prepared_layout(req.requests[0]), **ar_kv_kwargs, ) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py b/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py new file mode 100644 index 000000000..d09c20c52 --- /dev/null +++ b/vllm_omni/diffusion/models/hunyuan_image3/request_layout.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest +from vllm_omni.diffusion.request import OmniDiffusionRequest + +from .hunyuan_image3_tokenizer import TokenizerEncodeOutput +from .hunyuan_image3_transformer import ImageInfo, JointImageInfo +from .system_prompt import get_system_prompt + +if TYPE_CHECKING: + from transformers.generation.configuration_utils import GenerationConfig + + from .hunyuan_image3_tokenizer import TokenizerWrapper + from .hunyuan_image3_transformer import HunyuanImage3ImageProcessor + + +def _to_python_scalar(value: Any) -> Any: + if isinstance(value, np.generic): + return value.item() + return value + + +def _to_tensor_if_needed(value: Any) -> Any: + if isinstance(value, np.generic): + return value.item() + if isinstance(value, list): + return torch.tensor(value) + return value + + +def _image_info_to_payload(image_info: ImageInfo) -> dict[str, Any]: + return { + "image_type": image_info.image_type, + "image_tensor": image_info.image_tensor, + "image_width": _to_python_scalar(image_info.image_width), + "image_height": _to_python_scalar(image_info.image_height), + "token_width": _to_python_scalar(image_info.token_width), + "token_height": _to_python_scalar(image_info.token_height), + "image_token_length": _to_python_scalar(image_info.image_token_length), + "base_size": _to_python_scalar(image_info.base_size), + "ratio_index": _to_python_scalar(image_info.ratio_index), + "add_timestep_token": image_info.add_timestep_token, + "add_guidance_token": image_info.add_guidance_token, + "use_front_boi_token": image_info.use_front_boi_token, + "add_image_shape_token": image_info.add_image_shape_token, + } + + +def _image_info_from_payload(payload: dict[str, Any]) -> ImageInfo: + return ImageInfo( + image_type=payload.get("image_type"), + image_tensor=_to_tensor_if_needed(payload.get("image_tensor")), + image_width=payload.get("image_width"), + image_height=payload.get("image_height"), + token_width=payload.get("token_width"), + token_height=payload.get("token_height"), + image_token_length=payload.get("image_token_length"), + base_size=payload.get("base_size"), + ratio_index=payload.get("ratio_index"), + add_timestep_token=payload.get("add_timestep_token", True), + add_guidance_token=payload.get("add_guidance_token", False), + use_front_boi_token=payload.get("use_front_boi_token", True), + add_image_shape_token=payload.get("add_image_shape_token", True), + ) + + +def joint_image_info_to_payload(joint_image_info: JointImageInfo) -> dict[str, Any]: + return { + "type": "joint_image_info", + "vae_image_info": _image_info_to_payload(joint_image_info.vae_image_info), + "vision_image_info": _image_info_to_payload(joint_image_info.vision_image_info), + "vision_encoder_kwargs": joint_image_info.vision_encoder_kwargs, + } + + +def _joint_image_info_from_payload(payload: Any) -> JointImageInfo: + if isinstance(payload, JointImageInfo): + return payload + if not isinstance(payload, dict): + raise TypeError(f"Expected dict or JointImageInfo for conditional image payload, got {type(payload)}.") + + vae_image_info = _image_info_from_payload(payload["vae_image_info"]) + vision_image_info = _image_info_from_payload(payload["vision_image_info"]) + vision_encoder_kwargs = payload.get("vision_encoder_kwargs") or {} + if isinstance(vision_encoder_kwargs, dict): + vision_encoder_kwargs = {key: _to_tensor_if_needed(value) for key, value in vision_encoder_kwargs.items()} + return JointImageInfo( + vae_image_info=vae_image_info, + vision_image_info=vision_image_info, + vision_encoder_kwargs=vision_encoder_kwargs, + ) + + +def normalize_hunyuan_single_stage_bot_task(bot_task: Any) -> str: + if isinstance(bot_task, str) and bot_task.lower() == "none": + bot_task = None + tokenizer_bot_task = bot_task + if tokenizer_bot_task == "think_recaption": + tokenizer_bot_task = "think" + elif tokenizer_bot_task == "vanilla": + tokenizer_bot_task = "image" + supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"} + if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks: + raise ValueError( + f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. " + f"Supported values are: {sorted(supported_bot_tasks)}." + ) + return tokenizer_bot_task or "auto" + + +def normalize_hunyuan_cot_text(cot: str | None) -> str | None: + """Restore an AR generation trigger tag omitted from generated text.""" + + if not cot: + return cot + if "" in cot and not cot.startswith(""): + return "" + cot + if "" in cot and not cot.startswith(""): + return "" + cot + return cot + + +def extract_hunyuan_prompt_inputs( + prompts: list[Any], + extra_args: dict[str, Any], + *, + request_id: str, + allow_cond_image: bool, +) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]: + """Normalize request prompt fields shared by planning and execution.""" + + is_dummy_warmup = OmniDiffusionRequest.is_dummy_run_request_id(request_id) + bot_task = extra_args.get("bot_task") + use_system_prompt = extra_args.get("use_system_prompt") + system_prompt = extra_args.get("system_prompt") + + first_prompt = prompts[0] if prompts else None + if isinstance(first_prompt, dict): + if bot_task is None: + bot_task = first_prompt.get("bot_task") + if use_system_prompt is None: + use_system_prompt = first_prompt.get("use_system_prompt") + if system_prompt is None: + system_prompt = first_prompt.get("system_prompt") + tokenizer_bot_task = normalize_hunyuan_single_stage_bot_task(bot_task) + if use_system_prompt is not None: + system_prompt_bot_task = "image" if tokenizer_bot_task == "auto" else tokenizer_bot_task + system_prompt = get_system_prompt(use_system_prompt, system_prompt_bot_task, system_prompt) + system_prompt = system_prompt.strip() if system_prompt is not None else "" + + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in prompts] + cot_text_list = [ + (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in prompts + ] + + batch_cond_image_info: list[list[JointImageInfo]] | None = None + if any(not isinstance(p, str) for p in prompts): + batch_cond_image_info = [] + for prompt_item in prompts: + if isinstance(prompt_item, str): + batch_cond_image_info.append([]) + continue + additional_info = prompt_item.get("additional_information") or {} + cond_infos = additional_info.get("batch_cond_image_info", []) + if isinstance(cond_infos, JointImageInfo | dict): + cond_infos = [cond_infos] + if cond_infos is None: + cond_infos = [] + batch_cond_image_info.append([_joint_image_info_from_payload(item) for item in cond_infos]) + + has_cond_image = [len(cond_infos) > 0 for cond_infos in batch_cond_image_info] + if any(has_cond_image) and (not allow_cond_image) and not is_dummy_warmup: + raise ValueError("HunyuanImage3 step execution does not support image editing requests yet.") + if allow_cond_image and any(has_cond_image) and not all(has_cond_image): + raise ValueError("When batching Hunyuan image editing requests, every prompt must include input image(s).") + if not allow_cond_image or not any(has_cond_image): + batch_cond_image_info = None + + return prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task + + +def resolve_hunyuan_guidance_scale(sampling: Any, default_scale: float = 5.0) -> float: + if getattr(sampling, "guidance_scale_provided", False): + return sampling.guidance_scale + return default_scale + + +def hunyuan_num_image_tokens(image_info: ImageInfo) -> int: + """Return the generated-image span overwritten on every denoise step.""" + + return int(image_info.image_token_length) + int(image_info.add_timestep_token) + int(image_info.add_guidance_token) + + +def build_hunyuan_batch_rope_image_info( + output: TokenizerEncodeOutput, + sections: list[list[dict[str, Any]]], +) -> list[list[tuple[slice, tuple[int, int]]]]: + if output.all_image_slices is None: + raise ValueError("Hunyuan tokenizer output is missing all_image_slices.") + if len(output.all_image_slices) != len(sections): + raise ValueError( + "Hunyuan image-slice rows do not match template section rows: " + f"image_slices={len(output.all_image_slices)}, sections={len(sections)}" + ) + rope_image_info: list[list[tuple[slice, tuple[int, int]]]] = [] + for image_slices, sections_i in zip(output.all_image_slices, sections): + image_shapes: list[tuple[int, int]] = [] + for section in sections_i: + if "image" not in section["type"]: + continue + if isinstance(section["token_height"], list): + if len(section["token_height"]) != len(section["token_width"]): + raise ValueError( + "token_height and token_width should have the same length, " + f"but got {len(section['token_height'])} and {len(section['token_width'])}" + ) + image_shapes.extend( + (int(token_height), int(token_width)) + for token_height, token_width in zip(section["token_height"], section["token_width"]) + ) + else: + image_shapes.append((int(section["token_height"]), int(section["token_width"]))) + if len(image_slices) != len(image_shapes): + raise ValueError(f"Image slices({len(image_slices)}) do not match image shapes({len(image_shapes)}).") + rope_image_info.append(list(zip(image_slices, image_shapes))) + return rope_image_info + + +def prepare_hunyuan_layout( + request: OmniDiffusionRequest, + *, + tokenizer_wrapper: TokenizerWrapper, + image_processor: HunyuanImage3ImageProcessor, + generation_config: GenerationConfig, + image_base_size: int, +) -> HunyuanPreparedLayout: + """Build the CPU token/image layout reused by Scheduler and Worker.""" + + sampling = request.sampling_params + extra_args = getattr(sampling, "extra_args", {}) or {} + prompt, cot_text_list, system_prompt, batch_cond_image_info, tokenizer_bot_task = extract_hunyuan_prompt_inputs( + [request.prompt], + extra_args, + request_id=request.request_id, + allow_cond_image=True, + ) + cot_text = ( + [normalize_hunyuan_cot_text(text) for text in cot_text_list] + if any(text is not None for text in cot_text_list) + else None + ) + height = sampling.height or 1024 + width = sampling.width or 1024 + guidance_scale = resolve_hunyuan_guidance_scale(sampling) + generated_image_info = image_processor.build_image_info((height, width)) + result = tokenizer_wrapper.apply_chat_template( + batch_prompt=prompt, + mode="gen_image", + batch_gen_image_info=[generated_image_info], + batch_cond_image_info=batch_cond_image_info, + batch_system_prompt=[system_prompt], + batch_cot_text=cot_text, + max_length=None, + bot_task=tokenizer_bot_task, + image_base_size=image_base_size, + sequence_template=getattr(generation_config, "sequence_template", "pretrain"), + cfg_factor=1 + int(guidance_scale > 1.0), + drop_think=getattr(generation_config, "drop_think", False), + ) + tokenizer_output = result["output"] + return HunyuanPreparedLayout( + tokenizer_output=tokenizer_output, + rope_image_info=build_hunyuan_batch_rope_image_info(tokenizer_output, result["sections"]), + generated_image_info=generated_image_info, + ) + + +@dataclass(frozen=True) +class HunyuanPreparedLayout: + """CPU execution layout prepared once before Scheduler admission.""" + + tokenizer_output: TokenizerEncodeOutput + rope_image_info: list[list[tuple[slice, tuple[int, int]]]] + generated_image_info: ImageInfo + + def __post_init__(self) -> None: + tokens = self.tokenizer_output.tokens + if tokens is None or tokens.ndim != 2: + shape = None if tokens is None else tuple(tokens.shape) + raise ValueError(f"Hunyuan prepared tokenizer tokens must be 2-D, got {shape}") + if len(self.rope_image_info) != int(tokens.shape[0]): + raise ValueError( + "Hunyuan prepared RoPE rows do not match tokenizer rows: " + f"rope_rows={len(self.rope_image_info)}, tokens={int(tokens.shape[0])}" + ) + prefix_positions = self.tokenizer_output.gen_timestep_scatter_index + if prefix_positions is None or prefix_positions.ndim != 2 or prefix_positions.shape[1] == 0: + raise ValueError("Hunyuan prepared layout requires a non-empty 2-D gen_timestep_scatter_index") + if prefix_positions.shape[0] != int(tokens.shape[0]): + raise ValueError( + "Hunyuan prepared prefix-position rows do not match tokenizer rows: " + f"prefix_positions={prefix_positions.shape[0]}, tokens={int(tokens.shape[0])}" + ) + real_pos = self.tokenizer_output.real_pos + if real_pos is None or real_pos.ndim != 2 or real_pos.shape[1] == 0: + raise ValueError("Hunyuan prepared layout requires a non-empty 2-D real_pos") + if real_pos.shape[0] != int(tokens.shape[0]): + raise ValueError( + "Hunyuan prepared valid-length rows do not match tokenizer rows: " + f"real_pos={real_pos.shape[0]}, tokens={int(tokens.shape[0])}" + ) + seq_lens = [int(row[-1].item()) for row in real_pos] + if any(seq_len <= 0 or seq_len > int(tokens.shape[1]) for seq_len in seq_lens): + raise ValueError( + "Hunyuan prepared valid sequence lengths must fit the padded tokenizer width: " + f"seq_lens={seq_lens}, padded_width={int(tokens.shape[1])}" + ) + + @property + def num_branches(self) -> int: + return int(self.tokenizer_output.tokens.shape[0]) + + +def build_hunyuan_diffusion_kv_requests( + request: OmniDiffusionRequest, + prepared_layout: HunyuanPreparedLayout, +) -> tuple[DiffusionKVRequest, ...]: + """Build one persistent Scheduler KV request per Hunyuan execution row.""" + + tokenizer_output = prepared_layout.tokenizer_output + cfg_factor = 1 + int(resolve_hunyuan_guidance_scale(request.sampling_params) > 1.0) + if prepared_layout.num_branches != cfg_factor: + raise ValueError( + "Hunyuan tokenizer sequence count does not match CFG execution: " + f"rows={prepared_layout.num_branches}, cfg_factor={cfg_factor}" + ) + prefix_positions = tokenizer_output.gen_timestep_scatter_index + real_pos = tokenizer_output.real_pos + assert prefix_positions is not None and real_pos is not None + + target_len = hunyuan_num_image_tokens(prepared_layout.generated_image_info) + return tuple( + DiffusionKVRequest( + f"{request.request_id}/diffusion-kv/{sequence_id}", + sequence_id=sequence_id, + # The generated-image timestep position terminates the reusable + # prompt/reference-image prefix for this execution row. + prefix_len=int(prefix_row[-1].item()), + target_len=target_len, + seq_len=int(valid_row[-1].item()), + # Prompt and reference-image tokens are already embedded in this + # row's primary self-attention sequence. Hunyuan therefore has no + # independently projected cross/joint-attention KV context. + kv_contexts=(), + ) + for sequence_id, (prefix_row, valid_row) in enumerate(zip(prefix_positions, real_pos)) + ) + + +def get_hunyuan_prepared_layout(source: Any) -> HunyuanPreparedLayout | None: + prepared_layout = getattr(source, "prepared_layout", None) + if prepared_layout is None: + return None + if not isinstance(prepared_layout, HunyuanPreparedLayout): + raise TypeError( + f"HunyuanImage3 expected prepared_layout to be HunyuanPreparedLayout, got {type(prepared_layout).__name__}" + ) + return prepared_layout diff --git a/vllm_omni/diffusion/request.py b/vllm_omni/diffusion/request.py index e1147c3d6..8834e408e 100644 --- a/vllm_omni/diffusion/request.py +++ b/vllm_omni/diffusion/request.py @@ -2,12 +2,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + import random from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType +if TYPE_CHECKING: + from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest + DUMMY_DIFFUSION_REQUEST_ID = "dummy_req_id" @@ -30,6 +35,14 @@ class OmniDiffusionRequest: sampling_params: OmniDiffusionSamplingParams request_id: str kv_sender_info: dict[str, Any] | None = None + # Optional opaque, model-owned input prepared before Scheduler admission. + # Model code validates its concrete type when consuming it on the Worker. + prepared_layout: Any | None = None + # Scheduler-only native-compatible KV requests produced by model + # preprocessing. + # BaseScheduler takes ownership and clears this field before the request is + # sent to a Worker. + diffusion_kv_requests: tuple[DiffusionKVRequest, ...] | None = None def __post_init__(self): """Initialize dependent fields after dataclass initialization.""" diff --git a/vllm_omni/diffusion/sched/base_scheduler.py b/vllm_omni/diffusion/sched/base_scheduler.py index 5ea91b167..042bed0d1 100644 --- a/vllm_omni/diffusion/sched/base_scheduler.py +++ b/vllm_omni/diffusion/sched/base_scheduler.py @@ -10,6 +10,7 @@ from dataclasses import fields from vllm.logger import init_logger from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import ( CachedRequestData, @@ -47,6 +48,7 @@ class BaseScheduler(ABC): self._finished_req_ids: set[str] = set() self.max_num_running_reqs: int = 1 self._prefetch_enabled: bool = False + self._diffusion_kv_enabled = False def initialize(self, od_config: OmniDiffusionConfig) -> None: self.od_config = od_config @@ -63,6 +65,10 @@ class BaseScheduler(ABC): self.max_num_running_reqs = 1 omni_kv = getattr(od_config, "omni_kv_config", None) or {} self._prefetch_enabled = bool(omni_kv.get("enable_kv_async_prefetch", False)) + self._diffusion_kv_enabled = ( + getattr(od_config, "diffusion_kv_mode", DiffusionKVCacheMode.DENSE_LEGACY) + is DiffusionKVCacheMode.PAGED_SCHEDULER + ) self._reset_scheduler_state() def add_request(self, request: OmniDiffusionRequest) -> str: @@ -251,12 +257,43 @@ class BaseScheduler(ABC): """Remove subclass-owned per-request state before popping request state.""" def _make_request_state(self, request_id: str, request: OmniDiffusionRequest) -> SchedulerRequestState: + kv_requests = request.diffusion_kv_requests or () + if self._diffusion_kv_enabled: + self._reject_legacy_dense_kv(request) + if not kv_requests: + raise ValueError("paged_scheduler request preprocessing did not produce DiffusionKVRequest state") + elif kv_requests: + raise ValueError("dense_legacy request unexpectedly contains Scheduler Diffusion KV requests") + + # DiffusionKVRequest objects are mutable Scheduler/native-KVCacheManager + # state and must never ride the normal request payload to a Worker. + request.diffusion_kv_requests = None return SchedulerRequestState( request_id=request_id, req=request, sampling_params_key=self._build_sampling_params_key(request), + diffusion_kv_requests=kv_requests, ) + @staticmethod + def _reject_legacy_dense_kv(request: OmniDiffusionRequest) -> None: + """Keep dense injected KV out of the Scheduler-owned paged path.""" + + populated_fields: list[str] = [] + for owner_name, owner in ( + ("request", request), + ("sampling_params", request.sampling_params), + ): + for field_name, value in vars(owner).items(): + if value is not None and (field_name == "past_key_values" or field_name.endswith("_past_key_values")): + populated_fields.append(f"{owner_name}.{field_name}") + if populated_fields: + fields_text = ", ".join(sorted(populated_fields)) + raise ValueError( + "paged_scheduler Diffusion KV does not accept legacy dense KV payloads; " + f"clear these fields before admission: {fields_text}" + ) + def _can_schedule_waiting(self, state: SchedulerRequestState) -> bool: if not self._running: return True diff --git a/vllm_omni/diffusion/sched/interface.py b/vllm_omni/diffusion/sched/interface.py index d292084cd..a08108b79 100644 --- a/vllm_omni/diffusion/sched/interface.py +++ b/vllm_omni/diffusion/sched/interface.py @@ -6,10 +6,13 @@ from __future__ import annotations import enum from dataclasses import dataclass from functools import cached_property -from typing import Any, TypedDict +from typing import TYPE_CHECKING, Any, TypedDict from vllm_omni.diffusion.request import OmniDiffusionRequest +if TYPE_CHECKING: + from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest + class DiffusionRequestStatus(enum.IntEnum): """Request status tracked by diffusion scheduler.""" @@ -129,6 +132,7 @@ class SchedulerRequestState: request_id: str req: OmniDiffusionRequest sampling_params_key: StepBatchSamplingParamsKey | RequestBatchSamplingParamsKey | None = None + diffusion_kv_requests: tuple[DiffusionKVRequest, ...] = () status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING error: str | None = None diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index ade911620..52c2c646d 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -641,6 +641,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): sampling=copy.deepcopy(sched_new_req.req.sampling_params), prompt=sched_new_req.req.prompt, kv_sender_info=sched_new_req.req.kv_sender_info, + prepared_layout=getattr(sched_new_req.req, "prepared_layout", None), ) state_req = copy.copy(sched_new_req.req) state_req.sampling_params = new_state.sampling diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index 40ed550d0..62def6a0c 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/vllm_omni/diffusion/worker/utils.py @@ -77,6 +77,9 @@ class StepRequestState: sampling: OmniDiffusionSamplingParams prompt: OmniPromptType | None = None kv_sender_info: dict | None = None + # Opaque model-owned preprocessing result. Scheduler-owned Diffusion KV + # request state is removed before this Worker state is built. + prepared_layout: Any | None = None # ── Encoded prompts (set once by prepare_encode) ── prompt_embeds: torch.Tensor | None = None