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 15, 06:00 UTC
[Kernel] Support deep_gemm for linear methods (#19085)
Signed-off-by: artetaout <lulala341@gmail.com>
artetaout committed
Jun 11, 2025 at 15:14 UTC
b8e809a057765c574726a6077fd124db5077ce1f
3 files changed
+124
-1
vllm/model_executor/layers/quantization/deepgemm.py
new
+84
@@ -0,0 +1,84 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+import importlib.util
3
+import logging
4
+
5
+import torch
6
+
7
+from vllm.platforms import current_platform
8
+from vllm.triton_utils import triton
9
+from vllm.utils import direct_register_custom_op
10
+
11
+has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
12
+if has_deep_gemm:
13
+ import deep_gemm
14
+
15
+logger = logging.getLogger(__name__)
16
+
17
+
18
+def prepare_block_fp8_matmul_inputs(
19
+ A: torch.Tensor,
20
+ B: torch.Tensor,
21
+ As: torch.Tensor,
22
+ Bs: torch.Tensor,
23
+ block_size: list[int],
24
+ output_dtype: torch.dtype = torch.float16,
25
+) -> tuple[int, int, int, torch.Tensor]:
26
+ assert len(block_size) == 2
27
+ block_n, block_k = block_size[0], block_size[1]
28
+
29
+ assert A.shape[-1] == B.shape[-1]
30
+ assert A.shape[:-1] == As.shape[:-1]
31
+ assert A.is_contiguous()
32
+ assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
33
+
34
+ M = A.numel() // A.shape[-1]
35
+
36
+ assert B.ndim == 2
37
+ assert B.is_contiguous()
38
+ assert Bs.ndim == 2
39
+ N, K = B.shape
40
+ assert triton.cdiv(N, block_n) == Bs.shape[0]
41
+ assert triton.cdiv(K, block_k) == Bs.shape[1]
42
+
43
+ C_shape = A.shape[:-1] + (N, )
44
+ C = A.new_empty(C_shape, dtype=output_dtype)
45
+
46
+ return M, N, K, C
47
+
48
+
49
+def w8a8_block_fp8_matmul_deepgemm(
50
+ A: torch.Tensor,
51
+ B: torch.Tensor,
52
+ As: torch.Tensor,
53
+ Bs: torch.Tensor,
54
+ block_size: list[int],
55
+ output_dtype: torch.dtype,
56
+) -> torch.Tensor:
57
+ M, N, K, C = prepare_block_fp8_matmul_inputs(A, B, As, Bs, block_size,
58
+ output_dtype)
59
+ # Deepgemm only supports output tensor type as bfloat16
60
+ assert C.dtype == torch.bfloat16
61
+ deep_gemm.gemm_fp8_fp8_bf16_nt((A, As), (B, Bs), C)
62
+ return C
63
+
64
+
65
+def w8a8_block_fp8_matmul_deepgemm_fake(
66
+ A: torch.Tensor,
67
+ B: torch.Tensor,
68
+ As: torch.Tensor,
69
+ Bs: torch.Tensor,
70
+ block_size: list[int],
71
+ output_dtype: torch.dtype,
72
+) -> torch.Tensor:
73
+ M, N, K, C = prepare_block_fp8_matmul_inputs(A, B, As, Bs, block_size,
74
+ output_dtype)
75
+ return C
76
+
77
+
78
+direct_register_custom_op(
79
+ op_name="w8a8_block_fp8_matmul_deepgemm",
80
+ op_func=w8a8_block_fp8_matmul_deepgemm,
81
+ mutates_args=[],
82
+ fake_impl=w8a8_block_fp8_matmul_deepgemm_fake,
83
+ dispatch_key=current_platform.dispatch_key,
84
+)
vllm/model_executor/layers/quantization/fp8.py
+1
@@ -402,6 +402,7 @@ class Fp8LinearMethod(LinearMethodBase):
402
403
if self.block_quant:
404
assert self.quant_config.weight_block_size is not None
405
+
406
return torch.ops.vllm.apply_w8a8_block_fp8_linear(
407
input=x,
408
weight=layer.weight,
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+39
-1
@@ -3,12 +3,14 @@
3
4
# Adapted from https://github.com/sgl-project/sglang/pull/2575
5
import functools
6
+import importlib.util
7
import json
8
import os
9
from typing import Any, Callable, Optional, Union
10
11
import torch
12
13
+import vllm.envs as envs
14
from vllm import _custom_ops as ops
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -20,6 +22,7 @@ from vllm.triton_utils import tl, triton
22
from vllm.utils import direct_register_custom_op
23
24
logger = init_logger(__name__)
25
+has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
26
27
28
def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:
@@ -98,6 +101,19 @@ def dispatch_w8a8_blockscale_func(
101
return w8a8_block_fp8_matmul
102
103
104
+def should_use_deepgemm(output_dtype: torch.dtype, weight: torch.Tensor):
105
+ """
106
+ Check if DeepGEMM should be used based on the output dtype and weight shape.
107
+ DeepGEMM is only supported for bfloat16 output dtype and weights with shape
108
+ divisible by 128.
109
+ """
110
+
111
+ return (current_platform.is_cuda()
112
+ and current_platform.is_device_capability(90) and has_deep_gemm
113
+ and envs.VLLM_USE_DEEP_GEMM and output_dtype == torch.bfloat16
114
+ and weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)
115
+
116
+
117
# TODO fix ROCm->Triton custom path:
118
# https://github.com/vllm-project/vllm/issues/14397
119
def apply_w8a8_block_fp8_linear(
@@ -114,6 +130,29 @@ def apply_w8a8_block_fp8_linear(
130
# View input as 2D matrix for fp8 methods
131
input_2d = input.view(-1, input.shape[-1])
132
output_shape = [*input.shape[:-1], weight.shape[0]]
133
+ output_dtype = input.dtype
134
+
135
+ if should_use_deepgemm(output_dtype, weight):
136
+
137
+ input_2d = input.view(-1, input.shape[-1])
138
+ output_shape = [*input.shape[:-1], weight.shape[0]]
139
+
140
+ q_input, x_scale = per_token_group_quant_fp8(
141
+ input_2d,
142
+ block_size[1],
143
+ column_major_scales=True,
144
+ )
145
+
146
+ output = torch.ops.vllm.w8a8_block_fp8_matmul_deepgemm(
147
+ q_input,
148
+ weight,
149
+ x_scale,
150
+ weight_scale,
151
+ block_size,
152
+ output_dtype=output_dtype)
153
+ if bias is not None:
154
+ output += bias
155
+ return output.to(dtype=output_dtype).view(*output_shape)
156
157
if current_platform.is_cuda():
158
if current_platform.has_device_capability(100):
@@ -134,7 +173,6 @@ def apply_w8a8_block_fp8_linear(
173
174
w8a8_blockscale_func = dispatch_w8a8_blockscale_func(
175
use_cutlass, use_aiter_and_is_supported)
137
-
176
if use_cutlass:
177
q_input, x_scale = per_token_group_quant_fp8(
178
input_2d, block_size[1], column_major_scales=use_cutlass)