diff --git a/tests/diffusion/kernels/mot/test_mot_gemm.py b/tests/diffusion/kernels/mot/test_mot_gemm.py new file mode 100644 index 000000000..7dc6d6265 --- /dev/null +++ b/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 diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py index 49622aecd..30f087174 100644 --- a/tests/diffusion/kernels/mot/test_mot_linear.py +++ b/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 # ========================================================================= diff --git a/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py b/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py index f1bffbbac..c088dfd85 100644 --- a/vllm_omni/diffusion/layers/mot/ops/mot_gemm.py +++ b/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