A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Canonicalize FP8 weight layout to (K, N) at the source (#44735)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jun 8, 2026 at 16:37 UTC
6afa25000c42eea892a779a1a9e25d8c447b8891
3 files changed
+11
-29
vllm/model_executor/kernels/linear/scaled_mm/marlin.py
+1
-19
@@ -75,25 +75,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
75
# Update layer with new values
76
replace_parameter(layer, "weight", weight.data)
77
replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
78
- else:
79
- w_q, *_ = self._get_layer_params(layer)
80
- # Compressed tensors transposes the weight to (K, N)
81
- # for channel and tensor quant strategies.
82
- # So we can skip the transpose if the layout is
83
- # already (K, N).
84
- # TODO: Remove this check once the layouts have been
85
- # canonicalized to a standard (N, K) dimension. See issue
86
- # #33314 for more details.
87
- if w_q.shape != (
88
- layer.input_size_per_partition,
89
- layer.output_size_per_partition,
90
- ):
91
- # transpose the weights to (K,N)
92
- replace_parameter(
93
- layer,
94
- "weight",
95
- w_q.t(),
96
- )
78
+ # Non-block: callers must pass weight in (K, N) layout.
79
80
layer.input_scale = None
81
prepare_fp8_layer_for_marlin(
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py
+2
@@ -137,6 +137,8 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
137
"weight_scale",
138
convert_to_channelwise(layer.weight_scale, layer.logical_widths),
139
)
140
+ # Canonicalize to (K, N) for the kernel.
141
+ replace_parameter(layer, "weight", layer.weight.t())
142
143
self.linear_kernel.process_weights_after_loading(layer)
144
vllm/model_executor/layers/quantization/fp8.py
+8
-10
@@ -384,6 +384,9 @@ class Fp8LinearMethod(LinearMethodBase):
384
385
def process_weights_after_loading(self, layer: RoutedExperts) -> None:
386
if self.use_marlin:
387
+ if not self.block_quant:
388
+ # Canonicalize to (K, N) for the kernel.
389
+ replace_parameter(layer, "weight", layer.weight.t())
390
# Only Marlin kernels support `marlin_input_dtype`; guard to avoid
391
# AttributeError if backend selection changes.
392
if hasattr(self.fp8_linear, "marlin_input_dtype"):
@@ -537,20 +540,15 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
540
layer.input_scale = None
541
qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
542
540
- # Update layer with new values.
541
- replace_parameter(layer, "weight", qweight.data)
543
+ # Update layer with new values. Canonicalize to (K, N) for the kernel.
544
+ replace_parameter(layer, "weight", qweight.t().data)
545
replace_parameter(layer, "weight_scale", weight_scale.data)
546
544
- if self.use_marlin:
547
+ if self.use_marlin and hasattr(self.fp8_linear, "marlin_input_dtype"):
548
# Only Marlin kernels support `marlin_input_dtype`; guard to avoid
549
# AttributeError if backend selection changes.
547
- if hasattr(self.fp8_linear, "marlin_input_dtype"):
548
- self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
549
- self.fp8_linear.process_weights_after_loading(layer)
550
- else:
551
- weight = qweight.t()
552
- replace_parameter(layer, "weight", weight.data)
553
- self.fp8_linear.process_weights_after_loading(layer)
550
+ self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
551
+ self.fp8_linear.process_weights_after_loading(layer)
552
553
# Prevent duplicate processing (e.g., during weight reload)
554
layer._already_called_process_weights_after_loading = True