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

RMSNorm

Bases: CustomOp

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, dtype=weight_dtype)
)

forward_cuda 

forward_cuda(
    x: Tensor, residual: Tensor | None = None
) -> Tensor | tuple[Tensor, Tensor]

forward_hip 

forward_hip(
    x: Tensor, residual: Tensor | None = None
) -> Tensor | tuple[Tensor, Tensor]

forward_musa 

forward_musa(
    x: Tensor, residual: Tensor | None = None
) -> Tensor | tuple[Tensor, Tensor]

forward_native 

forward_native(
    x: Tensor, residual: Tensor | None = None
) -> Tensor | tuple[Tensor, Tensor]

forward_npu 

forward_npu(
    x: Tensor, residual: Tensor | None = None
) -> Tensor | tuple[Tensor, Tensor]