feat: support tile-misaligned GPTQ Marlin shapes

gptqmodel/nn_modules/qlinear/marlin.py

@@ -32,11 +32,17 @@
     _marlin_capability_supported,
     _transform_param,
     apply_gptq_marlin_linear,
+    apply_gptq_marlin_linear_padded,
     gptq_marlin_repack,
     marlin_import_exception,
+    marlin_is_tile_aligned,
     marlin_is_k_full,
     marlin_make_empty_g_idx,
     marlin_make_workspace_new,
+    marlin_pad_dim,
+    marlin_pad_qweight,
+    marlin_pad_scales,
+    marlin_padded_nk,
     marlin_permute_bias,
     marlin_permute_scales,
     marlin_repeat_scales_on_all_ranks,
@@ -62,9 +68,10 @@ class MarlinLinear(GPTQQuantLinear):
     SUPPORTS_SYM = [True]
     SUPPORTS_SHARDS = True
     SUPPORTS_TRAINING = False
+    # Tile padding is handled below; group boundaries must still divide K.
     SUPPORTS_AUTO_PADDING = False
     SUPPORTS_IN_FEATURES_DIVISIBLE_BY = [1]
-    SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [64]
+    SUPPORTS_OUT_FEATURES_DIVISIBLE_BY = [1]
 
     SUPPORTS_DEVICES = [DEVICE.CUDA]
     SUPPORTS_PLATFORM = [PLATFORM.LINUX]
