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 5, 23:30 UTC
[Bugfix][CI] Fix Marlin FP8 Linear Kernel for Compressed Tensors Format (#38092)
Signed-off-by: BadrBasowid <Badr.Basowid@gmail.com> Signed-off-by: BadrBasowid <61441185+BadrBasowid@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
BadrBasowid committed
Mar 26, 2026 at 12:11 UTC
e6bf9f15ec9d8bc75a583cd9477159b1232b635b
3 files changed
+25
-2
vllm/model_executor/kernels/linear/scaled_mm/marlin.py
+19
-2
@@ -76,8 +76,25 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
76
replace_parameter(layer, "weight", weight.data)
77
replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
78
else:
79
- weight = layer.weight.t()
80
- replace_parameter(layer, "weight", weight.data)
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
+ )
97
+
98
layer.input_scale = None
99
prepare_fp8_layer_for_marlin(
100
layer, self.size_k_first, input_dtype=self.marlin_input_dtype
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+3
@@ -188,6 +188,9 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
188
if self.strategy == QuantizationStrategy.BLOCK:
189
maybe_post_process_fp8_weight_block(layer)
190
191
+ if hasattr(self, "fp8_linear"):
192
+ self.fp8_linear.process_weights_after_loading(layer)
193
+
194
def apply_weights(
195
self,
196
layer: torch.nn.Module,
vllm/model_executor/layers/quantization/modelopt.py
+3
@@ -705,6 +705,9 @@ class ModelOptFp8PbWoLinearMethod(LinearMethodBase):
705
706
layer.weight_scale = Parameter(scale.contiguous(), requires_grad=False)
707
708
+ if hasattr(self, "fp8_linear"):
709
+ self.fp8_linear.process_weights_after_loading(layer)
710
+
711
def apply(
712
self,
713
layer: torch.nn.Module,