[Bugfix] Fix DeepGEMM warmup when using `FlashInferFp8DeepGEMMDynamicBlockScaledKernel` (#49467)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 22, 2026 at 19:28 UTC
917fdb5bf7d95c6e2ec321924e7d4620759e9674
1 file changed
+14
-4
vllm/model_executor/warmup/deep_gemm_warmup.py
+14
-4
@@ -133,6 +133,18 @@ def _extract_data_from_fused_moe_module(
133
return w13, w13_s, w2, w2_s, num_topk
134
135
136
+def _is_deep_gemm_backed_kernel(fp8_linear: object) -> bool:
137
+ """
138
+ Return True if the selected linear kernel dispatches to DeepGEMM, either
139
+ directly or as the fallback branch of a dynamic wrapper.
140
+ """
141
+ if isinstance(fp8_linear, DeepGemmFp8BlockScaledMMKernel):
142
+ return True
143
+ return isinstance(
144
+ getattr(fp8_linear, "fallback", None), DeepGemmFp8BlockScaledMMKernel
145
+ )
146
+
147
+
148
def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
149
"""
150
Return True if the input module/layer could be processed with DeepGEMM.
@@ -147,10 +159,8 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
159
):
160
return False
161
150
- if not isinstance(
151
- getattr(module.quant_method, "fp8_linear", None),
152
- DeepGemmFp8BlockScaledMMKernel,
153
- ):
162
+ fp8_linear = getattr(module.quant_method, "fp8_linear", None)
163
+ if not _is_deep_gemm_backed_kernel(fp8_linear):
164
return False
165
166
block_size = get_mk_alignment_for_contiguous_layout()[0]