@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 228 lines 6.94 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch
5 import torch.nn.functional as F
6 from nvfp4_utils import (
7 FLOAT4_E2M1_MAX,
8 FLOAT8_E4M3_MAX,
9 convert_swizzled_to_linear,
10 dequantize_nvfp4_to_dtype,
11 )
12
13 from vllm import _custom_ops as ops
14 from vllm.model_executor.kernels.linear.nvfp4 import NvFp4LinearLayerConfig
15 from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
16 FlashInferCuteDslNvFp4W4A16LinearKernel,
17 )
18 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
19 dequantize_to_dtype,
20 )
21 from vllm.platforms import current_platform
22 from vllm.utils.flashinfer import (
23 flashinfer_scaled_fp4_mm,
24 has_flashinfer_b12x_gemm,
25 )
26 from vllm.utils.torch_utils import set_random_seed
27
28 if not current_platform.has_device_capability(100):
29 pytest.skip(
30 reason="Nvfp4 Requires compute capability of 10 or above.",
31 allow_module_level=True,
32 )
33
34 DTYPES = [torch.float16, torch.bfloat16]
35 # m, n, k
36 SHAPES = [
37 (128, 128, 64),
38 (128, 128, 128),
39 (256, 128, 64),
40 (128, 256, 128),
41 (1, 128, 128),
42 ]
43 PAD_SHAPES = [(150, 128, 64), (128, 128, 96), (2, 128, 64), (3, 128, 96)]
44 SHAPES.extend(PAD_SHAPES)
45
46 SEEDS = [42]
47 CUDA_DEVICES = ["cuda:0"]
48
49
50 def get_ref_results(
51 a_fp4,
52 b_fp4,
53 a_sf,
54 b_sf,
55 a_global_scale,
56 b_global_scale,
57 m,
58 n,
59 dtype,
60 block_size,
61 device,
62 is_sf_128x4_layout,
63 ):
64 _, m_k = a_fp4.shape
65 _, n_k = b_fp4.shape
66 assert m_k == n_k
67 a_in_dtype = dequantize_nvfp4_to_dtype(
68 a_fp4,
69 a_sf,
70 a_global_scale,
71 dtype=dtype,
72 device=device,
73 block_size=block_size,
74 is_sf_128x4_layout=is_sf_128x4_layout,
75 )
76 b_in_dtype = dequantize_nvfp4_to_dtype(
77 b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size
78 )
79 return torch.matmul(a_in_dtype, b_in_dtype.t())
80
81
82 @pytest.mark.parametrize("dtype", DTYPES)
83 @pytest.mark.parametrize("shape", SHAPES)
84 @pytest.mark.parametrize("seed", SEEDS)
85 @pytest.mark.parametrize("device", CUDA_DEVICES)
86 @pytest.mark.parametrize("backend", ["cute-dsl", "cutlass", "cudnn", "trtllm", "b12x"])
87 @pytest.mark.parametrize("autotune", [False, True])
88 @torch.inference_mode()
89 def test_flashinfer_nvfp4_gemm(
90 dtype: torch.dtype,
91 shape: tuple[int, int, int],
92 seed: int,
93 device: str,
94 backend: str,
95 autotune: bool,
96 ) -> None:
97 if "trtllm" in backend and dtype == torch.float16:
98 pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")
99 if backend == "cute-dsl" and not current_platform.is_device_capability_family(100):
100 pytest.skip("FlashInfer cutedsl backend is only supported on SM10x")
101 if backend == "b12x" and not current_platform.has_device_capability(120):
102 pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)")
103 if backend == "b12x" and not has_flashinfer_b12x_gemm():
104 pytest.skip("b12x FP4 GEMM backend not available in installed FlashInfer")
105
106 set_random_seed(seed)
107 m, n, packed_k = shape
108 k = packed_k * 2
109 block_size = 16
110 a_dtype = torch.randn((m, k), dtype=dtype, device=device)
111 b_dtype = torch.randn((n, k), dtype=dtype, device=device)
112
113 a_global_scale = (
114 (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
115 ).to(torch.float32)
116 b_global_scale = (
117 (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
118 ).to(torch.float32)
119 alpha = 1.0 / (a_global_scale * b_global_scale)
120
121 # ops.scaled_fp4_quant returns swizzled scales, while weights
122 # from checkpoints are in linear scales.
123 # cutlass and b12x use swizzled scales directly; trtllm needs them unswizzled.
124 a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
125 a_dtype, a_global_scale, is_sf_swizzled_layout=True, backend=backend
126 )
127 is_sf_128x4_layout = not (backend == "trtllm" and m <= 32)
128
129 b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(
130 b_dtype, b_global_scale, is_sf_swizzled_layout=True
131 )
132
133 # get_ref_results unswizzles the scales internally.
134 expected_out = get_ref_results(
135 a_fp4,
136 b_fp4,
137 a_scale_interleaved,
138 b_scale_interleaved,
139 a_global_scale,
140 b_global_scale,
141 m,
142 n,
143 dtype,
144 block_size,
145 device,
146 is_sf_128x4_layout,
147 )
148
149 import flashinfer
150
151 if "trtllm" in backend:
152 epilogue_tile_m = 128
153 b_fp4 = flashinfer.shuffle_matrix_a(b_fp4.view(torch.uint8), epilogue_tile_m)
154 b_scale_interleaved = convert_swizzled_to_linear(
155 b_scale_interleaved, n, k, block_size
156 )
157 b_scale_interleaved = (
158 flashinfer.shuffle_matrix_sf_a(
159 b_scale_interleaved.view(torch.uint8), epilogue_tile_m
160 )
161 .reshape(b_scale_interleaved.shape)
162 .view(torch.float8_e4m3fn)
163 )
164
165 with flashinfer.autotune(autotune):
166 out = flashinfer_scaled_fp4_mm(
167 a_fp4,
168 b_fp4,
169 a_scale_interleaved,
170 b_scale_interleaved,
171 alpha,
172 dtype,
173 backend=backend,
174 )
175
176 torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)
177
178
179 @pytest.mark.parametrize("shape", SHAPES)
180 @pytest.mark.parametrize("seed", SEEDS)
181 @pytest.mark.parametrize("device", CUDA_DEVICES)
182 @torch.inference_mode()
183 def test_flashinfer_cutedsl_nvfp4_w4a16_linear(
184 shape: tuple[int, int, int],
185 seed: int,
186 device: str,
187 ) -> None:
188 supported, reason = FlashInferCuteDslNvFp4W4A16LinearKernel.is_supported()
189 if not supported:
190 pytest.skip(reason)
191
192 set_random_seed(seed)
193 m, n, packed_k = shape
194 k = packed_k * 2
195 x = torch.randn((m, k), dtype=torch.bfloat16, device=device)
196 weight = torch.randn((n, k), dtype=torch.bfloat16, device=device)
197 weight_quant_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / weight.abs().max()).to(
198 torch.float32
199 )
200 weight_global_scale = weight_quant_scale.reciprocal()
201 weight_fp4, weight_scale = ops.scaled_fp4_quant(
202 weight,
203 weight_quant_scale,
204 is_sf_swizzled_layout=False,
205 )
206 weight_ref = dequantize_to_dtype(
207 weight_fp4,
208 weight_scale,
209 weight_global_scale,
210 dtype=torch.bfloat16,
211 swizzle=False,
212 )
213
214 layer = torch.nn.Module()
215 layer.output_size_per_partition = n
216 layer.weight = torch.nn.Parameter(weight_fp4, requires_grad=False)
217 layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
218 layer.weight_global_scale = torch.nn.Parameter(
219 weight_global_scale, requires_grad=False
220 )
221 kernel = FlashInferCuteDslNvFp4W4A16LinearKernel(NvFp4LinearLayerConfig())
222
223 kernel.process_weights_after_loading(layer)
224 output = kernel.apply_weights(layer, x)
225
226 expected = F.linear(x, weight_ref)
227 assert output.shape == (m, n)
228 torch.testing.assert_close(output, expected, atol=1e-1, rtol=1e-1)