[Enhancement] Introduce High-Performance MoT (Mixture-of-Tokens) Kernels: Triton Implementation (#3960)

Signed-off-by: Yufeng Zhu <yzhu802@gatech.edu>
Signed-off-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: princepride <wangzhipeng628@gmail.com>
Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Co-authored-by: Yufeng Zhu <yzhu802@gatech.edu>
Co-authored-by: Zhengyuan Su <su.zhengyuan@u.nus.edu>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -0,0 +1,689 @@
+# ruff: noqa: N803, E741
+"""Layer-level correctness & performance test for MoT parallel linear layers.
+
+Compares two equivalent computation paths:
+  - Reference: 2x standard vLLM parallel linear layers + PyTorch index
+    scatter/gather (text_linear(x[text_idx]) + vae_linear(x[vae_idx]))
+  - Target:    1x MoT fused parallel linear layer
+    (mot_linear(x, text_indices, vae_indices))
+
+The reference path uses cuBLAS GEMM (always auto-tuned by cuBLAS).
+The MoT path uses a fused Triton kernel whose tile config is loaded
+from a JSON file matched by ``device + dtype``.  If no tuned
+config is found for the current GPU, the kernel falls back to a conservative
+default and a warning is printed — the correctness test still passes
+but the performance comparison is NOT representative.
+
+Usage::
+    pytest tests/diffusion/kernels/mot/test_mot_linear.py -v -s
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from dataclasses import dataclass
+
+import pytest
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.linear import (
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+
+from vllm_omni.diffusion.distributed.parallel_state import (
+    destroy_distributed_env,
+    init_distributed_environment,
+    initialize_model_parallel,
+    model_parallel_is_initialized,
+)
+from vllm_omni.diffusion.layers.mot.mot_qkv_parallel_linear import (
+    MoTQKVParallelLinear,
+)
+from vllm_omni.diffusion.layers.mot.mot_row_parallel_linear import (
+    MoTRowParallelLinear,
+)
+from vllm_omni.diffusion.layers.mot.ops.mot_gemm import (
+    get_best_mot_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+# BAGEL-7B-MoT architecture parameters
+_BAGEL_HEAD_SIZE = 128
+_BAGEL_TOTAL_NUM_HEADS = 28
+_BAGEL_TOTAL_NUM_KV_HEADS = 4
+_VAE_CHUNK_SIZE = 1024  # the token number of one image
+_IMAGE_NUM = [1, 2, 4, 8]
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _init_single_rank_tp_env():
+    """Initialize single-rank distributed/TP env for vLLM linear params."""
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29501")
+
+    if not torch.distributed.is_initialized():
+        init_distributed_environment(world_size=1, rank=0, local_rank=0)
+
+    if not model_parallel_is_initialized():
+        initialize_model_parallel(
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            sequence_parallel_size=1,
+            ulysses_degree=1,
+            ring_degree=1,
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+        )
+
+    yield
+    destroy_distributed_env()
+
+
+# ---------------------------------------------------------------------------
+# DType configuration — extensible for future quantized types
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DTypeConfig:
+    """Parsed dtype configuration for a test case."""
+
+    torch_dtype: torch.dtype
+    use_fp8_w8a8: bool = False
+    use_int8_w8a16: bool = False
+    use_int4_w4a16: bool = False
+
+
+def _parse_dtype(dtype_str: str) -> DTypeConfig:
+    """Parse a dtype string into quantization flags and torch dtype.
+
+    Supported now:
+        "w16a16_bf16"  — BF16 weights & activations (no quantization)
+        "w16a16_fp16"  — FP16 weights & activations (no quantization)
+    Reserved for future:
+        "fp8_w8a8"     — FP8 W8A8 quantization
+        "int8_w8a16"   — INT8 weight-only quantization
+        "int4_w4a16"   — INT4 weight-only quantization
+    """
+    supported: dict[str, DTypeConfig] = {
+        "w16a16_bf16": DTypeConfig(torch_dtype=torch.bfloat16),
+        "w16a16_fp16": DTypeConfig(torch_dtype=torch.float16),
+    }
+    if dtype_str in supported:
+        return supported[dtype_str]
+    pytest.skip(f"Quantized dtype '{dtype_str}' not yet implemented in layer test")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _report_mot_config(K: int, N: int, M: int):
+    """Print which Triton tile config the MoT kernel will use."""
+    loaded_m_key, config = get_best_mot_config(M, N, K, None)
+    if loaded_m_key == -1:
+        print(
+            "  [config] WARNING: No tuned config found — "
+            "using conservative default. "
+            "Performance numbers are NOT representative. "
+            "Run mot_linear_benchmarks.py --tune to generate configs."
+        )
+    else:
+        print(f"  [config] Tuned config loaded (actual M={M}, loaded M={loaded_m_key}) config = {config})")
+
+
+def _make_indices(image_num: int, vae_chunk_size: int, device: str = "cuda") -> tuple[torch.Tensor, torch.Tensor, int]:
+    """
+    Simulate exact Bagel-MoT distributions for image generation:
+    Pattern per image like: [1 Text] + [4096 VAE] + [1 Text]
+    Returns text_indices, vae_indices, and the exact total M.
+    """
+    text_idx_list = []
+    vae_idx_list = []
+
+    current_idx = 0
+    for _ in range(image_num):
+        text_idx_list.append(current_idx)
+        current_idx += 1
+
+        vae_idx_list.extend(range(current_idx, current_idx + vae_chunk_size))
+        current_idx += vae_chunk_size
+
+        text_idx_list.append(current_idx)
+        current_idx += 1
+
+    text_indices = torch.tensor(text_idx_list, dtype=torch.long, device=device)
+    vae_indices = torch.tensor(vae_idx_list, dtype=torch.long, device=device)
+
+    exact_M = current_idx  # exact_M = image_num * (vae_chunk_size + 2)
+
+    return text_indices, vae_indices, exact_M
+
+
+def _benchmark(fn, warmup: int = 20, iters: int = 100) -> float:
+    """Return mean latency in milliseconds."""
+    cache_flusher = torch.empty(int(256 * 1024 * 1024 / 4), dtype=torch.int32, device="cuda")
+    for _ in range(warmup):
+        fn()
+    torch.accelerator.synchronize()
+
+    t_flush_start = time.perf_counter()
+    for _ in range(iters):
+        cache_flusher.zero_()
+    torch.accelerator.synchronize()
+    flush_time_total = time.perf_counter() - t_flush_start
+
+    #  3. Measure the total time of "flush cache + operator execution"
+    torch.accelerator.synchronize()
+    t_total_start = time.perf_counter()
+    for _ in range(iters):
+        cache_flusher.zero_()
+        fn()
+    torch.accelerator.synchronize()
+    total_time = time.perf_counter() - t_total_start
+
+    # 4. Asynchronous subtraction separation
+    # The E2E time of the pure operator = total time - flush time
+    pure_fn_time_total = total_time - flush_time_total
+    avg_ms = (pure_fn_time_total / iters) * 1000.0
+
+    del cache_flusher
+    return max(avg_ms, 0.001)
+
+
+def _sync_weights(ref_text, ref_vae, mot_layer):
+    """Assign the same random weights to reference layers and MoT layer."""
+    with torch.no_grad():
+        W_text = torch.randn_like(ref_text.weight) * 0.02
+        W_vae = torch.randn_like(ref_vae.weight) * 0.02
+        ref_text.weight.copy_(W_text)
+        ref_vae.weight.copy_(W_vae)
+        mot_layer.weight.copy_(W_text)
+        mot_layer.gen_exp.weight.copy_(W_vae)
+        if ref_text.bias is not None and mot_layer.bias is not None:
+            b_text = torch.randn_like(ref_text.bias) * 0.02
+            ref_text.bias.copy_(b_text)
+            mot_layer.bias.copy_(b_text)
+        if ref_vae.bias is not None and mot_layer.gen_exp.bias is not None:
+            b_vae = torch.randn_like(ref_vae.bias) * 0.02
+            ref_vae.bias.copy_(b_vae)
+            mot_layer.gen_exp.bias.copy_(b_vae)
+
+
+def _reference_forward(x, text_indices, vae_indices, text_linear, vae_linear):
+    """Reference path: index-gather → 2x standard linear → index-scatter."""
+    M = x.size(0)
+    out_text = text_linear(x[text_indices])
+    out_vae = vae_linear(x[vae_indices])
+    if isinstance(out_text, tuple):
+        out_text = out_text[0]
+    if isinstance(out_vae, tuple):
+        out_vae = out_vae[0]
+    N = out_text.size(-1)
+    output = torch.empty(M, N, dtype=x.dtype, device=x.device)
+    output[text_indices] = out_text
+    output[vae_indices] = out_vae
+    return output
+
+
+def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
+    """Compare outputs, print metrics, and assert correctness.
+
+    Both ``ref`` and ``mot`` are in the original compute dtype (e.g. bf16).
+    We upcast to fp32 solely for computing error metrics with higher
+    arithmetic precision — the actual layer outputs remain bf16.
+    """
+    # Upcast for numerically stable error computation only
+    ref_hp = ref.float()
+    mot_hp = mot.float()
+
+    abs_err = (ref_hp - mot_hp).abs()
+    max_abs = abs_err.max().item()
+
+    # Mixed metric: relative error where |ref| >= 1, absolute error otherwise
+    denom = ref_hp.abs().clamp(min=1.0)
+    max_rel = (abs_err / denom).max().item()
+
+    cos_sim = (
+        torch.nn.functional.cosine_similarity(
+            ref_hp,
+            mot_hp,
+            dim=-1,
+        )
+        .min()
+        .item()
+    )
+
+    print(f"\n  [{tag}]  max_abs={max_abs:.4e}  max_rel={max_rel:.4e}  min_cos_sim={cos_sim:.6f}")
+
+    # Cosine similarity is the primary correctness gate: robust to scale
+    # and accumulation-order differences between cuBLAS and Triton.
+    # For bf16 GEMM with K up to ~20k, cos_sim > 0.99 is easily achieved.
+    assert cos_sim > 0.98, f"Cosine similarity too low: {cos_sim:.6f}"
+    # Supplementary per-element check (generous to avoid flaky failures
+    # on extreme K dimensions like 18944)
+    assert max_rel < 0.1, f"Max relative error too large: {max_rel:.4e}"
+
+
+def _run_timing(
+    ref_fn,
+    mot_fn,
+    tag: str,
+    warmup: int = 20,
+    iters: int = 100,
+):
+    """Benchmark both paths and print timing comparison."""
+    ref_ms = _benchmark(ref_fn, warmup=warmup, iters=iters)
+    mot_ms = _benchmark(mot_fn, warmup=warmup, iters=iters)
+    speedup = ref_ms / mot_ms if mot_ms > 0 else float("inf")
+    print(f"  [{tag}]  Ref(2x linear): {ref_ms:.3f} ms  |  MoT(fused): {mot_ms:.3f} ms  |  Speedup: {speedup:.2f}x")
+
+
+# =========================================================================
+#  Test: qkv proj
+# =========================================================================
+
+
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize(
+    "image_num, K, N, dtype",
+    [(num, 3584, 4608, "w16a16_bf16") for num in _IMAGE_NUM] + [(num, 3584, 4608, "w16a16_fp16") for num in _IMAGE_NUM],
+    ids=[f"img{num}_K3584_N4608_bf16" for num in _IMAGE_NUM] + [f"img{num}_K3584_N4608_fp16" for num in _IMAGE_NUM],
+)
+def test_mot_qkv_parallel(image_num: int, K: int, N: int, dtype: str, bias: bool):
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = QKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = QKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTQKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            vae_bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        assert text_linear.output_size_per_partition == N, (
+            f"Expected output_size_per_partition={N}, "
+            f"got {text_linear.output_size_per_partition}. "
+            f"Check head parameters."
+        )
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+
+        text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        tag = f"QKVParallel M={M} K={K} N={N}"
+        _report_mot_config(K, N, M)
+
+        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, tag)
+
+        with torch.no_grad():
+            _run_timing(
+                lambda: _reference_forward(
+                    x,
+                    text_idx,
+                    vae_idx,
+                    text_linear,
+                    vae_linear,
+                ),
+                lambda: mot_linear(x, text_idx, vae_idx),
+                tag,
+            )
+
+
+# =========================================================================
+#  Test: o proj
+# =========================================================================
+
+
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize(
+    "image_num, K, N, dtype",
+    [(num, 3584, 3584, "w16a16_bf16") for num in _IMAGE_NUM] + [(num, 3584, 3584, "w16a16_fp16") for num in _IMAGE_NUM],
+    ids=[f"img{num}_K3584_N3584_bf16" for num in _IMAGE_NUM] + [f"img{num}_K3584_N3584_fp16" for num in _IMAGE_NUM],
+)
+def test_mot_o_proj(
+    image_num: int,
+    K: int,
+    N: int,
+    dtype: str,
+    bias: bool,
+):
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTRowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            vae_bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+
+        text_idx, vae_idx, M = _make_indices(image_num, _VAE_CHUNK_SIZE)
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        tag = f"O Proj M={M} K={K} N={N}"
+        _report_mot_config(K, N, M)
+
+        # Correctness (also warms up Triton JIT compilation)
+        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, tag)
+
+        # Performance
+        with torch.no_grad():
+            _run_timing(
+                lambda: _reference_forward(
+                    x,
+                    text_idx,
+                    vae_idx,
+                    text_linear,
+                    vae_linear,
+                ),
+                lambda: mot_linear(x, text_idx, vae_idx),
+                tag,
+            )
+
+
+# =========================================================================
+#  Test: und-mode (text_indices=None) — falls back to standard forward
+# =========================================================================
+
+
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize(
+    "K, N, dtype",
+    [
+        (3584, 4608, "w16a16_bf16"),
+        (3584, 4608, "w16a16_fp16"),
+    ],
+    ids=["QKV_K3584_N4608_bf16", "QKV_K3584_N4608_fp16"],
+)
+def test_mot_qkv_und_mode(K: int, N: int, dtype: str, bias: bool):
+    """und-mode: text_indices=None should produce same output as standard QKVParallelLinear."""
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+    M = 1026
+
+    with set_current_vllm_config(VllmConfig()):
+        ref_linear = QKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTQKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=bias,
+            vae_bias=bias,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        with torch.no_grad():
+            mot_linear.weight.copy_(ref_linear.weight)
+            if bias and ref_linear.bias is not None:
+                mot_linear.bias.copy_(ref_linear.bias)
+
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        with torch.no_grad():
+            ref_out, _ = ref_linear(x)
+            mot_out, _ = mot_linear(x, text_indices=None, vae_indices=None)
+
+        _check_and_report(ref_out, mot_out, f"QKV und-mode M={M} K={K}")
+
+
+@pytest.mark.parametrize("bias", [True, False])
+@pytest.mark.parametrize(
+    "K, N, dtype",
+    [
+        (3584, 3584, "w16a16_bf16"),
+        (3584, 3584, "w16a16_fp16"),
+    ],
+    ids=["Row_K3584_N3584_bf16", "Row_K3584_N3584_fp16"],
+)
+def test_mot_row_und_mode(K: int, N: int, dtype: str, bias: bool):
+    """und-mode: text_indices=None should produce same output as standard RowParallelLinear."""
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+    M = 1026
+
+    with set_current_vllm_config(VllmConfig()):
+        ref_linear = RowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTRowParallelLinear(
+            K,
+            N,
+            bias=bias,
+            vae_bias=bias,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        with torch.no_grad():
+            mot_linear.weight.copy_(ref_linear.weight)
+            if bias and ref_linear.bias is not None:
+                mot_linear.bias.copy_(ref_linear.bias)
+
+        x = torch.randn(M, K, dtype=dcfg.torch_dtype, device="cuda")
+
+        with torch.no_grad():
+            ref_out, _ = ref_linear(x)
+            mot_out, _ = mot_linear(x, text_indices=None, vae_indices=None)
+
+        _check_and_report(ref_out, mot_out, f"Row und-mode M={M} K={K}")
+
+
+# =========================================================================
+#  Test: boundary cases (all-text, all-VAE)
+# =========================================================================
+
+
+@pytest.mark.parametrize(
+    "boundary_mode",
+    ["all_text", "all_vae"],
+    ids=["all_text", "all_vae"],
+)
+@pytest.mark.parametrize("dtype", ["w16a16_bf16", "w16a16_fp16"])
+def test_mot_qkv_boundary(boundary_mode: str, dtype: str):
+    """Boundary: all tokens routed to one expert."""
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+    M = 512
+    K = 3584
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = QKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=False,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = QKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=False,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTQKVParallelLinear(
+            hidden_size=K,
+            head_size=_BAGEL_HEAD_SIZE,
+            total_num_heads=_BAGEL_TOTAL_NUM_HEADS,
+            total_num_kv_heads=_BAGEL_TOTAL_NUM_KV_HEADS,
+            bias=False,
+            vae_bias=False,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+
+        all_indices = torch.arange(M, dtype=torch.long, device="cuda")
+        empty_indices = torch.empty(0, dtype=torch.long, device="cuda")
+
+        if boundary_mode == "all_text":
+            text_idx, vae_idx = all_indices, empty_indices
+        else:
+            text_idx, vae_idx = empty_indices, all_indices
+
+        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"QKV boundary={boundary_mode} M={M}")
+
+
+@pytest.mark.parametrize(
+    "boundary_mode",
+    ["all_text", "all_vae"],
+    ids=["all_text", "all_vae"],
+)
+@pytest.mark.parametrize("dtype", ["w16a16_bf16", "w16a16_fp16"])
+def test_mot_row_boundary(boundary_mode: str, dtype: str):
+    """Boundary: all tokens routed to one expert for RowParallel."""
+    dcfg = _parse_dtype(dtype)
+    torch.manual_seed(42)
+    M = 512
+    K = 3584
+    N = 3584
+
+    with set_current_vllm_config(VllmConfig()):
+        text_linear = RowParallelLinear(
+            K,
+            N,
+            bias=False,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        vae_linear = RowParallelLinear(
+            K,
+            N,
+            bias=False,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+        mot_linear = MoTRowParallelLinear(
+            K,
+            N,
+            bias=False,
+            vae_bias=False,
+            input_is_parallel=True,
+            params_dtype=dcfg.torch_dtype,
+            disable_tp=True,
+        ).cuda()
+
+        _sync_weights(text_linear, vae_linear, mot_linear)
+
+        all_indices = torch.arange(M, dtype=torch.long, device="cuda")
+        empty_indices = torch.empty(0, dtype=torch.long, device="cuda")
+
+        if boundary_mode == "all_text":
+            text_idx, vae_idx = all_indices, empty_indices
+        else:
+            text_idx, vae_idx = empty_indices, all_indices
+
+        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 boundary={boundary_mode} M={M}")

tests/diffusion/kernels/mot/test_mot_norm.py

@@ -0,0 +1,357 @@
+# ruff: noqa: N803, E741
+"""Layer-level correctness & performance test for MoTRMSNorm.
+
+Compares two equivalent computation paths:
+  - Reference: 2x vLLM IR RMSNorm op + PyTorch index scatter/gather
+    (ir.ops.rms_norm(x[text_idx], text_w) + ir.ops.rms_norm(x[vae_idx], vae_w))
+  - Target:    1x MoTRMSNorm fused Triton kernel
+    (mot_norm(x, text_indices, vae_indices))
+
+Usage::
+    pytest tests/diffusion/kernels/mot/test_mot_norm.py -v -s
+"""
+
+from __future__ import annotations
+
+import time
+
+import pytest
+import torch
+from vllm import ir
+from vllm.config import VllmConfig, set_current_vllm_config
+
+from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+
+pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu]
+
+# ---------------------------------------------------------------------------
+# Constants
+# ---------------------------------------------------------------------------
+
+_EPS = 1e-6
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_indices(
+    M: int,
+    text_ratio: float,
+    device: str = "cuda",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    M_text = max(1, int(M * text_ratio))
+    perm = torch.randperm(M, device=device)
+    text_indices = perm[:M_text].sort().values
+    vae_indices = perm[M_text:].sort().values
+    return text_indices, vae_indices
+
+
+def _benchmark(fn, warmup: int = 20, iters: int = 200) -> float:
+    """Return mean latency in milliseconds."""
+    for _ in range(warmup):
+        fn()
+    torch.accelerator.synchronize()
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        fn()
+    torch.accelerator.synchronize()
+    return (time.perf_counter() - t0) / iters * 1000.0
+
+
+def _reference_forward(
+    x: torch.Tensor,
+    text_indices: torch.Tensor,
+    vae_indices: torch.Tensor,
+    text_weight: torch.Tensor,
+    vae_weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Reference: index-gather -> 2x vLLM IR RMSNorm op -> scatter."""
+    output = torch.empty_like(x)
+    output[text_indices] = ir.ops.rms_norm(x[text_indices], text_weight, eps)
+    output[vae_indices] = ir.ops.rms_norm(x[vae_indices], vae_weight, eps)
+    return output
+
+
+def _reference_forward_head_norm(
+    x: torch.Tensor,
+    text_indices: torch.Tensor,
+    vae_indices: torch.Tensor,
+    text_weight: torch.Tensor,
+    vae_weight: torch.Tensor,
+    eps: float,
+) -> torch.Tensor:
+    """Reference for head_norm=True path.
+
+    The MoT routing happens on token dimension while RMSNorm is applied
+    independently on each head's last dimension.
+    """
+    output = torch.empty_like(x)
+    hidden_size = x.shape[-1]
+    output[text_indices] = ir.ops.rms_norm(
+        x[text_indices].reshape(-1, hidden_size),
+        text_weight,
+        eps,
+    ).reshape_as(x[text_indices])
+    output[vae_indices] = ir.ops.rms_norm(
+        x[vae_indices].reshape(-1, hidden_size),
+        vae_weight,
+        eps,
+    ).reshape_as(x[vae_indices])
+    return output
+
+
+def _check_and_report(ref: torch.Tensor, mot: torch.Tensor, tag: str):
+    """Compare outputs, print metrics, and assert correctness.
+
+    Both ``ref`` and ``mot`` are in the original compute dtype (e.g. bf16).
+    We upcast to fp32 solely for computing error metrics with higher
+    arithmetic precision — the actual layer outputs remain bf16.
+    """
+    ref_hp = ref.float()
+    mot_hp = mot.float()
+
+    abs_err = (ref_hp - mot_hp).abs()
+    max_abs = abs_err.max().item()
+
+    denom = ref_hp.abs().clamp(min=1.0)
+    max_rel = (abs_err / denom).max().item()
+
+    cos_sim = (
+        torch.nn.functional.cosine_similarity(
+            ref_hp,
+            mot_hp,
+            dim=-1,
+        )
+        .min()
+        .item()
+    )
+
+    print(f"\n  [{tag}]  max_abs={max_abs:.4e}  max_rel={max_rel:.4e}  min_cos_sim={cos_sim:.6f}")
+
+    # RMSNorm is element-wise (no cross-element accumulation like GEMM),
+    # so the error between two fp32-accumulating implementations should
+    # be very small — well within bf16 rounding.
+    assert cos_sim > 0.99, f"Cosine similarity too low: {cos_sim:.6f}"
+    assert max_rel < 0.05, f"Max relative error too large: {max_rel:.4e}"
+
+
+def _run_timing(
+    ref_fn,
+    mot_fn,
+    tag: str,
+    warmup: int = 20,
+    iters: int = 200,
+):
+    """Benchmark both paths and print timing comparison."""
+    ref_ms = _benchmark(ref_fn, warmup=warmup, iters=iters)
+    mot_ms = _benchmark(mot_fn, warmup=warmup, iters=iters)
+    speedup = ref_ms / mot_ms if mot_ms > 0 else float("inf")
+    print(f"  [{tag}]  Ref(2x rms_norm): {ref_ms:.3f} ms  |  MoT(fused): {mot_ms:.3f} ms  |  Speedup: {speedup:.2f}x")
+
+
+# =========================================================================
+#  Test: MoTRMSNorm vs 2x vLLM rms_norm
+# =========================================================================
+
+
+@pytest.mark.parametrize(
+    "M, hidden_size, text_ratio, dtype",
+    [
+        (2048, 3584, 0.01, torch.bfloat16),
+        (8192, 3584, 0.01, torch.bfloat16),
+        (2048, 128, 0.01, torch.bfloat16),
+        (8192, 128, 0.01, torch.bfloat16),
+        (2048, 3584, 0.01, torch.float16),
+        (8192, 3584, 0.01, torch.float16),
+    ],
+    ids=[
+        "M2048_H3584_layernorm_bf16",
+        "M8192_H3584_layernorm_bf16",
+        "M2048_H128_qknorm_bf16",
+        "M8192_H128_qknorm_bf16",
+        "M2048_H3584_layernorm_fp16",
+        "M8192_H3584_layernorm_fp16",
+    ],
+)
+def test_mot_rms_norm(M: int, hidden_size: int, text_ratio: float, dtype: torch.dtype):
+    torch.manual_seed(42)
+
+    # --- Build MoT layer ---
+    mot_norm = MoTRMSNorm(hidden_size, eps=_EPS).cuda()
+    with torch.no_grad():
+        W_text = torch.randn(hidden_size, dtype=dtype, device="cuda")
+        W_vae = torch.randn(hidden_size, dtype=dtype, device="cuda")
+        mot_norm.weight.data.copy_(W_text)
+        mot_norm.gen_weight.data.copy_(W_vae)
+
+    # --- Build reference weights (same data) ---
+    ref_text_weight = W_text.clone()
+    ref_vae_weight = W_vae.clone()
+
+    # --- Inputs ---
+    x = torch.randn(M, hidden_size, dtype=dtype, device="cuda")
+    text_idx, vae_idx = _make_indices(M, text_ratio)
+
+    tag = f"RMSNorm M={M} H={hidden_size}"
+
+    # vLLM's IR RMSNorm op may inspect global config; wrap in VllmConfig
+    # context for safety.
+    with set_current_vllm_config(VllmConfig()):
+        # --- Correctness (also warms up Triton JIT) ---
+        with torch.no_grad():
+            ref = _reference_forward(
+                x,
+                text_idx,
+                vae_idx,
+                ref_text_weight,
+                ref_vae_weight,
+                _EPS,
+            )
+            mot_out = mot_norm(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, tag)
+
+        # --- Performance ---
+        with torch.no_grad():
+            _run_timing(
+                lambda: _reference_forward(
+                    x,
+                    text_idx,
+                    vae_idx,
+                    ref_text_weight,
+                    ref_vae_weight,
+                    _EPS,
+                ),
+                lambda: mot_norm(x, text_idx, vae_idx),
+                tag,
+            )
+
+
+@pytest.mark.parametrize(
+    "M, num_heads, head_dim, text_ratio",
+    [
+        (2048, 28, 128, 0.01),
+    ],
+    ids=[
+        "M2048_NH28_HD128_qknorm_head_norm",
+    ],
+)
+def test_mot_rms_norm_head_norm(
+    M: int,
+    num_heads: int,
+    head_dim: int,
+    text_ratio: float,
+):
+    torch.manual_seed(42)
+    dtype = torch.bfloat16
+
+    # --- Build MoT layer (head_norm path) ---
+    mot_norm = MoTRMSNorm(head_dim, head_norm=True, eps=_EPS).cuda()
+    with torch.no_grad():
+        W_text = torch.randn(head_dim, dtype=dtype, device="cuda")
+        W_vae = torch.randn(head_dim, dtype=dtype, device="cuda")
+        mot_norm.weight.data.copy_(W_text)
+        mot_norm.gen_weight.data.copy_(W_vae)
+
+    # --- Build reference weights (same data) ---
+    ref_text_weight = W_text.clone()
+    ref_vae_weight = W_vae.clone()
+
+    # --- Inputs ---
+    x = torch.randn(M, num_heads, head_dim, dtype=dtype, device="cuda")
+    text_idx, vae_idx = _make_indices(M, text_ratio)
+
+    tag = f"RMSNorm(head_norm=True) M={M} NH={num_heads} HD={head_dim}"
+
+    with set_current_vllm_config(VllmConfig()):
+        # --- Correctness (also warms up Triton JIT) ---
+        with torch.no_grad():
+            ref = _reference_forward_head_norm(
+                x,
+                text_idx,
+                vae_idx,
+                ref_text_weight,
+                ref_vae_weight,
+                _EPS,
+            )
+            mot_out = mot_norm(x, text_idx, vae_idx)
+
+        _check_and_report(ref, mot_out, tag)
+
+
+# =========================================================================
+#  Test: und-mode (text_indices=None) — degrades to standard RMSNorm
+# =========================================================================
+
+
+@pytest.mark.parametrize(
+    "M, hidden_size, dtype",
+    [
+        (2048, 3584, torch.bfloat16),
+        (2048, 3584, torch.float16),
+        (2048, 128, torch.bfloat16),
+    ],
+    ids=["M2048_H3584_bf16", "M2048_H3584_fp16", "M2048_H128_bf16"],
+)
+def test_mot_rms_norm_und_mode(M: int, hidden_size: int, dtype: torch.dtype):
+    """und-mode: text_indices=None should apply self.weight to all tokens (standard RMSNorm)."""
+    torch.manual_seed(42)
+
+    mot_norm = MoTRMSNorm(hidden_size, eps=_EPS).cuda()
+    with torch.no_grad():
+        W_text = torch.randn(hidden_size, dtype=dtype, device="cuda")
+        mot_norm.weight.data.copy_(W_text)
+
+    x = torch.randn(M, hidden_size, dtype=dtype, device="cuda")
+
+    with set_current_vllm_config(VllmConfig()):
+        with torch.no_grad():
+            ref = ir.ops.rms_norm(x, W_text, _EPS)
+            mot_out = mot_norm(x, text_indices=None, vae_indices=None)
+
+    _check_and_report(ref, mot_out, f"RMSNorm und-mode M={M} H={hidden_size}")
+
+
+# =========================================================================
+#  Test: boundary cases (all-text, all-VAE)
+# =========================================================================
+
+
+@pytest.mark.parametrize(
+    "boundary_mode",
+    ["all_text", "all_vae"],
+    ids=["all_text", "all_vae"],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"])
+def test_mot_rms_norm_boundary(boundary_mode: str, dtype: torch.dtype):
+    """Boundary: all tokens routed to one expert."""
+    torch.manual_seed(42)
+    M = 1024
+    hidden_size = 3584
+
+    mot_norm = MoTRMSNorm(hidden_size, eps=_EPS).cuda()
+    with torch.no_grad():
+        W_text = torch.randn(hidden_size, dtype=dtype, device="cuda")
+        W_vae = torch.randn(hidden_size, dtype=dtype, device="cuda")
+        mot_norm.weight.data.copy_(W_text)
+        mot_norm.gen_weight.data.copy_(W_vae)
+
+    x = torch.randn(M, hidden_size, dtype=dtype, device="cuda")
+    all_indices = torch.arange(M, dtype=torch.long, device="cuda")
+    empty_indices = torch.empty(0, dtype=torch.long, device="cuda")
+
+    if boundary_mode == "all_text":
+        text_idx, vae_idx = all_indices, empty_indices
+        expected_weight = W_text
+    else:
+        text_idx, vae_idx = empty_indices, all_indices
+        expected_weight = W_vae
+
+    with set_current_vllm_config(VllmConfig()):
+        with torch.no_grad():
+            ref = ir.ops.rms_norm(x, expected_weight, _EPS)
+            mot_out = mot_norm(x, text_idx, vae_idx)
+
+    _check_and_report(ref, mot_out, f"RMSNorm boundary={boundary_mode} M={M}")