From 37dc45d392960e58acecea7c1486df2d09fa37a4 Mon Sep 17 00:00:00 2001 From: ZX-ModelCloud Date: Tue, 25 Aug 2026 13:18:56 +0800 Subject: [PATCH] feat: support tile-misaligned GPTQ Marlin shapes --- gptqmodel/nn_modules/qlinear/marlin.py | 167 +++++++++++-- gptqmodel/utils/marlin.py | 121 +++++++++ tests/test_marlin_jit.py | 328 +++++++++++++++++++++++++ 3 files changed, 589 insertions(+), 27 deletions(-) diff --git a/gptqmodel/nn_modules/qlinear/marlin.py b/gptqmodel/nn_modules/qlinear/marlin.py index 165711939..e670688e6 100644 --- a/gptqmodel/nn_modules/qlinear/marlin.py +++ b/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) diff --git a/gptqmodel/utils/marlin.py b/gptqmodel/utils/marlin.py index 18f2571ea..31f135121 100644 --- a/gptqmodel/utils/marlin.py +++ b/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, diff --git a/tests/test_marlin_jit.py b/tests/test_marlin_jit.py index 44e72f8c2..9adfb32d4 100644 --- a/tests/test_marlin_jit.py +++ b/tests/test_marlin_jit.py @@ -13,6 +13,7 @@ import gptqmodel.nn_modules.qlinear.marlin_awq as marlin_awq_qlinear_module import gptqmodel.utils.marlin as marlin_utils from gptqmodel import extension as extension_api +from gptqmodel.utils.backend import BACKEND from gptqmodel.utils import cpp as cpp_module from gptqmodel.utils.marlin_scalar_type import scalar_types @@ -230,6 +231,133 @@ def test_marlin_capability_checks_allow_sm75_but_reject_sm70(monkeypatch): assert marlin_utils._validate_marlin_device_support() is False +@pytest.mark.parametrize( + "shape,group_size,expected", + [ + ((64, 128), -1, (64, 128)), + ((128, 64), -1, (128, 64)), + ((200, 288), 32, (256, 320)), + ((256, 208), -1, (256, 256)), + ((200, 384), 128, (256, 384)), + ], +) +def test_marlin_padded_nk_selects_minimal_thread_tile(shape, group_size, expected): + size_n, size_k = shape + + padded_n, padded_k = marlin_utils.marlin_padded_nk( + size_n, size_k, group_size + ) + + assert (padded_n, padded_k) == expected + assert marlin_utils.marlin_is_tile_aligned(padded_n, padded_k) + if group_size > 0: + assert padded_k % group_size == 0 + + +def test_marlin_tile_padding_helpers_preserve_values_and_shapes(): + size_n, size_k, group_size = 200, 288, 32 + padded_n, padded_k = marlin_utils.marlin_padded_nk( + size_n, size_k, group_size + ) + + qweight = torch.ones((size_k // 8, size_n), dtype=torch.int32) + padded_qweight = marlin_utils.marlin_pad_qweight( + qweight, size_n, size_k, padded_n, padded_k + ) + assert padded_qweight.shape == (padded_k // 8, padded_n) + assert torch.equal(padded_qweight[: qweight.size(0), :size_n], qweight) + assert torch.count_nonzero(padded_qweight[:, size_n:]) == 0 + assert torch.count_nonzero(padded_qweight[qweight.size(0) :, :]) == 0 + + scales = torch.ones((size_k // group_size, size_n)) + padded_scales = marlin_utils.marlin_pad_scales( + scales, size_n, size_k, padded_n, padded_k, group_size + ) + assert padded_scales.shape == (padded_k // group_size, padded_n) + assert torch.equal(padded_scales[: scales.size(0), :size_n], scales) + assert torch.count_nonzero(padded_scales[:, size_n:]) == 0 + + +def test_marlin_quant_linear_validation_limits_tile_padding_to_non_act_order(monkeypatch): + monkeypatch.setattr(marlin_qlinear_module, "marlin_import_exception", None) + common = { + "bits": 4, + "group_size": 32, + "sym": True, + "in_features": 288, + "out_features": 200, + "pack_dtype": torch.int32, + "dtype": torch.float16, + "dynamic": None, + "device": None, + "trainable": False, + "adapter": None, + } + + ok, err = marlin_qlinear_module.MarlinLinear._validate( + **common, desc_act=False + ) + assert ok is True + assert err is None + + ok, err = marlin_qlinear_module.MarlinLinear._validate( + **common, desc_act=True + ) + assert ok is False + assert "activation-order" in str(err) + + channelwise = dict(common, group_size=-1) + ok, err = marlin_qlinear_module.MarlinLinear._validate( + **channelwise, desc_act=True + ) + assert ok is True + assert err is None + + explicit_channelwise = dict(common, group_size=common["in_features"]) + ok, err = marlin_qlinear_module.MarlinLinear._validate( + **explicit_channelwise, desc_act=True + ) + assert ok is True + assert err is None + + aligned = dict(common, in_features=64, out_features=128) + ok, err = marlin_qlinear_module.MarlinLinear._validate( + **aligned, desc_act=True + ) + assert ok is True + assert err is None + + +def test_marlin_auto_selection_keeps_tile_padding_opt_in(monkeypatch): + monkeypatch.setattr(marlin_qlinear_module, "marlin_import_exception", None) + kwargs = { + "bits": 4, + "group_size": 32, + "desc_act": False, + "sym": True, + "in_features": 288, + "out_features": 200, + "bias": False, + "dtype": torch.float16, + } + + with pytest.raises(NotImplementedError, match="request GPTQ_MARLIN explicitly"): + marlin_qlinear_module.MarlinLinear(**kwargs, backend=BACKEND.AUTO) + + explicit = marlin_qlinear_module.MarlinLinear( + **kwargs, backend=BACKEND.GPTQ_MARLIN + ) + assert explicit.in_features == 288 + assert explicit.out_features == 200 + + aligned = marlin_qlinear_module.MarlinLinear( + **dict(kwargs, in_features=128, out_features=64), + backend=BACKEND.AUTO, + ) + assert aligned.in_features == 128 + assert aligned.out_features == 64 + + def test_marlin_quant_linear_validate_device_allows_sm75(monkeypatch): monkeypatch.setattr(marlin_qlinear_module, "IS_ROCM", False) monkeypatch.setattr(torch.cuda, "device_count", lambda: 2) @@ -427,6 +555,115 @@ def test_marlin_quant_linear_post_init_uses_compute_dtype_for_repack(monkeypatch module.post_init() assert captured == {"dtype": torch.bfloat16, "shape": tuple(module.qweight.shape)} + assert module._marlin_tile_padding is None + + +def test_marlin_quant_linear_post_init_pads_weight_scales_and_bias(monkeypatch): + captured = {} + + monkeypatch.setattr(marlin_qlinear_module, "marlin_import_exception", None) + monkeypatch.setattr(marlin_qlinear_module, "marlin_runtime_available", lambda dtype: True) + monkeypatch.setattr(marlin_qlinear_module, "marlin_runtime_error", lambda dtype: "") + monkeypatch.setattr( + marlin_qlinear_module, + "marlin_make_workspace_new", + lambda device: torch.zeros(128, dtype=torch.int32, device=device), + ) + + def fake_repack(b_q_weight, perm, size_k, size_n, num_bits, dtype=None): + captured["qweight"] = (tuple(b_q_weight.shape), size_k, size_n, dtype) + pack_factor = 32 // num_bits + return torch.zeros( + (size_k // 16, size_n * 16 // pack_factor), + dtype=torch.int32, + device=b_q_weight.device, + ) + + def fake_permute_scales(scales, size_k, size_n, group_size): + captured["scales"] = ( + tuple(scales.shape), + size_k, + size_n, + group_size, + ) + return scales + + monkeypatch.setattr(marlin_qlinear_module, "gptq_marlin_repack", fake_repack) + monkeypatch.setattr( + marlin_qlinear_module, "marlin_permute_scales", fake_permute_scales + ) + monkeypatch.setattr(marlin_qlinear_module, "marlin_permute_bias", lambda bias: bias) + + module = marlin_qlinear_module.MarlinLinear( + bits=4, + group_size=32, + desc_act=False, + sym=True, + in_features=288, + out_features=200, + bias=True, + dtype=torch.float16, + ) + module.post_init() + + assert module.in_features == 288 + assert module.out_features == 200 + assert module.qweight.shape == (20, 512) + assert module.scales.shape == (10, 256) + assert module.bias.shape == (256,) + assert module._marlin_tile_padding == (256, 320) + assert captured == { + "qweight": ((40, 256), 320, 256, torch.float16), + "scales": ((10, 256), 320, 256, 32), + } + + +def test_apply_gptq_marlin_linear_pads_input_and_slices_output(monkeypatch): + captured = {} + + def fake_gemm(a, _c, _weight, bias, _scales, _global_scale, + _weight_zp, _g_idx, _sort_indices, _workspace, _wtype, + **kwargs): + captured.update( + { + "input_shape": tuple(a.shape), + "bias_shape": tuple(bias.shape), + "size_m": kwargs["size_m"], + "size_n": kwargs["size_n"], + "size_k": kwargs["size_k"], + } + ) + return torch.ones( + (kwargs["size_m"], kwargs["size_n"]), dtype=a.dtype + ) + + monkeypatch.setattr(marlin_utils, "gptq_marlin_gemm", fake_gemm) + + output = marlin_utils.apply_gptq_marlin_linear_padded( + input=torch.randn(2, 3, 288, dtype=torch.float16), + weight=torch.zeros((20, 512), dtype=torch.int32), + weight_scale=torch.ones((10, 256), dtype=torch.float16), + weight_zp=torch.empty(0, dtype=torch.int32), + g_idx=torch.empty(0, dtype=torch.int32), + g_idx_sort_indices=torch.empty(0, dtype=torch.int32), + workspace=torch.zeros(128, dtype=torch.int32), + wtype=scalar_types.uint4b8, + output_size_per_partition=200, + input_size_per_partition=288, + is_k_full=True, + bias=torch.zeros(256, dtype=torch.float16), + tile_padding=(256, 320), + ) + + assert captured == { + "input_shape": (6, 320), + "bias_shape": (256,), + "size_m": 6, + "size_n": 256, + "size_k": 320, + } + assert output.shape == (2, 3, 200) + assert output.is_contiguous() def test_marlin_quant_linear_registers_runtime_buffers_in_compute_dtype(monkeypatch): @@ -622,6 +859,97 @@ def test_marlin_cuda_smoke_build_and_forward(monkeypatch, tmp_path): assert out.dtype == dtype +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("bits", [4, 8]) +@pytest.mark.parametrize( + "in_features,out_features,group_size", + [ + (256, 200, 32), # N tail + (208, 256, -1), # K tail with channelwise scales + (208, 256, 208), # Explicit K-sized channelwise group + (288, 200, 32), # N and K tails + ], +) +def test_marlin_cuda_padded_shape_matches_dequantized_reference( + dtype, bits, in_features, out_features, group_size +): + capability = torch.cuda.get_device_capability() + if capability[0] < 8 and dtype == torch.bfloat16: + pytest.skip("Marlin BF16 requires compute capability >= 8.0") + if not marlin_utils.marlin_runtime_available(dtype): + pytest.skip(marlin_utils.marlin_runtime_error(dtype)) + + torch.manual_seed(17) + device = torch.device("cuda:0") + pack_factor = 32 // bits + num_groups = 1 if group_size == -1 else in_features // group_size + scale_group_size = in_features if group_size == -1 else group_size + + codes = torch.randint( + 1, + 1 << bits, + (in_features, out_features), + dtype=torch.int32, + device=device, + ) + qweight = torch.zeros( + (in_features // pack_factor, out_features), + dtype=torch.int32, + device=device, + ) + for lane in range(pack_factor): + qweight.bitwise_or_(codes[lane::pack_factor] << (lane * bits)) + + scales = ( + torch.rand( + (num_groups, out_features), + device=device, + dtype=torch.float32, + ) + * 0.02 + + 0.002 + ).to(dtype) + bias = (torch.randn(out_features, device=device) * 0.01).to(dtype) + + module = marlin_qlinear_module.MarlinLinear( + bits=bits, + group_size=group_size, + desc_act=False, + sym=True, + in_features=in_features, + out_features=out_features, + bias=True, + dtype=dtype, + ).to(device) + with torch.no_grad(): + module.qweight.copy_(qweight) + module.scales.copy_(scales) + module.g_idx.copy_( + torch.arange(in_features, device=device, dtype=torch.int32) + // scale_group_size + ) + module.qzeros.zero_() + module.bias.copy_(bias) + module.post_init() + + x = torch.randn((8, in_features), device=device, dtype=dtype) / in_features**0.5 + dense_weight = (codes.to(dtype) - (1 << (bits - 1))) * scales.repeat_interleave( + scale_group_size, dim=0 + ) + expected = x @ dense_weight + bias + with torch.inference_mode(): + actual = module(x) + repeated = module(x) + torch.cuda.synchronize(device) + + assert actual.shape == (8, out_features) + assert actual.dtype == dtype + torch.testing.assert_close(actual, expected, rtol=5e-2, atol=5e-2) + torch.testing.assert_close(repeated, expected, rtol=5e-2, atol=5e-2) + + def test_marlin_include_paths_use_wheel_headers_when_local_cuda_is_incomplete(monkeypatch, tmp_path): root = tmp_path / "marlin" local_cuda_include = tmp_path / "local_cuda_include"