@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 5.25 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 tests.kernels.moe.utils import make_test_quant_config
8 from tests.kernels.quant_utils import (
9 native_per_token_group_quant_int8,
10 native_w8a8_block_matmul,
11 )
12 from vllm.config import VllmConfig, set_current_vllm_config
13 from vllm.model_executor.layers.activation import SiluAndMul
14 from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
15 from vllm.model_executor.layers.quantization.utils.int8_utils import (
16 per_token_group_quant_int8,
17 )
18 from vllm.platforms import current_platform
19
20 if current_platform.get_device_capability() < (7, 0):
21 pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True)
22
23 vllm_config = VllmConfig()
24
25 DTYPES = [torch.bfloat16]
26
27 MNK_FACTORS = [
28 (1, 128, 128),
29 (1, 128, 7168),
30 (1, 1024, 7168),
31 (1, 4096, 512),
32 (1, 4096, 7168),
33 (33, 512, 512),
34 (33, 128, 7168),
35 (33, 1024, 7168),
36 (33, 4096, 128),
37 (33, 4096, 7168),
38 (128, 128, 128),
39 (128, 1024, 7168),
40 (128, 4096, 512),
41 (128, 4096, 7168),
42 (222, 512, 512),
43 (222, 1024, 7168),
44 (222, 4096, 7168),
45 (2048, 128, 128),
46 (2048, 1024, 7168),
47 (2048, 4096, 4096),
48 ]
49
50 E = [8, 24]
51 TOP_KS = [2, 6]
52 # BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
53 BLOCK_SIZE = [[128, 128]]
54 SEEDS = [0]
55
56
57 # For test
58 def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
59 """This function performs fused moe with block-wise quantization using
60 native torch."""
61 B, D = a.shape
62 a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
63 out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
64 score = torch.softmax(score, dim=-1, dtype=torch.float32)
65 topk_weight, topk_ids = torch.topk(score, topk)
66 topk_weight = topk_weight.view(-1)
67 topk_ids = topk_ids.view(-1)
68
69 _, block_k = block_shape[0], block_shape[1]
70 a_q, a_s = native_per_token_group_quant_int8(a, block_k, round_to_nearest=False)
71 for i in range(w1.shape[0]):
72 mask = topk_ids == i
73 if mask.sum():
74 inter_out = native_w8a8_block_matmul(
75 a_q[mask], w1[i], a_s[mask], w1_s[i], block_shape, output_dtype=a.dtype
76 )
77 act_out = SiluAndMul().forward_native(inter_out)
78 act_out_q, act_out_s = native_per_token_group_quant_int8(
79 act_out, block_k, round_to_nearest=False
80 )
81 act_out = act_out.to(torch.float32)
82 out[mask] = native_w8a8_block_matmul(
83 act_out_q, w2[i], act_out_s, w2_s[i], block_shape, output_dtype=a.dtype
84 )
85 return (
86 out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
87 ).sum(dim=1)
88
89
90 @pytest.fixture(autouse=True)
91 def setup_cuda():
92 """Sets the default CUDA device before each test in this module."""
93 torch.set_default_device("cuda")
94
95
96 @pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS)
97 @pytest.mark.parametrize("E", E)
98 @pytest.mark.parametrize("topk", TOP_KS)
99 @pytest.mark.parametrize("block_size", BLOCK_SIZE)
100 @pytest.mark.parametrize("dtype", DTYPES)
101 @pytest.mark.parametrize("seed", SEEDS)
102 @torch.inference_mode()
103 def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
104 """Tests the fused_moe kernel with W8A8 INT8 block quantization against a
105 native torch reference."""
106 torch.manual_seed(seed)
107
108 a = torch.randn((M, K), dtype=dtype) / 10
109 score = torch.randn((M, E), dtype=dtype)
110 topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
111
112 w1, w2, quant_config = make_test_quant_config(
113 E,
114 N,
115 K,
116 dtype,
117 quant_dtype=torch.int8,
118 per_act_token_quant=False,
119 block_shape=block_size,
120 )
121
122 # Set the context to avoid lots of warning spam.
123 with set_current_vllm_config(vllm_config):
124 out = fused_experts(
125 a, w1, w2, topk_weights, topk_ids, quant_config=quant_config
126 )
127 ref_out = torch_w8a8_block_int8_moe(
128 a,
129 w1,
130 w2,
131 quant_config.w1_scale,
132 quant_config.w2_scale,
133 score,
134 topk,
135 block_size,
136 )
137
138 # Check results
139 torch.testing.assert_close(out, ref_out, atol=0.065, rtol=0.065)
140
141
142 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
143 def test_block_int8_quantization_truncates_fractional_values(dtype):
144 """Block INT8 quantization truncates, unlike per-token round-to-nearest."""
145 device = current_platform.device_type
146 x = torch.tensor(
147 [[-127, -1.75, -0.75, 0.75, 1.75, 127] + [0] * 122],
148 dtype=dtype,
149 device=device,
150 )
151 expected = torch.tensor(
152 [[-127, -1, 0, 0, 1, 127] + [0] * 122], dtype=torch.int8, device=device
153 )
154
155 out_q, scale = per_token_group_quant_int8(x, group_size=128)
156 ref_q, ref_s = native_per_token_group_quant_int8(
157 x, group_size=128, round_to_nearest=False
158 )
159
160 torch.testing.assert_close(out_q, expected, atol=0, rtol=0)
161 torch.testing.assert_close(ref_q, expected, atol=0, rtol=0)
162 torch.testing.assert_close(scale, torch.ones_like(scale), atol=0, rtol=0)
163 torch.testing.assert_close(ref_s, scale, atol=0, rtol=0)