b354734d17 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -1,11 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""

4 -Test batch-invariant RMS normalization against standard implementations.

5 -

6 -This test compares the Triton-based batch-invariant RMS norm implementation

7 -with the standard CUDA-based implementation to ensure numerical accuracy.

8 -"""

3 +"""Test batch-invariant RMS normalization against a PyTorch reference."""

4

5 import pytest

6 import torch

@@ -14,18 +9,30 @@ from utils import skip_if_not_cuda, skip_unsupported

9 from vllm.model_executor.layers.batch_invariant import (

10 rms_norm_batch_invariant,

11 )

17 -from vllm.model_executor.layers.layernorm import RMSNorm

12 from vllm.platforms import current_platform

13

14 DEVICE_TYPE = current_platform.device_type

15

16

17 +def _rms_norm_reference(

18 + input_tensor: torch.Tensor,

19 + weight: torch.Tensor,

20 + eps: float,

21 +) -> torch.Tensor:

22 + """Compute RMSNorm independently using PyTorch operations."""

23 + input_fp32 = input_tensor.float()

24 + output = input_fp32 * torch.rsqrt(

25 + input_fp32.square().mean(dim=-1, keepdim=True) + eps

26 + )

27 + return (output * weight.float()).to(input_tensor.dtype)

28 +

29 +

30 @skip_if_not_cuda

31 @pytest.mark.parametrize("batch_size", [1, 4, 16, 64])

32 @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])

33 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

34 @pytest.mark.parametrize("eps", [1e-6, 1e-5])

28 -def test_rms_norm_batch_invariant_vs_standard(

35 +def test_rms_norm_batch_invariant_vs_reference(

36 default_vllm_config,

37 batch_size: int,

38 hidden_size: int,

@@ -33,11 +40,10 @@ def test_rms_norm_batch_invariant_vs_standard(

40 eps: float,

41 ):

42 """

36 - Compare batch-invariant Triton RMS norm against standard CUDA implementation.

43 + Compare batch-invariant Triton RMS norm against a PyTorch reference.

44

45 Tests that the Triton-based batch-invariant RMS norm produces numerically

39 - equivalent results to the standard CUDA implementation across various

40 - configurations.

46 + equivalent results to an independent implementation across configurations.

