96 lines
3.3 KiB
Python
96 lines
3.3 KiB
Python
|
# 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}"
|