[Feat] Add support and testcase for HunyuanImage-3.0-Instruct-Distil (#4048)

Signed-off-by: zengchuang <zengchuang3@huawei.com>
Signed-off-by: zengchuang-hw <zengchuang1005@gmail.com>

tests/diffusion/models/hunyuan_image3/test_diffusion_kv_request.py

@@ -24,6 +24,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.request_layout import (
     build_hunyuan_diffusion_kv_requests,
     extract_hunyuan_prompt_inputs,
     hunyuan_num_image_tokens,
+    hunyuan_num_special_tokens,
     normalize_hunyuan_cot_text,
     prepare_hunyuan_layout,
 )
@@ -114,7 +115,7 @@ class _FakeImageProcessor:
         self.image_sizes: list[tuple[int, int]] = []
         self.vision_encoder_processor = SimpleNamespace(patch_size=1)
 
-    def build_image_info(self, image_size):
+    def build_image_info(self, image_size, **kwargs):
         self.image_sizes.append(image_size)
         return ImageInfo(
             image_type="gen_image",
@@ -125,6 +126,7 @@ class _FakeImageProcessor:
             image_token_length=self.image_token_length,
             base_size=1024,
             ratio_index=0,
+            **kwargs,
         )
 
 
@@ -166,6 +168,29 @@ def _request(
     )
 
 
+def test_distilled_layout_uses_embedded_cfg_and_meanflow_tokens() -> None:
+    tokenizer, image_processor = _components([12])
+    request = _request(guidance_scale=2.5)
+
+    prepared_layout = prepare_hunyuan_layout(
+        request,
+        tokenizer_wrapper=tokenizer,
+        image_processor=image_processor,
+        generation_config=SimpleNamespace(sequence_template="instruct", drop_think=False),
+        image_base_size=1024,
+        cfg_distilled=True,
+        use_meanflow=True,
+    )
+
+    image_info = prepared_layout.generated_image_info
+    assert image_info.add_guidance_token
+    assert image_info.add_timestep_r_token
+    assert tokenizer.calls[0]["cfg_factor"] == 1
+    assert hunyuan_num_special_tokens(image_info) == 3
+    assert hunyuan_num_image_tokens(image_info) == 19
+    assert len(build_hunyuan_diffusion_kv_requests(request, prepared_layout)) == 1
+
+
 def test_builds_kv_request_lengths_without_model_execution() -> None:
     tokenizer, image_processor = _components([12])
     request = _request(guidance_scale=1.0)
@@ -385,9 +410,13 @@ def test_prepared_model_inputs_match_local_tokenization(
     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)
+    image_info = prepared_layout.generated_image_info
+    generation_token_counts = {
+        "num_image_tokens": hunyuan_num_image_tokens(image_info),
+        "num_special_tokens": hunyuan_num_special_tokens(image_info),
+    }
+    local_inputs.update(attention_mask=local_mask, **generation_token_counts)
+    prepared_inputs.update(attention_mask=prepared_mask, **generation_token_counts)
     local_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), local_inputs)
     prepared_step_inputs = pipeline._update_model_kwargs_for_generation(ModelOutput(), prepared_inputs)
 

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_distil_inputs.py

