From 162dc3fd95f0c4275e84e0d64333b51f630c34c6 Mon Sep 17 00:00:00 2001 From: Mike G <180722391+mikekg@users.noreply.github.com> Date: Tue, 9 Jun 2026 00:05:51 -0700 Subject: [PATCH 1/2] Pad FP8 Marlin weights to valid thread tiles FP8 and MXFP8 Marlin repack and GEMM require rank-local N/K extents that match a supported Marlin thread tile family. Padding everything to N64/K128 is valid but can overpad cases where the N128/K64 family is sufficient. Select the lower-overhead rank-local padded size from the two non-dominated tile families that are valid for both small and large batches: N multiple 128 with K multiple 64, or N multiple 64 with K multiple 128. Include the quantization group size in the K multiple when needed so block-scale layouts remain consistent. Use the selected extents consistently when repacking weights and scales and when launching the Marlin GEMM. Pad the activation K dimension with zeros, slice padded output columns after GEMM, and pad bias before Marlin bias permutation when bias is present. This padding is applied after tensor-parallel partitioning via output_size_per_partition and input_size_per_partition, so global checkpoint tensors are not expanded before TP slicing. Signed-off-by: Michael Gschwind Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com> --- .../kernels/linear/mxfp8/marlin.py | 2 + .../kernels/linear/scaled_mm/marlin.py | 2 + .../quantization/utils/marlin_utils_fp8.py | 124 +++++++++++++++--- 3 files changed, 109 insertions(+), 19 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mxfp8/marlin.py b/vllm/model_executor/kernels/linear/mxfp8/marlin.py index bec54cd942ed..ff5c21f407b9 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/marlin.py +++ b/vllm/model_executor/kernels/linear/mxfp8/marlin.py @@ -50,4 +50,6 @@ def apply_weights( size_n=layer.output_size_per_partition, size_k=layer.input_size_per_partition, bias=bias, + padded_size_n=getattr(layer, "marlin_padded_size_n", None), + padded_size_k=getattr(layer, "marlin_padded_size_k", None), ) diff --git a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py index 66a03b4d205b..7b343edfd5ef 100644 --- a/vllm/model_executor/kernels/linear/scaled_mm/marlin.py +++ b/vllm/model_executor/kernels/linear/scaled_mm/marlin.py @@ -102,6 +102,8 @@ def apply_weights( size_k=layer.input_size_per_partition, input_dtype=self.marlin_input_dtype, bias=bias, + padded_size_n=getattr(layer, "marlin_padded_size_n", None), + padded_size_k=getattr(layer, "marlin_padded_size_k", None), ) def apply_scaled_mm( diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 6e2ae5c91a36..760a11db73a1 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from math import gcd + import torch import vllm._custom_ops as ops @@ -21,6 +23,31 @@ logger = init_logger(__name__) +def _round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +def _lcm(a: int, b: int) -> int: + return a * b // gcd(a, b) + + +def _get_fp8_marlin_padded_sizes( + size_n: int, size_k: int, group_size: int = -1 +) -> tuple[int, int]: + group_k = group_size if group_size > 0 else 1 + candidates = ( + (_round_up(size_n, 128), _round_up(size_k, _lcm(64, group_k))), + (_round_up(size_n, 64), _round_up(size_k, _lcm(128, group_k))), + ) + return min( + candidates, + key=lambda sizes: ( + sizes[0] * sizes[1], + sizes[0] - size_n + sizes[1] - size_k, + ), + ) + + def is_fp8_marlin_supported(): return current_platform.has_device_capability(75) @@ -49,15 +76,23 @@ def apply_fp8_marlin_linear( bias: torch.Tensor | None, input_dtype: torch.dtype | None = None, use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, + padded_size_n: int | None = None, + padded_size_k: int | None = None, ) -> torch.Tensor: # For GPUs that lack FP8 hardware support, we can leverage the # Marlin kernel for fast weight-only FP8 quantization reshaped_x = input.reshape(-1, input.shape[-1]) out_shape = input.shape[:-1] + (size_n,) + padded_size_n = padded_size_n if padded_size_n is not None else size_n + padded_size_k = padded_size_k if padded_size_k is not None else size_k use_atomic_add = should_use_atomic_add_reduce( - m=reshaped_x.size(0), n=size_n, k=size_k, device=input.device, dtype=input.dtype + m=reshaped_x.size(0), + n=padded_size_n, + k=padded_size_k, + device=input.device, + dtype=input.dtype, ) inputs = reshaped_x @@ -66,6 +101,9 @@ def apply_fp8_marlin_linear( # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn) raise RuntimeError("Marlin W8A8 is not supported.") + if padded_size_k != size_k: + inputs = torch.nn.functional.pad(inputs, (0, padded_size_k - size_k)) + output = ops.marlin_gemm( a=inputs, c=None, @@ -80,12 +118,14 @@ def apply_fp8_marlin_linear( workspace=workspace, b_q_type=scalar_types.float8_e4m3fn, size_m=reshaped_x.size(0), - size_n=size_n, - size_k=size_k, + size_n=padded_size_n, + size_k=padded_size_k, use_atomic_add=use_atomic_add, use_fp32_reduce=use_fp32_reduce, ) + if padded_size_n != size_n: + output = output[..., :size_n].contiguous() return output.reshape(out_shape) @@ -123,12 +163,20 @@ def prepare_fp8_layer_for_marlin( qweight = pack_fp8_to_int32(layer.weight, size_k_first) if not size_k_first: qweight = qweight.T.contiguous() + group_size = -1 if weight_block_size is None else weight_block_size[1] + padded_part_size_n, padded_part_size_k = _get_fp8_marlin_padded_sizes( + part_size_n, part_size_k, group_size + ) + n_pad = padded_part_size_n - part_size_n + k_pad_i32 = (padded_part_size_k - part_size_k) // 4 + if n_pad or k_pad_i32: + qweight = torch.nn.functional.pad(qweight, (0, n_pad, 0, k_pad_i32)) marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, perm=perm, - size_k=part_size_k, - size_n=part_size_n, + size_k=padded_part_size_k, + size_n=padded_part_size_n, num_bits=8, ) replace_parameter(layer, "weight", marlin_qweight) @@ -140,8 +188,6 @@ def prepare_fp8_layer_for_marlin( elif "weight_scale_inv" in dir(layer): scales = layer.weight_scale_inv.to(layer.orig_dtype) - group_size = -1 if weight_block_size is None else weight_block_size[1] - # marlin kernel only support channel-wise and group-wise quantization # we need to convert the scales if weight_block_size is None: @@ -181,9 +227,17 @@ def prepare_fp8_layer_for_marlin( scales = scales.repeat_interleave(block_n, 1) # size_n may not divisible by block_size[0] scales = scales[:, :part_size_n] + k_pad_groups = 0 + if group_size != -1: + k_pad_groups = padded_part_size_k // group_size - part_size_k // group_size + if n_pad or k_pad_groups: + scales = torch.nn.functional.pad(scales, (0, n_pad, 0, k_pad_groups)) marlin_scales = marlin_permute_scales( - s=scales, size_k=part_size_k, size_n=part_size_n, group_size=group_size + s=scales, + size_k=padded_part_size_k, + size_n=padded_part_size_n, + group_size=group_size, ) if input_dtype != torch.float8_e4m3fn: marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales) @@ -194,9 +248,15 @@ def prepare_fp8_layer_for_marlin( if hasattr(layer, "bias") and layer.bias is not None: assert layer.bias.shape == (part_size_n,) - bias = marlin_permute_bias(layer.bias) + bias = layer.bias + if n_pad: + bias = torch.nn.functional.pad(bias, (0, n_pad)) + bias = marlin_permute_bias(bias) replace_parameter(layer, "bias", bias) + layer.marlin_padded_size_n = padded_part_size_n + layer.marlin_padded_size_k = padded_part_size_k + def prepare_fp8_moe_layer_for_marlin( layer: torch.nn.Module, @@ -355,20 +415,28 @@ def apply_mxfp8_marlin_linear( size_k: int, bias: torch.Tensor | None = None, use_fp32_reduce: bool = USE_FP32_REDUCE_DEFAULT, + padded_size_n: int | None = None, + padded_size_k: int | None = None, ) -> torch.Tensor: reshaped_x = input.reshape(-1, input.shape[-1]) out_shape = input.shape[:-1] + (size_n,) + padded_size_n = padded_size_n if padded_size_n is not None else size_n + padded_size_k = padded_size_k if padded_size_k is not None else size_k use_atomic_add = should_use_atomic_add_reduce( m=reshaped_x.size(0), - n=size_n, - k=size_k, + n=padded_size_n, + k=padded_size_k, device=input.device, dtype=input.dtype, ) + inputs = reshaped_x + if padded_size_k != size_k: + inputs = torch.nn.functional.pad(inputs, (0, padded_size_k - size_k)) + output = ops.marlin_gemm( - a=reshaped_x, + a=inputs, c=None, b_q_weight=weight, b_bias=bias, @@ -381,12 +449,14 @@ def apply_mxfp8_marlin_linear( workspace=workspace, b_q_type=scalar_types.float8_e4m3fn, size_m=reshaped_x.size(0), - size_n=size_n, - size_k=size_k, + size_n=padded_size_n, + size_k=padded_size_k, use_atomic_add=use_atomic_add, use_fp32_reduce=use_fp32_reduce, ) + if padded_size_n != size_n: + output = output[..., :size_n].contiguous() return output.reshape(out_shape) @@ -411,12 +481,19 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: perm = torch.empty(0, dtype=torch.int, device=device) qweight = pack_fp8_to_int32(layer.weight, size_k_first=False) qweight = qweight.T.contiguous() + padded_part_size_n, padded_part_size_k = _get_fp8_marlin_padded_sizes( + part_size_n, part_size_k, group_size + ) + n_pad = padded_part_size_n - part_size_n + k_pad_i32 = (padded_part_size_k - part_size_k) // 4 + if n_pad or k_pad_i32: + qweight = torch.nn.functional.pad(qweight, (0, n_pad, 0, k_pad_i32)) marlin_qweight = ops.gptq_marlin_repack( b_q_weight=qweight, perm=perm, - size_k=part_size_k, - size_n=part_size_n, + size_k=padded_part_size_k, + size_n=padded_part_size_n, num_bits=8, ) replace_parameter(layer, "weight", marlin_qweight) @@ -429,12 +506,15 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: scales = scales.contiguous() scales = scales.view(torch.float8_e8m0fnu).to(param_dtype) scales = scales.T.contiguous() + k_pad_groups = padded_part_size_k // group_size - part_size_k // group_size + if n_pad or k_pad_groups: + scales = torch.nn.functional.pad(scales, (0, n_pad, 0, k_pad_groups)) # Permute scales to Marlin layout marlin_scales = marlin_permute_scales( s=scales, - size_k=part_size_k, - size_n=part_size_n, + size_k=padded_part_size_k, + size_n=padded_part_size_n, group_size=group_size, ) @@ -445,9 +525,15 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None: # BIAS if hasattr(layer, "bias") and layer.bias is not None: assert layer.bias.shape == (part_size_n,) - bias = marlin_permute_bias(layer.bias) + bias = layer.bias + if n_pad: + bias = torch.nn.functional.pad(bias, (0, n_pad)) + bias = marlin_permute_bias(bias) replace_parameter(layer, "bias", bias) + layer.marlin_padded_size_n = padded_part_size_n + layer.marlin_padded_size_k = padded_part_size_k + def prepare_mxfp8_moe_layer_for_marlin( layer: torch.nn.Module, From b9729634a680d52687891479c760e3552273203b Mon Sep 17 00:00:00 2001 From: Mike G <180722391+mikekg@users.noreply.github.com> Date: Tue, 9 Jun 2026 00:05:51 -0700 Subject: [PATCH 2/2] Add FP8 Marlin padding unit test Signed-off-by: Michael Gschwind Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com> --- .../model_executor/test_fp8_marlin_padding.py | 41 +++++++++++++++++++ .../quantization/utils/marlin_utils_fp8.py | 15 ++----- 2 files changed, 45 insertions(+), 11 deletions(-) create mode 100644 tests/model_executor/test_fp8_marlin_padding.py diff --git a/tests/model_executor/test_fp8_marlin_padding.py b/tests/model_executor/test_fp8_marlin_padding.py new file mode 100644 index 000000000000..0fe2772494a1 --- /dev/null +++ b/tests/model_executor/test_fp8_marlin_padding.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( + _get_fp8_marlin_padded_sizes, +) + + +def _is_valid_fp8_marlin_thread_tile(size_n: int, size_k: int) -> bool: + return (size_n % 128 == 0 and size_k % 64 == 0) or ( + size_n % 64 == 0 and size_k % 128 == 0 + ) + + +@pytest.mark.parametrize( + ("size_n", "size_k", "group_size", "expected"), + [ + (4640, 4096, -1, (4672, 4096)), + (200, 129, -1, (256, 192)), + (129, 65, -1, (192, 128)), + (200, 257, 128, (256, 384)), + ], +) +def test_fp8_marlin_padding_maps_invalid_shapes_to_valid_thread_tiles( + size_n: int, + size_k: int, + group_size: int, + expected: tuple[int, int], +) -> None: + assert not _is_valid_fp8_marlin_thread_tile(size_n, size_k) + + padded_n, padded_k = _get_fp8_marlin_padded_sizes(size_n, size_k, group_size) + + assert (padded_n, padded_k) == expected + assert _is_valid_fp8_marlin_thread_tile(padded_n, padded_k) + assert padded_n >= size_n + assert padded_k >= size_k + if group_size > 0: + assert padded_k % group_size == 0 diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py index 760a11db73a1..15ac56f2613b 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from math import gcd +from math import lcm import torch @@ -19,25 +19,18 @@ from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.scalar_type import scalar_types +from vllm.utils.math_utils import round_up logger = init_logger(__name__) -def _round_up(value: int, multiple: int) -> int: - return (value + multiple - 1) // multiple * multiple - - -def _lcm(a: int, b: int) -> int: - return a * b // gcd(a, b) - - def _get_fp8_marlin_padded_sizes( size_n: int, size_k: int, group_size: int = -1 ) -> tuple[int, int]: group_k = group_size if group_size > 0 else 1 candidates = ( - (_round_up(size_n, 128), _round_up(size_k, _lcm(64, group_k))), - (_round_up(size_n, 64), _round_up(size_k, _lcm(128, group_k))), + (round_up(size_n, 128), round_up(size_k, lcm(64, group_k))), + (round_up(size_n, 64), round_up(size_k, lcm(128, group_k))), ) return min( candidates,