| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Producer side of the QuantizedActivation contract for activation layers. |
| 4 | |
| 5 | Given an activation module and the downstream linear it feeds, fuse the |
| 6 | activation with that linear's input quantization into a single kernel when the |
| 7 | linear advertises a consumable input quantization key (see quant_activation.py). |
| 8 | Falls back to the plain activation when nothing matches, so a model forward can |
| 9 | always call maybe_fused_act_quant unconditionally. |
| 10 | |
| 11 | This is the manual-fusion counterpart to ActivationQuantFusionPass: when fusion |
| 12 | fires here, the activation and quantization are already consumed, so a compiler |
| 13 | pass cannot fuse the same boundary again. |
| 14 | """ |
| 15 | |
| 16 | from collections.abc import Callable |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | from vllm.model_executor.layers.activation import ReLUSquaredActivation, SiluAndMul |
| 21 | from vllm.model_executor.layers.fusion.quant_activation import ( |
| 22 | QuantizedActivation, |
| 23 | get_input_quant_key, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.fusion.relu2_fp8_quant import ( |
| 26 | relu_squared_static_fp8_quant, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.linear import LinearBase |
| 29 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 30 | QuantKey, |
| 31 | kFp8Dynamic128Sym, |
| 32 | kFp8StaticTensorSym, |
| 33 | kNvfp4Dynamic, |
| 34 | ) |
| 35 | from vllm.platforms import current_platform |
| 36 | from vllm.utils.math_utils import round_up |
| 37 | |
| 38 | FP8_DTYPE = current_platform.fp8_dtype() |
| 39 | FP4_DTYPE = torch.uint8 |
| 40 | |
| 41 | |
| 42 | def _silu_and_mul_fp8_static( |
| 43 | x: torch.Tensor, linear: LinearBase |
| 44 | ) -> QuantizedActivation: |
| 45 | """SiluAndMul + FP8 static per-tensor quantization.""" |
| 46 | d = x.shape[-1] // 2 |
| 47 | out_shape = x.shape[:-1] + (d,) |
| 48 | result = torch.empty(out_shape, dtype=FP8_DTYPE, device=x.device) |
| 49 | # TODO(mgoin): read the consumer scale via the contract instead of reaching |
| 50 | # into the kernel-specific input_scale attribute. |
| 51 | scale = linear.input_scale |
| 52 | torch.ops._C.silu_and_mul_quant(result, x, scale) |
| 53 | return QuantizedActivation( |
| 54 | data=result, |
| 55 | scale=scale, |
| 56 | orig_dtype=x.dtype, |
| 57 | orig_shape=out_shape, |
| 58 | quant_key=kFp8StaticTensorSym, |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | def _silu_and_mul_fp8_dynamic_block( |
| 63 | x: torch.Tensor, linear: LinearBase, group_size: int, quant_key: QuantKey |
| 64 | ) -> QuantizedActivation: |
| 65 | """SiluAndMul + FP8 dynamic per-block quantization.""" |
| 66 | assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" |
| 67 | |
| 68 | d = x.shape[-1] // 2 |
| 69 | out_shape = x.shape[:-1] + (d,) |
| 70 | num_tokens = x.shape[0] |
| 71 | num_groups = d // group_size |
| 72 | |
| 73 | result = torch.empty((num_tokens, d), dtype=FP8_DTYPE, device=x.device) |
| 74 | scales = torch.empty((num_tokens, num_groups), dtype=torch.float32, device=x.device) |
| 75 | |
| 76 | torch.ops._C.silu_and_mul_per_block_quant( |
| 77 | out=result, |
| 78 | input=x, |
| 79 | scales=scales, |
| 80 | group_size=group_size, |
| 81 | scale_ub=None, |
| 82 | is_scale_transposed=False, |
| 83 | ) |
| 84 | |
| 85 | return QuantizedActivation( |
| 86 | data=result.view(out_shape), |
| 87 | scale=scales.view(out_shape[:-1] + (num_groups,)), |
| 88 | orig_dtype=x.dtype, |
| 89 | orig_shape=out_shape, |
| 90 | quant_key=quant_key, |
| 91 | ) |
| 92 | |
| 93 | |
| 94 | def _silu_and_mul_fp8_dynamic_128( |
| 95 | x: torch.Tensor, linear: LinearBase |
| 96 | ) -> QuantizedActivation: |
| 97 | """SiluAndMul + FP8 dynamic per-group (group=128) quantization.""" |
| 98 | return _silu_and_mul_fp8_dynamic_block(x, linear, 128, kFp8Dynamic128Sym) |
| 99 | |
| 100 | |
| 101 | def _silu_and_mul_nvfp4_dynamic( |
| 102 | x: torch.Tensor, linear: LinearBase |
| 103 | ) -> QuantizedActivation: |
| 104 | """SiluAndMul + NVFP4 dynamic quantization.""" |
| 105 | assert x.ndim == 2, f"Input must be 2D [batch, hidden*2], got {x.shape}" |
| 106 | |
| 107 | d = x.shape[-1] // 2 |
| 108 | out_shape = x.shape[:-1] + (d,) |
| 109 | num_tokens = x.shape[0] |
| 110 | |
| 111 | # NVFP4 packs 2 values into 1 byte |
| 112 | result = torch.empty((num_tokens, d // 2), dtype=FP4_DTYPE, device=x.device) |
| 113 | |
| 114 | # Block scale output: swizzled tensor-core layout |
| 115 | # [num_m_tiles, num_k_tiles, 32, 4, 4] of int32-packed FP8 scales, so the |
| 116 | # row/col extents must be padded to the 128x4 tile like |
| 117 | # scaled_fp4_quant's allocator (create_fp4_scale_tensor). Each group of 16 |
| 118 | # elements shares one FP8 scale. |
| 119 | rounded_m = round_up(num_tokens, 128) |
| 120 | rounded_n = round_up(d // 16, 4) |
| 121 | block_scale = torch.empty( |
| 122 | (rounded_m, rounded_n // 4), dtype=torch.int32, device=x.device |
| 123 | ).view(FP8_DTYPE) |
| 124 | |
| 125 | # The kernel folds the global scale into the block scales, and the |
| 126 | # consumer GEMM's alpha (= input_global_scale * weight_global_scale) |
| 127 | # divides it back out, so quantize with the reciprocal like the unfused |
| 128 | # scaled_fp4_quant path does. |
| 129 | input_global_scale_inv = getattr(linear, "input_global_scale_inv", None) |
| 130 | assert input_global_scale_inv is not None, ( |
| 131 | "input_global_scale_inv is required for NVFP4 quantization" |
| 132 | ) |
| 133 | |
| 134 | torch.ops._C.silu_and_mul_nvfp4_quant( |
| 135 | result, block_scale, x, input_global_scale_inv |
| 136 | ) |
| 137 | |
| 138 | return QuantizedActivation( |
| 139 | data=result.view(out_shape[:-1] + (d // 2,)), |
| 140 | scale=block_scale, |
| 141 | orig_dtype=x.dtype, |
| 142 | orig_shape=out_shape, |
| 143 | quant_key=kNvfp4Dynamic, |
| 144 | ) |
| 145 | |
| 146 | |
| 147 | # (activation module type, consumer input quantization key) -> fused producer. |
| 148 | # Add a row for each supported manual producer. |
| 149 | _FUSED_ACT_QUANT: dict[tuple[type, QuantKey], Callable] = { |
| 150 | (SiluAndMul, kFp8StaticTensorSym): _silu_and_mul_fp8_static, |
| 151 | } |
| 152 | |
| 153 | |
| 154 | def _relu_squared_static_fp8_quant_supported( |
| 155 | act_fn: torch.nn.Module, |
| 156 | x: torch.Tensor, |
| 157 | linear: LinearBase, |
| 158 | ) -> bool: |
| 159 | """Return whether the ReLU2 static-FP8 producer can consume this input.""" |
| 160 | scale = getattr(linear, "input_scale", None) |
| 161 | return ( |
| 162 | isinstance(act_fn, ReLUSquaredActivation) |
| 163 | and x.is_cuda |
| 164 | and x.dtype == torch.bfloat16 |
| 165 | and x.is_contiguous() |
| 166 | and isinstance(scale, torch.Tensor) |
| 167 | and scale.dtype == torch.float32 |
| 168 | and scale.device == x.device |
| 169 | and scale.numel() == 1 |
| 170 | ) |
| 171 | |
| 172 | |
| 173 | # Optional per-entry predicates preserve the unconditional fallback contract |
| 174 | # for producers whose supported inputs are narrower than their registry key. |
| 175 | _FUSED_ACT_QUANT_SUPPORT: dict[ |
| 176 | tuple[type, QuantKey], Callable[[torch.nn.Module, torch.Tensor, LinearBase], bool] |
| 177 | ] = {} |
| 178 | |
| 179 | # This path is validated on SM90 and newer CUDA devices. |
| 180 | if current_platform.is_cuda() and current_platform.has_device_capability(90): |
| 181 | _RELU2_STATIC_FP8_KEY = (ReLUSquaredActivation, kFp8StaticTensorSym) |
| 182 | _FUSED_ACT_QUANT[_RELU2_STATIC_FP8_KEY] = relu_squared_static_fp8_quant |
| 183 | _FUSED_ACT_QUANT_SUPPORT[_RELU2_STATIC_FP8_KEY] = ( |
| 184 | _relu_squared_static_fp8_quant_supported |
| 185 | ) |
| 186 | |
| 187 | # Add CUDA-specific entries for dynamic block quantization |
| 188 | if current_platform.is_cuda_alike(): |
| 189 | _FUSED_ACT_QUANT[(SiluAndMul, kFp8Dynamic128Sym)] = _silu_and_mul_fp8_dynamic_128 |
| 190 | |
| 191 | # Add NVFP4 if supported (requires SM100+) |
| 192 | if current_platform.is_cuda() and hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant"): |
| 193 | _FUSED_ACT_QUANT[(SiluAndMul, kNvfp4Dynamic)] = _silu_and_mul_nvfp4_dynamic |
| 194 | |
| 195 | |
| 196 | def maybe_fused_act_quant( |
| 197 | act_fn: torch.nn.Module, |
| 198 | x: torch.Tensor, |
| 199 | linear: LinearBase, |
| 200 | ) -> "torch.Tensor | QuantizedActivation": |
| 201 | """Apply act_fn, fusing the downstream linear's input quant when possible. |
| 202 | |
| 203 | Returns a QuantizedActivation when a fused kernel matches the activation and |
| 204 | the consumer's effective input quantization key, else the plain activation. |
| 205 | """ |
| 206 | key = get_input_quant_key(linear) |
| 207 | if key is not None: |
| 208 | registry_key = (type(act_fn), key) |
| 209 | producer = _FUSED_ACT_QUANT.get(registry_key) |
| 210 | support = _FUSED_ACT_QUANT_SUPPORT.get(registry_key) |
| 211 | if producer is not None and (support is None or support(act_fn, x, linear)): |
| 212 | return producer(x, linear) |
| 213 | return act_fn(x) |