@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 121 lines 3.83 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 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8 from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
9 from vllm.model_executor.layers.fused_moe.experts.batched_deep_gemm_moe import (
10 BatchedDeepGemmExperts,
11 )
12 from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
13 BatchedTritonExperts,
14 )
15 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEKernel
16 from vllm.model_executor.layers.fused_moe.prepare_finalize.batched import (
17 BatchedPrepareAndFinalize,
18 )
19 from vllm.utils.deep_gemm import calc_diff, is_deep_gemm_supported
20
21 from .test_deepgemm import make_block_quant_fp8_weights
22 from .utils import make_dummy_moe_config
23
24 BLOCK_SIZE = [128, 128]
25
26
27 @pytest.mark.skipif(not is_deep_gemm_supported(), reason="Requires deep_gemm kernels")
28 @pytest.mark.parametrize("E", [16, 32]) # number of experts
29 @pytest.mark.parametrize("T", [256, 512]) # tokens per expert
30 @pytest.mark.parametrize("K", [128, 256]) # hidden dim
31 @pytest.mark.parametrize("N", [512, 1024]) # intermediate dim per expert
32 @pytest.mark.parametrize("topk", [2, 4])
33 def test_batched_deepgemm_vs_triton(
34 E: int, T: int, K: int, N: int, topk: int, monkeypatch, workspace_init
35 ):
36 """Compare BatchedDeepGemmExperts to BatchedTritonExperts."""
37 monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1")
38
39 device = "cuda"
40 w1, w2, w1_s, w2_s = make_block_quant_fp8_weights(E, N, K, BLOCK_SIZE)
41
42 M = E * T # total tokens
43 a = torch.randn(M, K, device=device, dtype=torch.bfloat16) / 10.0
44 fp8_info = torch.finfo(torch.float8_e4m3fn)
45 a.clamp_(fp8_info.min, fp8_info.max)
46
47 # random router outputs → top-k indices / weights
48 router_logits = torch.randn(M, E, device=device, dtype=torch.float32)
49 topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
50 topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
51
52 # token number for each expert
53 cnt = torch.bincount(topk_ids.flatten(), minlength=E)
54 max_cnt = int(cnt.max().item())
55 # next power of 2 for max token number
56 max_num_tokens = 1 << (max_cnt - 1).bit_length()
57
58 prep_finalize = BatchedPrepareAndFinalize(
59 max_num_tokens=max_num_tokens,
60 num_local_experts=E,
61 num_dispatchers=1,
62 rank=0,
63 )
64
65 quant_config = fp8_w8a8_moe_quant_config(
66 w1_scale=w1_s,
67 w2_scale=w2_s,
68 per_act_token_quant=False,
69 block_shape=BLOCK_SIZE,
70 )
71
72 # triton (reference)
73 triton_experts = BatchedTritonExperts(
74 max_num_tokens=max_num_tokens,
75 num_dispatchers=1,
76 quant_config=quant_config,
77 moe_config=make_dummy_moe_config(),
78 )
79 mk_triton = FusedMoEKernel(
80 prep_finalize,
81 triton_experts,
82 )
83
84 out_triton = mk_triton.apply(
85 hidden_states=a,
86 w1=w1,
87 w2=w2,
88 topk_weights=topk_weights,
89 topk_ids=topk_ids,
90 activation=MoEActivation.SILU,
91 global_num_experts=E,
92 expert_map=None,
93 apply_router_weight_on_input=False,
94 )
95
96 # deepgemm
97 deepgemm_experts = BatchedDeepGemmExperts(
98 max_num_tokens=max_num_tokens,
99 num_dispatchers=1,
100 quant_config=quant_config,
101 moe_config=make_dummy_moe_config(),
102 )
103 mk_deepgemm = FusedMoEKernel(
104 prep_finalize,
105 deepgemm_experts,
106 )
107
108 out_deepgemm = mk_deepgemm.apply(
109 hidden_states=a,
110 w1=w1,
111 w2=w2,
112 topk_weights=topk_weights,
113 topk_ids=topk_ids,
114 activation=MoEActivation.SILU,
115 global_num_experts=E,
116 expert_map=None,
117 apply_router_weight_on_input=False,
118 )
119
120 diff = calc_diff(out_deepgemm, out_triton)
121 assert diff < 1e-3, f"Output diff too large: {diff}"