47 """

48 device = torch.device(DEVICE_TYPE)

49

@@ -46,11 +52,7 @@ def test_rms_norm_batch_invariant_vs_standard(

52 input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)

53 weight = torch.randn(hidden_size, dtype=dtype, device=device)

54

49 - # Standard implementation (CUDA ops)

50 - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)

51 - rms_norm_layer.weight.data = weight.clone()

52 -

53 - standard_output = rms_norm_layer.forward_cuda(input_tensor)

55 + reference_output = _rms_norm_reference(input_tensor, weight, eps)

56

57 # Batch-invariant implementation (Triton)

58 triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

@@ -64,7 +66,7 @@ def test_rms_norm_batch_invariant_vs_standard(

66

67 torch.testing.assert_close(

68 triton_output,

67 - standard_output,

69 + reference_output,

70 rtol=rtol,

71 atol=atol,

72 msg=f"RMS norm mismatch for batch_size={batch_size}, "

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

129 )

130

131 merged_single = x_single + residual_single

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

132 + ref_out = _rms_norm_reference(merged_single, weight, eps)

133

134 torch.testing.assert_close(

135 residual_out_single,

@@ -162,7 +164,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(

164 rtol=rtol,

165 atol=atol,

166 msg="Fused add RMSNorm output should stay numerically close to the "

165 - "batch-invariant RMSNorm reference",

167 + "PyTorch RMSNorm reference",

168 )

169

170

@@ -189,10 +191,7 @@ def test_rms_norm_3d_input(

191 )

192 weight = torch.randn(hidden_size, dtype=dtype, device=device)

193

192 - # Standard implementation

193 - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)

194 - rms_norm_layer.weight.data = weight.clone()

195 - standard_output = rms_norm_layer.forward_cuda(input_tensor)

194 + reference_output = _rms_norm_reference(input_tensor, weight, eps)

195

196 # Batch-invariant implementation

197 triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

@@ -202,7 +201,7 @@ def test_rms_norm_3d_input(

201

202 torch.testing.assert_close(

203 triton_output,

205 - standard_output,

204 + reference_output,

205 rtol=rtol,

206 atol=atol,

207 msg=f"RMS norm mismatch for 3D input with batch_size={batch_size}, "

@@ -238,20 +237,17 @@ def test_rms_norm_numerical_stability(default_vllm_config):

237 weight = torch.ones(hidden_size, dtype=dtype, device=device)

238

239 for idx, input_tensor in enumerate(test_cases):

241 - # Standard implementation

242 - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)

243 - rms_norm_layer.weight.data = weight.clone()

244 - standard_output = rms_norm_layer.forward_cuda(input_tensor)

240 + reference_output = _rms_norm_reference(input_tensor, weight, eps)

241

242 # Batch-invariant implementation

243 triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

244

245 # Check for NaN or Inf

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

251 - f"Standard RMS norm produced NaN for test case {idx}"

246 + assert not torch.isnan(reference_output).any(), (

247 + f"Reference RMS norm produced NaN for test case {idx}"

248 )

253 - assert not torch.isinf(standard_output).any(), (

254 - f"Standard RMS norm produced Inf for test case {idx}"

249 + assert not torch.isinf(reference_output).any(), (

250 + f"Reference RMS norm produced Inf for test case {idx}"

251 )

252 assert not torch.isnan(triton_output).any(), (

253 f"Triton RMS norm produced NaN for test case {idx}"

@@ -263,7 +259,7 @@ def test_rms_norm_numerical_stability(default_vllm_config):

259 # Compare outputs - very lenient for extreme values with float16

260 torch.testing.assert_close(

261 triton_output,

266 - standard_output,

262 + reference_output,

263 rtol=2e-1, # 20% tolerance for extreme values

264 atol=2e-1,

265 msg=f"RMS norm mismatch for extreme value test case {idx}",

@@ -321,10 +317,7 @@ def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):

317 input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)

318 weight = torch.randn(hidden_size, dtype=dtype, device=device)

319

324 - # Standard implementation

325 - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)

326 - rms_norm_layer.weight.data = weight.clone()

327 - standard_output = rms_norm_layer.forward_cuda(input_tensor)

320 + reference_output = _rms_norm_reference(input_tensor, weight, eps)

321

322 # Batch-invariant implementation

323 triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

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

327

328 torch.testing.assert_close(

329 triton_output,

337 - standard_output,

330 + reference_output,

331 rtol=rtol,

332 atol=atol,

333 msg=f"RMS norm mismatch for hidden_size={hidden_size}",

@@ -420,21 +413,18 @@ if __name__ == "__main__":

413 input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)

414 weight = torch.randn(hidden_size, dtype=dtype, device=device)

415

423 - # Standard implementation

424 - rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)

425 - rms_norm_layer.weight.data = weight.clone()

426 - standard_output = rms_norm_layer.forward_cuda(input_tensor)

416 + reference_output = _rms_norm_reference(input_tensor, weight, eps)

417

418 # Batch-invariant implementation

419 triton_output = rms_norm_batch_invariant(input_tensor, weight, eps=eps)

420

421 # Compare

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

433 - mean_diff = (triton_output - standard_output).abs().mean().item()

422 + max_diff = (triton_output - reference_output).abs().max().item()

423 + mean_diff = (triton_output - reference_output).abs().mean().item()

424

425 print(f"Max difference: {max_diff:.6e}")

426 print(f"Mean difference: {mean_diff:.6e}")

437 - print(f"Standard output sample: {standard_output[0, :5].tolist()}")

427 + print(f"Reference output sample: {reference_output[0, :5].tolist()}")

428 print(f"Triton output sample: {triton_output[0, :5].tolist()}")

429

430 if max_diff < 1e-3:

[Bug] Fix batch invariance rms norm comparison (#49603)