fix pre commit

Signed-off-by: princepride <wangzhipeng628@gmail.com>

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -79,6 +79,7 @@ Usage:
 import argparse
 import ast
 import asyncio
+import base64
 import glob
 import json
 import logging
@@ -1013,6 +1014,73 @@ def calculate_metrics(
     return metrics
 
 
+def _save_generated_outputs(
+    outputs: list[RequestFuncOutput],
+    requests_list: list[RequestFuncInput],
+    save_dir: str,
+) -> None:
+    """Decode and save base64 images/videos from successful responses."""
+    os.makedirs(save_dir, exist_ok=True)
+    saved = 0
+    failed = 0
+
+    for idx, (req, out) in enumerate(zip(requests_list, outputs)):
+        if not out.success or not out.response_body:
+            continue
+
+        media_urls: list[str] = []
+
+        # Chat-completions style: choices[*].message.content[*].image_url.url
+        choices = out.response_body.get("choices", [])
+        if isinstance(choices, list):
+            for choice in choices:
+                content = (choice or {}).get("message", {}).get("content")
+                if not isinstance(content, list):
+                    continue
+                for item in content:
+                    if not isinstance(item, dict) or item.get("type") != "image_url":
+                        continue
+                    url = (item.get("image_url") or {}).get("url", "")
+                    if isinstance(url, str) and url.startswith("data:"):
+                        media_urls.append(url)
+
+        # Images endpoint style: data[*].b64_json
+        data_items = out.response_body.get("data", [])
+        if isinstance(data_items, list):
+            for data_item in data_items:
+                if not isinstance(data_item, dict):
+                    continue
+                b64_json = data_item.get("b64_json", "")
+                if isinstance(b64_json, str) and b64_json:
+                    media_urls.append(f"data:image/png;base64,{b64_json}")
+
+        for img_idx, url in enumerate(media_urls):
+            if "," not in url:
+                continue
+
+            try:
+                header, b64_data = url.split(",", 1)
+                ext = "png"
+                if "image/jpeg" in header:
+                    ext = "jpg"
+                elif "image/webp" in header:
+                    ext = "webp"
+                elif "video/mp4" in header:
+                    ext = "mp4"
+
+                img_bytes = base64.b64decode(b64_data)
+                fname = f"req_{idx:04d}_{img_idx}.{ext}"
+                fpath = os.path.join(save_dir, fname)
+                with open(fpath, "wb") as f:
+                    f.write(img_bytes)
+                saved += 1
+            except Exception as e:
+                failed += 1
+                logger.warning(f"Failed to save image for request {idx}: {e}", exc_info=True)
+
+    logger.info(f"Saved {saved} generated image(s) to {save_dir}. Failed to save {failed} image(s).")
+
+
 def wait_for_service(base_url: str, timeout: int = 120) -> None:
     print(f"Waiting for service at {base_url}...")
     start_time = time.time()
@@ -1202,6 +1270,9 @@ async def benchmark(args):
 
     print("\n" + "=" * 60)
 
+    if args.save_dir:
+        _save_generated_outputs(outputs, requests_list, args.save_dir)
+
     if args.output_file:
         with open(args.output_file, "w") as f:
             json.dump(metrics, f, indent=2)
@@ -1278,11 +1349,18 @@ if __name__ == "__main__":
         default=1,
         help="Number of warmup requests to run before measurement.",
     )
+    # NOTE Changed default from 1 to 2 because some models (e.g., Bagel) run
+    # `num_timesteps - 1` denoising iterations. A default of 1 results in 0 steps,
+    # which causes errors.
+    # TODO If this slightly longer warmup causes regression issues for other
+    # diffusion pipelines in the future, consider implementing model-specific
+    # overrides instead of a global default.
     parser.add_argument(
         "--warmup-num-inference-steps",
         type=int,
-        default=1,
-        help="num_inference_steps used for warmup requests.",
+        default=2,
+        help="Number of inference steps used for warmup requests. "
+        "Default is 2 to ensure at least one denoising step is executed.",
     )
     parser.add_argument(
         "--warmup-concurrency",
@@ -1323,6 +1401,13 @@ if __name__ == "__main__":
         default=3.0,
         help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
     )
+    parser.add_argument(
+        "--save-dir",
+        type=str,
+        default=None,
+        help="Directory to save generated images/outputs for visual inspection. "
+        "If not set, generated outputs are discarded after metric collection.",
+    )
     parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
     parser.add_argument(
         "--enable-negative-prompt",

docs/user_guide/diffusion/mot_config.md

@@ -0,0 +1,50 @@
+This directory contains auto-tuned Triton kernel configurations for the
+MoT (Mixture-of-Tokens) GEMM and RMSNorm operators used by BAGEL and other
+MoT-architecture diffusion models.
+
+File naming convention:
+    device_name=<GPU_NAME>,dtype=<DTYPE>.json
+
+For example:
+    device_name=NVIDIA_A100-SXM4-80GB,dtype=w16a16.json
+
+Each JSON file maps (K, N) matrix shapes to a dictionary of batch sizes (M)
+and their optimal Triton tile configurations:
+
+    {
+        "3584_4608": {              // K=3584, N=4608 (QKV projection)
+            "1024": {
+                "BLOCK_SIZE_M": 64,
+                "BLOCK_SIZE_N": 128,
+                "BLOCK_SIZE_K": 64,
+                "GROUP_SIZE_M": 8,
+                "num_warps": 4,
+                "num_stages": 3
+            },
+            ...
+        },
+        "3584_3584": { ... }        // K=3584, N=3584 (output projection)
+    }
+
+Config loading order (3-tier, see ops/mot_gemm.py):
+    1. $VLLM_TUNED_CONFIG_FOLDER/<filename>  (env override)
+    2. This directory: vllm_omni/diffusion/layers/mot/configs/<filename>
+    3. Conservative default config (compiles everywhere, sub-optimal perf)
+
+If no config file matches the current device, a warning is printed with
+instructions to run the auto-tuning benchmark.
+
+To generate configs for your hardware:
+
+    python benchmarks/kernels/mot_linear_benchmarks.py \
+        --model ByteDance-Seed/BAGEL-7B-MoT \
+        --tp-size 1 --dtype w16a16 --tune \
+        --save-dir vllm_omni/diffusion/layers/mot/configs/
+
+For multi-GPU tuning (uses Ray for parallel search):
+
+    python benchmarks/kernels/mot_linear_benchmarks.py \
+        --model ByteDance-Seed/BAGEL-7B-MoT \
+        --tp-size 2 --tune
+
+See benchmarks/kernels/mot_linear_benchmarks.py for full options.

tests/diffusion/kernels/mot/test_mot_linear.py

@@ -0,0 +1,459 @@
+# 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],
+    ids=[f"img{num}_K3584_N4608" 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],
+    ids=[f"img{num}_K3584_N3584" 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,
+            )

tests/diffusion/kernels/mot/test_mot_norm.py

@@ -0,0 +1,278 @@
+# 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",
+    [
+        (2048, 3584, 0.01),
+        (8192, 3584, 0.01),
+        (2048, 128, 0.01),
+        (8192, 128, 0.01),
+    ],
+    ids=[
+        "M2048_H3584_layernorm",
+        "M8192_H3584_layernorm",
+        "M2048_H128_qknorm",
+        "M8192_H128_qknorm",
+    ],
+)
+def test_mot_rms_norm(M: int, hidden_size: int, text_ratio: float):
+    torch.manual_seed(42)
+    dtype = torch.bfloat16
+
+    # --- 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)

vllm_omni/diffusion/layers/mot/mot_layernorm.py

