| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import math |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 10 | kFp8DynamicTensorSym, |
| 11 | kFp8DynamicTokenSym, |
| 12 | kFp8StaticChannelSym, |
| 13 | kFp8StaticTensorSym, |
| 14 | kFp8StaticTokenSym, |
| 15 | ) |
| 16 | from vllm.model_executor.utils import replace_parameter |
| 17 | from vllm.platforms import current_platform |
| 18 | |
| 19 | from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel |
| 20 | from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig |
| 21 | |
| 22 | |
| 23 | class XPUW8A8FP8LinearKernel(FP8ScaledMMLinearKernel): |
| 24 | _SUPPORTED_ACT_QUANT_KEYS = { |
| 25 | kFp8DynamicTensorSym, |
| 26 | kFp8DynamicTokenSym, |
| 27 | kFp8StaticTensorSym, |
| 28 | kFp8StaticTokenSym, |
| 29 | } |
| 30 | _SUPPORTED_WEIGHT_QUANT_KEYS = { |
| 31 | kFp8StaticChannelSym, |
| 32 | kFp8StaticTensorSym, |
| 33 | } |
| 34 | |
| 35 | @classmethod |
| 36 | def is_supported( |
| 37 | cls, compute_capability: int | None = None |
| 38 | ) -> tuple[bool, str | None]: |
| 39 | if not current_platform.is_xpu(): |
| 40 | return False, "XPUW8A8FP8Linear only support on XPU" |
| 41 | return True, None |
| 42 | |
| 43 | @classmethod |
| 44 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 45 | if c.weight_quant_key not in cls._SUPPORTED_WEIGHT_QUANT_KEYS: |
| 46 | return ( |
| 47 | False, |
| 48 | "XPUW8A8FP8Linear only support per-channel and per-tensor quantization", |
| 49 | ) |
| 50 | if c.activation_quant_key not in cls._SUPPORTED_ACT_QUANT_KEYS: |
| 51 | return ( |
| 52 | False, |
| 53 | "XPUW8A8FP8Linear only support per-tensor and per-token activation " |
| 54 | "quantization", |
| 55 | ) |
| 56 | if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}: |
| 57 | return False, "XPUW8A8FP8Linear only support FP8 weight dtype" |
| 58 | if c.activation_quant_key.dtype not in { |
| 59 | torch.float8_e5m2, |
| 60 | torch.float8_e4m3fn, |
| 61 | }: |
| 62 | return False, "XPUW8A8FP8Linear only support FP8 activation dtype" |
| 63 | return True, None |
| 64 | |
| 65 | def __init__( |
| 66 | self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str] |
| 67 | ) -> None: |
| 68 | super().__init__(c, layer_param_names) |
| 69 | |
| 70 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 71 | """Ensure weight is stored as C-contiguous [K, N] (KN layout). |
| 72 | |
| 73 | Checkpoints store weight as [N, K]; fp8_gemm requires [K, N], |
| 74 | C-contiguous. Three incoming layouts are possible: |
| 75 | • [N, K] C-contiguous ← direct checkpoint → .t().contiguous() |
| 76 | • [K, N] Fortran-order ← fp8.py's weight.t() → .contiguous() |
| 77 | • [K, N] C-contiguous ← already correct → no-op |
| 78 | |
| 79 | For square weights (K == N) the shape is ambiguous; contiguity is used |
| 80 | as a proxy: C-contiguous ≡ checkpoint [N, K] (needs transpose); |
| 81 | Fortran-order ≡ fp8.py already transposed (needs only contiguous). |
| 82 | """ |
| 83 | K = getattr(layer, "input_size_per_partition", self.config.weight_shape[1]) |
| 84 | N = getattr(layer, "output_size_per_partition", self.config.weight_shape[0]) |
| 85 | w = layer.weight |
| 86 | |
| 87 | if w.shape not in {(K, N), (N, K)}: |
| 88 | raise ValueError( |
| 89 | f"XPUFP8ScaledMM expects weight shape ({K},{N}) or ({N},{K}), " |
| 90 | f"but got {tuple(w.shape)}" |
| 91 | ) |
| 92 | |
| 93 | needs_transpose = w.shape == (N, K) if K != N else w.is_contiguous() |
| 94 | layer_weight = w.t() if needs_transpose else w |
| 95 | replace_parameter(layer, "weight", layer_weight) |
| 96 | ws = layer.weight_scale |
| 97 | if ws.numel() == 1: |
| 98 | replace_parameter(layer, "weight_scale", ws.reshape(1)) |
| 99 | |
| 100 | def apply_scaled_mm( |
| 101 | self, |
| 102 | *, |
| 103 | A: torch.Tensor, |
| 104 | B: torch.Tensor, |
| 105 | out_dtype: torch.dtype, |
| 106 | As: torch.Tensor, |
| 107 | Bs: torch.Tensor, |
| 108 | bias: torch.Tensor | None, |
| 109 | output_shape: list, |
| 110 | ) -> torch.Tensor: |
| 111 | # B is C-contiguous [K, N] from process_weights_after_loading. |
| 112 | # fp8_gemm routes on scale dtype (float32) and numel: |
| 113 | # As [1] → per-tensor (numel==1 branch) |
| 114 | # As [M,1] → per-token (group={1,K} branch, broadcast across K) |
| 115 | # Bs [1] → per-tensor |
| 116 | # Bs [N] → per-channel (mask=bit1 branch) |
| 117 | # No shape manipulation needed here. |
| 118 | output = torch.ops._xpu_C.fp8_gemm(A, B, out_dtype, As, Bs, bias) |
| 119 | return output.view(*output_shape) |
| 120 | |
| 121 | |
| 122 | class XPUW8A16FP8LinearKernel(FP8ScaledMMLinearKernel): |
| 123 | @classmethod |
| 124 | def is_supported( |
| 125 | cls, compute_capability: int | None = None |
| 126 | ) -> tuple[bool, str | None]: |
| 127 | if not current_platform.is_xpu(): |
| 128 | return False, "XPUW8A16FP8Linear only support on XPU" |
| 129 | return True, None |
| 130 | |
| 131 | @classmethod |
| 132 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 133 | if c.weight_quant_key not in {kFp8StaticChannelSym, kFp8StaticTensorSym}: |
| 134 | return ( |
| 135 | False, |
| 136 | "XPUW8A16FP8Linear only support per-channel and per-tensor " |
| 137 | "quantization", |
| 138 | ) |
| 139 | if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}: |
| 140 | return False, "XPUW8A16FP8Linear only support FP8 weight dtype" |
| 141 | return True, None |
| 142 | |
| 143 | def __init__( |
| 144 | self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str] |
| 145 | ) -> None: |
| 146 | assert self.can_implement(c)[0] |
| 147 | assert self.is_supported()[0] |
| 148 | self.config = c |
| 149 | self.layer_param_names = layer_param_names |
| 150 | |
| 151 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 152 | # fp8_gemm_w8a16 expects weight in [in, out] layout. |
| 153 | # Transpose if weight is still in [out, in] layout. |
| 154 | # For square matrices, use contiguity as tie-breaker: |
| 155 | # checkpoint weights are contiguous, .t() views are not. |
| 156 | weight = layer.weight |
| 157 | out_features, in_features = self.config.weight_shape |
| 158 | |
| 159 | if weight.shape == (out_features, in_features) and ( |
| 160 | in_features != out_features or weight.is_contiguous() |
| 161 | ): |
| 162 | replace_parameter(layer, "weight", weight.data.t()) |
| 163 | # else: already in [in, out] layout — no-op |
| 164 | |
| 165 | weight_scale = layer.weight_scale.t().contiguous() |
| 166 | replace_parameter(layer, "weight_scale", weight_scale.data) |
| 167 | |
| 168 | def apply_weights( |
| 169 | self, |
| 170 | layer: torch.nn.Module, |
| 171 | x: torch.Tensor, |
| 172 | bias: torch.Tensor | None = None, |
| 173 | ) -> torch.Tensor: |
| 174 | weight = layer.weight |
| 175 | weight_scale = layer.weight_scale |
| 176 | return torch.ops._xpu_C.fp8_gemm_w8a16(x, weight, weight_scale, bias) |
| 177 | |
| 178 | def apply_scaled_mm( |
| 179 | self, |
| 180 | *, |
| 181 | A: torch.Tensor, |
| 182 | B: torch.Tensor, |
| 183 | out_dtype: torch.dtype, |
| 184 | As: torch.Tensor, |
| 185 | Bs: torch.Tensor, |
| 186 | bias: torch.Tensor | None, |
| 187 | output_shape: list, |
| 188 | ) -> torch.Tensor: |
| 189 | pass |
| 190 | |
| 191 | |
| 192 | class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): |
| 193 | @classmethod |
| 194 | def is_supported( |
| 195 | cls, compute_capability: int | None = None |
| 196 | ) -> tuple[bool, str | None]: |
| 197 | if not current_platform.is_xpu(): |
| 198 | return False, "XPUFp8BlockScaledMM only support on XPU" |
| 199 | return True, None |
| 200 | |
| 201 | def process_weights_after_loading(self, layer: torch.nn.Module): |
| 202 | super().process_weights_after_loading(layer) |
| 203 | scale_attr = ( |
| 204 | "weight_scale_inv" if hasattr(layer, "weight_scale_inv") else "weight_scale" |
| 205 | ) |
| 206 | scale = getattr(layer, scale_attr) |
| 207 | |
| 208 | # Ragged N (N % block_n != 0): oneDNN needs n_blocks to divide N. |
| 209 | # Weight untouched; only repeat scale rows to a finer N-group gn that |
| 210 | # divides both N and block_n (gn = gcd(N, block_n)): |
| 211 | # scale [ceil(N/block_n), K/block_k] --> [N/gn, K/block_k] |
| 212 | # oneDNN only accepts gn that is a multiple of 16, and gcd(N, block_n) |
| 213 | # is a power of two (block_n=128), so gn must be >= 16. No-op when |
| 214 | # N % block_n == 0. |
| 215 | block_n, block_k = self.weight_group_shape |
| 216 | N, K = layer.weight.shape |
| 217 | if N % block_n != 0: |
| 218 | gn = math.gcd(N, block_n) |
| 219 | assert gn % 16 == 0, ( |
| 220 | f"XPU block-scaled FP8: N ({N}) yields group width {gn}, but " |
| 221 | f"oneDNN only supports multiples of 16; this weight shape is " |
| 222 | f"unsupported." |
| 223 | ) |
| 224 | col_start = torch.arange(N // gn, device=scale.device) * gn |
| 225 | src_idx = torch.div(col_start, block_n, rounding_mode="floor") |
| 226 | scale = scale.index_select(0, src_idx).contiguous() |
| 227 | |
| 228 | # Ragged K needs the runtime activation scale expanded too, which we |
| 229 | # don't handle; DeepSeek/GLM keep K block-aligned, so fail loudly. |
| 230 | assert K % block_k == 0, ( |
| 231 | f"XPU block-scaled FP8 requires K ({K}) to be a multiple of the " |
| 232 | f"weight block size ({block_k}); ragged-K weights are unsupported." |
| 233 | ) |
| 234 | |
| 235 | # Checkpoint scale is [n_blocks, k_blocks] (one value per block tile). |
| 236 | # oneDNN fp8_gemm requires contiguous [k_blocks, n_blocks] layout. |
| 237 | # We store the transposed contiguous buffer as a .t() view so that: |
| 238 | # - MLA's scaled_dequantize still sees [n_blocks, k_blocks] shape |
| 239 | # - apply_block_scaled_mm recovers the contiguous buffer via .t() |
| 240 | scale_kn = scale.data.t().contiguous() # [k_blocks, n_blocks] |
| 241 | replace_parameter(layer, scale_attr, scale_kn.t()) # view: [n_blocks, k_blocks] |
| 242 | |
| 243 | if getattr(layer, "is_bmm", False): |
| 244 | self._prepare_bmm_params(layer, scale_kn) |
| 245 | |
| 246 | def _prepare_bmm_params( |
| 247 | self, layer: torch.nn.Module, scale_kn: torch.Tensor |
| 248 | ) -> None: |
| 249 | """Precompute batched weight and scale for grouped fp8_bmm (e.g. wo_a). |
| 250 | |
| 251 | Splits scale [k_blocks, n_blocks] into [G, k_blocks, n_blocks_per_group] |
| 252 | and weight [N_total, K] into [G, K, N_per_group] for batch GEMM. |
| 253 | """ |
| 254 | batch = layer.bmm_batch_size |
| 255 | k_blocks, n_blocks = scale_kn.shape |
| 256 | layer.bmm_scale = ( |
| 257 | scale_kn.reshape(k_blocks, batch, n_blocks // batch) |
| 258 | .permute(1, 0, 2) |
| 259 | .contiguous() |
| 260 | ) |
| 261 | w = layer.weight |
| 262 | N_total, K = w.shape |
| 263 | layer.bmm_weight = w.reshape(batch, N_total // batch, K).permute( |
| 264 | 0, 2, 1 |
| 265 | ) # [G, K, N_per_group] |
| 266 | |
| 267 | def apply_block_scaled_mm( |
| 268 | self, |
| 269 | A: torch.Tensor, |
| 270 | B: torch.Tensor, |
| 271 | As: torch.Tensor, |
| 272 | Bs: torch.Tensor, |
| 273 | ) -> torch.Tensor: |
| 274 | # B is [N, K]; .t() gives [K, N] view (no copy). |
| 275 | # Bs is stored as [n_blocks, k_blocks] view; .t() recovers the |
| 276 | # contiguous [k_blocks, n_blocks] buffer that oneDNN expects. |
| 277 | return torch.ops._xpu_C.fp8_gemm( |
| 278 | A, |
| 279 | B.t(), |
| 280 | self.config.out_dtype, |
| 281 | As, |
| 282 | Bs.t(), |
| 283 | torch.Tensor(), |
| 284 | ) |