A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 02:45 UTC
[CustomOp] CustomOp FusedRMSNormGated (#35877)
Signed-off-by: Elias Ellison <elias.ellison@gmail.com> Signed-off-by: eellison <elias.ellison@gmail.com>
eellison committed
Mar 6, 2026 at 13:53 UTC
f3c6c9c9d794fac5e74b59bc75da6e9d1921eeac
2 files changed
+133
-2
tests/kernels/core/test_fused_rms_norm_gated.py
new
+103
@@ -0,0 +1,103 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Tests that FusedRMSNormGated decomposes correctly under torch.compile,
5
+matching the eager triton kernel output."""
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.model_executor.layers.fla.ops.kda import FusedRMSNormGated
11
+from vllm.utils.torch_utils import set_random_seed
12
+
13
+DTYPES = [torch.bfloat16]
14
+HIDDEN_SIZES = [128, 512]
15
+NUM_TOKENS = [64, 128]
16
+ACTIVATIONS = ["swish", "sigmoid"]
17
+ELEMENTWISE_AFFINE = [True, False]
18
+SEEDS = [0]
19
+
20
+
21
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
22
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
23
+@pytest.mark.parametrize("activation", ACTIVATIONS)
24
+@pytest.mark.parametrize("elementwise_affine", ELEMENTWISE_AFFINE)
25
+@pytest.mark.parametrize("dtype", DTYPES)
26
+@pytest.mark.parametrize("seed", SEEDS)
27
+@torch.inference_mode()
28
+def test_compiled_vs_eager(
29
+ default_vllm_config,
30
+ num_tokens: int,
31
+ hidden_size: int,
32
+ activation: str,
33
+ elementwise_affine: bool,
34
+ dtype: torch.dtype,
35
+ seed: int,
36
+) -> None:
37
+ """forward_native decomposition matches forward_cuda triton kernel."""
38
+ torch._dynamo.reset()
39
+ set_random_seed(seed)
40
+ device = torch.device("cuda:0")
41
+
42
+ module = FusedRMSNormGated(
43
+ hidden_size,
44
+ elementwise_affine=elementwise_affine,
45
+ eps=1e-5,
46
+ activation=activation,
47
+ device=device,
48
+ dtype=dtype,
49
+ )
50
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
51
+ g = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
52
+
53
+ # forward_cuda may modify x in-place, so clone inputs
54
+ cuda_out = module.forward_cuda(x.clone(), g.clone())
55
+ compiled_native = torch.compile(module.forward_native, fullgraph=True)
56
+ native_out = compiled_native(x.clone(), g.clone())
57
+
58
+ torch.testing.assert_close(native_out, cuda_out, atol=1e-3, rtol=1e-2)
59
+
60
+
61
+@pytest.mark.parametrize(
62
+ "shape",
63
+ [
64
+ (1, 16, 32, 128),
65
+ (2, 8, 16, 64),
66
+ ],
67
+)
68
+@pytest.mark.parametrize("activation", ACTIVATIONS)
69
+@pytest.mark.parametrize("elementwise_affine", ELEMENTWISE_AFFINE)
70
+@pytest.mark.parametrize("dtype", DTYPES)
71
+@pytest.mark.parametrize("seed", SEEDS)
72
+@torch.inference_mode()
73
+def test_compiled_vs_eager_multidim(
74
+ default_vllm_config,
75
+ shape: tuple,
76
+ activation: str,
77
+ elementwise_affine: bool,
78
+ dtype: torch.dtype,
79
+ seed: int,
80
+) -> None:
81
+ """forward_native decomposition handles multi-dimensional inputs."""
82
+ torch._dynamo.reset()
83
+ set_random_seed(seed)
84
+ device = torch.device("cuda:0")
85
+ head_dim = shape[-1]
86
+
87
+ module = FusedRMSNormGated(
88
+ head_dim,
89
+ elementwise_affine=elementwise_affine,
90
+ eps=1e-5,
91
+ activation=activation,
92
+ device=device,
93
+ dtype=dtype,
94
+ )
95
+ x = torch.randn(*shape, dtype=dtype, device=device)
96
+ g = torch.randn(*shape, dtype=dtype, device=device)
97
+
98
+ # forward_cuda may modify x in-place, so clone inputs
99
+ cuda_out = module.forward_cuda(x.clone(), g.clone())
100
+ compiled_native = torch.compile(module.forward_native, fullgraph=True)
101
+ native_out = compiled_native(x.clone(), g.clone())
102
+
103
+ torch.testing.assert_close(native_out, cuda_out, atol=1e-3, rtol=1e-2)
vllm/model_executor/layers/fla/ops/kda.py
+30
-2
@@ -12,6 +12,7 @@
12
import torch
13
import torch.nn as nn
14
15
+from vllm.model_executor.custom_op import CustomOp
16
from vllm.triton_utils import tl, triton
17
from vllm.utils.math_utils import cdiv, next_power_of_2
18
@@ -431,7 +432,8 @@ def rms_norm_gated(
432
return y if not prenorm else (y, residual_out.reshape(x_shape_og))
433
434
434
-class FusedRMSNormGated(nn.Module):
435
+@CustomOp.register("fused_rms_norm_gated")
436
+class FusedRMSNormGated(CustomOp):
437
def __init__(
438
self,
439
hidden_size: int,
@@ -458,7 +460,33 @@ class FusedRMSNormGated(nn.Module):
460
self.register_parameter("weight", None)
461
self.register_parameter("bias", None)
462
461
- def forward(
463
+ def forward_native(
464
+ self,
465
+ x: torch.Tensor,
466
+ g: torch.Tensor,
467
+ residual: torch.Tensor | None = None,
468
+ prenorm: bool = False,
469
+ residual_in_fp32: bool = False,
470
+ ) -> torch.Tensor:
471
+ """Decomposed PyTorch ops for torch.compile/inductor fusion."""
472
+ # TODO(https://github.com/vllm-project/vllm/issues/36175): implement
473
+ # native residual/prenorm path and unify with RMSNormGated.
474
+ # For now, fall back to the triton kernel.
475
+ if residual is not None or prenorm:
476
+ return self.forward_cuda(x, g, residual, prenorm, residual_in_fp32)
477
+ x_float = x.float()
478
+ variance = x_float.pow(2).mean(dim=-1, keepdim=True)
479
+ x_normed = x_float * torch.rsqrt(variance + self.eps)
480
+ if self.weight is not None:
481
+ x_normed = x_normed * self.weight.float()
482
+ g_float = g.float()
483
+ if self.activation in ("swish", "silu"):
484
+ out = x_normed * g_float * torch.sigmoid(g_float)
485
+ else: # sigmoid
486
+ out = x_normed * torch.sigmoid(g_float)
487
+ return out.to(x.dtype)
488
+
489
+ def forward_cuda(
490
self,
491
x: torch.Tensor,
492
g: torch.Tensor,