@@ -0,0 +1,96 @@
+# ruff: noqa: N803, E741
+"""Mixture-of-Tokens (MoT) RMS Normalization layer.
+
+Holds two sets of weights (text / gen) and routes tokens to the
+appropriate weight based on indices.  When text_indices is None the
+layer degrades to a standard RMSNorm using self.weight (und mode).
+"""
+
+import torch
+import torch.nn as nn
+from vllm import ir
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+
+class MoTRMSNorm(CustomOp):
+    """Mixture-of-Tokens RMS Normalization.
+
+    In *und* mode (``text_indices is None``), every token is normalised
+    with ``self.weight`` – exactly like a vanilla RMSNorm.
+
+    In *gen* mode, text tokens are normalised with ``self.weight`` and
+    gen tokens are normalised with ``self.gen_weight``, using a single
+    fused Triton kernel that avoids the gather / scatter overhead.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        head_norm: bool = False,
+        eps: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.variance_epsilon = eps
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.gen_weight = nn.Parameter(torch.ones(hidden_size))
+        self.head_norm = head_norm
+
+    # ------------------------------------------------------------------
+    # Native (pure-PyTorch) fallback
+    # ------------------------------------------------------------------
+    def forward_native(
+        self,
+        x: torch.Tensor,
+        text_indices: torch.Tensor | None = None,
+        vae_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if text_indices is None:
+            return self._rms_norm_native(x, self.weight)
+
+        output = torch.empty_like(x)
+        output[text_indices] = self._rms_norm_native(x[text_indices], self.weight)
+        output[vae_indices] = self._rms_norm_native(x[vae_indices], self.gen_weight)
+        return output
+
+    # ------------------------------------------------------------------
+    # CUDA fast-path
+    # ------------------------------------------------------------------
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+        text_indices: torch.Tensor | None = None,
+        vae_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if text_indices is None:
+            # und mode – delegate to vllm's highly-optimised CUDA kernel
+            return ir.ops.rms_norm(x, self.weight.data, self.variance_epsilon)
+
+        # gen mode – fused MoT Triton kernel
+        from vllm_omni.diffusion.layers.mot.ops.mot_rms_norm import (
+            mot_rms_norm,
+        )
+
+        return mot_rms_norm(
+            x,
+            self.weight.data,
+            self.gen_weight.data,
+            text_indices,
+            vae_indices,
+            head_norm=self.head_norm,
+            eps=self.variance_epsilon,
+        )
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+    def _rms_norm_native(self, x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+        orig_dtype = x.dtype
+        x = x.float()
+        variance = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.variance_epsilon)
+        return (x * weight.float()).to(orig_dtype)
+
+    def extra_repr(self) -> str:
+        return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}"

vllm_omni/diffusion/layers/mot/mot_qkv_parallel_linear.py

@@ -0,0 +1,341 @@
+from __future__ import annotations
+
+import torch
+from torch.nn.parameter import Parameter
+from vllm.distributed import (
+    tensor_model_parallel_all_gather,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import (
+    WEIGHT_LOADER_V2_SUPPORTED,
+    QKVParallelLinear,
+)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+)
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform
+
+from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm
+
+logger = init_logger(__name__)
+
+
+class MoTQKVParallelLinear(QKVParallelLinear):
+    """QKVParallelLinear with Mixture-of-Tokens routing.
+
+    Text weights: stored directly on self (self.weight, self.weight_scale, ...),
+              created through the standard QKVParallelLinear.__init__ process.
+
+    VAE weights: stored in the permanent submodule self.gen_exp
+             (self.gen_exp.weight, ...),
+             created via quant_method.create_weights(self.gen_exp, ...).
+             gen_exp.quant_method points to the same quant_method, so that
+             the vLLM framework’s process_weights_after_loading
+             can automatically detect and process it.
+
+    Forward behavior:
+        - und mode (text_indices is None): fully reuse super().forward()
+        - gen mode: call the MoT fused GEMM kernel
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        head_size: int,
+        total_num_heads: int,
+        total_num_kv_heads: int | None = None,
+        bias: bool = True,
+        vae_bias: bool = False,
+        skip_bias_add: bool = False,
+        params_dtype: torch.dtype | None = None,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        *,
+        return_bias: bool = True,
+        disable_tp: bool = False,
+        v_head_size: int | None = None,
+    ):
+        # ---- 1) Parent class creates text weights ----
+        # super().__init__  will do the following:
+        # quant_method.create_weights(self, ...) → self.weight, self.weight_scale, ...
+        # QKVParallelLinear hardcodes gather_output=False
+        super().__init__(
+            hidden_size,
+            head_size,
+            total_num_heads,
+            total_num_kv_heads,
+            bias,
+            skip_bias_add,
+            params_dtype,
+            quant_config,
+            prefix,
+            return_bias=return_bias,
+            disable_tp=disable_tp,
+            v_head_size=v_head_size,
+        )
+
+        # ---- 2) Create vae weights (permanent submodule) ----
+        assert self.quant_method is not None
+
+        # NOTE: We instantiate a bare torch.nn.Module() here as a lightweight namespace
+        # container to hold the secondary VAE weights.
+        # This is a design choice for dynamically attaching `quant_method` to the bare module.
+        # This is slightly unconventional but necessary. It allows vLLM's framework
+        # (specifically `process_weights_after_loading`) to automatically discover
+        # and process these secondary weights without requiring a dedicated,
+        # boilerplate Module subclass just for the VAE pathway.
+        self.gen_exp = torch.nn.Module()
+
+        # Use the same weight_loader as text
+        vae_weight_loader = (
+            self.weight_loader_v2
+            if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
+            else self.weight_loader
+        )
+        self.quant_method.create_weights(
+            layer=self.gen_exp,
+            input_size_per_partition=self.input_size_per_partition,
+            output_partition_sizes=self.output_partition_sizes,
+            input_size=self.input_size,
+            output_size=self.output_size,
+            params_dtype=self.params_dtype,
+            weight_loader=vae_weight_loader,
+        )
+
+        # Make gen_exp discoverable by vLLM framework's process_weights_after_loading
+        self.gen_exp.quant_method = self.quant_method
+
+        # ---- 3) vae bias ----
+        if vae_bias:
+            self.gen_exp.bias = Parameter(torch.empty(self.output_size_per_partition, dtype=self.params_dtype))
+            set_weight_attrs(
+                self.gen_exp.bias,
+                {"output_dim": 0, "weight_loader": self.weight_loader},
+            )
+        else:
+            self.gen_exp.register_parameter("bias", None)
+
+        self.update_param_tp_status()
+
+    # ==================================================================
+    #  Forward
+    # ==================================================================
+    def forward(
+        self,
+        input_: torch.Tensor,
+        text_indices: torch.Tensor | None = None,
+        vae_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        # ---- und mode: fully reuse parent class (only text path) ----
+        if text_indices is None:
+            return super().forward(input_)
+
+        # ---- gen mode: fuse MoT GEMM ----
+        output_parallel = self._mot_gemm_dispatch(input_, text_indices, vae_indices)
+
+        # QKVParallelLinear hardcodes gather_output=False, this branch never enters;
+        # retained for future subclass changes
+        if self.gather_output and self.tp_size > 1:
+            output = tensor_model_parallel_all_gather(output_parallel)
+        else:
+            output = output_parallel
+
+        if not self.return_bias:
+            return output
+
+        if self.skip_bias_add and text_indices is not None and (self.bias is not None or self.gen_exp.bias is not None):
+            merged_bias = torch.zeros(
+                output.size(0),
+                self.output_size_per_partition,
+                dtype=output.dtype,
+                device=output.device,
+            )
+            if self.bias is not None:
+                merged_bias[text_indices] = self.bias
+            if self.gen_exp.bias is not None:
+                merged_bias[vae_indices] = self.gen_exp.bias
+            return output, merged_bias
+
+        output_bias = self.bias if self.skip_bias_add else None
+        return output, output_bias
+
+    # ==================================================================
+    #  MoT GEMM dispatcher
+    # ==================================================================
+    def _mot_gemm_dispatch(
+        self,
+        x: torch.Tensor,
+        text_indices: torch.Tensor,
+        vae_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        # Any other backend (ROCm, XPU, TPU, CPU) uses the safe vllm fallback.
+        if not current_platform.is_cuda():
+            return self._mot_fallback(x, text_indices, vae_indices)
+
+        """Dispatch to different MoT kernel paths based on weight dtype / quant attributes."""
+        bias_text = self.bias if not self.skip_bias_add else None
+        bias_vae = self.gen_exp.bias if not self.skip_bias_add else None
+
+        w_text = self.weight
+        w_vae = self.gen_exp.weight
+        assert w_text.dtype.is_floating_point == w_vae.dtype.is_floating_point, (
+            "weight of text expert and image expert should be the same dtype."
+        )
+
+        # w_text.dtype.itemsize >= 2 means bytes_per_element >= 2 (16bits or 32bits)
+        if w_text.dtype.is_floating_point and w_text.dtype.itemsize >= 2:
+            return self._mot_gemm_unquantized(
+                x,
+                text_indices,
+                vae_indices,
+                bias_text,
+                bias_vae,
+            )
+        elif w_text.dtype == torch.float8_e4m3fn:
+            return self._mot_gemm_fp8_w8a8(
+                x,
+                text_indices,
+                vae_indices,
+                bias_text,
+                bias_vae,
+            )
+        elif w_text.dtype == torch.int8:
+            return self._mot_gemm_weight_only(
+                x,
+                text_indices,
+                vae_indices,
+                bias_text,
+                bias_vae,
+            )
+        else:
+            return self._mot_fallback(
+                x,
+                text_indices,
+                vae_indices,
+            )
+
+    # ==================================================================
+    #  Implementation of each quantization path
+    # ==================================================================
+    def _mot_gemm_unquantized(self, x, text_idx, vae_idx, bias_t, bias_v):
+        """BF16/FP16 path."""
+        N = self.output_size_per_partition
+        C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device)
+        invoke_mot_gemm(
+            A=x,
+            B_text=self.weight.data.t(),
+            B_vae=self.gen_exp.weight.data.t(),
+            C=C,
+            bias_text=bias_t,
+            bias_vae=bias_v,
+            text_indices=text_idx,
+            vae_indices=vae_idx,
+            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,
+        )
+        return C
+
+    def _mot_gemm_fp8_w8a8(self, x, text_idx, vae_idx, bias_t, bias_v):
+        """FP8 W8A8 path.
+
+        1) Activation quantization: reuse vllm's ops.scaled_fp8_quant
+        2) MoT GEMM: text/vae use different fp8 weights and weight_scale
+        3) De-quantization: done by MoT kernel's epilogue
+        """
+        from vllm import _custom_ops as ops
+
+        x_2d = x.view(-1, x.shape[-1])
+        input_scale = getattr(self, "input_scale", None)
+        x_fp8, x_scale = ops.scaled_fp8_quant(
+            x_2d,
+            input_scale,
+            use_per_token_if_dynamic=True,
+        )
+
+        N = self.output_size_per_partition
+        C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device)
+
+        invoke_mot_gemm(
+            A=x_fp8,
+            B_text=self.weight.data,
+            B_vae=self.gen_exp.weight.data,
+            C=C,
+            bias_text=bias_t,
+            bias_vae=bias_v,
+            text_indices=text_idx,
+            vae_indices=vae_idx,
+            A_scale=x_scale,
+            B_text_scale=self.weight_scale.data,
+            B_vae_scale=self.gen_exp.weight_scale.data,
+            use_fp8_w8a8=True,
+            use_int8_w8a8=False,
+            use_int8_w8a16=False,
+            use_int4_w4a16=False,
+            A_per_channel_quant=True,
+            B_per_channel_quant=False,
+        )
+        return C
+
+    def _mot_gemm_weight_only(self, x, text_idx, vae_idx, bias_t, bias_v):
+        """Weight-Only W8A16 path.
+
+        Activation values kept as bf16/fp16, weights are int8 + per-channel scale.
+        De-quantization done by MoT kernel's epilogue.
+        """
+        N = self.output_size_per_partition
+        C = torch.zeros(x.size(0), N, dtype=x.dtype, device=x.device)
+        invoke_mot_gemm(
+            A=x,
+            B_text=self.weight.data.t(),
+            B_vae=self.gen_exp.weight.data.t(),
+            C=C,
+            bias_text=bias_t,
+            bias_vae=bias_v,
+            text_indices=text_idx,
+            vae_indices=vae_idx,
+            A_scale=None,
+            B_text_scale=self.weight_scale.data,
+            B_vae_scale=self.gen_exp.weight_scale.data,
+            use_fp8_w8a8=False,
+            use_int8_w8a8=False,
+            use_int8_w8a16=True,
+            use_int4_w4a16=False,
+            A_per_channel_quant=False,
+            B_per_channel_quant=True,
+        )
+        return C
+
+    def _mot_fallback(self, x, text_idx, vae_idx):
+        """Fallback: fall back to gather/scatter + quant_method.apply.
+
+        For unsupported quantization types, call standard forward for text/vae tokens.
+        """
+        assert self.quant_method is not None
+
+        bias_text = self.bias if not self.skip_bias_add else None
+        bias_vae = self.gen_exp.bias if not self.skip_bias_add else None
+
+        output = torch.zeros(
+            x.size(0),
+            self.output_size_per_partition,
+            dtype=x.dtype,
+            device=x.device,
+        )
+        output_text = self.quant_method.apply(self, x[text_idx], bias_text)
+
+        output_vae = self.quant_method.apply(
+            self.gen_exp,
+            x[vae_idx],
+            bias_vae,
+        )
+        output[text_idx] = output_text
+        output[vae_idx] = output_vae
+        return output

