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, 03:45 UTC
[Refactor] Refactor for `DeepGemmQuantScaleFMT` using cache (#30898)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Dec 19, 2025 at 15:50 UTC
3bd8335bd015d695819e2313d09c6d449ef8248e
2 files changed
+29
-9
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+2
-2
@@ -31,6 +31,7 @@ from vllm.model_executor.utils import replace_parameter
31
from vllm.platforms import current_platform
32
from vllm.triton_utils import tl, triton
33
from vllm.utils.deep_gemm import (
34
+ DeepGemmQuantScaleFMT,
35
fp8_gemm_nt,
36
is_deep_gemm_e8m0_used,
37
is_deep_gemm_supported,
@@ -247,7 +248,6 @@ class W8A8BlockFp8LinearOp:
248
self.act_quant_group_shape = act_quant_group_shape
249
self.is_deep_gemm_supported = is_deep_gemm_supported()
250
self.is_hopper = current_platform.is_device_capability(90)
250
- self.is_blackwell = current_platform.is_device_capability_family(100)
251
self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()
252
253
# Get the correct blockscale mul and input quant operations.
@@ -303,7 +303,7 @@ class W8A8BlockFp8LinearOp:
303
weight: torch.Tensor,
304
weight_scale: torch.Tensor,
305
) -> torch.Tensor:
306
- if self.use_deep_gemm_e8m0 and self.is_blackwell:
306
+ if DeepGemmQuantScaleFMT.from_oracle() == DeepGemmQuantScaleFMT.UE8M0:
307
q_input, input_scale = per_token_group_quant_fp8_packed_for_deepgemm(
308
input_2d,
309
group_size=self.act_quant_group_shape.col,
vllm/utils/deep_gemm.py
+27
-7
@@ -32,16 +32,35 @@ class DeepGemmQuantScaleFMT(Enum):
32
# element contains 4 scale values.
33
UE8M0 = 2
34
35
- @staticmethod
36
- def from_oracle() -> "DeepGemmQuantScaleFMT":
37
- if not is_deep_gemm_e8m0_used():
38
- return DeepGemmQuantScaleFMT.FLOAT32
39
- return (
40
- DeepGemmQuantScaleFMT.UE8M0
35
+ @classmethod
36
+ def init_oracle_cache(cls) -> None:
37
+ """Initialize the oracle decision and store it in the class cache"""
38
+ cached = getattr(cls, "_oracle_cache", None)
39
+ if cached is not None:
40
+ return
41
+
42
+ use_e8m0 = (
43
+ envs.VLLM_USE_DEEP_GEMM_E8M0
44
+ and is_deep_gemm_supported()
45
+ and (_fp8_gemm_nt_impl is not None)
46
+ )
47
+ if not use_e8m0:
48
+ cls._oracle_cache = cls.FLOAT32 # type: ignore
49
+ return
50
+
51
+ cls._oracle_cache = ( # type: ignore
52
+ cls.UE8M0
53
if current_platform.is_device_capability_family(100)
42
- else DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0
54
+ else cls.FLOAT32_CEIL_UE8M0
55
)
56
57
+ @classmethod
58
+ def from_oracle(cls) -> "DeepGemmQuantScaleFMT":
59
+ """Return the pre-initialized oracle decision"""
60
+ cached = getattr(cls, "_oracle_cache", None)
61
+ assert cached is not None, "DeepGemmQuantScaleFMT oracle cache not initialized"
62
+ return cached
63
+
64
65
@functools.cache
66
def is_deep_gemm_supported() -> bool:
@@ -149,6 +168,7 @@ def _lazy_init() -> None:
168
_transform_sf_into_required_layout_impl = getattr(
169
_dg, "transform_sf_into_required_layout", None
170
)
171
+ DeepGemmQuantScaleFMT.init_oracle_cache()
172
173
174
def get_num_sms() -> int: