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
[Refactor] Remove DeepGEMM OP Register (#25710)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 25, 2025 at 20:13 UTC
9fe4c2bdb9859c14ad7f7479e1db7e01083bada3
2 files changed
+5
-90
vllm/model_executor/layers/quantization/deepgemm.py
deleted
-78
@@ -1,78 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import logging
4
-
5
-import torch
6
-
7
-from vllm.triton_utils import triton
8
-from vllm.utils import direct_register_custom_op
9
-from vllm.utils.deep_gemm import fp8_gemm_nt
10
-
11
-logger = logging.getLogger(__name__)
12
-
13
-
14
-def prepare_block_fp8_matmul_inputs(
15
- A: torch.Tensor,
16
- B: torch.Tensor,
17
- As: torch.Tensor,
18
- Bs: torch.Tensor,
19
- block_size: list[int],
20
- output_dtype: torch.dtype = torch.float16,
21
-) -> tuple[int, int, int, torch.Tensor]:
22
- assert len(block_size) == 2
23
- block_n, block_k = block_size[0], block_size[1]
24
-
25
- assert A.shape[-1] == B.shape[-1]
26
- assert A.shape[:-1] == As.shape[:-1]
27
- assert A.is_contiguous()
28
- assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
29
-
30
- M = A.numel() // A.shape[-1]
31
-
32
- assert B.ndim == 2
33
- assert B.is_contiguous()
34
- assert Bs.ndim == 2
35
- N, K = B.shape
36
- assert triton.cdiv(N, block_n) == Bs.shape[0]
37
- assert triton.cdiv(K, block_k) == Bs.shape[1]
38
-
39
- C_shape = A.shape[:-1] + (N, )
40
- C = A.new_empty(C_shape, dtype=output_dtype)
41
-
42
- return M, N, K, C
43
-
44
-
45
-def w8a8_block_fp8_matmul_deepgemm(
46
- A: torch.Tensor,
47
- B: torch.Tensor,
48
- As: torch.Tensor,
49
- Bs: torch.Tensor,
50
- block_size: list[int],
51
- output_dtype: torch.dtype,
52
-) -> torch.Tensor:
53
- M, N, K, C = prepare_block_fp8_matmul_inputs(A, B, As, Bs, block_size,
54
- output_dtype)
55
- # Deepgemm only supports output tensor type as bfloat16
56
- assert C.dtype == torch.bfloat16
57
- fp8_gemm_nt((A, As), (B, Bs), C)
58
- return C
59
-
60
-
61
-def w8a8_block_fp8_matmul_deepgemm_fake(
62
- A: torch.Tensor,
63
- B: torch.Tensor,
64
- As: torch.Tensor,
65
- Bs: torch.Tensor,
66
- block_size: list[int],
67
- output_dtype: torch.dtype,
68
-) -> torch.Tensor:
69
- M, N, K, C = prepare_block_fp8_matmul_inputs(A, B, As, Bs, block_size,
70
- output_dtype)
71
- return C
72
-
73
-
74
-direct_register_custom_op(
75
- op_name="w8a8_block_fp8_matmul_deepgemm",
76
- op_func=w8a8_block_fp8_matmul_deepgemm,
77
- fake_impl=w8a8_block_fp8_matmul_deepgemm_fake,
78
-)
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+5
-12
@@ -23,7 +23,7 @@ from vllm.model_executor.parameter import (BlockQuantScaleParameter,
23
from vllm.platforms import current_platform
24
from vllm.triton_utils import tl, triton
25
from vllm.utils import direct_register_custom_op
26
-from vllm.utils.deep_gemm import (is_deep_gemm_e8m0_used,
26
+from vllm.utils.deep_gemm import (fp8_gemm_nt, is_deep_gemm_e8m0_used,
27
should_use_deepgemm_for_fp8_linear)
28
29
logger = init_logger(__name__)
@@ -141,17 +141,10 @@ def apply_w8a8_block_fp8_linear(
141
block_size[1],
142
column_major_scales=True,
143
)
144
-
145
- # ensure DeepGEMM-backed custom op is registered before use
146
- import vllm.model_executor.layers.quantization.deepgemm # noqa: F401
147
-
148
- output = torch.ops.vllm.w8a8_block_fp8_matmul_deepgemm(
149
- q_input,
150
- weight,
151
- x_scale,
152
- weight_scale,
153
- block_size,
154
- output_dtype=output_dtype)
144
+ output = torch.empty((q_input.shape[0], weight.shape[0]),
145
+ dtype=torch.bfloat16,
146
+ device=q_input.device)
147
+ fp8_gemm_nt((q_input, x_scale), (weight, weight_scale), output)
148
if bias is not None:
149
output += bias
150
return output.to(dtype=output_dtype).view(*output_shape)