@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 187 lines 6.21 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_int8_kernel.py
5 import itertools
6
7 import pytest
8 import torch
9
10 from tests.kernels.quant_utils import native_per_token_group_quant_int8
11 from vllm.model_executor.layers.activation import SiluAndMul
12 from vllm.model_executor.layers.fused_moe import fused_experts
13 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
14 from vllm.model_executor.layers.quantization.utils.int8_utils import (
15 per_token_quant_int8,
16 )
17 from vllm.platforms import current_platform
18
19 if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
20 pytest.skip(
21 "INT8 Triton kernels require a CUDA-alike or XPU device",
22 allow_module_level=True,
23 )
24
25 if current_platform.is_cuda_alike() and not current_platform.has_device_capability(
26 (7, 0)
27 ):
28 pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True)
29
30
31 def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
32 """Matrix multiplication function that supports per-token input
33 quantization and per-column weight quantization"""
34 A = A.to(torch.float32)
35 B = B.to(torch.float32)
36
37 assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
38 assert B.ndim == 2 and B.is_contiguous(), "B must be a 2D contiguous tensor"
39
40 # Reshape input
41 M = A.numel() // A.shape[-1]
42 B = B.t() # Transpose weight matrix
43 N, K = B.shape
44 origin_C_shape = A.shape[:-1] + (K,)
45 A = A.reshape(M, N)
46
47 # As is per-token [M, 1], Bs is per-column [1, K]
48 C = torch.matmul(A, B) # [M, K]
49 C = As * C * Bs.view(1, -1) # Broadcast per-column scale
50
51 return C.reshape(origin_C_shape).to(output_dtype)
52
53
54 def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids):
55 """This function performs fused moe with per-column int8 quantization
56 using native torch."""
57 B, D = a.shape
58 # Perform per-token quantization
59 a_q, a_s = per_token_quant_int8(a)
60 # Repeat tokens to match topk
61 a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
62 # Also repeat the scale
63 a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1) # [B*topk, 1]
64
65 out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
66
67 # Calculate routing
68 topk_weight = topk_weight.view(-1)
69 topk_ids = topk_ids.view(-1)
70 # Process each expert
71 for i in range(w1.shape[0]):
72 mask = topk_ids == i
73 if mask.sum():
74 # First MLP layer: note that a_s is now per-token
75 inter_out = native_w8a8_per_token_matmul(
76 a_q[mask], w1[i], a_s[mask], w1_s[i], output_dtype=a.dtype
77 )
78 # Activation function
79 act_out = SiluAndMul().forward_native(inter_out)
80 # Quantize activation output with per-token
81 act_out_q, act_out_s = per_token_quant_int8(act_out)
82
83 # Second MLP layer
84 out[mask] = native_w8a8_per_token_matmul(
85 act_out_q, w2[i], act_out_s, w2_s[i], output_dtype=a.dtype
86 )
87 # Apply routing weights and sum
88 return (
89 out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
90 ).sum(dim=1)
91
92
93 DTYPES = [torch.half, torch.bfloat16]
94 M = [1, 33]
95 N = [128, 1024]
96 K = [256, 4096]
97 E = [8]
98 TOP_KS = [2, 6]
99 SEEDS = [0]
100
101
102 @pytest.mark.parametrize(
103 "M, N, K, E, topk, dtype, seed",
104 itertools.product(M, N, K, E, TOP_KS, DTYPES, SEEDS),
105 )
106 @torch.inference_mode()
107 def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
108 torch.manual_seed(seed)
109 device = current_platform.device_type
110 # Initialize int8 quantization parameters
111 factor_for_scale = 1e-2
112 int8_max = 127
113 int8_min = -128
114
115 # Input tensor
116 # M * K
117 a = torch.randn((M, K), dtype=dtype, device=device) / 10
118
119 # Generate int8 weights
120 w1_fp32 = (
121 torch.rand(
122 (E, 2 * N, K),
123 dtype=torch.float32,
124 device=device,
125 )
126 - 0.5
127 ) * 2
128 w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
129
130 w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32, device=device) - 0.5) * 2
131 w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
132
133 # Generate scale for each column (per-column quantization)
134 w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
135 w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
136 score = torch.randn((M, E), dtype=dtype, device=device)
137 score = torch.softmax(score, dim=-1, dtype=torch.float32)
138 topk_weights, topk_ids = torch.topk(score, topk)
139
140 ref_out = torch_w8a8_per_column_moe(
141 a, w1, w2, w1_s, w2_s, topk, topk_weights, topk_ids
142 )
143
144 quant_config = FusedMoEQuantConfig.make(
145 torch.int8,
146 per_act_token_quant=True,
147 block_shape=None,
148 w1_scale=w1_s,
149 w2_scale=w2_s,
150 )
151
152 out = fused_experts(
153 a,
154 w1,
155 w2,
156 topk_weights,
157 topk_ids,
158 quant_config=quant_config,
159 )
160
161 # Check results
162 rel_diff = torch.mean(
163 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
164 ) / torch.mean(torch.abs(ref_out.to(torch.float32)))
165 assert rel_diff < 0.05
166
167
168 @pytest.mark.parametrize("num_tokens", [1, 33, 64])
169 @pytest.mark.parametrize("hidden_dim", [128, 1024])
170 @pytest.mark.parametrize("dtype", DTYPES)
171 @torch.inference_mode()
172 def test_per_token_quant_int8(num_tokens, hidden_dim, dtype):
173 """Covers the round_int8 Triton kernel via per_token_quant_int8."""
174 torch.manual_seed(0)
175 device = current_platform.device_type
176 x = torch.randn((num_tokens, hidden_dim), dtype=dtype, device=device) * 8
177
178 q, s = per_token_quant_int8(x)
179 # Per-token int8 quant is per-token-group quant with group_size = hidden_dim.
180 ref_q, ref_s = native_per_token_group_quant_int8(x, hidden_dim)
181
182 assert q.dtype == torch.int8
183 assert q.shape == (num_tokens, hidden_dim)
184 assert s.shape == (num_tokens, 1)
185 assert torch.allclose(s.float(), ref_s.float(), atol=1e-3, rtol=1e-3)
186 # Quantized values may differ by at most 1 ULP due to rounding.
187 assert torch.all((q.float() - ref_q.float()).abs() <= 1)