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 11, 13:45 UTC
[ROCm] Enable Triton ScaledMM fallback + kernel selection fix (#26668)
Signed-off-by: Shivam <shivampr.dev@gmail.com> Signed-off-by: Shivam <shivamprasad91@gmail.com>
shivampr committed
Dec 12, 2025 at 10:28 UTC
cd7740ac5c3906b2913d58ade61f231ec3a93296
9 files changed
+193
-69
.buildkite/test-pipeline.yaml
+1
-1
@@ -836,7 +836,7 @@ steps:
836
- tests/models/multimodal
837
no_gpu: true
838
commands:
839
- - pip install git+https://github.com/TIGER-AI-Lab/Mantis.git
839
+ - "pip install git+https://github.com/TIGER-AI-Lab/Mantis.git || echo 'Mantis installation skipped (decord not available on CPU-only environment)'"
840
- pytest -v -s models/multimodal/processing --ignore models/multimodal/processing/test_tensor_schema.py
841
842
- label: Multi-Modal Processor Test
tests/kernels/quantization/test_scaled_mm_kernel_selection.py
new
+91
@@ -0,0 +1,91 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for ScaledMM kernel selection logic (CPU-only)
4
+
5
+Run `pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py`.
6
+"""
7
+
8
+import inspect
9
+from abc import ABC
10
+
11
+import pytest
12
+
13
+from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
14
+ ScaledMMLinearLayerConfig,
15
+)
16
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (
17
+ AiterScaledMMLinearKernel,
18
+)
19
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.cpu import (
20
+ CPUScaledMMLinearKernel,
21
+)
22
+from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
23
+ ScaledMMLinearKernel,
24
+)
25
+
26
+pytestmark = pytest.mark.cpu_test
27
+
28
+
29
+def test_is_supported_is_abstract():
30
+ """Test that is_supported() is properly defined as abstract."""
31
+ assert issubclass(ScaledMMLinearKernel, ABC)
32
+ assert hasattr(ScaledMMLinearKernel, "is_supported")
33
+
34
+
35
+def test_cpu_kernel_implements_is_supported():
36
+ """Test that CPUScaledMMLinearKernel implements is_supported() method."""
37
+ assert hasattr(CPUScaledMMLinearKernel, "is_supported"), (
38
+ "CPUScaledMMLinearKernel missing is_supported() method"
39
+ )
40
+ # Verify it's a classmethod by checking if it can be called with the class
41
+ # and by checking the method type
42
+ assert inspect.ismethod(CPUScaledMMLinearKernel.is_supported) or inspect.isfunction(
43
+ CPUScaledMMLinearKernel.is_supported
44
+ ), "CPUScaledMMLinearKernel.is_supported() should be a classmethod"
45
+ # Verify it can be called as a classmethod
46
+ result, reason = CPUScaledMMLinearKernel.is_supported()
47
+ assert isinstance(result, bool), "is_supported() should return a bool"
48
+ assert reason is None or isinstance(reason, str), "reason should be str or None"
49
+
50
+
51
+def test_aiter_kernel_implements_is_supported():
52
+ """Test that AiterScaledMMLinearKernel implements is_supported() method."""
53
+ assert hasattr(AiterScaledMMLinearKernel, "is_supported"), (
54
+ "AiterScaledMMLinearKernel missing is_supported() method"
55
+ )
56
+ # Verify it's a classmethod by checking if it can be called with the class
57
+ # and by checking the method type
58
+ assert inspect.ismethod(
59
+ AiterScaledMMLinearKernel.is_supported
60
+ ) or inspect.isfunction(AiterScaledMMLinearKernel.is_supported), (
61
+ "AiterScaledMMLinearKernel.is_supported() should be a classmethod"
62
+ )
63
+ # Verify it can be called as a classmethod
64
+ # (will return False on CPU, which is expected)
65
+ result, reason = AiterScaledMMLinearKernel.is_supported()
66
+ assert isinstance(result, bool), "is_supported() should return a bool"
67
+ assert reason is None or isinstance(reason, str), "reason should be str or None"
68
+ # On CPU, it should return False with a reason about requiring ROCm
69
+ # This validates the method works correctly even on non-ROCm platforms
70
+
71
+
72
+def test_cpu_kernel_accepts_all_configs():
73
+ """Test that CPUScaledMMLinearKernel accepts all config combinations."""
74
+ configs = [
75
+ ScaledMMLinearLayerConfig(
76
+ is_channelwise=False,
77
+ is_static_input_scheme=True,
78
+ input_symmetric=True,
79
+ ),
80
+ ScaledMMLinearLayerConfig(
81
+ is_channelwise=True,
82
+ is_static_input_scheme=False,
83
+ input_symmetric=False,
84
+ ),
85
+ ]
86
+
87
+ for config in configs:
88
+ can_impl, reason = CPUScaledMMLinearKernel.can_implement(config)
89
+ assert can_impl, (
90
+ f"CPUScaledMMLinearKernel should accept config {config}: {reason}"
91
+ )
vllm/model_executor/layers/quantization/kernels/scaled_mm/ScaledMMLinearKernel.py
+4
-1
@@ -17,7 +17,9 @@ class ScaledMMLinearLayerConfig:
17
class ScaledMMLinearKernel(ABC):
18
@classmethod
19
@abstractmethod
20
- def get_min_capability(cls) -> int:
20
+ def is_supported(
21
+ cls, compute_capability: int | None = None
22
+ ) -> tuple[bool, str | None]:
23
raise NotImplementedError
24
25
@classmethod
@@ -35,6 +37,7 @@ class ScaledMMLinearKernel(ABC):
37
azp_adj_param_name: str,
38
) -> None:
39
assert self.can_implement(c)
40
+ assert self.is_supported()
41
self.config = c
42
self.w_q_name = w_q_param_name
43
self.w_s_name = w_s_param_name
vllm/model_executor/layers/quantization/kernels/scaled_mm/__init__.py
+12
-28
@@ -27,7 +27,7 @@ from vllm.platforms import PlatformEnum, current_platform
27
# in priority/performance order (when available)
28
_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[ScaledMMLinearKernel]]] = {
29
PlatformEnum.CPU: [CPUScaledMMLinearKernel],
30
- PlatformEnum.CUDA: [CutlassScaledMMLinearKernel],
30
+ PlatformEnum.CUDA: [CutlassScaledMMLinearKernel, TritonScaledMMLinearKernel],
31
PlatformEnum.ROCM: [AiterScaledMMLinearKernel, TritonScaledMMLinearKernel],
32
PlatformEnum.TPU: [XLAScaledMMLinearKernel],
33
}
@@ -55,41 +55,25 @@ def choose_scaled_mm_linear_kernel(
55
type[ScaledMMLinearKernel]: Chosen kernel.
56
"""
57
58
- if compute_capability is None:
59
- _cc = current_platform.get_device_capability()
60
- if _cc is not None:
61
- compute_capability = _cc[0] * 10 + _cc[1]
62
-
58
failure_reasons = []
59
for kernel in _POSSIBLE_KERNELS[current_platform._enum]:
60
if kernel.__name__ in os.environ.get("VLLM_DISABLED_KERNELS", "").split(","):
66
- failure_reasons.append(
67
- f" {kernel.__name__} disabled by environment variable"
68
- )
61
+ failure_reasons.append(f"{kernel.__name__}: disabled by env var")
62
continue
63
64
# If the current platform uses compute_capability,
65
# make sure the kernel supports the compute cability.
73
- if compute_capability is not None:
74
- kernel_min_capability = kernel.get_min_capability()
75
- if (
76
- kernel_min_capability is not None
77
- and kernel_min_capability > compute_capability
78
- ):
79
- failure_reasons.append(
80
- f"{kernel.__name__} requires capability "
81
- f"{kernel_min_capability}, current compute capability "
82
- f"is {compute_capability}"
83
- )
84
- continue
66
+ is_supported, reason = kernel.is_supported(compute_capability)
67
+ if not is_supported:
68
+ failure_reasons.append(f"{kernel.__name__}: {reason}")
69
+ continue
70
+
71
+ can_implement, reason = kernel.can_implement(config)
72
+ if not can_implement:
73
+ failure_reasons.append(f"{kernel.__name__}: {reason}")
74
+ continue
75
86
- can_implement, failure_reason = kernel.can_implement(config)
87
- if can_implement:
88
- return kernel
89
- else:
90
- failure_reasons.append(
91
- f" {kernel.__name__} cannot implement due to: {failure_reason}"
92
- )
76
+ return kernel
77
78
raise ValueError(
79
"Failed to find a kernel that can implement the "
vllm/model_executor/layers/quantization/kernels/scaled_mm/aiter.py
+15
-7
@@ -14,17 +14,21 @@ from .ScaledMMLinearKernel import ScaledMMLinearLayerConfig
14
15
class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
16
@classmethod
17
- def get_min_capability(cls) -> int:
18
- return 90
19
-
20
- @classmethod
21
- def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
17
+ def is_supported(
18
+ cls, compute_capability: int | None = None
19
+ ) -> tuple[bool, str | None]:
20
if not current_platform.is_rocm():
21
return (
22
False,
23
"AiterScaledMMLinearKernel requires `aiter` which is not "
24
+ "currently supported on non-ROCm platform.",
25
)
26
+ if compute_capability is None:
27
+ _cc = current_platform.get_device_capability()
28
+ if _cc is not None:
29
+ compute_capability = _cc.major * 10 + _cc.minor
30
+ if compute_capability is not None and compute_capability < 90:
31
+ return False, f"requires capability 90, got {compute_capability}"
32
33
try:
34
import aiter # noqa: F401 # deliberately attempt to import aiter
@@ -34,8 +38,8 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
38
"AiterScaledMMLinearKernel requires `aiter` which is not "
39
+ "installed on ROCm.",
40
)
37
- # Check if rocm_aiter_gemm_w8a8_scaled_mm is enabled
38
- if not (rocm_aiter_ops.is_linear_enabled()):
41
+
42
+ if not rocm_aiter_ops.is_linear_enabled():
43
return (
44
False,
45
"AiterScaledMMLinearKernel is disabled. "
@@ -44,6 +48,10 @@ class AiterScaledMMLinearKernel(CutlassScaledMMLinearKernel):
48
+ "`VLLM_ROCM_USE_AITER_LINEAR` default is True.",
49
)
50
51
+ return True, None
52
+
53
+ @classmethod
54
+ def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
55
if not c.input_symmetric:
56
return (
57
False,
vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
+6
-5
@@ -19,14 +19,15 @@ from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfi
19
20
class CPUScaledMMLinearKernel(ScaledMMLinearKernel):
21
@classmethod
22
- def get_min_capability(cls) -> int:
23
- return 75
22
+ def is_supported(
23
+ cls, compute_capability: int | None = None
24
+ ) -> tuple[bool, str | None]:
25
+ if not current_platform.is_cpu():
26
+ return False, "Requires CPU."
27
+ return True, None
28
29
@classmethod
30
def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
27
- if not current_platform.is_cpu():
28
- return False, "CPUScaledMM requires running on CPU."
29
-
31
return True, None
32
33
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/layers/quantization/kernels/scaled_mm/cutlass.py
+12
-5
@@ -16,14 +16,21 @@ from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfi
16
17
class CutlassScaledMMLinearKernel(ScaledMMLinearKernel):
18
@classmethod
19
- def get_min_capability(cls) -> int:
20
- return 75
19
+ def is_supported(
20
+ cls, compute_capability: int | None = None
21
+ ) -> tuple[bool, str | None]:
22
+ if not current_platform.is_cuda():
23
+ return False, "Requires CUDA."
24
+ if compute_capability is None:
25
+ _cc = current_platform.get_device_capability()
26
+ if _cc is not None:
27
+ compute_capability = _cc.major * 10 + _cc.minor
28
+ if compute_capability is not None and compute_capability < 75:
29
+ return False, f"requires capability 75, got {compute_capability}"
30
+ return True, None
31
32
@classmethod
33
def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
24
- if not current_platform.is_cuda():
25
- return False, "CutlassScaledMM requires running on CUDA."
26
-
34
return True, None
35
36
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
vllm/model_executor/layers/quantization/kernels/scaled_mm/triton.py
+46
-17
@@ -4,34 +4,53 @@
4
5
import torch
6
7
+from vllm import _custom_ops as ops
8
+from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import ( # noqa: E501
9
+ triton_scaled_mm,
10
+)
11
+from vllm.model_executor.layers.quantization.utils import replace_parameter
12
from vllm.platforms import current_platform
13
9
-from .cutlass import CutlassScaledMMLinearKernel
10
-from .ScaledMMLinearKernel import ScaledMMLinearLayerConfig
14
+from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfig
15
16
13
-class TritonScaledMMLinearKernel(CutlassScaledMMLinearKernel):
17
+class TritonScaledMMLinearKernel(ScaledMMLinearKernel):
18
@classmethod
15
- def get_min_capability(cls) -> int:
16
- return 75
19
+ def is_supported(
20
+ cls, compute_capability: int | None = None
21
+ ) -> tuple[bool, str | None]:
22
+ if current_platform.is_cuda_alike():
23
+ return True, None
24
+ return False, "Requires ROCm or CUDA."
25
26
@classmethod
27
def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
20
- if current_platform.is_cpu():
21
- return (
22
- False,
23
- "TritonScaledMMLinearKernel requires Triton which is not "
24
- + "currently supported on CPU.",
25
- )
28
if not c.input_symmetric:
27
- return (
28
- False,
29
- "TritonScaledMMLinearKernel only supports symmetric " + "quantization.",
30
- )
29
+ return False, "Only symmetric input is supported."
30
return True, None
31
32
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
34
- super().process_weights_after_loading(layer)
33
+ weight = getattr(layer, self.w_q_name)
34
+ replace_parameter(
35
+ layer,
36
+ self.w_q_name,
37
+ torch.nn.Parameter(weight.t().data, requires_grad=False),
38
+ )
39
+
40
+ # INPUT SCALE
41
+ if self.config.is_static_input_scheme:
42
+ input_scale = getattr(layer, self.i_s_name)
43
+ replace_parameter(
44
+ layer,
45
+ self.i_s_name,
46
+ torch.nn.Parameter(input_scale.max(), requires_grad=False),
47
+ )
48
+ setattr(layer, self.i_zp_name, None)
49
+ else:
50
+ setattr(layer, self.i_s_name, None)
51
+ setattr(layer, self.i_zp_name, None)
52
+
53
+ setattr(layer, self.azp_adj_name, None)
54
55
def apply_weights(
56
self,
@@ -39,4 +58,14 @@ class TritonScaledMMLinearKernel(CutlassScaledMMLinearKernel):
58
x: torch.Tensor,
59
bias: torch.Tensor | None = None,
60
) -> torch.Tensor:
42
- return super().apply_weights(layer, x, bias)
61
+ w_q, w_s, i_s, i_zp, azp_adj = self._get_weight_params(layer)
62
+
63
+ x_q, x_s, x_zp = ops.scaled_int8_quant(
64
+ x.contiguous(), i_s, i_zp, symmetric=True
65
+ )
66
+
67
+ assert x_zp is None, "Triton kernel only supports symmetric quantization"
68
+
69
+ return triton_scaled_mm(
70
+ x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
71
+ )
vllm/model_executor/layers/quantization/kernels/scaled_mm/xla.py
+6
-5
@@ -17,11 +17,12 @@ from .ScaledMMLinearKernel import ScaledMMLinearKernel, ScaledMMLinearLayerConfi
17
18
class XLAScaledMMLinearKernel(ScaledMMLinearKernel):
19
@classmethod
20
- def get_min_capability(cls) -> int:
21
- raise NotImplementedError(
22
- "TPU platform does have a concept of compute capability, "
23
- "this method should not be called."
24
- )
20
+ def is_supported(
21
+ cls, compute_capability: int | None = None
22
+ ) -> tuple[bool, str | None]:
23
+ if not current_platform.is_tpu():
24
+ return False, "Requires TPU."
25
+ return True, None
26
27
@classmethod
28
def can_implement(cls, c: ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: