[Refactor][Diffusion] Move single-model layers into their model directories (#5908)

Signed-off-by: congw729 <115451386+congw729@users.noreply.github.com>
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

vllm_omni/diffusion/models/bagel/bagel_transformer.py

@@ -1,6 +1,7 @@
 # Copyright 2025 Bytedance Ltd. and/or its affiliates.
 # Copyright (c) 2024 The Qwen Team and The HuggingFace Inc. team.
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 #
 # This file has been modified by ByteDance Ltd. and/or its affiliates.
 #
@@ -48,10 +49,10 @@ 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.diffusion.models.bagel.mot.mot_layernorm import MoTRMSNorm
+from vllm_omni.diffusion.models.bagel.mot.mot_qkv_parallel_linear import MoTQKVParallelLinear
+from vllm_omni.diffusion.models.bagel.mot.mot_row_parallel_linear import MoTRowParallelLinear
 from vllm_omni.diffusion.utils.kv_utils import left_pad_stack
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding
 

vllm_omni/diffusion/models/bagel/mot/mot_layernorm.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+# 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.models.bagel.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,
+        )
+
+    # ------------------------------------------------------------------
+    # NPU fast-path (not implemented yet; reserved for future optimization)
+    # ------------------------------------------------------------------
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+        text_indices: torch.Tensor | None = None,
+        vae_indices: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.forward_native(x, text_indices, vae_indices)
+
+    # ------------------------------------------------------------------
+    # 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/models/bagel/mot/mot_qkv_parallel_linear.py

@@ -0,0 +1,367 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+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.models.bagel.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) ----
+        if self.quant_method is None:
+            raise ValueError(
+                f"quant_method must not be None for MoTQKVParallelLinear (prefix={prefix!r}). "
+                "Ensure a valid QuantizationConfig is provided or the default UnquantizedLinearMethod is set."
+            )
+
+        # 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()
+
+        # ``gen_exp`` participates in the same tensor-parallel layout as this
+        # layer.  vLLM's online quantizers inspect these attributes while
+        # processing weights (for example, FP8 uses them to share scales across
+        # TP ranks), so the lightweight container must expose the parallel
+        # linear metadata rather than only holding parameters.
+        for attr in (
+            "tp_rank",
+            "tp_size",
+            "input_size",
+            "input_size_per_partition",
+            "output_size",
+            "output_size_per_partition",
+        ):
+            setattr(self.gen_exp, attr, getattr(self, attr))
+
+        # 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.
+        """
+        if self.quant_method is None:
+            raise ValueError(
+                "quant_method must not be None in MoTQKVParallelLinear._mot_fallback. "
+                "This indicates an initialization error."
+            )
+
+        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/models/bagel/mot/mot_row_parallel_linear.py

@@ -0,0 +1,379 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+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.models.bagel.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) ----
+        if self.quant_method is None:
+            raise ValueError(
+                f"quant_method must not be None for MoTRowParallelLinear (prefix={prefix!r}). "
+                "Ensure a valid QuantizationConfig is provided or the default UnquantizedLinearMethod is set."
+            )
+
+        # 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()
+
+        # ``gen_exp`` participates in the same tensor-parallel layout as this
+        # layer.  vLLM's online quantizers inspect these attributes while
+        # processing weights (for example, FP8 uses them to share scales across
+        # TP ranks), so the lightweight container must expose the parallel
+        # linear metadata rather than only holding parameters.
+        for attr in (
+            "tp_rank",
+            "tp_size",
+            "input_size",
+            "input_size_per_partition",
+            "output_size",
+            "output_size_per_partition",
+        ):
+            setattr(self.gen_exp, attr, getattr(self, attr))
+
+        # 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.
+        """
+        if self.quant_method is None:
+            raise ValueError(
+                "quant_method must not be None in MoTRowParallelLinear._mot_fallback. "
+                "This indicates an initialization error."
+            )
+
+        # 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/models/bagel/mot/ops/mot_gemm.py

@@ -0,0 +1,911 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+# 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/models/bagel/mot/ops/mot_rms_norm.py

@@ -0,0 +1,286 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+# 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[-1], (
+        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/wan2_2/norm.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Normalization layer used by the Wan VAE patch."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from vllm_omni.diffusion.layers.custom_op import CustomOp
+
+
+class RMSNormVAE(CustomOp):
+    """Root Mean Square Layer Normalization for Channel-First or Last"""
+
+    def __init__(
+        self,
+        dim: int,
+        channel_first: bool = True,
+        images: bool = True,
+        bias: bool = False,
+        epsilon: float = 1e-6,
+    ) -> None:
+        super().__init__()
+        broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+        shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+        self.channel_first = channel_first
+        self.scale = dim**0.5
+        self.gamma = nn.Parameter(torch.ones(shape))
+        self.bias = nn.Parameter(torch.zeros(shape)) if bias else None
+        self.epsilon = epsilon
+
+        self.gamma_rmsnorm = None
+
+    def forward_cuda(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_hip(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x)
+
+    def forward_npu(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        import torch_npu
+
+        if self.gamma_rmsnorm is None:
+            self.gamma_rmsnorm = self.gamma.reshape(-1)
+
+        if self.channel_first:
+            x = x.transpose(1, -1)
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1)
+        else:
+            out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0]
+
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+    def forward_native(
+        self,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        out = (
+            F.normalize(
+                x,
+                dim=(1 if self.channel_first else -1),
+                eps=self.epsilon,
+            )
+            * self.scale
+            * self.gamma
+        )
+        if self.bias is not None:
+            out = out + self.bias
+        return out

vllm_omni/diffusion/models/wan2_2/patch_diffusers.py

@@ -1,9 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import sys
 
-from vllm_omni.diffusion.layers.norm import RMSNormVAE
+from vllm_omni.diffusion.models.wan2_2.norm import RMSNormVAE
 
 
 def patch_wan_rms_norm():

vllm_omni/diffusion/models/wan2_2/rope.py

@@ -0,0 +1,142 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Wan S2V rotary positional embedding layers."""
+
+import torch
+
+from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan
+from vllm_omni.platforms import current_omni_platform
+
+
+class WanS2VRotaryPosEmbed(torch.nn.Module):
+    """Precompute complex-valued RoPE embeddings for S2V multi-grid positions.
+
+    Owns the base frequency buffer and provides forward() to compute position
+    embeddings given hidden_states (for shape) and grid_sizes.
+    """
+
+    def __init__(self, num_heads: int, head_dim: int, max_seq_len: int = 1024):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        d = head_dim
+        freqs = torch.cat(
+            [
+                self._rope_params(max_seq_len, d - 4 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+            ],
+            dim=1,
+        )
+        self.register_buffer("freqs", freqs.to(torch.complex64), persistent=False)
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def _rope_params(max_seq_len, dim, theta=10000):
+        if dim % 2 != 0:
+            raise ValueError(f"dim ({dim}) must be even")
+        freqs = torch.outer(
+            torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim))
+        )
+        return torch.polar(torch.ones_like(freqs), freqs)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        grid_sizes: list,
+        trainable_freqs: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Precompute RoPE embeddings for the given grid layout.
+
+        Args:
+            hidden_states: Tensor [B, S, ...] (used for batch/seq shape and device)
+            grid_sizes: Grid specification (list of [offsets, sizes, totals])
+            trainable_freqs: Optional trainable frequency overrides for t_f < 0
+
+        Returns:
+            Complex tensor [B, S, 1, head_dim//2] of precomputed position embeddings
+        """
+        b, s = hidden_states.shape[0], hidden_states.shape[1]
+        c = self.head_dim // 2
+        device = hidden_states.device
+
+        freqs = self.freqs.to(device)
+        if trainable_freqs is not None:
+            freqs_input = [freqs, trainable_freqs]
+        else:
+            freqs_input = freqs
+
+        if isinstance(freqs_input, list):
+            trainable_f = freqs_input[1]
+            freqs_split = freqs_input[0]
+        else:
+            trainable_f = None
+            freqs_split = freqs_input
+        freqs_split = freqs_split.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+        output = torch.empty((b, s, 1, c), device=device, dtype=torch.complex64)
+        seq_bucket = [0]
+        if not isinstance(grid_sizes, list):
+            grid_sizes = [grid_sizes]
+        for g in grid_sizes:
+            if not isinstance(g, list):
+                g = [torch.zeros_like(g), g]
+            batch_size = g[0].shape[0]
+            for i in range(batch_size):
+                f_o, h_o, w_o = g[0][i]
+                f, h, w = g[1][i]
+                t_f, t_h, t_w = g[2][i]
+                seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
+                seq_len = int(seq_f * seq_h * seq_w)
+                if seq_len > 0:
+                    if t_f > 0:
+                        assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
+                        seq_f_int = int(seq_f)
+                        seq_h_int = int(seq_h)
+                        seq_w_int = int(seq_w)
+
+                        if f_o >= 0:
+                            f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long()
+                        else:
+                            f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long()
+                        h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
+                        w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
+
+                        freqs_0 = torch.index_select(freqs_split[0] if f_o >= 0 else freqs_split[0].conj(), 0, f_sam)
+                        freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
+
+                        freqs_i = torch.cat(
+                            [
+                                freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                torch.index_select(freqs_split[1], 0, h_sam)
+                                .view(1, seq_h_int, 1, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                torch.index_select(freqs_split[2], 0, w_sam)
+                                .view(1, 1, seq_w_int, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                            ],
+                            dim=-1,
+                        ).reshape(seq_len, 1, -1)
+                    elif t_f < 0:
+                        freqs_i = trainable_f.unsqueeze(1)
+                    output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i
+            seq_bucket.append(seq_bucket[-1] + seq_len)
+        return output
+
+
+class RotaryEmbeddingWanS2V(RotaryEmbeddingWan):
+    """Apply RoPE using precomputed complex freqs for Wan S2V main transformer.
+
+    Converts complex freqs (from WanS2VRotaryPosEmbed) to cos/sin and delegates
+    to RotaryEmbeddingWan for platform-optimized application (float32 kernel).
+    Under TP, freqs has 1 head — broadcasts automatically via cos/sin.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(is_neox_style=False, half_head_dim=True)
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        freqs_sliced = freqs[:, : x.size(1)]
+        cos = freqs_sliced.real.to(x.dtype)
+        sin = freqs_sliced.imag.to(x.dtype)
+        return super().forward(x, cos, sin)

vllm_omni/diffusion/models/wan2_2/wan2_2_s2v_transformer.py

@@ -32,7 +32,8 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from vllm_omni.diffusion.attention.layer import Attention
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
-from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid, RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed
+from vllm_omni.diffusion.layers.rope import RotaryEmbeddingS2VGrid
+from vllm_omni.diffusion.models.wan2_2.rope import RotaryEmbeddingWanS2V, WanS2VRotaryPosEmbed
 from vllm_omni.platforms import current_omni_platform
 
 from .wan2_2_transformer import DistributedRMSNorm, WanFeedForward