| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | import vllm._custom_ops as ops |
| 8 | from vllm.logger import init_logger |
| 9 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 10 | from vllm.model_executor.layers.quantization.utils.marlin_utils import ( |
| 11 | USE_FP32_REDUCE_DEFAULT, |
| 12 | get_marlin_input_dtype, |
| 13 | marlin_pad_dim, |
| 14 | marlin_pad_qweight, |
| 15 | marlin_pad_scales, |
| 16 | marlin_padded_nk, |
| 17 | marlin_permute_bias, |
| 18 | marlin_permute_scales, |
| 19 | marlin_quant_input, |
| 20 | marlin_repacked_nk, |
| 21 | marlin_unpad_output, |
| 22 | should_use_atomic_add_reduce, |
| 23 | ) |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.scalar_type import scalar_types |
| 26 | from vllm.utils.math_utils import round_up |
| 27 | |
| 28 | FP4_MARLIN_SUPPORTED_GROUP_SIZES = [16] |
| 29 | |
| 30 | logger = init_logger(__name__) |
| 31 | |
| 32 | |
| 33 | def is_fp4_marlin_supported(): |
| 34 | return current_platform.is_cuda() and current_platform.has_device_capability(75) |
| 35 | |
| 36 | |
| 37 | def _nvfp4_compute_scale_factor( |
| 38 | marlin_scales: torch.Tensor, |
| 39 | a_dtype: torch.dtype | None = None, |
| 40 | ) -> float: |
| 41 | """Compute the power-of-2 scale_factor needed so that all non-zero |
| 42 | values in marlin_scales * 2^7 are >= 2 after rescaling. |
| 43 | Returns a Python float (power of 2, >= 1.0).""" |
| 44 | # Since half has a smaller dynamic range compared to bfloat16, |
| 45 | # no rescaling is applied here if active dtype is half. |
| 46 | if a_dtype is not None and a_dtype == torch.half: |
| 47 | return 1.0 |
| 48 | |
| 49 | ws_float = marlin_scales.float() * (2**7) |
| 50 | nonzero_mask = ws_float > 0 |
| 51 | if nonzero_mask.any(): |
| 52 | max_val = ws_float[nonzero_mask].max() |
| 53 | if max_val < 448 * (2**7): |
| 54 | sf = (448 * (2**7) / max_val).log2().floor().exp2() |
| 55 | return sf.item() |
| 56 | return 1.0 |
| 57 | |
| 58 | |
| 59 | def nvfp4_marlin_process_scales( |
| 60 | marlin_scales: torch.Tensor, |
| 61 | scale_factor: float | None = None, |
| 62 | a_dtype: torch.dtype | None = None, |
| 63 | ) -> tuple[torch.Tensor, float]: |
| 64 | """Process NVFP4 weight scales into the special S0E5M3 format for Marlin. |
| 65 | |
| 66 | Args: |
| 67 | marlin_scales: Weight scales tensor in half precision, already |
| 68 | permuted for the Marlin kernel layout. |
| 69 | scale_factor: Optional power-of-2 rescaling factor. If None, the |
| 70 | factor is computed automatically so that every non-zero scale |
| 71 | satisfies ``scale * 2^7 >= 2`` (i.e., the MSB of the S0E5M3 |
| 72 | representation is always 1). When provided (e.g., for MoE |
| 73 | layers where all experts must share the same factor), the |
| 74 | given value is used directly. The caller is responsible for |
| 75 | dividing ``global_scale`` by the returned ``scale_factor`` to |
| 76 | preserve numerical correctness. |
| 77 | a_dtype: dtype of the activations the scales will be used with. None |
| 78 | keeps the dtype of ``marlin_scales``. |
| 79 | |
| 80 | Returns: |
| 81 | A tuple of (processed_scales, scale_factor). |
| 82 | |
| 83 | """ |
| 84 | if not (marlin_scales >= 0).all(): |
| 85 | logger.warning_once( |
| 86 | "NVFP4 Marlin assumes the scales to be >=0, but has encountered " |
| 87 | "negative scales. Accuracy will likely be degraded. This is " |
| 88 | "because it changes the scales from FP8-S1E4M3 to a special " |
| 89 | "FP8-S0E5M3 format to speedup the dequantization." |
| 90 | ) |
| 91 | |
| 92 | # convert to half first, we would convert to fp8 later |
| 93 | marlin_scales = marlin_scales.to(torch.half) |
| 94 | |
| 95 | # fit the layout of fp8 dequantization |
| 96 | marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view( |
| 97 | marlin_scales.size(0), -1 |
| 98 | ) |
| 99 | |
| 100 | # We assume that weight_scale (FP8-S1E4M3) is always greater |
| 101 | # than or equal to 0. So we can convert |
| 102 | # (weight_scale * (2 ** 7) to a special FP8-S0E5M3 format. |
| 103 | # After multiplying by 2 ** 7, the top bit of FP8-S0E5M3 would always be 1 |
| 104 | # when weight_scale > 0. This allows us to have an exponent bias |
| 105 | # closer to zero after dequantization. |
| 106 | |
| 107 | # Rescale weight_scale so that all non-zero values have MSB=1 |
| 108 | # after multiplying by 2^7 (i.e., weight_scale * 2^7 >= 2). |
| 109 | # This is needed for models whose E4M3 scales were not normalized |
| 110 | # to fully utilize the E4M3 dynamic range (e.g., global_scale=1). |
| 111 | # The caller must compensate by dividing global_scale by scale_factor. |
| 112 | if scale_factor is None: |
| 113 | scale_factor = _nvfp4_compute_scale_factor(marlin_scales, a_dtype) |
| 114 | if scale_factor > 1.0: |
| 115 | marlin_scales = (marlin_scales.float() * scale_factor).to(torch.half) |
| 116 | |
| 117 | marlin_scales = marlin_scales * (2**7) |
| 118 | marlin_scales[marlin_scales < 2] = 0 |
| 119 | marlin_scales = marlin_scales.view(torch.int16) << 1 |
| 120 | marlin_scales = marlin_scales.view(torch.float8_e4m3fn) |
| 121 | marlin_scales = marlin_scales[:, 1::2].contiguous() |
| 122 | |
| 123 | return marlin_scales, scale_factor |
| 124 | |
| 125 | |
| 126 | def mxfp4_marlin_process_scales(marlin_scales, input_dtype=None): |
| 127 | # fit the layout of fp8 dequantization |
| 128 | if input_dtype is None or input_dtype.itemsize == 2: |
| 129 | marlin_scales = marlin_scales.view(-1, 4)[:, [0, 2, 1, 3]].view( |
| 130 | marlin_scales.size(0), -1 |
| 131 | ) |
| 132 | |
| 133 | marlin_scales = marlin_scales.to(torch.float8_e8m0fnu) |
| 134 | if input_dtype == torch.float8_e4m3fn: |
| 135 | marlin_scales = marlin_scales.view(torch.uint8) |
| 136 | assert marlin_scales.max() <= 249 |
| 137 | # exponent_bias (fp4->fp8) = 2 ** 3 - 2 ** 1 = 6 |
| 138 | marlin_scales = marlin_scales + 6 |
| 139 | marlin_scales = marlin_scales.view(torch.float8_e8m0fnu) |
| 140 | return marlin_scales |
| 141 | |
| 142 | |
| 143 | def nvfp4_marlin_process_global_scale(global_scale, a_dtype: torch.dtype | None = None): |
| 144 | if a_dtype is None: |
| 145 | a_dtype = global_scale.dtype |
| 146 | assert a_dtype in [torch.half, torch.bfloat16] |
| 147 | fp4_exponent = 2 |
| 148 | if a_dtype == torch.half: |
| 149 | target_exponent = 5 |
| 150 | elif a_dtype == torch.bfloat16: |
| 151 | target_exponent = 8 |
| 152 | # exponent_bias_fp16 = 2 ** 4 - 2 ** 1 = 14 |
| 153 | # exponent_bias_bf16 = 2 ** 7 - 2 ** 1 = 126 |
| 154 | exponent_bias = 2 ** (target_exponent - 1) - 2 ** (fp4_exponent - 1) |
| 155 | return global_scale * (2.0 ** (exponent_bias - 7)) |
| 156 | |
| 157 | |
| 158 | def apply_fp4_marlin_linear( |
| 159 | input: torch.Tensor, |
| 160 | weight: torch.Tensor, |
| 161 | weight_scale: torch.Tensor, |
| 162 | weight_global_scale: torch.Tensor | None, |
| 163 | workspace: torch.Tensor | None, |
| 164 | size_n: int, |
| 165 | size_k: int, |
| 166 | bias: torch.Tensor | None = None, |
| 167 | input_dtype: torch.dtype | None = None, |
| 168 | use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, |
| 169 | ) -> torch.Tensor: |
| 170 | # For GPUs that lack FP4 hardware support, we can leverage the |
| 171 | # Marlin kernel for fast weight-only FP4 quantization |
| 172 | |
| 173 | reshaped_x = input.reshape(-1, input.shape[-1]) |
| 174 | out_shape = input.shape[:-1] + (size_n,) |
| 175 | |
| 176 | padded_n, padded_k = marlin_repacked_nk(weight, num_bits=4) |
| 177 | reshaped_x = marlin_pad_dim(reshaped_x, size_k, padded_k) |
| 178 | |
| 179 | use_atomic_add = should_use_atomic_add_reduce( |
| 180 | m=reshaped_x.size(0), |
| 181 | n=padded_n, |
| 182 | k=padded_k, |
| 183 | device=input.device, |
| 184 | dtype=input.dtype, |
| 185 | ) |
| 186 | |
| 187 | inputs = reshaped_x |
| 188 | a_scales = None |
| 189 | is_nvfp4 = weight_global_scale is not None |
| 190 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 191 | if is_nvfp4: |
| 192 | raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.") |
| 193 | elif input_dtype != torch.float8_e4m3fn: |
| 194 | raise RuntimeError("MXFP4 weight + INT8 activation is not supported.") |
| 195 | |
| 196 | inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) |
| 197 | |
| 198 | output = torch.ops.vllm.marlin_gemm( |
| 199 | a=inputs, |
| 200 | b_q_weight=weight, |
| 201 | b_bias=bias, |
| 202 | b_scales=weight_scale, |
| 203 | a_scales=a_scales, |
| 204 | global_scale=weight_global_scale, |
| 205 | b_zeros=None, |
| 206 | workspace=workspace, |
| 207 | b_q_type_id=scalar_types.float4_e2m1f.id, |
| 208 | size_m=reshaped_x.size(0), |
| 209 | size_n=padded_n, |
| 210 | size_k=padded_k, |
| 211 | use_atomic_add=use_atomic_add, |
| 212 | use_fp32_reduce=use_fp32_reduce, |
| 213 | ) |
| 214 | |
| 215 | output = marlin_unpad_output(output, size_n, padded_n) |
| 216 | return output.reshape(out_shape) |
| 217 | |
| 218 | |
| 219 | def prepare_fp4_layer_for_marlin( |
| 220 | layer: torch.nn.Module, input_dtype: torch.dtype | None = None |
| 221 | ) -> None: |
| 222 | is_nvfp4 = hasattr(layer, "weight_global_scale") |
| 223 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 224 | if is_nvfp4: |
| 225 | raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.") |
| 226 | elif input_dtype != torch.float8_e4m3fn: |
| 227 | raise RuntimeError("MXFP4 weight + INT8 activation is not supported.") |
| 228 | |
| 229 | group_size = 16 if is_nvfp4 else 32 |
| 230 | |
| 231 | part_size_n = layer.output_size_per_partition |
| 232 | part_size_k = layer.input_size_per_partition |
| 233 | padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size) |
| 234 | param_dtype = layer.params_dtype |
| 235 | |
| 236 | assert layer.weight.shape == (part_size_n, part_size_k // 2) |
| 237 | |
| 238 | # WEIGHT |
| 239 | # Repack weights to marlin format |
| 240 | qweight = layer.weight.view(torch.int32).T.contiguous() |
| 241 | qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k) |
| 242 | |
| 243 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 244 | marlin_qweight = ops.gptq_marlin_repack( |
| 245 | b_q_weight=qweight, |
| 246 | size_k=padded_k, |
| 247 | size_n=padded_n, |
| 248 | num_bits=4, |
| 249 | is_a_8bit=is_a_8bit, |
| 250 | ) |
| 251 | layer.weight = torch.nn.Parameter(marlin_qweight, requires_grad=False) |
| 252 | |
| 253 | # WEIGHT SCALES |
| 254 | # Permute scales |
| 255 | weight_scale = layer.weight_scale.T.contiguous() |
| 256 | |
| 257 | if not is_nvfp4: |
| 258 | weight_scale = weight_scale.view(torch.float8_e8m0fnu) |
| 259 | |
| 260 | weight_scale = weight_scale.to(param_dtype) |
| 261 | weight_scale = marlin_pad_scales( |
| 262 | weight_scale, part_size_n, part_size_k, padded_n, padded_k, group_size |
| 263 | ) |
| 264 | weight_scale = marlin_permute_scales( |
| 265 | s=weight_scale, |
| 266 | size_k=padded_k, |
| 267 | size_n=padded_n, |
| 268 | group_size=group_size, |
| 269 | is_a_8bit=is_a_8bit, |
| 270 | ) |
| 271 | |
| 272 | if is_nvfp4: |
| 273 | weight_scale, scale_factor = nvfp4_marlin_process_scales( |
| 274 | weight_scale, a_dtype=param_dtype |
| 275 | ) |
| 276 | layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) |
| 277 | |
| 278 | weight_global_scale = layer.weight_global_scale.to(torch.float32) |
| 279 | weight_global_scale = nvfp4_marlin_process_global_scale( |
| 280 | weight_global_scale, param_dtype |
| 281 | ) |
| 282 | weight_global_scale = weight_global_scale / scale_factor |
| 283 | layer.weight_global_scale = torch.nn.Parameter( |
| 284 | weight_global_scale, requires_grad=False |
| 285 | ) |
| 286 | else: |
| 287 | weight_scale = mxfp4_marlin_process_scales( |
| 288 | weight_scale, input_dtype=input_dtype |
| 289 | ) |
| 290 | layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) |
| 291 | |
| 292 | if hasattr(layer, "bias") and layer.bias is not None: |
| 293 | assert layer.bias.shape == (part_size_n,) |
| 294 | bias = marlin_permute_bias(marlin_pad_dim(layer.bias, part_size_n, padded_n)) |
| 295 | layer.bias = torch.nn.Parameter(bias, requires_grad=False) |
| 296 | |
| 297 | return |
| 298 | |
| 299 | |
| 300 | def _repack_marlin_experts( |
| 301 | weight: torch.Tensor, |
| 302 | size_n: int, |
| 303 | size_k: int, |
| 304 | is_a_8bit: bool, |
| 305 | ) -> torch.Tensor: |
| 306 | """Repack each expert to marlin format into a preallocated output.""" |
| 307 | num_experts = weight.shape[0] |
| 308 | out: torch.Tensor | None = None |
| 309 | for i in range(num_experts): |
| 310 | qweight = weight[i].view(torch.int32).T.contiguous() |
| 311 | marlin_qweight = ops.gptq_marlin_repack( |
| 312 | b_q_weight=qweight, |
| 313 | size_k=size_k, |
| 314 | size_n=size_n, |
| 315 | num_bits=4, |
| 316 | is_a_8bit=is_a_8bit, |
| 317 | ) |
| 318 | if out is None: |
| 319 | out = torch.empty( |
| 320 | (num_experts, *marlin_qweight.shape), |
| 321 | dtype=marlin_qweight.dtype, |
| 322 | device=marlin_qweight.device, |
| 323 | ) |
| 324 | out[i] = marlin_qweight |
| 325 | assert out is not None |
| 326 | return out |
| 327 | |
| 328 | |
| 329 | def prepare_nvfp4_moe_layer_for_marlin( |
| 330 | layer: RoutedExperts, |
| 331 | w13: torch.Tensor, |
| 332 | w13_scale: torch.Tensor, |
| 333 | w13_scale_2: torch.Tensor, |
| 334 | w2: torch.Tensor, |
| 335 | w2_scale: torch.Tensor, |
| 336 | w2_scale_2: torch.Tensor, |
| 337 | is_act_and_mul: bool, |
| 338 | ) -> tuple[ |
| 339 | torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor |
| 340 | ]: |
| 341 | logger.warning_once( |
| 342 | "Your GPU does not have native support for FP4 computation but " |
| 343 | "FP4 quantization is being used. Weight-only FP4 compression will " |
| 344 | "be used leveraging the Marlin kernel. This may degrade " |
| 345 | "performance for compute-heavy workloads." |
| 346 | ) |
| 347 | |
| 348 | input_dtype = get_marlin_input_dtype(prefix="") |
| 349 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 350 | raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.") |
| 351 | |
| 352 | GROUP_SIZE = 16 |
| 353 | E = layer.num_experts |
| 354 | K = layer.hidden_size |
| 355 | N = layer.intermediate_size_per_partition |
| 356 | num_shards = 2 if is_act_and_mul else 1 |
| 357 | |
| 358 | # Pad the rank-local intermediate size to satisfy Marlin thread tiles: |
| 359 | # N is an output extent of w13 (per gate/up shard) and the input extent |
| 360 | # of w2, so the padded region never reaches the MoE output. |
| 361 | if K % 128 == 0: |
| 362 | padded_N = round_up(N, 64) |
| 363 | else: |
| 364 | assert K % 64 == 0, f"hidden_size = {K} unsupported by Marlin tiles" |
| 365 | padded_N = round_up(N, 128) |
| 366 | |
| 367 | def pad_w13(x: torch.Tensor) -> torch.Tensor: |
| 368 | """Zero-pad each gate/up shard of a (E, num_shards * N, cols) |
| 369 | tensor to padded_N rows.""" |
| 370 | if padded_N == N: |
| 371 | return x |
| 372 | x = x.view(E, num_shards, N, x.size(-1)) |
| 373 | x = torch.nn.functional.pad(x, (0, 0, 0, padded_N - N)) |
| 374 | return x.reshape(E, num_shards * padded_N, -1) |
| 375 | |
| 376 | def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor: |
| 377 | """Zero-pad the packed N (last) dim of a (E, K, N / packing) |
| 378 | tensor.""" |
| 379 | if padded_N == N: |
| 380 | return x |
| 381 | return torch.nn.functional.pad(x, (0, (padded_N - N) // packing)) |
| 382 | |
| 383 | param_dtype = layer.params_dtype |
| 384 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 385 | |
| 386 | # WEIGHT |
| 387 | # Repack weights to marlin format |
| 388 | def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: |
| 389 | if "w13" in name: |
| 390 | size_n, size_k = N * num_shards, K |
| 391 | assert weight.shape == (E, size_n, size_k // 2) |
| 392 | weight = pad_w13(weight) |
| 393 | size_n = padded_N * num_shards |
| 394 | else: |
| 395 | size_n, size_k = K, N |
| 396 | assert weight.shape == (E, size_n, size_k // 2) |
| 397 | weight = pad_w2(weight, packing=2) |
| 398 | size_k = padded_N |
| 399 | |
| 400 | return _repack_marlin_experts(weight, size_n, size_k, is_a_8bit) |
| 401 | |
| 402 | w13 = repack_weight(w13, "w13") |
| 403 | w2 = repack_weight(w2, "w2") |
| 404 | |
| 405 | # WEIGHT SCALES |
| 406 | # Permute scales |
| 407 | def permute_scales( |
| 408 | scales: torch.Tensor, g_scales: torch.Tensor, name: str |
| 409 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 410 | scales = scales.to(param_dtype) |
| 411 | |
| 412 | tensor_list = [] |
| 413 | if "w13" in name: |
| 414 | scales = pad_w13(scales) |
| 415 | size_n, size_k = padded_N * num_shards, K |
| 416 | else: |
| 417 | scales = pad_w2(scales, packing=GROUP_SIZE) |
| 418 | size_n, size_k = K, padded_N |
| 419 | |
| 420 | # All experts share one global_scale, so compute the max |
| 421 | # scale_factor across all experts first, then apply uniformly. |
| 422 | combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype) |
| 423 | |
| 424 | for i in range(E): |
| 425 | scale = scales[i].T |
| 426 | marlin_scales = marlin_permute_scales( |
| 427 | s=scale, |
| 428 | size_k=size_k, |
| 429 | size_n=size_n, |
| 430 | group_size=GROUP_SIZE, |
| 431 | is_a_8bit=is_a_8bit, |
| 432 | ) |
| 433 | marlin_scales, _ = nvfp4_marlin_process_scales( |
| 434 | marlin_scales, scale_factor=combined_scale_factor, a_dtype=param_dtype |
| 435 | ) |
| 436 | tensor_list.append(marlin_scales) |
| 437 | |
| 438 | scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 439 | g_scales = nvfp4_marlin_process_global_scale(g_scales, param_dtype) |
| 440 | g_scales = g_scales / combined_scale_factor |
| 441 | return scales, g_scales |
| 442 | |
| 443 | w13_scale, w13_scale_2 = permute_scales(w13_scale, w13_scale_2, "w13") |
| 444 | w2_scale, w2_scale_2 = permute_scales(w2_scale, w2_scale_2, "w2") |
| 445 | |
| 446 | return w13, w13_scale, w13_scale_2, w2, w2_scale, w2_scale_2 |
| 447 | |
| 448 | |
| 449 | def prepare_moe_fp4_layer_for_marlin( |
| 450 | layer: torch.nn.Module, input_dtype: torch.dtype | None = None |
| 451 | ) -> None: |
| 452 | is_nvfp4 = hasattr(layer, "w13_weight_scale_2") |
| 453 | if input_dtype is not None and input_dtype.itemsize == 1: |
| 454 | if is_nvfp4: |
| 455 | raise RuntimeError("NVFP4 weight + INT8/FP8 activation is not supported.") |
| 456 | elif input_dtype != torch.float8_e4m3fn: |
| 457 | raise RuntimeError("MXFP4 weight + INT8 activation is not supported.") |
| 458 | |
| 459 | group_size = 16 if is_nvfp4 else 32 |
| 460 | |
| 461 | # Use the per-rank (local) expert count: under expert parallelism the |
| 462 | # w13_weight/w2_weight tensors only hold this rank's experts (created with |
| 463 | # local_num_experts), whereas moe_config.num_experts is the global count. |
| 464 | # With no EP the two are equal, so the non-EP path is unchanged. |
| 465 | e = layer.moe_config.num_local_experts |
| 466 | k = layer.moe_config.hidden_dim |
| 467 | n = layer.moe_config.intermediate_size_per_partition |
| 468 | |
| 469 | param_dtype = layer.params_dtype |
| 470 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 471 | |
| 472 | # WEIGHT |
| 473 | # Repack weights to marlin format |
| 474 | for name in ["w13_weight", "w2_weight"]: |
| 475 | weight = getattr(layer, name) |
| 476 | if "w13" in name: |
| 477 | size_n, size_k = n * 2, k |
| 478 | else: |
| 479 | size_n, size_k = k, n |
| 480 | |
| 481 | assert weight.shape == (e, size_n, size_k // 2) |
| 482 | |
| 483 | weight = _repack_marlin_experts(weight, size_n, size_k, is_a_8bit) |
| 484 | weight = torch.nn.Parameter(weight, requires_grad=False) |
| 485 | |
| 486 | setattr(layer, name, weight) |
| 487 | |
| 488 | # WEIGHT SCALES |
| 489 | # Permute scales |
| 490 | for name in ["w13", "w2"]: |
| 491 | scales = getattr(layer, name + "_weight_scale") |
| 492 | if not is_nvfp4: |
| 493 | scales = scales.view(torch.float8_e8m0fnu) |
| 494 | scales = scales.to(param_dtype) |
| 495 | if is_nvfp4: |
| 496 | global_scale = getattr(layer, name + "_weight_scale_2") |
| 497 | |
| 498 | tensor_list = [] |
| 499 | if "w13" in name: |
| 500 | size_n, size_k = n * 2, k |
| 501 | else: |
| 502 | size_n, size_k = k, n |
| 503 | |
| 504 | # For NVFP4: compute unified scale_factor across all experts |
| 505 | combined_scale_factor = None |
| 506 | if is_nvfp4: |
| 507 | combined_scale_factor = _nvfp4_compute_scale_factor(scales, param_dtype) |
| 508 | |
| 509 | for i in range(e): |
| 510 | scale = scales[i].T |
| 511 | |
| 512 | marlin_scales = marlin_permute_scales( |
| 513 | s=scale, |
| 514 | size_k=size_k, |
| 515 | size_n=size_n, |
| 516 | group_size=group_size, |
| 517 | is_a_8bit=is_a_8bit, |
| 518 | ) |
| 519 | if is_nvfp4: |
| 520 | marlin_scales, _ = nvfp4_marlin_process_scales( |
| 521 | marlin_scales, |
| 522 | scale_factor=combined_scale_factor, |
| 523 | a_dtype=param_dtype, |
| 524 | ) |
| 525 | else: |
| 526 | marlin_scales = mxfp4_marlin_process_scales( |
| 527 | marlin_scales, input_dtype=input_dtype |
| 528 | ) |
| 529 | tensor_list.append(marlin_scales) |
| 530 | |
| 531 | scales = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 532 | scales = torch.nn.Parameter(scales, requires_grad=False) |
| 533 | setattr(layer, name + "_weight_scale", scales) |
| 534 | |
| 535 | if is_nvfp4: |
| 536 | assert combined_scale_factor is not None |
| 537 | global_scale = nvfp4_marlin_process_global_scale(global_scale, param_dtype) |
| 538 | global_scale = global_scale / combined_scale_factor |
| 539 | global_scale = torch.nn.Parameter(global_scale, requires_grad=False) |
| 540 | setattr(layer, name + "_weight_scale_2", global_scale) |
| 541 | |
| 542 | # BIAS |
| 543 | # Permute bias |
| 544 | for name in ["w13_bias", "w2_bias"]: |
| 545 | if not hasattr(layer, name): |
| 546 | continue |
| 547 | bias = getattr(layer, name).to(param_dtype) |
| 548 | |
| 549 | tensor_list = [] |
| 550 | for i in range(e): |
| 551 | expert_bias = bias[i] |
| 552 | |
| 553 | tensor_list.append(marlin_permute_bias(expert_bias)) |
| 554 | |
| 555 | bias = torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 556 | bias = torch.nn.Parameter(bias, requires_grad=False) |
| 557 | setattr(layer, name, bias) |
| 558 | |
| 559 | |
| 560 | def prepare_moe_mxfp4_layer_for_marlin( |
| 561 | layer: torch.nn.Module, |
| 562 | w13: torch.Tensor, |
| 563 | w2: torch.Tensor, |
| 564 | w13_scale: torch.Tensor, |
| 565 | w2_scale: torch.Tensor, |
| 566 | w13_bias: torch.Tensor | None, |
| 567 | w2_bias: torch.Tensor | None, |
| 568 | ) -> tuple[ |
| 569 | torch.Tensor, |
| 570 | torch.Tensor, |
| 571 | torch.Tensor, |
| 572 | torch.Tensor, |
| 573 | torch.Tensor | None, |
| 574 | torch.Tensor | None, |
| 575 | ]: |
| 576 | """Pure-function version of prepare_moe_fp4_layer_for_marlin for MXFP4. |
| 577 | |
| 578 | Takes weight tensors as inputs and returns transformed tensors. |
| 579 | Does NOT modify the layer in-place. |
| 580 | """ |
| 581 | input_dtype = get_marlin_input_dtype() |
| 582 | if ( |
| 583 | input_dtype is not None |
| 584 | and input_dtype.itemsize == 1 |
| 585 | and input_dtype != torch.float8_e4m3fn |
| 586 | ): |
| 587 | raise RuntimeError("MXFP4 weight + INT8 activation is not supported.") |
| 588 | |
| 589 | group_size = 32 # MXFP4 block size |
| 590 | |
| 591 | # Derive dimensions from actual weight shapes to handle rounded/padded |
| 592 | # sizes correctly (e.g., Mxfp4MoEMethod rounds up hidden_dim). |
| 593 | # w13 shape: (E, 2*N, K//2) |
| 594 | e = w13.shape[0] |
| 595 | n = w13.shape[1] // 2 # intermediate_size_per_partition |
| 596 | k = w13.shape[2] * 2 # hidden_size |
| 597 | |
| 598 | param_dtype = layer.params_dtype |
| 599 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 600 | |
| 601 | # WEIGHT: Repack weights to marlin format |
| 602 | def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor: |
| 603 | if "w13" in name: |
| 604 | size_n, size_k = n * 2, k |
| 605 | else: |
| 606 | size_n, size_k = k, n |
| 607 | |
| 608 | assert weight.shape == (e, size_n, size_k // 2) |
| 609 | |
| 610 | return _repack_marlin_experts(weight, size_n, size_k, is_a_8bit) |
| 611 | |
| 612 | w13 = repack_weight(w13, "w13") |
| 613 | w2 = repack_weight(w2, "w2") |
| 614 | |
| 615 | # WEIGHT SCALES: Permute scales |
| 616 | def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor: |
| 617 | scales = scales.view(torch.float8_e8m0fnu) |
| 618 | scales = scales.to(param_dtype) |
| 619 | |
| 620 | tensor_list = [] |
| 621 | if "w13" in name: |
| 622 | size_n, size_k = n * 2, k |
| 623 | else: |
| 624 | size_n, size_k = k, n |
| 625 | |
| 626 | for i in range(e): |
| 627 | scale = scales[i].T |
| 628 | marlin_scales = marlin_permute_scales( |
| 629 | s=scale, |
| 630 | size_k=size_k, |
| 631 | size_n=size_n, |
| 632 | group_size=group_size, |
| 633 | is_a_8bit=is_a_8bit, |
| 634 | ) |
| 635 | marlin_scales = mxfp4_marlin_process_scales( |
| 636 | marlin_scales, input_dtype=input_dtype |
| 637 | ) |
| 638 | tensor_list.append(marlin_scales) |
| 639 | return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 640 | |
| 641 | w13_scale = permute_scales(w13_scale, "w13") |
| 642 | w2_scale = permute_scales(w2_scale, "w2") |
| 643 | |
| 644 | # BIAS: Permute bias |
| 645 | def permute_bias(bias: torch.Tensor | None) -> torch.Tensor | None: |
| 646 | if bias is None: |
| 647 | return None |
| 648 | bias = bias.to(param_dtype) |
| 649 | tensor_list = [] |
| 650 | for i in range(e): |
| 651 | tensor_list.append(marlin_permute_bias(bias[i])) |
| 652 | return torch.cat([x.unsqueeze(0) for x in tensor_list], 0) |
| 653 | |
| 654 | w13_bias = permute_bias(w13_bias) |
| 655 | w2_bias = permute_bias(w2_bias) |
| 656 | |
| 657 | return w13, w2, w13_scale, w2_scale, w13_bias, w2_bias |
| 658 | |
| 659 | |
| 660 | def rand_marlin_weight_nvfp4_like(weight, group_size, input_dtype=None): |
| 661 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 662 | |
| 663 | assert not is_a_8bit, "NVFP4 weight + INT8/FP8 activation is not supported." |
| 664 | assert group_size > 0 |
| 665 | size_n, size_k = weight.shape |
| 666 | scales = weight.view(size_n, -1, group_size).abs().max(-1)[0] / 6 |
| 667 | global_scale = scales.max() / 448 |
| 668 | scales = (scales / global_scale).to(torch.float8_e4m3fn) |
| 669 | |
| 670 | fp4_weight = torch.randint( |
| 671 | 0, 256, (size_n, size_k // 2), dtype=torch.uint8, device=weight.device |
| 672 | ) |
| 673 | fp4_weight_part_1 = (fp4_weight & 0b10000000) | ((fp4_weight & 0b01110000) >> 2) |
| 674 | fp4_weight_part_1 = fp4_weight_part_1.view(torch.float8_e4m3fn) |
| 675 | fp4_weight_part_1 = fp4_weight_part_1.to(weight.dtype) * (2**6) |
| 676 | |
| 677 | fp4_weight2 = fp4_weight << 4 |
| 678 | fp4_weight_part_2 = (fp4_weight2 & 0b10000000) | ((fp4_weight2 & 0b01110000) >> 2) |
| 679 | fp4_weight_part_2 = fp4_weight_part_2.view(torch.float8_e4m3fn) |
| 680 | fp4_weight_part_2 = fp4_weight_part_2.to(weight.dtype) * (2**6) |
| 681 | |
| 682 | weight_ref = torch.cat( |
| 683 | [fp4_weight_part_2.unsqueeze(2), fp4_weight_part_1.unsqueeze(2)], 2 |
| 684 | ).view(size_n, size_k) |
| 685 | weight_ref = ( |
| 686 | weight_ref |
| 687 | * global_scale.to(weight.dtype) |
| 688 | * scales.repeat_interleave(group_size, 1).to(weight.dtype) |
| 689 | ) |
| 690 | |
| 691 | marlin_qweight = ops.gptq_marlin_repack( |
| 692 | b_q_weight=fp4_weight.view(torch.int32).T.contiguous(), |
| 693 | size_k=size_k, |
| 694 | size_n=size_n, |
| 695 | num_bits=4, |
| 696 | is_a_8bit=is_a_8bit, |
| 697 | ) |
| 698 | |
| 699 | marlin_scales = marlin_permute_scales( |
| 700 | s=scales.T.to(weight.dtype), |
| 701 | size_k=size_k, |
| 702 | size_n=size_n, |
| 703 | group_size=group_size, |
| 704 | is_a_8bit=is_a_8bit, |
| 705 | ) |
| 706 | marlin_scales, scale_factor = nvfp4_marlin_process_scales(marlin_scales) |
| 707 | |
| 708 | global_scale = nvfp4_marlin_process_global_scale(global_scale).to(torch.float32) |
| 709 | global_scale = global_scale / scale_factor |
| 710 | |
| 711 | return weight_ref.T, marlin_qweight, marlin_scales, global_scale |
| 712 | |
| 713 | |
| 714 | def rand_marlin_weight_mxfp4_like(weight, group_size, input_dtype=None): |
| 715 | is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1 |
| 716 | if is_a_8bit: |
| 717 | assert input_dtype == torch.float8_e4m3fn, ( |
| 718 | "MXFP4 weight + INT8 activation is not supported." |
| 719 | ) |
| 720 | |
| 721 | assert group_size > 0 |
| 722 | size_n, size_k = weight.shape |
| 723 | scales = torch.randint( |
| 724 | 110, |
| 725 | 120, |
| 726 | (size_n, size_k // group_size), |
| 727 | dtype=torch.uint8, |
| 728 | device=weight.device, |
| 729 | ) |
| 730 | scales = scales.view(torch.float8_e8m0fnu) |
| 731 | |
| 732 | fp4_weight = torch.randint( |
| 733 | 0, 256, (size_n, size_k // 2), dtype=torch.uint8, device=weight.device |
| 734 | ) |
| 735 | fp4_weight_part_1 = (fp4_weight & 0b10000000) | ((fp4_weight & 0b01110000) >> 2) |
| 736 | fp4_weight_part_1 = fp4_weight_part_1.view(torch.float8_e4m3fn) |
| 737 | fp4_weight_part_1 = fp4_weight_part_1.to(weight.dtype) * (2**6) |
| 738 | |
| 739 | fp4_weight2 = fp4_weight << 4 |
| 740 | fp4_weight_part_2 = (fp4_weight2 & 0b10000000) | ((fp4_weight2 & 0b01110000) >> 2) |
| 741 | fp4_weight_part_2 = fp4_weight_part_2.view(torch.float8_e4m3fn) |
| 742 | fp4_weight_part_2 = fp4_weight_part_2.to(weight.dtype) * (2**6) |
| 743 | |
| 744 | weight_ref = torch.cat( |
| 745 | [fp4_weight_part_2.unsqueeze(2), fp4_weight_part_1.unsqueeze(2)], 2 |
| 746 | ).view(size_n, size_k) |
| 747 | weight_ref = weight_ref * scales.repeat_interleave(group_size, 1).to(weight.dtype) |
| 748 | |
| 749 | fp4_weight = fp4_weight.view(torch.int32).T.contiguous() |
| 750 | marlin_qweight = ops.gptq_marlin_repack( |
| 751 | b_q_weight=fp4_weight, |
| 752 | size_k=size_k, |
| 753 | size_n=size_n, |
| 754 | num_bits=4, |
| 755 | is_a_8bit=is_a_8bit, |
| 756 | ) |
| 757 | |
| 758 | marlin_scales = marlin_permute_scales( |
| 759 | s=scales.T.to(weight.dtype), |
| 760 | size_k=size_k, |
| 761 | size_n=size_n, |
| 762 | group_size=group_size, |
| 763 | is_a_8bit=is_a_8bit, |
| 764 | ) |
| 765 | |
| 766 | marlin_scales = mxfp4_marlin_process_scales(marlin_scales, input_dtype=input_dtype) |
| 767 | |
| 768 | return weight_ref.T, marlin_qweight, marlin_scales.to(torch.float8_e8m0fnu) |