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 4, 21:15 UTC
[Perf] Optimize batch invariant with fused rms norm, 2.1% E2E latency improvement (#40413)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 21, 2026 at 15:51 UTC
16688b26a6fbabe0100c440462874ad4e4c78b16
3 files changed
+89
-5
tests/v1/determinism/test_rms_norm_batch_invariant.py
+88
-1
@@ -12,7 +12,7 @@ import torch
12
from utils import skip_unsupported
13
14
from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
15
-from vllm.model_executor.layers.layernorm import RMSNorm
15
+from vllm.model_executor.layers.layernorm import RMSNorm, fused_add_rms_norm
16
from vllm.platforms import current_platform
17
18
DEVICE_TYPE = current_platform.device_type
@@ -71,6 +71,93 @@ def test_rms_norm_batch_invariant_vs_standard(
71
)
72
73
74
+@skip_unsupported
75
+@pytest.mark.parametrize("hidden_size", [512, 4096])
76
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
77
+@pytest.mark.parametrize("eps", [1e-6])
78
+def test_fused_add_rms_norm_batch_invariant_residual_path(
79
+ hidden_size: int,
80
+ dtype: torch.dtype,
81
+ eps: float,
82
+):
83
+ """
84
+ Test the batch-invariant fused residual-add + RMSNorm helper directly.
85
+ """
86
+ device = torch.device(DEVICE_TYPE)
87
+
88
+ torch.manual_seed(42)
89
+ x_single = torch.randn(1, hidden_size, dtype=dtype, device=device)
90
+ residual_single = torch.randn(1, hidden_size, dtype=dtype, device=device)
91
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
92
+
93
+ x_batch = torch.cat(
94
+ [
95
+ x_single,
96
+ torch.randn(3, hidden_size, dtype=dtype, device=device),
97
+ ],
98
+ dim=0,
99
+ )
100
+ residual_batch = torch.cat(
101
+ [
102
+ residual_single,
103
+ torch.randn(3, hidden_size, dtype=dtype, device=device),
104
+ ],
105
+ dim=0,
106
+ )
107
+
108
+ out_single, residual_out_single = fused_add_rms_norm(
109
+ x_single.clone(),
110
+ residual_single.clone(),
111
+ weight,
112
+ eps,
113
+ )
114
+ out_batch, residual_out_batch = fused_add_rms_norm(
115
+ x_batch.clone(),
116
+ residual_batch.clone(),
117
+ weight,
118
+ eps,
119
+ )
120
+
121
+ merged_single = x_single + residual_single
122
+ ref_out = triton_rms_norm(merged_single, weight, eps=eps)
123
+
124
+ torch.testing.assert_close(
125
+ residual_out_single,
126
+ merged_single,
127
+ rtol=0.0,
128
+ atol=0.0,
129
+ msg="Residual output should equal x + residual exactly",
130
+ )
131
+ torch.testing.assert_close(
132
+ residual_out_batch[:1],
133
+ merged_single,
134
+ rtol=0.0,
135
+ atol=0.0,
136
+ msg="Residual output should be batch invariant",
137
+ )
138
+ torch.testing.assert_close(
139
+ out_single,
140
+ out_batch[:1],
141
+ rtol=0.0,
142
+ atol=0.0,
143
+ msg="Fused add RMSNorm output should be batch invariant",
144
+ )
145
+
146
+ if dtype == torch.bfloat16:
147
+ rtol, atol = 1e-1, 1e-1
148
+ else:
149
+ rtol, atol = 1e-2, 1e-2
150
+
151
+ torch.testing.assert_close(
152
+ out_single,
153
+ ref_out,
154
+ rtol=rtol,
155
+ atol=atol,
156
+ msg="Fused add RMSNorm output should stay numerically close to the "
157
+ "batch-invariant RMSNorm reference",
158
+ )
159
+
160
+
161
@skip_unsupported
162
@pytest.mark.parametrize("batch_size", [1, 16, 128])
163
@pytest.mark.parametrize("seq_len", [1, 32, 512])
vllm/_custom_ops.py
+1
@@ -420,6 +420,7 @@ def rms_norm(
420
def fused_add_rms_norm(
421
input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, epsilon: float
422
) -> None:
423
+ # Note: this func is batch invariant
424
torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
425
426
vllm/model_executor/layers/layernorm.py
-4
@@ -61,10 +61,6 @@ def fused_add_rms_norm(
61
) -> tuple[torch.Tensor, torch.Tensor]:
62
from vllm import _custom_ops as ops
63
64
- if envs.VLLM_BATCH_INVARIANT:
65
- return rms_norm_batch_invariant(
66
- x + residual, weight, variance_epsilon
67
- ), x + residual
64
ops.fused_add_rms_norm(
65
x,
66
residual,