@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 131 lines 4.24 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Test batch-invariant matmul against torch.matmul for various shape combinations.
4
5 Tests correctness (matches torch.matmul) and batch invariance (result for one
6 item doesn't change based on other items in the batch).
7 """
8
9 import pytest
10 import torch
11 from utils import skip_unsupported
12
13 from vllm.model_executor.determinism.batch_invariant import matmul_batch_invariant
14 from vllm.model_executor.determinism.batch_invariant_configs import (
15 _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS,
16 _get_tuned_matmul_arch_family,
17 )
18 from vllm.platforms import current_platform
19
20 DEVICE_TYPE = current_platform.device_type
21
22
23 @skip_unsupported
24 @pytest.mark.parametrize(
25 "a_shape,b_shape",
26 [
27 # 2D x 2D
28 ((32, 64), (64, 16)),
29 # 2D x 3D
30 ((64, 16), (4, 16, 32)),
31 # 3D x 2D
32 ((4, 32, 64), (64, 16)),
33 # 4D x 2D
34 ((1, 4, 32, 64), (64, 16)),
35 # 3D x 3D
36 ((4, 32, 64), (4, 64, 16)),
37 # 3D x 4D
38 ((2, 32, 64), (1, 2, 64, 16)),
39 # 4D x 3D (Gemma4 pattern)
40 ((1, 2, 32, 64), (2, 64, 16)),
41 # 4D x 4D
42 ((1, 2, 32, 64), (4, 2, 64, 16)),
43 # 2D x 4D
44 ((32, 64), (1, 2, 64, 16)),
45 # 2D x 5D
46 ((32, 64), (1, 2, 2, 64, 16)),
47 # 5D x 2D
48 ((1, 2, 2, 32, 64), (64, 16)),
49 # 5D x 5D
50 ((1, 2, 4, 32, 64), (1, 2, 4, 64, 16)),
51 ],
52 )
53 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
54 def test_matmul_correctness(a_shape, b_shape, dtype):
55 """Compare matmul_batch_invariant against torch.matmul for various shapes."""
56 device = torch.device(DEVICE_TYPE)
57
58 torch.manual_seed(42)
59 a = torch.rand(a_shape, dtype=dtype, device=device)
60 b = torch.rand(b_shape, dtype=dtype, device=device)
61
62 # Standard implementation (CUDA ops)
63 standard_output = torch.matmul(a, b)
64
65 # Batch-invariant implementation (Triton)
66 triton_output = matmul_batch_invariant(a, b)
67
68 # Compare outputs
69 # Use looser tolerance for bfloat16 due to its lower precision
70 if dtype == torch.bfloat16:
71 rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16
72 else:
73 rtol, atol = 1e-2, 1e-2 # 1% for float16/float32
74
75 torch.testing.assert_close(
76 triton_output,
77 standard_output,
78 rtol=rtol,
79 atol=atol,
80 msg=f"matmul mismatch for a ndim={a.ndim}, b ndim={b.ndim},",
81 )
82
83
84 @skip_unsupported
85 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
86 def test_matmul_batch_invariance(dtype):
87 """Verify that the result for one item is bitwise identical regardless
88 of what other items are in the batch.
89 """
90 device = torch.device(DEVICE_TYPE)
91
92 torch.manual_seed(42)
93 a_single = torch.rand((1, 64, 32), dtype=dtype, device=device)
94 b = torch.rand((32, 128), dtype=dtype, device=device)
95
96 standard_output = matmul_batch_invariant(a_single, b)
97
98 a_batch = torch.rand((8, 64, 32), dtype=dtype, device=device)
99 a_batch[3] = a_single[0]
100
101 batch_output = matmul_batch_invariant(a_batch, b)
102 batch_output_a = batch_output[3]
103
104 assert torch.equal(standard_output[0], batch_output_a)
105
106
107 @skip_unsupported
108 @pytest.mark.parametrize("m", [8, 32, 256, 2048])
109 @pytest.mark.parametrize("transpose_b", [False, True], ids=["contiguous", "transposed"])
110 def test_matmul_batch_invariance_across_tuned_m_buckets(m, transpose_b):
111 # Tuned M buckets must preserve each row's K-reduction order.
112 capability = (
113 current_platform.get_device_capability() if current_platform.is_cuda() else None
114 )
115 arch_family = _get_tuned_matmul_arch_family(capability)
116 if arch_family not in _BATCH_INVARIANT_MATMUL_TUNED_CONFIGS:
117 pytest.skip("No tuned persistent matmul config for this architecture")
118
119 device = torch.device(DEVICE_TYPE)
120 n = k = 2048
121 torch.manual_seed(42)
122 a = torch.rand((m, k), dtype=torch.bfloat16, device=device)
123 if transpose_b:
124 b = torch.rand((n, k), dtype=torch.bfloat16, device=device).t()
125 else:
126 b = torch.rand((k, n), dtype=torch.bfloat16, device=device)
127
128 single_output = matmul_batch_invariant(a[:1], b)
129 batch_output = matmul_batch_invariant(a, b)
130
131 assert torch.equal(single_output[0], batch_output[0])