[Bugfix][Quantization] Support CT block FP8 with Marlin (#52966)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Misha Goin committed
Aug 19, 2026 at 15:41 UTC
541c6d64c19b64b970edc9122a56027d41b7a0be
2 files changed
+41
-17
tests/quantization/test_compressed_tensors.py
+30
-13
@@ -21,6 +21,9 @@ from tests.models.utils import check_logprobs_close
21
from vllm.model_executor.kernels.linear import (
22
Fp8BlockScaledMMLinearKernel,
23
)
24
+from vllm.model_executor.kernels.linear.scaled_mm import (
25
+ MarlinFP8ScaledMMLinearKernel,
26
+)
27
from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
28
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
29
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
@@ -527,9 +530,14 @@ def test_compressed_tensors_transforms_perplexity(
530
assert perplexity <= exp_perplexity
531
532
530
-def test_compressed_tensors_fp8_block_enabled(vllm_runner):
533
+@pytest.mark.parametrize(
534
+ "linear_backend", ["auto", "marlin"] if current_platform.is_cuda() else ["auto"]
535
+)
536
+def test_compressed_tensors_fp8_block_enabled(vllm_runner, linear_backend):
537
model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK"
532
- with vllm_runner(model_path, enforce_eager=True) as llm:
538
+ with vllm_runner(
539
+ model_path, enforce_eager=True, linear_backend=linear_backend
540
+ ) as llm:
541
fp8_dtype = current_platform.fp8_dtype()
542
543
def check_model(model):
@@ -538,20 +546,29 @@ def test_compressed_tensors_fp8_block_enabled(vllm_runner):
546
qkv_proj = layer.self_attn.qkv_proj
547
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
548
assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8)
541
- assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel)
542
-
543
- assert qkv_proj.weight.dtype is fp8_dtype
544
- assert qkv_proj.weight_scale.dtype is torch.float32
549
+ if linear_backend == "marlin":
550
+ assert isinstance(
551
+ qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel
552
+ )
553
+ assert qkv_proj.weight.dtype is torch.int32
554
+ assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype
555
+ else:
556
+ assert isinstance(
557
+ qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel
558
+ )
559
+ assert qkv_proj.weight.dtype is fp8_dtype
560
+ assert qkv_proj.weight_scale.dtype is torch.float32
561
assert len(qkv_proj.weight.shape) == 2
562
assert len(qkv_proj.weight_scale.shape) == 2
563
548
- input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8
549
- assert isinstance(input_quant_op, QuantFP8)
550
- assert input_quant_op._forward_method in (
551
- input_quant_op.forward_cuda,
552
- input_quant_op.forward_hip,
553
- input_quant_op.forward_xpu,
554
- )
564
+ if linear_backend == "auto":
565
+ input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8
566
+ assert isinstance(input_quant_op, QuantFP8)
567
+ assert input_quant_op._forward_method in (
568
+ input_quant_op.forward_cuda,
569
+ input_quant_op.forward_hip,
570
+ input_quant_op.forward_xpu,
571
+ )
572
573
llm.apply_model(check_model)
574
vllm/model_executor/kernels/linear/scaled_mm/marlin.py
+11
-4
@@ -57,14 +57,21 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
57
self.block_quant = self.config.weight_quant_key in {kFp8Static128BlockSym}
58
self.size_k_first = not self.block_quant
59
60
+ @staticmethod
61
+ def _block_scale_name(layer: torch.nn.Module) -> str:
62
+ if getattr(layer, "weight_scale_inv", None) is not None:
63
+ return "weight_scale_inv"
64
+ return "weight_scale"
65
+
66
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
67
if self.block_quant:
62
- weight, weight_scale_inv = process_fp8_weight_block_strategy(
63
- layer.weight, layer.weight_scale_inv
68
+ scale_name = self._block_scale_name(layer)
69
+ weight, weight_scale = process_fp8_weight_block_strategy(
70
+ layer.weight, getattr(layer, scale_name)
71
)
72
# Update layer with new values
73
replace_parameter(layer, "weight", weight.data)
67
- replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
74
+ replace_parameter(layer, scale_name, weight_scale.data)
75
# Non-block: callers must pass weight in (K, N) layout.
76
77
layer.input_scale = None
@@ -80,7 +87,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
87
bias: torch.Tensor | None = None,
88
) -> torch.Tensor:
89
if self.block_quant:
83
- weight_scale = layer.weight_scale_inv
90
+ weight_scale = getattr(layer, self._block_scale_name(layer))
91
else:
92
weight_scale = layer.weight_scale
93
return apply_fp8_marlin_linear(