@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 684 lines 19.4 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import dataclass
5
6 import pytest
7 import torch
8
9 from tests.kernels.moe.utils import (
10 batched_moe,
11 make_dummy_moe_config,
12 make_quantized_test_activations,
13 make_test_weights,
14 naive_batched_moe,
15 )
16 from tests.kernels.quant_utils import native_batched_masked_quant_matmul
17 from tests.kernels.utils import torch_experts
18 from vllm.config import VllmConfig, set_current_vllm_config
19 from vllm.model_executor.layers.fused_moe import fused_topk
20 from vllm.model_executor.layers.fused_moe.activation import (
21 ApplyMoEActivationConfig,
22 MoEActivation,
23 apply_moe_activation,
24 )
25 from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
26 BatchedTritonExperts,
27 invoke_moe_batched_triton_kernel,
28 )
29 from vllm.platforms import current_platform
30 from vllm.triton_utils import tl
31 from vllm.utils.torch_utils import set_random_seed
32
33 DEVICE = current_platform.device_type
34
35 pytestmark = pytest.mark.skipif(
36 not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
37 reason="Triton MoE kernels require CUDA/ROCm/XPU.",
38 )
39
40 MNK_FACTORS = [
41 (1, 128, 128),
42 (1, 512, 512),
43 (1, 1024, 2048),
44 (32, 128, 128),
45 (32, 512, 512),
46 (32, 1024, 2048),
47 (45, 128, 2048),
48 (45, 1024, 128),
49 (64, 512, 512),
50 (64, 1024, 2048),
51 (222, 128, 2048),
52 (222, 1024, 2048),
53 ]
54 NUM_EXPERTS = [8, 64]
55 TOP_KS = [1, 2, 6]
56
57 DTYPES = [torch.bfloat16]
58
59 if not current_platform.is_fp8_fnuz():
60 DTYPES.append(torch.float8_e4m3fn)
61
62 vllm_config = VllmConfig()
63
64
65 @dataclass
66 class BatchedMMConfig:
67 in_dtype: torch.dtype
68 quant_dtype: torch.dtype | None
69 out_dtype: torch.dtype
70 num_experts: int
71 max_tokens_per_expert: int
72 K: int
73 N: int
74
75
76 @dataclass
77 class BatchedMMTensors:
78 A: torch.Tensor # [E, max_tokens, K]
79 B: torch.Tensor # [E, K, N] - column major
80 C: torch.Tensor # [E, max_tokens, N]
81 num_expert_tokens: torch.Tensor # [E]
82
83 @staticmethod
84 def make_tensors(config: BatchedMMConfig):
85 A = (
86 torch.randn(
87 (config.num_experts, config.max_tokens_per_expert, config.K),
88 device=DEVICE,
89 dtype=config.in_dtype,
90 )
91 / 10
92 )
93 B = torch.randn(
94 (config.num_experts, config.N, config.K),
95 device=DEVICE,
96 dtype=config.in_dtype,
97 )
98 C = torch.zeros(
99 (config.num_experts, config.max_tokens_per_expert, config.N),
100 device=DEVICE,
101 dtype=config.out_dtype,
102 )
103
104 num_expert_tokens = torch.randint(
105 low=0,
106 high=config.max_tokens_per_expert,
107 size=(config.num_experts,),
108 device=DEVICE,
109 dtype=torch.int32,
110 )
111
112 return BatchedMMTensors(A, B, C, num_expert_tokens)
113
114
115 @pytest.mark.parametrize("num_experts", [8, 32])
116 @pytest.mark.parametrize("max_tokens_per_expert", [32, 224, 512])
117 @pytest.mark.parametrize("K", [128, 1024])
118 @pytest.mark.parametrize("N", [128, 1024])
119 @pytest.mark.parametrize("dtype", DTYPES)
120 @pytest.mark.parametrize("block_shape", [None, [128, 128]])
121 @pytest.mark.parametrize("per_act_token_quant", [False, True])
122 def test_batched_mm(
123 num_experts: int,
124 max_tokens_per_expert: int,
125 K: int,
126 N: int,
127 dtype: torch.dtype,
128 block_shape: list[int] | None,
129 per_act_token_quant: bool,
130 ):
131 """Note: float8_e4m3fn is not supported on CUDA architecture < 89,
132 and those tests will be skipped on unsupported hardware."""
133 set_random_seed(7)
134
135 use_fp8_w8a8 = dtype == torch.float8_e4m3fn
136
137 if (
138 dtype == torch.float8_e4m3fn
139 and current_platform.is_cuda_alike()
140 and not current_platform.has_device_capability(89)
141 ):
142 pytest.skip(
143 "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
144 )
145
146 if (per_act_token_quant or block_shape is not None) and not use_fp8_w8a8:
147 pytest.skip("Don't test blocking for non-quantized types.")
148
149 if per_act_token_quant and block_shape is not None:
150 pytest.skip("Skip illegal quantization test.")
151
152 if dtype.itemsize == 1:
153 act_dtype = torch.bfloat16
154 quant_dtype = dtype
155 else:
156 act_dtype = dtype
157 quant_dtype = None
158
159 num_expert_tokens = torch.randint(
160 low=0,
161 high=max_tokens_per_expert,
162 size=(num_experts,),
163 device=DEVICE,
164 dtype=torch.int32,
165 )
166
167 A, A_q, A_scale = make_quantized_test_activations(
168 num_experts,
169 max_tokens_per_expert,
170 K,
171 in_dtype=act_dtype,
172 quant_dtype=quant_dtype,
173 block_shape=block_shape,
174 per_act_token_quant=per_act_token_quant,
175 )
176
177 (B, B_q, B_scale, _), _ = make_test_weights(
178 num_experts,
179 N // 2,
180 K,
181 in_dtype=act_dtype,
182 quant_dtype=quant_dtype,
183 block_shape=block_shape,
184 per_out_ch_quant=per_act_token_quant,
185 )
186
187 out_shape = (num_experts, max_tokens_per_expert, N)
188 test_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE)
189 ref_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE)
190 q_ref_output = torch.zeros(out_shape, dtype=act_dtype, device=DEVICE)
191
192 compute_tl_dtype = {
193 torch.float16: tl.float16,
194 torch.bfloat16: tl.bfloat16,
195 torch.float32: tl.float32,
196 }[test_output.dtype]
197
198 assert A_q.dtype == B_q.dtype
199
200 invoke_moe_batched_triton_kernel(
201 A_q,
202 B_q,
203 test_output,
204 num_expert_tokens,
205 compute_tl_dtype,
206 # Quantization data
207 A_scale,
208 B_scale,
209 None,
210 # Quantization schemes
211 use_fp8_w8a8,
212 False,
213 False,
214 config={
215 "BLOCK_SIZE_M": 16,
216 "BLOCK_SIZE_N": 16,
217 "BLOCK_SIZE_K": 16 if dtype.itemsize > 1 else 32,
218 },
219 per_act_token_quant=per_act_token_quant,
220 block_shape=block_shape,
221 )
222
223 ref_output = native_batched_masked_quant_matmul(
224 A,
225 B,
226 ref_output,
227 num_expert_tokens,
228 )
229
230 q_ref_output = native_batched_masked_quant_matmul(
231 A_q,
232 B_q,
233 q_ref_output,
234 num_expert_tokens,
235 A_scale,
236 B_scale,
237 block_shape,
238 per_act_token_quant,
239 )
240
241 rtol, atol = {
242 torch.float16: (6e-2, 6e-2),
243 torch.bfloat16: (6e-2, 6e-2),
244 torch.float32: (1e-2, 1e-2),
245 }[test_output.dtype]
246
247 torch.testing.assert_close(ref_output, q_ref_output, atol=atol, rtol=rtol)
248 torch.testing.assert_close(test_output, q_ref_output, atol=atol, rtol=rtol)
249
250
251 @pytest.mark.parametrize(("m", "n", "k"), MNK_FACTORS)
252 @pytest.mark.parametrize("e", NUM_EXPERTS)
253 @pytest.mark.parametrize("topk", TOP_KS)
254 @pytest.mark.parametrize("dtype", DTYPES)
255 @pytest.mark.parametrize("per_act_token_quant", [False, True])
256 @pytest.mark.parametrize("block_shape", [None, [128, 128]])
257 @pytest.mark.parametrize("input_scales", [False])
258 def test_fused_moe_batched_experts(
259 m: int,
260 n: int,
261 k: int,
262 e: int,
263 topk: int,
264 dtype: torch.dtype,
265 per_act_token_quant: bool,
266 block_shape: list[int] | None,
267 input_scales: bool,
268 workspace_init,
269 ):
270 """Note: float8_e4m3fn is not supported on CUDA architecture < 89,
271 and those tests will be skipped on unsupported hardware."""
272 set_random_seed(7)
273
274 use_fp8_w8a8 = dtype == torch.float8_e4m3fn
275
276 if (
277 dtype == torch.float8_e4m3fn
278 and current_platform.is_cuda_alike()
279 and not current_platform.has_device_capability(89)
280 ):
281 pytest.skip(
282 "Triton limitation: fp8e4nv data type is not supported on CUDA arch < 89"
283 )
284
285 if topk > e:
286 pytest.skip("topk > e")
287
288 if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
289 pytest.skip("Skip quantization test for non-quantized type")
290
291 if per_act_token_quant and block_shape is not None:
292 pytest.skip("Skip illegal quantization test.")
293
294 a = torch.randn((m, k), device=DEVICE, dtype=torch.bfloat16) / 10
295 score = torch.randn((m, e), device=DEVICE, dtype=torch.bfloat16)
296
297 if dtype.itemsize == 1:
298 act_dtype = torch.bfloat16
299 quant_dtype = dtype
300 else:
301 act_dtype = dtype
302 quant_dtype = None
303
304 (w1_16, w1, w1_s, _), (w2_16, w2, w2_s, _) = make_test_weights(
305 e,
306 n,
307 k,
308 block_shape=block_shape,
309 in_dtype=act_dtype,
310 quant_dtype=quant_dtype,
311 per_out_ch_quant=per_act_token_quant,
312 )
313
314 if input_scales and quant_dtype is not None:
315 a1_scale = torch.tensor(1, device=DEVICE, dtype=torch.float32)
316 a2_scale = torch.tensor(1, device=DEVICE, dtype=torch.float32)
317 else:
318 a1_scale = None
319 a2_scale = None
320
321 with set_current_vllm_config(vllm_config):
322 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
323
324 baseline_output = torch_experts(
325 a,
326 w1,
327 w2,
328 topk_weight,
329 topk_ids,
330 w1_scale=w1_s,
331 w2_scale=w2_s,
332 a1_scale=a1_scale,
333 a2_scale=a2_scale,
334 quant_dtype=quant_dtype,
335 per_act_token_quant=per_act_token_quant,
336 block_shape=block_shape,
337 )
338
339 batched_output = naive_batched_moe(
340 a,
341 w1,
342 w2,
343 topk_weight,
344 topk_ids,
345 w1_scale=w1_s,
346 w2_scale=w2_s,
347 a1_scale=a1_scale,
348 a2_scale=a2_scale,
349 quant_dtype=quant_dtype,
350 per_act_token_quant=per_act_token_quant,
351 block_shape=block_shape,
352 )
353
354 triton_output = batched_moe(
355 a,
356 w1,
357 w2,
358 topk_weight,
359 topk_ids,
360 w1_scale=w1_s,
361 w2_scale=w2_s,
362 a1_scale=a1_scale,
363 a2_scale=a2_scale,
364 quant_dtype=quant_dtype,
365 per_act_token_quant=per_act_token_quant,
366 block_shape=block_shape,
367 )
368
369 torch.testing.assert_close(batched_output, baseline_output, atol=3e-2, rtol=2e-2)
370
371 torch.testing.assert_close(triton_output, batched_output, atol=2e-2, rtol=2e-2)
372
373
374 # USE_TD in moe_mmk
375
376 # K % 8 == 0 (bf16, 16-byte alignment)
377 _TD_SHAPES = [
378 (4, 1, 128, 256),
379 (8, 64, 512, 512),
380 (4, 256, 1024, 2048),
381 ]
382
383
384 def _td_supported() -> bool:
385 return current_platform.is_xpu() or (
386 current_platform.is_cuda() and current_platform.has_device_capability(90)
387 )
388
389
390 def _run_td(A, B, num_expert_tokens, use_td: bool):
391 out = torch.zeros(
392 A.shape[0],
393 A.shape[1],
394 B.shape[1],
395 dtype=torch.bfloat16,
396 device=A.device,
397 )
398 import triton
399
400 from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
401 batched_triton_kernel,
402 )
403
404 if use_td:
405 from vllm.triton_utils.allocation import set_triton_allocator
406
407 set_triton_allocator(A.device)
408
409 E = A.shape[0]
410 max_tokens = A.shape[1]
411 K = A.shape[2]
412 N = B.shape[1]
413 BM = BN = BK = 64
414 grid = (E, triton.cdiv(max_tokens, BM) * triton.cdiv(N, BN))
415 batched_triton_kernel[grid](
416 A,
417 B,
418 out,
419 num_expert_tokens,
420 tl.bfloat16,
421 max_tokens,
422 K,
423 N,
424 None,
425 None,
426 None,
427 A.stride(0),
428 A.stride(1),
429 A.stride(2),
430 B.stride(0),
431 B.stride(2),
432 B.stride(1),
433 out.stride(0),
434 out.stride(1),
435 out.stride(2),
436 0,
437 0,
438 0,
439 0,
440 0,
441 0,
442 0,
443 0,
444 False,
445 False,
446 False,
447 BLOCK_M=BM,
448 BLOCK_N=BN,
449 BLOCK_K=BK,
450 USE_TD=use_td,
451 )
452 return out
453
454
455 @pytest.mark.parametrize("num_experts,max_tokens_per_expert,K,N", _TD_SHAPES)
456 def test_batched_mm_td_matches_plain(num_experts, max_tokens_per_expert, K, N):
457 if not _td_supported():
458 pytest.skip("TD requires XPU or CUDA sm_90+")
459 set_random_seed(42)
460 device = current_platform.device_type
461 A = (
462 torch.randn(
463 num_experts, max_tokens_per_expert, K, device=device, dtype=torch.bfloat16
464 )
465 / 10
466 )
467 B = torch.randn(num_experts, N, K, device=device, dtype=torch.bfloat16)
468 num_expert_tokens = torch.randint(
469 1,
470 max_tokens_per_expert + 1,
471 size=(num_experts,),
472 device=device,
473 dtype=torch.int32,
474 )
475
476 out_plain = _run_td(A, B, num_expert_tokens, False)
477 out_td = _run_td(A, B, num_expert_tokens, True)
478
479 torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2)
480
481
482 def test_batched_mm_td_zero_expert_tokens():
483 if not _td_supported():
484 pytest.skip("TD requires XPU or CUDA sm_90+")
485 set_random_seed(42)
486 device = current_platform.device_type
487 E, M, K, N = 8, 32, 256, 256
488 A = torch.randn(E, M, K, device=device, dtype=torch.bfloat16) / 10
489 B = torch.randn(E, N, K, device=device, dtype=torch.bfloat16)
490 num_expert_tokens = torch.zeros(E, device=device, dtype=torch.int32)
491 num_expert_tokens[::2] = M
492
493 out_plain = _run_td(A, B, num_expert_tokens, False)
494 out_td = _run_td(A, B, num_expert_tokens, True)
495
496 for e in range(E):
497 if num_expert_tokens[e].item() == 0:
498 assert out_plain[e].abs().max().item() == 0.0
499 assert out_td[e].abs().max().item() == 0.0
500
501 torch.testing.assert_close(out_td, out_plain, atol=6e-2, rtol=6e-2)
502
503
504 # BatchedTritonExperts device enablement (XPU)
505 def test_batched_triton_experts_supports_current_device():
506 assert BatchedTritonExperts._supports_current_device()
507
508
509 @pytest.mark.parametrize(
510 "activation",
511 [
512 MoEActivation.SITU,
513 MoEActivation.SWIGLUOAI_UNINTERLEAVE,
514 MoEActivation.SWIGLUSTEP,
515 ],
516 )
517 def test_batched_triton_supports_masked_activations(
518 monkeypatch, activation: MoEActivation
519 ):
520 monkeypatch.setattr(current_platform, "is_xpu", lambda: False)
521 assert BatchedTritonExperts._supports_activation(activation)
522
523 monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
524 assert not BatchedTritonExperts._supports_activation(activation)
525
526
527 def _torch_experts_with_shared_activation(
528 a: torch.Tensor,
529 w1: torch.Tensor,
530 w2: torch.Tensor,
531 topk_weights: torch.Tensor,
532 topk_ids: torch.Tensor,
533 activation: MoEActivation,
534 activation_config: ApplyMoEActivationConfig,
535 ) -> torch.Tensor:
536 num_tokens, hidden_dim = a.shape
537 topk = topk_ids.shape[1]
538 routed_input = a[:, None, :].expand(-1, topk, -1).reshape(-1, hidden_dim)
539 routed_experts = topk_ids.reshape(-1)
540 routed_output = torch.zeros_like(routed_input)
541
542 for expert in range(w1.shape[0]):
543 expert_mask = routed_experts == expert
544 if not expert_mask.any():
545 continue
546 projected = routed_input[expert_mask] @ w1[expert].T
547 activated = torch.empty(
548 projected.shape[0], w2.shape[2], dtype=a.dtype, device=a.device
549 )
550 apply_moe_activation(
551 activation,
552 activated,
553 projected,
554 activation_config=activation_config,
555 )
556 routed_output[expert_mask] = activated @ w2[expert].T
557
558 return (
559 (
560 routed_output.view(num_tokens, topk, hidden_dim).float()
561 * topk_weights[..., None]
562 )
563 .sum(dim=1)
564 .to(a.dtype)
565 )
566
567
568 _MASKED_BATCHED_ACTIVATION_CASES = [
569 pytest.param(
570 MoEActivation.SITU,
571 ApplyMoEActivationConfig(
572 activation_situ_beta=1.5,
573 activation_situ_linear_beta=2.0,
574 ),
575 id="situ",
576 ),
577 pytest.param(
578 MoEActivation.SWIGLUOAI_UNINTERLEAVE,
579 ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5),
580 id="swigluoai-uninterleave",
581 ),
582 pytest.param(
583 MoEActivation.SWIGLUSTEP,
584 ApplyMoEActivationConfig(),
585 id="swiglustep",
586 ),
587 ]
588
589
590 @pytest.mark.parametrize(
591 ("activation", "activation_config"), _MASKED_BATCHED_ACTIVATION_CASES
592 )
593 @torch.inference_mode()
594 def test_batched_experts_masked_activation_end_to_end(
595 activation: MoEActivation,
596 activation_config: ApplyMoEActivationConfig,
597 ):
598 if current_platform.is_xpu():
599 pytest.skip("Masked activations are not enabled on XPU")
600
601 from vllm.v1.worker.workspace import init_workspace_manager
602
603 set_random_seed(7)
604 device = current_platform.device_type
605 init_workspace_manager(torch.device(f"{device}:0"))
606 m, n, k, e, topk = 16, 128, 128, 8, 2
607
608 a = torch.randn((m, k), device=device, dtype=torch.bfloat16) / 10
609 score = torch.randn((m, e), device=device, dtype=torch.bfloat16)
610 w1 = torch.randn((e, 2 * n, k), device=device, dtype=torch.bfloat16) / 15
611 w2 = torch.randn((e, k, n), device=device, dtype=torch.bfloat16) / 15
612 moe_config = make_dummy_moe_config(
613 num_experts=e,
614 experts_per_token=topk,
615 hidden_dim=k,
616 intermediate_size=n,
617 in_dtype=a.dtype,
618 activation=activation,
619 )
620 moe_config.swiglu_limit = activation_config.clamp_limit
621 moe_config.swiglu_alpha = activation_config.alpha
622 moe_config.swiglu_beta = activation_config.beta
623 moe_config.activation_situ_beta = activation_config.activation_situ_beta
624 moe_config.activation_situ_linear_beta = (
625 activation_config.activation_situ_linear_beta
626 )
627
628 routing_weights = torch.softmax(score.float(), dim=-1)
629 topk_weights, topk_ids = torch.topk(routing_weights, topk, dim=-1)
630 topk_ids = topk_ids.to(torch.int32)
631 with set_current_vllm_config(vllm_config):
632 baseline_output = _torch_experts_with_shared_activation(
633 a,
634 w1,
635 w2,
636 topk_weights,
637 topk_ids,
638 activation,
639 activation_config,
640 )
641 triton_output = batched_moe(
642 a, w1, w2, topk_weights, topk_ids, moe_config=moe_config
643 )
644
645 torch.testing.assert_close(triton_output, baseline_output, atol=3e-2, rtol=2e-2)
646
647
648 @pytest.mark.parametrize("m,n,k,e,topk", [(32, 512, 512, 8, 2), (45, 1024, 128, 8, 1)])
649 def test_batched_experts_end_to_end(m, n, k, e, topk):
650 """End-to-end BatchedTritonExperts via the reference (no-comms)
651 BatchedPrepareAndFinalize, validated against a torch reference. Exercises
652 the device-enablement path."""
653 from vllm.v1.worker.workspace import init_workspace_manager
654
655 set_random_seed(7)
656 device = current_platform.device_type
657 init_workspace_manager(torch.device(f"{device}:0"))
658
659 a = torch.randn((m, k), device=device, dtype=torch.bfloat16) / 10
660 score = torch.randn((m, e), device=device, dtype=torch.bfloat16)
661 # w1 is the fused gate+up projection [E, 2N, K]; w2 is [E, K, N].
662 w1 = torch.randn((e, 2 * n, k), device=device, dtype=torch.bfloat16) / 15
663 w2 = torch.randn((e, k, n), device=device, dtype=torch.bfloat16) / 15
664
665 with set_current_vllm_config(vllm_config):
666 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
667
668 baseline_output = torch_experts(a, w1, w2, topk_weight, topk_ids)
669 triton_output = batched_moe(a, w1, w2, topk_weight, topk_ids)
670
671 torch.testing.assert_close(triton_output, baseline_output, atol=3e-2, rtol=2e-2)
672
673
674 def test_batched_triton_backend_mapping():
675 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (
676 UnquantizedMoeBackend,
677 map_unquantized_backend,
678 )
679
680 assert (
681 map_unquantized_backend("batched_triton")
682 == UnquantizedMoeBackend.BATCHED_TRITON
683 )
684 assert map_unquantized_backend("triton") == UnquantizedMoeBackend.TRITON