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
 # =========================================================================