@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 201 lines 6.9 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the FlashInfer TRTLLM BF16 MoE backend
4 (`TrtLlmBf16ExpertsModular`).
5
6 This mirrors the TRTLLM NvFP4 modular test shape: construct the modular
7 expert wrapper directly, pass production-format BlockMajorK weights, and
8 compare against a torch MoE reference using the original BF16 weights.
9 """
10
11 import pytest
12 import torch
13
14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15 from tests.kernels.moe.utils import make_dummy_moe_config
16 from tests.kernels.utils import torch_moe
17 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
18 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19 from vllm.model_executor.layers.fused_moe.all2all_utils import (
20 maybe_make_prepare_finalize,
21 )
22 from vllm.model_executor.layers.fused_moe.config import (
23 FUSED_MOE_UNQUANTIZED_CONFIG,
24 FusedMoEConfig,
25 FusedMoEParallelConfig,
26 RoutingMethodType,
27 )
28 from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (
29 TrtLlmBf16ExpertsModular,
30 )
31 from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (
32 TrtLlmBf16LoRAExperts,
33 )
34 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
35 UnquantizedMoeBackend,
36 convert_to_unquantized_kernel_format,
37 )
38 from vllm.platforms import current_platform
39 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
40 from vllm.utils.math_utils import next_power_of_2
41 from vllm.utils.torch_utils import set_random_seed
42
43 if pytest and (
44 not has_flashinfer_trtllm_fused_moe()
45 or not current_platform.is_device_capability_family(100)
46 ):
47 pytest.skip(
48 "Requires flashinfer TRTLLM fused MoE BF16 backend (SM100)",
49 allow_module_level=True,
50 )
51
52 # (m, n, k) = (tokens, intermediate_size_per_partition, hidden_dim).
53 MNK_FACTORS = [
54 (2, 160, 2560),
55 (2, 1024, 1024),
56 (64, 2048, 1536),
57 (64, 1024, 4096),
58 ]
59
60
61 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)
62 @pytest.mark.parametrize("e", [128])
63 @pytest.mark.parametrize("topk", [8])
64 @pytest.mark.parametrize("dtype", [torch.bfloat16])
65 @torch.inference_mode()
66 def test_trtllm_bf16_moe_modular_no_graph(
67 m: int,
68 n: int,
69 k: int,
70 e: int,
71 topk: int,
72 dtype: torch.dtype,
73 workspace_init,
74 ):
75 set_random_seed(7)
76 with set_current_vllm_config(
77 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
78 ):
79 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
80 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
81 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
82 # The FlashInfer conversion may rewrite unpadded input storage in place.
83 # Preserve the original layout for the independent torch reference.
84 reference_w1 = w1.clone()
85 reference_w2 = w2.clone()
86 score = torch.randn((m, e), device="cuda", dtype=dtype)
87 scores = torch.softmax(score, dim=-1, dtype=torch.float32)
88 topk_weights, topk_ids = torch.topk(scores, topk)
89 topk_weights = topk_weights.contiguous()
90 topk_ids = topk_ids.to(torch.int32).contiguous()
91
92 moe_config = FusedMoEConfig(
93 num_experts=e,
94 experts_per_token=topk,
95 hidden_dim=k,
96 intermediate_size=n,
97 num_local_experts=e,
98 num_logical_experts=e,
99 activation=MoEActivation.SILU,
100 device="cuda",
101 moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
102 in_dtype=dtype,
103 routing_method=RoutingMethodType.TopK,
104 max_num_tokens=next_power_of_2(m),
105 )
106
107 trtllm_w1, trtllm_w2 = convert_to_unquantized_kernel_format(
108 UnquantizedMoeBackend.FLASHINFER_TRTLLM,
109 moe_config,
110 w1,
111 w2,
112 )
113 expected_n = (n + 127) // 128 * 128
114 assert moe_config.intermediate_size_per_partition == expected_n
115 assert trtllm_w2.numel() == e * k * expected_n
116
117 trtllm_experts = mk.FusedMoEKernel(
118 maybe_make_prepare_finalize(
119 moe=moe_config,
120 quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
121 allow_new_interface=True,
122 use_monolithic=False,
123 ),
124 TrtLlmBf16ExpertsModular(
125 moe_config=moe_config,
126 quant_config=FUSED_MOE_UNQUANTIZED_CONFIG,
127 ),
128 )
129
130 trtllm_output = trtllm_experts.apply(
131 hidden_states=a,
132 w1=trtllm_w1,
133 w2=trtllm_w2,
134 topk_weights=topk_weights,
135 topk_ids=topk_ids,
136 activation=MoEActivation.SILU,
137 global_num_experts=e,
138 expert_map=None,
139 apply_router_weight_on_input=False,
140 )
141
142 torch_output = torch_moe(
143 a,
144 reference_w1,
145 reference_w2,
146 score,
147 topk,
148 activation=MoEActivation.SILU,
149 )
150
151 torch.testing.assert_close(
152 torch_output,
153 trtllm_output,
154 atol=1e-1,
155 rtol=2e-1,
156 )
157
158
159 @pytest.mark.parametrize("has_lora_delta", [False, True])
160 @torch.inference_mode()
161 def test_trtllm_bf16_lora_accepts_checkpoint_shaped_weights(has_lora_delta):
162 """LoRA dispatch accepts both 3D parameters and legacy 4D packed views."""
163 config = make_dummy_moe_config(
164 num_experts=128, hidden_dim=256, intermediate_size=128
165 )
166 experts = TrtLlmBf16LoRAExperts(config, FUSED_MOE_UNQUANTIZED_CONFIG)
167 x = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) / 10
168 w1, w2 = convert_to_unquantized_kernel_format(
169 UnquantizedMoeBackend.FLASHINFER_TRTLLM,
170 config,
171 torch.randn(128, 256, 256, device="cuda", dtype=torch.bfloat16) / 10,
172 torch.randn(128, 256, 128, device="cuda", dtype=torch.bfloat16) / 10,
173 )
174 assert w1.ndim == w2.ndim == 3
175 topk = (
176 torch.tensor([[0], [1]], device="cuda", dtype=torch.int32),
177 torch.ones(2, 1, device="cuda", dtype=torch.float32),
178 )
179 delta = torch.randn(2, 1, 256, device="cuda", dtype=torch.bfloat16)
180
181 def invoke(w1, w2):
182 result = experts.invoke_routed_moe(
183 hidden_states=x,
184 w1=w1,
185 w2=w2,
186 topk_ids_and_weights=topk,
187 gemm1_lora_delta=delta if has_lora_delta else None,
188 global_num_experts=128,
189 a1q_scale=None,
190 output=torch.empty_like(x),
191 )
192 if has_lora_delta:
193 # Only compare rows belonging to real tokens, excluding padding.
194 indices = result[2].flatten().long()
195 return result[0][indices], result[3][indices]
196 return result
197
198 actual = invoke(w1, w2)
199 expected = invoke(w1.view(128, 4, 256, 64), w2.view(128, 2, 256, 64))
200 for actual_tensor, expected_tensor in zip(actual, expected):
201 torch.testing.assert_close(actual_tensor, expected_tensor)