| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Sequence |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | import vllm.envs as envs |
| 9 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 10 | process_fp8_weight_block_strategy, |
| 11 | ) |
| 12 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 13 | apply_fp8_marlin_linear, |
| 14 | is_fp8_marlin_supported, |
| 15 | prepare_fp8_layer_for_marlin, |
| 16 | ) |
| 17 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 18 | kFp8Static128BlockSym, |
| 19 | ) |
| 20 | from vllm.model_executor.utils import replace_parameter |
| 21 | from vllm.platforms import current_platform |
| 22 | |
| 23 | from .ScaledMMLinearKernel import ( |
| 24 | FP8ScaledMMLinearKernel, |
| 25 | FP8ScaledMMLinearLayerConfig, |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 30 | """FP8 Marlin kernel for GPUs that lack FP8 hardware support. |
| 31 | Leverages the Marlin kernel for fast weight-only FP8 quantization. |
| 32 | """ |
| 33 | |
| 34 | @classmethod |
| 35 | def is_supported( |
| 36 | cls, compute_capability: int | None = None |
| 37 | ) -> tuple[bool, str | None]: |
| 38 | if not current_platform.is_cuda(): |
| 39 | return False, "requires CUDA." |
| 40 | # Check if platform supports FP8 Marlin |
| 41 | if not is_fp8_marlin_supported(): |
| 42 | return False, "FP8 Marlin requires compute capability 7.5 or higher" |
| 43 | if envs.VLLM_BATCH_INVARIANT: |
| 44 | return False, "FP8 Marlin not supported for batch invariant execution." |
| 45 | return True, None |
| 46 | |
| 47 | @classmethod |
| 48 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 49 | return True, None |
| 50 | |
| 51 | def __init__( |
| 52 | self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str] |
| 53 | ) -> None: |
| 54 | super().__init__(c, layer_param_names) |
| 55 | self.marlin_input_dtype = None |
| 56 | self.block_quant = self.config.weight_quant_key in {kFp8Static128BlockSym} |
| 57 | self.size_k_first = not self.block_quant |
| 58 | |
| 59 | @staticmethod |
| 60 | def _block_scale_name(layer: torch.nn.Module) -> str: |
| 61 | if getattr(layer, "weight_scale_inv", None) is not None: |
| 62 | return "weight_scale_inv" |
| 63 | return "weight_scale" |
| 64 | |
| 65 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 66 | if self.block_quant: |
| 67 | scale_name = self._block_scale_name(layer) |
| 68 | weight, weight_scale = process_fp8_weight_block_strategy( |
| 69 | layer.weight, getattr(layer, scale_name) |
| 70 | ) |
| 71 | # Update layer with new values |
| 72 | replace_parameter(layer, "weight", weight.data) |
| 73 | replace_parameter(layer, scale_name, weight_scale.data) |
| 74 | # Non-block: callers must pass weight in (K, N) layout. |
| 75 | |
| 76 | layer.input_scale = None |
| 77 | prepare_fp8_layer_for_marlin( |
| 78 | layer, self.size_k_first, input_dtype=self.marlin_input_dtype |
| 79 | ) |
| 80 | del layer.input_scale |
| 81 | |
| 82 | def apply_weights( |
| 83 | self, |
| 84 | layer: torch.nn.Module, |
| 85 | x: torch.Tensor, |
| 86 | bias: torch.Tensor | None = None, |
| 87 | ) -> torch.Tensor: |
| 88 | if self.block_quant: |
| 89 | weight_scale = getattr(layer, self._block_scale_name(layer)) |
| 90 | else: |
| 91 | weight_scale = layer.weight_scale |
| 92 | return apply_fp8_marlin_linear( |
| 93 | input=x, |
| 94 | weight=layer.weight, |
| 95 | weight_scale=weight_scale, |
| 96 | workspace=None, |
| 97 | size_n=layer.output_size_per_partition, |
| 98 | size_k=layer.input_size_per_partition, |
| 99 | input_dtype=self.marlin_input_dtype, |
| 100 | bias=bias, |
| 101 | ) |
| 102 | |
| 103 | def apply_scaled_mm( |
| 104 | self, |
| 105 | *, |
| 106 | A: torch.Tensor, |
| 107 | B: torch.Tensor, |
| 108 | out_dtype: torch.dtype, |
| 109 | As: torch.Tensor, |
| 110 | Bs: torch.Tensor, |
| 111 | bias: torch.Tensor | None, |
| 112 | output_shape: list, |
| 113 | ) -> torch.Tensor: |
| 114 | pass |