vllm_omni/diffusion/layers/mot/mot_row_parallel_linear.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+import torch
+from torch.nn.parameter import Parameter
+from vllm.distributed import (
+    split_tensor_along_last_dim,
+    tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.linear import (
+    WEIGHT_LOADER_V2_SUPPORTED,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig,
+)
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform
+
+from vllm_omni.diffusion.layers.mot.ops.mot_gemm import invoke_mot_gemm
+
+
+class MoTRowParallelLinear(RowParallelLinear):
+    """RowParallelLinear with Mixture-of-Tokens routing.
+
+    text weights: directly on self (self.weight, self.weight_scale, ...),
+                   created by RowParallelLinear.__init__ standard process.
+    vae  weights: on permanent submodule self.gen_exp (self.gen_exp.weight, ...),
+                   created by quant_method.create_weights(self.gen_exp, ...).
+                   gen_exp.quant_method points to the same quant_method, enabling
+                   vLLM framework's process_weights_after_loading to automatically
+                   discover and process it.
+
+    Forward behavior:
+        - und mode (text_indices is None): fully reuse super().forward()
+        - gen mode: call MoT fused GEMM kernel, then execute TP all-reduce
+    """
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        bias: bool = True,
+        vae_bias: bool = False,
+        input_is_parallel: bool = True,
+        skip_bias_add: bool = False,
+        params_dtype: torch.dtype | None = None,
+        reduce_results: bool = True,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+        *,
+        return_bias: bool = True,
+        disable_tp: bool = False,
+    ):
+        # ---- Step 1: Parent class creates text weights ----
+        # super().__init__ internally calls:
+        #   quant_method.create_weights(self, ...) → self.weight, self.weight_scale, ...
+        super().__init__(
+            input_size,
+            output_size,
+            bias,
+            input_is_parallel,
+            skip_bias_add,
+            params_dtype,
+            reduce_results,
+            quant_config,
+            prefix,
+            return_bias=return_bias,
+            disable_tp=disable_tp,
+        )
+
+        # ---- Step 2: Create vae weights (permanent submodule) ----
+        assert self.quant_method is not None
+
+        # NOTE: We instantiate a bare torch.nn.Module() here as a lightweight namespace
+        # container to hold the secondary VAE weights.
+        # This is a design choice for dynamically attaching `quant_method` to the bare module.
+        # This is slightly unconventional but necessary. It allows vLLM's framework
+        # (specifically `process_weights_after_loading`) to automatically discover
+        # and process these secondary weights without requiring a dedicated,
+        # boilerplate Module subclass just for the VAE pathway.
+        self.gen_exp = torch.nn.Module()
+
+        # Select weight_loader consistent with text
+        vae_weight_loader = (
+            self.weight_loader_v2
+            if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED
+            else self.weight_loader
+        )
+        self.quant_method.create_weights(
+            layer=self.gen_exp,
+            input_size_per_partition=self.input_size_per_partition,
+            output_partition_sizes=self.output_partition_sizes,
+            input_size=self.input_size,
+            output_size=self.output_size,
+            params_dtype=self.params_dtype,
+            weight_loader=vae_weight_loader,
+        )
+
+        # Enable gen_exp to be automatically discovered by vLLM framework's process_weights_after_loading
+        self.gen_exp.quant_method = self.quant_method
+
+        # ---- Step 3: vae bias ----
+        # RowParallelLinear's bias is full output_size (not sharded)
+        if vae_bias:
+            self.gen_exp.bias = Parameter(torch.empty(self.output_size, dtype=self.params_dtype))
+            set_weight_attrs(
+                self.gen_exp.bias,
+                {"output_dim": 0, "weight_loader": self.weight_loader},
+            )
+        else:
+            self.gen_exp.register_parameter("bias", None)
+
+        self.update_param_tp_status()
+
+    # ==================================================================
+    #  Forward
+    # ==================================================================
+    def forward(
+        self,
+        input_: torch.Tensor,
+        text_indices: torch.Tensor | None = None,
+        vae_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, Parameter | None]:
+        # ---- und mode: fully reuse parent class (only text path) ----
+        if text_indices is None:
+            return super().forward(input_)
+
+        # ---- gen mode ----
+        # Handle input_is_parallel (consistent with parent class logic)
+        if self.input_is_parallel:
+            input_parallel = input_
+        else:
+            split_input = split_tensor_along_last_dim(input_, num_partitions=self.tp_size)
+            input_parallel = split_input[self.tp_rank].contiguous()
+
+        # Fused MoT GEMM
+        output_parallel = self._mot_gemm_dispatch(input_parallel, text_indices, vae_indices)
+
+        # ---- TP communication: all-reduce (consistent with parent class) ----
+        if self.reduce_results and self.tp_size > 1:
+            output = tensor_model_parallel_all_reduce(output_parallel)
+        else:
+            output = output_parallel
+
+        if not self.return_bias:
+            return output
+
+        if self.skip_bias_add and text_indices is not None and (self.bias is not None or self.gen_exp.bias is not None):
+            # Construct per-token mixed bias
+            merged_bias = torch.zeros(
+                output.size(0),
+                self.output_size_per_partition,
+                dtype=output.dtype,
+                device=output.device,
+            )
+            if self.bias is not None:
+                merged_bias[text_indices] = self.bias
+            if self.gen_exp.bias is not None:
+                merged_bias[vae_indices] = self.gen_exp.bias
+            return output, merged_bias
+
+        output_bias = self.bias if self.skip_bias_add else None
+        return output, output_bias
+
+    # ==================================================================
+    #  MoT GEMM Dispatcher
+    # ==================================================================
+    def _mot_gemm_dispatch(
+        self,
+        x: torch.Tensor,
+        text_indices: torch.Tensor,
+        vae_indices: torch.Tensor,
+    ) -> torch.Tensor:
+        """Dispatch to different MoT kernel paths based on weight dtype / quant attributes."""
+
+        # Any other backend (ROCm, XPU, TPU, CPU) uses the safe vllm fallback.
+        if not current_platform.is_cuda():
+            return self._mot_fallback(x, text_indices, vae_indices)
+
+        # RowParallelLinear: bias only fused into GEMM at rank 0,
+        # other ranks pass None to avoid duplicate accumulation
+        bias_text = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+        bias_vae = None if (self.tp_rank > 0 or self.skip_bias_add) else self.gen_exp.bias
+
+        # Determine quantization type by weight dtype, avoid isinstance coupling to specific quant_method
+        # TODO: Currently does not support online quantization fp8, does not support quantization types involving int4
+        w_text = self.weight
+        w_vae = self.gen_exp.weight
+        assert w_text.dtype.is_floating_point == w_vae.dtype.is_floating_point, (
+            "weight of text expert and image expert should be the same dtype."
+        )
+
+        # w_text.dtype.itemsize >= 2 means bytes_per_element >= 2 (16bits or 32bits)
+        if w_text.dtype.is_floating_point and w_text.dtype.itemsize >= 2:
+            # ---- Path 0: BF16 / FP16 (unquantized) ----
+            return self._mot_gemm_unquantized(
+                x,
+                text_indices,
+                vae_indices,
+                bias_text,
+                bias_vae,
+            )
+        elif w_text.dtype == torch.float8_e4m3fn:
+            # ---- Path 1: FP8 W8A8 ----
+            return self._mot_gemm_fp8_w8a8(
+                x,
+                text_indices,
+                vae_indices,
+                bias_text,
+                bias_vae,
+            )
+        elif w_text.dtype == torch.int8:
+            # ---- Path 2: Weight-Only INT8 W8A16 ----
+            return self._mot_gemm_weight_only(
+                x,
+                text_indices,
+                vae_indices,
+                bias_text,
+                bias_vae,
+            )
+        else:
+            # ---- Fallback: gather/scatter + quant_method.apply ----
+            return self._mot_fallback(
+                x,
+                text_indices,
+                vae_indices,
+            )
+
+    # ==================================================================
+    #  Specific implementations for each quantization path
+    # ==================================================================
+    def _mot_gemm_unquantized(self, x, text_idx, vae_idx, bias_t, bias_v):
+        """BF16/FP16 path."""
+        N = self.output_size_per_partition
+        C = torch.empty(x.size(0), N, dtype=x.dtype, device=x.device)
+        invoke_mot_gemm(
+            A=x,
+            B_text=self.weight.data.t(),
+            B_vae=self.gen_exp.weight.data.t(),
+            C=C,
+            bias_text=bias_t,
+            bias_vae=bias_v,
+            text_indices=text_idx,
+            vae_indices=vae_idx,
+            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,
+        )
+        return C
+
+    def _mot_gemm_fp8_w8a8(self, x, text_idx, vae_idx, bias_t, bias_v):
+        """FP8 W8A8 path.
+
+        1) Activation quantization: reuse vllm's ops.scaled_fp8_quant
+        2) MoT GEMM: text/vae each use different fp8 weights and weight_scale
+        3) Dequantization: completed by MoT kernel internal epilogue
+        """
+        from vllm import _custom_ops as ops
+
+        x_2d = x.view(-1, x.shape[-1])
+        input_scale = getattr(self, "input_scale", None)
+        x_fp8, x_scale = ops.scaled_fp8_quant(
+            x_2d,
+            input_scale,
+            use_per_token_if_dynamic=True,
+        )
+
+        N = self.output_size_per_partition
+        C = torch.empty(x.size(0), N, dtype=x.dtype, device=x.device)
+
+        # weight has been transposed to (K, N) in process_weights_after_loading
+        invoke_mot_gemm(
+            A=x_fp8,
+            B_text=self.weight.data,
+            B_vae=self.gen_exp.weight.data,
+            C=C,
+            bias_text=bias_t,
+            bias_vae=bias_v,
+            text_indices=text_idx,
+            vae_indices=vae_idx,
+            A_scale=x_scale,
+            B_text_scale=self.weight_scale.data,
+            B_vae_scale=self.gen_exp.weight_scale.data,
+            use_fp8_w8a8=True,
+            use_int8_w8a8=False,
+            use_int8_w8a16=False,
+            use_int4_w4a16=False,
+            A_per_channel_quant=True,
+            B_per_channel_quant=False,
+        )
+        return C
+
+    def _mot_gemm_weight_only(self, x, text_idx, vae_idx, bias_t, bias_v):
+        """Weight-Only W8A16 path.
+
+        Activations remain bf16/fp16, weights are int8 + per-channel scale.
+        MoT kernel internally performs immediate dequantization.
+        """
+        N = self.output_size_per_partition
+        C = torch.empty(x.size(0), N, dtype=x.dtype, device=x.device)
+        invoke_mot_gemm(
+            A=x,
+            B_text=self.weight.data.t(),
+            B_vae=self.gen_exp.weight.data.t(),
+            C=C,
+            bias_text=bias_t,
+            bias_vae=bias_v,
+            text_indices=text_idx,
+            vae_indices=vae_idx,
+            A_scale=None,
+            B_text_scale=self.weight_scale.data,
+            B_vae_scale=self.gen_exp.weight_scale.data,
+            use_fp8_w8a8=False,
+            use_int8_w8a8=False,
+            use_int8_w8a16=True,
+            use_int4_w4a16=False,
+            A_per_channel_quant=False,
+            B_per_channel_quant=True,
+        )
+        return C
+
+    def _mot_fallback(self, x, text_idx, vae_idx):
+        """Fallback: degrade to gather/scatter + quant_method.apply.
+
+        For unsupported quantization types, call standard forward separately for text/vae tokens.
+        """
+        assert self.quant_method is not None
+
+        # RowParallelLinear: bias only fused at rank 0
+        bias_text = None if (self.tp_rank > 0 or self.skip_bias_add) else self.bias
+        bias_vae = None if (self.tp_rank > 0 or self.skip_bias_add) else self.gen_exp.bias
+
+        output = torch.empty(
+            x.size(0),
+            self.output_size_per_partition,
+            dtype=x.dtype,
+            device=x.device,
+        )
+        # text tokens → standard quant_method (operate on weights on self)
+        output[text_idx] = self.quant_method.apply(self, x[text_idx], bias_text)
+        # vae tokens → same quant_method (operate on weights on gen_exp)
+        output[vae_idx] = self.quant_method.apply(
+            self.gen_exp,
+            x[vae_idx],
+            bias_vae,
+        )
+        return output

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

