Skip to content

vllm_omni.diffusion.layers.norm

logger module-attribute 

logger = init_logger(__name__)

LayerNorm

Bases: LayerNorm, CustomOp

LayerNorm implementation that inherits from both nn.LayerNorm and CustomOp. NPU: Uses mindiesd.fast_layernorm(self, x) when MindIE-SD is installed. CUDA / HIP / XPU / native: Falls back to FP32 nn.LayerNorm implementation.

forward 

forward(x: Tensor) -> Tensor

forward_cuda 

forward_cuda(x: Tensor) -> Tensor

forward_hip 

forward_hip(x: Tensor) -> Tensor

forward_native 

forward_native(x: Tensor) -> Tensor

forward_npu 

forward_npu(x: Tensor) -> Tensor

forward_xpu 

forward_xpu(x: Tensor) -> Tensor

RMSNorm

Bases: CustomOp

variance_epsilon instance-attribute 

variance_epsilon = eps

weight instance-attribute 

weight = Parameter(ones(hidden_size))

forward_cuda 

forward_cuda(x: Tensor) -> Tensor

forward_hip 

forward_hip(x: Tensor) -> Tensor

forward_native 

forward_native(x: Tensor) -> Tensor

forward_npu 

forward_npu(x: Tensor) -> Tensor

forward_xpu 

forward_xpu(x: Tensor) -> Tensor

RMSNormVAE

Bases: CustomOp

Root Mean Square Layer Normalization for Channel-First or Last

bias instance-attribute 

bias = Parameter(zeros(shape)) if bias else None

channel_first instance-attribute 

channel_first = channel_first

epsilon instance-attribute 

epsilon = epsilon

gamma instance-attribute 

gamma = Parameter(ones(shape))

gamma_rmsnorm instance-attribute 

gamma_rmsnorm = None

scale instance-attribute 

scale = dim ** 0.5

forward_cuda 

forward_cuda(x: Tensor) -> Tensor

forward_hip 

forward_hip(x: Tensor) -> Tensor

forward_native 

forward_native(x: Tensor) -> Tensor

forward_npu 

forward_npu(x: Tensor) -> Tensor

forward_xpu 

forward_xpu(x: Tensor) -> Tensor