| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | from typing import ClassVar |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | import vllm.envs as envs |
| 10 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 11 | per_token_group_quant_fp8, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 14 | GroupShape, |
| 15 | QuantKey, |
| 16 | kFp8StaticTensorSym, |
| 17 | ) |
| 18 | from vllm.platforms import current_platform |
| 19 | from vllm.utils.flashinfer import ( |
| 20 | flashinfer_fp8_blockscale_gemm, |
| 21 | flashinfer_scaled_fp8_mm, |
| 22 | has_flashinfer, |
| 23 | is_flashinfer_fp8_blockscale_gemm_supported, |
| 24 | should_use_flashinfer_for_blockscale_fp8_gemm, |
| 25 | ) |
| 26 | from vllm.utils.torch_utils import direct_register_custom_op |
| 27 | |
| 28 | from .BlockScaledMMLinearKernel import ( |
| 29 | Fp8BlockScaledDynamicMMLinearKernel, |
| 30 | Fp8BlockScaledMMLinearKernel, |
| 31 | ) |
| 32 | from .deep_gemm import DeepGemmFp8BlockScaledMMKernel, fp8_gemm_nt |
| 33 | from .ScaledMMLinearKernel import ( |
| 34 | FP8ScaledMMLinearKernel, |
| 35 | FP8ScaledMMLinearLayerConfig, |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | class FlashInferFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): |
| 40 | @classmethod |
| 41 | def is_supported( |
| 42 | cls, compute_capability: int | None = None |
| 43 | ) -> tuple[bool, str | None]: |
| 44 | if not current_platform.is_cuda(): |
| 45 | return False, "requires CUDA." |
| 46 | |
| 47 | if not has_flashinfer(): |
| 48 | return False, "requires FlashInfer to be installed." |
| 49 | |
| 50 | if compute_capability is not None and compute_capability < 100: |
| 51 | return False, "requires compute capability 100 and above." |
| 52 | |
| 53 | return True, None |
| 54 | |
| 55 | @classmethod |
| 56 | def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]: |
| 57 | per_tensor_activation_scales = ( |
| 58 | c.activation_quant_key.scale.group_shape.is_per_tensor() |
| 59 | ) |
| 60 | per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor() |
| 61 | |
| 62 | if not (per_tensor_activation_scales and per_tensor_weight_scales): |
| 63 | return False, "requires per tensor activation and weight scales." |
| 64 | |
| 65 | return True, None |
| 66 | |
| 67 | def input_quant_key(self) -> QuantKey | None: |
| 68 | if self.config.activation_quant_key == kFp8StaticTensorSym: |
| 69 | return kFp8StaticTensorSym |
| 70 | return None |
| 71 | |
| 72 | def apply_scaled_mm( |
| 73 | self, |
| 74 | *, |
| 75 | A: torch.Tensor, |
| 76 | B: torch.Tensor, |
| 77 | out_dtype: torch.dtype, |
| 78 | As: torch.Tensor, |
| 79 | Bs: torch.Tensor, |
| 80 | bias: torch.Tensor | None, |
| 81 | output_shape: list, |
| 82 | ) -> torch.Tensor: |
| 83 | return flashinfer_scaled_fp8_mm( |
| 84 | A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias |
| 85 | ) |
| 86 | |
| 87 | |
| 88 | class FlashInferFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): |
| 89 | # FlashInfer accepts BF16 input and handles FP8 conversion internally. |
| 90 | apply_input_quant: ClassVar[bool] = False |
| 91 | |
| 92 | def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None: |
| 93 | super().__init__(config) |
| 94 | |
| 95 | @classmethod |
| 96 | def can_implement(cls, config: FP8ScaledMMLinearLayerConfig): |
| 97 | can_implement_base, reason = super().can_implement(config) |
| 98 | if not can_implement_base: |
| 99 | return can_implement_base, reason |
| 100 | |
| 101 | act_quant_desc = config.activation_quant_key.scale |
| 102 | if act_quant_desc.group_shape != GroupShape(1, 128): |
| 103 | return ( |
| 104 | False, |
| 105 | "Supports only dynamic per token group activation " |
| 106 | "quantization with group_shape=(1,128).", |
| 107 | ) |
| 108 | |
| 109 | if not should_use_flashinfer_for_blockscale_fp8_gemm( |
| 110 | is_flashinfer_fp8_blockscale_gemm_supported(), |
| 111 | config.out_dtype, |
| 112 | config.input_dtype, |
| 113 | config.weight_quant_key.dtype, |
| 114 | config.weight_shape, |
| 115 | ): |
| 116 | return ( |
| 117 | False, |
| 118 | "The provided metadata is not supported.", |
| 119 | ) |
| 120 | |
| 121 | return True, None |
| 122 | |
| 123 | @classmethod |
| 124 | def is_supported(cls, compute_capability=None): |
| 125 | if not current_platform.is_cuda(): |
| 126 | return False, "only cuda devices are supported." |
| 127 | |
| 128 | if not is_flashinfer_fp8_blockscale_gemm_supported(): |
| 129 | return False, "FlashInfer block-scale FP8 GEMM is not available." |
| 130 | |
| 131 | return True, None |
| 132 | |
| 133 | def apply_block_scaled_mm( |
| 134 | self, |
| 135 | A: torch.Tensor, |
| 136 | B: torch.Tensor, |
| 137 | As: torch.Tensor, |
| 138 | Bs: torch.Tensor, |
| 139 | ) -> torch.Tensor: |
| 140 | # A is BF16 — FlashInfer handles FP8 conversion internally. |
| 141 | # As is a placeholder (apply_input_quant=False) and is not used here. |
| 142 | return torch.ops.vllm.flashinfer_fp8_blockscale_gemm( |
| 143 | A, # BF16 input |
| 144 | B, # FP8 weight |
| 145 | Bs, # Weight scales |
| 146 | ) |
| 147 | |
| 148 | |
| 149 | class FlashInferFp8DeepGEMMDynamicBlockScaledKernel( |
| 150 | Fp8BlockScaledDynamicMMLinearKernel |
| 151 | ): |
| 152 | """Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM. |
| 153 | |
| 154 | Dispatches between two kernels based on input batch size: |
| 155 | - Small batches (M < 32): FlashInfer's swapAB trick for better utilisation. |
| 156 | - Large batches (M >= 32): DeepGEMM for peak throughput. |
| 157 | |
| 158 | apply_input_quant is False because FlashInfer accepts BF16 input and |
| 159 | handles FP8 conversion internally. The DeepGEMM branch therefore |
| 160 | quantises BF16→FP8 inside apply_mm via a closure before dispatching to |
| 161 | the DeepGEMM kernel — keeping both branches compatible with the single |
| 162 | BF16 tensor operand list passed by torch.cond. |
| 163 | """ |
| 164 | |
| 165 | base_type: ClassVar[type[FlashInferFp8BlockScaledMMKernel]] = ( |
| 166 | FlashInferFp8BlockScaledMMKernel |
| 167 | ) |
| 168 | fallback_type: ClassVar[type[DeepGemmFp8BlockScaledMMKernel]] = ( |
| 169 | DeepGemmFp8BlockScaledMMKernel |
| 170 | ) |
| 171 | apply_input_quant: ClassVar[bool] = False |
| 172 | |
| 173 | def __init__(self, config: FP8ScaledMMLinearLayerConfig): |
| 174 | super().__init__(config) |
| 175 | self.base: FlashInferFp8BlockScaledMMKernel |
| 176 | self.fallback: DeepGemmFp8BlockScaledMMKernel |
| 177 | |
| 178 | def process_weights_after_loading(self, layer: torch.nn.Module): |
| 179 | # DeepGEMM need post-processing; both kernels share the same |
| 180 | # parameter tensor layout so processing once is sufficient. |
| 181 | self.fallback.process_weights_after_loading(layer) |
| 182 | |
| 183 | def apply_block_scaled_mm( |
| 184 | self, |
| 185 | A: torch.Tensor, |
| 186 | B: torch.Tensor, |
| 187 | As: torch.Tensor, |
| 188 | Bs: torch.Tensor, |
| 189 | ) -> torch.Tensor: |
| 190 | group_size = self.weight_group_shape.col |
| 191 | use_deep_gemm_e8m0 = self.fallback.use_deep_gemm_e8m0 |
| 192 | |
| 193 | return torch.ops.vllm.dynamic_flashinfer_deepgemm_blockscale_gemm( |
| 194 | A, B, Bs, group_size, use_deep_gemm_e8m0 |
| 195 | ) |
| 196 | |
| 197 | |
| 198 | def _flashinfer_fp8_blockscale_gemm_impl( |
| 199 | input: torch.Tensor, |
| 200 | weight: torch.Tensor, |
| 201 | weight_scale: torch.Tensor, |
| 202 | ) -> torch.Tensor: |
| 203 | return flashinfer_fp8_blockscale_gemm( |
| 204 | input=input, |
| 205 | weight=weight, |
| 206 | weight_scale=weight_scale, |
| 207 | out_dtype=torch.bfloat16, |
| 208 | ) |
| 209 | |
| 210 | |
| 211 | def _flashinfer_fp8_blockscale_gemm_fake( |
| 212 | input: torch.Tensor, |
| 213 | weight: torch.Tensor, |
| 214 | weight_scale: torch.Tensor, |
| 215 | ) -> torch.Tensor: |
| 216 | """Required fake/meta implementation for torch.compile graph tracing.""" |
| 217 | return torch.empty( |
| 218 | input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device |
| 219 | ) |
| 220 | |
| 221 | |
| 222 | direct_register_custom_op( |
| 223 | "flashinfer_fp8_blockscale_gemm", |
| 224 | _flashinfer_fp8_blockscale_gemm_impl, |
| 225 | fake_impl=_flashinfer_fp8_blockscale_gemm_fake, |
| 226 | ) |
| 227 | |
| 228 | |
| 229 | def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl( |
| 230 | input: torch.Tensor, |
| 231 | weight: torch.Tensor, |
| 232 | weight_scale: torch.Tensor, |
| 233 | group_size: int, |
| 234 | use_deep_gemm_e8m0: bool, |
| 235 | ) -> torch.Tensor: |
| 236 | """Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection. |
| 237 | |
| 238 | This function switches between two optimized kernels based on the input batch size: |
| 239 | - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization. |
| 240 | - For larger batches (M >= 32): Uses the official DeepGEMM kernel. |
| 241 | |
| 242 | The conditional logic must use torch.cond() instead of a simple if-else statement |
| 243 | to maintain compatibility with torch.compile graph compilation. |
| 244 | |
| 245 | This batch-size-dependent selection is essential for maintaining model accuracy. |
| 246 | Benchmarks on GSM8K show a significant accuracy gap (88% vs 95%) for DeepSeek-V3.1 |
| 247 | when using FlashInfer's DeepGEMM on M>=32. The M < 32 strategy fixes the accuracy |
| 248 | drop. |
| 249 | |
| 250 | Args: |
| 251 | input: Input tensor of shape (batch_size, input_dim) in FP8 format |
| 252 | weight: Weight tensor of shape (output_dim, input_dim) in FP8 format |
| 253 | weight_scale: Scale factors for weight quantization (per-group) |
| 254 | group_size: Quantization group size for the weight tensor |
| 255 | use_deep_gemm_e8m0: Whether to use the E8M0 format in DeepGEMM quantization |
| 256 | |
| 257 | Returns: |
| 258 | Output tensor of shape (batch_size, output_dim) in bfloat16 format |
| 259 | |
| 260 | """ |
| 261 | |
| 262 | def run_flashinfer_deepgemm_swapAB( |
| 263 | input: torch.Tensor, |
| 264 | weight: torch.Tensor, |
| 265 | weight_scale: torch.Tensor, |
| 266 | ) -> torch.Tensor: |
| 267 | return flashinfer_fp8_blockscale_gemm( |
| 268 | input=input, |
| 269 | weight=weight, |
| 270 | weight_scale=weight_scale, |
| 271 | out_dtype=torch.bfloat16, |
| 272 | ) |
| 273 | |
| 274 | def run_deepgemm( |
| 275 | input: torch.Tensor, |
| 276 | weight: torch.Tensor, |
| 277 | weight_scale: torch.Tensor, |
| 278 | ) -> torch.Tensor: |
| 279 | q_input, input_scale = per_token_group_quant_fp8( |
| 280 | input, |
| 281 | group_size=group_size, |
| 282 | column_major_scales=True, |
| 283 | use_ue8m0=use_deep_gemm_e8m0, |
| 284 | ) |
| 285 | output = torch.empty( |
| 286 | (q_input.shape[0], weight.shape[0]), |
| 287 | dtype=torch.bfloat16, |
| 288 | device=q_input.device, |
| 289 | ) |
| 290 | fp8_gemm_nt( |
| 291 | (q_input, input_scale), |
| 292 | (weight, weight_scale), |
| 293 | output, |
| 294 | is_deep_gemm_e8m0_used=use_deep_gemm_e8m0, |
| 295 | ) |
| 296 | return output |
| 297 | |
| 298 | if envs.VLLM_BATCH_INVARIANT: |
| 299 | return run_deepgemm(input, weight, weight_scale) |
| 300 | |
| 301 | condition = input.shape[0] < 32 |
| 302 | |
| 303 | # PyTorch's torch.compile cannot handle input-dependent control flow in standard |
| 304 | # Python conditionals. torch.cond() explicitly registers both code paths in the |
| 305 | # computation graph, allowing torch.compile to capture both branches. |
| 306 | # without torch.cond, the M < 32 condition won't be able to be captured by torch |
| 307 | # compile |
| 308 | return torch.cond( |
| 309 | condition, |
| 310 | run_flashinfer_deepgemm_swapAB, |
| 311 | run_deepgemm, |
| 312 | (input, weight, weight_scale), |
| 313 | ) |
| 314 | |
| 315 | |
| 316 | def _dynamic_flashinfer_deepgemm_blockscale_gemm_fake( |
| 317 | input: torch.Tensor, |
| 318 | weight: torch.Tensor, |
| 319 | weight_scale: torch.Tensor, |
| 320 | group_size: int, |
| 321 | use_deep_gemm_e8m0: bool, |
| 322 | ) -> torch.Tensor: |
| 323 | """Required fake/meta implementation for torch.compile graph tracing.""" |
| 324 | return torch.empty( |
| 325 | input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device |
| 326 | ) |
| 327 | |
| 328 | |
| 329 | direct_register_custom_op( |
| 330 | "dynamic_flashinfer_deepgemm_blockscale_gemm", |
| 331 | _dynamic_flashinfer_deepgemm_blockscale_gemm_impl, |
| 332 | fake_impl=_dynamic_flashinfer_deepgemm_blockscale_gemm_fake, |
| 333 | ) |