| 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.quantization.utils.mxfp8_utils import ( |
| 8 | MXFP8_BLOCK_SIZE, |
| 9 | MXFP8_SCALE_DTYPE, |
| 10 | dequant_mxfp8_to_bf16, |
| 11 | ) |
| 12 | |
| 13 | from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig |
| 14 | |
| 15 | |
| 16 | class EmulationMxfp8LinearKernel(Mxfp8LinearKernel): |
| 17 | """Software emulation fallback for MXFP8 (dequant to BF16).""" |
| 18 | |
| 19 | supports_pre_processed_weights = True |
| 20 | |
| 21 | @classmethod |
| 22 | def is_supported( |
| 23 | cls, compute_capability: int | None = None |
| 24 | ) -> tuple[bool, str | None]: |
| 25 | return True, None |
| 26 | |
| 27 | @classmethod |
| 28 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 29 | return True, None |
| 30 | |
| 31 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 32 | weight = layer.weight.data # [N, K] |
| 33 | N, K = weight.shape |
| 34 | scale_k = K // MXFP8_BLOCK_SIZE |
| 35 | |
| 36 | weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous() |
| 37 | |
| 38 | # Dequantize MXFP8 -> BF16 ONCE here, at load time, so apply_weights runs |
| 39 | # a plain BF16 linear with no per-step dequant -- i.e. run as if from a |
| 40 | # BF16 checkpoint. The 1-byte MXFP8 weight is replaced by BF16 (2x its |
| 41 | # size, but linear weights are small vs the MoE experts); the tiny E8M0 |
| 42 | # scale is kept for the dtype/ndim asserts but is otherwise unused. |
| 43 | # Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the MXFP8 |
| 44 | # weight and dequant per-step in apply_weights instead. |
| 45 | import vllm.envs as envs |
| 46 | |
| 47 | if envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD: |
| 48 | weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale) |
| 49 | layer.weight = Parameter(weight.contiguous(), requires_grad=False) |
| 50 | layer.weight_scale = Parameter(weight_scale, requires_grad=False) |
| 51 | |
| 52 | def apply_weights( |
| 53 | self, |
| 54 | layer: torch.nn.Module, |
| 55 | x: torch.Tensor, |
| 56 | bias: torch.Tensor | None = None, |
| 57 | ) -> torch.Tensor: |
| 58 | weight = layer.weight |
| 59 | # Load-time dequant path: weights are already BF16/FP16 (>= 2-byte), so |
| 60 | # run a plain linear -- no per-step dequant. (MXFP8 weights are 1-byte.) |
| 61 | if weight.element_size() >= 2: |
| 62 | # F.linear requires x and weight share a dtype; .to() is a no-op when |
| 63 | # they already match (e.g. both BF16). |
| 64 | output = torch.nn.functional.linear(x, weight.to(x.dtype), bias) |
| 65 | return output.to(x.dtype) |
| 66 | |
| 67 | # Fallback: weights still in MXFP8 -- dequant on the fly (other archs / |
| 68 | # if a future caller skips the load-time conversion above). |
| 69 | weight_scale = layer.weight_scale |
| 70 | if weight_scale.dtype != MXFP8_SCALE_DTYPE: |
| 71 | raise ValueError( |
| 72 | f"Emulation backend requires {MXFP8_SCALE_DTYPE} " |
| 73 | f"weight_scale dtype, got {weight_scale.dtype}." |
| 74 | ) |
| 75 | if weight_scale.ndim != 2: |
| 76 | raise ValueError( |
| 77 | f"Emulation backend requires 2D weight_scale, " |
| 78 | f"got {weight_scale.ndim}D. " |
| 79 | f"Ensure process_weights_after_loading was called." |
| 80 | ) |
| 81 | |
| 82 | # Cast to x's dtype: dequant yields BF16, but F.linear needs both operands |
| 83 | # to match (e.g. an FP16 model). No-op when x is already BF16. |
| 84 | weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale).to(x.dtype) |
| 85 | output = torch.nn.functional.linear(x, weight_bf16, bias) |
| 86 | return output.to(x.dtype) |