[Bugfix][HunyuanImage3] Fix DiT LoRA binding: namespace alias + GQA QKV de-interleave (#6802)

Signed-off-by: Asthenia <asthenia0412@gmail.com>
Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>
Co-authored-by: Asthenia <asthenia0412@gmail.com>
Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>

.buildkite/common/ci_source_file_dependencies.yml

@@ -398,6 +398,8 @@ source_file_dependencies:
   diffusion_hunyuan_image3_accuracy:
     - *hunyuan_image3
     - tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py
+    - tests/e2e/accuracy/test_hunyuan_image3_lora.py
+    - vllm_omni/diffusion/lora/
 
   diffusion_hunyuan_image3_perf:
     - *hunyuan_image3

.buildkite/cuda/test-nightly.yml

@@ -263,7 +263,7 @@ steps:
         commands:
           - export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
           - export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
-          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and H100 and B200 and cards_4" --run-level full_model
+          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py tests/e2e/accuracy/test_hunyuan_image3_lora.py -m "full_model and H100 and B200 and cards_4" --run-level full_model
 
       - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image · Single-GPU"
         source_file_dependencies: diffusion_qwen_image_perf

tests/diffusion/distributed/test_wan_vae_fastpath_install.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """CPU tests for the Wan VAE decoder fast path installer and its exact PyTorch fallbacks."""
 
 from __future__ import annotations
 
 import importlib
+from collections.abc import Callable
 from functools import wraps
 from types import SimpleNamespace
 
@@ -13,6 +14,10 @@ import torch
 from diffusers.models.autoencoders import AutoencoderKLWan
 from torch import nn
 
+from tests.diffusion.distributed.wan_vae_fastpath_helpers import (  # noqa: F401
+    original_wan_rms_norm,
+    unpatched_wan_rms_norm,
+)
 from vllm_omni.diffusion import registry as registry_module
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
@@ -25,7 +30,12 @@ from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
 )
 from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fastpath_forwards
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cpu,
+    pytest.mark.diffusion,
+    pytest.mark.usefixtures("unpatched_wan_rms_norm"),
+]
 
 TINY_RESIDUAL = dict(
     base_dim=8,
@@ -201,7 +211,7 @@ def test_failed_install_restores_original_state(failure_stage: str, monkeypatch)
     with monkeypatch.context() as patch:
         if failure_stage == "binding":
             conv_type = type(vae.post_quant_conv)
-            original_setattr = conv_type.__setattr__
+            original_setattr: Callable[..., None] = conv_type.__setattr__
 
             def fail_binding(module, name, value):
                 original_setattr(module, name, value)
@@ -534,13 +544,30 @@ def test_upsample_forward_only_fuses_nearest_2x() -> None:
             assert torch.equal(actual, expected)
 
 
-def test_rms_norm_vae_substitute_is_not_matched() -> None:
+@torch.no_grad()
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_rms_norm_vae_substitute_is_not_matched(dtype: torch.dtype, monkeypatch) -> None:
+    from diffusers.models.autoencoders import autoencoder_kl_wan
+
     from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
 
     assert not fastpath_forwards.is_diffusers_rms_norm(RMSNormVAE(8, images=False))
-    _, vae = _build_pair(TINY_RESIDUAL, torch.float32)
-    norm = vae.decoder.norm_out
-    assert fastpath_forwards.is_diffusers_rms_norm(norm)
+    _, vae = _build_pair(TINY_RESIDUAL, dtype)
+    assert fastpath_forwards.is_diffusers_rms_norm(vae.decoder.norm_out)
+
+    # Production Wan pipelines use this substitute. The installer must leave
+    # its different epsilon/upcast behavior intact while optimizing the rest.
+    monkeypatch.setattr(autoencoder_kl_wan, "WanRMS_norm", RMSNormVAE)
+    reference, candidate = _build_pair(TINY_RESIDUAL, dtype)
+    report = install_wan_vae_fastpath(candidate)
+    assert report.installed
+    assert "WanRMS_norm" not in report.patched
+    assert "RMSNormVAE" not in report.patched
+    assert "forward" not in candidate.decoder.norm_out.__dict__
+    latents = torch.randn(1, 4, 3, 6, 8).to(dtype)
+    expected = reference.decode(latents, return_dict=False)[0]
+    actual = candidate.decode(latents, return_dict=False)[0]
+    assert torch.equal(actual, expected)
 
 
 @torch.no_grad()

tests/diffusion/distributed/test_wan_vae_fastpath_kernels.py

@@ -18,6 +18,10 @@ from diffusers.models.autoencoders.autoencoder_kl_wan import (
 )
 from torch import nn
 
+from tests.diffusion.distributed.wan_vae_fastpath_helpers import (  # noqa: F401
+    original_wan_rms_norm,
+    unpatched_wan_rms_norm,
+)
 from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import forwards as fp
 from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import (
     install_wan_vae_fastpath,
@@ -28,6 +32,7 @@ from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton
 from vllm_omni.diffusion.distributed.autoencoders.wan_vae_fastpath import triton_upsample as up
 
 pytestmark = [
+    pytest.mark.usefixtures("unpatched_wan_rms_norm"),
     pytest.mark.core_model,
     pytest.mark.cuda,
     pytest.mark.diffusion,

tests/diffusion/distributed/wan_vae_fastpath_helpers.py

@@ -0,0 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Isolate the upstream numerical reference from Wan's global norm patch."""
+
+from __future__ import annotations
+
+import importlib.util
+
+import pytest
+from diffusers.models.autoencoders import autoencoder_kl_wan
+from torch import nn
+
+
+@pytest.fixture(scope="session")
+def original_wan_rms_norm() -> type[nn.Module]:
+    # Wan's package import replaces the public class, including aliases held
+    # by already collected tests. Load the installed upstream source under a
+    # private name instead of reloading (and changing) the live module/classes.
+    spec = importlib.util.spec_from_file_location(
+        "diffusers.models.autoencoders._wan_fastpath_reference", autoencoder_kl_wan.__file__
+    )
+    assert spec is not None and spec.loader is not None
+    reference = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(reference)
+    return reference.WanRMS_norm
+
+
+@pytest.fixture
+def unpatched_wan_rms_norm(monkeypatch, request, original_wan_rms_norm):
+    # Import the patching package before temporarily restoring the reference,
+    # so importing RMSNormVAE in a test cannot change it halfway through.
+    from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
+
+    assert original_wan_rms_norm is not RMSNormVAE
+    monkeypatch.setattr(autoencoder_kl_wan, "WanRMS_norm", original_wan_rms_norm)
+    if hasattr(request.module, "WanRMS_norm"):
+        monkeypatch.setattr(request.module, "WanRMS_norm", original_wan_rms_norm)

tests/diffusion/lora/test_lora_manager.py

@@ -22,6 +22,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummyLoRALayer:
+    base_layer: torch.nn.Module
+
     def __init__(self, n_slices: int, output_slices: tuple[int, ...], tp_size: int = 1):
         self.n_slices = n_slices
         self.output_slices = output_slices
@@ -550,6 +552,146 @@ def test_lora_manager_warns_when_all_adapters_pinned(monkeypatch):
     assert set(manager.list_adapters()) == {1, 2}
 
 
+def _make_hunyuan_image3_pipeline(num_heads=4, num_kv_heads=2, head_dim=2, hidden_size=8):
+    """Use production hooks and the production base-weight splitter without
+    allocating the full checkpoint or initializing distributed workers.
+    """
+    from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import HunyuanImage3Model
+    from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import HunyuanImage3Pipeline
+
+    model = HunyuanImage3Model.__new__(HunyuanImage3Model)
+    torch.nn.Module.__init__(model)
+    model.config = SimpleNamespace(
+        num_attention_heads=num_heads,
+        num_key_value_heads=num_kv_heads,
+        attention_head_dim=head_dim,
+        hidden_size=hidden_size,
+    )
+    pipeline = HunyuanImage3Pipeline.__new__(HunyuanImage3Pipeline)
+    torch.nn.Module.__init__(pipeline)
+    pipeline.model = model
+    pipeline.transformer = model
+    return pipeline
+
+
+def _make_hunyuan_image3_qkv_base():
+    """An un-initialized QKVParallelLinear instance for isinstance() checks.
+
+    Supply the type and checkpoint head metadata for binding-only tests.
+    The PEFT forward test below constructs the real distributed projection.
+    """
+    from vllm.model_executor.layers.linear import QKVParallelLinear
+
+    layer = object.__new__(QKVParallelLinear)
+    layer.total_num_heads = 4
+    layer.total_num_kv_heads = 2
+    layer.head_size = layer.v_head_size = 2
+    return layer
+
+
+def test_lora_manager_activates_hunyuan_image3_fused_qkv_lora():
+    """A fused HI3 ``qkv_proj`` LoRA-B (GQA-interleaved in the adapter) is
+    de-interleaved to the block [Q; K; V] layout the QKV output slices expect,
+    with the namespace alias resolving the PEFT adapter tensors."""
+    manager = DiffusionLoRAManager(
+        pipeline=_make_hunyuan_image3_pipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    wrapper = "transformer.layers.0.self_attn.qkv_proj"
+    packed_layer = _DummyLoRALayer(
+        n_slices=3,
+        output_slices=(8, 4, 4),
+        tp_size=1,
+    )
+    # Use QKV type identity and checkpoint head metadata for the manager's
+    # de-interleave branch; the forward test below constructs real layers.
+    packed_layer.base_layer = _make_hunyuan_image3_qkv_base()
+    manager._lora_modules = {wrapper: packed_layer}
+
+    rank = 3
+    # GQA-interleaved fused QKV rows: [Q-group0, K0, V0, Q-group1, K1, V1].
+    interleaved = torch.arange(16, dtype=torch.bfloat16).unsqueeze(1).repeat(1, rank)
+    lora = LoRALayerWeights(
+        module_name="model.layers.0.self_attn.qkv_proj",
+        rank=rank,
+        lora_alpha=rank,
+        lora_a=torch.ones((rank, 8)),
+        lora_b=interleaved,
+    )
+    manager._registered_adapters = {
+        11: type(
+            "LM",
+            (),
+            {
+                "id": 11,
+                "loras": {"model.layers.0.self_attn.qkv_proj": lora},
+                "get_lora": lambda self, k: self.loras.get(k),
+            },
+        )()
+    }
+
+    manager._activate_adapter(11, 0.5)
+
+    assert packed_layer.reset_calls == 0
+    assert len(packed_layer.set_calls) == 1
+    lora_a_list, lora_b_list = packed_layer.set_calls[0]
+    assert len(lora_a_list) == 3 and len(lora_b_list) == 3
+    # De-interleaved full [Q; K; V] slices, scaled by the lora scale.
+    expected_q = torch.cat([interleaved[0:4], interleaved[8:12]], dim=0) * 0.5
+    expected_k = torch.cat([interleaved[4:6], interleaved[12:14]], dim=0) * 0.5
+    expected_v = torch.cat([interleaved[6:8], interleaved[14:16]], dim=0) * 0.5
+    assert torch.equal(lora_b_list[0], expected_q)
+    assert torch.equal(lora_b_list[1], expected_k)
+    assert torch.equal(lora_b_list[2], expected_v)
+
+
+def test_lora_manager_rejects_hunyuan_image3_qkv_lora_with_unexpected_rows():
+    """If the fused HI3 QKV layout cannot be established (wrong row count), the
+    manager must fail closed instead of applying the LoRA to wrong slices."""
+    manager = DiffusionLoRAManager(
+        pipeline=_make_hunyuan_image3_pipeline(),
+        device=torch.device("cpu"),
+        dtype=torch.bfloat16,
+        max_cached_adapters=1,
+    )
+
+    wrapper = "transformer.layers.0.self_attn.qkv_proj"
+    packed_layer = _DummyLoRALayer(n_slices=3, output_slices=(8, 4, 4))
+    packed_layer.base_layer = _make_hunyuan_image3_qkv_base()
+    manager._lora_modules = {wrapper: packed_layer}
+
+    rank = 3
+    # Row count (12) does not match the HI3 QKV layout (16): fail closed.
+    lora = LoRALayerWeights(
+        module_name="model.layers.0.self_attn.qkv_proj",
+        rank=rank,
+        lora_alpha=rank,
+        lora_a=torch.ones((rank, 8)),
+        lora_b=torch.arange(12 * rank, dtype=torch.bfloat16).view(-1, rank),
+    )
+    manager._registered_adapters = {
+        12: type(
+            "LM",
+            (),
+            {
+                "id": 12,
+                "loras": {"model.layers.0.self_attn.qkv_proj": lora},
+                "get_lora": lambda self, k: self.loras.get(k),
+            },
+        )()
+    }
+
+    with pytest.raises(ValueError, match="applies to no layer"):
+        manager._activate_adapter(12, 0.5)
+
+    assert packed_layer.set_calls == []
+    assert packed_layer.reset_calls >= 1
+    assert manager._active_adapter_id is None
+
+
 def test_lora_manager_applies_multiple_scales_correctly(monkeypatch):
     """Ensure that the LoRA manager applies scales correctly when the
     active adapter receives a different scale, i.e., the rank is unchanged.
@@ -894,3 +1036,96 @@ def test_removing_suspended_adapter_drops_the_upload(monkeypatch):
     assert manager._suspended_adapter_id is None
     assert layer.reset_calls == 1, "the upload must be torn down"
     assert layer.suspended_slices is None
+
+
+@pytest.mark.parametrize("tp_size", [1, 2, 4])
+def test_hunyuan_image3_peft_qkv_forward(tmp_path, monkeypatch, tp_size):
+    """Disk PEFT -> real manager/hooks -> vLLM sharding -> numerical forward.
+
+    TP4 also covers replicated KV heads. Only distributed rank discovery is
+    stubbed: projection construction, adapter loading and all LoRA operations
+    execute production code. The oracle indexes checkpoint rows explicitly.
+    """
+    import json
+
+    from safetensors.torch import save_file
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor import parameter
+    from vllm.model_executor.layers import linear
+
+    adapter_dir = tmp_path / "adapter"
+    adapter_dir.mkdir()
+    rank = 2
+    a = torch.arange(16, dtype=torch.float32).reshape(rank, 8) / 16
+    b = torch.arange(32, dtype=torch.float32).reshape(16, rank) / 32
+    prefix = "base_model.model.model.layers.0.self_attn.qkv_proj"
+    save_file(
+        {f"{prefix}.lora_A.weight": a, f"{prefix}.lora_B.weight": b},
+        str(adapter_dir / "adapter_model.safetensors"),
+    )
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": rank, "lora_alpha": rank, "target_modules": ["qkv_proj"]})
+    )
+    request = LoRARequest(lora_name="hi3", lora_int_id=6411, lora_path=str(adapter_dir))
+    x = torch.arange(24, dtype=torch.float32).reshape(3, 8) / 24
+    q_rows = [0, 1, 2, 3, 8, 9, 10, 11]
+    k_rows = [4, 5, 12, 13]
+    v_rows = [6, 7, 14, 15]
+    for tp_rank in range(tp_size):
+        monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: tp_size)
+        monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: tp_rank)
+        monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: tp_rank)
+        monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: tp_size)
+        with set_current_vllm_config(VllmConfig()):
+            pipeline = _make_hunyuan_image3_pipeline()
+            base = linear.QKVParallelLinear(8, 2, 4, 2, bias=False, params_dtype=torch.float32)
+            base.weight.data.zero_()
+            block = torch.nn.Module()
+            block.self_attn = torch.nn.Module()
+            block.self_attn.qkv_proj = base
+            pipeline.model.layers = torch.nn.ModuleList([block])
+            manager = DiffusionLoRAManager(pipeline, torch.device("cpu"), torch.float32)
+            manager.set_active_adapter(request, lora_scale=0.5)
+            wrapped = pipeline.model.layers[0].self_attn.qkv_proj
+            assert wrapped is not base
+            q_size = 8 // tp_size
+            kv_size = 4 // min(tp_size, 2)
+            kv_rank = tp_rank // max(tp_size // 2, 1)
+            rows = (
+                q_rows[tp_rank * q_size : (tp_rank + 1) * q_size]
+                + k_rows[kv_rank * kv_size : (kv_rank + 1) * kv_size]
+                + v_rows[kv_rank * kv_size : (kv_rank + 1) * kv_size]
+            )
+            expected = (x @ a.T) @ b[rows].T * 0.5
+            with torch.inference_mode():
+                torch.testing.assert_close(wrapped.apply(x), expected)
+                manager.set_active_adapter(None)
+                torch.testing.assert_close(wrapped.apply(x), torch.zeros_like(expected))
+                manager.set_active_adapter(request, lora_scale=0.5)
+                torch.testing.assert_close(wrapped.apply(x), expected)
+
+
+@pytest.mark.parametrize("projection", ["qkv_proj", "o_proj"])
+def test_hunyuan_image3_lora_namespace_is_model_scoped(projection):
+    from vllm.lora.lora_model import LoRAModel
+
+    name = f"model.layers.0.self_attn.{projection}"
+    weights = LoRALayerWeights(name, 2, 2, torch.ones(2, 8), torch.ones(8, 2))
+    adapter = LoRAModel(6411, 2, {name: weights})
+    wrapper = f"transformer.layers.0.self_attn.{projection}"
+    hi3 = DiffusionLoRAManager(_make_hunyuan_image3_pipeline(), torch.device("cpu"), torch.float32)
+    other = DiffusionLoRAManager(_DummyPipeline(), torch.device("cpu"), torch.float32)
+    assert hi3._get_lora_weights(adapter, wrapper) is weights
+    assert other._get_lora_weights(adapter, wrapper) is None
+
+
+@pytest.mark.parametrize("head_dim_key", ["head_dim", "attention_head_dim", None])
+def test_hunyuan_image3_real_qkv_hook_head_dimension(head_dim_key):
+    pipeline = _make_hunyuan_image3_pipeline()
+    del pipeline.model.config.attention_head_dim
+    if head_dim_key is not None:
+        setattr(pipeline.model.config, head_dim_key, 2)
+    b = torch.arange(32, dtype=torch.float32).reshape(16, 2)
+    expected_rows = [0, 1, 2, 3, 8, 9, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15]
+    torch.testing.assert_close(pipeline._deinterleave_fused_qkv_lora_b(b), b[expected_rows])
+    assert pipeline._deinterleave_fused_qkv_lora_b(b[:-1]) is None

tests/e2e/accuracy/test_hunyuan_image3_lora.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""HI3 DiT generation with an on-disk PEFT adapter (issue #6411).
+
+Run with --run-level full_model on four GPUs with at least 80 GiB each.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+import yaml
+from huggingface_hub import hf_hub_download
+from safetensors.torch import save_file
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.lora.request import LoRARequest
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+
+@hardware_test(res={"cuda": ["H100", "B200"]}, num_cards=4)
+def test_hunyuan_image3_dit_lora_generation(tmp_path: Path):
+    model = "tencent/HunyuanImage-3.0-Instruct"
+    config = json.loads(Path(hf_hub_download(model, "config.json")).read_text())
+    hidden = config["hidden_size"]
+    heads = config["num_attention_heads"]
+    kv_heads = config["num_key_value_heads"]
+    head_dim = config.get("head_dim", config["attention_head_dim"])
+    rows = (heads + 2 * kv_heads) * head_dim
+    rank = 8
+    generator = torch.Generator().manual_seed(6411)
+    adapter_dir = tmp_path / "adapter"
+    adapter_dir.mkdir()
+    tensors = {}
+    # Exercise both affected projections in several layers, with different
+    # updates for every interleaved Q/K/V row, rather than a constant delta
+    # that would be invariant under a row permutation.
+    for layer in (0, 1):
+        for projection, out_dim in (("qkv_proj", rows), ("o_proj", hidden)):
+            prefix = f"base_model.model.model.layers.{layer}.self_attn.{projection}"
+            tensors[f"{prefix}.lora_A.weight"] = torch.randn(rank, hidden, generator=generator) * 0.02
+            tensors[f"{prefix}.lora_B.weight"] = torch.randn(out_dim, rank, generator=generator) * 0.1
+    save_file(tensors, str(adapter_dir / "adapter_model.safetensors"))
+    (adapter_dir / "adapter_config.json").write_text(
+        json.dumps({"r": rank, "lora_alpha": rank, "target_modules": ["qkv_proj", "o_proj"]})
+    )
+    tp = 4
+    deploy = tmp_path / "deploy.yaml"
+    deploy.write_text(
+        yaml.safe_dump(
+            {
+                "pipeline": "hunyuan_image3_dit",
+                "async_chunk": False,
+                "trust_remote_code": True,
+                "stages": [
+                    {
+                        "stage_id": 0,
+                        "devices": ",".join(map(str, range(tp))),
+                        "max_num_seqs": 1,
+                        "enforce_eager": True,
+                        "trust_remote_code": True,
+                        "parallel_config": {"tensor_parallel_size": tp},
+                    }
+                ],
+            }
+        )
+    )
+    request = LoRARequest(lora_name="hi3", lora_int_id=6411, lora_path=str(adapter_dir))
+    with OmniRunner(
+        model,
+        trust_remote_code=True,
+        deploy_config=str(deploy),
+        stage_init_timeout=1800,
+        init_timeout=2400,
+    ) as runner:
+
+        def generate(label, lora_request=None, scale=1.0):
+            outputs = runner.omni.generate(
+                "A red ceramic teapot on a wooden table.",
+                OmniDiffusionSamplingParams(
+                    height=512,
+                    width=512,
+                    seed=42,
+                    num_inference_steps=2,
+                    guidance_scale=1.0,
+                    lora_request=lora_request,
+                    lora_scale=scale,
+                ),
+            )
+            image = outputs[0].images[0]
+            assert image.size == (512, 512)
+            image.save(tmp_path / f"{label}.png")
+            return np.asarray(image).copy()
+
+        baseline = generate("baseline")
+        adapted = generate("adapted", request)
+        restored = generate("restored")
+        repeated = generate("repeated", request)
+        zero_scale = generate("zero_scale", request, scale=0.0)
+        assert np.abs(adapted.astype(float) - baseline).mean() > 0.1
+        np.testing.assert_array_equal(restored, baseline)
+        np.testing.assert_array_equal(repeated, adapted)
+        np.testing.assert_array_equal(zero_scale, baseline)

vllm_omni/diffusion/lora/manager.py

@@ -548,7 +548,21 @@ class DiffusionLoRAManager:
             return lora_weights
 
         module_suffix = full_module_name.split(".")[-1]
-        return lora_model.get_lora(module_suffix)
+        lora_weights = lora_model.get_lora(module_suffix)
+        if lora_weights is not None:
+            return lora_weights
+
+        # Model-scoped namespace aliases. HunyuanImage-3 registers the DiT
+        # under ``transformer.layers.*`` (``self.transformer`` aliases
+        # ``self.model``) while PEFT adapters use ``model.layers.*``.
+        name_aliases = getattr(self.pipeline, "_get_diffusion_lora_name_aliases", None)
+        if callable(name_aliases):
+            for alias in name_aliases(full_module_name) or []:
+                lora_weights = lora_model.get_lora(alias)
+                if lora_weights is not None:
+                    return lora_weights
+
+        return None
 
     def _is_active_at_scale(self, adapter_id: int, scale: float) -> bool:
         """True if the adapter_id is active and the current scale matches."""
@@ -646,18 +660,46 @@ class DiffusionLoRAManager:
                     lora_layer.reset_lora(0)
                     continue
 
-                total = sum(output_slices)
-                if lora_weights.lora_b.shape[0] != total:
-                    logger.warning(
-                        "Skipping LoRA for %s due to shape mismatch: lora_b[0]=%d != sum(output_slices)=%d",
-                        full_module_name,
-                        lora_weights.lora_b.shape[0],
-                        total,
+                # HunyuanImage-3 fused ``qkv_proj`` LoRA-B rows are
+                # GQA-interleaved in the checkpoint (per-KV-head
+                # [Q-group, K, V]); de-interleave them to the block layout
+                # [all Q, all K, all V] the QKV output slices expect. The
+                # checkpoint sizes exclude replicated KV heads; set_lora()
+                # selects the appropriate Q shard and shared KV shard.
+                deinterleave = getattr(self.pipeline, "_deinterleave_fused_qkv_lora_b", None)
+                if isinstance(getattr(lora_layer, "base_layer", None), QKVParallelLinear) and callable(deinterleave):
+                    deinterleaved = deinterleave(lora_weights.lora_b)
+                    base = lora_layer.base_layer
+                    output_sizes = (
+                        base.total_num_heads * base.head_size,
+                        base.total_num_kv_heads * base.head_size,
+                        base.total_num_kv_heads * base.v_head_size,
                     )
-                    lora_layer.reset_lora(0)
-                    continue
+                    if (
+                        deinterleaved is None
+                        or len(output_sizes) != n_slices
+                        or deinterleaved.shape[0] != sum(output_sizes)
+                    ):
+                        logger.warning(
+                            "Skipping LoRA for %s: cannot establish HunyuanImage-3 fused-QKV layout",
+                            full_module_name,
+                        )
+                        lora_layer.reset_lora(0)
+                        continue
+                    b_splits = list(torch.split(deinterleaved, list(output_sizes), dim=0))
+                else:
+                    total = sum(output_slices)
+                    if lora_weights.lora_b.shape[0] != total:
+                        logger.warning(
+                            "Skipping LoRA for %s due to shape mismatch: lora_b[0]=%d != sum(output_slices)=%d",
+                            full_module_name,
+                            lora_weights.lora_b.shape[0],
+                            total,
+                        )
+                        lora_layer.reset_lora(0)
+                        continue
+                    b_splits = list(torch.split(lora_weights.lora_b, list(output_slices), dim=0))
 
-                b_splits = list(torch.split(lora_weights.lora_b, list(output_slices), dim=0))
                 lora_a_list = [lora_weights.lora_a] * n_slices
                 lora_b_list = [b * scale for b in b_splits]
                 lora_layer.set_lora(index=0, lora_a=lora_a_list, lora_b=lora_b_list)

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -440,6 +440,38 @@ class HunyuanImage3Pipeline(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler,
         )
 
+    def _get_diffusion_lora_name_aliases(self, full_module_name: str) -> list[str]:
+        """Map LoRA wrapper names to the PEFT adapter namespace.
+
+        ``self.transformer`` aliases ``self.model``, so the generic diffusion
+        LoRA manager registers the DiT layers under ``transformer.layers.*``.
+        PEFT adapters store the same weights under the ``model.layers.*``
+        namespace (matching the HF checkpoint), so expose the alias to let both
+        names resolve to the same adapter tensors.
+        """
+        if full_module_name.startswith("transformer."):
+            return [f"model.{full_module_name[len('transformer.') :]}"]
+        return []
+
+    def _deinterleave_fused_qkv_lora_b(self, lora_b: torch.Tensor) -> torch.Tensor | None:
+        """Re-order a fused ``qkv_proj`` LoRA-B from the checkpoint's
+        GQA-interleaved layout (per-KV-head ``[Q-group, K, V]``) into the block
+        layout ``[all Q, all K, all V]`` expected by the QKV output slices.
+
+        Reuses the exact same split as the base-weight loader
+        (``HunyuanImage3Model._split_qkv_weight``) so the LoRA deltas land on
+        the same output rows as the checkpoint weights. Returns ``None`` if the
+        HI3 fused-QKV layout cannot be established, in which case the caller
+        must fail closed rather than apply the rows to the wrong slices.
+        """
+        split_qkv = getattr(self.transformer, "_split_qkv_weight", None)
+        if not callable(split_qkv):
+            return None
+        try:
+            return split_qkv(lora_b)
+        except (RuntimeError, ValueError):
+            return None
+
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         skip_prefixes = ["lm_head."] if self.hf_config.tie_word_embeddings else []
         # List of unexpected keywords in weight names