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