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
[Kernel] Centralize platform kernel import in `current_platform.import_kernels` (#26286)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Oct 8, 2025 at 22:25 UTC
4ebc9108a70ef6226a0ed7853b0f6681cd5dc18b
4 files changed
+13
-15
vllm/_custom_ops.py
+2
-3
@@ -12,8 +12,7 @@ from vllm.scalar_type import ScalarType
12
13
logger = init_logger(__name__)
14
15
-current_platform.import_core_kernels()
16
-supports_moe_ops = current_platform.try_import_moe_kernels()
15
+current_platform.import_kernels()
16
17
if TYPE_CHECKING:
18
@@ -1921,7 +1920,7 @@ def moe_wna16_marlin_gemm(
1920
)
1921
1922
1924
-if supports_moe_ops and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
1923
+if hasattr(torch.ops, "_moe_C") and hasattr(torch.ops._moe_C, "marlin_gemm_moe"):
1924
1925
@register_fake("_moe_C::marlin_gemm_moe")
1926
def marlin_gemm_moe_fake(
vllm/platforms/interface.py
+1
-8
@@ -170,22 +170,15 @@ class Platform:
170
return device_id
171
172
@classmethod
173
- def import_core_kernels(cls) -> None:
173
+ def import_kernels(cls) -> None:
174
"""Import any platform-specific C kernels."""
175
try:
176
import vllm._C # noqa: F401
177
except ImportError as e:
178
logger.warning("Failed to import from vllm._C: %r", e)
179
-
180
- @classmethod
181
- def try_import_moe_kernels(cls) -> bool:
182
- """Import any platform-specific MoE kernels."""
179
with contextlib.suppress(ImportError):
180
import vllm._moe_C # noqa: F401
181
186
- return True
187
- return False
188
-
182
@classmethod
183
def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> "_Backend":
184
from vllm.attention.backends.registry import _Backend
vllm/platforms/tpu.py
+5
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import contextlib
5
from typing import TYPE_CHECKING, Optional, Union, cast
6
7
import torch
@@ -45,8 +46,10 @@ class TpuPlatform(Platform):
46
additional_env_vars: list[str] = ["TPU_CHIPS_PER_HOST_BOUNDS", "TPU_HOST_BOUNDS"]
47
48
@classmethod
48
- def import_core_kernels(cls) -> None:
49
- pass
49
+ def import_kernels(cls) -> None:
50
+ # Do not import vllm._C
51
+ with contextlib.suppress(ImportError):
52
+ import vllm._moe_C # noqa: F401
53
54
@classmethod
55
def get_attn_backend_cls(
vllm/platforms/xpu.py
+5
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import contextlib
5
import os
6
from typing import TYPE_CHECKING, Optional
7
@@ -35,8 +36,10 @@ class XPUPlatform(Platform):
36
device_control_env_var: str = "ZE_AFFINITY_MASK"
37
38
@classmethod
38
- def import_core_kernels(cls) -> None:
39
- pass
39
+ def import_kernels(cls) -> None:
40
+ # Do not import vllm._C
41
+ with contextlib.suppress(ImportError):
42
+ import vllm._moe_C # noqa: F401
43
44
@classmethod
45
def get_attn_backend_cls(