| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 7 | MXFP8_BLOCK_SIZE, |
| 8 | ) |
| 9 | |
| 10 | from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig |
| 11 | |
| 12 | |
| 13 | class MarlinMxfp8LinearKernel(Mxfp8LinearKernel): |
| 14 | """MXFP8 W8A16 GEMM via Marlin (SM80+).""" |
| 15 | |
| 16 | @classmethod |
| 17 | def is_supported( |
| 18 | cls, compute_capability: int | None = None |
| 19 | ) -> tuple[bool, str | None]: |
| 20 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 21 | is_fp8_marlin_supported, |
| 22 | ) |
| 23 | |
| 24 | if is_fp8_marlin_supported(): |
| 25 | return True, None |
| 26 | return False, "Marlin FP8 not available" |
| 27 | |
| 28 | @classmethod |
| 29 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 30 | N, K = c.weight_shape |
| 31 | if N <= 0 or K <= 0: |
| 32 | return False, f"Marlin MXFP8 requires positive N and K, got N={N}, K={K}." |
| 33 | if K % MXFP8_BLOCK_SIZE != 0: |
| 34 | return ( |
| 35 | False, |
| 36 | f"Marlin MXFP8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, " |
| 37 | f"got K={K}.", |
| 38 | ) |
| 39 | # Tile misalignment is fixed by zero-padding at weight prep. |
| 40 | return True, None |
| 41 | |
| 42 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 43 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 44 | prepare_mxfp8_layer_for_marlin, |
| 45 | ) |
| 46 | |
| 47 | prepare_mxfp8_layer_for_marlin(layer) |
| 48 | |
| 49 | def apply_weights( |
| 50 | self, |
| 51 | layer: torch.nn.Module, |
| 52 | x: torch.Tensor, |
| 53 | bias: torch.Tensor | None = None, |
| 54 | ) -> torch.Tensor: |
| 55 | from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( |
| 56 | apply_mxfp8_marlin_linear, |
| 57 | ) |
| 58 | |
| 59 | return apply_mxfp8_marlin_linear( |
| 60 | input=x, |
| 61 | weight=layer.weight, |
| 62 | weight_scale=layer.weight_scale, |
| 63 | workspace=None, |
| 64 | size_n=layer.output_size_per_partition, |
| 65 | size_k=layer.input_size_per_partition, |
| 66 | bias=bias, |
| 67 | ) |