Pad FP8 Marlin weights to valid thread tiles

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

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),
         )

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(

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 lcm
+
 import torch
 
 import vllm._custom_ops as ops
@@ -17,10 +19,28 @@
 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 _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 +69,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 +94,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 +111,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 +156,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 +181,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 +220,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 +241,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 +408,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 +442,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 +474,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 +499,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 +518,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,