diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index a301e8d9cf..304ef1efb0 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/tests/quantization/test_compressed_tensors.py @@ -21,6 +21,9 @@ from tests.models.utils import check_logprobs_close from vllm.model_executor.kernels.linear import ( Fp8BlockScaledMMLinearKernel, ) +from vllm.model_executor.kernels.linear.scaled_mm import ( + MarlinFP8ScaledMMLinearKernel, +) from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 @@ -527,9 +530,14 @@ def test_compressed_tensors_transforms_perplexity( assert perplexity <= exp_perplexity -def test_compressed_tensors_fp8_block_enabled(vllm_runner): +@pytest.mark.parametrize( + "linear_backend", ["auto", "marlin"] if current_platform.is_cuda() else ["auto"] +) +def test_compressed_tensors_fp8_block_enabled(vllm_runner, linear_backend): model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK" - with vllm_runner(model_path, enforce_eager=True) as llm: + with vllm_runner( + model_path, enforce_eager=True, linear_backend=linear_backend + ) as llm: fp8_dtype = current_platform.fp8_dtype() def check_model(model): @@ -538,20 +546,29 @@ def test_compressed_tensors_fp8_block_enabled(vllm_runner): qkv_proj = layer.self_attn.qkv_proj assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8) - assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel) - - assert qkv_proj.weight.dtype is fp8_dtype - assert qkv_proj.weight_scale.dtype is torch.float32 + if linear_backend == "marlin": + assert isinstance( + qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel + ) + assert qkv_proj.weight.dtype is torch.int32 + assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype + else: + assert isinstance( + qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel + ) + assert qkv_proj.weight.dtype is fp8_dtype + assert qkv_proj.weight_scale.dtype is torch.float32 assert len(qkv_proj.weight.shape) == 2 assert len(qkv_proj.weight_scale.shape) == 2 - input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8 - assert isinstance(input_quant_op, QuantFP8) - assert input_quant_op._forward_method in ( - input_quant_op.forward_cuda, - input_quant_op.forward_hip, - input_quant_op.forward_xpu, - ) + if linear_backend == "auto": + input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8 + assert isinstance(input_quant_op, QuantFP8) + assert input_quant_op._forward_method in ( + input_quant_op.forward_cuda, + input_quant_op.forward_hip, + input_quant_op.forward_xpu, + ) llm.apply_model(check_model) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 53bdc724bb..9e148ca308 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -57,14 +57,21 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): self.block_quant = self.config.weight_quant_key in {kFp8Static128BlockSym} self.size_k_first = not self.block_quant + @staticmethod + def _block_scale_name(layer: torch.nn.Module) -> str: + if getattr(layer, "weight_scale_inv", None) is not None: + return "weight_scale_inv" + return "weight_scale" + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: if self.block_quant: - weight, weight_scale_inv = process_fp8_weight_block_strategy( - layer.weight, layer.weight_scale_inv + scale_name = self._block_scale_name(layer) + weight, weight_scale = process_fp8_weight_block_strategy( + layer.weight, getattr(layer, scale_name) ) # Update layer with new values replace_parameter(layer, "weight", weight.data) - replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data) + replace_parameter(layer, scale_name, weight_scale.data) # Non-block: callers must pass weight in (K, N) layout. layer.input_scale = None @@ -80,7 +87,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): bias: torch.Tensor | None = None, ) -> torch.Tensor: if self.block_quant: - weight_scale = layer.weight_scale_inv + weight_scale = getattr(layer, self._block_scale_name(layer)) else: weight_scale = layer.weight_scale return apply_fp8_marlin_linear(