diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index c1ba96a792..85ec4903ee 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random +from unittest.mock import patch import pytest import torch @@ -236,6 +237,10 @@ def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None: torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2) +@pytest.mark.parametrize( + ("alpha", "beta"), + [(1.0, 0.0), (1.702, 0.0), (1.0, 1.0)], +) @pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS) @pytest.mark.parametrize("num_tokens", NUM_TOKENS) @pytest.mark.parametrize("d", D) @@ -245,6 +250,8 @@ def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None: @torch.inference_mode() def test_silu_and_mul_with_clamp( default_vllm_config, + alpha: float, + beta: float, swiglu_limit: float, num_tokens: int, d: int, @@ -258,8 +265,38 @@ def test_silu_and_mul_with_clamp( # Use large values to ensure clamping is exercised. x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2 - layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False) - out = layer(x) + default_vllm_config.compilation_config.custom_ops = [ + "none", + "+silu_and_mul_with_clamp", + ] + layer = SiluAndMulWithClamp( + swiglu_limit, + alpha=alpha, + beta=beta, + compile_native=False, + ) + if current_platform.is_rocm(): + # forward_hip is always dispatched; the alpha/beta gate is checked + # inside it at call time rather than picked at construction time, so + # verify the actual routing by spying on the two candidate methods. + assert layer._forward_method == layer.forward_hip + with ( + patch.object(layer, "forward_cuda", wraps=layer.forward_cuda) as cuda_spy, + patch.object( + layer, "forward_native", wraps=layer.forward_native + ) as native_spy, + ): + out = layer(x) + if alpha == 1.0 and beta == 0.0: + cuda_spy.assert_called_once() + native_spy.assert_not_called() + else: + native_spy.assert_called_once() + cuda_spy.assert_not_called() + else: + assert layer._forward_method == layer.forward_cuda + out = layer(x) + ref_out = layer.forward_native(x) rtol = { @@ -273,10 +310,11 @@ def test_silu_and_mul_with_clamp( # Verify clamping is actually being applied: the clamped output should # differ from the unclamped SiluAndMul output when inputs are large. - unclamped_out = SiluAndMul.forward_native(x) - assert not torch.equal(ref_out.float(), unclamped_out.float()), ( - "Input was not large enough to exercise the clamp; increase scale" - ) + if alpha == 1.0 and beta == 0.0: + unclamped_out = SiluAndMul.forward_native(x) + assert not torch.equal(ref_out.float(), unclamped_out.float()), ( + "Input was not large enough to exercise the clamp; increase scale" + ) # Verify gate clamping semantics with a controlled scalar case. # gate=large_val is clamped to limit first, then silu(limit) * 1.0. @@ -309,7 +347,10 @@ def test_silu_and_mul_with_clamp( # opcheck out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device) - opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit)) + opcheck( + torch.ops._C.silu_and_mul_with_clamp, + (out_buf, x, swiglu_limit, layer.alpha, layer.beta), + ) @pytest.mark.parametrize("linear_beta", [-1.0, 2.0]) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index 35a6fdf9b0..808d6e7112 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -233,12 +233,8 @@ class SiluAndMulWithClamp(CustomOp): self.swiglu_limit = float(swiglu_limit) self.alpha = float(alpha) self.beta = float(beta) - if current_platform.is_rocm() or current_platform.is_xpu(): - self._forward_method = self.forward_native - elif current_platform.is_cuda_alike(): + if current_platform.is_cuda_alike(): self.op = torch.ops._C.silu_and_mul_with_clamp - elif current_platform.is_cpu(): - self._forward_method = self.forward_native def forward_native(self, x: torch.Tensor) -> torch.Tensor: d = x.shape[-1] // 2 @@ -253,7 +249,11 @@ class SiluAndMulWithClamp(CustomOp): self.op(out, x, self.swiglu_limit, self.alpha, self.beta) return out - def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + def forward_hip(self, x: torch.Tensor) -> torch.Tensor: + # Limit the ROCm _C path to the default case to avoid the + # precision loss MiniMax saw without fp32 intermediates. + if self.alpha == 1.0 and self.beta == 0.0: + return self.forward_cuda(x) return self.forward_native(x) def extra_repr(self) -> str: