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 3, 04:15 UTC
[Bug] Fix ci issue `assert output_size is not None` AssertionError (#43261)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Wentao Ye committed
May 21, 2026 at 04:58 UTC
68e07d59161a8d268b773c181fab17994a7c5d0a
5 files changed
+34
-3
tests/kernels/quantization/test_cutlass_scaled_mm.py
+9
@@ -206,7 +206,16 @@ def test_cutlass_fp8_gemm_padded(
206
207
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
208
209
+ # process_weights_after_loading pad b to 16
210
+ pad_k = (16 - k % 16) % 16
211
+ pad_n = (16 - n % 16) % 16
212
+ if pad_k > 0 or pad_n > 0:
213
+ b = torch.nn.functional.pad(b.t().contiguous(), (0, pad_k, 0, pad_n)).t()
214
+ if pad_n > 0 and scale_b.numel() > 1:
215
+ scale_b = torch.nn.functional.pad(scale_b, (0, pad_n), value=1.0)
216
+
217
kernel = object.__new__(CutlassFP8ScaledMMLinearKernel)
218
+ kernel.logical_output_size = n
219
out = kernel.apply_scaled_mm(
220
A=a,
221
B=b,
tests/quantization/test_fp8.py
+1
-3
@@ -465,11 +465,9 @@ def test_fp8_reloading(
465
method.process_weights_after_loading(layer)
466
467
# test reloading works after loading
468
- # assuming that no reshaping occurred
469
- for name, shape, original_weight_loader in original_metadata:
468
+ for name, shape, _ in original_metadata:
469
param = getattr(layer, name)
470
weight_loader = getattr(param, "weight_loader", default_weight_loader)
472
- assert weight_loader is original_weight_loader
471
weight_loader(param, torch.zeros(shape)) # cannot use empty
472
473
method.process_weights_after_loading(layer)
vllm/model_executor/kernels/linear/scaled_mm/cutlass.py
+21
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
16
CUTLASS_BLOCK_FP8_SUPPORTED,
17
convert_to_channelwise,
18
)
19
+from vllm.model_executor.utils import set_weight_attrs
20
from vllm.platforms import current_platform
21
from vllm.utils.torch_utils import direct_register_custom_op
22
@@ -184,6 +185,14 @@ class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
185
pad_spec[-(2 * dim + 1)] = pad_size
186
return torch.nn.functional.pad(x, pad_spec, value=value)
187
188
+ @staticmethod
189
+ def padded_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
190
+ if loaded_weight.shape != param.shape:
191
+ slices = tuple(slice(0, s) for s in loaded_weight.shape)
192
+ param.data[slices].copy_(loaded_weight)
193
+ else:
194
+ param.data.copy_(loaded_weight.view(param.shape))
195
+
196
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
197
weight_name, weight_scale_name, _, _ = self.layer_param_names
198
weight = getattr(layer, weight_name)
@@ -202,6 +211,12 @@ class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
211
(0, pad_k, 0, pad_n),
212
).t()
213
replace_parameter(layer, weight_name, padded_weight.data)
214
+ set_weight_attrs(
215
+ getattr(layer, weight_name),
216
+ {
217
+ "weight_loader": self.padded_weight_loader,
218
+ },
219
+ )
220
221
weight_scale = getattr(layer, weight_scale_name, None)
222
if weight_scale is not None and pad_n > 0 and weight_scale.numel() > 1:
@@ -210,6 +225,12 @@ class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
225
flat_scale, dim=0, alignment=16, value=1.0
226
).view(-1, *weight_scale.shape[1:])
227
replace_parameter(layer, weight_scale_name, padded_scale.data)
228
+ set_weight_attrs(
229
+ getattr(layer, weight_name),
230
+ {
231
+ "weight_loader": self.padded_weight_loader,
232
+ },
233
+ )
234
235
def apply_scaled_mm(
236
self,
vllm/model_executor/layers/quantization/fp8.py
+1
@@ -557,6 +557,7 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
557
else:
558
weight = qweight.t()
559
replace_parameter(layer, "weight", weight.data)
560
+ self.fp8_linear.process_weights_after_loading(layer)
561
562
# Prevent duplicate processing (e.g., during weight reload)
563
layer._already_called_process_weights_after_loading = True
vllm/model_executor/layers/quantization/online/fp8.py
+2
@@ -153,6 +153,8 @@ class Fp8PerTensorOnlineLinearMethod(_Fp8OnlineLinearBase):
153
replace_parameter(layer, "weight", qweight.t().data)
154
replace_parameter(layer, "weight_scale", weight_scale.data)
155
156
+ self.fp8_linear.process_weights_after_loading(layer)
157
+
158
# Prevent duplicate processing (e.g., during weight reload)
159
layer._already_called_process_weights_after_loading = True
160