Skip to content

vllm_omni.diffusion.models.bagel.mot.mot_layernorm

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).

MoTRMSNorm

Bases: 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.

gen_weight instance-attribute 

gen_weight = nn.Parameter(torch.ones(hidden_size))

head_norm instance-attribute 

head_norm = head_norm

hidden_size instance-attribute 

hidden_size = hidden_size

variance_epsilon instance-attribute 

variance_epsilon = eps

weight instance-attribute 

weight = nn.Parameter(torch.ones(hidden_size))

extra_repr 

extra_repr() -> str

forward_cuda 

forward_cuda(
    x: Tensor,
    text_indices: Tensor | None = None,
    vae_indices: Tensor | None = None,
) -> Tensor

forward_native 

forward_native(
    x: Tensor,
    text_indices: Tensor | None = None,
    vae_indices: Tensor | None = None,
) -> Tensor

forward_npu 

forward_npu(
    x: Tensor,
    text_indices: Tensor | None = None,
    vae_indices: Tensor | None = None,
) -> Tensor