| 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 | xpu_mxfp8_quantize as quant_mxfp8, |
| 8 | ) |
| 9 | from vllm.model_executor.utils import replace_parameter |
| 10 | from vllm.platforms import current_platform |
| 11 | |
| 12 | from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig |
| 13 | |
| 14 | |
| 15 | class XPUMxFp8LinearKernel(Mxfp8LinearKernel): |
| 16 | """MXFP8 W8A8 GEMM on XPU.""" |
| 17 | |
| 18 | @classmethod |
| 19 | def is_supported( |
| 20 | cls, compute_capability: int | None = None |
| 21 | ) -> tuple[bool, str | None]: |
| 22 | if not current_platform.is_xpu(): |
| 23 | return False, "XPUMxFp8 only support on XPU" |
| 24 | return True, None |
| 25 | |
| 26 | @classmethod |
| 27 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 28 | return True, None |
| 29 | |
| 30 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 31 | # Checkpoint scale is [N, K//32] (one E8M0 scale per 32 elements). |
| 32 | # oneDNN fp8_gemm requires contiguous [K//32, N] layout. Store the |
| 33 | # transposed contiguous buffer as a .t() view so that: |
| 34 | # - dequantization consumers still see the checkpoint shape |
| 35 | # - apply_weights recovers the oneDNN layout via .t() at zero cost |
| 36 | weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu) |
| 37 | scale_kn = weight_scale.data.t().contiguous() |
| 38 | replace_parameter(layer, "weight_scale", scale_kn.t()) |
| 39 | |
| 40 | if getattr(layer, "is_bmm", False): |
| 41 | self._prepare_bmm_params(layer, scale_kn) |
| 42 | |
| 43 | def _prepare_bmm_params( |
| 44 | self, layer: torch.nn.Module, scale_kn: torch.Tensor |
| 45 | ) -> None: |
| 46 | """Precompute batched weight and scale for grouped fp8_bmm (e.g. wo_a). |
| 47 | |
| 48 | Splits scale [K//32, N_total] into [G, K//32, N_per_group] and weight |
| 49 | [N_total, K] into contiguous [G, K, N_per_group] for batch GEMM. |
| 50 | """ |
| 51 | batch = layer.bmm_batch_size |
| 52 | k_blocks, n_blocks = scale_kn.shape |
| 53 | layer.bmm_scale = ( |
| 54 | scale_kn.reshape(k_blocks, batch, n_blocks // batch) |
| 55 | .permute(1, 0, 2) |
| 56 | .contiguous() |
| 57 | ) |
| 58 | w = layer.weight |
| 59 | n_total, k = w.shape |
| 60 | layer.bmm_weight = ( |
| 61 | w.reshape(batch, n_total // batch, k).permute(0, 2, 1).contiguous() |
| 62 | ) |
| 63 | |
| 64 | def apply_weights( |
| 65 | self, |
| 66 | layer: torch.nn.Module, |
| 67 | x: torch.Tensor, |
| 68 | bias: torch.Tensor | None = None, |
| 69 | ) -> torch.Tensor: |
| 70 | out_dtype = x.dtype |
| 71 | x_fp8, x_scale = quant_mxfp8(x) |
| 72 | # Weight is [N, K]; .t() gives a [K, N] view without copying. |
| 73 | # Scale is stored as [N, K//32]; .t() recovers the contiguous |
| 74 | # [K//32, N] buffer that oneDNN expects. |
| 75 | return torch.ops._xpu_C.fp8_gemm( |
| 76 | x_fp8, |
| 77 | layer.weight.t(), |
| 78 | out_dtype, |
| 79 | x_scale, |
| 80 | layer.weight_scale.t(), |
| 81 | bias, |
| 82 | ) |