228 lines
6.8 KiB
Python
228 lines
6.8 KiB
Python
from importlib.util import find_spec
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from vllm.logger import init_logger
|
|
|
|
from vllm_omni.diffusion.layers.custom_op import CustomOp
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
_HAS_MINDIESD = find_spec("mindiesd") is not None
|
|
|
|
|
|
class LayerNorm(nn.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.
|
|
"""
|
|
|
|
def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True):
|
|
super().__init__(normalized_shape=dim, eps=eps, elementwise_affine=elementwise_affine)
|
|
# CustomOp.__init__ cannot be called here because it would re-run
|
|
# nn.Module initialization and clear LayerNorm parameters.
|
|
self._forward_method = CustomOp.dispatch_forward(self)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self._forward_method(x)
|
|
|
|
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_xpu(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.forward_native(x)
|
|
|
|
def forward_npu(self, x: torch.Tensor) -> torch.Tensor:
|
|
if _HAS_MINDIESD:
|
|
try:
|
|
from mindiesd import fast_layernorm
|
|
|
|
return fast_layernorm(self, x)
|
|
except ImportError as e:
|
|
logger.warning_once(
|
|
"mindiesd.fast_layernorm import failed, falling back to FP32 layer_norm: %s",
|
|
e,
|
|
)
|
|
|
|
return self.forward_native(x)
|
|
|
|
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
|
|
origin_dtype = x.dtype
|
|
return F.layer_norm(
|
|
x.float(),
|
|
self.normalized_shape,
|
|
self.weight.float() if self.weight is not None else None,
|
|
self.bias.float() if self.bias is not None else None,
|
|
self.eps,
|
|
).to(origin_dtype)
|
|
|
|
|
|
class RMSNorm(CustomOp):
|
|
def __init__(self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.float32) -> None:
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(hidden_size, dtype=dtype))
|
|
self.hidden_size = hidden_size
|
|
self.variance_epsilon = eps
|
|
|
|
def _forward_fused(self, x: torch.Tensor) -> torch.Tensor:
|
|
from vllm._custom_ops import rms_norm as fused_rms_norm
|
|
|
|
orig_shape = x.shape
|
|
hidden_size = orig_shape[-1]
|
|
x_2d = x.reshape(-1, hidden_size)
|
|
out = torch.empty_like(x_2d)
|
|
fused_rms_norm(out, x_2d, self.weight.data, self.variance_epsilon)
|
|
return out.reshape(orig_shape)
|
|
|
|
def forward_cuda(
|
|
self,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
# During torch.compile tracing, fused_rms_norm writes to `out` in-place
|
|
# (returns None) and accesses self.weight.data, which is a DTensor under
|
|
# HSDP. Both patterns confuse inductor's compute_ancestors scheduler.
|
|
# Fall back to forward_native so inductor can fuse the pure-PyTorch ops
|
|
# itself.
|
|
if torch.compiler.is_compiling():
|
|
return self.forward_native(x)
|
|
try:
|
|
return self._forward_fused(x)
|
|
except Exception:
|
|
return self.forward_native(x)
|
|
|
|
def forward_hip(
|
|
self,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
if torch.compiler.is_compiling():
|
|
return self.forward_native(x)
|
|
try:
|
|
return self._forward_fused(x)
|
|
except Exception:
|
|
return self.forward_native(x)
|
|
|
|
def forward_musa(
|
|
self,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
# Preserve the aten::rms_norm graph so dynamic Inductor can fuse the
|
|
# H3 Q/K norm with its inline RoPE path on MUSA.
|
|
return F.rms_norm(
|
|
x,
|
|
(self.hidden_size,),
|
|
self.weight,
|
|
self.variance_epsilon,
|
|
)
|
|
|
|
def forward_npu(
|
|
self,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
import torch_npu
|
|
|
|
output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0]
|
|
|
|
return output
|
|
|
|
def forward_xpu(
|
|
self,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return self.forward_native(x)
|
|
|
|
def forward_native(
|
|
self,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
input_dtype = x.dtype
|
|
x = x.to(torch.float32)
|
|
variance = x.pow(2).mean(-1, keepdim=True)
|
|
out = x * torch.rsqrt(variance + self.variance_epsilon)
|
|
out = self.weight.to(torch.float32) * out
|
|
return out.to(input_dtype)
|
|
|
|
|
|
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_xpu(
|
|
self,
|
|
x: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
return self.forward_native(x)
|
|
|
|
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
|