diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index 55bfbab00..e42041528 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/tests/diffusion/layers/test_norm.py @@ -451,97 +451,3 @@ 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 2f320c2a6..f397c1a85 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -70,33 +70,17 @@ 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: - try: - return self._forward_fused(x) - except Exception: - return self.forward_native(x) + return self.forward_native(x) def forward_hip( self, x: torch.Tensor, ) -> torch.Tensor: - try: - return self._forward_fused(x) - except Exception: - return self.forward_native(x) + 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 527c472eb..80daec6b8 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -25,6 +25,7 @@ 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, @@ -33,8 +34,6 @@ 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,