| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | from torch.nn.parameter import Parameter |
| 6 | |
| 7 | from vllm.model_executor.layers.fusion.quant_activation import ( |
| 8 | QuantizedActivation, |
| 9 | as_quantized_activation, |
| 10 | ) |
| 11 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 12 | MXFP8_BLOCK_SIZE, |
| 13 | mxfp8_e4m3_quantize, |
| 14 | swizzle_mxfp8_scale, |
| 15 | ) |
| 16 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 17 | QuantKey, |
| 18 | kMxfp8Dynamic, |
| 19 | ) |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.utils import flashinfer as vllm_flashinfer |
| 22 | from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutedsl |
| 23 | |
| 24 | from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig |
| 25 | |
| 26 | |
| 27 | def _check_mm_mxfp8_shape(weight_shape: tuple[int, int]) -> tuple[bool, str | None]: |
| 28 | N, K = weight_shape |
| 29 | if K < 128: |
| 30 | return False, f"mm_mxfp8 requires K >= 128, got K={K}." |
| 31 | if K % MXFP8_BLOCK_SIZE != 0: |
| 32 | return ( |
| 33 | False, |
| 34 | f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}.", |
| 35 | ) |
| 36 | if N < 128: |
| 37 | return False, f"mm_mxfp8 requires N >= 128, got N={N}." |
| 38 | return True, None |
| 39 | |
| 40 | |
| 41 | class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel): |
| 42 | """MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+).""" |
| 43 | |
| 44 | supports_pre_processed_weights = True |
| 45 | |
| 46 | @classmethod |
| 47 | def is_supported( |
| 48 | cls, compute_capability: int | None = None |
| 49 | ) -> tuple[bool, str | None]: |
| 50 | if not ( |
| 51 | current_platform.is_cuda() and current_platform.has_device_capability(100) |
| 52 | ): |
| 53 | return False, "requires >=sm_100 (Blackwell)" |
| 54 | if not has_flashinfer(): |
| 55 | return False, "requires FlashInfer" |
| 56 | return True, None |
| 57 | |
| 58 | @classmethod |
| 59 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 60 | can_implement, reason = _check_mm_mxfp8_shape(c.weight_shape) |
| 61 | if not can_implement: |
| 62 | return False, reason |
| 63 | N, _ = c.weight_shape |
| 64 | if current_platform.is_device_capability_family(120) and N % 32 != 0: |
| 65 | return ( |
| 66 | False, |
| 67 | f"mm_mxfp8 on SM12x requires N to be divisible by 32, got N={N}.", |
| 68 | ) |
| 69 | return True, None |
| 70 | |
| 71 | def input_quant_key(self) -> QuantKey: |
| 72 | # Activations use FlashInfer's F8_128x4 swizzled scale layout. |
| 73 | return kMxfp8Dynamic |
| 74 | |
| 75 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 76 | weight = layer.weight.data # [N, K] |
| 77 | N, K = weight.shape |
| 78 | |
| 79 | scale_k = K // MXFP8_BLOCK_SIZE |
| 80 | weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous() |
| 81 | weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K) |
| 82 | |
| 83 | layer.weight = Parameter(weight.contiguous(), requires_grad=False) |
| 84 | layer.weight_scale = Parameter( |
| 85 | weight_scale_swizzled.contiguous(), requires_grad=False |
| 86 | ) |
| 87 | |
| 88 | def apply_weights( |
| 89 | self, |
| 90 | layer: torch.nn.Module, |
| 91 | x: torch.Tensor | QuantizedActivation, |
| 92 | bias: torch.Tensor | None = None, |
| 93 | ) -> torch.Tensor: |
| 94 | weight = layer.weight |
| 95 | weight_scale = layer.weight_scale |
| 96 | N, K = weight.shape |
| 97 | |
| 98 | qa = as_quantized_activation(x, self.input_quant_key()) |
| 99 | if qa is not None: |
| 100 | input_mxfp8, input_scale = qa.data, qa.scale |
| 101 | out_dtype, input_shape = qa.orig_dtype, qa.orig_shape |
| 102 | else: |
| 103 | assert isinstance(x, torch.Tensor) |
| 104 | out_dtype, input_shape = x.dtype, x.shape |
| 105 | input_mxfp8, input_scale = mxfp8_e4m3_quantize( |
| 106 | x.view(-1, K), is_sf_swizzled_layout=True |
| 107 | ) |
| 108 | |
| 109 | if not weight.is_contiguous(): |
| 110 | weight = weight.contiguous() |
| 111 | |
| 112 | output = vllm_flashinfer.mm_mxfp8( |
| 113 | input_mxfp8, |
| 114 | weight.t(), |
| 115 | input_scale, |
| 116 | weight_scale, |
| 117 | out_dtype=out_dtype, |
| 118 | backend="cutlass", |
| 119 | ) |
| 120 | |
| 121 | if bias is not None: |
| 122 | output = output + bias |
| 123 | |
| 124 | output_shape = (*input_shape[:-1], N) |
| 125 | return output.view(output_shape) |
| 126 | |
| 127 | |
| 128 | class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel): |
| 129 | """MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103).""" |
| 130 | |
| 131 | supports_pre_processed_weights = True |
| 132 | |
| 133 | @classmethod |
| 134 | def is_supported( |
| 135 | cls, compute_capability: int | None = None |
| 136 | ) -> tuple[bool, str | None]: |
| 137 | if not ( |
| 138 | current_platform.is_cuda() |
| 139 | and current_platform.is_device_capability_family(100) |
| 140 | ): |
| 141 | return False, "requires sm_100/sm_103 (Blackwell)" |
| 142 | if not has_flashinfer_cutedsl(): |
| 143 | return False, "requires FlashInfer CuTe-DSL module" |
| 144 | return True, None |
| 145 | |
| 146 | @classmethod |
| 147 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 148 | return _check_mm_mxfp8_shape(c.weight_shape) |
| 149 | |
| 150 | def input_quant_key(self) -> QuantKey: |
| 151 | # Activations use FlashInfer's F8_128x4 swizzled scale layout. |
| 152 | return kMxfp8Dynamic |
| 153 | |
| 154 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 155 | weight = layer.weight.data # [N, K] |
| 156 | N, K = weight.shape |
| 157 | |
| 158 | scale_k = K // MXFP8_BLOCK_SIZE |
| 159 | weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous() |
| 160 | weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K) |
| 161 | |
| 162 | # Store weight column-major [K, N] as mm_mxfp8 expects for operand B. |
| 163 | layer.weight = Parameter(weight.contiguous().t(), requires_grad=False) |
| 164 | layer.weight_scale = Parameter( |
| 165 | weight_scale_swizzled.contiguous(), requires_grad=False |
| 166 | ) |
| 167 | |
| 168 | def apply_weights( |
| 169 | self, |
| 170 | layer: torch.nn.Module, |
| 171 | x: torch.Tensor | QuantizedActivation, |
| 172 | bias: torch.Tensor | None = None, |
| 173 | ) -> torch.Tensor: |
| 174 | weight = layer.weight # [K, N], column-major |
| 175 | weight_scale = layer.weight_scale |
| 176 | K, N = weight.shape |
| 177 | |
| 178 | qa = as_quantized_activation(x, self.input_quant_key()) |
| 179 | if qa is not None: |
| 180 | input_mxfp8, input_scale = qa.data, qa.scale |
| 181 | out_dtype, input_shape = qa.orig_dtype, qa.orig_shape |
| 182 | else: |
| 183 | assert isinstance(x, torch.Tensor) |
| 184 | out_dtype, input_shape = x.dtype, x.shape |
| 185 | input_mxfp8, input_scale = mxfp8_e4m3_quantize( |
| 186 | x.view(-1, K), is_sf_swizzled_layout=True |
| 187 | ) |
| 188 | |
| 189 | output = vllm_flashinfer.mm_mxfp8( |
| 190 | input_mxfp8, |
| 191 | weight, |
| 192 | input_scale, |
| 193 | weight_scale, |
| 194 | out_dtype=out_dtype, |
| 195 | backend="cute-dsl", |
| 196 | ) |
| 197 | |
| 198 | if bias is not None: |
| 199 | output = output + bias |
| 200 | |
| 201 | output_shape = (*input_shape[:-1], N) |
| 202 | return output.view(output_shape) |
| 203 | |
| 204 | |
| 205 | class FlashInferTrtllmMxfp8LinearKernel(Mxfp8LinearKernel): |
| 206 | """MXFP8 W8A8 GEMM via FlashInfer's TensorRT-LLM wrapper.""" |
| 207 | |
| 208 | @classmethod |
| 209 | def is_supported( |
| 210 | cls, compute_capability: int | None = None |
| 211 | ) -> tuple[bool, str | None]: |
| 212 | if not ( |
| 213 | current_platform.is_cuda() |
| 214 | and current_platform.is_device_capability_family(100) |
| 215 | ): |
| 216 | return False, "requires SM100-family GPU" |
| 217 | if not has_flashinfer(): |
| 218 | return False, "requires FlashInfer" |
| 219 | return True, None |
| 220 | |
| 221 | @classmethod |
| 222 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 223 | return True, None |
| 224 | |
| 225 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 226 | from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a |
| 227 | |
| 228 | if hasattr(layer, "_mxfp8_trtllm_output_size") and layer.weight_scale.ndim == 1: |
| 229 | return |
| 230 | |
| 231 | weight = layer.weight.data # [N, K] |
| 232 | N, K = weight.shape |
| 233 | if K % 256 != 0: |
| 234 | raise ValueError( |
| 235 | f"FlashInfer TRTLLM MXFP8 requires K to be divisible by 256, got K={K}." |
| 236 | ) |
| 237 | |
| 238 | scale_k = K // MXFP8_BLOCK_SIZE |
| 239 | weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() |
| 240 | padded_n = ((N + 127) // 128) * 128 |
| 241 | if padded_n != N: |
| 242 | padded_weight = weight.new_zeros((padded_n, K)) |
| 243 | padded_weight[:N] = weight |
| 244 | weight = padded_weight |
| 245 | |
| 246 | padded_scale = weight_scale.new_zeros((padded_n, scale_k)) |
| 247 | padded_scale[:N] = weight_scale |
| 248 | weight_scale = padded_scale |
| 249 | else: |
| 250 | weight = weight.contiguous() |
| 251 | |
| 252 | layer.weight = Parameter( |
| 253 | shuffle_matrix_a(weight, 128).reshape(padded_n, K), |
| 254 | requires_grad=False, |
| 255 | ) |
| 256 | layer.weight_scale = Parameter( |
| 257 | shuffle_matrix_sf_a( |
| 258 | weight_scale, |
| 259 | 128, |
| 260 | num_elts_per_sf=MXFP8_BLOCK_SIZE, |
| 261 | ).reshape(-1), |
| 262 | requires_grad=False, |
| 263 | ) |
| 264 | layer._mxfp8_trtllm_output_size = N |
| 265 | |
| 266 | def apply_weights( |
| 267 | self, |
| 268 | layer: torch.nn.Module, |
| 269 | x: torch.Tensor, |
| 270 | bias: torch.Tensor | None = None, |
| 271 | ) -> torch.Tensor: |
| 272 | assert x.dtype == torch.bfloat16, ( |
| 273 | f"FlashInfer TRTLLM MXFP8 requires bfloat16 activations, got {x.dtype}." |
| 274 | ) |
| 275 | |
| 276 | weight = layer.weight # shuffled [padded N, K] |
| 277 | weight_scale = layer.weight_scale |
| 278 | _, K = weight.shape |
| 279 | output_size = layer._mxfp8_trtllm_output_size |
| 280 | input_shape = x.shape |
| 281 | input_2d = x.view(-1, K) |
| 282 | |
| 283 | input_mxfp8, input_scale = vllm_flashinfer.flashinfer_mxfp8_quantize_8x4( |
| 284 | input_2d |
| 285 | ) |
| 286 | output = vllm_flashinfer.mm_mxfp8( |
| 287 | input_mxfp8, |
| 288 | weight.t(), |
| 289 | input_scale, |
| 290 | weight_scale, |
| 291 | out_dtype=x.dtype, |
| 292 | backend="trtllm", |
| 293 | use_8x4_sf_layout=True, |
| 294 | ) |
| 295 | if output.shape[-1] != output_size: |
| 296 | output = output[:, :output_size].contiguous() |
| 297 | |
| 298 | if bias is not None: |
| 299 | output = output + bias |
| 300 | |
| 301 | output_shape = (*input_shape[:-1], output_size) |
| 302 | return output.view(output_shape) |