@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 170 lines 5.54 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Platform-agnostic eligibility tests for GateLinear's ROCm router GEMMs.
4
5 These assert the dispatch flags directly, so they run device-free by mocking
6 the platform predicates. ``allow_cublas_router_gemm`` selects the
7 bf16xbf16->fp32 ``torch.mm`` epilogue, while ``allow_fp32_router_gemm`` selects
8 the gfx950 low-M kernel with fp32 weights and output.
9
10 The ROCm branch is guarded on ``not bias`` because ``torch.mm`` has no bias
11 term; a biased gate must fall back so the bias is not silently dropped.
12 """
13
14 import pytest
15 import torch
16
17 import vllm.model_executor.layers.fused_moe.router.gate_linear as gate_linear_mod
18 from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear
19
20
21 def _make_gate(
22 monkeypatch,
23 *,
24 is_rocm: bool,
25 is_cuda: bool = False,
26 bias: bool = False,
27 params_dtype: torch.dtype = torch.bfloat16,
28 out_dtype: torch.dtype | None = torch.float32,
29 input_size: int = 2048,
30 output_size: int = 64,
31 on_gfx950: bool = False,
32 parallel_world_size: int = 1,
33 ) -> GateLinear:
34 """Build a GateLinear with platform predicates mocked, no GPU needed."""
35 for target in (
36 "vllm.model_executor.layers.linear",
37 "vllm.model_executor.parameter",
38 ):
39 monkeypatch.setattr(
40 f"{target}.get_tensor_model_parallel_rank",
41 lambda: 0,
42 )
43 monkeypatch.setattr(
44 f"{target}.get_tensor_model_parallel_world_size",
45 lambda: parallel_world_size,
46 )
47
48 platform = gate_linear_mod.current_platform
49 monkeypatch.setattr(platform, "is_cuda", lambda: is_cuda)
50 monkeypatch.setattr(platform, "is_rocm", lambda: is_rocm)
51 monkeypatch.setattr(platform, "is_device_capability", lambda *a, **k: False)
52 monkeypatch.setattr(platform, "is_device_capability_family", lambda *a, **k: False)
53 if is_rocm:
54 import vllm.platforms.rocm as rocm_platform
55
56 monkeypatch.setattr(rocm_platform, "on_gfx950", lambda: on_gfx950)
57
58 return GateLinear(
59 input_size=input_size,
60 output_size=output_size,
61 bias=bias,
62 out_dtype=out_dtype,
63 params_dtype=params_dtype,
64 )
65
66
67 def test_rocm_no_bias_bf16_fp32_enables_fused_gemm(monkeypatch):
68 gate = _make_gate(monkeypatch, is_rocm=True, bias=False)
69 assert not gate.allow_specialized_router_gemm
70 assert gate.allow_cublas_router_gemm
71
72
73 def test_rocm_bias_disables_fused_gemm(monkeypatch):
74 # torch.mm cannot add a bias, so a biased gate must not take the fused path.
75 gate = _make_gate(monkeypatch, is_rocm=True, bias=True)
76 assert not gate.allow_cublas_router_gemm
77
78
79 def test_rocm_fp32_weight_disables_fused_gemm(monkeypatch):
80 gate = _make_gate(monkeypatch, is_rocm=True, params_dtype=torch.float32)
81 assert not gate.allow_cublas_router_gemm
82
83
84 def test_rocm_non_fp32_out_dtype_disables_fused_gemm(monkeypatch):
85 gate = _make_gate(monkeypatch, is_rocm=True, out_dtype=torch.bfloat16)
86 assert not gate.allow_cublas_router_gemm
87
88
89 def test_non_rocm_non_cuda_disables_fused_gemm(monkeypatch):
90 # Neither the CUDA specialized path nor the ROCm branch applies.
91 gate = _make_gate(monkeypatch, is_rocm=False, is_cuda=False)
92 assert not gate.allow_cublas_router_gemm
93
94
95 def test_rocm_set_out_dtype_enables_fused_gemm(monkeypatch):
96 gate = _make_gate(monkeypatch, is_rocm=True, bias=False, out_dtype=None)
97 assert not gate.allow_cublas_router_gemm
98 gate.set_out_dtype(torch.float32)
99 assert gate.allow_cublas_router_gemm
100
101
102 def test_rocm_set_out_dtype_respects_bias_guard(monkeypatch):
103 gate = _make_gate(monkeypatch, is_rocm=True, bias=True, out_dtype=None)
104 gate.set_out_dtype(torch.float32)
105 assert not gate.allow_cublas_router_gemm
106
107
108 @pytest.mark.parametrize(
109 ("input_size", "output_size"),
110 [(3072, 256), (4096, 8), (4096, 192), (6144, 128), (6144, 256)],
111 )
112 def test_rocm_gfx950_enables_fp32_router_gemm(
113 monkeypatch, input_size: int, output_size: int
114 ) -> None:
115 gate = _make_gate(
116 monkeypatch,
117 is_rocm=True,
118 params_dtype=torch.float32,
119 input_size=input_size,
120 output_size=output_size,
121 on_gfx950=True,
122 )
123 assert gate.allow_fp32_router_gemm
124 assert not gate.allow_cublas_router_gemm
125
126
127 @pytest.mark.parametrize("ep_size", [2, 4, 8])
128 def test_rocm_fp32_router_gemm_is_replicated_across_ep_sizes(
129 monkeypatch, ep_size: int
130 ) -> None:
131 gate = _make_gate(
132 monkeypatch,
133 is_rocm=True,
134 params_dtype=torch.float32,
135 input_size=6144,
136 output_size=128,
137 on_gfx950=True,
138 parallel_world_size=ep_size,
139 )
140 assert gate.weight.shape == (128, 6144)
141 assert gate.allow_fp32_router_gemm
142
143
144 @pytest.mark.parametrize(
145 ("input_size", "output_size", "params_dtype", "bias", "on_gfx950"),
146 [
147 pytest.param(6144, 128, torch.float32, False, False, id="gfx942"),
148 pytest.param(2048, 64, torch.float32, False, True, id="shape"),
149 pytest.param(6144, 128, torch.bfloat16, False, True, id="bf16-weight"),
150 pytest.param(6144, 128, torch.float32, True, True, id="bias"),
151 ],
152 )
153 def test_rocm_fp32_router_gemm_rejects_unsupported_configs(
154 monkeypatch,
155 input_size: int,
156 output_size: int,
157 params_dtype: torch.dtype,
158 bias: bool,
159 on_gfx950: bool,
160 ) -> None:
161 gate = _make_gate(
162 monkeypatch,
163 is_rocm=True,
164 params_dtype=params_dtype,
165 bias=bias,
166 input_size=input_size,
167 output_size=output_size,
168 on_gfx950=on_gfx950,
169 )
170 assert not gate.allow_fp32_router_gemm