| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from __future__ import annotations |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 9 | MXFP8_BLOCK_SIZE, |
| 10 | MXFP8_SCALE_DTYPE, |
| 11 | MXFP8_VALUE_DTYPE, |
| 12 | ) |
| 13 | from vllm.model_executor.utils import replace_parameter |
| 14 | from vllm.platforms import current_platform |
| 15 | from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage |
| 16 | from vllm.utils.b12x import ( |
| 17 | get_b12x_mxfp8_linear as _import_b12x_mxfp8, |
| 18 | ) |
| 19 | from vllm.utils.torch_utils import current_stream |
| 20 | |
| 21 | from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig |
| 22 | |
| 23 | |
| 24 | def _apply_b12x_mxfp8_packed_linear( |
| 25 | layer: torch.nn.Module, |
| 26 | x: torch.Tensor, |
| 27 | bias: torch.Tensor | None, |
| 28 | ) -> torch.Tensor: |
| 29 | packed_weight = layer.b12x_mxfp8_packed_weight |
| 30 | |
| 31 | input_2d = x.reshape(-1, x.shape[-1]).contiguous() |
| 32 | output_shape = [*x.shape[:-1], int(packed_weight.out_features)] |
| 33 | |
| 34 | mxfp8 = _import_b12x_mxfp8() |
| 35 | assert mxfp8 is not None |
| 36 | output = mxfp8.mm( |
| 37 | input_2d, |
| 38 | packed_weight, |
| 39 | bias=bias, |
| 40 | expected_m=max(1, int(input_2d.shape[0])), |
| 41 | ) |
| 42 | return output.view(*output_shape) |
| 43 | |
| 44 | |
| 45 | class B12xMxfp8LinearKernel(Mxfp8LinearKernel): |
| 46 | """ModelOpt MXFP8 linear through the native b12x SM120 dense GEMM path.""" |
| 47 | |
| 48 | @classmethod |
| 49 | def is_supported( |
| 50 | cls, |
| 51 | compute_capability: int | None = None, |
| 52 | ) -> tuple[bool, str | None]: |
| 53 | del compute_capability |
| 54 | if not current_platform.is_cuda(): |
| 55 | return False, "b12x MXFP8 kernels are only available on CUDA" |
| 56 | if not current_platform.is_device_capability_family(120): |
| 57 | return False, "b12x MXFP8 kernels require a Blackwell 12x device" |
| 58 | mxfp8 = _import_b12x_mxfp8() |
| 59 | if mxfp8 is None: |
| 60 | return False, "Install the B12X backend with `pip install vllm[b12x]`" |
| 61 | if not mxfp8.is_supported(): |
| 62 | return False, "b12x.gemm.mxfp8_linear is not supported" |
| 63 | return True, None |
| 64 | |
| 65 | @classmethod |
| 66 | def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: |
| 67 | del c |
| 68 | return True, None |
| 69 | |
| 70 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 71 | weight = layer.weight.data |
| 72 | assert weight.dtype == MXFP8_VALUE_DTYPE, ( |
| 73 | f"b12x MXFP8 requires {MXFP8_VALUE_DTYPE}, got {weight.dtype}" |
| 74 | ) |
| 75 | assert weight.ndim == 2, f"b12x MXFP8 weight must be 2D, got {weight.ndim}D" |
| 76 | assert hasattr(layer, "weight_scale"), "b12x MXFP8 linear requires weight_scale" |
| 77 | |
| 78 | out_features, in_features = map(int, weight.shape) |
| 79 | assert in_features % MXFP8_BLOCK_SIZE == 0, ( |
| 80 | "b12x MXFP8 requires input features divisible by " |
| 81 | f"{MXFP8_BLOCK_SIZE}, got {in_features}" |
| 82 | ) |
| 83 | weight_scale = layer.weight_scale.data |
| 84 | assert weight_scale.dtype == MXFP8_SCALE_DTYPE, ( |
| 85 | f"b12x MXFP8 requires {MXFP8_SCALE_DTYPE} weight_scale, " |
| 86 | f"got {weight_scale.dtype}" |
| 87 | ) |
| 88 | assert weight_scale.ndim == 2, ( |
| 89 | f"b12x MXFP8 weight_scale must be 2D, got {weight_scale.ndim}D" |
| 90 | ) |
| 91 | |
| 92 | mxfp8 = _import_b12x_mxfp8() |
| 93 | assert mxfp8 is not None |
| 94 | scale_k = in_features // MXFP8_BLOCK_SIZE |
| 95 | packed_weight = mxfp8.pack_weight( |
| 96 | weight[:out_features, :in_features].detach(), |
| 97 | weight_scale[:out_features, :scale_k].detach(), |
| 98 | ) |
| 99 | layer.b12x_mxfp8_packed_weight = reuse_packed_weight_storage( |
| 100 | getattr(layer, "b12x_mxfp8_packed_weight", None), |
| 101 | packed_weight, |
| 102 | ) |
| 103 | replace_parameter(layer, "weight", weight.new_empty((0,))) |
| 104 | replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,))) |
| 105 | layer.b12x_warmup_provider = self |
| 106 | |
| 107 | def get_b12x_warmup_unit( |
| 108 | self, |
| 109 | layer: torch.nn.Module, |
| 110 | token_counts: tuple[int, ...], |
| 111 | output_dtype: torch.dtype, |
| 112 | ) -> B12xWarmupUnit: |
| 113 | packed_weight = layer.b12x_mxfp8_packed_weight |
| 114 | device = torch.device(packed_weight.weight.values.device) |
| 115 | |
| 116 | def compile() -> None: |
| 117 | mxfp8 = _import_b12x_mxfp8() |
| 118 | assert mxfp8 is not None |
| 119 | for tokens in token_counts: |
| 120 | source = torch.zeros( |
| 121 | (tokens, int(packed_weight.in_features)), |
| 122 | dtype=output_dtype, |
| 123 | device=device, |
| 124 | ) |
| 125 | mxfp8.mm( |
| 126 | source, |
| 127 | packed_weight, |
| 128 | expected_m=max(1, int(tokens)), |
| 129 | stream=current_stream().cuda_stream, |
| 130 | ) |
| 131 | |
| 132 | return B12xWarmupUnit( |
| 133 | name="MXFP8", |
| 134 | key=( |
| 135 | type(self), |
| 136 | device, |
| 137 | int(packed_weight.in_features), |
| 138 | int(packed_weight.padded_in_features), |
| 139 | int(packed_weight.out_features), |
| 140 | output_dtype, |
| 141 | ), |
| 142 | compile=compile, |
| 143 | ) |
| 144 | |
| 145 | def apply_weights( |
| 146 | self, |
| 147 | layer: torch.nn.Module, |
| 148 | x: torch.Tensor, |
| 149 | bias: torch.Tensor | None = None, |
| 150 | ) -> torch.Tensor: |
| 151 | return _apply_b12x_mxfp8_packed_linear(layer, x, bias) |