@@ -0,0 +1,908 @@
+# ruff: noqa: N803, E741
+import functools
+import json
+import logging
+import os
+import pathlib
+import re
+from typing import Any
+
+import torch
+from vllm.triton_utils import tl, triton
+
+logger = logging.getLogger(__name__)
+
+# =====================================================================
+#  MoT GEMM Config Loading (3-tier: env → built-in → default)
+#
+#  Usage pattern (mirrors vLLM fused_moe):
+#    1. Layer invokes invoke_mot_gemm(...) without extra config context.
+#    2. invoke_mot_gemm lazily calls
+#         get_mot_configs(K, N, dtype_str)
+#       which is @lru_cache'd and lazily loads the JSON on first hit.
+#    3. If get_mot_configs returns None, fall back to
+#         get_mot_default_config(M, N, K, ...)
+# =====================================================================
+
+_CONFIGS_DIR = pathlib.Path(__file__).resolve().parent.parent / "configs"
+_ENV_CONFIG_FOLDER = "VLLM_TUNED_CONFIG_FOLDER"
+
+
+def get_device_name() -> str:
+    """Sanitized GPU device name, matching mot_linear_benchmarks.py output."""
+    raw = torch.cuda.get_device_name(0)
+    name = re.sub(r"[^a-zA-Z0-9]", "", raw.replace(" ", ""))
+    for prefix in ("NVIDIA", "AMD"):
+        if name.startswith(prefix):
+            name = name[len(prefix) :]
+    # --- Device Aliasing Patch ---
+    alias_map = {
+        "A800": "A100",
+        "H800": "H100",
+    }
+    # map A800/H800 to A100/H100 for chinese market
+    for key, target in alias_map.items():
+        if key in name:
+            name = name.replace(key, target)
+            break
+    return name
+
+
+def build_config_filename(device_name: str, dtype_str: str) -> str:
+    return f"device_name={device_name},dtype={dtype_str}.json"
+
+
+def _try_load_json(filepath: str) -> dict | None:
+    if os.path.isfile(filepath):
+        with open(filepath) as f:
+            return json.load(f)
+    return None
+
+
+@functools.lru_cache
+def _load_mot_config_file(dtype_str: str) -> dict | None:
+    """Load and cache the full MoT config JSON (one file per device/dtype).
+
+    Search order:
+      1. ``$VLLM_TUNED_CONFIG_FOLDER/device_name=...,dtype=....json``
+      2. ``vllm_omni/.../mot/configs/device_name=...,dtype=....json``
+      3. Return ``None`` (caller falls back to ``get_mot_default_config``).
+    """
+    device_name = get_device_name()
+    filename = build_config_filename(device_name, dtype_str)
+
+    config_file_paths: list[str] = []
+    env_dir = os.environ.get(_ENV_CONFIG_FOLDER)
+    if env_dir:
+        config_file_paths.append(str(pathlib.Path(env_dir) / filename))
+    config_file_paths.append(str(_CONFIGS_DIR / filename))
+
+    for path in config_file_paths:
+        data = _try_load_json(path)
+        if data is not None:
+            logger.info("MoT config loaded from %s", path)
+            return data
+
+    logger.warning(
+        f"\n{'=' * 80}\n"
+        f" ⚠️  [WARNING] No tuned MoT config found.\n"
+        f" Searched paths: {', '.join(config_file_paths)}\n"
+        f" Using conservative default configs which are NOT optimal.\n"
+        f" Run `python benchmarks/kernels/mot_linear_benchmarks.py --tune` \n"
+        f" to generate hardware-specific optimal configs.\n"
+        f"{'=' * 80}\n"
+    )
+    return None
+
+
+@functools.lru_cache
+def get_mot_configs(
+    K: int,
+    N: int,
+    dtype_str: str | None = None,
+) -> dict[int, dict[str, int]] | None:
+    """Return ``{M: tile_config}`` for a given (K, N) shape, or ``None``.
+
+    The return value maps an irregular grid of batch sizes (M) to Triton
+    tile configurations.  The caller should pick the entry whose M is
+    closest to the actual batch size.
+
+    Config file is selected by ``device_name + dtype``.
+    """
+    file_data = _load_mot_config_file(dtype_str or "w16a16")
+    if file_data is None:
+        return None
+
+    shape_entry = file_data.get(f"{K}_{N}")
+    if shape_entry is None:
+        logger.warning(
+            f"\n{'=' * 80}\n"
+            f" ⚠️  [WARNING] MoT config file found, but NO tuned entry for shape K={K}, N={N}.\n"
+            f" Using conservative default configs which are NOT optimal for this specific shape.\n"
+            f" Run `python benchmarks/kernels/mot_linear_benchmarks.py --tune` \n"
+            f" to generate hardware-specific optimal configs.\n"
+            f"{'=' * 80}\n"
+        )
+        return None
+
+    return {int(k): dict(v) for k, v in shape_entry.items() if k != "_comment"}
+
+
+def get_mot_default_config(
+    M: int,
+    N: int,
+    K: int,
+    dtype: str | None = None,
+    block_quant_shape: list[int] | None = None,
+) -> dict[str, int]:
+    """Conservative fallback config guaranteed to compile on all hardware.
+
+    Trades peak performance for universal compatibility (T4 / V100 / A100 /
+    H100, CUDA & ROCm).
+    """
+    # FP8 block-wise quantization requires strict alignment
+    if dtype == "fp8_w8a8" and block_quant_shape is not None:
+        return {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": block_quant_shape[0],
+            "BLOCK_SIZE_K": block_quant_shape[1],
+            "GROUP_SIZE_M": 8,
+            "num_warps": 4,
+            "num_stages": 2,
+        }
+
+    # Very small M (tail batches, final feature concat, etc.)
+    # NOTE If M is not a multiple of the block size (e.g., M=3 with BLOCK_SIZE_M=16),
+    # the kernel will allocate a full tile. Correctness is guaranteed by boundary masks.
+    # We tolerate this minor compute/SRAM padding to avoid generating excessive
+    # Triton compiled variants for extremely small sequence lengths.
+    if M <= 64:
+        return {
+            "BLOCK_SIZE_M": 16 if M <= 16 else 32,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 32,
+            "GROUP_SIZE_M": 8,
+            "num_warps": 4,
+            "num_stages": 2,
+        }
+
+    # Standard fallback for typical image-generation M (2048 / 4096 / …)
+    # SRAM usage: (64*32 + 64*32) * 2 * 2 = 16 KB — safe on decade-old GPUs
+    return {
+        "BLOCK_SIZE_M": 64,
+        "BLOCK_SIZE_N": 64,
+        "BLOCK_SIZE_K": 32,
+        "GROUP_SIZE_M": 8,
+        "num_warps": 4,
+        "num_stages": 2,
+    }
+
+
+def get_best_mot_config(M: int, N: int, K: int, dtype_str: str | None = None) -> tuple[int, dict[str, int]]:
+    configs = get_mot_configs(K, N, dtype_str)
+    if configs:
+        loaded_m_key = min(configs.keys(), key=lambda x: abs(x - M))
+        return loaded_m_key, configs[loaded_m_key]
+    else:
+        return -1, get_mot_default_config(M, N, K, dtype=dtype_str)
+
+
+# =================================================================
+#  Part 1: The Router (Routing Component)
+#  Responsibilities: Handle PID mapping, Text/VAE distribution, indirect index loading, pointer calculation
+# =================================================================
+@triton.jit
+def _get_mot_pointers(
+    # System Inputs
+    pid,
+    # Matrix Pointers
+    a_ptr,
+    b_text_ptr,
+    b_vae_ptr,
+    bias_text_ptr,
+    bias_vae_ptr,
+    scale_a_ptr,
+    scale_b_text_ptr,
+    scale_b_vae_ptr,
+    # Indices & Meta
+    text_indices_ptr,
+    vae_indices_ptr,
+    M_text,
+    M_vae,
+    N,
+    # Strides (need to select based on Text/VAE)
+    stride_bk_text,
+    stride_bn_text,
+    stride_bk_vae,
+    stride_bn_vae,
+    # Block Config
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    # 1. Calculate Text/VAE task boundaries
+    num_pid_m_text = tl.cdiv(M_text, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+
+    # 2. MoT Routing
+    # Initialize VAE variables first and overwrite for Text
+
+    # VAE Path
+    cur_pid = pid - (num_pid_m_text * num_pid_n)
+    # Select VAE Pointers
+    cur_b_ptr = b_vae_ptr
+    cur_bias_ptr = bias_vae_ptr
+    cur_scale_b_ptr = scale_b_vae_ptr
+    cur_indices_ptr = vae_indices_ptr
+    # Select VAE Strides / Limits
+    cur_stride_bk = stride_bk_vae
+    cur_stride_bn = stride_bn_vae
+    M_limit = M_vae
+
+    # Text Path
+    if pid < num_pid_m_text * num_pid_n:
+        cur_pid = pid
+        # Select Text Pointers
+        cur_b_ptr = b_text_ptr
+        cur_bias_ptr = bias_text_ptr
+        cur_scale_b_ptr = scale_b_text_ptr
+        cur_indices_ptr = text_indices_ptr
+        # Select Text Strides / Limits
+        cur_stride_bk = stride_bk_text
+        cur_stride_bn = stride_bn_text
+        M_limit = M_text
+
+    # 3. Calculate Grid coordinates(grouping)
+    cur_num_pid_m = tl.cdiv(M_limit, BLOCK_SIZE_M)
+
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = cur_pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m_adj = tl.minimum(cur_num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + (cur_pid % group_size_m_adj)
+    pid_n = (cur_pid % num_pid_in_group) // group_size_m_adj
+
+    # 4. Load indirect indices (Indirect Indexing for A)
+    # Calculate the M range covered by current Block [0, BLOCK_SIZE_M]
+    offs_m_idx = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+    m_mask = offs_m_idx < M_limit
+    # Load real row indices from the index array
+    # for very big K (eg. ffn down proj*TP=1,K=16k), if M is also huge(M>130K)
+    # loading int32 indices may result in integer overflow when we compute offs_m
+    real_row_idxs = tl.load(cur_indices_ptr + offs_m_idx, mask=m_mask, other=0).to(tl.int64)
+
+    # 5. Calculate N-dimension Offsets
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+    n_mask = offs_n < N
+
+    return (
+        pid_m,
+        pid_n,  # Grid Coordinates
+        real_row_idxs,
+        m_mask,  # M-dim info (Indirect)
+        offs_n,
+        n_mask,  # N-dim info
+        M_limit,  # Boundary
+        cur_b_ptr,
+        cur_bias_ptr,  # Selected Pointers
+        cur_scale_b_ptr,  # Selected Scale Pointer
+        cur_stride_bk,
+        cur_stride_bn,  # Selected Strides
+    )
+
+
+# =================================================================
+#  Part 2: Compute Cores
+#  Responsibilities: Execute specific Loop structures based on QUANT_TYPE
+# =================================================================
+
+
+# Core A: Standard GEMM (for BF16/FP16 and W8A8)
+# Feature: No dequantization inside Loop, Scale is applied after Loop ends
+@triton.jit
+def _core_standard_gemm(
+    # Pointers
+    a_ptr,
+    b_ptr,
+    # Offsets & Masks
+    real_row_idxs,
+    m_mask,
+    offs_n,
+    n_mask,
+    offs_k,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    # Quant-Related
+    scale_a_ptr,
+    scale_b_ptr,
+    stride_scale_a,
+    stride_scale_b,
+    # Loop Info
+    K,
+    BLOCK_SIZE_K: tl.constexpr,
+    # Configs
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    ACCUMULATOR_DTYPE: tl.constexpr,
+    IS_W8A8: tl.constexpr,
+    # Accelerate Configs
+    EVEN_K: tl.constexpr,
+    EVEN_N: tl.constexpr,
+    STRIDE_AK_IS_1: tl.constexpr,
+    STRIDE_BK_IS_1: tl.constexpr,
+    STRIDE_BN_IS_1: tl.constexpr,
+):
+    # 1. Stride optimizations (Bypassing compiler limitations for unit strides)
+    _stride_ak = 1 if STRIDE_AK_IS_1 else stride_ak
+    _stride_bk = 1 if STRIDE_BK_IS_1 else stride_bk
+    _stride_bn = 1 if STRIDE_BN_IS_1 else stride_bn
+
+    # Pointer initialization (A uses indirect indexing, B uses standard striding)
+    a_ptrs = a_ptr + (stride_am * real_row_idxs[:, None] + _stride_ak * offs_k[None, :])
+    b_ptrs = b_ptr + (_stride_bk * offs_k[:, None] + _stride_bn * offs_n[None, :])
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=ACCUMULATOR_DTYPE)
+
+    # 2. Unified Main Loop
+    # Triton evaluates `constexpr` conditions at compile time, ensuring zero runtime overhead
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # Compute K-mask only if needed
+        if not EVEN_K:
+            mask_k = offs_k < K - k * BLOCK_SIZE_K
+
+        # Load A
+        if EVEN_K:
+            a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
+        else:
+            a = tl.load(a_ptrs, mask=m_mask[:, None] & mask_k[None, :], other=0.0)
+
+        # Load B
+        if EVEN_K and EVEN_N:
+            b = tl.load(b_ptrs)
+        elif not EVEN_K and EVEN_N:
+            b = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0)
+        elif EVEN_K and not EVEN_N:
+            b = tl.load(b_ptrs, mask=n_mask[None, :], other=0.0)
+        else:
+            b = tl.load(b_ptrs, mask=mask_k[:, None] & n_mask[None, :], other=0.0)
+
+        # Compute & Advance
+        accumulator = tl.dot(a, b, accumulator, out_dtype=ACCUMULATOR_DTYPE)
+        a_ptrs += BLOCK_SIZE_K * _stride_ak
+        b_ptrs += BLOCK_SIZE_K * _stride_bk
+
+    # 3. Epilogue (only needed for W8A8)
+    if IS_W8A8:
+        accumulator = accumulator.to(tl.float32)
+        # Load Scale A
+        scale_a_ptrs = scale_a_ptr + real_row_idxs * stride_scale_a
+        sa = tl.load(scale_a_ptrs, mask=m_mask, other=1.0)
+        accumulator = accumulator * sa[:, None]
+
+        # Load Scale B
+        scale_b_ptrs = scale_b_ptr + offs_n * stride_scale_b
+        sb = tl.load(scale_b_ptrs, mask=n_mask, other=1.0)
+        accumulator = accumulator * sb[None, :]
+
+    return accumulator
+
+
+# Core B: Weight Only GEMM (for W4A16 / W8A16)
+# Feature: Dequantization inside Loop (Dequantize-on-the-fly)
+# Currently only supports W8A16
+@triton.jit
+def _core_weight_only_gemm(
+    # Pointers
+    a_ptr,
+    b_ptr,
+    # Offsets & Masks
+    real_row_idxs,
+    m_mask,
+    offs_n,
+    n_mask,
+    offs_k,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    # Quant-Related
+    scale_b_ptr,
+    stride_scale_b,  # 1(per-channel) or 0(per-tensor)
+    # Loop Info
+    K,
+    BLOCK_SIZE_K: tl.constexpr,
+    # Configs
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    ACCUMULATOR_DTYPE: tl.constexpr,
+    COMPUTE_DTYPE: tl.constexpr,  # bf16 or fp16
+    WEIGHT_BITS: tl.constexpr,  # 4 or 8
+    # Accelerate Configs
+    EVEN_K: tl.constexpr,
+    EVEN_N: tl.constexpr,
+    STRIDE_AK_IS_1: tl.constexpr,
+    STRIDE_BK_IS_1: tl.constexpr,
+    STRIDE_BN_IS_1: tl.constexpr,
+):
+    # Compile-time Check
+    tl.static_assert(WEIGHT_BITS == 8, "For weight-only, we only support W8A16 at this point")
+
+    # 1. Stride optimizations
+    _stride_ak = 1 if STRIDE_AK_IS_1 else stride_ak
+    _stride_bk = 1 if STRIDE_BK_IS_1 else stride_bk
+    _stride_bn = 1 if STRIDE_BN_IS_1 else stride_bn
+
+    a_ptrs = a_ptr + (stride_am * real_row_idxs[:, None] + _stride_ak * offs_k[None, :])
+    b_ptrs = b_ptr + (_stride_bk * offs_k[:, None] + _stride_bn * offs_n[None, :])
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=ACCUMULATOR_DTYPE)
+
+    # 2. Unified Main Loop
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        if not EVEN_K:
+            mask_k = offs_k < K - k * BLOCK_SIZE_K
+
+        # Load A
+        if EVEN_K:
+            a = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0)
+        else:
+            a = tl.load(a_ptrs, mask=m_mask[:, None] & mask_k[None, :], other=0.0)
+
+        # Load B (Int)
+        if EVEN_K and EVEN_N:
+            b_int = tl.load(b_ptrs)
+        elif not EVEN_K and EVEN_N:
+            b_int = tl.load(b_ptrs, mask=mask_k[:, None], other=0.0)
+        elif EVEN_K and not EVEN_N:
+            b_int = tl.load(b_ptrs, mask=n_mask[None, :], other=0.0)
+        else:
+            b_int = tl.load(b_ptrs, mask=mask_k[:, None] & n_mask[None, :], other=0.0)
+
+        # --- Dequantize Logic (Type Cast Only) ---
+        # No multiplication here, correct as long as per-tensor/per-channel scaling (scale_b.shape=(N,))
+        b_compute = b_int.to(COMPUTE_DTYPE)
+
+        # Compute & Advance
+        accumulator = tl.dot(a, b_compute, accumulator, out_dtype=ACCUMULATOR_DTYPE)
+
+        a_ptrs += BLOCK_SIZE_K * _stride_ak
+        b_ptrs += BLOCK_SIZE_K * _stride_bk
+
+    # 3. Epilogue: Apply Scale B safely outside the loop
+    # Load Scale B
+    scale_b_ptrs = scale_b_ptr + offs_n * stride_scale_b
+    sb = tl.load(scale_b_ptrs, mask=n_mask, other=1.0)  # shape=(N，)
+
+    accumulator = accumulator * sb[None, :]
+    return accumulator
+
+
+# =================================================================
+#  Part 3: Unified Entry Kernel
+#  Responsibilities: Call Router, statically dispatch Core, store results
+# =================================================================
+@triton.jit
+def mot_unified_gemm_kernel(
+    # Inputs
+    a_ptr,
+    b_text_ptr,
+    b_vae_ptr,
+    c_ptr,
+    bias_text_ptr,
+    bias_vae_ptr,
+    text_indices_ptr,
+    vae_indices_ptr,
+    # Dimensions & Strides
+    M_text,
+    M_vae,
+    N,
+    K,
+    stride_am,
+    stride_ak,
+    stride_bk_text,
+    stride_bn_text,
+    stride_bk_vae,
+    stride_bn_vae,
+    stride_cm,
+    stride_cn,
+    # Scales (pass 0 if None)
+    scale_a_ptr,
+    scale_b_text_ptr,
+    scale_b_vae_ptr,
+    stride_scale_a,
+    stride_scale_b,  # 1(per-channel) or 0(per-tensor)
+    # Metas
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    # accelerate config
+    EVEN_K: tl.constexpr,
+    EVEN_N: tl.constexpr,
+    STRIDE_AK_IS_1: tl.constexpr,
+    STRIDE_BK_IS_1: tl.constexpr,
+    STRIDE_BN_IS_1: tl.constexpr,
+    # Quant-related dtypes
+    ACCUMULATOR_DTYPE: tl.constexpr,
+    COMPUTE_DTYPE: tl.constexpr,
+    OUTPUT_DTYPE: tl.constexpr,
+    # Quant Control
+    # 0=None, 1=W8A8, 2=W8A16, 3=W4A16
+    QUANT_TYPE: tl.constexpr,
+    HAS_BIAS: tl.constexpr = False,
+):
+    pid = tl.program_id(axis=0)
+
+    # -----------------------------------------------------------
+    # 1. Routing Phase (General)
+    # -----------------------------------------------------------
+    (
+        pid_m,
+        pid_n,
+        real_row_idxs,
+        m_mask,
+        offs_n,
+        n_mask,
+        M_limit,
+        cur_b_ptr,
+        cur_bias_ptr,
+        cur_scale_b_ptr,
+        cur_stride_bk,
+        cur_stride_bn,
+    ) = _get_mot_pointers(
+        pid,
+        a_ptr,
+        b_text_ptr,
+        b_vae_ptr,
+        bias_text_ptr,
+        bias_vae_ptr,
+        scale_a_ptr,
+        scale_b_text_ptr,
+        scale_b_vae_ptr,
+        text_indices_ptr,
+        vae_indices_ptr,
+        M_text,
+        M_vae,
+        N,
+        stride_bk_text,
+        stride_bn_text,
+        stride_bk_vae,
+        stride_bn_vae,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+        GROUP_SIZE_M,
+    )
+
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+    # -----------------------------------------------------------
+    # 2. Compute Phase (Static Dispatch)
+    # -----------------------------------------------------------
+    if QUANT_TYPE == 0:  # FP16 / BF16 Standard
+        c = _core_standard_gemm(
+            a_ptr,
+            cur_b_ptr,
+            real_row_idxs,
+            m_mask,
+            offs_n,
+            n_mask,
+            offs_k,
+            stride_am,
+            stride_ak,
+            cur_stride_bk,
+            cur_stride_bn,
+            0,
+            0,
+            0,
+            0,
+            K,
+            BLOCK_SIZE_K,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            ACCUMULATOR_DTYPE,
+            IS_W8A8=False,
+            EVEN_K=EVEN_K,
+            EVEN_N=EVEN_N,
+            STRIDE_AK_IS_1=STRIDE_AK_IS_1,
+            STRIDE_BK_IS_1=STRIDE_BK_IS_1,
+            STRIDE_BN_IS_1=STRIDE_BN_IS_1,
+        )
+    elif QUANT_TYPE == 1:  # W8A8
+        c = _core_standard_gemm(
+            a_ptr,
+            cur_b_ptr,
+            real_row_idxs,
+            m_mask,
+            offs_n,
+            n_mask,
+            offs_k,
+            stride_am,
+            stride_ak,
+            cur_stride_bk,
+            cur_stride_bn,
+            scale_a_ptr,
+            cur_scale_b_ptr,
+            stride_scale_a,
+            stride_scale_b,
+            K,
+            BLOCK_SIZE_K,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            ACCUMULATOR_DTYPE,
+            IS_W8A8=True,
+            EVEN_K=EVEN_K,
+            EVEN_N=EVEN_N,
+            STRIDE_AK_IS_1=STRIDE_AK_IS_1,
+            STRIDE_BK_IS_1=STRIDE_BK_IS_1,
+            STRIDE_BN_IS_1=STRIDE_BN_IS_1,
+        )
+    elif QUANT_TYPE == 2 or QUANT_TYPE == 3:  # Weight Only
+        bits = 8 if QUANT_TYPE == 2 else 4
+        c = _core_weight_only_gemm(
+            a_ptr,
+            cur_b_ptr,
+            real_row_idxs,
+            m_mask,
+            offs_n,
+            n_mask,
+            offs_k,
+            stride_am,
+            stride_ak,
+            cur_stride_bk,
+            cur_stride_bn,
+            cur_scale_b_ptr,
+            stride_scale_b,
+            K,
+            BLOCK_SIZE_K,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            ACCUMULATOR_DTYPE,
+            COMPUTE_DTYPE,
+            WEIGHT_BITS=bits,
+            EVEN_K=EVEN_K,
+            EVEN_N=EVEN_N,
+            STRIDE_AK_IS_1=STRIDE_AK_IS_1,
+            STRIDE_BK_IS_1=STRIDE_BK_IS_1,
+            STRIDE_BN_IS_1=STRIDE_BN_IS_1,
+        )
+
+    # -----------------------------------------------------------
+    # 3. Store Phase (General)
+    # -----------------------------------------------------------
+
+    # Bias Add
+    if HAS_BIAS:
+        bias = tl.load(cur_bias_ptr + offs_n, mask=n_mask, other=0.0)
+        c = c + bias[None, :]  # reshape to (1, BLOCK_SIZE_N)
+
+    # Cast C into the output dtype
+    c = c.to(OUTPUT_DTYPE)
+
+    # Store
+    offs_cm = real_row_idxs
+    c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_n[None, :]
+
+    # Need to be careful with the Mask here: real_row_idxs may contain 0
+    # (out-of-bounds padding), but c_mask needs the real boundary
+    store_mask = m_mask[:, None] & n_mask[None, :]
+    tl.store(c_ptrs, c, mask=store_mask)
+
+
+# Define is_weak_contiguous
+def is_weak_contiguous(x: torch.Tensor):
+    strides = x.stride()
+    sizes = x.shape
+    is_not_transpose = strides[0] == 1 and (strides[1] >= max(1, sizes[0]))
+    is_transpose = strides[1] == 1 and (strides[0] >= max(1, sizes[1]))
+    return is_transpose or is_not_transpose
+
+
+def invoke_mot_gemm(
+    # Inputs
+    A: torch.Tensor,
+    B_text: torch.Tensor,
+    B_vae: torch.Tensor,
+    C: torch.Tensor,
+    bias_text: torch.Tensor | None,
+    bias_vae: torch.Tensor | None,
+    # Indices
+    text_indices: torch.Tensor,
+    vae_indices: torch.Tensor,
+    # Quant Scales (None if disabled)
+    A_scale: torch.Tensor | None,
+    B_text_scale: torch.Tensor | None,
+    B_vae_scale: torch.Tensor | None,
+    # Quant Flags
+    use_fp8_w8a8: bool,
+    use_int8_w8a8: bool,
+    use_int8_w8a16: bool,
+    use_int4_w4a16: bool,
+    # Quant Type
+    A_per_channel_quant: bool,
+    B_per_channel_quant: bool,
+    # Config
+    config: dict[str, Any] | None = None,
+):
+    # ------ 1. Basic Assertions ------
+    M = A.size(0)
+    K = B_text.size(0)
+    N = B_text.size(1)
+
+    if config is None:
+        if use_fp8_w8a8:
+            _dtype_str: str | None = "fp8_w8a8"
+        elif use_int8_w8a16:
+            _dtype_str = "int8_w8a16"
+        else:
+            _dtype_str = None
+
+        loaded_m_key, config = get_best_mot_config(M, N, K, _dtype_str)
+
+    assert len(A.shape) == 2 and len(C.shape) == 2, (
+        "The input tensor and output tensor should be flattened to (batch_size*seq_len, hidden_dim)"
+    )
+
+    assert K == A.size(1), "the weights' first dimension should matchinputtensor's last dimension (hidden_dim)"
+
+    assert K == B_vae.size(0) and N == B_vae.size(1), (
+        "the weights dimension for text expert andimage expert should be the same"
+    )
+
+    assert C.size(0) == M and C.size(1) == N, "the output tensor shape is not correct"
+
+    M_text = text_indices.size(0)
+    M_vae = vae_indices.size(0)
+    assert M_text + M_vae == M, "the length sum of text and image indices should match input tensor's first dimension"
+    if bias_text is not None:
+        assert bias_text.dtype == C.dtype, "the bias tensor dtype should match the output tensor dtype"
+    if bias_vae is not None:
+        assert bias_vae.dtype == C.dtype, "the bias tensor dtype should match the output tensor dtype"
+
+    assert is_weak_contiguous(A)
+    assert is_weak_contiguous(B_text)
+    assert is_weak_contiguous(B_vae)
+
+    # --- 2. Quantization Logic Translation ---
+
+    def triton_dtype(torch_dtype):
+        if torch_dtype == torch.float8_e4m3fn:
+            return getattr(tl, "float8e4m3fn", tl.float8e4nv)
+        elif torch_dtype == torch.float8_e5m2:
+            return tl.float8e5
+        return {
+            torch.float16: tl.float16,
+            torch.bfloat16: tl.bfloat16,
+            torch.float32: tl.float32,
+            torch.int8: tl.int8,
+            torch.float8_e4m3fn: tl.float8e4nv,
+            torch.float8_e5m2: tl.float8e5,
+        }[torch_dtype]
+
+    # Determine QUANT_TYPE
+    # 0=None, 1=W8A8, 2=W8A16, 3=W4A16
+    quant_type = 0
+    ACCUMULATOR_DTYPE = tl.float32
+    COMPUTE_DTYPE = triton_dtype(A.dtype)
+    OUTPUT_DTYPE = triton_dtype(C.dtype)
+
+    if use_int8_w8a8 or use_fp8_w8a8:
+        quant_type = 1
+        assert A_scale is not None, "W8A8 requires A_scale"
+        assert B_text_scale is not None and B_vae_scale is not None, "W8A8 requires B_text_scale and B_vae_scale"
+        if use_int8_w8a8:
+            ACCUMULATOR_DTYPE = tl.int32
+            assert (
+                A.dtype == torch.int8
+                and B_text.dtype == torch.int8
+                and B_vae.dtype == torch.int8
+                and C.dtype in [torch.float16, torch.bfloat16]
+            ), "if you want to use INT8_W8A8, A should be INT8, B should be INT8, C should be FP16/BF16"
+        else:
+            ACCUMULATOR_DTYPE = tl.float32
+            assert (
+                A.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]
+                and B_text.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]
+                and B_vae.dtype in [torch.float8_e4m3fn, torch.float8_e5m2]
+                and C.dtype in [torch.float16, torch.bfloat16]
+            ), "if you want to use FP8_W8A8, A should be FP8, B should be FP8, C should be FP16/BF16"
+    elif use_int8_w8a16:
+        quant_type = 2
+        assert B_text_scale is not None and B_vae_scale is not None, "W8A16 requires B_text_scale and B_vae_scale"
+        ACCUMULATOR_DTYPE = tl.float32
+        assert (
+            A.dtype in [torch.float16, torch.bfloat16]
+            and B_text.dtype == torch.int8
+            and B_vae.dtype == torch.int8
+            and C.dtype in [torch.float16, torch.bfloat16]
+        ), "if you want to use INT8_W8A16, A should be FP16/BF16, B should be INT8, C should be FP16/BF16"
+
+    elif use_int4_w4a16:
+        raise NotImplementedError("For weight-only, we only support W8A16 at this point")
+        # quant_type = 3
+        # ACCUMULATOR_DTYPE=tl.float32
+
+    # accelerate config
+    EVEN_K = K % config["BLOCK_SIZE_K"] == 0
+    EVEN_N = N % config["BLOCK_SIZE_N"] == 0
+    STRIDE_AK_IS_1 = A.stride(1) == 1
+    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
+
+    # --- 3. Grid Calculation ---
+    def grid(META):
+        return (
+            (triton.cdiv(M_text, META["BLOCK_SIZE_M"]) + triton.cdiv(M_vae, META["BLOCK_SIZE_M"]))
+            * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+
+    # --- 4. Launch ---
+    run_config = config.copy()
+    run_config.update(
+        {
+            "QUANT_TYPE": quant_type,
+            "ACCUMULATOR_DTYPE": ACCUMULATOR_DTYPE,
+            "COMPUTE_DTYPE": COMPUTE_DTYPE,
+            "OUTPUT_DTYPE": OUTPUT_DTYPE,
+            "HAS_BIAS": has_bias,
+            "EVEN_K": EVEN_K,
+            "EVEN_N": EVEN_N,
+            "STRIDE_AK_IS_1": STRIDE_AK_IS_1,
+            "STRIDE_BK_IS_1": STRIDE_BK_IS_1,
+            "STRIDE_BN_IS_1": STRIDE_BN_IS_1,
+        }
+    )
+
+    # Pointers (Handle None -> 0)
+    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
+
+    # Quantization granularity
+    stride_scale_a = 1 if A_per_channel_quant else 0
+    stride_scale_b = 1 if B_per_channel_quant else 0
+
+    mot_unified_gemm_kernel[grid](
+        # Inputs
+        A,
+        B_text,
+        B_vae,
+        C,
+        p_bias_text,
+        p_bias_vae,
+        text_indices,
+        vae_indices,
+        # Dimensions
+        M_text,
+        M_vae,
+        N,
+        K,
+        # Strides
+        A.stride(0),
+        A.stride(1),
+        B_text.stride(0),
+        B_text.stride(1),
+        B_vae.stride(0),
+        B_vae.stride(1),
+        C.stride(0),
+        C.stride(1),
+        # Scales
+        p_a_scale,
+        p_b_text_scale,
+        p_b_vae_scale,
+        stride_scale_a,
+        stride_scale_b,
+        # Config
+        **run_config,
+    )

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

@@ -0,0 +1,283 @@
+# ruff: noqa: N803, E741
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _mot_rms_norm_kernel(
+    input_ptr,
+    text_weight_ptr,
+    vae_weight_ptr,
+    text_indices_ptr,  # MoT Routing Info
+    vae_indices_ptr,  # MoT Routing Info
+    M_text,  # MoT Routing Info
+    output_ptr,
+    input_row_stride,
+    output_row_stride,
+    n_cols,
+    eps,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """
+    input shape: (batch_size*seq_len, hidden_size)
+    RMSNorm: y = x / sqrt(mean(x^2) + eps) * weight
+    """
+    # Step 0: MoT Routing
+    pid = tl.program_id(0).to(tl.int64)
+
+    # dummy init (must be scalar int64[] to match tl.load return type)
+    row_idx = tl.cast(0, tl.int64)
+    weight_ptr = text_weight_ptr
+
+    if pid < M_text:
+        # --- Text Path ---
+        row_idx = tl.load(text_indices_ptr + pid)
+        weight_ptr = text_weight_ptr
+    else:
+        # --- VAE Path ---
+        vae_pid = pid - M_text
+        row_idx = tl.load(vae_indices_ptr + vae_pid)
+        weight_ptr = vae_weight_ptr
+
+    row_start_ptr = input_ptr + row_idx * input_row_stride
+    output_row_start_ptr = output_ptr + row_idx * output_row_stride
+
+    # Step 1: Compute sum of squares in float32 to avoid overflow
+    sum_sq = tl.zeros([1], dtype=tl.float32)
+    for col_offset in range(0, n_cols, BLOCK_SIZE):
+        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
+        mask = col_idx < n_cols
+
+        vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
+        # Convert to float32 for accumulation to prevent overflow
+        vals_f32 = vals.to(tl.float32)
+        sq_vals = vals_f32 * vals_f32
+        sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0))
+
+    # Step 2: Compute RMS (root mean square) in float32
+    mean_sq = sum_sq / n_cols
+    rms = tl.sqrt(mean_sq + eps)
+    inv_rms = 1.0 / rms
+
+    # Step 3: Normalize and apply weight
+    for col_offset in range(0, n_cols, BLOCK_SIZE):
+        col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
+        mask = col_idx < n_cols
+        vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
+        weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0)
+        # Compute in float32 then convert back to input dtype
+        vals_f32 = vals.to(tl.float32)
+        weight_f32 = weight.to(tl.float32)
+        output_f32 = vals_f32 * inv_rms * weight_f32
+        output = output_f32.to(vals.dtype)
+        tl.store(output_row_start_ptr + col_idx, output, mask=mask)
+
+
+@triton.jit
+def _mot_rms_norm_qk_kernel(
+    input_ptr,
+    text_weight_ptr,
+    vae_weight_ptr,
+    text_indices_ptr,
+    vae_indices_ptr,
+    M_text,
+    output_ptr,
+    stride_tok,
+    stride_head,
+    out_stride_tok,
+    out_stride_head,
+    num_heads,
+    head_dim,
+    eps,
+    BLOCK_DIM: tl.constexpr,  # for head_dim=128,should be 128
+    SHARED_WEIGHT: tl.constexpr,  # weights are shared across heads
+):
+    pid_tok = tl.program_id(0).to(tl.int64)
+
+    # MoT Routing
+    if pid_tok < M_text:
+        row_idx = tl.load(text_indices_ptr + pid_tok)
+        weight_ptr = text_weight_ptr
+    else:
+        vae_pid = pid_tok - M_text
+        row_idx = tl.load(vae_indices_ptr + vae_pid)
+        weight_ptr = vae_weight_ptr
+
+    dim_offsets = tl.arange(0, BLOCK_DIM)
+    mask = dim_offsets < head_dim
+
+    # input.shape=(4098,28,128)=(seq_len,num_heads,head_dim)
+    # weight.shape=(28,128)=(num_heads,head_dim) or (128,)=(head_dim)
+    tok_in_ptr = input_ptr + row_idx * stride_tok
+    tok_out_ptr = output_ptr + row_idx * out_stride_tok
+
+    # if weight.shape=(128,)
+    if SHARED_WEIGHT:
+        weight = tl.load(weight_ptr + dim_offsets, mask=mask, other=1.0)
+        weight_f32 = weight.to(tl.float32)
+
+    # one program loop over all heads
+    for h in range(num_heads):
+        head_in_ptr = tok_in_ptr + h * stride_head
+        head_out_ptr = tok_out_ptr + h * out_stride_head
+
+        # if weight.shape=(28,128)
+        if not SHARED_WEIGHT:
+            weight_offset = h * head_dim + dim_offsets
+            weight = tl.load(weight_ptr + weight_offset, mask=mask, other=1.0)
+            weight_f32 = weight.to(tl.float32)
+
+        vals = tl.load(head_in_ptr + dim_offsets, mask=mask, other=0.0)
+        vals_f32 = vals.to(tl.float32)
+
+        sq_vals = vals_f32 * vals_f32
+        sum_sq = tl.sum(tl.where(mask, sq_vals, 0.0))
+        mean_sq = sum_sq / head_dim
+        rms = tl.sqrt(mean_sq + eps)
+        inv_rms = 1.0 / rms
+
+        output_f32 = vals_f32 * inv_rms * weight_f32
+        tl.store(head_out_ptr + dim_offsets, output_f32.to(vals.dtype), mask=mask)
+
+
+def mot_rms_norm(
+    input: torch.Tensor,
+    text_weight: torch.Tensor,
+    vae_weight: torch.Tensor,
+    text_indices: torch.Tensor,
+    vae_indices: torch.Tensor,
+    head_norm: bool = False,
+    eps: float = 1e-6,
+    block_size: int | None = None,
+) -> torch.Tensor:
+    """
+    Compute RMS normalization using Triton kernel.
+
+    RMS Norm normalizes the input by the root mean square and scales by weight:
+    output = input / sqrt(mean(input^2) + eps) * weight
+
+    Args:
+        input: Input tensor of shape  (batch_size*seq_len, hidden_size) or (batch_size,seq_len, hidden_size)
+        text_weight: Weight for text tokens tensor of shape (hidden_size,)
+        vae_weight: Weight for vae tokens tensor of shape (hidden_size,)
+        text_indices: indices of text tokens, (batch_size*2,)
+        vae_indices: indices of vae tokens, (batch_size*(seq_len-2),)
+        eps: Small constant for numerical stability
+
+    Returns:
+        Tensor with RMS normalization applied along the last dimension
+    """
+
+    assert input.shape[-1] == text_weight.shape[-1], (
+        f"Input last dimension ({input.shape[-1]}) must match Text weight dimension ({text_weight.shape[-1]})"
+    )
+    assert input.shape[-1] == vae_weight.shape[0], (
+        f"Input last dimension ({input.shape[-1]}) must match VAE weight dimension ({vae_weight.shape[-1]})"
+    )
+
+    original_shape = input.shape
+    text_indices = text_indices.reshape(-1)
+    vae_indices = vae_indices.reshape(-1)
+    M_text = text_indices.shape[0]
+    M_vae = vae_indices.shape[0]
+    num_tokens = M_text + M_vae
+
+    if not head_norm:
+        input_2d = input.reshape(-1, input.shape[-1])
+        assert input_2d.shape[0] == num_tokens, (
+            f"batch_size={input_2d.shape[0]}, len(text_indices)={M_text}, len(vae_indices)={M_vae}"
+            f"for layer norm, batched_token_length should match the sum of indices_length"
+        )
+        input_2d = input_2d.contiguous()
+        text_indices = text_indices.contiguous()
+        vae_indices = vae_indices.contiguous()
+
+        text_weight = text_weight.contiguous()
+        vae_weight = vae_weight.contiguous()
+
+        n_rows, n_cols = input_2d.shape
+
+        output = torch.empty_like(input_2d)
+
+        if block_size is None:
+            block_size = triton.next_power_of_2(n_cols)
+            block_size = min(block_size, 4096)
+
+        num_warps = 4
+        if block_size >= 2048:
+            num_warps = 8
+        elif block_size >= 1024:
+            num_warps = 4
+        else:
+            num_warps = 2
+
+        grid = (n_rows,)
+        _mot_rms_norm_kernel[grid](
+            input_2d,
+            text_weight,
+            vae_weight,
+            text_indices,
+            vae_indices,
+            M_text,
+            output,
+            input_2d.stride(0),
+            output.stride(0),
+            n_cols,
+            eps,
+            BLOCK_SIZE=block_size,
+            num_warps=num_warps,
+        )
+        return output.reshape(original_shape)
+    else:
+        # qk norm scenarios:
+        # input.shape=(batch_size, seq_len,head_num, head_dim) or
+        # input.shape=(batch_size*seq_len,head_num, head_dim)
+        assert len(original_shape) > 2, (
+            "If head_norm=True,input shape be 3D or 3D, last 2 dimensions should be head_num and head_dim"
+        )
+        num_heads = input.shape[-2]
+        head_dim = input.shape[-1]
+
+        is_shared_weight = text_weight.dim() == 1
+        if not is_shared_weight:
+            assert num_heads == text_weight.shape[0] and num_heads == vae_weight.shape[0], (
+                "when weights are not shared across heads, the first dimension of "
+                "weights should be num of heads and match input.shape[-2]"
+            )
+
+        # reshape to 3D
+        input_3d = input.view(-1, num_heads, head_dim)
+
+        input_3d = input_3d.contiguous()
+        text_indices = text_indices.contiguous()
+        vae_indices = vae_indices.contiguous()
+        text_weight = text_weight.contiguous()
+        vae_weight = vae_weight.contiguous()
+
+        output_3d = torch.empty_like(input_3d)
+
+        block_dim = triton.next_power_of_2(head_dim)
+        num_warps = 4 if block_dim > 128 else 2
+        num_stages = 3
+
+        _mot_rms_norm_qk_kernel[(num_tokens,)](
+            input_3d,
+            text_weight,
+            vae_weight,
+            text_indices,
+            vae_indices,
+            M_text,
+            output_3d,
+            input_3d.stride(0),
+            input_3d.stride(1),
+            output_3d.stride(0),
+            output_3d.stride(1),
+            num_heads,
+            head_dim,
+            eps,
+            BLOCK_DIM=block_dim,
+            SHARED_WEIGHT=is_shared_weight,
+            num_warps=num_warps,
+            num_stages=num_stages,
+        )
+        return output_3d.view(original_shape)

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -21,11 +21,11 @@ from transformers.models.qwen2.modeling_qwen2 import (
 )
 from transformers.utils import ModelOutput
 from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
     MergedColumnParallelLinear,
