diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index e42041528..55bfbab00 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/tests/diffusion/layers/test_norm.py @@ -451,3 +451,97 @@ def test_rmsnorm_with_single_element_batch(): out = norm(x) assert out.shape == (1, 1, hidden_size) + + +# ── Fused vs Native correctness tests ── + + +def test_rmsnorm_fused_native_parity_cuda(): + """Verify fused and native RMSNorm produce identical results on CUDA. + + The fused kernel uses vllm._custom_ops.rms_norm while native uses pure PyTorch. + Both should produce numerically equivalent results. + """ + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 128 + eps = 1e-6 + torch.manual_seed(42) + + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + torch.manual_seed(42) + norm = RMSNorm(hidden_size, eps=eps).cuda() + norm.weight.data = norm.weight.data.to(dtype) + + x = torch.randn(4, 8, hidden_size, dtype=dtype, device="cuda") + + out_fused = norm._forward_fused(x) + out_native = norm.forward_native(x) + + torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3) + + +def test_rmsnorm_fused_native_parity_large_tensor(): + """Verify fused and native RMSNorm parity with large tensors. + + Tests realistic tensor sizes used in diffusion models. + """ + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 4096 + batch_size = 16 + seq_len = 1024 + eps = 1e-6 + + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + torch.manual_seed(123) + norm = RMSNorm(hidden_size, eps=eps).cuda() + norm.weight.data = norm.weight.data.to(dtype) + + x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") + + out_fused = norm._forward_fused(x) + out_native = norm.forward_native(x) + + torch.testing.assert_close(out_fused, out_native, atol=1e-2, rtol=1e-2) + + +def test_rmsnorm_fused_native_parity_2d_input(): + """Verify fused and native RMSNorm parity with 2D input tensors.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 256 + batch_size = 32 + eps = 1e-6 + + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + torch.manual_seed(456) + norm = RMSNorm(hidden_size, eps=eps).cuda() + norm.weight.data = norm.weight.data.to(dtype) + + x = torch.randn(batch_size, hidden_size, dtype=dtype, device="cuda") + + out_fused = norm._forward_fused(x) + out_native = norm.forward_native(x) + + torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3) + + +def test_rmsnorm_fused_native_parity_custom_weight(): + """Verify fused and native RMSNorm parity with non-default weights.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + eps = 1e-6 + + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + torch.manual_seed(789) + norm = RMSNorm(hidden_size, eps=eps).cuda() + norm.weight.data = torch.randn(hidden_size, dtype=dtype, device="cuda") + + x = torch.randn(2, 4, hidden_size, dtype=dtype, device="cuda") + + out_fused = norm._forward_fused(x) + out_native = norm.forward_native(x) + + torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3) diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index f397c1a85..2f320c2a6 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -70,17 +70,33 @@ class RMSNorm(CustomOp): self.weight = nn.Parameter(torch.ones(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: - 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: - return self.forward_native(x) + try: + return self._forward_fused(x) + except Exception: + return self.forward_native(x) def forward_npu( self, diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 80daec6b8..527c472eb 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -25,7 +25,6 @@ import torch.nn as nn from torch.nn.utils.rnn import pad_sequence from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( MergedColumnParallelLinear, QKVParallelLinear, @@ -34,6 +33,8 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.layers.norm import RMSNorm + if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig,