A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 10:00 UTC
[Quantization][Refactor] use platform dict to choose kernel (#33130)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com> Signed-off-by: zofia <110436990+zufangzhu@users.noreply.github.com>
zofia committed
Jan 29, 2026 at 18:44 UTC
a5aa4d5c0f31bba0491a2d9328785dd39dac33c0
1 file changed
+23
-13
vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py
+23
-13
@@ -33,20 +33,30 @@ from vllm.model_executor.layers.quantization.kernels.mixed_precision.MPLinearKer
33
from vllm.model_executor.layers.quantization.kernels.mixed_precision.xpu import ( # noqa: E501
34
XPUwNa16LinearKernel,
35
)
36
-from vllm.platforms import current_platform
36
+from vllm.platforms import PlatformEnum, current_platform
37
38
# in priority/performance order (when available)
39
-_POSSIBLE_KERNELS: list[type[MPLinearKernel]] = [
40
- CutlassW4A8LinearKernel,
41
- MacheteLinearKernel,
42
- AllSparkLinearKernel,
43
- MarlinLinearKernel,
44
- Dynamic4bitLinearKernel,
45
- ConchLinearKernel,
46
- ExllamaLinearKernel,
47
- XPUwNa16LinearKernel,
48
- CPUWNA16LinearKernel,
49
-]
39
+_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
40
+ PlatformEnum.CUDA: [
41
+ CutlassW4A8LinearKernel,
42
+ MacheteLinearKernel,
43
+ AllSparkLinearKernel,
44
+ MarlinLinearKernel,
45
+ ConchLinearKernel,
46
+ ExllamaLinearKernel,
47
+ ],
48
+ PlatformEnum.ROCM: [
49
+ ConchLinearKernel,
50
+ ExllamaLinearKernel,
51
+ ],
52
+ PlatformEnum.XPU: [
53
+ XPUwNa16LinearKernel,
54
+ ],
55
+ PlatformEnum.CPU: [
56
+ Dynamic4bitLinearKernel,
57
+ CPUWNA16LinearKernel,
58
+ ],
59
+}
60
61
62
def choose_mp_linear_kernel(
@@ -78,7 +88,7 @@ def choose_mp_linear_kernel(
88
compute_capability = _cc[0] * 10 + _cc[1]
89
90
failure_reasons = []
81
- for kernel in _POSSIBLE_KERNELS:
91
+ for kernel in _POSSIBLE_KERNELS[current_platform._enum]:
92
if kernel.__name__ in envs.VLLM_DISABLED_KERNELS:
93
failure_reasons.append(
94
f" {kernel.__name__} disabled by environment variable"