| 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/pull/2575 |
| 5 | import itertools |
| 6 | import types |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from tests.kernels.quant_utils import ( |
| 12 | native_per_token_group_quant_fp8, |
| 13 | native_w8a8_block_matmul, |
| 14 | ) |
| 15 | from tests.kernels.utils import fp8_ulp_distance |
| 16 | from vllm.config import VllmConfig |
| 17 | from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( |
| 18 | B12xFp8BlockScaledMMKernel, |
| 19 | _run_b12x_fp8_block_scaled_mm, |
| 20 | ) |
| 21 | from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm |
| 22 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 23 | per_token_group_quant_fp8, |
| 24 | w8a8_triton_block_scaled_mm, |
| 25 | ) |
| 26 | from vllm.platforms import current_platform |
| 27 | from vllm.utils.deep_gemm import ( |
| 28 | fp8_gemm_nt, |
| 29 | get_tma_aligned_size, |
| 30 | is_deep_gemm_e8m0_used, |
| 31 | per_block_cast_to_fp8, |
| 32 | should_use_deepgemm_for_fp8_linear, |
| 33 | ) |
| 34 | from vllm.utils.flashinfer import ( |
| 35 | flashinfer_fp8_blockscale_gemm, |
| 36 | has_flashinfer_fp8_blockscale_gemm, |
| 37 | ) |
| 38 | from vllm.utils.import_utils import has_deep_gemm |
| 39 | |
| 40 | capability = current_platform.get_device_capability() |
| 41 | if capability is None or capability < (9, 0): |
| 42 | pytest.skip("FP8 Triton requires CUDA 9.0 or higher", allow_module_level=True) |
| 43 | |
| 44 | vllm_config = VllmConfig() |
| 45 | |
| 46 | # Test configurations |
| 47 | DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32] |
| 48 | # Quantization test configs |
| 49 | NUM_TOKENS = [7, 2050] |
| 50 | D = [512, 4096, 5120, 13824] |
| 51 | GROUP_SIZE = [64, 128, 512] |
| 52 | COLUMN_MAJOR_SCALES = [True, False] |
| 53 | TMA_ALIGNED_SCALES = [True, False] |
| 54 | # Matmul test configs |
| 55 | M = [1, 7, 8, 83, 4096] |
| 56 | N = [128, 512, 576, 7168, 13824] |
| 57 | K = [256, 3884, 4096, 13824, 16384] |
| 58 | # Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8 |
| 59 | # and its hidden size is 7168. |
| 60 | BLOCK_SIZE = [[128, 128]] |
| 61 | OUT_DTYPES = [torch.bfloat16] # [torch.float32, torch.half, torch.bfloat16] |
| 62 | SEEDS = [0] |
| 63 | |
| 64 | # Skip all tests if CUDA is not available |
| 65 | pytest.importorskip("torch.cuda") |
| 66 | |
| 67 | |
| 68 | @pytest.fixture(autouse=True) |
| 69 | def setup_cuda(): |
| 70 | torch.set_default_device("cuda") |
| 71 | |
| 72 | |
| 73 | @pytest.mark.skipif( |
| 74 | current_platform.is_fp8_fnuz(), |
| 75 | reason="This platform supports e4m3fnuz, not e4m3fn.", |
| 76 | ) |
| 77 | @pytest.mark.parametrize( |
| 78 | "num_tokens,d,dtype,group_size,column_major_scales,tma_aligned_scales,seed", |
| 79 | itertools.product( |
| 80 | NUM_TOKENS, |
| 81 | D, |
| 82 | DTYPES, |
| 83 | GROUP_SIZE, |
| 84 | COLUMN_MAJOR_SCALES, |
| 85 | TMA_ALIGNED_SCALES, |
| 86 | SEEDS, |
| 87 | ), |
| 88 | ) |
| 89 | @torch.inference_mode() |
| 90 | def test_per_token_group_quant_fp8( |
| 91 | num_tokens, d, dtype, group_size, column_major_scales, tma_aligned_scales, seed |
| 92 | ): |
| 93 | torch.manual_seed(seed) |
| 94 | x = torch.rand(num_tokens, d, dtype=dtype) |
| 95 | |
| 96 | ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size) |
| 97 | out, scale = per_token_group_quant_fp8( |
| 98 | x, |
| 99 | group_size, |
| 100 | column_major_scales=column_major_scales, |
| 101 | tma_aligned_scales=tma_aligned_scales, |
| 102 | ) |
| 103 | |
| 104 | if current_platform.is_rocm(): |
| 105 | # On gfx950 the Triton and PyTorch FP8 kernels can round in opposite |
| 106 | # directions when an element lands at the midpoint between two adjacent |
| 107 | # e4m3fn values (1-ULP tie-breaking). Verify: (1) no element is more |
| 108 | # than 1 FP8 ULP away, and (2) fewer than 0.05% of elements have any |
| 109 | # mismatch. Observed worst case across all parameter combos: 0.049%, |
| 110 | # max ULP = 1. |
| 111 | ulp = fp8_ulp_distance(out, ref_out) |
| 112 | assert (ulp <= 1).all(), ( |
| 113 | f"FP8 mismatch > 1 ULP: {int((ulp > 1).sum())} elements" |
| 114 | ) |
| 115 | assert float((ulp > 0).float().mean()) < 5e-4, ( |
| 116 | f"Too many 1-ULP mismatches: {int((ulp > 0).sum())}/{ulp.numel()}" |
| 117 | ) |
| 118 | else: |
| 119 | assert torch.allclose( |
| 120 | out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15 |
| 121 | ) |
| 122 | assert torch.allclose(scale, ref_scale) |
| 123 | |
| 124 | if column_major_scales: |
| 125 | assert scale.stride()[-2] == 1 |
| 126 | if tma_aligned_scales: |
| 127 | assert scale.stride()[-1] == get_tma_aligned_size(num_tokens, 4) |
| 128 | |
| 129 | |
| 130 | @pytest.mark.parametrize( |
| 131 | "M,N,K,block_size,out_dtype,seed", |
| 132 | itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS), |
| 133 | ) |
| 134 | @torch.inference_mode() |
| 135 | def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed): |
| 136 | torch.manual_seed(seed) |
| 137 | factor_for_scale = 1e-2 |
| 138 | fp8_info = torch.finfo(current_platform.fp8_dtype()) |
| 139 | fp8_max, fp8_min = fp8_info.max, fp8_info.min |
| 140 | |
| 141 | A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 142 | A_fp8 = A_fp32.clamp(min=fp8_min, max=fp8_max).to(current_platform.fp8_dtype()) |
| 143 | |
| 144 | B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 145 | B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(current_platform.fp8_dtype()) |
| 146 | |
| 147 | block_n, block_k = block_size[0], block_size[1] |
| 148 | n_tiles = (N + block_n - 1) // block_n |
| 149 | k_tiles = (K + block_k - 1) // block_k |
| 150 | |
| 151 | As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale |
| 152 | Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale |
| 153 | |
| 154 | ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) |
| 155 | out = w8a8_triton_block_scaled_mm(A_fp8, B_fp8, As, Bs, block_size, out_dtype) |
| 156 | |
| 157 | rel_diff = torch.mean( |
| 158 | torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) |
| 159 | ) / torch.mean(torch.abs(ref_out.to(torch.float32))) |
| 160 | assert rel_diff < 0.001 |
| 161 | |
| 162 | |
| 163 | @pytest.mark.skipif( |
| 164 | not current_platform.is_cuda(), reason="CUTLASS only supported on CUDA platform." |
| 165 | ) |
| 166 | @pytest.mark.parametrize( |
| 167 | # 65/66/67 cover all M%4 residue classes above the SM100 swapAB |
| 168 | # threshold (m <= 64); 1026 crosses multiple 128-row SF atoms. |
| 169 | "M", |
| 170 | [32, 65, 66, 67, 1026], |
| 171 | ) |
| 172 | @torch.inference_mode() |
| 173 | def test_w8a8_block_fp8_cutlass_matmul(M): |
| 174 | # Test simple case where weight.shape % 128 != 0, |
| 175 | # like in DSV3 kv_a_proj_with_mqa |
| 176 | N = 576 |
| 177 | K = 7168 |
| 178 | block_size = [128, 128] |
| 179 | out_dtype = torch.bfloat16 |
| 180 | seed = 0 |
| 181 | |
| 182 | torch.manual_seed(seed) |
| 183 | factor_for_scale = 1e-2 |
| 184 | fp8_info = torch.finfo(torch.float8_e4m3fn) |
| 185 | fp8_max, fp8_min = fp8_info.max, fp8_info.min |
| 186 | |
| 187 | A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 188 | |
| 189 | B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 190 | B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) |
| 191 | |
| 192 | block_n, block_k = block_size[0], block_size[1] |
| 193 | n_tiles = (N + block_n - 1) // block_n |
| 194 | k_tiles = (K + block_k - 1) // block_k |
| 195 | |
| 196 | Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale |
| 197 | |
| 198 | A_fp8, As = per_token_group_quant_fp8( |
| 199 | A_fp32, block_size[1], column_major_scales=False |
| 200 | ) |
| 201 | # CUTLASS uses column-major format for scales |
| 202 | A_fp8_cutlass, As_cutlass = per_token_group_quant_fp8( |
| 203 | A_fp32, block_size[1], column_major_scales=True |
| 204 | ) |
| 205 | |
| 206 | ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) |
| 207 | out = cutlass_scaled_mm(A_fp8_cutlass, B_fp8, As_cutlass, Bs, block_size, out_dtype) |
| 208 | |
| 209 | rel_diff = torch.mean( |
| 210 | torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) |
| 211 | ) / torch.mean(torch.abs(ref_out.to(torch.float32))) |
| 212 | assert rel_diff < 0.001 |
| 213 | |
| 214 | |
| 215 | @pytest.mark.skipif( |
| 216 | not (current_platform.is_cuda() and current_platform.has_device_capability(90)), |
| 217 | reason="torch._scaled_mm DeepSeek-style block scaling only supports SM90.", |
| 218 | ) |
| 219 | def test_w8a8_block_fp8_torch_scaled_mm_matmul(): |
| 220 | # BlockWiseTorchFP8ScaledMMLinearKernel: 1x128 activation + 128x128 weight |
| 221 | # block scaling routed through torch._scaled_mm. M=83 is not a multiple of |
| 222 | # 4 so this also exercises the M-padding path. |
| 223 | from vllm.model_executor.kernels.linear.scaled_mm.pytorch import ( |
| 224 | BlockWiseTorchFP8ScaledMMLinearKernel, |
| 225 | ) |
| 226 | |
| 227 | M = 83 |
| 228 | N = 576 |
| 229 | K = 7168 |
| 230 | block_size = [128, 128] |
| 231 | out_dtype = torch.bfloat16 |
| 232 | seed = 0 |
| 233 | |
| 234 | torch.manual_seed(seed) |
| 235 | factor_for_scale = 1e-2 |
| 236 | fp8_info = torch.finfo(torch.float8_e4m3fn) |
| 237 | fp8_max, fp8_min = fp8_info.max, fp8_info.min |
| 238 | |
| 239 | A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 240 | B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 241 | B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn) |
| 242 | |
| 243 | block_n, block_k = block_size[0], block_size[1] |
| 244 | n_tiles = (N + block_n - 1) // block_n |
| 245 | k_tiles = (K + block_k - 1) // block_k |
| 246 | Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale |
| 247 | |
| 248 | # Reference uses row-major activation scales. |
| 249 | A_fp8, As = per_token_group_quant_fp8( |
| 250 | A_fp32, block_size[1], column_major_scales=False |
| 251 | ) |
| 252 | ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) |
| 253 | |
| 254 | # The kernel expects column-major activation scales on CUDA. |
| 255 | A_fp8_cuda, As_cuda = per_token_group_quant_fp8( |
| 256 | A_fp32, block_size[1], column_major_scales=True |
| 257 | ) |
| 258 | |
| 259 | stub = BlockWiseTorchFP8ScaledMMLinearKernel.__new__( |
| 260 | BlockWiseTorchFP8ScaledMMLinearKernel |
| 261 | ) |
| 262 | stub.config = types.SimpleNamespace(out_dtype=out_dtype) |
| 263 | out = stub.apply_block_scaled_mm( |
| 264 | A_fp8_cuda.cuda(), B_fp8.cuda(), As_cuda.cuda(), Bs.cuda() |
| 265 | ) |
| 266 | |
| 267 | ref_out = ref_out.cuda() |
| 268 | rel_diff = torch.mean( |
| 269 | torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) |
| 270 | ) / torch.mean(torch.abs(ref_out.to(torch.float32))) |
| 271 | assert rel_diff < 0.001 |
| 272 | |
| 273 | |
| 274 | @pytest.mark.skipif( |
| 275 | current_platform.is_fp8_fnuz(), |
| 276 | reason="This platform supports e4m3fnuz, not e4m3fn.", |
| 277 | ) |
| 278 | @pytest.mark.parametrize( |
| 279 | "M,N,K,block_size,out_dtype,seed", |
| 280 | itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS), |
| 281 | ) |
| 282 | @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGemm kernels not available.") |
| 283 | @torch.inference_mode() |
| 284 | def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed): |
| 285 | torch.manual_seed(seed) |
| 286 | fp8_info = torch.finfo(torch.float8_e4m3fn) |
| 287 | fp8_max = fp8_info.max |
| 288 | |
| 289 | A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 290 | B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max |
| 291 | |
| 292 | # only aligned sizes are supported by deepgemm |
| 293 | if not should_use_deepgemm_for_fp8_linear( |
| 294 | output_dtype=out_dtype, weight_shape=B_fp32.shape, supports_deep_gemm=True |
| 295 | ): |
| 296 | pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}") |
| 297 | |
| 298 | A_fp8, As_fp8 = per_token_group_quant_fp8( |
| 299 | A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True |
| 300 | ) |
| 301 | B_fp8, Bs_fp8 = per_block_cast_to_fp8( |
| 302 | B_fp32, |
| 303 | block_size=block_size, |
| 304 | use_ue8m0=is_deep_gemm_e8m0_used(), |
| 305 | ) |
| 306 | |
| 307 | As = As_fp8.to(torch.float32) |
| 308 | Bs = Bs_fp8.to(torch.float32) |
| 309 | |
| 310 | ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) |
| 311 | |
| 312 | out = torch.zeros((M, N), device="cuda", dtype=out_dtype) |
| 313 | |
| 314 | assert As_fp8.shape == (M, (K + 127) // 128), ( |
| 315 | f"{As_fp8.shape} != {(M, (K + 127) // 128)}" |
| 316 | ) |
| 317 | |
| 318 | fp8_gemm_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out) |
| 319 | |
| 320 | rel_diff = torch.mean( |
| 321 | torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)) |
| 322 | ) / torch.mean(torch.abs(ref_out.to(torch.float32))) |
| 323 | assert rel_diff < 0.001 |
| 324 | |
| 325 | |
| 326 | @pytest.mark.skipif( |
| 327 | current_platform.is_fp8_fnuz(), |
| 328 | reason="This platform supports e4m3fnuz, not e4m3fn.", |
| 329 | ) |
| 330 | @pytest.mark.parametrize( |
| 331 | "M,N,K,block_size,out_dtype,seed", |
| 332 | itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS), |
| 333 | ) |
| 334 | @torch.inference_mode() |
| 335 | def test_w8a8_block_fp8_flashinfer_matmul(M, N, K, block_size, out_dtype, seed): |
| 336 | if not has_flashinfer_fp8_blockscale_gemm(): |
| 337 | pytest.skip( |
| 338 | "FlashInfer block GEMM not available (requires SM90+ and FlashInfer)" |
| 339 | ) |
| 340 | # only aligned sizes |
| 341 | if K % 128 != 0 or N % 64 != 0: |
| 342 | pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}") |
| 343 | |
| 344 | torch.manual_seed(seed) |
| 345 | fp8_info = torch.finfo(torch.float8_e4m3fn) |
| 346 | fp8_max = fp8_info.max |
| 347 | |
| 348 | A_bf16 = (torch.rand(M, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max |
| 349 | B_bf16 = (torch.rand(N, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max |
| 350 | |
| 351 | A_fp8, As_fp8 = per_token_group_quant_fp8(A_bf16, block_size[1], use_ue8m0=False) |
| 352 | B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_bf16, block_size, use_ue8m0=False) |
| 353 | |
| 354 | As = As_fp8.to(torch.float32) |
| 355 | Bs = Bs_fp8.to(torch.float32) |
| 356 | |
| 357 | ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype) |
| 358 | |
| 359 | out = flashinfer_fp8_blockscale_gemm( |
| 360 | input=A_bf16, |
| 361 | weight=B_fp8, |
| 362 | input_scale=None, |
| 363 | weight_scale=Bs, |
| 364 | out_dtype=out_dtype, |
| 365 | ) |
| 366 | |
| 367 | rel_diff = torch.mean( |
| 368 | torch.abs(out.to(torch.bfloat16) - ref_out.to(torch.bfloat16)) |
| 369 | ) / torch.mean(torch.abs(ref_out.to(torch.bfloat16))) |
| 370 | assert rel_diff < 0.001 |
| 371 | |
| 372 | |
| 373 | @pytest.mark.parametrize( |
| 374 | "M,N,K", |
| 375 | [(1, 128, 256), (8, 256, 512), (129, 256, 256), (2, 4096, 4096)], |
| 376 | ) |
| 377 | @torch.inference_mode() |
| 378 | def test_w8a8_block_fp8_b12x_matmul(M, N, K): |
| 379 | supported, reason = B12xFp8BlockScaledMMKernel.is_supported() |
| 380 | if not supported: |
| 381 | pytest.skip(reason) |
| 382 | |
| 383 | torch.manual_seed(M) |
| 384 | fp8_max = torch.finfo(torch.float8_e4m3fn).max |
| 385 | A_bf16 = (torch.rand(M, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max |
| 386 | B_bf16 = (torch.rand(N, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max |
| 387 | A_fp8, As = per_token_group_quant_fp8(A_bf16, 128, use_ue8m0=False) |
| 388 | B_fp8, Bs = per_block_cast_to_fp8( |
| 389 | B_bf16, |
| 390 | block_size=[128, 128], |
| 391 | use_ue8m0=False, |
| 392 | ) |
| 393 | As = As.float() |
| 394 | Bs = Bs.float() |
| 395 | |
| 396 | ref_out = native_w8a8_block_matmul( |
| 397 | A_fp8, |
| 398 | B_fp8, |
| 399 | As, |
| 400 | Bs, |
| 401 | [128, 128], |
| 402 | torch.bfloat16, |
| 403 | ) |
| 404 | out = _run_b12x_fp8_block_scaled_mm( |
| 405 | A_fp8, |
| 406 | B_fp8, |
| 407 | As, |
| 408 | Bs, |
| 409 | torch.bfloat16, |
| 410 | ) |
| 411 | |
| 412 | rel_diff = torch.mean(torch.abs(out.float() - ref_out.float())) / torch.mean( |
| 413 | torch.abs(ref_out.float()) |
| 414 | ) |
| 415 | cosine = torch.nn.functional.cosine_similarity( |
| 416 | out.float().flatten(), |
| 417 | ref_out.float().flatten(), |
| 418 | dim=0, |
| 419 | ) |
| 420 | # Four-way split-K uses atomic BF16 reductions, so nondeterministic atomic |
| 421 | # ordering can flap an output between adjacent BF16 values one ULP apart. |
| 422 | assert rel_diff < 0.003 |
| 423 | assert cosine >= 0.99999 |