@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 183 lines 6.79 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5 import torch
6
7 import vllm.kernels # noqa: F401
8 from tests.ir.ir_test_utils import assert_close, clone_args
9 from vllm import ir
10 from vllm.platforms import current_platform
11 from vllm.triton_utils import HAS_TRITON
12
13 gelu_and_mul_sparse_native = ir.ops.gelu_and_mul_sparse.impls["native"].impl_fn
14
15
16 def test_gelu_and_mul_sparse_registration():
17 expected = {
18 "native": True,
19 "triton": HAS_TRITON and current_platform.is_cuda(),
20 }
21 actual = {
22 provider: impl.supported
23 for provider, impl in ir.ops.gelu_and_mul_sparse.impls.items()
24 }
25 assert actual == expected
26
27
28 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
29 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
30 @pytest.mark.parametrize("num_tokens", [1, 8, 64])
31 @pytest.mark.parametrize("intermediate_size", [2048, 4096, 8192, 13824, 16384, 32768])
32 def test_gelu_and_mul_sparse_triton(
33 dtype: torch.dtype, num_tokens: int, intermediate_size: int
34 ):
35 torch.manual_seed(0)
36 torch.set_default_device("cuda")
37 args = ir.ops.gelu_and_mul_sparse.generate_inputs(
38 num_tokens=num_tokens,
39 hidden_size=intermediate_size,
40 dtype=dtype,
41 )
42 impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
43 assert impl.supports_args(*args)
44
45 expected = gelu_and_mul_sparse_native(*clone_args(args))
46 actual = impl.impl_fn(*clone_args(args))
47 assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
48
49 mask_mismatch = torch.count_nonzero((actual == 0) != (expected == 0))
50 assert mask_mismatch / actual.numel() < 1e-3
51
52
53 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
54 def test_gelu_and_mul_sparse_triton_second_seed():
55 torch.manual_seed(1)
56 torch.set_default_device("cuda")
57 args = ir.ops.gelu_and_mul_sparse.generate_inputs(
58 num_tokens=32,
59 hidden_size=8192,
60 dtype=torch.bfloat16,
61 )
62
63 actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args))
64 expected = gelu_and_mul_sparse_native(*clone_args(args))
65 assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
66
67
68 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
69 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
70 @pytest.mark.parametrize("intermediate_size", [1, 31, 128])
71 def test_gelu_and_mul_sparse_triton_small_sizes(
72 dtype: torch.dtype, intermediate_size: int
73 ):
74 torch.manual_seed(0)
75 torch.set_default_device("cuda")
76 args = ir.ops.gelu_and_mul_sparse.generate_inputs(
77 num_tokens=5,
78 hidden_size=intermediate_size,
79 dtype=dtype,
80 )
81
82 actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*clone_args(args))
83 expected = gelu_and_mul_sparse_native(*clone_args(args))
84 assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
85 assert torch.equal(actual == 0, expected == 0)
86
87
88 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
89 def test_gelu_and_mul_sparse_triton_nd_and_empty():
90 torch.set_default_device("cuda")
91 impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
92
93 x = torch.randn(2, 3, 16384, dtype=torch.bfloat16)
94 args = (x, 1.6448536269514722, "tanh")
95 actual = impl.impl_fn(*args)
96 expected = gelu_and_mul_sparse_native(*args)
97 assert actual.shape == (2, 3, 8192)
98 assert_close(ir.ops.gelu_and_mul_sparse, actual, expected)
99
100 empty = torch.empty(0, 16384, dtype=torch.bfloat16)
101 empty_out = impl.impl_fn(empty, 1.6448536269514722, "tanh")
102 assert empty_out.shape == (0, 8192)
103 assert empty_out.dtype == empty.dtype
104
105
106 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
107 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
108 def test_gelu_and_mul_sparse_triton_constant_rows(dtype: torch.dtype):
109 torch.set_default_device("cuda")
110 gate = torch.full((4, 8192), 3.0, dtype=dtype)
111 up = torch.randn_like(gate)
112 x = torch.cat((gate, up), dim=-1)
113 args = (x, 1.6448536269514722, "tanh")
114
115 actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args)
116 expected = gelu_and_mul_sparse_native(*args)
117 torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
118 assert torch.count_nonzero(actual) == 0
119
120
121 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
122 @pytest.mark.parametrize("nonfinite", [float("nan"), float("inf")])
123 def test_gelu_and_mul_sparse_triton_nonfinite_gate(nonfinite: float):
124 torch.set_default_device("cuda")
125 x = torch.randn(2, 16384, dtype=torch.bfloat16)
126 x[0, 0] = nonfinite
127 args = (x, 1.6448536269514722, "tanh")
128
129 actual = ir.ops.gelu_and_mul_sparse.impls["triton"].impl_fn(*args)
130 expected = gelu_and_mul_sparse_native(*args)
131 torch.testing.assert_close(actual, expected, equal_nan=True)
132
133
134 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
135 def test_gelu_and_mul_sparse_triton_fallback_conditions():
136 torch.set_default_device("cuda")
137 impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
138 x = torch.randn(4, 16384, dtype=torch.bfloat16)
139
140 assert impl.supports_args(x, 1.0, "tanh")
141 assert not impl.supports_args(x, 1.0, "none")
142 assert not impl.supports_args(x[:, ::2], 1.0, "tanh")
143 assert not impl.supports_args(x[:, :-1], 1.0, "tanh")
144 assert not impl.supports_args(
145 torch.randn(1, 65538, dtype=torch.bfloat16), 1.0, "tanh"
146 )
147
148
149 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
150 def test_gelu_and_mul_sparse_triton_dispatch_and_opcheck():
151 torch.set_default_device("cuda")
152 args = ir.ops.gelu_and_mul_sparse.generate_inputs(
153 num_tokens=8,
154 hidden_size=4096,
155 dtype=torch.bfloat16,
156 )
157 impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
158
159 with ir.ops.gelu_and_mul_sparse.set_priority(["triton", "native"]):
160 dispatched = ir.ops.gelu_and_mul_sparse(*args)
161 torch.library.opcheck(torch.ops.vllm_ir.gelu_and_mul_sparse, args)
162
163 direct = impl.impl_fn(*args)
164 torch.testing.assert_close(dispatched, direct, rtol=0.0, atol=0.0)
165
166
167 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA-only provider")
168 def test_gelu_and_mul_sparse_triton_cudagraph():
169 torch.set_default_device("cuda")
170 x = torch.randn(16, 32768, dtype=torch.bfloat16)
171 impl = ir.ops.gelu_and_mul_sparse.impls["triton"]
172
173 for _ in range(3):
174 expected = impl.impl_fn(x, 1.6448536269514722, "tanh")
175 torch.accelerator.synchronize()
176
177 graph = torch.cuda.CUDAGraph()
178 with torch.cuda.graph(graph):
179 actual = impl.impl_fn(x, 1.6448536269514722, "tanh")
180 graph.replay()
181 torch.accelerator.synchronize()
182
183 torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)