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
[Feature] Enable DeepGEMM Linear on B200; 1.5% E2E throughput improvement (#23351)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Aug 22, 2025 at 00:01 UTC
394591e34371e4a1ad23a85401148df31d8de451
2 files changed
+13
-16
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+6
-16
@@ -19,8 +19,9 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
19
CUTLASS_BLOCK_FP8_SUPPORTED)
20
from vllm.platforms import current_platform
21
from vllm.triton_utils import tl, triton
22
-from vllm.utils import cdiv, direct_register_custom_op, has_deep_gemm
23
-from vllm.utils.deep_gemm import is_blackwell_deep_gemm_e8m0_used
22
+from vllm.utils import cdiv, direct_register_custom_op
23
+from vllm.utils.deep_gemm import (is_blackwell_deep_gemm_e8m0_used,
24
+ should_use_deepgemm_for_fp8_linear)
25
26
logger = init_logger(__name__)
27
@@ -108,19 +109,6 @@ def dispatch_w8a8_blockscale_func(
109
return w8a8_block_fp8_matmul
110
111
111
-def should_use_deepgemm(output_dtype: torch.dtype, weight: torch.Tensor):
112
- """
113
- Check if DeepGEMM should be used based on the output dtype and weight shape.
114
- DeepGEMM is only supported for bfloat16 output dtype and weights with shape
115
- divisible by 128.
116
- """
117
-
118
- return (current_platform.is_cuda()
119
- and current_platform.is_device_capability(90) and has_deep_gemm()
120
- and envs.VLLM_USE_DEEP_GEMM and output_dtype == torch.bfloat16
121
- and weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)
122
-
123
-
112
# TODO fix ROCm->Triton custom path:
113
# https://github.com/vllm-project/vllm/issues/14397
114
def apply_w8a8_block_fp8_linear(
@@ -139,7 +127,7 @@ def apply_w8a8_block_fp8_linear(
127
output_shape = [*input.shape[:-1], weight.shape[0]]
128
output_dtype = input.dtype
129
142
- if should_use_deepgemm(output_dtype, weight):
130
+ if should_use_deepgemm_for_fp8_linear(output_dtype, weight):
131
132
input_2d = input.view(-1, input.shape[-1])
133
output_shape = [*input.shape[:-1], weight.shape[0]]
@@ -150,7 +138,9 @@ def apply_w8a8_block_fp8_linear(
138
column_major_scales=True,
139
)
140
141
+ # ensure DeepGEMM-backed custom op is registered before use
142
import vllm.model_executor.layers.quantization.deepgemm # noqa: F401
143
+
144
output = torch.ops.vllm.w8a8_block_fp8_matmul_deepgemm(
145
q_input,
146
weight,
vllm/utils/deep_gemm.py
+7
@@ -202,6 +202,12 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
202
return 1 - sim
203
204
205
+def should_use_deepgemm_for_fp8_linear(output_dtype: torch.dtype,
206
+ weight: torch.Tensor):
207
+ return (is_deep_gemm_supported() and output_dtype == torch.bfloat16
208
+ and weight.shape[0] % 128 == 0 and weight.shape[1] % 128 == 0)
209
+
210
+
211
__all__ = [
212
"calc_diff",
213
"fp8_gemm_nt",
@@ -210,4 +216,5 @@ __all__ = [
216
"per_block_cast_to_fp8",
217
"is_blackwell_deep_gemm_e8m0_used",
218
"is_deep_gemm_supported",
219
+ "should_use_deepgemm_for_fp8_linear",
220
]