@@ -105,11 +112,22 @@ def __init__(
         # self.original_in_features = in_features
         # self.original_out_features = out_features
 
-        if desc_act and group_size == -1:
+        if desc_act and group_size in (-1, in_features):
             # In this case, act_order == True is the same as act_order == False
             # (since we have only one group per output channel)
             desc_act = False
 
+        selected_backend = kwargs.pop("backend", BACKEND.GPTQ_MARLIN)
+        # Padding adds work to every forward, so automatic selection stays conservative.
+        if selected_backend in (BACKEND.AUTO, BACKEND.AUTO_TRAINABLE) and not marlin_is_tile_aligned(
+            out_features, in_features
+        ):
+            raise NotImplementedError(
+                "Automatic Marlin selection keeps tile-misaligned shapes on "
+                "the next compatible backend; request GPTQ_MARLIN explicitly "
+                "to enable runtime tile padding."
+            )
+
         self.compute_dtype = kwargs.get("dtype") or torch.float16
         self.fp32 = env_flag("GPTQMODEL_MARLIN_USE_FP32", default=True)
 
@@ -122,7 +140,7 @@ def __init__(
             out_features=out_features,
             bias=bias,
             pack_dtype=pack_dtype,
-            backend=kwargs.pop("backend", BACKEND.GPTQ_MARLIN),
+            backend=selected_backend,
             adapter=adapter,
             register_buffers=False, # do not register buffers in super()
             **kwargs)
@@ -225,6 +243,40 @@ def validate_once(cls) -> Tuple[bool, Optional[Exception]]:
             return False, ImportError(marlin_import_exception)
         return True, None
 
+    @classmethod
+    def _validate(cls, **args) -> Tuple[bool, Optional[Exception]]:
+        ok, err = super()._validate(**args)
+        if not ok:
+            return ok, err
+
+        bits = args.get("bits", 4)
+        in_features = args.get("in_features")
+        out_features = args.get("out_features")
+        desc_act = args.get("desc_act", False)
+        group_size = args.get("group_size", -1)
+        if in_features is None or out_features is None:
+            return True, None
+
+        pack_factor = 32 // bits
+        # Tile padding cannot repair a partially packed int32 row or column.
+        if in_features % pack_factor != 0 or out_features % pack_factor != 0:
+            return False, NotImplementedError(
+                "Marlin packed dimensions must be divisible by "
+                f"pack_factor={pack_factor}; got K={in_features}, "
+                f"N={out_features}."
+            )
+
+        effective_desc_act = desc_act and group_size not in (-1, in_features)
+        # Act-order indices only describe the original K dimension.
+        if effective_desc_act and not marlin_is_tile_aligned(
+            out_features, in_features
+        ):
+            return False, NotImplementedError(
+                "Marlin activation-order weights require an aligned thread "
+                f"tile; got K={in_features}, N={out_features}."
+            )
+
+        return True, None
 
     @classmethod
     def validate_device(cls, device: DEVICE):
@@ -257,20 +309,59 @@ def post_init(self):
         # Allocate marlin workspace.
         self.workspace = marlin_make_workspace_new(device)
 
+        # GPTQModel also accepts group_size=K as channelwise quantization.
+        marlin_group_size = (
+            -1
+            if self.requested_group_size == self.in_features
+            else self.requested_group_size
+        )
+        # Validation keeps act-order shapes aligned; other shapes may use zero padding.
+        if self.desc_act:
+            padded_n, padded_k = self.out_features, self.in_features
+        else:
+            padded_n, padded_k = marlin_padded_nk(
+                self.out_features,
+                self.in_features,
+                marlin_group_size,
+            )
+        self._marlin_tile_padding = (
+            None
+            if (padded_n, padded_k) == (self.out_features, self.in_features)
+            else (padded_n, padded_k)
+        )
+
         def transform_w_q(x):
-            x.data = gptq_marlin_repack(x.data.contiguous(),
+            # Pad in GPTQ layout before converting to Marlin layout.
+            padded = marlin_pad_qweight(
+                x.data.contiguous(),
+                self.out_features,
+                self.in_features,
+                padded_n,
+                padded_k,
+            )
+            x.data = gptq_marlin_repack(padded,
                                         perm=self.g_idx_sort_indices,
-                                        size_k=self.in_features,
-                                        size_n=self.out_features,
+                                        size_k=padded_k,
+                                        size_n=padded_n,
                                         num_bits=self.bits,
                                         dtype=self.compute_dtype)
             return x
 
         def transform_w_s(x):
-            x.data = marlin_permute_scales(x.data.contiguous(),
-                                           size_k=self.in_features,
-                                           size_n=self.out_features,
-                                           group_size=self.group_size)
+            padded = marlin_pad_scales(
+                x.data.contiguous(),
+                self.out_features,
+                self.in_features,
+                padded_n,
+                padded_k,
+                marlin_group_size,
+            )
+            x.data = marlin_permute_scales(
+                padded,
+                size_k=padded_k,
+                size_n=padded_n,
+                group_size=marlin_group_size,
+            )
             return x
 
         # Handle sorting for activation reordering if needed.
@@ -288,7 +379,9 @@ def transform_w_s(x):
         _transform_param(self, "scales", transform_w_s)
 
         if hasattr(self, "bias") and self.bias is not None:
-            self.bias.data = marlin_permute_bias(self.bias)
+            self.bias.data = marlin_permute_bias(
+                marlin_pad_dim(self.bias, self.out_features, padded_n)
+            )
 
         super().post_init()
 
@@ -314,22 +407,42 @@ def forward(self, x: torch.Tensor):
         if self.bias is not None and self.bias.dtype != x.dtype:
             self.bias.data = self.bias.data.to(dtype=x.dtype)
 
-        out = apply_gptq_marlin_linear(
-            input=x.contiguous() if self.is_lm_head else x,
-            weight=self.qweight,
-            weight_scale=self.scales,
-            weight_zp=self.qzeros,
-            g_idx=self.g_idx,
-            g_idx_sort_indices=self.g_idx_sort_indices,
-            workspace=self.workspace,
-            wtype=self.weight_type,
-            output_size_per_partition=self.out_features,
-            input_size_per_partition=self.in_features,
-            is_k_full=self.is_k_full,
-            bias=self.bias,
-            use_fp32_reduce=self.fp32,
-            use_atomics=False, # reduces accuracy with slightly faster performance
-        )
+        # Keep aligned layers on the original decode-sensitive call path.
+        if self._marlin_tile_padding is None:
+            out = apply_gptq_marlin_linear(
+                input=x.contiguous() if self.is_lm_head else x,
+                weight=self.qweight,
+                weight_scale=self.scales,
+                weight_zp=self.qzeros,
+                g_idx=self.g_idx,
+                g_idx_sort_indices=self.g_idx_sort_indices,
+                workspace=self.workspace,
+                wtype=self.weight_type,
+                output_size_per_partition=self.out_features,
+                input_size_per_partition=self.in_features,
+                is_k_full=self.is_k_full,
+                bias=self.bias,
+                use_fp32_reduce=self.fp32,
+                use_atomics=False, # reduces accuracy with slightly faster performance
+            )
+        else:
+            out = apply_gptq_marlin_linear_padded(
+                tile_padding=self._marlin_tile_padding,
+                input=x.contiguous() if self.is_lm_head else x,
+                weight=self.qweight,
+                weight_scale=self.scales,
+                weight_zp=self.qzeros,
+                g_idx=self.g_idx,
+                g_idx_sort_indices=self.g_idx_sort_indices,
+                workspace=self.workspace,
+                wtype=self.weight_type,
+                output_size_per_partition=self.out_features,
+                input_size_per_partition=self.in_features,
+                is_k_full=self.is_k_full,
+                bias=self.bias,
+                use_fp32_reduce=self.fp32,
+                use_atomics=False,
+            )
 
         if self.adapter:
             out = self.adapter.apply(x=x, out=out)

gptqmodel/utils/marlin.py

@@ -4,6 +4,7 @@
 # Contact: qubitium@modelcloud.ai, x.com/qubitium
 from __future__ import annotations
 
+import math
 import subprocess
 import sys
 import threading
@@ -319,6 +320,86 @@ def marlin_make_workspace_new(device: torch.device,
                        requires_grad=False)
 
 
+def _round_up(value: int, multiple: int) -> int:
+    """Round value up to the next multiple."""
+    return ((value + multiple - 1) // multiple) * multiple
+
+
+# Marlin accepts either orientation of its 64 x 128 thread tile.
+def marlin_is_tile_aligned(size_n: int, size_k: int) -> bool:
+    return (
+        size_n % 64 == 0 and size_k % 128 == 0
+    ) or (
+        size_n % 128 == 0 and size_k % 64 == 0
+    )
+
+
+def marlin_padded_nk(size_n: int, size_k: int,
+                     group_size: int = -1) -> Tuple[int, int]:
+    """Return the smallest N/K pair supported by a Marlin thread tile.
+
+    Padded K consumes zero activations; padded N uses zero scales, so neither
+    region changes the logical output.
+    """
+    group = group_size if group_size > 0 else 1
+    # Try both tile orientations and keep the one with the least padded work.
+    candidates = (
+        (_round_up(size_n, 64), _round_up(size_k, math.lcm(128, group))),
+        (_round_up(size_n, 128), _round_up(size_k, math.lcm(64, group))),
+    )
+    padded_nk = min(candidates, key=lambda nk: (nk[0] * nk[1], nk[0] + nk[1]))
+    if padded_nk != (size_n, size_k):
+        log.warn.once(
+            "Marlin is padding a tile-misaligned weight shape. Activations "
+            "and outputs for this layer are padded and sliced on each forward; "
+            "performance may be degraded."
+        )
+    return padded_nk
+
+
+def marlin_pad_qweight(qweight: torch.Tensor, size_n: int, size_k: int,
+                       padded_n: int, padded_k: int) -> torch.Tensor:
+    """Zero-pad a GPTQ-layout packed weight before Marlin repacking."""
+    if (padded_n, padded_k) == (size_n, size_k):
+        return qweight
+    # Each packed row stores pack_factor consecutive K values.
+    pack_factor = size_k // qweight.size(0)
+    return torch.nn.functional.pad(
+        qweight,
+        (0, padded_n - size_n, 0, (padded_k - size_k) // pack_factor),
+    )
+
+
+def marlin_pad_scales(scales: torch.Tensor, size_n: int, size_k: int,
+                      padded_n: int, padded_k: int,
+                      group_size: int) -> torch.Tensor:
+    """Zero-pad scale rows and columns to the padded Marlin shape."""
+    if (padded_n, padded_k) == (size_n, size_k):
+        return scales
+    # Extra K groups need zero scales so padded weights stay inactive.
+    pad_rows = padded_k // group_size - scales.size(0) if group_size > 0 else 0
+    if pad_rows < 0:
+        raise ValueError("Padded Marlin K cannot contain fewer scale groups.")
+    return torch.nn.functional.pad(
+        scales, (0, padded_n - size_n, 0, pad_rows)
+    )
+
+
+def marlin_pad_dim(x: torch.Tensor, size: int, padded: int) -> torch.Tensor:
+    """Zero-pad the last tensor dimension when a Marlin tile requires it."""
+    if padded == size:
+        return x
+    return torch.nn.functional.pad(x, (0, padded - size))
+
+
+def marlin_unpad_output(output: torch.Tensor, size_n: int,
+                        padded_n: int) -> torch.Tensor:
+    """Slice a padded Marlin result back to its logical output width."""
+    if padded_n == size_n:
+        return output
+    return output[..., :size_n].contiguous()
+
+
 def update_tensor_inplace(dst: torch.Tensor, src: torch.Tensor):
     assert dst.dtype == src.dtype, "Tensors must have the same dtype"
 
@@ -516,6 +597,46 @@ def apply_gptq_marlin_linear(
     return output.reshape(out_shape)
 
 
+def apply_gptq_marlin_linear_padded(
+        *,
+        tile_padding: Tuple[int, int],
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+        weight_zp: torch.Tensor,
+        g_idx: torch.Tensor,
+        g_idx_sort_indices: torch.Tensor,
+        workspace: torch.Tensor,
+        wtype: ScalarType,
+        output_size_per_partition: int,
+        input_size_per_partition: int,
+        is_k_full: bool,
+        bias: Optional[torch.Tensor] = None,
+        use_fp32_reduce: bool = True,
+        use_atomics: bool = False,
+) -> torch.Tensor:
+    """Pad one tile-misaligned GEMM around the unchanged Marlin hot path."""
+    padded_n, padded_k = tile_padding
+    padded_input = marlin_pad_dim(input, input_size_per_partition, padded_k)
+    output = apply_gptq_marlin_linear(
+        input=padded_input,
+        weight=weight,
+        weight_scale=weight_scale,
+        weight_zp=weight_zp,
+        g_idx=g_idx,
+        g_idx_sort_indices=g_idx_sort_indices,
+        workspace=workspace,
+        wtype=wtype,
+        output_size_per_partition=padded_n,
+        input_size_per_partition=padded_k,
+        is_k_full=is_k_full,
+        bias=bias,
+        use_fp32_reduce=use_fp32_reduce,
+        use_atomics=use_atomics,
+    )
+    return marlin_unpad_output(output, output_size_per_partition, padded_n)
+
+
 def apply_awq_marlin_linear(
         input: torch.Tensor,
         weight: torch.Tensor,