@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dad95e34d8

Fork 0

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 2, 01:00 UTC

[Feature] Support batch invariant rms norm with residual (#42453)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Jun 3, 2026 at 15:22 UTC dad95e34d896d75badd1d67c203021fb59374b75

3 files changed +39 -41

tests/v1/determinism/test_rms_norm_batch_invariant.py

+11 -9

@@ -11,7 +11,9 @@ import pytest

11 import torch

12 from utils import skip_unsupported

14 -from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm

14 +from vllm.model_executor.layers.batch_invariant import (

15 + rms_norm_batch_invariant,

16 +)

17 from vllm.model_executor.layers.layernorm import RMSNorm

18 from vllm.platforms import current_platform

@@ -51,7 +53,7 @@ def test_rms_norm_batch_invariant_vs_standard(

53 standard_output = rms_norm_layer.forward_cuda(input_tensor)

55 # Batch-invariant implementation (Triton)

54 - triton_output = triton_rms_norm(input_tensor, weight, eps=eps)

56 + triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

58 # Compare outputs

59 # Use looser tolerance for bfloat16 due to its lower precision

@@ -125,7 +127,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(

127 )

128

129 merged_single = x_single + residual_single

128 - ref_out = triton_rms_norm(merged_single, weight, eps=eps)

130 + ref_out = rms_norm_batch_invariant(merged_single, weight, eps=eps)

131

132 torch.testing.assert_close(

133 residual_out_single,

@@ -193,7 +195,7 @@ def test_rms_norm_3d_input(

195 standard_output = rms_norm_layer.forward_cuda(input_tensor)

196

197 # Batch-invariant implementation

196 - triton_output = triton_rms_norm(input_tensor, weight, eps=eps)

198 + triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

199

200 # Use looser tolerance for bfloat16

201 rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16

@@ -242,7 +244,7 @@ def test_rms_norm_numerical_stability(default_vllm_config):

244 standard_output = rms_norm_layer.forward_cuda(input_tensor)

245

246 # Batch-invariant implementation

245 - triton_output = triton_rms_norm(input_tensor, weight, eps=eps)

247 + triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

248

249 # Check for NaN or Inf

250 assert not torch.isnan(standard_output).any(), (

@@ -289,7 +291,7 @@ def test_rms_norm_formula(default_vllm_config):

291 expected_output = input_tensor * torch.rsqrt(variance + eps) * weight

292

293 # Batch-invariant implementation

292 - triton_output = triton_rms_norm(input_tensor, weight, eps=eps)

294 + triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

295

296 # Compare against formula

297 torch.testing.assert_close(

@@ -325,7 +327,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):

327 standard_output = rms_norm_layer.forward_cuda(input_tensor)

328

329 # Batch-invariant implementation

328 - triton_output = triton_rms_norm(input_tensor, weight, eps=eps)

330 + triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

331

332 # Use looser tolerance for bfloat16

333 rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16

@@ -360,7 +362,7 @@ def test_rms_norm_determinism(default_vllm_config):

362 # Run multiple times

363 outputs = []

364 for _ in range(5):

363 - output = triton_rms_norm(input_tensor.clone(), weight, eps=eps)

365 + output = rms_norm_batch_invariant(input_tensor.clone(), weight, eps=eps)

366 outputs.append(output)

367

368 # All outputs should be identical

@@ -395,7 +397,7 @@ if __name__ == "__main__":

397 standard_output = rms_norm_layer.forward_cuda(input_tensor)

398

399 # Batch-invariant implementation

398 - triton_output = triton_rms_norm(input_tensor, weight, eps=eps)

400 + triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

401

402 # Compare

403 max_diff = (triton_output - standard_output).abs().max().item()

vllm/model_executor/layers/batch_invariant.py

+18 -26

@@ -822,23 +822,35 @@ def _rms_norm_kernel(

822 tl.store(output_row_start_ptr + col_idx, output, mask=mask)

823

824

825 -def rms_norm(

826 - input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6

827 -) -> torch.Tensor:

825 +def rms_norm_batch_invariant(

826 + input: torch.Tensor,

827 + weight: torch.Tensor,

828 + eps: float = 1e-6,

829 + residual: torch.Tensor | None = None,

830 +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

831 """

832 Compute RMS normalization using Triton kernel.

833

831 - RMS Norm normalizes the input by the root mean square and scales by weight:

832 - output = input / sqrt(mean(input^2) + eps) * weight

834

835 Args:

836 input: Input tensor of shape (..., hidden_size)

837 weight: Weight tensor of shape (hidden_size,)

838 eps: Small constant for numerical stability

839 + residual: Optional residual tensor fused into the normalization path

840

841 Returns:

840 - Tensor with RMS normalization applied along the last dimension

842 + RMS normalized tensor, or ``(output, residual_out)`` when ``residual``

843 + is provided

844 """

845 + if residual is not None:

846 + assert input.shape == residual.shape, (

847 + f"Input shape {input.shape} must match residual shape {residual.shape}"

848 + )

849 + import vllm._custom_ops as ops

850 +

851 + ops.fused_add_rms_norm(input, residual, weight, eps)

852 + return input, residual

853 +

854 assert weight.dim() == 1, "Weight must be 1-dimensional"

855 assert input.shape[-1] == weight.shape[0], (

856 f"Input last dimension ({input.shape[-1]}) must match "

@@ -869,26 +881,6 @@ def rms_norm(

881 return output.reshape(original_shape)

882

883

872 -def rms_norm_batch_invariant(

873 - input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6

874 -) -> torch.Tensor:

875 - """

876 - Batch-invariant wrapper for RMS normalization.

877 -

878 - This function provides a deterministic, batch-invariant implementation

879 - of RMS normalization for use with the batch_invariant mode.

880 -

881 - Args:

882 - input: Input tensor of shape (..., hidden_size)

883 - weight: Weight tensor of shape (hidden_size,)

884 - eps: Small constant for numerical stability

885 -

886 - Returns:

887 - RMS normalized tensor

888 - """

889 - return rms_norm(input, weight, eps=eps)

890 -

891 -

884 def linear_batch_invariant(input, weight, bias=None):

885 output = matmul_batch_invariant(input, weight.t())

886

vllm/model_executor/layers/layernorm.py

+10 -6

@@ -106,12 +106,16 @@ class RMSNorm(CustomOp):

106 x: torch.Tensor,

107 residual: torch.Tensor | None = None,

108 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

109 - if (

110 - envs.VLLM_BATCH_INVARIANT

111 - and residual is None

112 - and self.variance_size_override is None

113 - ):

114 - return rms_norm_batch_invariant(x, self.weight.data, self.variance_epsilon)

109 + if envs.VLLM_BATCH_INVARIANT:

110 + assert self.variance_size_override is None, (

111 + "Batch invariance is not supported for variance_size_override"

112 + )

113 + return rms_norm_batch_invariant(

114 + x,

115 + self.weight.data,

116 + self.variance_epsilon,

117 + residual=residual,

118 + )

119

120 return self.forward_native(x, residual)

121