@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 135 lines 4.91 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for MXFP8 linear kernel selection logic (CPU-only)
4
5 Run `pytest tests/kernels/quantization/test_mxfp8_kernel_selection.py`.
6 """
7
8 from contextlib import contextmanager
9 from unittest.mock import patch
10
11 import pytest
12
13 from vllm.model_executor.kernels.linear import (
14 FlashInferCutedslMxfp8LinearKernel,
15 FlashInferCutlassMxfp8LinearKernel,
16 MarlinMxfp8LinearKernel,
17 Mxfp8LinearLayerConfig,
18 init_mxfp8_linear_kernel,
19 )
20 from vllm.platforms import PlatformEnum
21
22 pytestmark = pytest.mark.cpu_test
23
24 # Kernels backed by FlashInfer mm_mxfp8, which requires N, K >= 128 and
25 # K % 32 == 0.
26 _MM_MXFP8_KERNELS = [
27 FlashInferCutedslMxfp8LinearKernel,
28 FlashInferCutlassMxfp8LinearKernel,
29 ]
30
31 _SUPPORTED_SHAPE = (4096, 4096)
32 # (N, K): N < 128, K < 128, K % 32 != 0.
33 _UNSUPPORTED_SHAPES = [(64, 4096), (4096, 64), (4096, 4112)]
34
35 # (N, K) with N < 128 or K < 128, which only Marlin runs.
36 _MARLIN_ONLY_SHAPES = [(32, 256), (64, 256), (96, 256), (256, 64)]
37 # (N, K): N <= 0, K <= 0, K % 32 != 0.
38 _MARLIN_UNSUPPORTED_SHAPES = [(0, 256), (256, 0), (256, 100)]
39
40 _SELECTION_CASES = [
41 (100, (128, 128), FlashInferCutedslMxfp8LinearKernel),
42 (100, (130, 256), FlashInferCutedslMxfp8LinearKernel),
43 (103, (128, 128), FlashInferCutedslMxfp8LinearKernel),
44 (103, (130, 256), FlashInferCutedslMxfp8LinearKernel),
45 (120, (128, 128), FlashInferCutlassMxfp8LinearKernel),
46 (120, (130, 256), MarlinMxfp8LinearKernel),
47 (121, (128, 128), FlashInferCutlassMxfp8LinearKernel),
48 (121, (130, 256), MarlinMxfp8LinearKernel),
49 *[
50 (capability, shape, MarlinMxfp8LinearKernel)
51 for capability in (100, 103, 120, 121)
52 for shape in _MARLIN_ONLY_SHAPES
53 ],
54 ]
55
56
57 @pytest.mark.parametrize("kernel_cls", _MM_MXFP8_KERNELS)
58 def test_mm_mxfp8_kernels_accept_supported_shape(kernel_cls):
59 config = Mxfp8LinearLayerConfig(weight_shape=_SUPPORTED_SHAPE)
60 can_implement, reason = kernel_cls.can_implement(config)
61 assert can_implement, reason
62
63
64 @pytest.mark.parametrize("kernel_cls", _MM_MXFP8_KERNELS)
65 @pytest.mark.parametrize("weight_shape", _UNSUPPORTED_SHAPES)
66 def test_mm_mxfp8_kernels_reject_unsupported_shape(kernel_cls, weight_shape):
67 config = Mxfp8LinearLayerConfig(weight_shape=weight_shape)
68 can_implement, reason = kernel_cls.can_implement(config)
69 assert not can_implement
70 assert reason
71
72
73 @pytest.mark.parametrize("weight_shape", _MARLIN_UNSUPPORTED_SHAPES)
74 def test_marlin_mxfp8_rejects_unsupported_shape(weight_shape):
75 config = Mxfp8LinearLayerConfig(weight_shape=weight_shape)
76 can_implement, reason = MarlinMxfp8LinearKernel.can_implement(config)
77 assert not can_implement
78 assert reason
79
80
81 @contextmanager
82 def _patch_cuda_platform(capability: int):
83 """Patch the platform so kernel selection runs as on a CUDA GPU with the
84 given compute capability. The FlashInfer and Marlin kernels report as
85 supported, CuTe-DSL only on SM100/103."""
86 with (
87 patch("vllm.model_executor.kernels.linear.current_platform") as platform,
88 patch(
89 "vllm.model_executor.kernels.linear.mxfp8.flashinfer.current_platform."
90 "is_device_capability_family",
91 side_effect=lambda cap, device_id=0: cap // 10 == capability // 10,
92 ),
93 patch.object(
94 FlashInferCutedslMxfp8LinearKernel,
95 "is_supported",
96 return_value=(capability in (100, 103), None),
97 ),
98 patch.object(
99 FlashInferCutlassMxfp8LinearKernel,
100 "is_supported",
101 return_value=(True, None),
102 ),
103 patch.object(
104 MarlinMxfp8LinearKernel, "is_supported", return_value=(True, None)
105 ),
106 ):
107 platform._enum = PlatformEnum.CUDA
108 yield
109
110
111 @pytest.mark.parametrize(
112 ("capability", "expected"),
113 [(100, True), (103, True), (120, False), (121, False)],
114 )
115 def test_cutlass_mxfp8_requires_aligned_n_only_on_sm12x(capability, expected):
116 """(130, 256) passes the generic mm_mxfp8 limits. Only the SM12x CUTLASS
117 backend needs N % 32 == 0, so SM100/103 must keep accepting it."""
118 config = Mxfp8LinearLayerConfig(weight_shape=(130, 256))
119 with _patch_cuda_platform(capability):
120 can_implement, _ = FlashInferCutlassMxfp8LinearKernel.can_implement(config)
121 assert can_implement == expected
122
123
124 @pytest.mark.parametrize(
125 ("capability", "weight_shape", "expected_kernel_cls"), _SELECTION_CASES
126 )
127 def test_init_mxfp8_linear_kernel(capability, weight_shape, expected_kernel_cls):
128 """A layer that mm_mxfp8 cannot handle must fall through to the next
129 kernel in the CUDA priority list instead of being selected and failing
130 in apply_weights."""
131 with _patch_cuda_platform(capability):
132 kernel = init_mxfp8_linear_kernel(weight_shape=weight_shape)
133
134 assert isinstance(kernel, expected_kernel_cls)
135 assert kernel.config == Mxfp8LinearLayerConfig(weight_shape=weight_shape)