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, 01:00 UTC
[Bug] DeepGemm: Fix Cuda Init Error (#21312)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 22, 2025 at 02:36 UTC
e69a92a1cea23b36803caac2d251d906789eed1d
1 file changed
+32
-22
vllm/utils/deep_gemm.py
+32
-22
@@ -45,30 +45,36 @@ def _resolve_symbol(module, new: str, old: str) -> Callable[..., Any] | None:
45
return None
46
47
48
-if not has_deep_gemm():
49
- _fp8_gemm_nt_impl: Callable[..., Any] | None = None
50
- _grouped_impl: Callable[..., Any] | None = None
51
- _grouped_masked_impl: Callable[..., Any] | None = None
52
- _per_block_cast_impl: Callable[..., Any] | None = None
53
-else:
54
- _dg = importlib.import_module("deep_gemm") # type: ignore
55
-
56
- _fp8_gemm_nt_impl = _resolve_symbol(
57
- _dg,
58
- "fp8_gemm_nt",
59
- "gemm_fp8_fp8_bf16_nt",
60
- )
48
+_fp8_gemm_nt_impl: Callable[..., Any] | None = None
49
+_grouped_impl: Callable[..., Any] | None = None
50
+_grouped_masked_impl: Callable[..., Any] | None = None
51
+_per_block_cast_impl: Callable[..., Any] | None = None
52
+
53
+
54
+def _lazy_init() -> None:
55
+ """Import deep_gemm and resolve symbols on first use."""
56
+ global _fp8_gemm_nt_impl, _grouped_impl, _grouped_masked_impl, \
57
+ _per_block_cast_impl
58
+
59
+ # fast path
60
+ if (_fp8_gemm_nt_impl is not None or _grouped_impl is not None
61
+ or _grouped_masked_impl is not None
62
+ or _per_block_cast_impl is not None):
63
+ return
64
+
65
+ if not has_deep_gemm():
66
+ return
67
+
68
+ _dg = importlib.import_module("deep_gemm")
69
+
70
+ _fp8_gemm_nt_impl = _resolve_symbol(_dg, "fp8_gemm_nt",
71
+ "gemm_fp8_fp8_bf16_nt")
72
_grouped_impl = _resolve_symbol(
62
- _dg,
63
- "m_grouped_fp8_gemm_nt_contiguous",
64
- "m_grouped_gemm_fp8_fp8_bf16_nt_contiguous",
65
- )
73
+ _dg, "m_grouped_fp8_gemm_nt_contiguous",
74
+ "m_grouped_gemm_fp8_fp8_bf16_nt_contiguous")
75
_grouped_masked_impl = _resolve_symbol(
67
- _dg,
68
- "fp8_m_grouped_gemm_nt_masked",
69
- "m_grouped_gemm_fp8_fp8_bf16_nt_masked",
70
- )
71
-
76
+ _dg, "fp8_m_grouped_gemm_nt_masked",
77
+ "m_grouped_gemm_fp8_fp8_bf16_nt_masked")
78
# Try to get per_token_cast_to_fp8 from DeepGEMM math utils.
79
try:
80
_math_mod = importlib.import_module(
@@ -80,24 +86,28 @@ else:
86
87
88
def fp8_gemm_nt(*args, **kwargs):
89
+ _lazy_init()
90
if _fp8_gemm_nt_impl is None:
91
return _missing(*args, **kwargs)
92
return _fp8_gemm_nt_impl(*args, **kwargs)
93
94
95
def m_grouped_fp8_gemm_nt_contiguous(*args, **kwargs):
96
+ _lazy_init()
97
if _grouped_impl is None:
98
return _missing(*args, **kwargs)
99
return _grouped_impl(*args, **kwargs)
100
101
102
def fp8_m_grouped_gemm_nt_masked(*args, **kwargs):
103
+ _lazy_init()
104
if _grouped_masked_impl is None:
105
return _missing(*args, **kwargs)
106
return _grouped_masked_impl(*args, **kwargs)
107
108
109
def per_block_cast_to_fp8(x, *args, **kwargs):
110
+ _lazy_init()
111
if _per_block_cast_impl is not None and is_blackwell_deep_gemm_used():
112
return _per_block_cast_impl(x, use_ue8m0=True)
113
# TODO: refactor the `per_block_cast_to_fp8` from tests to vllm utils