[Bugfix] Fix DeepGEMM FP8 warmup coverage (#56452)
Signed-off-by: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Ayushman Singh committed
Sep 12, 2026 at 14:20 UTC
7bcca16729b71294bd1b67ecde9aca584618234d
4 files changed
+140
-99
tests/model_executor/test_deep_gemm_warmup.py
new
+90
@@ -0,0 +1,90 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import (
10
+ DeepGemmFp8BlockScaledMMKernel,
11
+)
12
+from vllm.model_executor.warmup import deep_gemm_warmup
13
+
14
+
15
+def _block_fp8_layer(n: int, k: int, scale_name: str = "weight_scale"):
16
+ layer = SimpleNamespace(
17
+ weight=torch.empty((n, k), dtype=torch.float8_e4m3fn),
18
+ weight_block_size=[128, 128],
19
+ deep_gemm_warmup_provider=object.__new__(DeepGemmFp8BlockScaledMMKernel),
20
+ )
21
+ setattr(layer, scale_name, torch.empty((n // 128, k // 128), dtype=torch.float32))
22
+ return layer
23
+
24
+
25
+def _run_warmup(monkeypatch, layers) -> list[dict]:
26
+ calls: list[dict] = []
27
+ monkeypatch.setattr(
28
+ deep_gemm_warmup,
29
+ "get_mk_alignment_for_contiguous_layout",
30
+ lambda: [128, 128],
31
+ )
32
+ monkeypatch.setattr(
33
+ deep_gemm_warmup,
34
+ "_deepgemm_fp8_gemm_nt_warmup",
35
+ lambda **kwargs: calls.append(kwargs),
36
+ )
37
+ model = SimpleNamespace(modules=lambda: iter(layers))
38
+ deep_gemm_warmup.deepgemm_fp8_gemm_nt_warmup(model, max_tokens=16)
39
+ return calls
40
+
41
+
42
+@pytest.mark.parametrize("scale_name", ["weight_scale", "weight_scale_inv"])
43
+def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> None:
44
+ """Any layer stamped by the DeepGEMM kernel is warmed, regardless of the
45
+ quantization method that owns it or the name of its scale parameter."""
46
+ layer = _block_fp8_layer(256, 128, scale_name)
47
+
48
+ calls = _run_warmup(monkeypatch, [layer])
49
+
50
+ assert len(calls) == 1
51
+ assert calls[0]["w"] is layer.weight
52
+ assert calls[0]["ws"] is getattr(layer, scale_name)
53
+
54
+
55
+def test_n_multiple_of_64_matches_kernel_selection(monkeypatch) -> None:
56
+ """DeepGEMM accepts N % 64 == 0 (e.g. DeepSeek kv_a_proj_with_mqa, N=576),
57
+ so warmup must not require N % 128 == 0."""
58
+ assert len(_run_warmup(monkeypatch, [_block_fp8_layer(576, 256)])) == 1
59
+
60
+
61
+def test_unstamped_and_mismatched_block_layers_are_skipped(monkeypatch) -> None:
62
+ mismatched = _block_fp8_layer(256, 128)
63
+ mismatched.weight_block_size = [1, 32]
64
+ layers = [SimpleNamespace(), mismatched, _block_fp8_layer(256, 128)]
65
+
66
+ calls = _run_warmup(monkeypatch, layers)
67
+
68
+ assert len(calls) == 1
69
+ assert calls[0]["w"] is layers[-1].weight
70
+
71
+
72
+@pytest.mark.parametrize("is_bmm", [False, True])
73
+def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None:
74
+ """bmm layers bypass the kernel at runtime, so they must not be warmed."""
75
+ kernel = object.__new__(DeepGemmFp8BlockScaledMMKernel)
76
+ kernel.is_deep_gemm_supported = False
77
+ layer = torch.nn.Module()
78
+ layer.weight = torch.nn.Parameter(
79
+ torch.empty((256, 128), dtype=torch.float8_e4m3fn), requires_grad=False
80
+ )
81
+ layer.weight_scale = torch.nn.Parameter(
82
+ torch.empty((2, 1), dtype=torch.float32), requires_grad=False
83
+ )
84
+ layer.weight_block_size = [128, 128]
85
+ layer.is_bmm = is_bmm
86
+
87
+ kernel.process_weights_after_loading(layer)
88
+
89
+ provider = getattr(layer, "deep_gemm_warmup_provider", None)
90
+ assert provider is (None if is_bmm else kernel)
vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py
+19
-19
@@ -38,6 +38,22 @@ class FP8BlockParams(FP8Params):
38
input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
39
)
40
41
+ @property
42
+ def block_scale_attr(self) -> str:
43
+ """Fp8LinearMethod registers the block scale as ``weight_scale_inv``,
44
+ compressed-tensors as ``weight_scale``."""
45
+ return (
46
+ self.WEIGHT_SCALE
47
+ if self.weight_scale_inv is None
48
+ else self.WEIGHT_SCALE_INV
49
+ )
50
+
51
+ @property
52
+ def block_scale(self) -> torch.Tensor:
53
+ scale = getattr(self, self.block_scale_attr)
54
+ assert scale is not None
55
+ return scale
56
+
57
58
class Fp8BlockScaledMMLinearKernel(
59
MMLinearKernel[FP8ScaledMMLinearLayerConfig, FP8BlockParams], ABC
@@ -74,25 +90,13 @@ class Fp8BlockScaledMMLinearKernel(
90
91
def process_weights_after_loading(self, layer: torch.nn.Module):
92
params = self._get_layer_params(layer)
77
- # Fp8LinearMethod registered weight scale
78
- # buffer as weight_scale_inv unlike compressed tensors.
79
- weight_scale = (
80
- params.weight_scale
81
- if params.weight_scale_inv is None
82
- else params.weight_scale_inv
83
- )
84
- scale_attr_name = (
85
- params.WEIGHT_SCALE
86
- if params.weight_scale_inv is None
87
- else params.WEIGHT_SCALE_INV
88
- )
93
new_weight, new_weight_scale = process_fp8_weight_block_strategy(
94
params.weight,
91
- weight_scale,
95
+ params.block_scale,
96
)
97
98
replace_parameter(layer, params.WEIGHT, new_weight.data)
95
- replace_parameter(layer, scale_attr_name, new_weight_scale.data)
99
+ replace_parameter(layer, params.block_scale_attr, new_weight_scale.data)
100
101
def apply_weights(
102
self,
@@ -104,11 +108,7 @@ class Fp8BlockScaledMMLinearKernel(
108
out_dtype = self.config.out_dtype
109
params = self._get_layer_params(layer)
110
weight = params.weight
107
- weight_scale = (
108
- params.weight_scale
109
- if params.weight_scale_inv is None
110
- else params.weight_scale_inv
111
- )
111
+ weight_scale = params.block_scale
112
input_scale = params.input_scale
113
scale_up = params.input_scale_ub
114
vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py
+14
-11
@@ -86,25 +86,28 @@ class DeepGemmFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
86
params = self._get_layer_params(layer)
87
assert layer.weight_block_size is not None
88
89
+ is_bmm = getattr(layer, "is_bmm", False)
90
+
91
if self.is_deep_gemm_supported:
90
- weight_scale_invs = params.weight_scale_inv
91
- scale_attr = (
92
- params.WEIGHT_SCALE_INV
93
- if weight_scale_invs is not None
94
- else params.WEIGHT_SCALE
95
- )
92
dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(
93
wq=params.weight,
98
- ws=weight_scale_invs
99
- if weight_scale_invs is not None
100
- else params.weight_scale,
94
+ ws=params.block_scale,
95
quant_block_shape=tuple(layer.weight_block_size),
96
use_e8m0=self.use_deep_gemm_e8m0,
103
- is_bmm=getattr(layer, "is_bmm", False),
97
+ is_bmm=is_bmm,
98
bmm_batch_size=getattr(layer, "bmm_batch_size", 0),
99
)
100
replace_parameter(layer, params.WEIGHT, dg_weight)
107
- replace_parameter(layer, scale_attr, dg_weight_scale)
101
+ replace_parameter(layer, params.block_scale_attr, dg_weight_scale)
102
+ # bmm layers bypass this kernel and run through fp8_einsum.
103
+ if not is_bmm:
104
+ layer.deep_gemm_warmup_provider = self
105
+
106
+ def get_deep_gemm_warmup_weights(
107
+ self, layer: torch.nn.Module
108
+ ) -> tuple[torch.Tensor, torch.Tensor]:
109
+ params = self._get_layer_params(layer)
110
+ return params.weight, params.block_scale
111
112
def apply_block_scaled_mm(
113
self,
vllm/model_executor/warmup/deep_gemm_warmup.py
+17
-69
@@ -22,9 +22,6 @@ from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import DeepGemmE
22
from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import (
23
TritonOrDeepGemmExperts,
24
)
25
-from vllm.model_executor.layers.linear import LinearBase
26
-from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
27
-from vllm.model_executor.layers.quantization.online.mxfp8 import Mxfp8OnlineLinearMethod
25
from vllm.tracing import instrument
26
from vllm.utils.deep_gemm import (
27
fp8_gemm_nt,
@@ -82,28 +79,6 @@ def _generate_optimal_warmup_m_values(
79
return sorted(m_values)
80
81
85
-def _extract_data_from_linear_base_module(
86
- m: torch.nn.Module,
87
-) -> tuple[torch.Tensor, torch.Tensor, list[int]]:
88
- """
89
- Extract weights, weight scales and quantization block sizes from the given
90
- LinearBase module.
91
- """
92
- assert isinstance(m, LinearBase)
93
- assert isinstance(m.quant_method, Fp8LinearMethod)
94
- assert m.quant_method.block_quant
95
- assert m.quant_method.quant_config is not None
96
-
97
- w = m.weight
98
- ws = m.weight_scale_inv if hasattr(m, "weight_scale_inv") else m.weight_scale
99
- quant_block_size = m.quant_method.quant_config.weight_block_size
100
-
101
- assert isinstance(w, torch.Tensor)
102
- assert isinstance(ws, torch.Tensor)
103
- assert quant_block_size is not None
104
- return (w, ws, quant_block_size)
105
-
106
-
82
def _extract_data_from_fused_moe_module(
83
m_: torch.nn.Module,
84
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]:
@@ -133,45 +108,19 @@ def _extract_data_from_fused_moe_module(
108
return w13, w13_s, w2, w2_s, num_topk
109
110
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:
111
+def _deep_gemm_linear_data(
112
+ module: torch.nn.Module,
113
+) -> tuple[torch.Tensor, torch.Tensor] | None:
114
"""
150
- Return True if the input module/layer could be processed with DeepGEMM.
115
+ Return (weight, weight_scale) for a layer whose kernel stamped itself as
116
+ ``deep_gemm_warmup_provider`` in ``process_weights_after_loading``, else None.
117
"""
152
-
153
- if not (
154
- isinstance(module, LinearBase)
155
- and isinstance(module.quant_method, Fp8LinearMethod)
156
- and not isinstance(module.quant_method, Mxfp8OnlineLinearMethod)
157
- and getattr(module.quant_method, "block_quant", False)
158
- and not getattr(module.quant_method, "use_marlin", True)
159
- ):
160
- return False
161
-
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]
167
-
168
- w, _, block_sizes = _extract_data_from_linear_base_module(module)
169
- return (
170
- block_sizes == get_mk_alignment_for_contiguous_layout()
171
- and w.ndim == 2
172
- and w.shape[0] % block_size == 0
173
- and w.shape[1] % block_size == 0
174
- )
118
+ provider = getattr(module, "deep_gemm_warmup_provider", None)
119
+ if not isinstance(provider, DeepGemmFp8BlockScaledMMKernel):
120
+ return None
121
+ if module.weight_block_size != get_mk_alignment_for_contiguous_layout():
122
+ return None
123
+ return provider.get_deep_gemm_warmup_weights(module)
124
125
126
def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
@@ -361,11 +310,10 @@ def _deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
310
def deepgemm_fp8_gemm_nt_warmup(
311
model: torch.nn.Module, max_tokens: int, pbar: tqdm | None = None
312
):
364
- dg_modules = [m for m in model.modules() if _fp8_linear_may_use_deep_gemm(m)]
365
-
366
- for dgm in dg_modules:
367
- w, ws, _ = _extract_data_from_linear_base_module(dgm)
368
- _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar)
313
+ for m in model.modules():
314
+ if (data := _deep_gemm_linear_data(m)) is not None:
315
+ w, ws = data
316
+ _deepgemm_fp8_gemm_nt_warmup(w=w, ws=ws, max_tokens=max_tokens, pbar=pbar)
317
318
319
def deepgemm_grouped_fp8_gemm_nt_contiguous_warmup(
@@ -392,8 +340,8 @@ def _count_warmup_iterations(model: torch.nn.Module, max_tokens: int) -> int:
340
341
total = 0
342
for m in model.modules():
395
- if _fp8_linear_may_use_deep_gemm(m):
396
- w, _, _ = _extract_data_from_linear_base_module(m)
343
+ if (data := _deep_gemm_linear_data(m)) is not None:
344
+ w, _ = data
345
if w.size() not in seen_fp8_sizes:
346
total += len(_get_fp8_gemm_nt_m_values(w, max_tokens))
347
seen_fp8_sizes.add(w.size())