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 5, 18:30 UTC
[OOT] Add OOT support for linear kernel. (#37989)
Signed-off-by: menogrey <1299267905@qq.com>
zhangyiming committed
Mar 31, 2026 at 14:33 UTC
1ac66942975e163d3cc8a4bbb9b1832f9f7edc05
2 files changed
+76
tests/kernels/quantization/test_scaled_mm_kernel_selection.py
+42
@@ -7,15 +7,21 @@ Run `pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py`.
7
8
import inspect
9
from abc import ABC
10
+from unittest.mock import patch
11
12
import pytest
13
+import torch
14
15
from vllm.model_executor.kernels.linear import (
16
AiterInt8ScaledMMLinearKernel,
17
CPUInt8ScaledMMLinearKernel,
18
+ Int8ScaledMMLinearKernel,
19
Int8ScaledMMLinearLayerConfig,
20
ScaledMMLinearKernel,
21
+ init_int8_linear_kernel,
22
+ register_linear_kernel,
23
)
24
+from vllm.platforms import PlatformEnum
25
26
pytestmark = pytest.mark.cpu_test
27
@@ -85,3 +91,39 @@ def test_cpu_kernel_accepts_all_configs():
91
assert can_impl, (
92
f"CPUInt8ScaledMMLinearKernel should accept config {config}: {reason}"
93
)
94
+
95
+
96
+class OOTInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
97
+ @classmethod
98
+ def is_supported(
99
+ cls, compute_capability: int | None = None
100
+ ) -> tuple[bool, str | None]:
101
+ return True, None
102
+
103
+ @classmethod
104
+ def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
105
+ return True, None
106
+
107
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
108
+ pass
109
+
110
+ def apply_weights(
111
+ self,
112
+ layer: torch.nn.Module,
113
+ x: torch.Tensor,
114
+ bias: torch.Tensor | None = None,
115
+ ) -> torch.Tensor:
116
+ pass
117
+
118
+
119
+@patch("vllm.model_executor.kernels.linear.current_platform")
120
+def test_register_oot_linear_kernel(platform_mock):
121
+ """Test that the linear kernel registration works correctly."""
122
+ platform_mock._enum = PlatformEnum.OOT
123
+ register_linear_kernel(OOTInt8ScaledMMLinearKernel, PlatformEnum.OOT, "int8")
124
+
125
+ kernel = init_int8_linear_kernel(True, True, True, "module")
126
+
127
+ assert isinstance(kernel, OOTInt8ScaledMMLinearKernel), (
128
+ "init_int8_linear_kernel should return an instance of the registered kernel"
129
+ )
vllm/model_executor/kernels/linear/__init__.py
+34
@@ -367,10 +367,44 @@ def choose_mp_linear_kernel(
367
)
368
369
370
+def register_linear_kernel(
371
+ kernel_class: type,
372
+ platform: PlatformEnum,
373
+ kernel_type: str = "mp",
374
+) -> None:
375
+ """
376
+ Register a new linear kernel class to be considered in kernel selection.
377
+
378
+ Args:
379
+ kernel_class (type): The kernel class to register.
380
+ platform (PlatformEnum): The platform for which this kernel is applicable.
381
+ kernel_type (str): The type of the kernel, either "mp", "int8", or "fp8".
382
+ Defaults to "mp".
383
+
384
+ Raises:
385
+ ValueError: If the kernel_type is not recognized.
386
+ """
387
+ if kernel_type == "mp":
388
+ if platform not in _POSSIBLE_KERNELS:
389
+ _POSSIBLE_KERNELS[platform] = []
390
+ _POSSIBLE_KERNELS[platform].append(kernel_class)
391
+ elif kernel_type == "int8":
392
+ if platform not in _POSSIBLE_INT8_KERNELS:
393
+ _POSSIBLE_INT8_KERNELS[platform] = []
394
+ _POSSIBLE_INT8_KERNELS[platform].append(kernel_class)
395
+ elif kernel_type == "fp8":
396
+ if platform not in _POSSIBLE_FP8_KERNELS:
397
+ _POSSIBLE_FP8_KERNELS[platform] = []
398
+ _POSSIBLE_FP8_KERNELS[platform].append(kernel_class)
399
+ else:
400
+ raise ValueError(f"Unrecognized kernel type: {kernel_type}")
401
+
402
+
403
__all__ = [
404
"init_fp8_linear_kernel",
405
"init_int8_linear_kernel",
406
"choose_mp_linear_kernel",
407
+ "register_linear_kernel",
408
"FP8ScaledMMLinearKernel",
409
"Int8ScaledMMLinearKernel",
410
"ScaledMMLinearKernel",