Support one-sided MoT biases

Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>

tests/diffusion/kernels/mot/test_mot_gemm.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.layers.mot.ops import mot_gemm
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
+
+_CONFIG = {
+    "BLOCK_SIZE_M": 16,
+    "BLOCK_SIZE_N": 64,
+    "BLOCK_SIZE_K": 32,
+    "GROUP_SIZE_M": 8,
+    "num_warps": 4,
+    "num_stages": 2,
+}
+
+
+class _RecordingKernel:
+    def __init__(self) -> None:
+        self.calls: list[tuple[tuple[object, ...], dict[str, object]]] = []
+
+    def __getitem__(self, grid):
+        del grid
+        return self._launch
+
+    def _launch(self, *args, **kwargs) -> None:
+        self.calls.append((args, kwargs))
+
+
+@pytest.mark.parametrize(
+    ("bias_text", "bias_vae"),
+    [(None, None), ("text", None), (None, "vae"), ("text", "vae")],
+    ids=["no-bias", "text-only", "vae-only", "both"],
+)
+def test_invoke_mot_gemm_accepts_independent_biases(
+    bias_text: str | None,
+    bias_vae: str | None,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    kernel = _RecordingKernel()
+    monkeypatch.setattr(mot_gemm, "mot_unified_gemm_kernel", kernel)
+
+    A = torch.ones((3, 4), dtype=torch.float16)
+    B_text = torch.ones((4, 5), dtype=torch.float16)
+    B_vae = torch.ones((4, 5), dtype=torch.float16)
+    C = torch.empty((3, 5), dtype=torch.float16)
+    text_indices = torch.tensor([0, 2], dtype=torch.long)
+    vae_indices = torch.tensor([1], dtype=torch.long)
+    text_bias = torch.ones(5, dtype=torch.float16) if bias_text else None
+    vae_bias = torch.ones(5, dtype=torch.float16) if bias_vae else None
+
+    mot_gemm.invoke_mot_gemm(
+        A=A,
+        B_text=B_text,
+        B_vae=B_vae,
+        C=C,
+        bias_text=text_bias,
+        bias_vae=vae_bias,
+        text_indices=text_indices,
+        vae_indices=vae_indices,
+        A_scale=None,
+        B_text_scale=None,
+        B_vae_scale=None,
+        use_fp8_w8a8=False,
+        use_int8_w8a8=False,
+        use_int8_w8a16=False,
+        use_int4_w4a16=False,
+        A_per_channel_quant=False,
+        B_per_channel_quant=False,
+        config=_CONFIG,
+    )
+
+    assert len(kernel.calls) == 1
+    args, kwargs = kernel.calls[0]
+    assert kwargs["HAS_BIAS_TEXT"] is (text_bias is not None)
+    assert kwargs["HAS_BIAS_VAE"] is (vae_bias is not None)
+    assert "HAS_BIAS" not in kwargs
+
+    launched_text_bias, launched_vae_bias = args[4:6]
+    if text_bias is None and vae_bias is None:
+        assert launched_text_bias == launched_vae_bias == 0
+    elif text_bias is None:
+        assert launched_text_bias is vae_bias
+        assert launched_vae_bias is vae_bias
+    elif vae_bias is None:
+        assert launched_text_bias is text_bias
+        assert launched_vae_bias is text_bias
+    else:
+        assert launched_text_bias is text_bias
+        assert launched_vae_bias is vae_bias

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -513,6 +513,57 @@ def test_mot_o_proj(
             )
 
 
+@pytest.mark.parametrize(
+    ("bias_text", "bias_vae"),
+    [(True, False), (False, True)],
+    ids=["text-only", "vae-only"],
+)
+def test_mot_row_one_sided_bias(bias_text: bool, bias_vae: bool):
+    """The fused row projection must support a bias on only one expert."""
+    dcfg = _parse_dtype("w16a16_bf16")
+    K = 32
+    N = 64
+    M = 5
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias_text,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias_vae,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTRowParallelLinear(
+            K,
+            N,
+            bias=bias_text,
+            vae_bias=bias_vae,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+        text_idx = torch.tensor([0, 2], dtype=torch.long, device="cuda")
+        vae_idx = torch.tensor([1, 3, 4], dtype=torch.long, device="cuda")
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        with torch.no_grad():
+            ref = _reference_forward(x, text_idx, vae_idx, text_linear, vae_linear)
+            mot_out, _ = mot_linear(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, f"Row one-sided bias text={bias_text} vae={bias_vae}")
+
+
 # =========================================================================
 #  Test: und-mode (text_indices=None) — falls back to standard forward
 # =========================================================================

tests/model_executor/models/test_kv_cache_scale_mapper.py

@@ -67,11 +67,8 @@ class _FakeHunyuanModel(_FakeModel):
             use_cla=False,
         )
 
-    def get_expert_mapping(self) -> list[object]:
-        return []
-
-    def _get_expert_weights_remapping(self) -> dict[str, object]:
-        return {}
+    def get_expert_mapping(self) -> tuple[list[object], dict[str, object]]:
+        return [], {}
 
     def _split_qkv_weight(self, weight: torch.Tensor) -> torch.Tensor:
         return weight
@@ -84,49 +81,6 @@ def test_model_loaders_do_not_call_removed_get_cache_scale(rel_path: str) -> Non
     assert ".get_cache_scale(" not in source
 
 
-@pytest.mark.parametrize(
-    ("module_name", "class_name"),
-    [
-        (
-            "vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3",
-            "HunyuanModel",
-        ),
-        (
-            "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer",
-            "HunyuanImage3Model",
-        ),
-    ],
-)
-def test_hunyuan_expert_mapping_separates_public_and_local_contracts(
-    module_name: str,
-    class_name: str,
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
-    """The public expert mapping must be a flat list, not a loader-side pair."""
-    module = __import__(module_name, fromlist=[class_name])
-    expected_mapping = [("param", "weight", 0, "gate_proj")]
-    fake_model = SimpleNamespace(
-        config=SimpleNamespace(num_experts=1),
-        num_redundant_experts=0,
-    )
-    monkeypatch.setattr(module, "_is_moe", lambda config: True)
-    if hasattr(module, "fused_moe_make_expert_params_mapping"):
-        monkeypatch.setattr(module, "fused_moe_make_expert_params_mapping", lambda **kwargs: expected_mapping)
-    else:
-        monkeypatch.setattr(
-            module,
-            "FusedMoE",
-            SimpleNamespace(make_expert_params_mapping=lambda *args, **kwargs: expected_mapping),
-        )
-
-    model_cls = getattr(module, class_name)
-    assert model_cls.get_expert_mapping(fake_model) == expected_mapping
-    assert model_cls._get_expert_weights_remapping(fake_model) == {
-        "gate_proj": ("gate_and_up_proj", 1, 2),
-        "up_proj": ("gate_and_up_proj", 0, 2),
-    }
-
-
 @pytest.mark.parametrize(
     ("module_name", "class_name", "fake_model_cls"),
     [

vllm_omni/diffusion/layers/mot/ops/mot_gemm.py

@@ -238,6 +238,7 @@ def _get_mot_pointers(
     cur_stride_bk = stride_bk_vae
     cur_stride_bn = stride_bn_vae
     M_limit = M_vae
+    is_vae = 1
 
     # Text Path
     if pid < num_pid_m_text * num_pid_n:
@@ -251,6 +252,7 @@ def _get_mot_pointers(
         cur_stride_bk = stride_bk_text
         cur_stride_bn = stride_bn_text
         M_limit = M_text
+        is_vae = 0
 
     # 3. Calculate Grid coordinates(grouping)
     cur_num_pid_m = tl.cdiv(M_limit, BLOCK_SIZE_M)
@@ -283,6 +285,7 @@ def _get_mot_pointers(
         offs_n,
         n_mask,  # N-dim info
         M_limit,  # Boundary
+        is_vae,  # Selected expert kind: 0=text, 1=VAE
         cur_b_ptr,
         cur_bias_ptr,  # Selected Pointers
         cur_scale_b_ptr,  # Selected Scale Pointer
@@ -533,7 +536,8 @@ def mot_unified_gemm_kernel(
     # Quant Control
     # 0=None, 1=W8A8, 2=W8A16, 3=W4A16
     QUANT_TYPE: tl.constexpr,
-    HAS_BIAS: tl.constexpr = False,
+    HAS_BIAS_TEXT: tl.constexpr = False,
+    HAS_BIAS_VAE: tl.constexpr = False,
 ):
     pid = tl.program_id(axis=0)
 
@@ -548,6 +552,7 @@ def mot_unified_gemm_kernel(
         offs_n,
         n_mask,
         M_limit,
+        is_vae,
         cur_b_ptr,
         cur_bias_ptr,
         cur_scale_b_ptr,
@@ -675,8 +680,12 @@ def mot_unified_gemm_kernel(
     # -----------------------------------------------------------
 
     # Bias Add
-    if HAS_BIAS:
+    if HAS_BIAS_TEXT or HAS_BIAS_VAE:
         bias = tl.load(cur_bias_ptr + offs_n, mask=n_mask, other=0.0)
+        if HAS_BIAS_TEXT and not HAS_BIAS_VAE:
+            bias = tl.where(is_vae == 0, bias, 0.0)
+        elif HAS_BIAS_VAE and not HAS_BIAS_TEXT:
+            bias = tl.where(is_vae == 1, bias, 0.0)
         c = c + bias[None, :]  # reshape to (1, BLOCK_SIZE_N)
 
     # Cast C into the output dtype
@@ -832,11 +841,8 @@ def invoke_mot_gemm(
     STRIDE_BK_IS_1 = (B_text.stride(0) == 1) and (B_vae.stride(0) == 1)
     STRIDE_BN_IS_1 = (B_text.stride(1) == 1) and (B_vae.stride(1) == 1)
 
-    # bias check
-    assert (bias_text is None) == (bias_vae is None), (
-        "Bias must be provided for both Text and VAE simultaneously, or neither."
-    )
-    has_bias = bias_text is not None
+    has_bias_text = bias_text is not None
+    has_bias_vae = bias_vae is not None
 
     # --- 3. Grid Calculation ---
     def grid(META):
@@ -853,7 +859,8 @@ def invoke_mot_gemm(
             "ACCUMULATOR_DTYPE": ACCUMULATOR_DTYPE,
             "COMPUTE_DTYPE": COMPUTE_DTYPE,
             "OUTPUT_DTYPE": OUTPUT_DTYPE,
-            "HAS_BIAS": has_bias,
+            "HAS_BIAS_TEXT": has_bias_text,
+            "HAS_BIAS_VAE": has_bias_vae,
             "EVEN_K": EVEN_K,
             "EVEN_N": EVEN_N,
             "STRIDE_AK_IS_1": STRIDE_AK_IS_1,
@@ -866,8 +873,17 @@ def invoke_mot_gemm(
     p_a_scale = A_scale if A_scale is not None else 0
     p_b_text_scale = B_text_scale if B_text_scale is not None else 0
     p_b_vae_scale = B_vae_scale if B_vae_scale is not None else 0
-    p_bias_text = bias_text if bias_text is not None else 0
-    p_bias_vae = bias_vae if bias_vae is not None else 0
+    # Triton requires both route-selected bias arguments to have a pointer
+    # type. Alias the missing route to the available tensor; its compile-time
+    # bias flag and runtime route mask keep it from being read for that route.
+    if bias_text is None:
+        p_bias_text = bias_vae if bias_vae is not None else 0
+    else:
+        p_bias_text = bias_text
+    if bias_vae is None:
+        p_bias_vae = bias_text if bias_text is not None else 0
+    else:
+        p_bias_vae = bias_vae
 
     # Quantization granularity
     stride_scale_a = 1 if A_per_channel_quant else 0

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -2190,28 +2190,24 @@ class HunyuanImage3Model(nn.Module):
         return torch.concat((q, k, v))
 
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        if not _is_moe(self.config):
-            return []
-
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        return FusedMoE.make_expert_params_mapping(
-            self,
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.num_experts,
-            num_redundant_experts=self.num_redundant_experts,
-        )
-
-    def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]:
-        """Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
-        if not _is_moe(self.config):
-            return {}
-        return {
-            "gate_proj": ("gate_and_up_proj", 1, 2),
-            "up_proj": ("gate_and_up_proj", 0, 2),
-        }
+        if _is_moe(self.config):
+            # Params for weights, fp8 weight scales, fp8 activation scales
+            # (param_name, weight_name, expert_id, shard_id)
+            fused_moe_expert_mapping = FusedMoE.make_expert_params_mapping(
+                self,
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=self.config.num_experts,
+                num_redundant_experts=self.num_redundant_experts,
+            )
+            expert_weights_remapping = {
+                "gate_proj": ("gate_and_up_proj", 1, 2),
+                "up_proj": ("gate_and_up_proj", 0, 2),
+            }
+            return fused_moe_expert_mapping, expert_weights_remapping
+        else:
+            return [], {}
 
     # rename for delay load
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
@@ -2247,8 +2243,7 @@ class HunyuanImage3Model(nn.Module):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
-        expert_params_mapping = self.get_expert_mapping()
-        expert_weights_remapping = self._get_expert_weights_remapping()
+        expert_params_mapping, expert_weights_remapping = self.get_expert_mapping()
 
         # List of unexpected keywords in weight names
         unexpected_keywords = [

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -128,29 +128,25 @@ class HunyuanModel(HunYuanModel):
         v = v.reshape(-1, hidden_size)
         return torch.concat((q, k, v))
 
-    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
-        if not _is_moe(self.config):
-            return []
-
-        # Params for weights, fp8 weight scales, fp8 activation scales
-        # (param_name, weight_name, expert_id, shard_id)
-        return fused_moe_make_expert_params_mapping(
-            model=self,
-            ckpt_gate_proj_name="gate_proj",
-            ckpt_down_proj_name="down_proj",
-            ckpt_up_proj_name="up_proj",
-            num_experts=self.config.num_experts,
-            num_redundant_experts=self.num_redundant_experts,
-        )
-
-    def _get_expert_weights_remapping(self) -> dict[str, tuple[str, int, int]]:
-        """Return the vLLM-Omni-local remapping for Hunyuan checkpoint weights."""
-        if not _is_moe(self.config):
-            return {}
-        return {
-            "gate_proj": ("gate_and_up_proj", 1, 2),
-            "up_proj": ("gate_and_up_proj", 0, 2),
-        }
+    def get_expert_mapping(self) -> tuple[list[tuple[str, str, int, str]], dict[str, tuple[str, int, int]]]:
+        if _is_moe(self.config):
+            # Params for weights, fp8 weight scales, fp8 activation scales
+            # (param_name, weight_name, expert_id, shard_id)
+            fused_moe_expert_mapping = fused_moe_make_expert_params_mapping(
+                model=self,
+                ckpt_gate_proj_name="gate_proj",
+                ckpt_down_proj_name="down_proj",
+                ckpt_up_proj_name="up_proj",
+                num_experts=self.config.num_experts,
+                num_redundant_experts=self.num_redundant_experts,
+            )
+            expert_weights_remapping = {
+                "gate_proj": ("gate_and_up_proj", 1, 2),
+                "up_proj": ("gate_and_up_proj", 0, 2),
+            }
+            return fused_moe_expert_mapping, expert_weights_remapping
+        else:
+            return [], {}
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
         cla_factor = _get_cla_factor(self.config)
@@ -178,8 +174,7 @@ class HunyuanModel(HunYuanModel):
 
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
-        expert_params_mapping = self.get_expert_mapping()
-        expert_weights_remapping = self._get_expert_weights_remapping()
+        expert_params_mapping, expert_weights_remapping = self.get_expert_mapping()
 
         # List of unexpected keywords in weight names
         unexpected_keywords = [