@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6af03f2394
Fork 0

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 7, 20:15 UTC

[Refactor] [1/N] Reorganize kernel abstraction directory (#34055)

Signed-off-by: BadrBasowid <badr.basowid@gmail.com> Co-authored-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

BadrBasowid committed Feb 24, 2026 at 14:47 UTC 6af03f2394b54e95440945d2659a8120288b21ea
40 files changed +328 -210
tests/compile/passes/test_fusion.py
+5 -13
@@ -26,24 +26,16 @@ from vllm.config import (
26 PassConfig,
27 VllmConfig,
28 )
29 -from vllm.model_executor.layers.layernorm import RMSNorm
30 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
29 +from vllm.model_executor.kernels.linear import (
30 + ChannelWiseTorchFP8ScaledMMLinearKernel,
31 CutlassFP8ScaledMMLinearKernel,
32 -)
33 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
32 FlashInferFP8ScaledMMLinearKernel,
35 -)
36 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
37 - ChannelWiseTorchFP8ScaledMMLinearKernel,
33 + FP8ScaledMMLinearKernel,
34 PerTensorTorchFP8ScaledMMLinearKernel,
39 - RowWiseTorchFP8ScaledMMLinearKernel,
40 -)
41 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
35 ROCmFP8ScaledMMLinearKernel,
36 + RowWiseTorchFP8ScaledMMLinearKernel,
37 )
44 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
45 - FP8ScaledMMLinearKernel,
46 -)
38 +from vllm.model_executor.layers.layernorm import RMSNorm
39 from vllm.model_executor.layers.quantization.utils.quant_utils import (
40 GroupShape,
41 QuantKey,
tests/compile/passes/test_silu_mul_quant_fusion.py
+3 -11
@@ -26,22 +26,14 @@ from vllm.config import (
26 VllmConfig,
27 set_current_vllm_config,
28 )
29 -from vllm.model_executor.layers.activation import SiluAndMul
30 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
29 +from vllm.model_executor.kernels.linear import (
30 CutlassFP8ScaledMMLinearKernel,
32 -)
33 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
31 FlashInferFP8ScaledMMLinearKernel,
35 -)
36 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
32 + FP8ScaledMMLinearKernel,
33 PerTensorTorchFP8ScaledMMLinearKernel,
38 -)
39 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
34 ROCmFP8ScaledMMLinearKernel,
35 )
42 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
43 - FP8ScaledMMLinearKernel,
44 -)
36 +from vllm.model_executor.layers.activation import SiluAndMul
37 from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
38 from vllm.model_executor.layers.quantization.utils.quant_utils import (
39 GroupShape,
tests/kernels/quantization/test_scaled_mm_kernel_selection.py
+2 -8
@@ -10,16 +10,10 @@ from abc import ABC
10
11 import pytest
12
13 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
14 - Int8ScaledMMLinearLayerConfig,
15 -)
16 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (
13 +from vllm.model_executor.kernels.linear import (
14 AiterInt8ScaledMMLinearKernel,
18 -)
19 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.cpu import (
15 CPUInt8ScaledMMLinearKernel,
21 -)
22 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
16 + Int8ScaledMMLinearLayerConfig,
17 ScaledMMLinearKernel,
18 )
19
tests/utils.py
+2 -4
@@ -42,11 +42,9 @@ from vllm.distributed import (
42 )
43 from vllm.engine.arg_utils import AsyncEngineArgs
44 from vllm.entrypoints.cli.serve import ServeSubcommand
45 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
46 - init_fp8_linear_kernel,
47 -)
48 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
45 +from vllm.model_executor.kernels.linear import (
46 FP8ScaledMMLinearKernel,
47 + init_fp8_linear_kernel,
48 )
49 from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
50 from vllm.model_executor.layers.quantization.utils.quant_utils import (
vllm/model_executor/kernels/__init__.py renamed
vllm/model_executor/kernels/linear/__init__.py renamed
+177 -16
@@ -1,45 +1,89 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +"""
5 +This module re-exports linear kernel implementations to provide a
6 +stable import interface during an ongoing reorganization. Upcoming
7 +PRs will remove the scaled_mm and mixed_precision subdirectories
8 +and reorganize kernels by provider (aiter, cutlass, flashinfer, etc.)
9 +rather than by precision type. By centralizing exports here, we
10 +minimize the need to update imports across other modules when the
11 +internal structure changes. If you are adding a new kernel selector
12 +or kernel implementation, add it to this __init__.py to maintain
13 +import stability.
14 +"""
15 +
16 import os
17 from typing import TypeVar
18
19 import torch
20
21 +import vllm.envs as envs
22 from vllm.logger import init_logger
10 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.aiter import (
23 +from vllm.model_executor.kernels.linear.mixed_precision import (
24 + MPLinearKernel,
25 + MPLinearLayerConfig,
26 +)
27 +from vllm.model_executor.kernels.linear.mixed_precision.allspark import (
28 + AllSparkLinearKernel,
29 +)
30 +from vllm.model_executor.kernels.linear.mixed_precision.conch import (
31 + ConchLinearKernel,
32 +)
33 +from vllm.model_executor.kernels.linear.mixed_precision.cpu import (
34 + CPUWNA16LinearKernel,
35 +)
36 +from vllm.model_executor.kernels.linear.mixed_precision.cutlass import (
37 + CutlassW4A8LinearKernel,
38 +)
39 +from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import (
40 + Dynamic4bitLinearKernel,
41 +)
42 +from vllm.model_executor.kernels.linear.mixed_precision.exllama import (
43 + ExllamaLinearKernel,
44 +)
45 +from vllm.model_executor.kernels.linear.mixed_precision.machete import (
46 + MacheteLinearKernel,
47 +)
48 +from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
49 + MarlinLinearKernel,
50 +)
51 +from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
52 + XPUwNa16LinearKernel,
53 +)
54 +from vllm.model_executor.kernels.linear.scaled_mm import (
55 + FP8ScaledMMLinearKernel,
56 + FP8ScaledMMLinearLayerConfig,
57 + Int8ScaledMMLinearKernel,
58 + Int8ScaledMMLinearLayerConfig,
59 + ScaledMMLinearKernel,
60 + ScaledMMLinearLayerConfig,
61 +)
62 +from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
63 AiterInt8ScaledMMLinearKernel,
64 )
13 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.cpu import (
65 +from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
66 CPUInt8ScaledMMLinearKernel,
67 )
16 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.cutlass import (
68 +from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
69 CutlassFP8ScaledMMLinearKernel,
70 CutlassInt8ScaledMMLinearKernel,
71 )
20 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.flashinfer import (
72 +from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
73 FlashInferFP8ScaledMMLinearKernel,
74 )
23 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.pytorch import (
75 +from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
76 ChannelWiseTorchFP8ScaledMMLinearKernel,
77 PerTensorTorchFP8ScaledMMLinearKernel,
78 RowWiseTorchFP8ScaledMMLinearKernel,
79 )
28 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.rocm import (
80 +from vllm.model_executor.kernels.linear.scaled_mm.rocm import (
81 ROCmFP8ScaledMMLinearKernel,
82 )
31 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
32 - FP8ScaledMMLinearKernel,
33 - FP8ScaledMMLinearLayerConfig,
34 - Int8ScaledMMLinearKernel,
35 - Int8ScaledMMLinearLayerConfig,
36 - ScaledMMLinearKernel,
37 - ScaledMMLinearLayerConfig,
38 -)
39 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.triton import (
83 +from vllm.model_executor.kernels.linear.scaled_mm.triton import (
84 TritonInt8ScaledMMLinearKernel,
85 )
42 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.xpu import (
86 +from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
87 XPUFP8ScaledMMLinearKernel,
88 )
89 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
@@ -80,6 +124,29 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
124 ],
125 }
126
127 +# in priority/performance order (when available)
128 +_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
129 + PlatformEnum.CUDA: [
130 + CutlassW4A8LinearKernel,
131 + MacheteLinearKernel,
132 + AllSparkLinearKernel,
133 + MarlinLinearKernel,
134 + ConchLinearKernel,
135 + ExllamaLinearKernel,
136 + ],
137 + PlatformEnum.ROCM: [
138 + ConchLinearKernel,
139 + ExllamaLinearKernel,
140 + ],
141 + PlatformEnum.XPU: [
142 + XPUwNa16LinearKernel,
143 + ],
144 + PlatformEnum.CPU: [
145 + Dynamic4bitLinearKernel,
146 + CPUWNA16LinearKernel,
147 + ],
148 +}
149 +
150 _KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel)
151 _KernelConfigT = TypeVar("_KernelConfigT", bound=ScaledMMLinearLayerConfig)
152
@@ -234,3 +301,97 @@ def init_int8_linear_kernel(
301 "azp_adj",
302 ],
303 )
304 +
305 +
306 +def choose_mp_linear_kernel(
307 + config: MPLinearLayerConfig, compute_capability: int | None = None
308 +) -> type[MPLinearKernel]:
309 + """
310 + Choose an MPLinearKernel that can implement the given config for the given
311 + compute capability. Attempts to choose the best kernel in terms of
312 + performance.
313 +
314 + Args:
315 + config (MPLinearLayerConfig): Description of the linear layer to be
316 + implemented.
317 + compute_capability (Optional[int], optional): The compute capability of
318 + the target device, if None uses `current_platform` to get
319 + the compute capability. Defaults to None.
320 +
321 + Raises:
322 + ValueError: If no kernel can implement the given config.
323 +
324 + Returns:
325 + type[MPLinearKernel]: Chosen kernel.
326 + """
327 + if compute_capability is None:
328 + if current_platform is None:
329 + raise ValueError("Cannot determine compute capability")
330 + _cc = current_platform.get_device_capability()
331 + if _cc is not None:
332 + compute_capability = _cc[0] * 10 + _cc[1]
333 +
334 + failure_reasons = []
335 + for kernel in _POSSIBLE_KERNELS[current_platform._enum]:
336 + if kernel.__name__ in envs.VLLM_DISABLED_KERNELS:
337 + failure_reasons.append(
338 + f" {kernel.__name__} disabled by environment variable"
339 + )
340 + continue
341 + if (
342 + compute_capability is not None
343 + and kernel.get_min_capability() > compute_capability
344 + ):
345 + failure_reasons.append(
346 + f"{kernel.__name__} requires capability "
347 + f"{kernel.get_min_capability()}, current compute "
348 + f" capability is {compute_capability}"
349 + )
350 + continue
351 +
352 + can_implement, failure_reason = kernel.can_implement(config)
353 + if can_implement:
354 + return kernel
355 + else:
356 + failure_reasons.append(
357 + f" {kernel.__name__} cannot implement due to: {failure_reason}"
358 + )
359 +
360 + raise ValueError(
361 + "Failed to find a kernel that can implement the "
362 + "WNA16 linear layer. Reasons: \n" + "\n".join(failure_reasons)
363 + )
364 +
365 +
366 +__all__ = [
367 + "init_fp8_linear_kernel",
368 + "init_int8_linear_kernel",
369 + "choose_mp_linear_kernel",
370 + "FP8ScaledMMLinearKernel",
371 + "Int8ScaledMMLinearKernel",
372 + "ScaledMMLinearKernel",
373 + "FP8ScaledMMLinearLayerConfig",
374 + "Int8ScaledMMLinearLayerConfig",
375 + "ScaledMMLinearLayerConfig",
376 + "AiterInt8ScaledMMLinearKernel",
377 + "CPUInt8ScaledMMLinearKernel",
378 + "CutlassFP8ScaledMMLinearKernel",
379 + "CutlassInt8ScaledMMLinearKernel",
380 + "FlashInferFP8ScaledMMLinearKernel",
381 + "ChannelWiseTorchFP8ScaledMMLinearKernel",
382 + "PerTensorTorchFP8ScaledMMLinearKernel",
383 + "RowWiseTorchFP8ScaledMMLinearKernel",
384 + "ROCmFP8ScaledMMLinearKernel",
385 + "TritonInt8ScaledMMLinearKernel",
386 + "MPLinearKernel",
387 + "MPLinearLayerConfig",
388 + "AllSparkLinearKernel",
389 + "ConchLinearKernel",
390 + "CPUWNA16LinearKernel",
391 + "CutlassW4A8LinearKernel",
392 + "Dynamic4bitLinearKernel",
393 + "ExllamaLinearKernel",
394 + "MacheteLinearKernel",
395 + "MarlinLinearKernel",
396 + "XPUwNa16LinearKernel",
397 +]
vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py renamed
vllm/model_executor/kernels/linear/mixed_precision/__init__.py new
+48
@@ -0,0 +1,48 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.model_executor.kernels.linear.mixed_precision.allspark import (
5 + AllSparkLinearKernel,
6 +)
7 +from vllm.model_executor.kernels.linear.mixed_precision.conch import (
8 + ConchLinearKernel,
9 +)
10 +from vllm.model_executor.kernels.linear.mixed_precision.cpu import (
11 + CPUWNA16LinearKernel,
12 +)
13 +from vllm.model_executor.kernels.linear.mixed_precision.cutlass import (
14 + CutlassW4A8LinearKernel,
15 +)
16 +from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import (
17 + Dynamic4bitLinearKernel,
18 +)
19 +from vllm.model_executor.kernels.linear.mixed_precision.exllama import (
20 + ExllamaLinearKernel,
21 +)
22 +from vllm.model_executor.kernels.linear.mixed_precision.machete import (
23 + MacheteLinearKernel,
24 +)
25 +from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
26 + MarlinLinearKernel,
27 +)
28 +from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
29 + MPLinearKernel,
30 + MPLinearLayerConfig,
31 +)
32 +from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
33 + XPUwNa16LinearKernel,
34 +)
35 +
36 +__all__ = [
37 + "MPLinearKernel",
38 + "MPLinearLayerConfig",
39 + "AllSparkLinearKernel",
40 + "ConchLinearKernel",
41 + "CPUWNA16LinearKernel",
42 + "CutlassW4A8LinearKernel",
43 + "Dynamic4bitLinearKernel",
44 + "ExllamaLinearKernel",
45 + "MacheteLinearKernel",
46 + "MarlinLinearKernel",
47 + "XPUwNa16LinearKernel",
48 +]
vllm/model_executor/kernels/linear/mixed_precision/allspark.py renamed
vllm/model_executor/kernels/linear/mixed_precision/conch.py renamed
vllm/model_executor/kernels/linear/mixed_precision/cpu.py renamed
vllm/model_executor/kernels/linear/mixed_precision/cutlass.py renamed
vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py renamed
vllm/model_executor/kernels/linear/mixed_precision/exllama.py renamed
vllm/model_executor/kernels/linear/mixed_precision/machete.py renamed
vllm/model_executor/kernels/linear/mixed_precision/marlin.py renamed
vllm/model_executor/kernels/linear/mixed_precision/xpu.py renamed
vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py renamed
vllm/model_executor/kernels/linear/scaled_mm/__init__.py new
+54
@@ -0,0 +1,54 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
5 + AiterInt8ScaledMMLinearKernel,
6 +)
7 +from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
8 + CPUInt8ScaledMMLinearKernel,
9 +)
10 +from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
11 + CutlassFP8ScaledMMLinearKernel,
12 + CutlassInt8ScaledMMLinearKernel,
13 +)
14 +from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
15 + FlashInferFP8ScaledMMLinearKernel,
16 +)
17 +from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
18 + ChannelWiseTorchFP8ScaledMMLinearKernel,
19 + PerTensorTorchFP8ScaledMMLinearKernel,
20 + RowWiseTorchFP8ScaledMMLinearKernel,
21 +)
22 +from vllm.model_executor.kernels.linear.scaled_mm.rocm import (
23 + ROCmFP8ScaledMMLinearKernel,
24 +)
25 +from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
26 + FP8ScaledMMLinearKernel,
27 + FP8ScaledMMLinearLayerConfig,
28 + Int8ScaledMMLinearKernel,
29 + Int8ScaledMMLinearLayerConfig,
30 + ScaledMMLinearKernel,
31 + ScaledMMLinearLayerConfig,
32 +)
33 +from vllm.model_executor.kernels.linear.scaled_mm.triton import (
34 + TritonInt8ScaledMMLinearKernel,
35 +)
36 +
37 +__all__ = [
38 + "FP8ScaledMMLinearKernel",
39 + "FP8ScaledMMLinearLayerConfig",
40 + "Int8ScaledMMLinearKernel",
41 + "Int8ScaledMMLinearLayerConfig",
42 + "ScaledMMLinearKernel",
43 + "ScaledMMLinearLayerConfig",
44 + "AiterInt8ScaledMMLinearKernel",
45 + "CPUInt8ScaledMMLinearKernel",
46 + "CutlassFP8ScaledMMLinearKernel",
47 + "CutlassInt8ScaledMMLinearKernel",
48 + "FlashInferFP8ScaledMMLinearKernel",
49 + "ChannelWiseTorchFP8ScaledMMLinearKernel",
50 + "PerTensorTorchFP8ScaledMMLinearKernel",
51 + "RowWiseTorchFP8ScaledMMLinearKernel",
52 + "ROCmFP8ScaledMMLinearKernel",
53 + "TritonInt8ScaledMMLinearKernel",
54 +]
vllm/model_executor/kernels/linear/scaled_mm/aiter.py renamed
vllm/model_executor/kernels/linear/scaled_mm/cpu.py renamed
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py renamed
vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py renamed
vllm/model_executor/kernels/linear/scaled_mm/pytorch.py renamed
vllm/model_executor/kernels/linear/scaled_mm/rocm.py renamed
vllm/model_executor/kernels/linear/scaled_mm/triton.py renamed
vllm/model_executor/kernels/linear/scaled_mm/xpu.py renamed
+1 -1
@@ -5,7 +5,7 @@ from collections.abc import Sequence
5
6 import torch
7
8 -from vllm.model_executor.layers.quantization.kernels.scaled_mm.ScaledMMLinearKernel import ( # noqa: E501
8 +from vllm.model_executor.kernels.linear import ( # noqa: E501
9 FP8ScaledMMLinearKernel,
10 FP8ScaledMMLinearLayerConfig,
11 )
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_fp8.py
+4 -4
@@ -7,13 +7,13 @@ import torch
7 from compressed_tensors.quantization import ActivationOrdering
8
9 from vllm.logger import init_logger
10 -from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11 - CompressedTensorsScheme,
12 -)
13 -from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
10 +from vllm.model_executor.kernels.linear import (
11 MPLinearLayerConfig,
12 choose_mp_linear_kernel,
13 )
14 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
15 + CompressedTensorsScheme,
16 +)
17 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
18 marlin_repeat_scales_on_all_ranks,
19 )
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_int.py
+4 -4
@@ -6,13 +6,13 @@ from collections.abc import Callable
6 import torch
7
8 from vllm.logger import init_logger
9 -from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
10 - CompressedTensorsScheme,
11 -)
12 -from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
9 +from vllm.model_executor.kernels.linear import (
10 MPLinearLayerConfig,
11 choose_mp_linear_kernel,
12 )
13 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
14 + CompressedTensorsScheme,
15 +)
16 from vllm.model_executor.parameter import (
17 ChannelQuantScaleParameter,
18 GroupQuantScaleParameter,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+3 -3
@@ -9,12 +9,12 @@ from torch.nn import Parameter
9
10 from vllm._aiter_ops import rocm_aiter_ops
11 from vllm.logger import init_logger
12 +from vllm.model_executor.kernels.linear import (
13 + init_fp8_linear_kernel,
14 +)
15 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
16 CompressedTensorsScheme,
17 )
15 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
16 - init_fp8_linear_kernel,
17 -)
18 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
19 W8A8BlockFp8LinearOp,
20 create_fp8_input_scale,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
+3 -3
@@ -7,12 +7,12 @@ import torch
7 from compressed_tensors.quantization import QuantizationStrategy
8
9 from vllm.logger import init_logger
10 +from vllm.model_executor.kernels.linear import (
11 + init_int8_linear_kernel,
12 +)
13 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
14 CompressedTensorsScheme,
15 )
13 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
14 - init_int8_linear_kernel,
15 -)
16 from vllm.model_executor.parameter import (
17 BasevLLMParameter,
18 ChannelQuantScaleParameter,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py
+4 -6
@@ -7,15 +7,13 @@ import torch
7 from compressed_tensors.quantization import ActivationOrdering
8
9 from vllm.logger import init_logger
10 -from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11 - CompressedTensorsScheme,
12 -)
13 -from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
10 +from vllm.model_executor.kernels.linear import (
11 + MarlinLinearKernel,
12 MPLinearLayerConfig,
13 choose_mp_linear_kernel,
14 )
17 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.marlin import (
18 - MarlinLinearKernel,
15 +from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
16 + CompressedTensorsScheme,
17 )
18 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
19 get_marlin_input_dtype,
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+3 -3
@@ -8,6 +8,9 @@ from torch.nn import Module
8 from torch.nn.parameter import Parameter
9
10 from vllm.logger import init_logger
11 +from vllm.model_executor.kernels.linear import (
12 + init_fp8_linear_kernel,
13 +)
14 from vllm.model_executor.layers.linear import (
15 LinearBase,
16 LinearMethodBase,
@@ -18,9 +21,6 @@ from vllm.model_executor.layers.quantization.base_config import (
21 QuantizationConfig,
22 QuantizeMethodBase,
23 )
21 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
22 - init_fp8_linear_kernel,
23 -)
24 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
25 apply_fp8_marlin_linear,
26 prepare_fp8_layer_for_marlin,
vllm/model_executor/layers/quantization/fp8.py
+3 -3
@@ -13,6 +13,9 @@ from vllm import _custom_ops as ops
13 from vllm._aiter_ops import rocm_aiter_ops
14 from vllm.distributed import get_tensor_model_parallel_world_size
15 from vllm.logger import init_logger
16 +from vllm.model_executor.kernels.linear import (
17 + init_fp8_linear_kernel,
18 +)
19 from vllm.model_executor.layers.attention import Attention
20 from vllm.model_executor.layers.batch_invariant import (
21 vllm_is_batch_invariant,
@@ -46,9 +49,6 @@ from vllm.model_executor.layers.quantization.base_config import (
49 QuantizationConfig,
50 QuantizeMethodBase,
51 )
49 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
50 - init_fp8_linear_kernel,
51 -)
52 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
53 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
54 apply_fi_trtllm_fp8_per_tensor_moe,
vllm/model_executor/layers/quantization/gptq_marlin.py
+4 -4
@@ -10,6 +10,10 @@ from safetensors.torch import _TYPES as _SAFETENSORS_TO_TORCH_DTYPE
10 import vllm.model_executor.layers.fused_moe # noqa
11 from vllm import _custom_ops as ops
12 from vllm.logger import init_logger
13 +from vllm.model_executor.kernels.linear import (
14 + MPLinearLayerConfig,
15 + choose_mp_linear_kernel,
16 +)
17 from vllm.model_executor.layers.fused_moe.config import (
18 FusedMoEConfig,
19 FusedMoEQuantConfig,
@@ -27,10 +31,6 @@ from vllm.model_executor.layers.quantization.base_config import (
31 QuantizationConfig,
32 QuantizeMethodBase,
33 )
30 -from vllm.model_executor.layers.quantization.kernels.mixed_precision import (
31 - MPLinearLayerConfig,
32 - choose_mp_linear_kernel,
33 -)
34 from vllm.model_executor.layers.quantization.utils import replace_parameter
35 from vllm.model_executor.layers.quantization.utils.gptq_utils import (
36 get_dynamic_override,
vllm/model_executor/layers/quantization/kernels/mixed_precision/__init__.py deleted
-119
@@ -1,119 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import vllm.envs as envs
5 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.allspark import ( # noqa: E501
6 - AllSparkLinearKernel,
7 -)
8 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.conch import ( # noqa: E501
9 - ConchLinearKernel,
10 -)
11 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.cpu import ( # noqa: E501
12 - CPUWNA16LinearKernel,
13 -)
14 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.cutlass import ( # noqa: E501
15 - CutlassW4A8LinearKernel,
16 -)
17 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.dynamic_4bit import ( # noqa: E501
18 - Dynamic4bitLinearKernel,
19 -)
20 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.exllama import ( # noqa: E501
21 - ExllamaLinearKernel,
22 -)
23 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.machete import ( # noqa: E501
24 - MacheteLinearKernel,
25 -)
26 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.marlin import ( # noqa: E501
27 - MarlinLinearKernel,
28 -)
29 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.MPLinearKernel import ( # noqa: E501
30 - MPLinearKernel,
31 - MPLinearLayerConfig,
32 -)
33 -from vllm.model_executor.layers.quantization.kernels.mixed_precision.xpu import ( # noqa: E501
34 - XPUwNa16LinearKernel,
35 -)
36 -from vllm.platforms import PlatformEnum, current_platform
37 -
38 -# in priority/performance order (when available)
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(
63 - config: MPLinearLayerConfig, compute_capability: int | None = None
64 -) -> type[MPLinearKernel]:
65 - """
66 - Choose an MPLinearKernel that can implement the given config for the given
67 - compute capability. Attempts to choose the best kernel in terms of
68 - performance.
69 -
70 - Args:
71 - config (MPLinearLayerConfig): Description of the linear layer to be
72 - implemented.
73 - compute_capability (Optional[int], optional): The compute capability of
74 - the target device, if None uses `current_platform` to get
75 - the compute capability. Defaults to None.
76 -
77 - Raises:
78 - ValueError: If no kernel can implement the given config.
79 -
80 - Returns:
81 - type[MPLinearKernel]: Chosen kernel.
82 - """
83 - if compute_capability is None:
84 - if current_platform is None:
85 - raise ValueError("Cannot determine compute capability")
86 - _cc = current_platform.get_device_capability()
87 - if _cc is not None:
88 - compute_capability = _cc[0] * 10 + _cc[1]
89 -
90 - failure_reasons = []
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"
95 - )
96 - continue
97 - if (
98 - compute_capability is not None
99 - and kernel.get_min_capability() > compute_capability
100 - ):
101 - failure_reasons.append(
102 - f"{kernel.__name__} requires capability "
103 - f"{kernel.get_min_capability()}, current compute "
104 - f" capability is {compute_capability}"
105 - )
106 - continue
107 -
108 - can_implement, failure_reason = kernel.can_implement(config)
109 - if can_implement:
110 - return kernel
111 - else:
112 - failure_reasons.append(
113 - f" {kernel.__name__} cannot implement due to: {failure_reason}"
114 - )
115 -
116 - raise ValueError(
117 - "Failed to find a kernel that can implement the "
118 - "WNA16 linear layer. Reasons: \n" + "\n".join(failure_reasons)
119 - )
vllm/model_executor/layers/quantization/modelopt.py
+3 -3
@@ -9,6 +9,9 @@ from torch.nn.parameter import Parameter
9
10 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 from vllm.logger import init_logger
12 +from vllm.model_executor.kernels.linear import (
13 + init_fp8_linear_kernel,
14 +)
15 from vllm.model_executor.layers.attention import Attention
16 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
17 from vllm.model_executor.layers.fused_moe.config import (
@@ -45,9 +48,6 @@ from vllm.model_executor.layers.quantization.base_config import (
48 QuantizationConfig,
49 QuantizeMethodBase,
50 )
48 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
49 - init_fp8_linear_kernel,
50 -)
51 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
52 from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
53 flashinfer_trtllm_fp4_moe,
vllm/model_executor/layers/quantization/ptpc_fp8.py
+3 -3
@@ -8,6 +8,9 @@ from torch.nn.parameter import Parameter
8
9 from vllm import _custom_ops as ops
10 from vllm.logger import init_logger
11 +from vllm.model_executor.kernels.linear import (
12 + init_fp8_linear_kernel,
13 +)
14 from vllm.model_executor.layers.attention import Attention
15 from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
16 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -17,9 +20,6 @@ from vllm.model_executor.layers.quantization.fp8 import (
20 Fp8KVCacheMethod,
21 Fp8LinearMethod,
22 )
20 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
21 - init_fp8_linear_kernel,
22 -)
23 from vllm.model_executor.layers.quantization.utils.quant_utils import (
24 is_layer_skipped,
25 kFp8DynamicTokenSym,
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+1 -1
@@ -8,7 +8,7 @@ import torch
8 from torch.nn import Parameter
9
10 from vllm.logger import init_logger
11 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
11 +from vllm.model_executor.kernels.linear import (
12 init_fp8_linear_kernel,
13 )
14 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_int8.py
+1 -1
@@ -6,7 +6,7 @@ from collections.abc import Callable
6 import torch
7
8 from vllm.logger import init_logger
9 -from vllm.model_executor.layers.quantization.kernels.scaled_mm import (
9 +from vllm.model_executor.kernels.linear import (
10 init_int8_linear_kernel,
11 )
12 from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme