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 14, 05:30 UTC
[Platform][Kernel] platform-specific kernel loading (#25823)
Signed-off-by: Hank <hcc.mayday@gmail.com>
Hank_ committed
Oct 5, 2025 at 19:25 UTC
17edd8a807019c8d1e58634aecb1de7984e8d467
4 files changed
+27
-11
vllm/_custom_ops.py
+2
-11
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import contextlib
4
from typing import TYPE_CHECKING, Optional, Union
5
6
import torch
@@ -13,16 +12,8 @@ from vllm.scalar_type import ScalarType
12
13
logger = init_logger(__name__)
14
16
-if not current_platform.is_tpu() and not current_platform.is_xpu():
17
- try:
18
- import vllm._C
19
- except ImportError as e:
20
- logger.warning("Failed to import from vllm._C with %r", e)
21
-
22
-supports_moe_ops = False
23
-with contextlib.suppress(ImportError):
24
- import vllm._moe_C # noqa: F401
25
- supports_moe_ops = True
15
+current_platform.import_core_kernels()
16
+supports_moe_ops = current_platform.try_import_moe_kernels()
17
18
if TYPE_CHECKING:
19
vllm/platforms/interface.py
+17
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import contextlib
4
import enum
5
import os
6
import platform
@@ -163,6 +164,22 @@ class Platform:
164
else:
165
return device_id
166
167
+ @classmethod
168
+ def import_core_kernels(cls) -> None:
169
+ """ Import any platform-specific C kernels. """
170
+ try:
171
+ import vllm._C # noqa: F401
172
+ except ImportError as e:
173
+ logger.warning("Failed to import from vllm._C: %r", e)
174
+
175
+ @classmethod
176
+ def try_import_moe_kernels(cls) -> bool:
177
+ """ Import any platform-specific MoE kernels. """
178
+ with contextlib.suppress(ImportError):
179
+ import vllm._moe_C # noqa: F401
180
+ return True
181
+ return False
182
+
183
@classmethod
184
def get_vit_attn_backend(cls, head_size: int,
185
dtype: torch.dtype) -> "_Backend":
vllm/platforms/tpu.py
+4
@@ -47,6 +47,10 @@ class TpuPlatform(Platform):
47
"TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS"
48
]
49
50
+ @classmethod
51
+ def import_core_kernels(cls) -> None:
52
+ pass
53
+
54
@classmethod
55
def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
56
dtype: torch.dtype, kv_cache_dtype: Optional[str],
vllm/platforms/xpu.py
+4
@@ -34,6 +34,10 @@ class XPUPlatform(Platform):
34
dist_backend: str = "ccl" # ccl | xccl
35
device_control_env_var: str = "ZE_AFFINITY_MASK"
36
37
+ @classmethod
38
+ def import_core_kernels(cls) -> None:
39
+ pass
40
+
41
@classmethod
42
def get_attn_backend_cls(cls, selected_backend: "_Backend", head_size: int,
43
dtype: torch.dtype, kv_cache_dtype: Optional[str],