[ROCm] Use silu_and_mul_with_clamp's torch._C op (#52052)
Signed-off-by: Tres Popp <tres.popp@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Shanshan Shen <467638484@qq.com> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Tres committed
Sep 23, 2026 at 01:30 UTC
d110c2f19c4f47b92b2dbce2ab24f2014589bf0c
2 files changed
+54
-13
tests/kernels/core/test_activation.py
+48
-7
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import random
5
+from unittest.mock import patch
6
7
import pytest
8
import torch
@@ -236,6 +237,10 @@ def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None:
237
torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
238
239
240
+@pytest.mark.parametrize(
241
+ ("alpha", "beta"),
242
+ [(1.0, 0.0), (1.702, 0.0), (1.0, 1.0)],
243
+)
244
@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
245
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
246
@pytest.mark.parametrize("d", D)
@@ -245,6 +250,8 @@ def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None:
250
@torch.inference_mode()
251
def test_silu_and_mul_with_clamp(
252
default_vllm_config,
253
+ alpha: float,
254
+ beta: float,
255
swiglu_limit: float,
256
num_tokens: int,
257
d: int,
@@ -258,8 +265,38 @@ def test_silu_and_mul_with_clamp(
265
# Use large values to ensure clamping is exercised.
266
x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
267
261
- layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
262
- out = layer(x)
268
+ default_vllm_config.compilation_config.custom_ops = [
269
+ "none",
270
+ "+silu_and_mul_with_clamp",
271
+ ]
272
+ layer = SiluAndMulWithClamp(
273
+ swiglu_limit,
274
+ alpha=alpha,
275
+ beta=beta,
276
+ compile_native=False,
277
+ )
278
+ if current_platform.is_rocm():
279
+ # forward_hip is always dispatched; the alpha/beta gate is checked
280
+ # inside it at call time rather than picked at construction time, so
281
+ # verify the actual routing by spying on the two candidate methods.
282
+ assert layer._forward_method == layer.forward_hip
283
+ with (
284
+ patch.object(layer, "forward_cuda", wraps=layer.forward_cuda) as cuda_spy,
285
+ patch.object(
286
+ layer, "forward_native", wraps=layer.forward_native
287
+ ) as native_spy,
288
+ ):
289
+ out = layer(x)
290
+ if alpha == 1.0 and beta == 0.0:
291
+ cuda_spy.assert_called_once()
292
+ native_spy.assert_not_called()
293
+ else:
294
+ native_spy.assert_called_once()
295
+ cuda_spy.assert_not_called()
296
+ else:
297
+ assert layer._forward_method == layer.forward_cuda
298
+ out = layer(x)
299
+
300
ref_out = layer.forward_native(x)
301
302
rtol = {
@@ -273,10 +310,11 @@ def test_silu_and_mul_with_clamp(
310
311
# Verify clamping is actually being applied: the clamped output should
312
# differ from the unclamped SiluAndMul output when inputs are large.
276
- unclamped_out = SiluAndMul.forward_native(x)
277
- assert not torch.equal(ref_out.float(), unclamped_out.float()), (
278
- "Input was not large enough to exercise the clamp; increase scale"
279
- )
313
+ if alpha == 1.0 and beta == 0.0:
314
+ unclamped_out = SiluAndMul.forward_native(x)
315
+ assert not torch.equal(ref_out.float(), unclamped_out.float()), (
316
+ "Input was not large enough to exercise the clamp; increase scale"
317
+ )
318
319
# Verify gate clamping semantics with a controlled scalar case.
320
# gate=large_val is clamped to limit first, then silu(limit) * 1.0.
@@ -309,7 +347,10 @@ def test_silu_and_mul_with_clamp(
347
348
# opcheck
349
out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
312
- opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
350
+ opcheck(
351
+ torch.ops._C.silu_and_mul_with_clamp,
352
+ (out_buf, x, swiglu_limit, layer.alpha, layer.beta),
353
+ )
354
355
356
@pytest.mark.parametrize("linear_beta", [-1.0, 2.0])
vllm/model_executor/layers/activation.py
+6
-6
@@ -233,12 +233,8 @@ class SiluAndMulWithClamp(CustomOp):
233
self.swiglu_limit = float(swiglu_limit)
234
self.alpha = float(alpha)
235
self.beta = float(beta)
236
- if current_platform.is_rocm() or current_platform.is_xpu():
237
- self._forward_method = self.forward_native
238
- elif current_platform.is_cuda_alike():
236
+ if current_platform.is_cuda_alike():
237
self.op = torch.ops._C.silu_and_mul_with_clamp
240
- elif current_platform.is_cpu():
241
- self._forward_method = self.forward_native
238
239
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
240
d = x.shape[-1] // 2
@@ -253,7 +249,11 @@ class SiluAndMulWithClamp(CustomOp):
249
self.op(out, x, self.swiglu_limit, self.alpha, self.beta)
250
return out
251
256
- def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
252
+ def forward_hip(self, x: torch.Tensor) -> torch.Tensor:
253
+ # Limit the ROCm _C path to the default case to avoid the
254
+ # precision loss MiniMax saw without fp32 intermediates.
255
+ if self.alpha == 1.0 and self.beta == 0.0:
256
+ return self.forward_cuda(x)
257
return self.forward_native(x)
258
259
def extra_repr(self) -> str: