@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 169 lines 5.54 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_triton_moe_channel_fp8_kernel.py
5 import itertools
6
7 import pytest
8 import torch
9
10 from tests.kernels.moe.utils import fused_moe
11 from vllm import _custom_ops as ops
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.config import fp8_w8a8_moe_quant_config
15 from vllm.platforms import current_platform
16
17 if current_platform.get_device_capability() < (9, 0):
18 pytest.skip("FP8 Triton requires CUDA 9.0 or higher", allow_module_level=True)
19
20 vllm_config = VllmConfig()
21
22 if current_platform.is_fp8_fnuz():
23 pytest.skip(
24 "Tests in this file require float8_e4m3fn and platform does not support",
25 allow_module_level=True,
26 )
27
28
29 def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
30 """Matrix multiplication function that supports per-token input
31 quantization and per-column weight quantization"""
32 A = A.to(torch.float32)
33 B = B.to(torch.float32)
34
35 assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
36 assert B.ndim == 2 and B.is_contiguous(), "B must be a 2D contiguous tensor"
37
38 # Reshape input
39 M = A.numel() // A.shape[-1]
40 B = B.t() # Transpose weight matrix
41 N, K = B.shape
42 origin_C_shape = A.shape[:-1] + (K,)
43 A = A.reshape(M, N)
44
45 # As is per-token [M, 1], Bs is per-column [1, K]
46 C = torch.matmul(A, B) # [M, K]
47 C = As * C * Bs.view(1, -1) # Broadcast per-column scale
48
49 return C.reshape(origin_C_shape).to(output_dtype)
50
51
52 def fp8_mask(a, mask):
53 dtype = a.dtype
54 return a.view(torch.int8)[mask].view(dtype)
55
56
57 def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
58 """This function performs fused moe with per-column int8
59 quantization using native torch."""
60 B, D = a.shape
61 # Perform per-token quantization
62 a_q, a_s = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True)
63 # Repeat tokens to match topk
64 a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
65 # Also repeat the scale
66 a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1) # [B*topk, 1]
67
68 out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
69
70 # Calculate routing
71 score = torch.softmax(score, dim=-1, dtype=torch.float32)
72 topk_weight, topk_ids = torch.topk(score, topk)
73 topk_weight = topk_weight.view(-1)
74 topk_ids = topk_ids.view(-1)
75 # Process each expert
76 for i in range(w1.shape[0]):
77 mask = topk_ids == i
78 if mask.sum():
79 # First MLP layer: note that a_s is now per-token
80 inter_out = native_w8a8_per_token_matmul(
81 fp8_mask(a_q, mask),
82 w1[i],
83 fp8_mask(a_s, mask),
84 w1_s[i],
85 output_dtype=a.dtype,
86 )
87 # Activation function
88 act_out = SiluAndMul().forward_native(inter_out)
89 # Quantize activation output with per-token
90 act_out_q, act_out_s = ops.scaled_fp8_quant(
91 act_out, use_per_token_if_dynamic=True
92 )
93
94 # Second MLP layer
95 out[mask] = native_w8a8_per_token_matmul(
96 act_out_q, w2[i], act_out_s, w2_s[i], output_dtype=a.dtype
97 )
98 # Apply routing weights and sum
99 return (
100 out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
101 ).sum(dim=1)
102
103
104 @pytest.fixture(autouse=True)
105 def setup_cuda():
106 """Sets the default CUDA device before each test in this module."""
107 torch.set_default_device("cuda")
108
109
110 DTYPES = [torch.half, torch.bfloat16]
111 M = [1, 33]
112 N = [128, 1024]
113 K = [256, 4096]
114 E = [8]
115 TOP_KS = [2, 6]
116 SEEDS = [0]
117
118
119 @pytest.mark.parametrize(
120 "M, N, K, E, topk, dtype, seed",
121 itertools.product(M, N, K, E, TOP_KS, DTYPES, SEEDS),
122 )
123 @torch.inference_mode()
124 def test_w8a8_fp8_fused_moe(M, N, K, E, topk, dtype, seed):
125 torch.manual_seed(seed)
126 # Initialize int8 quantization parameters
127 factor_for_scale = 1e-2
128 finfo = torch.finfo(torch.float8_e4m3fn)
129 fp8_max = finfo.max
130 fp8_min = finfo.min
131
132 # Input tensor
133 # M * K
134 a = torch.randn((M, K), dtype=dtype) / 10
135
136 # Generate int8 weights
137 w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
138 w1 = (w1_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
139
140 w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
141 w2 = (w2_fp32 * fp8_max).clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
142
143 # Generate scale for each column (per-column quantization)
144 w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
145 w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
146 score = torch.randn((M, E), dtype=dtype)
147
148 with set_current_vllm_config(vllm_config):
149 ref_out = torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk)
150 out = fused_moe(
151 a,
152 w1,
153 w2,
154 score,
155 topk,
156 renormalize=False,
157 quant_config=fp8_w8a8_moe_quant_config(
158 per_act_token_quant=True,
159 w1_scale=w1_s,
160 w2_scale=w2_s,
161 block_shape=None, # Not using block quantization
162 ),
163 )
164
165 # Check results
166 rel_diff = torch.mean(
167 torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
168 ) / torch.mean(torch.abs(ref_out.to(torch.float32)))
169 assert rel_diff < 0.05