| 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.quant_utils import ( |
| 9 | kMxfp4Dynamic, |
| 10 | ) |
| 11 | from vllm.model_executor.utils import replace_parameter |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils.b12x import B12xWarmupUnit |
| 14 | from vllm.utils.b12x import ( |
| 15 | get_b12x_blockscaled as _import_b12x_blockscaled, |
| 16 | ) |
| 17 | from vllm.utils.b12x import get_b12x_intrinsics as _import_b12x_intrinsics |
| 18 | |
| 19 | from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig |
| 20 | |
| 21 | |
| 22 | def _apply_b12x_mxfp4_linear( |
| 23 | x: torch.Tensor, |
| 24 | weight: torch.Tensor, |
| 25 | weight_scale_storage: torch.Tensor, |
| 26 | bias: torch.Tensor | None, |
| 27 | ) -> torch.Tensor: |
| 28 | from vllm.utils.flashinfer import flashinfer_mxfp4_quantize |
| 29 | |
| 30 | blockscaled = _import_b12x_blockscaled() |
| 31 | assert blockscaled is not None |
| 32 | |
| 33 | output_size = int(weight.shape[0]) |
| 34 | output_shape = [*x.shape[:-1], output_size] |
| 35 | x_2d = x.reshape(-1, x.shape[-1]).contiguous() |
| 36 | x_packed, x_scale_swizzled = flashinfer_mxfp4_quantize(x_2d, backend="cute-dsl") |
| 37 | output = blockscaled.mm_mxfp4( |
| 38 | x_packed, |
| 39 | x_scale_swizzled, |
| 40 | weight, |
| 41 | weight_scale_storage, |
| 42 | out_dtype=x.dtype, |
| 43 | ) |
| 44 | if bias is not None: |
| 45 | output = output + bias |
| 46 | return output.view(*output_shape) |
| 47 | |
| 48 | |
| 49 | class B12xMxFp4LinearKernel(MxFp4LinearKernel): |
| 50 | """MXFP4 linear through the native B12X SM120 dense GEMM.""" |
| 51 | |
| 52 | @classmethod |
| 53 | def is_supported( |
| 54 | cls, compute_capability: int | None = None |
| 55 | ) -> tuple[bool, str | None]: |
| 56 | del compute_capability |
| 57 | if not current_platform.is_cuda(): |
| 58 | return False, "B12X MXFP4 kernels are only available on CUDA" |
| 59 | if not current_platform.is_device_capability_family(120): |
| 60 | return False, "B12X MXFP4 kernels require a Blackwell 12x device" |
| 61 | blockscaled = _import_b12x_blockscaled() |
| 62 | if blockscaled is None or _import_b12x_intrinsics() is None: |
| 63 | return False, "Install the B12X backend with `pip install vllm[b12x]`" |
| 64 | if not blockscaled.is_supported(): |
| 65 | return False, "b12x native MXFP4 GEMM is not supported" |
| 66 | return True, None |
| 67 | |
| 68 | @classmethod |
| 69 | def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]: |
| 70 | if config.activation_quant_key != kMxfp4Dynamic: |
| 71 | return False, "B12X MXFP4 GEMM requires dynamic MXFP4 activations" |
| 72 | return True, None |
| 73 | |
| 74 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 75 | intrinsics = _import_b12x_intrinsics() |
| 76 | assert intrinsics is not None |
| 77 | replace_parameter( |
| 78 | layer, |
| 79 | "weight_scale", |
| 80 | intrinsics.swizzle_block_scale(layer.weight_scale.data), |
| 81 | ) |
| 82 | layer.b12x_warmup_provider = self |
| 83 | |
| 84 | def get_b12x_warmup_unit( |
| 85 | self, |
| 86 | layer: torch.nn.Module, |
| 87 | token_counts: tuple[int, ...], |
| 88 | output_dtype: torch.dtype, |
| 89 | ) -> B12xWarmupUnit: |
| 90 | weight = layer.weight |
| 91 | weight_scale = layer.weight_scale |
| 92 | n, packed_k = map(int, weight.shape) |
| 93 | k = packed_k * 2 |
| 94 | |
| 95 | def compile() -> None: |
| 96 | for tokens in token_counts: |
| 97 | source = torch.zeros( |
| 98 | (tokens, k), dtype=output_dtype, device=weight.device |
| 99 | ) |
| 100 | _apply_b12x_mxfp4_linear(source, weight, weight_scale, None) |
| 101 | |
| 102 | return B12xWarmupUnit( |
| 103 | name="MXFP4", |
| 104 | key=( |
| 105 | type(self), |
| 106 | weight.device, |
| 107 | n, |
| 108 | k, |
| 109 | weight.dtype, |
| 110 | weight_scale.dtype, |
| 111 | output_dtype, |
| 112 | ), |
| 113 | compile=compile, |
| 114 | ) |
| 115 | |
| 116 | def apply_weights( |
| 117 | self, |
| 118 | layer: torch.nn.Module, |
| 119 | x: torch.Tensor, |
| 120 | bias: torch.Tensor | None = None, |
| 121 | ) -> torch.Tensor: |
| 122 | return _apply_b12x_mxfp4_linear( |
| 123 | x, |
| 124 | layer.weight, |
| 125 | layer.weight_scale, |
| 126 | bias, |
| 127 | ) |
| 128 | |
| 129 | |
| 130 | __all__ = ["B12xMxFp4LinearKernel"] |