-    QKVParallelLinear,
     RowParallelLinear,
 )
 from vllm.model_executor.layers.quantization.base_config import (
@@ -47,6 +47,9 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     get_sp_group,
 )
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
+from vllm_omni.diffusion.layers.mot.mot_layernorm import MoTRMSNorm
+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.rope import RotaryEmbedding
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding
 
@@ -164,6 +167,13 @@ class BagelRotaryEmbedding(nn.Module):
 
 
 class BagelMLP(nn.Module):
+    """FFN with Mixture-of-Tokens routing via MoT parallel linear layers.
+
+    gate_proj + up_proj are fused into a single MoTMergedColumnParallelLinear.
+    down_proj uses MoTRowParallelLinear.  Both layers hold text weights on self
+    and vae weights on self.gen_exp, routing by text_indices / vae_indices.
+    """
+
     def __init__(
         self,
         hidden_size: int,
@@ -173,29 +183,35 @@ class BagelMLP(nn.Module):
         prefix: str = "",
     ) -> None:
         super().__init__()
+
+        self.intermediate_size = intermediate_size
+
         self.gate_up_proj = MergedColumnParallelLinear(
-            hidden_size,
-            [intermediate_size, intermediate_size],
+            input_size=hidden_size,
+            output_sizes=[intermediate_size, intermediate_size],
             bias=False,
+            gather_output=False,
             quant_config=quant_config,
             prefix=f"{prefix}.gate_up_proj",
         )
         self.down_proj = RowParallelLinear(
-            intermediate_size,
-            hidden_size,
-            input_is_parallel=True,
+            input_size=intermediate_size,
+            output_size=hidden_size,
             bias=False,
+            input_is_parallel=True,
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
         )
         if hidden_act != "silu":
             raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported.")
