@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 228 lines 7.91 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Correctness tests for the gfx950 low-M fp32 router GEMM."""
4
5 import pytest
6 import torch
7
8 from vllm.model_executor.layers.fused_moe.router.gate_linear import (
9 fp32_router_gemm_dispatch_impl,
10 )
11 from vllm.model_executor.layers.fused_moe.router.rocm_fp32_router_gemm import (
12 rocm_fp32_router_gemm,
13 )
14 from vllm.platforms import current_platform
15
16 SHAPES = [
17 (3072, 256),
18 (4096, 8),
19 (4096, 192),
20 (6144, 128),
21 (6144, 256),
22 ]
23 MAX_TOKENS = 32
24 ATOL = 5e-4
25 RTOL = 0.0
26
27
28 def _on_gfx950() -> bool:
29 if not current_platform.is_rocm():
30 return False
31 try:
32 from vllm.platforms.rocm import on_gfx950
33
34 return on_gfx950()
35 except Exception:
36 return False
37
38
39 pytestmark = pytest.mark.skipif(
40 not _on_gfx950(), reason="rocm_fp32_router_gemm requires ROCm gfx950"
41 )
42
43
44 def _reference(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
45 return torch.nn.functional.linear(x.float(), weight)
46
47
48 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
49 @pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES)
50 @pytest.mark.parametrize("num_tokens", range(MAX_TOKENS + 1))
51 @torch.inference_mode()
52 def test_rocm_fp32_router_gemm_matches_reference(
53 num_tokens: int,
54 hidden_size: int,
55 num_experts: int,
56 dtype: torch.dtype,
57 ) -> None:
58 torch.manual_seed(41 + num_tokens + hidden_size + num_experts)
59 device = torch.device("cuda")
60 x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
61 weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
62
63 output = rocm_fp32_router_gemm(x, weight)
64 expected = _reference(x, weight)
65
66 assert output.shape == (num_tokens, num_experts)
67 assert output.dtype == torch.float32
68 assert output.device == x.device
69 torch.testing.assert_close(output, expected, atol=ATOL, rtol=RTOL)
70
71
72 @pytest.mark.parametrize("num_tokens", [1, 4, 16, 24, 32])
73 @pytest.mark.parametrize(("hidden_size", "num_experts"), SHAPES)
74 @torch.inference_mode()
75 def test_rocm_fp32_router_gemm_preserves_topk(
76 num_tokens: int, hidden_size: int, num_experts: int
77 ) -> None:
78 torch.manual_seed(1000 + num_tokens + hidden_size + num_experts)
79 device = torch.device("cuda")
80 x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=device)
81 weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
82
83 top_k = 2 if (hidden_size, num_experts) == (4096, 8) else 8
84 actual = rocm_fp32_router_gemm(x, weight).topk(top_k, dim=-1).indices
85 reference = _reference(x, weight)
86 reference_values, expected = reference.topk(top_k, dim=-1)
87 for token in range(num_tokens):
88 actual_set = set(actual[token].tolist())
89 expected_set = set(expected[token].tolist())
90 if actual_set == expected_set:
91 continue
92 kth_value = reference_values[token, -1].item()
93 for expert in actual_set.symmetric_difference(expected_set):
94 gap = abs(reference[token, expert].item() - kth_value)
95 assert gap < 1e-3
96
97
98 @torch.inference_mode()
99 def test_rocm_fp32_router_gemm_rejects_invalid_inputs() -> None:
100 device = torch.device("cuda")
101 hidden_size, num_experts = 6144, 128
102 x = torch.randn(4, hidden_size, dtype=torch.bfloat16, device=device)
103 weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
104
105 with pytest.raises(ValueError, match="num_tokens"):
106 rocm_fp32_router_gemm(
107 torch.randn(
108 MAX_TOKENS + 1,
109 hidden_size,
110 dtype=torch.bfloat16,
111 device=device,
112 ),
113 weight,
114 )
115
116 with pytest.raises(ValueError, match="shape"):
117 rocm_fp32_router_gemm(x[:, :2048].contiguous(), weight[:, :2048].contiguous())
118
119 with pytest.raises(ValueError, match="contiguous"):
120 wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device)
121 rocm_fp32_router_gemm(wide_x[:, ::2], weight)
122
123 with pytest.raises(ValueError, match="float32"):
124 rocm_fp32_router_gemm(x, weight.to(torch.bfloat16))
125
126 with pytest.raises(ValueError, match="bfloat16 or float32"):
127 rocm_fp32_router_gemm(x.to(torch.float16), weight)
128
129
130 @torch.inference_mode()
131 def test_rocm_fp32_router_gemm_cuda_graph_observes_input_mutation() -> None:
132 torch.manual_seed(2026)
133 device = torch.device("cuda")
134 hidden_size, num_experts = 6144, 128
135 static_x = torch.randn(16, hidden_size, dtype=torch.bfloat16, device=device)
136 weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
137
138 rocm_fp32_router_gemm(static_x, weight)
139 torch.accelerator.synchronize()
140
141 graph = torch.cuda.CUDAGraph()
142 with torch.cuda.graph(graph):
143 captured_output = rocm_fp32_router_gemm(static_x, weight)
144 torch.accelerator.synchronize()
145
146 first_input = torch.randn_like(static_x)
147 static_x.copy_(first_input)
148 graph.replay()
149 torch.accelerator.synchronize()
150 first_output = captured_output.clone()
151
152 second_input = torch.randn_like(static_x)
153 static_x.copy_(second_input)
154 graph.replay()
155 torch.accelerator.synchronize()
156 second_output = captured_output.clone()
157
158 torch.testing.assert_close(
159 first_output, _reference(first_input, weight), atol=ATOL, rtol=RTOL
160 )
161 torch.testing.assert_close(
162 second_output, _reference(second_input, weight), atol=ATOL, rtol=RTOL
163 )
164 assert not torch.equal(first_output, second_output)
165
166
167 @pytest.mark.parametrize("num_tokens", [4, 32, 33])
168 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])
169 @torch.inference_mode()
170 def test_rocm_fp32_router_gemm_custom_op_dispatch(
171 num_tokens: int, dtype: torch.dtype
172 ) -> None:
173 torch.manual_seed(3100 + num_tokens)
174 device = torch.device("cuda")
175 hidden_size, num_experts = 6144, 128
176 x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
177 weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
178
179 output = torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
180 direct_output = fp32_router_gemm_dispatch_impl(x, weight, False)
181
182 torch.testing.assert_close(output, direct_output, atol=ATOL, rtol=RTOL)
183 torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL)
184
185
186 @torch.inference_mode()
187 def test_rocm_fp32_router_gemm_dispatch_accepts_noncontiguous_input(
188 monkeypatch: pytest.MonkeyPatch,
189 ) -> None:
190 torch.manual_seed(4100)
191 device = torch.device("cuda")
192 hidden_size, num_experts = 6144, 128
193 wide_x = torch.randn(4, hidden_size * 2, dtype=torch.bfloat16, device=device)
194 x = wide_x[:, ::2]
195 weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
196 assert not x.is_contiguous()
197
198 expected = _reference(x, weight)
199
200 def fail_fallback(*args, **kwargs):
201 raise AssertionError("noncontiguous low-M input used the linear fallback")
202
203 monkeypatch.setattr(torch.nn.functional, "linear", fail_fallback)
204 output = torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
205
206 torch.testing.assert_close(output, expected, atol=ATOL, rtol=RTOL)
207
208
209 @torch.inference_mode()
210 def test_rocm_fp32_router_gemm_dynamic_compile_dispatch() -> None:
211 torch.manual_seed(5100)
212 device = torch.device("cuda")
213 hidden_size, num_experts = 6144, 128
214 weight = torch.randn(num_experts, hidden_size, dtype=torch.float32, device=device)
215
216 def dispatch(x: torch.Tensor) -> torch.Tensor:
217 return torch.ops.vllm.fp32_router_gemm_dispatch(x, weight, False)
218
219 compiled_dispatch = torch.compile(dispatch, dynamic=True, fullgraph=True)
220 for num_tokens in (4, 16, 33, 5, 32):
221 x = torch.randn(
222 num_tokens,
223 hidden_size,
224 dtype=torch.bfloat16,
225 device=device,
226 )
227 output = compiled_dispatch(x)
228 torch.testing.assert_close(output, _reference(x, weight), atol=ATOL, rtol=RTOL)