@@ -0,0 +1,294 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HunyuanImage3 prepare_inputs_for_generation to ensure
+distilled model parameters (guidance, timesteps_r) are properly passed."""
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
+    HunyuanImage3Text2ImagePipeline,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class MockHunyuanImage3Pipeline:
+    """Minimal mock that replicates prepare_inputs_for_generation method."""
+
+    def __init__(self):
+        # Mock config for distilled models
+        self.config = type("Config", (), {})()
+        self.hf_config = type("HFConfig", (), {"cfg_distilled": True, "use_meanflow": True})()
+
+    # Bind the real method from the pipeline class
+    from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import (
+        HunyuanImage3Pipeline as _Real,
+    )
+
+    prepare_inputs_for_generation = _Real.prepare_inputs_for_generation
+    _update_model_kwargs_for_generation = _Real._update_model_kwargs_for_generation
+
+
+class TestPrepareInputsForGenerationDistilled:
+    """Tests that distilled model parameters are passed through
+    prepare_inputs_for_generation."""
+
+    def setup_method(self):
+        self.pipeline = MockHunyuanImage3Pipeline()
+
+    def _create_mock_kwargs(self, include_distilled_params: bool = True) -> dict:
+        """Create mock kwargs for prepare_inputs_for_generation."""
+        kwargs = {
+            "position_ids": torch.tensor([[0, 1, 2]]),
+            "custom_pos_emb": (torch.tensor([[0.0]]), torch.tensor([[1.0]])),
+            "mode": "gen_image",
+            "images": torch.randn(1, 4, 64, 64),
+            "image_mask": torch.ones(1, 64, 64),
+            "timestep": torch.tensor([0.5]),
+            "gen_timestep_scatter_index": torch.tensor([[5]]),
+            "cond_vae_images": None,
+            "cond_timestep": None,
+            "cond_vae_image_mask": None,
+            "cond_vit_images": None,
+            "cond_vit_image_mask": None,
+            "vit_kwargs": None,
+            "cond_timestep_scatter_index": None,
+            "query_lens": [10],
+            "seq_lens": [100],
+            "num_image_tokens": 4096,
+            "num_special_tokens": 3,
+            "ar_kv_reuse_len": 0,
+            "full_attn_spans": None,
+            "use_cache": True,
+        }
+        if include_distilled_params:
+            # CFG distilled model parameters
+            kwargs["guidance"] = torch.tensor([2500.0])
+            kwargs["guidance_scatter_index"] = torch.tensor([[4]])
+            # MeanFlow distilled model parameters
+            kwargs["timesteps_r"] = torch.tensor([0.3])
+            kwargs["timesteps_r_scatter_index"] = torch.tensor([[3]])
+        else:
+            kwargs["guidance"] = None
+            kwargs["guidance_scatter_index"] = None
+            kwargs["timesteps_r"] = None
+            kwargs["timesteps_r_scatter_index"] = None
+        return kwargs
+
+    def test_guidance_passed_to_model_inputs(self):
+        """Test that guidance is passed through prepare_inputs_for_generation."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "guidance" in model_inputs, "guidance should be in model_inputs"
+        assert model_inputs["guidance"] is not None, "guidance should not be None"
+        assert torch.equal(model_inputs["guidance"], kwargs["guidance"])
+
+    def test_guidance_scatter_index_passed_to_model_inputs(self):
+        """Test that guidance_scatter_index is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "guidance_scatter_index" in model_inputs
+        assert model_inputs["guidance_scatter_index"] is not None
+        assert torch.equal(model_inputs["guidance_scatter_index"], kwargs["guidance_scatter_index"])
+
+    def test_timesteps_r_passed_to_model_inputs(self):
+        """Test that timesteps_r is passed through for MeanFlow models."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "timesteps_r" in model_inputs, "timesteps_r should be in model_inputs"
+        assert model_inputs["timesteps_r"] is not None
+        assert torch.equal(model_inputs["timesteps_r"], kwargs["timesteps_r"])
+
+    def test_timesteps_r_scatter_index_passed_to_model_inputs(self):
+        """Test that timesteps_r_scatter_index is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "timesteps_r_scatter_index" in model_inputs
+        assert model_inputs["timesteps_r_scatter_index"] is not None
+        assert torch.equal(model_inputs["timesteps_r_scatter_index"], kwargs["timesteps_r_scatter_index"])
+
+    def test_num_special_tokens_passed_to_model_inputs(self):
+        """Test that num_special_tokens is passed through."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        assert "num_special_tokens" in model_inputs
+        assert model_inputs["num_special_tokens"] == kwargs["num_special_tokens"]
+
+    def test_decode_update_keeps_all_special_token_positions(self):
+        """Test decode step uses timestep, guidance, and timestep_r positions."""
+        model_kwargs = self._create_mock_kwargs(include_distilled_params=True)
+        model_kwargs.update(
+            {
+                "tokenizer_output": object(),
+                "attention_mask": torch.ones(1, 1, 7, 7, dtype=torch.bool),
+                "image_mask": torch.tensor([[False, False, False, False, True, True, True]]),
+                "gen_timestep_scatter_index": torch.tensor([[1]]),
+                "guidance_scatter_index": torch.tensor([[2]]),
+                "timesteps_r_scatter_index": torch.tensor([[3]]),
+            }
+        )
+
+        updated = self.pipeline._update_model_kwargs_for_generation({}, model_kwargs)
+
+        assert torch.equal(updated["position_ids"], torch.tensor([[1, 2, 3, 4, 5, 6]]))
+        assert updated["num_special_tokens"] == 3
+        assert updated["guidance_scatter_index"] is model_kwargs["guidance_scatter_index"]
+        assert updated["timesteps_r_scatter_index"] is model_kwargs["timesteps_r_scatter_index"]
+        assert updated["attention_mask"].shape == (1, 1, 6, 7)
+
+    def test_distilled_params_none_when_not_provided(self):
+        """Test that distilled params are None when not provided."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=False)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        # All distilled params should be present but None
+        assert "guidance" in model_inputs
+        assert model_inputs["guidance"] is None
+        assert "guidance_scatter_index" in model_inputs
+        assert model_inputs["guidance_scatter_index"] is None
+        assert "timesteps_r" in model_inputs
+        assert model_inputs["timesteps_r"] is None
+        assert "timesteps_r_scatter_index" in model_inputs
+        assert model_inputs["timesteps_r_scatter_index"] is None
+
+    def test_all_distilled_keys_present_in_model_inputs(self):
+        """Test that all 4 distilled keys are always present (even if None)."""
+        kwargs = self._create_mock_kwargs(include_distilled_params=False)
+        input_ids = torch.tensor([[1, 2, 3]])
+        attention_mask = torch.ones(1, 3)
+
+        model_inputs = self.pipeline.prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=None,
+            attention_mask=attention_mask,
+            inputs_embeds=None,
+            tokenizer_output=None,
+            batch_gen_image_info=None,
+            generator=None,
+            **kwargs,
+        )
+
+        required_distilled_keys = [
+            "guidance",
+            "guidance_scatter_index",
+            "timesteps_r",
+            "timesteps_r_scatter_index",
+        ]
+        for key in required_distilled_keys:
+            assert key in model_inputs, f"{key} must be in model_inputs dict"
+
+
+def test_ar_kv_reuse_rebases_distilled_scatter_indices():
+    pipeline = object.__new__(HunyuanImage3Text2ImagePipeline)
+    model_kwargs = {
+        "query_lens": [8],
+        "attention_mask": torch.ones(1, 1, 8, 12, dtype=torch.bool),
+        "position_ids": torch.arange(8).unsqueeze(0),
+        "image_mask": torch.ones(1, 8, dtype=torch.bool),
+        "gen_timestep_scatter_index": torch.tensor([[5]]),
+        "guidance_scatter_index": torch.tensor([[6]]),
+        "timesteps_r_scatter_index": torch.tensor([[7]]),
+    }
+
+    truncated = pipeline._truncate_reused_prefix(
+        input_ids=torch.arange(8).unsqueeze(0),
+        model_kwargs=model_kwargs,
+        positive_reuse_len=3,
+    )
+
+    assert truncated.tolist() == [[3, 4, 5, 6, 7]]
+    assert model_kwargs["gen_timestep_scatter_index"].tolist() == [[2]]
+    assert model_kwargs["guidance_scatter_index"].tolist() == [[3]]
+    assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[4]]
+
+
+def test_cfg_parallel_splits_meanflow_scatter_indices():
+    model_kwargs = {
+        "timesteps_r_scatter_index": torch.tensor([[1], [2], [3], [4]]),
+    }
+
+    HunyuanImage3Text2ImagePipeline._split_model_kwargs_for_cfg_parallel(
+        model_kwargs,
+        batch_size=2,
+        cfg_rank=1,
+    )
+
+    assert model_kwargs["timesteps_r_scatter_index"].tolist() == [[3], [4]]

tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_step_execution.py

@@ -38,6 +38,7 @@ def _pipeline():
         cache_backend=None,
         diffusion_kv_cache_skip_step_indices=None,
     )
+    pipeline.hf_config = SimpleNamespace(cfg_distilled=False, use_meanflow=False)
     pipeline._pipeline = SimpleNamespace()
     return pipeline
 
@@ -451,3 +452,44 @@ def test_denoise_step_uses_input_batch_group_order_and_splits_back(monkeypatch):
     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)]]
+
+
+def test_distilled_step_supplies_guidance_and_meanflow_timestep(monkeypatch):
+    pipeline = _pipeline()
+    pipeline.hf_config = SimpleNamespace(cfg_distilled=True, use_meanflow=True)
+    monkeypatch.setattr(HunyuanImage3Pipeline, "device", property(lambda self: torch.device("cpu")))
+    state = _state("distilled", 1)
+    state.latents = torch.zeros(1, 1)
+    state.extra[_STEP_GUIDANCE_SCALE] = 2.5
+    state.extra[_STEP_MODEL_KWARGS].update(
+        {
+            "attention_mask": torch.ones(1, 1, 2, 4, dtype=torch.bool),
+            "full_attn_spans": [[(2, 4)]],
+        }
+    )
+    state.extra[_STEP_PROMPT_KV] = [
+        {
+            "key": torch.zeros(1, 2, 1, 1),
+            "value": torch.zeros(1, 2, 1, 1),
+            "lens": torch.tensor([2]),
+        }
+    ]
+    state.scheduler = SimpleNamespace(get_timestep_r=lambda _timestep: torch.tensor(0.25))
+    captured = {}
+
+    pipeline._restore_prompt_kv_cache = lambda *_args: None
+
+    def fake_prepare_inputs(input_ids, images, timestep, **model_kwargs):
+        del input_ids, images, timestep
+        captured.update(model_kwargs)
+        return {"model_kwargs": model_kwargs}
+
+    pipeline.prepare_inputs_for_generation = fake_prepare_inputs
+    pipeline.forward_call = lambda **_kwargs: {"diffusion_prediction": torch.tensor([[1.0]])}
+    pipeline._update_model_kwargs_for_generation = lambda _output, model_kwargs: model_kwargs
+
+    output = pipeline.denoise_step(InputBatch.make_batch([state]))
+
+    torch.testing.assert_close(output, torch.tensor([[1.0]]))
+    torch.testing.assert_close(captured["guidance"], torch.tensor([2500.0], dtype=torch.bfloat16))
+    torch.testing.assert_close(captured["timesteps_r"], torch.tensor([0.25]))