-        self.act_fn = nn.SiLU()
+        self.act_fn = SiluAndMul()
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
         gate_up, _ = self.gate_up_proj(x)
-        gate, up = gate_up.chunk(2, dim=-1)
-        x = self.act_fn(gate) * up
+        x = self.act_fn(gate_up)
         x, _ = self.down_proj(x)
         return x
 
@@ -294,12 +310,9 @@ class BaseNavitOutputWithPast(ModelOutput):
 class PackedAttentionMoT(nn.Module):
     """Packed attention with Mixture-of-Tokens routing for understanding/generation.
 
-    Uses vLLM's QKVParallelLinear and RowParallelLinear for tensor parallelism
-    support, following the same pattern as vLLM's Qwen2Attention.
-
-    The q/k/v projections are stacked into a single QKVParallelLinear:
-      - qkv_proj      : stacks q_proj + k_proj + v_proj  (understanding + gen text)
-      - qkv_proj_moe_gen : stacks q_proj_moe_gen + k_proj_moe_gen + v_proj_moe_gen (gen vae)
+    Uses MoTQKVParallelLinear and MoTRowParallelLinear for tensor parallelism.
+    Text and vae weights are held within the same MoT layer (text on self,
+    vae on self.gen_exp).  Token routing is driven by text_indices / vae_indices.
     """
 
     def __init__(
@@ -329,47 +342,26 @@ class PackedAttentionMoT(nn.Module):
         self.q_size = self.num_heads * self.head_dim
         self.kv_size = self.num_kv_heads * self.head_dim
 
-        # Understanding mode projections (stacked q/k/v)
-        self.qkv_proj = QKVParallelLinear(
+        self.qkv_proj = MoTQKVParallelLinear(
             self.hidden_size,
             self.head_dim,
             self.total_num_heads,
             self.total_num_kv_heads,
             bias=True,
+            vae_bias=True,
             quant_config=quant_config,
             prefix=f"{prefix}.qkv_proj",
         )
-        self.o_proj = RowParallelLinear(
-            self.total_num_heads * self.head_dim,
-            self.hidden_size,
+        self.o_proj = MoTRowParallelLinear(
+            input_size=self.total_num_heads * self.head_dim,
+            output_size=self.hidden_size,
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
         )
 
-        # Generation mode MoE projections (stacked q/k/v)
-        self.qkv_proj_moe_gen = QKVParallelLinear(
-            self.hidden_size,
-            self.head_dim,
-            self.total_num_heads,
-            self.total_num_kv_heads,
-            bias=True,
-            quant_config=quant_config,
-            prefix=f"{prefix}.qkv_proj_moe_gen",
-        )
-        self.o_proj_moe_gen = RowParallelLinear(
-            self.total_num_heads * self.head_dim,
-            self.hidden_size,
-            bias=False,
-            quant_config=quant_config,
-            prefix=f"{prefix}.o_proj_moe_gen",
-        )
-
-        # QK normalization
-        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        self.q_norm_moe_gen = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
-        self.k_norm_moe_gen = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.q_norm = MoTRMSNorm(self.head_dim, head_norm=True, eps=config.rms_norm_eps)
+        self.k_norm = MoTRMSNorm(self.head_dim, head_norm=True, eps=config.rms_norm_eps)
 
         self.rotary_op = RotaryEmbedding(is_neox_style=True)
 
@@ -404,60 +396,45 @@ class PackedAttentionMoT(nn.Module):
     ) -> tuple[torch.Tensor, NaiveCache | None]:
         """SP-aware attention for gen mode denoising.
 
-        Converts packed format to batched (1, S, H, D) and uses the diffusion
-        Attention layer (Ulysses / Ring) with joint mechanism:
+        Uses MoT unified layers for projection/norm (fused text/vae routing),
+        then splits into text/vae for the DiffusionAttention (Ulysses / Ring)
+        with joint mechanism:
           - Main Q/K/V: VAE tokens (split across SP ranks)
           - Joint Q: text marker Q (replicated)
           - Joint K/V: KV cache K/V + text marker K/V (replicated)
         """
+        text_indices = packed_text_indexes
+        vae_indices = packed_vae_token_indexes
+
         packed_query_sequence = packed_query_sequence.to(torch.bfloat16)
 
-        packed_text_query_sequence = packed_query_sequence[packed_text_indexes]
-        packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes]
+        # MoT QKV projection — fused kernel routes text/vae to correct weights
+        qkv, _ = self.qkv_proj(packed_query_sequence, text_indices, vae_indices)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        q = q.view(-1, self.num_heads, self.head_dim)
+        k = k.view(-1, self.num_kv_heads, self.head_dim)
+        v = v.view(-1, self.num_kv_heads, self.head_dim)
 
-        # Project text tokens through base qkv
-        text_qkv, _ = self.qkv_proj(packed_text_query_sequence)
-        text_q, text_k, text_v = text_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        # MoT QK norms — routes text/vae to weight/gen_weight internally
+        q = self.q_norm(q.to(torch.float32), text_indices, vae_indices)
+        k = self.k_norm(k.to(torch.float32), text_indices, vae_indices)
 
-        # Project vae tokens through moe_gen qkv
-        vae_qkv, _ = self.qkv_proj_moe_gen(packed_vae_query_sequence)
-        vae_q, vae_k, vae_v = vae_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        # RoPE on full packed tensor
+        cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
+        q = self.rotary_op(q.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
+        k = self.rotary_op(k.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
 
-        # Reshape to (tokens, heads, head_dim)
-        text_q = text_q.view(-1, self.num_heads, self.head_dim)
-        text_k = text_k.view(-1, self.num_kv_heads, self.head_dim)
-        text_v = text_v.view(-1, self.num_kv_heads, self.head_dim)
-        vae_q = vae_q.view(-1, self.num_heads, self.head_dim)
-        vae_k = vae_k.view(-1, self.num_kv_heads, self.head_dim)
-        vae_v = vae_v.view(-1, self.num_kv_heads, self.head_dim)
+        q = q.to(torch.bfloat16)
+        k = k.to(torch.bfloat16)
+        v = v.to(torch.bfloat16)
 
-        # Apply QK norms
-        text_q = self.q_norm(text_q.to(torch.float32))
-        text_k = self.k_norm(text_k.to(torch.float32))
-        vae_q = self.q_norm_moe_gen(vae_q.to(torch.float32))
-        vae_k = self.k_norm_moe_gen(vae_k.to(torch.float32))
-
-        # Apply RoPE - need to build per-token cos/sin for text and vae separately
-        # packed_query_position_embeddings are ordered as the packed sequence
-        cos_full, sin_full = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
-
-        # Extract cos/sin for text and vae positions
-        text_cos = cos_full[packed_text_indexes]
-        text_sin = sin_full[packed_text_indexes]
-        vae_cos = cos_full[packed_vae_token_indexes]
-        vae_sin = sin_full[packed_vae_token_indexes]
-
-        text_q = self.rotary_op(text_q.to(text_cos.dtype).unsqueeze(0), text_cos, text_sin).squeeze(0)
-        text_k = self.rotary_op(text_k.to(text_cos.dtype).unsqueeze(0), text_cos, text_sin).squeeze(0)
-        vae_q = self.rotary_op(vae_q.to(vae_cos.dtype).unsqueeze(0), vae_cos, vae_sin).squeeze(0)
-        vae_k = self.rotary_op(vae_k.to(vae_cos.dtype).unsqueeze(0), vae_cos, vae_sin).squeeze(0)
-
-        text_q = text_q.to(torch.bfloat16)
-        text_k = text_k.to(torch.bfloat16)
-        text_v = text_v.to(torch.bfloat16)
-        vae_q = vae_q.to(torch.bfloat16)
-        vae_k = vae_k.to(torch.bfloat16)
-        vae_v = vae_v.to(torch.bfloat16)
+        # Split into text / vae for SP attention
+        text_q = q[text_indices]
+        text_k = k[text_indices]
+        text_v = v[text_indices]
+        vae_q = q[vae_indices]
+        vae_k = k[vae_indices]
+        vae_v = v[vae_indices]
 
         # Build joint K/V: [kv_cache, text_markers] (replicated across SP ranks)
         if past_key_values is not None and past_key_values.key_cache[self.layer_idx] is not None:
@@ -487,15 +464,19 @@ class PackedAttentionMoT(nn.Module):
                 joint_strategy="front",
             ),
         )
-        # attn_out: (1, text_len + local_vae_len, H, D)
+        # attn_out: (1, text_len + vae_len, H, D)
         text_len = text_q.shape[0]
-        attn_out = attn_out.squeeze(0)  # (text_len + local_vae_len, H, D)
+        attn_out = attn_out.squeeze(0)
         text_attn = attn_out[:text_len].reshape(text_len, self.q_size)
         vae_attn = attn_out[text_len:].reshape(-1, self.q_size)
 
-        # Apply output projections
-        text_out, _ = self.o_proj(text_attn)
-        vae_out, _ = self.o_proj_moe_gen(vae_attn)
+        # MoT output projection — construct local packed tensor with local indices
+        local_packed = torch.cat([text_attn, vae_attn], dim=0)
+        local_text_idx = torch.arange(text_len, device=local_packed.device)
+        local_vae_idx = torch.arange(text_len, text_len + vae_attn.shape[0], device=local_packed.device)
+        local_out, _ = self.o_proj(local_packed, local_text_idx, local_vae_idx)
+        text_out = local_out[:text_len]
+        vae_out = local_out[text_len:]
 
         # Merge back into packed format
         total_len = packed_query_sequence.shape[0]
@@ -537,56 +518,25 @@ class PackedAttentionMoT(nn.Module):
                 packed_text_indexes=packed_text_indexes,
             )
 
