diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 96c139c8f..f880455d4 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -64,7 +64,8 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea docker pull "${image_name}" remove_docker_container() { - docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true + docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true + docker image rm -f "${image_name}" || true } trap remove_docker_container EXIT @@ -107,7 +108,6 @@ if [[ $commands == *"--shard-id="* ]]; then --network=host \ --shm-size=16gb \ --group-add "$render_gid" \ - --rm \ -e MIOPEN_DEBUG_CONV_DIRECT=0 \ -e MIOPEN_DEBUG_CONV_GEMM=0 \ -e VLLM_ROCM_USE_AITER=0 \ @@ -140,7 +140,6 @@ else --network=host \ --shm-size=16gb \ --group-add "$render_gid" \ - --rm \ -e MIOPEN_DEBUG_CONV_DIRECT=0 \ -e MIOPEN_DEBUG_CONV_GEMM=0 \ -e VLLM_ROCM_USE_AITER=0 \ 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 new file mode 100644 index 000000000..7b6c94c62 --- /dev/null +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py @@ -0,0 +1,366 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 as hy3_module +from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec +from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import ( + _STEP_AR_KV, + _STEP_CFG_FACTOR, + _STEP_GENERATOR, + _STEP_GUIDANCE_SCALE, + _STEP_INPUT_IDS, + _STEP_MODEL_KWARGS, + _STEP_PROMPT_KV, + HunyuanImage3Pipeline, +) +from vllm_omni.diffusion.worker.input_batch import InputBatch +from vllm_omni.diffusion.worker.utils import DiffusionRequestState + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _pipeline(): + pipeline = object.__new__(HunyuanImage3Pipeline) + pipeline._tkwrapper = SimpleNamespace(pad_token_id=0) + pipeline.od_config = SimpleNamespace( + diffusion_attention_config=AttentionConfig(default=AttentionSpec(backend="TORCH_SDPA")), + parallel_config=SimpleNamespace(sequence_parallel_size=1, cfg_parallel_size=1), + cache_backend=None, + diffusion_kv_cache_skip_step_indices=None, + ) + pipeline._pipeline = SimpleNamespace() + return pipeline + + +def _state(request_id: str, step_index: int) -> DiffusionRequestState: + state = DiffusionRequestState( + request_id=request_id, + sampling=SimpleNamespace(), + prompts=["prompt"], + ) + state.step_index = step_index + state.timesteps = torch.tensor([1.0, 0.5, 0.25, 0.0]) + state.latents = torch.zeros(1, 4, 8, 8) + state.extra = { + _STEP_CFG_FACTOR: 1, + _STEP_AR_KV: None, + _STEP_INPUT_IDS: None, + _STEP_GUIDANCE_SCALE: 1.0, + _STEP_MODEL_KWARGS: { + "num_image_tokens": 17, + "ar_kv_reuse_len": 0, + }, + } + return state + + +def _sampling_params(**extra_args): + return SimpleNamespace( + timesteps=None, + sigmas=None, + num_outputs_per_prompt=None, + extra_args=extra_args, + height=512, + width=512, + num_inference_steps=4, + guidance_scale=1.0, + guidance_scale_provided=True, + guidance_rescale=0.0, + generator=None, + ) + + +def test_hunyuan_step_group_key_ignores_step_index_for_later_steps(): + pipeline = _pipeline() + states = [_state("req-0", 1), _state("req-1", 3)] + + groups = pipeline._split_step_groups(states) + + assert len(groups) == 1 + assert [state.request_id for state in groups[0]] == ["req-0", "req-1"] + + +@pytest.mark.parametrize( + ("sampling", "prompt_item", "expected_model_bot_task", "expected_system_bot_task"), + [ + pytest.param( + _sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"), + {"prompt": "prompt", "bot_task": "vanilla"}, + "think", + "think", + id="extra-args-precedence", + ), + pytest.param( + _sampling_params(use_system_prompt="dynamic"), + {"prompt": "prompt", "bot_task": "vanilla"}, + "image", + "image", + id="prompt-dict-fallback", + ), + pytest.param( + _sampling_params(use_system_prompt="dynamic"), + {"prompt": "prompt"}, + "auto", + "image", + id="default-auto-system-prompt", + ), + ], +) +def test_prepare_encode_preserves_normal_hunyuan_bot_task_semantics( + monkeypatch, + sampling, + prompt_item, + expected_model_bot_task, + expected_system_bot_task, +): + pipeline = _pipeline() + captured: dict[str, object] = {} + + def fake_get_system_prompt(sys_type, bot_task, system_prompt=None): + del sys_type, system_prompt + captured["system_prompt_bot_task"] = bot_task + return "system prompt" + + def fake_prepare_model_inputs(**kwargs): + captured.update(kwargs) + raise RuntimeError("stop after prepare_model_inputs") + + monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt) + pipeline.prepare_model_inputs = fake_prepare_model_inputs + state = DiffusionRequestState( + request_id="req-bot-task", + sampling=sampling, + prompts=[prompt_item], + ) + + with pytest.raises(RuntimeError, match="stop after prepare_model_inputs"): + pipeline.prepare_encode(state) + + assert captured["bot_task"] == expected_model_bot_task + assert captured["system_prompt_bot_task"] == expected_system_bot_task + + +def test_forward_uses_same_hunyuan_bot_task_semantics(monkeypatch): + pipeline = _pipeline() + captured: dict[str, object] = {} + + def fake_get_system_prompt(sys_type, bot_task, system_prompt=None): + del sys_type, system_prompt + captured["system_prompt_bot_task"] = bot_task + return "system prompt" + + def fake_prepare_model_inputs(**kwargs): + captured.update(kwargs) + raise RuntimeError("stop after prepare_model_inputs") + + monkeypatch.setattr(hy3_module, "get_system_prompt", fake_get_system_prompt) + pipeline.prepare_model_inputs = fake_prepare_model_inputs + req = SimpleNamespace( + request_id="req-forward-bot-task", + sampling_params=_sampling_params(bot_task="think_recaption", use_system_prompt="dynamic"), + prompts=[{"prompt": "prompt", "bot_task": "vanilla"}], + ) + + 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" + + +def test_grouped_denoise_rejects_non_sdpa_attention_backend(): + pipeline = _pipeline() + pipeline.od_config.diffusion_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")) + + with pytest.raises(ValueError, match="only supports TORCH_SDPA"): + pipeline._ensure_grouped_attention_backend_supported(2) + + +def test_single_denoise_allows_non_sdpa_attention_backend(): + pipeline = _pipeline() + pipeline.od_config.diffusion_attention_config = AttentionConfig(default=AttentionSpec(backend="FLASH_ATTN")) + + pipeline._ensure_grouped_attention_backend_supported(1) + + +def test_grouped_denoise_allows_sdpa_attention_backend(): + pipeline = _pipeline() + + pipeline._ensure_grouped_attention_backend_supported(2) + + +def test_step_scheduler_preserves_latent_dtype_for_mixed_progress_batches(): + pipeline = _pipeline() + pipeline._pipeline = SimpleNamespace(prepare_extra_func_kwargs=lambda step, kwargs: {}) + + class FakeScheduler: + def step(self, noise_pred, timestep, latents, **kwargs): + del timestep, kwargs + return (latents.float() + noise_pred.float(),) + + state = _state("req", 0) + state.timesteps = torch.tensor([1.0]) + state.scheduler = FakeScheduler() + state.latents = torch.zeros(1, 4, 8, 8, dtype=torch.bfloat16) + state.extra[_STEP_GENERATOR] = None + + pipeline.step_scheduler(state, torch.ones_like(state.latents, dtype=torch.float32)) + + assert state.latents.dtype == torch.bfloat16 + assert state.step_index == 1 + + +def test_later_step_merge_shifts_spans_without_polluting_request_state(): + pipeline = _pipeline() + states = [_state("short", 2), _state("long", 4)] + states[0].extra[_STEP_MODEL_KWARGS].update( + { + "attention_mask": torch.ones(1, 1, 3, 5, dtype=torch.bool), + "full_attn_spans": [[(2, 5)]], + } + ) + states[1].extra[_STEP_MODEL_KWARGS].update( + { + "attention_mask": torch.ones(1, 1, 3, 7, dtype=torch.bool), + "full_attn_spans": [[(4, 7)]], + } + ) + states[0].extra[_STEP_PROMPT_KV] = [{"lens": torch.tensor([2])}] + states[1].extra[_STEP_PROMPT_KV] = [{"lens": torch.tensor([4])}] + + row_state_indexes = [0, 1] + row_branches = [0, 0] + _, merged = pipeline._merge_step_model_inputs( + states, + row_state_indexes, + row_branches, + first_step=False, + ) + + assert merged["attention_mask"].shape == (2, 1, 3, 7) + assert merged["full_attn_spans"] == [[(4, 7)], [(4, 7)]] + + pipeline._split_merged_kwargs_to_states(states, merged, row_state_indexes, row_branches) + + assert states[0].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (1, 1, 3, 5) + assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (1, 1, 3, 7) + assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 5)]] + assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 7)]] + + +def test_later_step_merge_allows_request_local_step_counts_and_guidance_values(): + pipeline = _pipeline() + states = [_state("req-0", 1), _state("req-1", 3)] + for idx, state in enumerate(states): + state.extra[_STEP_MODEL_KWARGS].update( + { + "attention_mask": torch.ones(1, 1, 2, 4, dtype=torch.bool), + "full_attn_spans": [[(2, 4)]], + "guidance_scale": 3.0 + idx, + "num_inference_steps": 20 + idx, + } + ) + state.extra[_STEP_PROMPT_KV] = [{"lens": torch.tensor([2])}] + + _, merged = pipeline._merge_step_model_inputs( + states, + row_state_indexes=[0, 1], + row_branches=[0, 0], + first_step=False, + ) + + assert "guidance_scale" not in merged + assert "num_inference_steps" not in merged + + +@pytest.mark.parametrize( + ("request_id", "mutate_state", "error_match"), + [ + pytest.param( + "broken-req", + lambda state: state.extra.pop(_STEP_MODEL_KWARGS), + "broken-req", + id="missing-model-kwargs", + ), + pytest.param( + "bad-cfg", + lambda state: state.extra.__setitem__(_STEP_CFG_FACTOR, 3), + "bad-cfg", + id="unsupported-cfg-factor", + ), + ], +) +def test_denoise_step_reports_invalid_group_state_with_request_id(request_id, mutate_state, error_match): + pipeline = _pipeline() + state = _state(request_id, 0) + mutate_state(state) + + with pytest.raises(ValueError, match=error_match): + pipeline.denoise_step(InputBatch.make_batch([state])) + + +def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch): + pipeline = _pipeline() + monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu"))) + states = [_state("req-0", 1), _state("req-1", 3)] + for idx, state in enumerate(states): + prefix_len = 2 + idx * 2 + state.latents = torch.full((1, 1), float(idx)) + state.extra[_STEP_CFG_FACTOR] = 2 + state.extra[_STEP_GUIDANCE_SCALE] = 1.0 + state.extra[_STEP_INPUT_IDS] = None + state.extra[_STEP_MODEL_KWARGS].update( + { + "attention_mask": torch.ones(2, 1, 2, prefix_len + 2, dtype=torch.bool), + "full_attn_spans": [[(prefix_len, prefix_len + 2)], [(prefix_len, prefix_len + 2)]], + } + ) + state.extra[_STEP_PROMPT_KV] = [ + { + "key": torch.zeros(2, prefix_len, 1, 1), + "value": torch.zeros(2, prefix_len, 1, 1), + "lens": torch.tensor([prefix_len, prefix_len]), + } + ] + + captured = {} + + def fake_restore_prompt_kv_cache(states_arg, row_state_indexes, row_branches): + del states_arg + captured["row_state_indexes"] = list(row_state_indexes) + captured["row_branches"] = list(row_branches) + + def fake_prepare_inputs_for_generation(input_ids, images, timestep, **model_kwargs): + captured["input_ids"] = input_ids + captured["images"] = images.clone() + captured["timestep"] = timestep.clone() + captured["merged_attention_mask_shape"] = tuple(model_kwargs["attention_mask"].shape) + captured["merged_full_attn_spans"] = model_kwargs["full_attn_spans"] + return {"model_kwargs": model_kwargs} + + pipeline._restore_prompt_kv_cache = fake_restore_prompt_kv_cache + pipeline.prepare_inputs_for_generation = fake_prepare_inputs_for_generation + pipeline.forward_call = lambda **kwargs: {"diffusion_prediction": torch.tensor([[10.0], [20.0], [1.0], [2.0]])} + pipeline._update_model_kwargs_for_generation = lambda model_output, model_kwargs: model_kwargs + pipeline._pipeline = SimpleNamespace(cfg_operator=lambda cond, uncond, scale, step: cond + uncond) + + batch = InputBatch.make_batch(states) + out = pipeline.denoise_step(batch) + + assert captured["row_state_indexes"] == [0, 1, 0, 1] + assert captured["row_branches"] == [0, 0, 1, 1] + assert captured["input_ids"] is None + assert tuple(captured["images"].shape) == (4, 1) + assert captured["timestep"].tolist() == [0.5, 0.0, 0.5, 0.0] + assert captured["merged_attention_mask_shape"] == (4, 1, 2, 6) + assert captured["merged_full_attn_spans"] == [[(4, 6)], [(4, 6)], [(4, 6)], [(4, 6)]] + torch.testing.assert_close(out, torch.tensor([[11.0], [22.0]])) + assert states[0].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 4) + assert states[1].extra[_STEP_MODEL_KWARGS]["attention_mask"].shape == (2, 1, 2, 6) + assert states[0].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(2, 4)], [(2, 4)]] + assert states[1].extra[_STEP_MODEL_KWARGS]["full_attn_spans"] == [[(4, 6)], [(4, 6)]] diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index 029fea120..a21813d28 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/tests/diffusion/test_diffusion_step_pipeline.py @@ -39,7 +39,8 @@ from vllm_omni.diffusion.sched.interface import ( ) from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker -from vllm_omni.diffusion.worker.utils import RunnerOutput +from vllm_omni.diffusion.worker.input_batch import InputBatch +from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.platforms import current_omni_platform @@ -254,7 +255,9 @@ def _make_runner(): runner.cache_backend = None runner.offload_backend = None runner.state_cache = {} - runner.kv_transfer_manager = SimpleNamespace() + runner.kv_transfer_manager = SimpleNamespace( + receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None + ) return runner @@ -270,7 +273,9 @@ def _make_distributed_runner(mode: str, device: torch.device): runner.cache_backend = None runner.offload_backend = None runner.state_cache = {} - runner.kv_transfer_manager = SimpleNamespace() + runner.kv_transfer_manager = SimpleNamespace( + receive_multi_kv_cache_distributed=lambda req, cfg_kv_collect_func=None, target_device=None: None + ) return runner @@ -298,6 +303,17 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None): ) +def _make_input_batch_state(request_id: str, latent_value: float) -> DiffusionRequestState: + state = DiffusionRequestState( + request_id=request_id, + sampling=SimpleNamespace(), + prompts=None, + ) + state.latents = torch.tensor([[latent_value]]) + state.timesteps = torch.tensor([1.0]) + return state + + def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None): return DiffusionSchedulerOutput( step_id=step_id, @@ -374,6 +390,38 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master # --------------------------------------------------------------------------- +@pytest.mark.cpu +def test_input_batch_cached_repack_refreshes_state_references_without_prompt_embeds(): + first_state = _make_input_batch_state("req-1", 1.0) + batch = InputBatch.make_batch([first_state]) + assert batch.prompt_embeds is None + + replacement_state = _make_input_batch_state("req-1", 2.0) + repacked = InputBatch.make_batch([replacement_state], cached_batch=batch) + + assert repacked is batch + assert repacked.states[0] is replacement_state + assert repacked.prompt_embeds is None + torch.testing.assert_close(repacked.latents, replacement_state.latents) + + +@pytest.mark.cpu +def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_composition(): + first_state = _make_input_batch_state("req-1", 1.0) + first_state.prompt_embeds = torch.ones(1, 2, 3) + first_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool) + batch = InputBatch.make_batch([first_state]) + + replacement_state = _make_input_batch_state("req-1", 2.0) + replacement_state.prompt_embeds = torch.full((1, 2, 3), 2.0) + replacement_state.prompt_embeds_mask = torch.ones(1, 2, dtype=torch.bool) + repacked = InputBatch.make_batch([replacement_state], cached_batch=batch) + + assert repacked.states[0] is replacement_state + torch.testing.assert_close(repacked.latents, replacement_state.latents) + torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3)) + + @pytest.mark.cpu class TestRunner: """DiffusionModelRunner.execute_stepwise""" @@ -427,6 +475,36 @@ class TestRunner: result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output) assert len(result) == 2 + def test_receives_kv_payload_before_prepare_encode(self, monkeypatch): + runner = _make_runner() + captured: dict[str, object] = {} + kv_payload = object() + + class _CapturingStepPipeline(_StepPipeline): + def prepare_encode(self, state, **kwargs): + captured["past_key_values"] = getattr(state.sampling, "past_key_values", None) + return super().prepare_encode(state, **kwargs) + + class _KVTransferManager: + def receive_multi_kv_cache_distributed(self, req, cfg_kv_collect_func=None, target_device=None): + captured["cfg_kv_collect_func"] = cfg_kv_collect_func + captured["target_device"] = target_device + req.sampling_params.past_key_values = kv_payload + + runner.pipeline = _CapturingStepPipeline() + runner.pipeline.device = torch.device("cpu") + runner.od_config.cfg_kv_collect_func = "collect-cfg" + runner.kv_transfer_manager = _KVTransferManager() + monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context) + req = _make_step_request() + + DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0)) + + assert captured["past_key_values"] is kv_payload + assert captured["cfg_kv_collect_func"] == "collect-cfg" + assert captured["target_device"] == torch.device("cpu") + assert getattr(req.sampling_params, "past_key_values", None) is None + def test_rejects_missing_cached_state(self): runner = _make_runner() 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 e5fce6fce..6635168b9 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import copy import logging from collections.abc import Iterable -from typing import Any +from typing import TYPE_CHECKING, Any, ClassVar import numpy as np import torch @@ -21,6 +22,7 @@ from vllm.transformers_utils.config import get_config from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig 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 from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( @@ -30,7 +32,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer -from .hunyuan_image3_tokenizer import TokenizerWrapper +from .hunyuan_image3_tokenizer import TokenizerEncodeOutput, TokenizerWrapper from .hunyuan_image3_transformer import ( CausalMMOutputWithPast, HunyuanImage3ImageProcessor, @@ -45,14 +47,29 @@ from .hunyuan_image3_transformer import ( UNetUp, build_batch_2d_rope, real_batched_index_select, + retrieve_timesteps, ) from .system_prompt import get_system_prompt +if TYPE_CHECKING: + from vllm_omni.diffusion.worker.input_batch import InputBatch + from vllm_omni.diffusion.worker.utils import DiffusionRequestState + logger = logging.getLogger(__name__) BatchRaggedImages = torch.Tensor | list[torch.Tensor | list[torch.Tensor]] BatchRaggedTensor = torch.Tensor | list[torch.Tensor] +_STEP_MODEL_KWARGS = "hunyuan_model_kwargs" +_STEP_INPUT_IDS = "hunyuan_input_ids" +_STEP_GENERATOR = "hunyuan_generator" +_STEP_GUIDANCE_SCALE = "hunyuan_guidance_scale" +_STEP_CFG_FACTOR = "hunyuan_cfg_factor" +_STEP_OUTPUT_SIZE = "hunyuan_output_size" +_STEP_COT_TEXT_LIST = "hunyuan_cot_text_list" +_STEP_AR_KV = "hunyuan_ar_kv" +_STEP_PROMPT_KV = "hunyuan_prompt_kv" + def default(val, d): return val if val is not None else d @@ -132,6 +149,19 @@ def _to_python_scalar(value: Any) -> Any: return value +def _shift_full_attn_spans( + spans: list[tuple[int, int]], + prefix_len: int | None, + max_prefix_len: int | None, +) -> list[tuple[int, int]]: + if prefix_len is None or max_prefix_len is None: + return list(spans) + shift = max_prefix_len - prefix_len + if shift == 0: + return list(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, @@ -308,6 +338,7 @@ class HunyuanImage3Pipeline( SupportImageInput, DiffusionPipelineProfilerMixin, ): + supports_step_execution: ClassVar[bool] = True support_image_input = True hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ @@ -448,6 +479,402 @@ class HunyuanImage3Pipeline( self._pipeline = HunyuanImage3Text2ImagePipeline(model=self, scheduler=self.scheduler, vae=self.vae) return self._pipeline + def _validate_step_request(self, state: "DiffusionRequestState") -> None: + prompts = state.prompts or [] + sampling = state.sampling + if len(prompts) != 1: + raise ValueError("HunyuanImage3 step execution currently requires exactly one prompt per request.") + if sampling.timesteps is not None or sampling.sigmas is not None: + raise ValueError("HunyuanImage3 step execution does not support custom timesteps or sigmas yet.") + if sampling.num_outputs_per_prompt not in (None, 0, 1): + raise ValueError("HunyuanImage3 step execution supports only one output per prompt.") + parallel_config = self.od_config.parallel_config + if parallel_config.sequence_parallel_size > 1: + raise ValueError("HunyuanImage3 step execution does not support sequence parallel yet.") + if parallel_config.cfg_parallel_size > 1: + raise ValueError("HunyuanImage3 step execution does not support CFG parallel yet.") + cache_backend = getattr(self.od_config, "cache_backend", "none") + if cache_backend not in (None, "", "none"): + raise ValueError("HunyuanImage3 step execution does not support diffusion cache backends yet.") + 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: "DiffusionRequestState", + ) -> tuple[list[str], list[str | None], str | None, list[list[JointImageInfo]] | None, str]: + sampling = state.sampling + return self._extract_prompt_inputs( + state.prompts or [], + 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 + for layer in self.model.layers: + mgr = layer.self_attn.image_attn + injected = mgr._injected_ar_kv + if injected is None: + snapshot.append(None) + continue + found = True + snapshot.append([(k.detach().clone(), v.detach().clone()) for k, v in injected]) + mgr._injected_ar_kv = None + return snapshot if found else None + + def _restore_injected_ar_kv( + self, + states: list["DiffusionRequestState"], + row_state_indexes: list[int], + row_branches: list[int], + ) -> None: + snapshots = [state.extra.get(_STEP_AR_KV) for state in states] + if any(snapshot is None for snapshot in snapshots): + if any(snapshot is not None for snapshot in snapshots): + raise ValueError("Cannot mix Hunyuan AR-KV reuse and non-reuse requests in one step batch yet.") + for layer in self.model.layers: + layer.self_attn.image_attn._injected_ar_kv = None + return + + for layer_idx, layer in enumerate(self.model.layers): + injected_rows = [] + for state_idx, branch in zip(row_state_indexes, row_branches): + layer_snapshot = snapshots[state_idx][layer_idx] + if layer_snapshot is None: + raise ValueError("Hunyuan AR-KV snapshot is incomplete for grouped step execution.") + injected_rows.append(layer_snapshot[branch]) + layer.self_attn.image_attn._injected_ar_kv = injected_rows + + @staticmethod + def _pad_tensor_rows(rows: list[torch.Tensor], pad_value: bool | int | float = 0) -> torch.Tensor: + if not rows: + raise ValueError("Cannot merge an empty Hunyuan step batch.") + if all(tuple(row.shape[1:]) == tuple(rows[0].shape[1:]) for row in rows): + return torch.cat(rows, dim=0) + + max_shape = [max(int(row.shape[dim]) for row in rows) for dim in range(1, rows[0].ndim)] + padded = [] + for row in rows: + out = row.new_full((row.shape[0], *max_shape), pad_value) + slices = (slice(None),) + tuple(slice(0, int(size)) for size in row.shape[1:]) + out[slices] = row + padded.append(out) + return torch.cat(padded, dim=0) + + @staticmethod + def _merge_attention_masks( + rows: list[torch.Tensor], + prefix_lens: list[int] | None, + ) -> torch.Tensor: + if prefix_lens is None: + return HunyuanImage3Pipeline._pad_tensor_rows(rows, False) + + max_q = max(int(row.shape[-2]) for row in rows) + current_lens = [int(row.shape[-1]) - prefix_len for row, prefix_len in zip(rows, prefix_lens)] + max_prefix = max(prefix_lens) + max_current = max(current_lens) + padded = [] + for row, prefix_len, current_len in zip(rows, prefix_lens, current_lens): + out = row.new_zeros((row.shape[0], row.shape[1], max_q, max_prefix + max_current)) + q_len = int(row.shape[-2]) + out[:, :, :q_len, :prefix_len] = row[:, :, :q_len, :prefix_len] + out[:, :, :q_len, max_prefix : max_prefix + current_len] = row[ + :, + :, + :q_len, + prefix_len : prefix_len + current_len, + ] + padded.append(out) + return torch.cat(padded, dim=0) + + @staticmethod + def _shift_full_attn_spans( + spans: list[tuple[int, int]], + prefix_len: int | None, + max_prefix_len: int | None, + ) -> list[tuple[int, int]]: + return _shift_full_attn_spans(spans, prefix_len, max_prefix_len) + + def _row_from_value(self, value: Any, branch: int, pad_value: int = 0) -> Any: + if isinstance(value, torch.Tensor): + return value[branch : branch + 1] + if isinstance(value, TokenizerEncodeOutput): + return type(value)(**{key: self._row_from_value(value[key], branch, pad_value) for key in value.keys()}) + if isinstance(value, tuple): + return tuple(self._row_from_value(item, branch, pad_value) for item in value) + if isinstance(value, list): + return [value[branch]] + if isinstance(value, dict): + return {key: self._row_from_value(item, branch, pad_value) for key, item in value.items()} + return value + + def _merge_rows( + self, + values: list[Any], + branches: list[int], + key: str, + *, + first_step: bool, + prefix_lens: list[int] | None, + ) -> Any: + first = values[0] + if first is None: + return None + if isinstance(first, torch.Tensor): + rows = [self._row_from_value(value, branch) for value, branch in zip(values, branches)] + if key == "attention_mask": + return self._merge_attention_masks(rows, None if first_step else prefix_lens) + pad_value = self._tkwrapper.pad_token_id if key == "input_ids" else 0 + return self._pad_tensor_rows(rows, pad_value) + if isinstance(first, TokenizerEncodeOutput): + return type(first)( + **{ + output_key: self._merge_rows( + [value[output_key] for value in values], + branches, + output_key, + first_step=first_step, + prefix_lens=prefix_lens, + ) + for output_key in first.keys() + } + ) + if isinstance(first, tuple): + return tuple( + self._merge_rows( + [value[i] for value in values], + branches, + key, + first_step=first_step, + prefix_lens=prefix_lens, + ) + for i in range(len(first)) + ) + if isinstance(first, dict): + return { + sub_key: self._merge_rows( + [value[sub_key] for value in values], + branches, + sub_key, + first_step=first_step, + prefix_lens=prefix_lens, + ) + for sub_key in first + } + if isinstance(first, list): + return [value[branch] for value, branch in zip(values, branches)] + if all(value == first for value in values): + return first + raise ValueError(f"Cannot batch heterogeneous Hunyuan model kwarg {key!r}: {values}") + + def _prompt_kv_prefix_lens( + self, + states: list["DiffusionRequestState"], + row_state_indexes: list[int], + row_branches: list[int], + ) -> list[int] | None: + if not states or states[0].step_index == 0: + return None + prefix_lens: list[int] = [] + for state_idx, branch in zip(row_state_indexes, row_branches): + prompt_kv = states[state_idx].extra.get(_STEP_PROMPT_KV) + if prompt_kv is None: + raise ValueError( + f"Missing Hunyuan prompt KV cache for request {states[state_idx].request_id} " + "during later denoise step." + ) + prefix_lens.append(int(prompt_kv[0]["lens"][branch].item())) + return prefix_lens + + def _merge_step_model_inputs( + self, + states: list["DiffusionRequestState"], + row_state_indexes: list[int], + row_branches: list[int], + first_step: bool, + ) -> tuple[torch.Tensor | None, dict[str, Any]]: + prefix_lens = self._prompt_kv_prefix_lens(states, row_state_indexes, row_branches) + max_prefix_len = max(prefix_lens) if prefix_lens is not None else None + per_row_kwargs = [states[state_idx].extra[_STEP_MODEL_KWARGS] for state_idx in row_state_indexes] + merged: dict[str, Any] = {} + for key in per_row_kwargs[0]: + if key in { + "batch_gen_image_info", + "eos_token_id", + "generator", + "guidance_rescale", + "guidance_scale", + "max_new_tokens", + "num_inference_steps", + }: + continue + values = [kwargs.get(key) for kwargs in per_row_kwargs] + if key == "full_attn_spans" and values[0] is not None: + merged[key] = [ + self._shift_full_attn_spans(value[branch], prefix_len, max_prefix_len) + for value, branch, prefix_len in zip( + values, + row_branches, + prefix_lens or [None] * len(row_branches), + ) + ] + continue + merged[key] = self._merge_rows( + values, + row_branches, + key, + first_step=first_step, + prefix_lens=prefix_lens, + ) + + if "attention_mask" in merged: + b, _, q_len, seq_len = merged["attention_mask"].shape + merged["query_lens"] = [q_len] * b + merged["seq_lens"] = [seq_len] * b + + input_values = [states[state_idx].extra.get(_STEP_INPUT_IDS) for state_idx in row_state_indexes] + if input_values[0] is None: + input_ids = None + else: + rows = [value[branch : branch + 1] for value, branch in zip(input_values, row_branches)] + input_ids = self._pad_tensor_rows(rows, self._tkwrapper.pad_token_id) + return input_ids, merged + + def _restore_prompt_kv_cache( + self, + states: list["DiffusionRequestState"], + row_state_indexes: list[int], + row_branches: list[int], + ) -> None: + if states[0].step_index == 0: + for layer in self.model.layers: + mgr = layer.self_attn.image_attn + mgr.image_kv_cache_map = None + mgr.image_kv_cache_lens = None + return + + for layer_idx, layer in enumerate(self.model.layers): + rows_k, rows_v, lens = [], [], [] + for state_idx, branch in zip(row_state_indexes, row_branches): + cache = states[state_idx].extra[_STEP_PROMPT_KV][layer_idx] + rows_k.append(cache["key"][branch : branch + 1]) + rows_v.append(cache["value"][branch : branch + 1]) + lens.append(cache["lens"][branch : branch + 1]) + mgr = layer.self_attn.image_attn + mgr.image_kv_cache_map = ( + self._pad_tensor_rows(rows_k), + self._pad_tensor_rows(rows_v), + ) + mgr.image_kv_cache_lens = torch.cat(lens, dim=0).to(device=mgr.image_kv_cache_map[0].device) + + def _capture_prompt_kv_cache( + self, + states: list["DiffusionRequestState"], + row_state_indexes: list[int], + row_branches: list[int], + ) -> None: + by_state: dict[int, list[dict[str, torch.Tensor]]] = {i: [] for i in range(len(states))} + for layer in self.model.layers: + mgr = layer.self_attn.image_attn + if mgr.image_kv_cache_map is None or mgr.image_kv_cache_lens is None: + raise ValueError("Hunyuan first step did not produce prompt KV cache.") + key, value = mgr.image_kv_cache_map + lens = mgr.image_kv_cache_lens + for state_idx in by_state: + branch_rows = [ + row_idx for row_idx, (idx, _) in enumerate(zip(row_state_indexes, row_branches)) if idx == state_idx + ] + state_lens = lens[branch_rows] + max_len = int(state_lens.max().item()) + by_state[state_idx].append( + { + "key": key[branch_rows, :max_len].detach().clone(), + "value": value[branch_rows, :max_len].detach().clone(), + "lens": state_lens.detach().clone(), + } + ) + for state_idx, cache in by_state.items(): + states[state_idx].extra[_STEP_PROMPT_KV] = cache + @staticmethod def get_pos_emb(custom_pos_emb, position_ids): cos, sin = custom_pos_emb @@ -1387,6 +1814,426 @@ class HunyuanImage3Pipeline( ) return {"ar_kv_data": ar_kv_data} + def _extract_ar_kv_from_sampling(self, sampling: Any) -> dict[str, Any]: + kv = getattr(sampling, "past_key_values", None) + if kv is None: + return {} + key_cache = getattr(kv, "key_cache", None) + value_cache = getattr(kv, "value_cache", None) + if not key_cache or not value_cache: + return {} + ar_kv_data = { + i: {"key": k, "value": v} + for i, (k, v) in enumerate(zip(key_cache, value_cache)) + if k is not None and v is not None + } + return {"ar_kv_data": ar_kv_data} if ar_kv_data else {} + + def prepare_encode( + self, + state: "DiffusionRequestState", + **kwargs: Any, + ) -> "DiffusionRequestState": + del kwargs + self._validate_step_request(state) + pipe = self.pipeline + sampling = state.sampling + ( + prompt, + cot_text_list, + system_prompt, + batch_cond_image_info, + tokenizer_bot_task, + ) = self._extract_step_prompt_inputs(state) + cot_text = ( + [self._normalize_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 + 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 + if guidance_scale <= 1.0: + logger.info("HunyuanImage3.0 step execution runs without classifier-free guidance.") + pipe._guidance_scale = guidance_scale + pipe._guidance_rescale = getattr(sampling, "guidance_rescale", 0.0) + + model_kwargs = self.prepare_model_inputs( + prompt=prompt, + cot_text=cot_text, + system_prompt=system_prompt, + mode="gen_image", + generator=sampling.generator, + image_size=image_size, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + batch_cond_image_info=batch_cond_image_info, + bot_task=tokenizer_bot_task, + ) + model_kwargs.update(self._extract_ar_kv_from_sampling(sampling)) + model_kwargs["use_cache"] = False + + 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 + + timesteps, _ = retrieve_timesteps( + self.scheduler, + num_inference_steps, + self.device, + None, + None, + ) + pipe._num_timesteps = len(timesteps) + req_scheduler = copy.deepcopy(self.scheduler) + if hasattr(req_scheduler, "set_begin_index"): + req_scheduler.set_begin_index(0) + + latents = pipe.prepare_latents( + batch_size=1, + latent_channel=self.config.vae["latent_channels"], + image_size=[target_height, target_width], + dtype=torch.bfloat16, + device=self.device, + generator=model_kwargs["generator"], + ) + + attention_mask = self._prepare_attention_mask_for_generation( + input_ids, + self.generation_config, + model_kwargs=model_kwargs, + ) + bsz, _, query_len, seq_len = attention_mask.shape + model_kwargs["query_lens"] = [query_len] * bsz + model_kwargs["seq_lens"] = [seq_len] * bsz + model_kwargs["attention_mask"] = attention_mask.to(latents.device) + + pipe._guidance_scale = guidance_scale + pipe._guidance_rescale = 0.0 + input_ids, ar_kv_reuse_len = pipe._maybe_handle_ar_kv_reuse( + input_ids, + model_kwargs, + batch_size=1, + cfg_parallel_ready=False, + cfg_rank=None, + device=self.device, + ) + model_kwargs["ar_kv_reuse_len"] = ar_kv_reuse_len + + state.latents = latents + state.timesteps = timesteps + state.step_index = 0 + state.scheduler = req_scheduler + state.do_true_cfg = guidance_scale > 1.0 + state.extra = { + _STEP_MODEL_KWARGS: model_kwargs, + _STEP_INPUT_IDS: input_ids, + _STEP_GENERATOR: model_kwargs["generator"], + _STEP_GUIDANCE_SCALE: guidance_scale, + _STEP_CFG_FACTOR: 1 + int(guidance_scale > 1.0), + _STEP_OUTPUT_SIZE: (target_height, target_width), + _STEP_COT_TEXT_LIST: cot_text_list, + _STEP_AR_KV: self._snapshot_injected_ar_kv(), + } + return state + + def _step_group_key(self, state: "DiffusionRequestState") -> tuple[Any, ...]: + if _STEP_CFG_FACTOR not in state.extra: + raise ValueError(f"Missing Hunyuan CFG factor for request {state.request_id}.") + if _STEP_MODEL_KWARGS not in state.extra: + raise ValueError(f"Missing Hunyuan step model kwargs for request {state.request_id}.") + if state.latents is None: + raise ValueError(f"Missing Hunyuan latents for request {state.request_id}.") + model_kwargs = state.extra[_STEP_MODEL_KWARGS] + return ( + state.step_index == 0, + state.extra[_STEP_CFG_FACTOR], + tuple(state.latents.shape[1:]), + model_kwargs.get("num_image_tokens"), + model_kwargs.get("ar_kv_reuse_len", 0), + state.extra.get(_STEP_AR_KV) is not None, + ) + + def _ensure_grouped_attention_backend_supported(self, num_states: int) -> None: + if num_states <= 1: + return + spec, source = self.od_config.diffusion_attention_config.resolve_with_source(role="self") + backend = spec.backend.upper() if spec is not None else "AUTO" + if backend == "TORCH_SDPA": + return + source_msg = f" from {source}" if source is not None else "" + raise ValueError( + "HunyuanImage3 grouped DiT batching only supports TORCH_SDPA attention. " + f"Current diffusion attention backend is {backend}{source_msg}. " + "Set DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA or configure " + "diffusion_attention_config.default.backend=TORCH_SDPA when max_num_seqs > 1." + ) + + def _split_step_groups( + self, + states: list["DiffusionRequestState"], + ) -> list[list["DiffusionRequestState"]]: + grouped: dict[tuple[Any, ...], list[DiffusionRequestState]] = {} + for state in states: + grouped.setdefault(self._step_group_key(state), []).append(state) + return list(grouped.values()) + + @staticmethod + def _step_row_order( + states: list["DiffusionRequestState"], + cfg_factor: int, + ) -> tuple[list[int], list[int]]: + row_state_indexes = [state_idx for _ in range(cfg_factor) for state_idx in range(len(states))] + row_branches = [branch for branch in range(cfg_factor) for _ in states] + return row_state_indexes, row_branches + + @staticmethod + def _validate_step_group_states(states: list["DiffusionRequestState"]) -> tuple[bool, int]: + if not states: + raise ValueError("HunyuanImage3 denoise_step received an empty group.") + + first_step = states[0].step_index == 0 + if _STEP_CFG_FACTOR not in states[0].extra: + raise ValueError(f"Missing Hunyuan CFG factor for request {states[0].request_id}.") + cfg_factor = int(states[0].extra[_STEP_CFG_FACTOR]) + if cfg_factor not in (1, 2): + raise ValueError( + f"Unsupported Hunyuan CFG factor {cfg_factor} for request {states[0].request_id}; " + "expected no-CFG or cond/uncond CFG rows." + ) + + for state in states: + if (state.step_index == 0) != first_step: + raise ValueError("Hunyuan step group mixed first-step and later-step requests.") + if _STEP_CFG_FACTOR not in state.extra: + raise ValueError(f"Missing Hunyuan CFG factor for request {state.request_id}.") + if int(state.extra[_STEP_CFG_FACTOR]) != cfg_factor: + raise ValueError("Hunyuan step group mixed requests with different CFG factors.") + if _STEP_MODEL_KWARGS not in state.extra: + raise ValueError(f"Missing Hunyuan step model kwargs for request {state.request_id}.") + if state.latents is None: + raise ValueError(f"Missing Hunyuan latents for request {state.request_id}.") + if first_step: + if _STEP_AR_KV not in state.extra: + raise ValueError(f"Missing Hunyuan AR KV snapshot for request {state.request_id}.") + if _STEP_INPUT_IDS not in state.extra: + raise ValueError(f"Missing Hunyuan step input ids for request {state.request_id}.") + elif _STEP_PROMPT_KV not in state.extra: + raise ValueError(f"Missing Hunyuan prompt KV cache for request {state.request_id}.") + + return first_step, cfg_factor + + def _split_merged_kwargs_to_states( + self, + states: list["DiffusionRequestState"], + merged_kwargs: dict[str, Any], + row_state_indexes: list[int], + row_branches: list[int], + ) -> None: + cfg_factor = states[0].extra[_STEP_CFG_FACTOR] + for state_idx, state in enumerate(states): + state_rows = [ + row_idx for row_idx, (idx, _) in enumerate(zip(row_state_indexes, row_branches)) if idx == state_idx + ] + next_kwargs: dict[str, Any] = {} + for key, value in merged_kwargs.items(): + if key in {"query_lens", "seq_lens"}: + continue + if key == "full_attn_spans": + next_kwargs[key] = state.extra[_STEP_MODEL_KWARGS].get(key) + elif key == "attention_mask" and isinstance(value, torch.Tensor): + cache = state.extra.get(_STEP_PROMPT_KV) + if cache is None: + next_kwargs[key] = value[state_rows] + else: + compact_masks = [] + for row_idx in state_rows: + branch = row_branches[row_idx] + row_mask = value[row_idx : row_idx + 1] + query_len = int(row_mask.shape[-2]) + prefix_len = int(cache[0]["lens"][branch].item()) + max_prefix_len = int(row_mask.shape[-1]) - query_len + prefix_mask = row_mask[:, :, :, :prefix_len] + current_mask = row_mask[:, :, :, max_prefix_len : max_prefix_len + query_len] + compact_masks.append(torch.cat([prefix_mask, current_mask], dim=-1)) + next_kwargs[key] = torch.cat(compact_masks, dim=0) + elif isinstance(value, torch.Tensor): + next_kwargs[key] = value[state_rows] + elif isinstance(value, tuple): + next_kwargs[key] = tuple( + item[state_rows] if isinstance(item, torch.Tensor) else item for item in value + ) + elif isinstance(value, list) and len(value) == len(row_state_indexes): + next_kwargs[key] = [value[row_idx] for row_idx in state_rows] + else: + next_kwargs[key] = value + if len(state_rows) != cfg_factor: + raise ValueError("Hunyuan step split lost CFG rows while updating request state.") + if "attention_mask" in next_kwargs: + bsz, _, query_len, seq_len = next_kwargs["attention_mask"].shape + next_kwargs["query_lens"] = [query_len] * bsz + next_kwargs["seq_lens"] = [seq_len] * bsz + state.extra[_STEP_MODEL_KWARGS] = next_kwargs + state.extra[_STEP_INPUT_IDS] = None + + def _denoise_step_group(self, states: list["DiffusionRequestState"]) -> torch.Tensor: + first_step, cfg_factor = self._validate_step_group_states(states) + row_state_indexes, row_branches = self._step_row_order(states, cfg_factor) + latents = torch.cat([state.latents for state in states], dim=0) + latent_model_input = torch.cat([latents] * cfg_factor, dim=0) + timestep = torch.cat( + [ + state.current_timestep.reshape(1).to(device=latents.device) + for _ in range(cfg_factor) + for state in states + ], + dim=0, + ) + + input_ids, model_kwargs = self._merge_step_model_inputs( + states, + row_state_indexes, + row_branches, + first_step, + ) + if first_step: + self._restore_injected_ar_kv(states, row_state_indexes, row_branches) + else: + self._restore_prompt_kv_cache(states, row_state_indexes, row_branches) + + model_inputs = self.prepare_inputs_for_generation( + input_ids, + images=latent_model_input, + timestep=timestep, + **model_kwargs, + ) + + step_indexes = {state.step_index for state in states} + context_step_index = next(iter(step_indexes)) if len(step_indexes) == 1 else None + set_forward_context_denoise_step_idx(context_step_index) + try: + with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True): + model_output = self.forward_call(**model_inputs, first_step=first_step) + pred = model_output["diffusion_prediction"] + finally: + set_forward_context_denoise_step_idx(None) + pred = pred.to(dtype=torch.float32) + + if first_step: + self._capture_prompt_kv_cache(states, row_state_indexes, row_branches) + if cfg_factor > 1: + pred_cond, pred_uncond = pred.chunk(2) + pred = torch.cat( + [ + self.pipeline.cfg_operator( + pred_cond[state_idx : state_idx + 1], + pred_uncond[state_idx : state_idx + 1], + state.extra[_STEP_GUIDANCE_SCALE], + step=state.step_index, + ) + for state_idx, state in enumerate(states) + ], + dim=0, + ) + + updated_kwargs = self._update_model_kwargs_for_generation(model_output, model_kwargs) + self._split_merged_kwargs_to_states(states, updated_kwargs, row_state_indexes, row_branches) + return pred + + def denoise_step( + self, + input_batch: "InputBatch", + **kwargs: Any, + ) -> torch.Tensor | None: + del kwargs + if getattr(self, "interrupt", False): + return None + states = list(input_batch.states) + if not states: + raise ValueError("HunyuanImage3 denoise_step received an empty batch.") + self._ensure_grouped_attention_backend_supported(len(states)) + outputs: dict[str, torch.Tensor] = {} + for group in self._split_step_groups(states): + pred = self._denoise_step_group(group) + for state, state_pred in zip(group, pred): + outputs[state.request_id] = state_pred.unsqueeze(0) + return torch.cat([outputs[state.request_id] for state in states], dim=0) + + def step_scheduler( + self, + state: "DiffusionRequestState", + noise_pred: torch.Tensor, + **kwargs: Any, + ) -> None: + del kwargs + if getattr(self, "interrupt", False): + return + generator = state.extra.get(_STEP_GENERATOR) + step_kwargs = self.pipeline.prepare_extra_func_kwargs(state.scheduler.step, {"generator": generator}) + latent_dtype = state.latents.dtype + state.latents = state.scheduler.step( + noise_pred, + state.current_timestep, + state.latents, + **step_kwargs, + return_dict=False, + )[0].to(dtype=latent_dtype) + state.step_index += 1 + + def post_decode( + self, + state: "DiffusionRequestState", + **kwargs: Any, + ) -> DiffusionOutput: + output_type = kwargs.get("output_type", "pil") + generator = state.extra.get(_STEP_GENERATOR) + latents = state.latents + if output_type == "latent": + return DiffusionOutput( + output=latents, + custom_output={}, + stage_durations=getattr(self, "stage_durations", None), + ) + + if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor: + latents = latents / self.vae.config.scaling_factor + if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor: + latents = latents + self.vae.config.shift_factor + if hasattr(self.vae, "ffactor_temporal"): + latents = latents.unsqueeze(2) + + with torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=True): + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] + if hasattr(self.vae, "ffactor_temporal"): + assert image.shape[2] == 1, "image should have shape [B, C, T, H, W] and T should be 1" + image = image.squeeze(2) + image = self.pipeline.image_processor.postprocess( + image, + output_type=output_type, + do_denormalize=[True] * image.shape[0], + ) + + cot_text_list = state.extra.get(_STEP_COT_TEXT_LIST) or [] + custom_output = {} + if any(text is not None for text in cot_text_list): + custom_output["ar_generated_text"] = cot_text_list[0] + return DiffusionOutput( + output=image[0], + custom_output=custom_output, + stage_durations=getattr(self, "stage_durations", None), + ) + def forward( self, req: OmniDiffusionRequest, @@ -1400,56 +2247,23 @@ class HunyuanImage3Pipeline( **kwargs, ) -> DiffusionOutput: extra_args = getattr(getattr(req, "sampling_params", None), "extra_args", {}) or {} - bot_task = extra_args.get("bot_task") - use_system_prompt = extra_args.get("use_system_prompt") - system_prompt = extra_args.get("system_prompt") - # Fall back to per-prompt prompt kwargs forwarded by ar2diffusion - if req.prompts: - first_prompt = req.prompts[0] - 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") - if isinstance(bot_task, str) and bot_task.lower() == "none": - bot_task = None - if use_system_prompt is not None: - system_prompt = get_system_prompt(use_system_prompt, bot_task or "image", 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 req.prompts] or prompt - - cot_text_list = [ - (p.get("extra", {}).get("ar_generated_text") if isinstance(p, dict) else None) or None for p in req.prompts - ] + ( + prompt_from_req, + cot_text_list, + system_prompt, + batch_cond_image_info, + tokenizer_bot_task, + ) = self._extract_prompt_inputs( + req.prompts, + extra_args, + request_id=req.request_id, + allow_cond_image=True, + ) + 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 ) - batch_cond_image_info: list[list[JointImageInfo]] | None = None - if any(not isinstance(p, str) for p in req.prompts): - batch_cond_image_info = [] - for prompt_item in req.prompts: - if isinstance(prompt_item, str): - batch_cond_image_info.append([]) - continue - prompt_additional_information = prompt_item.get("additional_information") or {} - prompt_cond_infos = prompt_additional_information.get("batch_cond_image_info", []) - if isinstance(prompt_cond_infos, JointImageInfo | dict): - prompt_cond_infos = [prompt_cond_infos] - if prompt_cond_infos is None: - prompt_cond_infos = [] - batch_cond_image_info.append([_joint_image_info_from_payload(item) for item in prompt_cond_infos]) - - has_cond_image = [len(cond_infos) > 0 for cond_infos in batch_cond_image_info] - if 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 any(has_cond_image): - batch_cond_image_info = None - generator = req.sampling_params.generator or generator height = req.sampling_params.height or height width = req.sampling_params.width or width @@ -1463,18 +2277,6 @@ class HunyuanImage3Pipeline( # ---- AR KV Reuse: extract injected KV from request ---- ar_kv_kwargs = self._extract_ar_kv_from_request(req) - 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)}." - ) - model_inputs = self.prepare_model_inputs( prompt=prompt, cot_text=cot_text, @@ -1485,7 +2287,7 @@ class HunyuanImage3Pipeline( num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, batch_cond_image_info=batch_cond_image_info, - bot_task=tokenizer_bot_task or "auto", + bot_task=tokenizer_bot_task, **ar_kv_kwargs, ) diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index d4ace761b..c9d013ff3 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -16,6 +16,7 @@ if TYPE_CHECKING: import torch from vllm_omni.diffusion.data import DiffusionOutput + from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import DiffusionRequestState @@ -56,8 +57,8 @@ class SupportsStepExecution(Protocol): def prepare_encode(self, state: DiffusionRequestState, **kwargs: Any) -> DiffusionRequestState: """Prepare request-level inputs and return initialized state.""" - def denoise_step(self, state: DiffusionRequestState, **kwargs: Any) -> torch.Tensor | None: - """Run one denoise step.""" + def denoise_step(self, input_batch: InputBatch, **kwargs: Any) -> torch.Tensor | None: + """Run one denoise forward on the runner-assembled batch.""" def step_scheduler(self, state: DiffusionRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: """Run one scheduler step.""" diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 55e3dca1f..278f1df40 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -374,6 +374,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): sampling=copy.deepcopy(req.sampling_params), prompts=req.prompts, ) + state_req = copy.copy(req) + state_req.sampling_params = new_state.sampling + self.kv_transfer_manager.receive_multi_kv_cache_distributed( + state_req, + cfg_kv_collect_func=getattr(self.od_config, "cfg_kv_collect_func", None), + target_device=getattr(self.pipeline, "device", None), + ) self.state_cache[request_id] = new_state resolved.append(new_state) @@ -452,8 +459,8 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): if not self.supports_step_mode(): raise ValueError("Current pipeline does not support step execution.") # Stepwise mode only supports the basic state-driven denoise path for now. - # Request-mode extras such as cache backends, KV transfer, editing inputs, - # and similar features are not supported here yet. + # Request-mode extras such as cache backends, editing inputs, and + # similar features are not supported here yet. if self.od_config.cache_backend not in (None, "none"): raise ValueError("Step mode does not support cache_backend yet.") diff --git a/vllm_omni/diffusion/worker/input_batch.py b/vllm_omni/diffusion/worker/input_batch.py index 82033d9dc..ef502d0d4 100644 --- a/vllm_omni/diffusion/worker/input_batch.py +++ b/vllm_omni/diffusion/worker/input_batch.py @@ -12,7 +12,7 @@ step, and step outputs are scattered back into request states by from __future__ import annotations from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field import numpy as np import torch @@ -496,7 +496,7 @@ def _prepare_prompt_embeds( *, embeds_out: torch.Tensor | None = None, mask_out: torch.Tensor | None = None, -) -> tuple[torch.Tensor, torch.Tensor | None]: +) -> tuple[torch.Tensor | None, torch.Tensor | None]: prompt_embeds, prompt_embeds_mask = _prepare_padded_prompt_fields( states, embeds_attr="prompt_embeds", @@ -505,8 +505,6 @@ def _prepare_prompt_embeds( embeds_out=embeds_out, mask_out=mask_out, ) - if prompt_embeds is None: - raise ValueError("All requests must have `prompt_embeds` initialized.") return prompt_embeds, prompt_embeds_mask @@ -581,6 +579,10 @@ class InputBatch: source of truth. ``InputBatch`` only assembles a contiguous view for the current step and refreshes dynamic fields in-place when composition is unchanged. + + ``states`` is a narrow escape hatch for Hunyuan-style state-driven + pipelines that need request-private KV/cache metadata during denoise. + Other pipelines should continue to use the standard batch fields. """ request_ids: list[str] @@ -591,7 +593,7 @@ class InputBatch: latents: torch.Tensor timesteps: torch.Tensor - prompt_embeds: torch.Tensor + prompt_embeds: torch.Tensor | None prompt_embeds_mask: torch.Tensor | None negative_prompt_embeds: torch.Tensor | None negative_prompt_embeds_mask: torch.Tensor | None @@ -604,6 +606,7 @@ class InputBatch: img_shapes: list | None = None txt_seq_lens: list[int] | None = None negative_txt_seq_lens: list[int] | None = None + states: Sequence[DiffusionRequestState] = field(default_factory=tuple) def __post_init__(self) -> None: if len(self.request_ids) != int(self.idx_mapping.numel()): @@ -648,6 +651,9 @@ class InputBatch: self, selected_states: Sequence[DiffusionRequestState], ) -> None: + # Same-composition cache hits reuse static prompt fields; request ids + # must keep the same encoded prompt metadata for the batch lifetime. + self.states = tuple(selected_states) self._refresh_dynamic_fields(selected_states) def _rebuild( @@ -662,6 +668,7 @@ class InputBatch: self.num_reqs_after_padding = len(request_ids) self.idx_mapping = idx_mapping self.idx_mapping_np = idx_mapping_np + self.states = tuple(selected_states) self.latents = _prepare_latents(selected_states, out=self.latents) self.timesteps = _prepare_timesteps(selected_states, out=self.timesteps) self._refresh_static_fields(selected_states) @@ -718,6 +725,7 @@ class InputBatch: selected_states, "negative_txt_seq_lens", ), + states=tuple(selected_states), )