@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 465 lines 15.1 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Unit-test DeepGEMM FP8 and FP4 kernels (no DeepEP).
4 Compare DeepGEMM path against the Triton fallback inside vLLM's fused_experts.
5 """
6
7 import importlib
8 import math
9
10 import pytest
11 import torch
12
13 # vLLM fused-expert reference (Triton fallback + DeepGEMM option)
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 vllm.model_executor.layers.fused_moe.activation import (
17 MoEActivation,
18 )
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 FusedMoEQuantConfig,
24 FusedMoEQuantDesc,
25 fp8_w8a8_moe_quant_config,
26 )
27 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
28 deepgemm_moe_permute,
29 ep_gather,
30 )
31 from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
32 _fp8_workspace_shape,
33 )
34 from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
35 TritonOrDeepGemmExperts,
36 )
37 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
38 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
39 per_token_group_quant_fp8,
40 )
41 from vllm.utils.deep_gemm import (
42 calc_diff,
43 is_deep_gemm_supported,
44 per_block_cast_to_fp8,
45 )
46
47 BLOCK_SIZE = [128, 128]
48
49
50 def test_ep_gather_uses_64_bit_row_offsets():
51 hidden_size = 4096
52 source_row = 1 << 18
53 input_tensor = torch.empty(
54 (source_row + 1, hidden_size),
55 device="cuda",
56 dtype=torch.bfloat16,
57 )
58 input_tensor[source_row].fill_(1)
59 output = torch.empty((1, hidden_size), device="cuda", dtype=torch.bfloat16)
60
61 ep_gather(
62 input_tensor=input_tensor,
63 recv_topk_ids=torch.zeros((1, 1), device="cuda", dtype=torch.int64),
64 recv_topk_weight=torch.ones((1, 1), device="cuda", dtype=torch.float32),
65 input_index=torch.full((1, 1), source_row, device="cuda", dtype=torch.int32),
66 expert_map=None,
67 output_tensor=output,
68 )
69
70 torch.testing.assert_close(output, torch.ones_like(output), rtol=0, atol=0)
71
72
73 @pytest.mark.parametrize("workspace_dtype", [torch.float16, torch.bfloat16])
74 @pytest.mark.parametrize("num_columns", [2048, 6144, 6145])
75 def test_fp8_workspace_shape(workspace_dtype, num_columns):
76 num_rows = 17
77 shape = _fp8_workspace_shape(num_rows, num_columns, workspace_dtype)
78
79 allocated_bytes = math.prod(shape) * workspace_dtype.itemsize
80 required_bytes = num_rows * num_columns * torch.float8_e4m3fn.itemsize
81
82 assert allocated_bytes >= required_bytes
83 assert allocated_bytes - required_bytes < num_rows * workspace_dtype.itemsize
84
85
86 @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
87 def test_deepgemm_moe_permute_initializes_padding_scales(workspace_init):
88 hidden_states = torch.randn(2, 128, device="cuda", dtype=torch.bfloat16)
89 activations, scales = per_token_group_quant_fp8(
90 hidden_states,
91 group_size=128,
92 use_ue8m0=True,
93 )
94 topk_ids = torch.tensor([[0], [1]], device="cuda", dtype=torch.int64)
95
96 _, permuted_scales, expert_ids, _, _ = deepgemm_moe_permute(
97 aq=activations,
98 aq_scale=scales,
99 topk_ids=topk_ids,
100 local_num_experts=2,
101 expert_map=None,
102 expert_tokens_meta=None,
103 )
104
105 padding = expert_ids < 0
106 assert padding.any()
107 torch.testing.assert_close(
108 permuted_scales[padding],
109 torch.zeros_like(permuted_scales[padding]),
110 rtol=0,
111 atol=0,
112 )
113
114
115 def make_block_quant_fp8_weights(
116 e: int,
117 n: int,
118 k: int,
119 block_size: list[int],
120 ):
121 """Generate (w1, w2) expert weights and their per-block scale tensors
122 in FP8 block-quantized format.
123
124 w1 shape: (E, 2N, K)
125 w2 shape: (E, K, N)
126 """
127 dtype = torch.bfloat16
128 fp8_max, fp8_min = (
129 torch.finfo(torch.float8_e4m3fn).max,
130 torch.finfo(torch.float8_e4m3fn).min,
131 )
132
133 # bf16 reference weights
134 w1_bf16 = torch.randn(e, 2 * n, k, device="cuda", dtype=dtype) / 10
135 w2_bf16 = torch.randn(e, k, n, device="cuda", dtype=dtype) / 10
136 w1_bf16.clamp_(fp8_min, fp8_max)
137 w2_bf16.clamp_(fp8_min, fp8_max)
138
139 block_n, block_k = block_size
140 n_tiles_w1 = math.ceil((2 * n) / block_n)
141 k_tiles_w1 = math.ceil(k / block_k)
142 n_tiles_w2 = math.ceil(k / block_n)
143 k_tiles_w2 = math.ceil(n / block_k)
144
145 w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn)
146 w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn)
147 w1_s = torch.empty(e, n_tiles_w1, k_tiles_w1, device="cuda", dtype=torch.float32)
148 w2_s = torch.empty(e, n_tiles_w2, k_tiles_w2, device="cuda", dtype=torch.float32)
149
150 for i in range(e):
151 w1[i], w1_s[i] = per_block_cast_to_fp8(
152 w1_bf16[i], block_size=block_size, use_ue8m0=True
153 )
154 w2[i], w2_s[i] = per_block_cast_to_fp8(
155 w2_bf16[i], block_size=block_size, use_ue8m0=True
156 )
157
158 return w1, w2, w1_s, w2_s
159
160
161 def run_single_case(m, n, k, topk, num_experts, block_size):
162 """Run one (M,N,K) configuration on a single GPU and assert DeepGEMM ==
163 Triton baseline within tolerance.
164 """
165 tokens_bf16 = (
166 torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
167 .clamp_min_(-1)
168 .clamp_max_(1)
169 )
170 _, a1_scale = per_token_group_quant_fp8(tokens_bf16, block_size[1])
171
172 # expert weight tensors
173 w1, w2, w1_s, w2_s = make_block_quant_fp8_weights(num_experts, n, k, block_size)
174
175 router_logits = torch.randn(m, num_experts, device="cuda", dtype=torch.float32)
176 topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
177 topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
178
179 quant_config = fp8_w8a8_moe_quant_config(
180 w1_scale=w1_s,
181 w2_scale=w2_s,
182 a1_scale=a1_scale,
183 block_shape=block_size,
184 )
185 moe_config = make_dummy_moe_config()
186
187 deep_gemm_experts = mk.FusedMoEKernel(
188 prepare_finalize=maybe_make_prepare_finalize(
189 moe=moe_config,
190 quant_config=quant_config,
191 allow_new_interface=True,
192 use_monolithic=False,
193 ),
194 fused_experts=TritonOrDeepGemmExperts(
195 moe_config=moe_config,
196 quant_config=quant_config,
197 ),
198 )
199
200 # triton reference
201 out_triton = fused_experts(
202 hidden_states=tokens_bf16,
203 w1=w1,
204 w2=w2,
205 topk_weights=topk_weights,
206 topk_ids=topk_ids,
207 quant_config=quant_config,
208 )
209
210 # DeepGemm
211 out_deepgemm = deep_gemm_experts.apply(
212 hidden_states=tokens_bf16,
213 w1=w1,
214 w2=w2,
215 topk_weights=topk_weights,
216 topk_ids=topk_ids,
217 global_num_experts=num_experts,
218 activation=MoEActivation.SILU,
219 apply_router_weight_on_input=False,
220 expert_map=None,
221 )
222 diff = calc_diff(out_deepgemm, out_triton)
223 assert diff < 0.001, f"Diff exceeded 1%: {diff}"
224
225
226 # Note: N <= 512 will disable the deepgemm path due to performance issues.
227 MNKs = [
228 (1024, 768, 128),
229 (2048, 768, 512),
230 (512, 1024, 1024),
231 (4096, 4096, 1024),
232 ]
233
234 TOPKS = [2, 6]
235 NUM_EXPERTS = [32]
236
237
238 @pytest.mark.parametrize(("m", "n", "k"), MNKs)
239 @pytest.mark.parametrize("topk", TOPKS)
240 @pytest.mark.parametrize("num_experts", NUM_EXPERTS)
241 @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
242 def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch, workspace_init):
243 with monkeypatch.context() as mp:
244 mp.setenv("VLLM_USE_DEEP_GEMM", "1")
245
246 _DeepGemmExperts = importlib.import_module(
247 "vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe"
248 ).DeepGemmExperts
249
250 call_counter = {"cnt": 0}
251
252 orig_fn = _DeepGemmExperts.apply
253
254 def _spy_apply(*args, **kwargs):
255 call_counter["cnt"] += 1
256 return orig_fn(*args, **kwargs)
257
258 monkeypatch.setattr(_DeepGemmExperts, "apply", _spy_apply)
259 if topk > num_experts:
260 pytest.skip(f"topk={topk} > num_experts={num_experts}")
261
262 run_single_case(
263 m=m,
264 n=n,
265 k=k,
266 topk=topk,
267 num_experts=num_experts,
268 block_size=BLOCK_SIZE,
269 )
270
271 # ensure that the DeepGEMM path was indeed taken.
272 assert call_counter["cnt"] == 1, (
273 f"DeepGEMM path was not executed during the test. "
274 f"Call counter: {call_counter['cnt']}"
275 )
276
277
278 # ---------------------------------------------------------------------------
279 # FP4 weight tests (DeepGEMM m_grouped_fp8_fp4_gemm_nt_contiguous)
280 # ---------------------------------------------------------------------------
281
282
283 def make_mxfp4_weights(
284 e: int,
285 n: int,
286 k: int,
287 ):
288 """Generate (w1, w2) expert weights in MXFP4 packed format with float32 scales,
289 plus BF16 reference weights for validation.
290
291 w1 shape: (E, 2N, K//2) uint8 — packed FP4
292 w2 shape: (E, K, N//2) uint8 — packed FP4
293 w1_s shape: (E, 2N, K//32) float32 — per-row block-32 scales
294 w2_s shape: (E, K, N//32) float32 — per-row block-32 scales
295 w1_bf16: (E, 2N, K) — original BF16 for reference
296 w2_bf16: (E, K, N) — original BF16 for reference
297 """
298 from deep_gemm.utils.math import per_token_cast_to_fp4
299
300 dtype = torch.bfloat16
301 gran_k = 32 # MXFP4 block size
302
303 # bf16 reference weights — scale by 1/sqrt(dim) for numerical stability
304 w1_bf16 = torch.randn(e, 2 * n, k, device="cuda", dtype=dtype) * (k**-0.5)
305 w2_bf16 = torch.randn(e, k, n, device="cuda", dtype=dtype) * (n**-0.5)
306
307 # Quantize per-expert to FP4
308 w1 = torch.empty(e, 2 * n, k // 2, device="cuda", dtype=torch.uint8)
309 w2 = torch.empty(e, k, n // 2, device="cuda", dtype=torch.uint8)
310 w1_s = torch.empty(
311 e, 2 * n, math.ceil(k / gran_k), device="cuda", dtype=torch.float32
312 )
313 w2_s = torch.empty(e, k, math.ceil(n / gran_k), device="cuda", dtype=torch.float32)
314
315 for i in range(e):
316 w1[i], w1_s[i] = per_token_cast_to_fp4(
317 w1_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
318 )
319 w2[i], w2_s[i] = per_token_cast_to_fp4(
320 w2_bf16[i].float(), use_ue8m0=True, gran_k=gran_k
321 )
322
323 return w1, w2, w1_s, w2_s, w1_bf16, w2_bf16
324
325
326 def _bf16_moe_reference(x, w1, w2, topk_weights, topk_ids):
327 """BF16 token-loop MoE reference for correctness testing."""
328 import torch.nn.functional as F
329
330 num_tokens, hidden_size = x.shape
331 intermediate = w1.shape[1] // 2
332 top_k = topk_ids.shape[1]
333
334 output = torch.zeros(num_tokens, hidden_size, dtype=torch.float32, device=x.device)
335 for t in range(num_tokens):
336 for kk in range(top_k):
337 e = topk_ids[t, kk].item()
338 w = topk_weights[t, kk].item()
339 fc1 = x[t : t + 1].float() @ w1[e].float().T
340 linear = fc1[:, :intermediate]
341 gate = fc1[:, intermediate:]
342 act = F.silu(gate) * linear
343 fc2 = act @ w2[e].float().T
344 output[t] += w * fc2[0]
345 return output.to(torch.bfloat16)
346
347
348 def run_single_fp4_case(m, n, k, topk, num_experts):
349 """Run one (M,N,K) configuration with FP4 weights on DeepGEMM and assert
350 DeepGEMM FP4 == BF16 reference within tolerance.
351 """
352 tokens_bf16 = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
353
354 # FP4 expert weight tensors + BF16 originals for reference
355 w1, w2, w1_s, w2_s, w1_bf16, w2_bf16 = make_mxfp4_weights(num_experts, n, k)
356
357 router_logits = torch.randn(m, num_experts, device="cuda", dtype=torch.float32)
358 topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
359 topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
360
361 from vllm.model_executor.layers.quantization.utils.quant_utils import (
362 GroupShape,
363 )
364 from vllm.platforms import current_platform
365
366 _fp8_dtype = current_platform.fp8_dtype()
367 _block_shape = GroupShape(128, 128)
368 quant_config = FusedMoEQuantConfig(
369 _a1=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
370 _a2=FusedMoEQuantDesc(_fp8_dtype, _block_shape, None, None, None, None),
371 _w1=FusedMoEQuantDesc("mxfp4", None, w1_s, None, None, None),
372 _w2=FusedMoEQuantDesc("mxfp4", None, w2_s, None, None, None),
373 )
374 moe_config = make_dummy_moe_config()
375
376 from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import (
377 DeepGemmFP4Experts,
378 )
379
380 deep_gemm_fp4_experts = mk.FusedMoEKernel(
381 prepare_finalize=maybe_make_prepare_finalize(
382 moe=moe_config,
383 quant_config=quant_config,
384 allow_new_interface=True,
385 use_monolithic=False,
386 ),
387 fused_experts=DeepGemmFP4Experts(
388 moe_config=moe_config,
389 quant_config=quant_config,
390 ),
391 )
392
393 # DeepGEMM FP4 path
394 out_deepgemm_fp4 = deep_gemm_fp4_experts.apply(
395 hidden_states=tokens_bf16,
396 w1=w1,
397 w2=w2,
398 topk_weights=topk_weights,
399 topk_ids=topk_ids,
400 global_num_experts=num_experts,
401 activation=MoEActivation.SILU,
402 apply_router_weight_on_input=False,
403 expert_map=None,
404 )
405
406 # BF16 reference using the same original weights
407 out_ref = _bf16_moe_reference(tokens_bf16, w1_bf16, w2_bf16, topk_weights, topk_ids)
408
409 # FP4 vs BF16 reference: quantization error from FP4 weights + FP8 activations
410 diff = calc_diff(out_deepgemm_fp4, out_ref)
411 assert diff < 0.05, f"FP4 diff exceeded 5%: {diff}"
412
413
414 # DeepSeek V4 dims: H=4096, I=2048, so N=2*I=4096, K=H=4096.
415 # FP4 quantization with block_k=32 needs large K for good accuracy.
416 FP4_MNKs = [
417 (128, 4096, 4096), # DeepSeek V4 shape
418 (256, 2048, 2048), # Half-size variant
419 (128, 384, 3584), # Kimi-K3 TP8 latent-MoE shape
420 ]
421
422 FP4_TOPKS = [2]
423 FP4_NUM_EXPERTS = [8]
424
425
426 @pytest.mark.parametrize(("m", "n", "k"), FP4_MNKs)
427 @pytest.mark.parametrize("topk", FP4_TOPKS)
428 @pytest.mark.parametrize("num_experts", FP4_NUM_EXPERTS)
429 @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
430 def test_deepgemm_fp4_vs_triton(
431 m, n, k, topk, num_experts, monkeypatch, workspace_init
432 ):
433 pytest.importorskip("deep_gemm.utils.math")
434 with monkeypatch.context() as mp:
435 mp.setenv("VLLM_USE_DEEP_GEMM", "1")
436
437 _DeepGemmFP4Experts = importlib.import_module(
438 "vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe"
439 ).DeepGemmFP4Experts
440
441 call_counter = {"cnt": 0}
442
443 orig_fn = _DeepGemmFP4Experts.apply
444
445 def _spy_apply(*args, **kwargs):
446 call_counter["cnt"] += 1
447 return orig_fn(*args, **kwargs)
448
449 monkeypatch.setattr(_DeepGemmFP4Experts, "apply", _spy_apply)
450 if topk > num_experts:
451 pytest.skip(f"topk={topk} > num_experts={num_experts}")
452
453 run_single_fp4_case(
454 m=m,
455 n=n,
456 k=k,
457 topk=topk,
458 num_experts=num_experts,
459 )
460
461 # ensure that the DeepGEMM FP4 path was indeed taken.
462 assert call_counter["cnt"] == 1, (
463 f"DeepGEMM FP4 path was not executed during the test. "
464 f"Call counter: {call_counter['cnt']}"
465 )