-        if mode == "und":
-            qkv, _ = self.qkv_proj(packed_query_sequence)
-            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-            packed_query_states = q.view(-1, self.num_heads, self.head_dim)
-            packed_key_states = k.view(-1, self.num_kv_heads, self.head_dim)
-            packed_value_states = v.view(-1, self.num_kv_heads, self.head_dim)
-            packed_query_states = self.q_norm(packed_query_states)
-            packed_key_states = self.k_norm(packed_key_states)
-        elif mode == "gen":
+        text_indices = packed_text_indexes if mode == "gen" else None
+        vae_indices = packed_vae_token_indexes if mode == "gen" else None
+
+        if mode == "gen":
             packed_query_sequence = packed_query_sequence.to(torch.bfloat16)
 
-            packed_text_query_sequence = packed_query_sequence[packed_text_indexes]
-            packed_vae_query_sequence = packed_query_sequence[packed_vae_token_indexes]
-
-            # Project text tokens through base qkv
-            text_qkv, _ = self.qkv_proj(packed_text_query_sequence)
-            text_q, text_k, text_v = text_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-
-            # Project vae tokens through moe_gen qkv
-            vae_qkv, _ = self.qkv_proj_moe_gen(packed_vae_query_sequence)
-            vae_q, vae_k, vae_v = vae_qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-
-            # Merge into packed tensors
-            total_len = packed_query_sequence.shape[0]
-            packed_query_states = packed_query_sequence.new_zeros((total_len, self.q_size))
-            packed_key_states = packed_query_sequence.new_zeros((total_len, self.kv_size))
-            packed_value_states = packed_query_sequence.new_zeros((total_len, self.kv_size))
-
-            packed_query_states[packed_text_indexes] = text_q
-            packed_query_states[packed_vae_token_indexes] = vae_q
-            packed_key_states[packed_text_indexes] = text_k
-            packed_key_states[packed_vae_token_indexes] = vae_k
-            packed_value_states[packed_text_indexes] = text_v
-            packed_value_states[packed_vae_token_indexes] = vae_v
-
-            packed_query_states = packed_query_states.view(-1, self.num_heads, self.head_dim)
-            packed_key_states = packed_key_states.view(-1, self.num_kv_heads, self.head_dim)
-            packed_value_states = packed_value_states.view(-1, self.num_kv_heads, self.head_dim)
+        # QKV projection — MoT layer handles text/vae routing internally
+        qkv, _ = self.qkv_proj(packed_query_sequence, text_indices, vae_indices)
+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+        packed_query_states = q.view(-1, self.num_heads, self.head_dim)
+        packed_key_states = k.view(-1, self.num_kv_heads, self.head_dim)
+        packed_value_states = v.view(-1, self.num_kv_heads, self.head_dim)
 
+        if mode == "gen":
             packed_query_states = packed_query_states.to(torch.float32)
-            packed_query_states[packed_text_indexes] = self.q_norm(packed_query_states[packed_text_indexes])
-            packed_query_states[packed_vae_token_indexes] = self.q_norm_moe_gen(
-                packed_query_states[packed_vae_token_indexes]
-            )
-
             packed_key_states = packed_key_states.to(torch.float32)
-            packed_key_states[packed_text_indexes] = self.k_norm(packed_key_states[packed_text_indexes])
-            packed_key_states[packed_vae_token_indexes] = self.k_norm_moe_gen(
-                packed_key_states[packed_vae_token_indexes]
-            )
+
+        packed_query_states = self.q_norm(packed_query_states, text_indices, vae_indices)
+        packed_key_states = self.k_norm(packed_key_states, text_indices, vae_indices)
 
         cos, sin = [x[..., : self.head_dim // 2] for x in packed_query_position_embeddings]
         packed_query_states = self.rotary_op(packed_query_states.to(cos.dtype).unsqueeze(0), cos, sin).squeeze(0)
@@ -627,15 +577,9 @@ class PackedAttentionMoT(nn.Module):
             causal=is_causal,
         )
         packed_attn_output = packed_attn_output.reshape(-1, self.q_size)
-        if mode == "und":
-            packed_attn_output, _ = self.o_proj(packed_attn_output)
-        elif mode == "gen":
-            text_out, _ = self.o_proj(packed_attn_output[packed_text_indexes])
-            vae_out, _ = self.o_proj_moe_gen(packed_attn_output[packed_vae_token_indexes])
-            full_output = text_out.new_zeros((packed_attn_output.shape[0], self.hidden_size))
-            full_output[packed_text_indexes] = text_out
-            full_output[packed_vae_token_indexes] = vae_out
-            packed_attn_output = full_output
+
+        # Output projection — MoT layer handles text/vae routing internally
+        packed_attn_output, _ = self.o_proj(packed_attn_output, text_indices, vae_indices)
 
         if update_past_key_values:
             past_key_values.key_cache[self.layer_idx] = merged_key_states
@@ -655,12 +599,15 @@ class Qwen2MoTDecoderLayer(nn.Module):
         prefix: str = "",
     ):
         super().__init__()
+        self.layer_idx = layer_idx
         self.hidden_size = config.hidden_size
 
         self.self_attn = attn_module(
             config, layer_idx, parallel_config=parallel_config, quant_config=quant_config, prefix=f"{prefix}.self_attn"
         )
 
+        self.input_layernorm = MoTRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
         self.mlp = BagelMLP(
             config.hidden_size,
             config.intermediate_size,
@@ -675,8 +622,7 @@ class Qwen2MoTDecoderLayer(nn.Module):
             quant_config=quant_config,
             prefix=f"{prefix}.mlp_moe_gen",
         )
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.input_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
@@ -699,18 +645,12 @@ class Qwen2MoTDecoderLayer(nn.Module):
     ) -> BaseNavitOutputWithPast:
         if packed_query_sequence is None:
             packed_query_sequence = hidden_states
+
+        text_indices = packed_text_indexes if mode == "gen" else None
+        vae_indices = packed_vae_token_indexes if mode == "gen" else None
+
         residual = packed_query_sequence
-        if mode == "und":
-            packed_query_sequence = self.input_layernorm(packed_query_sequence)
-        elif mode == "gen":
-            packed_query_sequence_ = torch.zeros_like(packed_query_sequence)
-            packed_query_sequence_[packed_text_indexes] = self.input_layernorm(
-                packed_query_sequence[packed_text_indexes]
-            )
-            packed_query_sequence_[packed_vae_token_indexes] = self.input_layernorm_moe_gen(
-                packed_query_sequence[packed_vae_token_indexes]
-            )
-            packed_query_sequence = packed_query_sequence_
+        packed_query_sequence = self.input_layernorm(packed_query_sequence, text_indices, vae_indices)
 
         # Self Attention
         packed_query_sequence, past_key_values = self.self_attn(
@@ -741,7 +681,9 @@ class Qwen2MoTDecoderLayer(nn.Module):
             packed_vae_query_sequence = self.post_attention_layernorm_moe_gen(packed_vae_query_sequence).to(
                 torch.bfloat16
             )
-
+            packed_normed = torch.zeros_like(packed_query_sequence).to(torch.bfloat16)
+            packed_normed[packed_text_indexes] = packed_text_query_sequence
+            packed_normed[packed_vae_token_indexes] = packed_vae_query_sequence
             packed_query_sequence_ = torch.zeros_like(packed_query_sequence).to(torch.bfloat16)
             packed_query_sequence_[packed_text_indexes] = self.mlp(packed_text_query_sequence)
             packed_query_sequence_[packed_vae_token_indexes] = self.mlp_moe_gen(packed_vae_query_sequence)
@@ -788,9 +730,7 @@ class Qwen2MoTModel(Qwen2PreTrainedModel):
             ]
         )
 
-        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        if self.use_moe:
-            self.norm_moe_gen = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.norm = MoTRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = BagelRotaryEmbedding(config=config)
 
         # Initialize weights and apply final processing
@@ -856,18 +796,9 @@ class Qwen2MoTModel(Qwen2PreTrainedModel):
                 **extra_inputs,
             )
 
-        if self.use_moe:
-            if mode == "und":
-                packed_query_sequence = self.norm(packed_query_sequence)
-            elif mode == "gen":
-                packed_query_sequence_ = torch.zeros_like(packed_query_sequence)
-                packed_query_sequence_[packed_text_indexes] = self.norm(packed_query_sequence[packed_text_indexes])
-                packed_query_sequence_[packed_vae_token_indexes] = self.norm_moe_gen(
-                    packed_query_sequence[packed_vae_token_indexes]
-                )
-                packed_query_sequence = packed_query_sequence_
-        else:
-            packed_query_sequence = self.norm(packed_query_sequence)
+        text_indices = packed_text_indexes if self.use_moe and mode == "gen" else None
+        vae_indices = packed_vae_token_indexes if self.use_moe and mode == "gen" else None
+        packed_query_sequence = self.norm(packed_query_sequence, text_indices, vae_indices)
 
         return BaseNavitOutputWithPast(
             packed_query_sequence=packed_query_sequence,
@@ -950,57 +881,84 @@ class Qwen2MoTForCausalLM(Qwen2PreTrainedModel):
         return outputs
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        """Load weights for vLLM parallel layers.
+        """Load weights for MoT parallel layers.
 
-        Handles stacked parameter remapping for QKVParallelLinear:
-          - q_proj, k_proj, v_proj -> qkv_proj (shard ids: q, k, v)
-          - q_proj_moe_gen, k_proj_moe_gen, v_proj_moe_gen -> qkv_proj_moe_gen
-        Other parallel layers (gate_proj, up_proj, down_proj, embed_tokens, etc.)
-        keep HF checkpoint names and use weight_loader for TP sharding.
+        Stacked parameter remapping (checkpoint name → model parameter):
+          - q/k/v_proj       → qkv_proj          (text, shard q/k/v)
+          - q/k/v_proj_moe_gen → qkv_proj.gen_exp (gen,  shard q/k/v)
+
+        Direct remapping (no shard dimension):
+          - o_proj_moe_gen   → o_proj.gen_exp
+          - {norm}_moe_gen.weight → {norm}.gen_weight  (all MoTRMSNorm layers)
+
+        Text norm weights (input_layernorm.weight, q_norm.weight, etc.) and
+        other names (embed_tokens, lm_head) pass through unchanged.
         """
         stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            # More specific _moe_gen patterns FIRST to avoid substring
-            # ambiguity (`.q_proj` is a substring of `.q_proj_moe_gen`).
-            (".qkv_proj_moe_gen", ".q_proj_moe_gen", "q"),
-            (".qkv_proj_moe_gen", ".k_proj_moe_gen", "k"),
-            (".qkv_proj_moe_gen", ".v_proj_moe_gen", "v"),
+            # (param_name, weight_name, shard_id)
+            # _moe_gen patterns MUST come first — `.q_proj` is a substring
+            # of `.q_proj_moe_gen`, so the more specific pattern must match first.
+            (".qkv_proj.gen_exp", ".q_proj_moe_gen", "q"),
+            (".qkv_proj.gen_exp", ".k_proj_moe_gen", "k"),
+            (".qkv_proj.gen_exp", ".v_proj_moe_gen", "v"),
             (".qkv_proj", ".q_proj", "q"),
             (".qkv_proj", ".k_proj", "k"),
             (".qkv_proj", ".v_proj", "v"),
-            # MLP gate/up projections — fused into MergedColumnParallelLinear.
-            # HF checkpoints store separate gate_proj / up_proj weights;
-            # these entries remap them to the fused gate_up_proj parameter.
-            (".gate_up_proj", ".gate_proj", 0),
-            (".gate_up_proj", ".up_proj", 1),
+            (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.gate_proj", 0),
+            (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.up_proj", 1),
+            (".mlp.gate_up_proj", ".mlp.gate_proj", 0),
+            (".mlp.gate_up_proj", ".mlp.up_proj", 1),
         ]
-        self.stacked_params_mapping = stacked_params_mapping
+
+        direct_remap = [
+            (".o_proj_moe_gen.", ".o_proj.gen_exp."),
+            # Norm _moe_gen.weight → {norm_name}.gen_weight
+            (".input_layernorm_moe_gen.", ".input_layernorm.gen_"),
+            (".q_norm_moe_gen.", ".q_norm.gen_"),
+            (".k_norm_moe_gen.", ".k_norm.gen_"),
+            (".norm_moe_gen.", ".norm.gen_"),
+        ]
+
         params_dict = dict(self.named_parameters())
         loaded_params: set[str] = set()
 
+        def handle_weight(name, loaded_weight, shard_id=None):
+            param = params_dict.get(name)
+            if param is not None:
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                if shard_id is not None:
+                    weight_loader(param, loaded_weight, shard_id)
+                else:
+                    weight_loader(param, loaded_weight)
+                loaded_params.add(name)
+
         for name, loaded_weight in weights:
-            loaded = False
+            # match direct remap
+            handled = False
+            for old_substr, new_substr in direct_remap:
+                if old_substr in name:
+                    name = name.replace(old_substr, new_substr)
+                    handle_weight(name, loaded_weight)
+                    handled = True
+                    break
+
+            if handled:
+                continue
+
+            # match stacked params mapping
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
-                stacked_name = name.replace(weight_name, param_name)
-                param = params_dict.get(stacked_name)
-                if param is None:
-                    break
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight, shard_id)
-                name = stacked_name
-                loaded = True
+                name = name.replace(weight_name, param_name)
+                handle_weight(name, loaded_weight, shard_id)
+                handled = True
                 break
 
-            if not loaded:
-                param = params_dict.get(name)
-                if param is None:
-                    continue
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
+            if handled:
+                continue
 
-            loaded_params.add(name)
+            # no-name-match cases are handled here
+            handle_weight(name, loaded_weight)
 
         return loaded_params
 

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -845,18 +845,33 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
         tp_aware_params = {name for name, p in self.named_parameters() if hasattr(p, "weight_loader")}
 
         # Expand allowed/tp_aware_params with stacked param source names.
-        # QKVParallelLinear merges q_proj+k_proj+v_proj into qkv_proj; the
-        # checkpoint stores the original separate names.  We must recognise
-        # those names so _filtered_weights does not drop them.
+        # The model fuses several checkpoint projections into merged layers:
+        #   QKV: q/k/v_proj → qkv_proj, q/k/v_proj_moe_gen → qkv_proj.gen_exp
+        # and remaps non-stacked weights like:
+        #   {norm}_moe_gen.weight → {norm}.gen_weight  (MoTRMSNorm layers)
+        # We expand allowed names so _filtered_weights does not drop them.
         _stacked_expansions = [
+            # text QKV
             (".qkv_proj", ".q_proj"),
             (".qkv_proj", ".k_proj"),
             (".qkv_proj", ".v_proj"),
-            (".qkv_proj_moe_gen", ".q_proj_moe_gen"),
-            (".qkv_proj_moe_gen", ".k_proj_moe_gen"),
-            (".qkv_proj_moe_gen", ".v_proj_moe_gen"),
-            (".gate_up_proj", ".gate_proj"),
-            (".gate_up_proj", ".up_proj"),
+            # gen QKV
+            (".qkv_proj.gen_exp", ".q_proj_moe_gen"),
+            (".qkv_proj.gen_exp", ".k_proj_moe_gen"),
+            (".qkv_proj.gen_exp", ".v_proj_moe_gen"),
+            # gen o_proj (non-stacked, but still remapped)
+            (".o_proj.gen_exp", ".o_proj_moe_gen"),
+            # text FFN gate+up
+            (".mlp.gate_up_proj", ".mlp.gate_proj"),
+            (".mlp.gate_up_proj", ".mlp.up_proj"),
+            # gen FFN gate+up
+            (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.gate_proj"),
+            (".mlp_moe_gen.gate_up_proj", ".mlp_moe_gen.up_proj"),
+            # MoTRMSNorm gen_weight ← checkpoint _moe_gen.weight
+            (".input_layernorm.gen_", ".input_layernorm_moe_gen."),
+            (".q_norm.gen_", ".q_norm_moe_gen."),
+            (".k_norm.gen_", ".k_norm_moe_gen."),
+            (".norm.gen_", ".norm_moe_gen."),
         ]
         stacked_source_names: set[str] = set()
         for name in list(allowed):