@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9eaacb23ec
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Kernel] Consolidate Marlin thread-tile padding across all dense Marlin paths (#45295)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Jun 12, 2026 at 21:46 UTC 9eaacb23ec1826ddac31657e0eab699de6de3c59
7 files changed +770 -66
tests/kernels/quantization/test_marlin_tile_padding.py new
+470
@@ -0,0 +1,470 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for Marlin thread-tile padding of TP-sharded weight shapes.
4 +
5 +Run `pytest tests/kernels/quantization/test_marlin_tile_padding.py`.
6 +"""
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm import _custom_ops as ops
12 +from vllm.model_executor.layers.quantization.utils.marlin_utils import (
13 + GPTQ_MARLIN_TILE,
14 + apply_gptq_marlin_linear,
15 + marlin_make_empty_g_idx,
16 + marlin_make_workspace_new,
17 + marlin_pad_qweight,
18 + marlin_pad_scales,
19 + marlin_padded_nk,
20 + marlin_permute_scales,
21 + marlin_repacked_nk,
22 + marlin_zero_points,
23 +)
24 +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
25 + apply_fp4_marlin_linear,
26 + is_fp4_marlin_supported,
27 + prepare_fp4_layer_for_marlin,
28 +)
29 +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
30 + apply_fp8_marlin_linear,
31 + apply_mxfp8_marlin_linear,
32 + is_fp8_marlin_supported,
33 + prepare_fp8_layer_for_marlin,
34 + prepare_mxfp8_layer_for_marlin,
35 +)
36 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
37 + gptq_pack,
38 + gptq_quantize_weights,
39 + quantize_weights,
40 +)
41 +from vllm.platforms import current_platform
42 +from vllm.scalar_type import scalar_types
43 +
44 +# (size_n, size_k) rank-local shapes that violate Marlin tile alignment,
45 +# e.g. produced by TP-sharding dims that are valid at TP=1.
46 +ODD_SHAPES = [
47 + (200, 288), # N padded
48 + (256, 208), # K padded
49 + (200, 208), # both padded
50 + (4640, 512), # Nemotron-Super-120B q_proj shard at TP=4
51 +]
52 +ALIGNED_SHAPES = [(64, 128), (128, 64), (256, 256), (4608, 4096)]
53 +
54 +
55 +def _is_tile_aligned(size_n: int, size_k: int) -> bool:
56 + return (size_n % 64 == 0 and size_k % 128 == 0) or (
57 + size_n % 128 == 0 and size_k % 64 == 0
58 + )
59 +
60 +
61 +@pytest.mark.parametrize("shape", ODD_SHAPES + ALIGNED_SHAPES)
62 +@pytest.mark.parametrize("group_size", [-1, 16, 32, 64, 128])
63 +def test_marlin_padded_nk(shape, group_size):
64 + size_n, size_k = shape
65 + padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
66 +
67 + assert padded_n >= size_n and padded_k >= size_k
68 + assert _is_tile_aligned(padded_n, padded_k)
69 + if group_size > 0:
70 + assert padded_k % group_size == 0
71 +
72 + # Aligned shapes must pass through unchanged (zero hot-path cost).
73 + if _is_tile_aligned(size_n, size_k) and (
74 + group_size <= 0 or size_k % group_size == 0
75 + ):
76 + assert (padded_n, padded_k) == (size_n, size_k)
77 +
78 + # Minimal: no valid shape with a smaller padded area exists.
79 + area = padded_n * padded_k
80 + for cand_n in range(size_n, padded_n + 1):
81 + for cand_k in range(size_k, padded_k + 1):
82 + if (
83 + _is_tile_aligned(cand_n, cand_k)
84 + and (group_size <= 0 or cand_k % group_size == 0)
85 + and cand_n * cand_k < area
86 + ):
87 + pytest.fail(f"({cand_n}, {cand_k}) beats ({padded_n}, {padded_k})")
88 +
89 + # Apply-time derivation from the repacked-tensor shape must round-trip.
90 + for num_bits in (4, 8):
91 + pack_factor = 32 // num_bits
92 + repacked_shape = (
93 + padded_k // GPTQ_MARLIN_TILE,
94 + padded_n * GPTQ_MARLIN_TILE // pack_factor,
95 + )
96 + repacked = torch.empty(repacked_shape, device="meta")
97 + assert marlin_repacked_nk(repacked, num_bits) == (padded_n, padded_k)
98 +
99 +
100 +def test_marlin_pad_helpers_shapes():
101 + size_n, size_k, group_size = 200, 208, 16
102 + padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
103 +
104 + qweight = torch.zeros(size_k // 8, size_n, dtype=torch.int32)
105 + padded = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
106 + assert padded.shape == (padded_k // 8, padded_n)
107 +
108 + scales = torch.ones(size_k // group_size, size_n)
109 + padded = marlin_pad_scales(scales, size_n, size_k, padded_n, padded_k, group_size)
110 + assert padded.shape == (padded_k // group_size, padded_n)
111 + assert padded[:, size_n:].abs().sum() == 0
112 +
113 + channelwise = torch.ones(1, size_n)
114 + padded = marlin_pad_scales(channelwise, size_n, size_k, padded_n, padded_k, -1)
115 + assert padded.shape == (1, padded_n)
116 +
117 +
118 +def _gpu_marlin_unsupported() -> bool:
119 + return not (
120 + current_platform.is_cuda() and current_platform.has_device_capability(80)
121 + )
122 +
123 +
124 +@pytest.mark.skipif(
125 + _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
126 + reason="FP8 Marlin is not supported on this GPU type.",
127 +)
128 +@pytest.mark.parametrize("shape", ODD_SHAPES)
129 +@pytest.mark.parametrize("use_bias", [False, True])
130 +def test_fp8_marlin_padded_round_trip(shape, use_bias):
131 + size_n, size_k = shape
132 + dtype = torch.float16
133 + layer = torch.nn.Module()
134 + layer.output_size_per_partition = size_n
135 + layer.input_size_per_partition = size_k
136 + layer.orig_dtype = dtype
137 +
138 + weight = torch.randn(size_k, size_n, dtype=dtype, device="cuda") / size_k**0.5
139 + scale = weight.abs().max() / 448
140 + weight_fp8 = (weight / scale).to(torch.float8_e4m3fn)
141 + layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
142 + layer.weight_scale = torch.nn.Parameter(
143 + scale.to(torch.float32), requires_grad=False
144 + )
145 + bias = None
146 + if use_bias:
147 + bias = torch.randn(size_n, dtype=dtype, device="cuda")
148 + layer.bias = torch.nn.Parameter(bias.clone(), requires_grad=False)
149 +
150 + prepare_fp8_layer_for_marlin(layer, size_k_first=True)
151 +
152 + x = torch.randn(8, size_k, dtype=dtype, device="cuda")
153 + output = apply_fp8_marlin_linear(
154 + input=x,
155 + weight=layer.weight,
156 + weight_scale=layer.weight_scale,
157 + workspace=layer.workspace,
158 + size_n=size_n,
159 + size_k=size_k,
160 + bias=layer.bias if use_bias else None,
161 + )
162 + ref = x @ (weight_fp8.to(dtype) * scale.to(dtype))
163 + if use_bias:
164 + ref = ref + bias
165 +
166 + assert output.shape == (8, size_n)
167 + torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
168 +
169 +
170 +def _dequant_fp4(packed: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
171 + """Dequantize packed e2m1 nibbles (N, K // 2) -> (N, K) in dtype."""
172 + lo = (packed & 0b10000000) | ((packed & 0b01110000) >> 2)
173 + lo = lo.view(torch.float8_e4m3fn).to(dtype) * (2**6)
174 + hi_bits = packed << 4
175 + hi = (hi_bits & 0b10000000) | ((hi_bits & 0b01110000) >> 2)
176 + hi = hi.view(torch.float8_e4m3fn).to(dtype) * (2**6)
177 + return torch.cat([hi.unsqueeze(2), lo.unsqueeze(2)], 2).view(packed.size(0), -1)
178 +
179 +
180 +@pytest.mark.skipif(
181 + _gpu_marlin_unsupported() or not is_fp4_marlin_supported(),
182 + reason="FP4 Marlin is not supported on this GPU type.",
183 +)
184 +@pytest.mark.parametrize("shape", ODD_SHAPES)
185 +def test_nvfp4_marlin_padded_round_trip(shape):
186 + size_n, size_k = shape
187 + group_size = 16
188 + dtype = torch.float16
189 + layer = torch.nn.Module()
190 + layer.output_size_per_partition = size_n
191 + layer.input_size_per_partition = size_k
192 + layer.params_dtype = dtype
193 +
194 + packed = torch.randint(
195 + 0, 256, (size_n, size_k // 2), dtype=torch.uint8, device="cuda"
196 + )
197 + scales = (torch.rand(size_n, size_k // group_size, device="cuda") + 0.25).to(
198 + torch.float8_e4m3fn
199 + )
200 + global_scale = torch.tensor([0.002], dtype=torch.float32, device="cuda")
201 +
202 + ref_weight = (
203 + _dequant_fp4(packed, dtype)
204 + * scales.to(dtype).repeat_interleave(group_size, 1)
205 + * global_scale.to(dtype)
206 + )
207 +
208 + layer.weight = torch.nn.Parameter(packed, requires_grad=False)
209 + layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
210 + layer.weight_global_scale = torch.nn.Parameter(global_scale, requires_grad=False)
211 +
212 + prepare_fp4_layer_for_marlin(layer)
213 +
214 + x = torch.randn(8, size_k, dtype=dtype, device="cuda") / size_k**0.5
215 + output = apply_fp4_marlin_linear(
216 + input=x,
217 + weight=layer.weight,
218 + weight_scale=layer.weight_scale,
219 + weight_global_scale=layer.weight_global_scale,
220 + workspace=layer.workspace,
221 + size_n=size_n,
222 + size_k=size_k,
223 + )
224 + ref = x @ ref_weight.T
225 +
226 + assert output.shape == (8, size_n)
227 + torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
228 +
229 +
230 +@pytest.mark.skipif(
231 + _gpu_marlin_unsupported(),
232 + reason="Marlin is not supported on this GPU type.",
233 +)
234 +@pytest.mark.parametrize("shape", ODD_SHAPES)
235 +@pytest.mark.parametrize("group_size", [-1, 128])
236 +def test_gptq_marlin_padded_round_trip(shape, group_size):
237 + """Pad-then-repack a GPTQ int4 weight the way MarlinLinearKernel does and
238 + check the GEMM against the dequantized reference.
239 +
240 + Symmetric int4's quantized zero decodes to -8, so this exercises the
241 + zero-padded-scales cancellation, not just zero weights.
242 + """
243 + size_n, size_k = shape
244 + if group_size > 0 and size_k % group_size != 0:
245 + pytest.skip("group must divide the rank-local K (not fixable by padding)")
246 + dtype = torch.float16
247 + quant_type = scalar_types.uint4b8
248 + device = torch.device("cuda")
249 +
250 + weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5
251 + w_ref, q_w, s, _, _ = gptq_quantize_weights(
252 + weight, quant_type, group_size, act_order=False
253 + )
254 + qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
255 +
256 + padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
257 + qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
258 + marlin_qweight = ops.gptq_marlin_repack(
259 + b_q_weight=qweight,
260 + perm=torch.empty(0, dtype=torch.int, device=device),
261 + size_k=padded_k,
262 + size_n=padded_n,
263 + num_bits=quant_type.size_bits,
264 + )
265 + s = marlin_pad_scales(s, size_n, size_k, padded_n, padded_k, group_size)
266 + marlin_s = marlin_permute_scales(
267 + s, size_k=padded_k, size_n=padded_n, group_size=group_size
268 + )
269 +
270 + x = torch.randn(8, size_k, dtype=dtype, device=device)
271 + output = apply_gptq_marlin_linear(
272 + input=x,
273 + weight=marlin_qweight,
274 + weight_scale=marlin_s,
275 + weight_zp=marlin_make_empty_g_idx(device),
276 + g_idx=marlin_make_empty_g_idx(device),
277 + g_idx_sort_indices=marlin_make_empty_g_idx(device),
278 + workspace=marlin_make_workspace_new(device),
279 + wtype=quant_type,
280 + output_size_per_partition=size_n,
281 + input_size_per_partition=size_k,
282 + is_k_full=True,
283 + )
284 + ref = x @ w_ref
285 +
286 + assert output.shape == (8, size_n)
287 + torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
288 +
289 +
290 +@pytest.mark.skipif(
291 + _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
292 + reason="FP8 Marlin is not supported on this GPU type.",
293 +)
294 +@pytest.mark.parametrize("shape", [(200, 512), (4640, 512)])
295 +def test_fp8_block_marlin_padded_round_trip(shape):
296 + """Block-quantized FP8 (e.g. Nemotron NVFP4 checkpoints' FP8 layers):
297 + group_size=128 exercises the lcm K-alignment in marlin_padded_nk and the
298 + weight_scale_inv group-wise scale padding."""
299 + size_n, size_k = shape
300 + block = 128
301 + dtype = torch.float16
302 + layer = torch.nn.Module()
303 + layer.output_size_per_partition = size_n
304 + layer.input_size_per_partition = size_k
305 + layer.orig_dtype = dtype
306 + layer.weight_block_size = [block, block]
307 +
308 + weight = torch.randn(size_n, size_k, dtype=dtype, device="cuda") / size_k**0.5
309 + n_blocks, k_blocks = (size_n + block - 1) // block, size_k // block
310 + padded = torch.zeros(n_blocks * block, size_k, dtype=dtype, device="cuda")
311 + padded[:size_n] = weight
312 + scales = padded.view(n_blocks, block, k_blocks, block).abs().amax(dim=(1, 3)) / 448
313 + scales_expanded = scales.repeat_interleave(block, 0)[:size_n].repeat_interleave(
314 + block, 1
315 + )
316 + weight_fp8 = (weight / scales_expanded).to(torch.float8_e4m3fn)
317 +
318 + layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
319 + layer.weight_scale_inv = torch.nn.Parameter(
320 + scales.to(torch.float32), requires_grad=False
321 + )
322 +
323 + prepare_fp8_layer_for_marlin(layer, size_k_first=False)
324 +
325 + x = torch.randn(8, size_k, dtype=dtype, device="cuda")
326 + output = apply_fp8_marlin_linear(
327 + input=x,
328 + weight=layer.weight,
329 + weight_scale=layer.weight_scale_inv,
330 + workspace=layer.workspace,
331 + size_n=size_n,
332 + size_k=size_k,
333 + bias=None,
334 + )
335 + ref = x @ (weight_fp8.to(dtype) * scales_expanded.to(dtype)).T
336 +
337 + assert output.shape == (8, size_n)
338 + torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
339 +
340 +
341 +@pytest.mark.skipif(
342 + _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
343 + reason="FP8 Marlin is not supported on this GPU type.",
344 +)
345 +@pytest.mark.parametrize("shape", [(200, 288), (4640, 512)])
346 +def test_mxfp8_marlin_padded_round_trip(shape):
347 + """MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to
348 + 2^-127 instead of zero and must still contribute nothing."""
349 + size_n, size_k = shape
350 + group_size = 32
351 + # The e8m0-scale Marlin kernels are only instantiated for bf16 activations.
352 + dtype = torch.bfloat16
353 + layer = torch.nn.Module()
354 + layer.output_size_per_partition = size_n
355 + layer.input_size_per_partition = size_k
356 +
357 + weight_fp8 = (torch.randn(size_n, size_k, dtype=dtype, device="cuda") / 4).to(
358 + torch.float8_e4m3fn
359 + )
360 + # e8m0 exponents around 1.0 (127): scales in [2^-6, 2^0]
361 + scales = torch.randint(
362 + 121, 128, (size_n, size_k // group_size), dtype=torch.uint8, device="cuda"
363 + )
364 + ref_weight = weight_fp8.to(dtype) * (
365 + 2.0 ** (scales.to(dtype) - 127)
366 + ).repeat_interleave(group_size, 1)
367 +
368 + layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
369 + layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
370 +
371 + prepare_mxfp8_layer_for_marlin(layer)
372 +
373 + x = torch.randn(8, size_k, dtype=dtype, device="cuda") / size_k**0.5
374 + output = apply_mxfp8_marlin_linear(
375 + input=x,
376 + weight=layer.weight,
377 + weight_scale=layer.weight_scale,
378 + workspace=layer.workspace,
379 + size_n=size_n,
380 + size_k=size_k,
381 + )
382 + ref = x @ ref_weight.T
383 +
384 + assert output.shape == (8, size_n)
385 + torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
386 +
387 +
388 +@pytest.mark.skipif(
389 + _gpu_marlin_unsupported(),
390 + reason="Marlin is not supported on this GPU type.",
391 +)
392 +@pytest.mark.parametrize("shape", [(200, 512), (4640, 512)])
393 +def test_awq_zp_marlin_padded_round_trip(shape):
394 + """AWQ-style uint4 with runtime zero-points, padded the way
395 + MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0."""
396 + size_n, size_k = shape
397 + group_size = 128
398 + dtype = torch.float16
399 + quant_type = scalar_types.uint4
400 + device = torch.device("cuda")
401 +
402 + weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5
403 + w_ref, q_w, s, zp = quantize_weights(
404 + weight, quant_type, group_size, zero_points=True
405 + )
406 + qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
407 +
408 + padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
409 + qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
410 + marlin_qweight = ops.gptq_marlin_repack(
411 + b_q_weight=qweight,
412 + perm=torch.empty(0, dtype=torch.int, device=device),
413 + size_k=padded_k,
414 + size_n=padded_n,
415 + num_bits=quant_type.size_bits,
416 + )
417 + s = marlin_pad_scales(s, size_n, size_k, padded_n, padded_k, group_size)
418 + marlin_s = marlin_permute_scales(
419 + s, size_k=padded_k, size_n=padded_n, group_size=group_size
420 + )
421 + zp = marlin_pad_scales(zp, size_n, size_k, padded_n, padded_k, group_size)
422 + marlin_zp = marlin_zero_points(
423 + zp,
424 + size_k=padded_k // group_size,
425 + size_n=padded_n,
426 + num_bits=quant_type.size_bits,
427 + )
428 +
429 + x = torch.randn(8, size_k, dtype=dtype, device=device)
430 + output = apply_gptq_marlin_linear(
431 + input=x,
432 + weight=marlin_qweight,
433 + weight_scale=marlin_s,
434 + weight_zp=marlin_zp,
435 + g_idx=marlin_make_empty_g_idx(device),
436 + g_idx_sort_indices=marlin_make_empty_g_idx(device),
437 + workspace=marlin_make_workspace_new(device),
438 + wtype=quant_type,
439 + output_size_per_partition=size_n,
440 + input_size_per_partition=size_k,
441 + is_k_full=True,
442 + )
443 + ref = x @ w_ref
444 +
445 + assert output.shape == (8, size_n)
446 + torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
447 +
448 +
449 +class _FakeLinear:
450 + def __init__(self, size_n, size_k, input_size=None):
451 + self.output_size_per_partition = size_n
452 + self.input_size_per_partition = size_k
453 + self.output_size = size_n
454 + self.input_size = input_size if input_size is not None else size_k
455 +
456 +
457 +def test_check_marlin_supports_layer_allow_tile_padding():
458 + from vllm.model_executor.layers.quantization.utils.marlin_utils import (
459 + check_marlin_supports_layer,
460 + )
461 +
462 + # Tile-misaligned but group-aligned: rejected strictly, allowed w/ padding
463 + layer = _FakeLinear(4640, 512, input_size=2048)
464 + assert not check_marlin_supports_layer(layer, 128)
465 + assert check_marlin_supports_layer(layer, 128, allow_tile_padding=True)
466 + assert check_marlin_supports_layer(layer, -1, allow_tile_padding=True)
467 +
468 + # A group straddling the TP shard cannot be fixed by padding
469 + layer = _FakeLinear(4608, 4672, input_size=18688)
470 + assert not check_marlin_supports_layer(layer, 128, allow_tile_padding=True)
vllm/model_executor/kernels/linear/mixed_precision/marlin.py
+68 -23
@@ -13,6 +13,10 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
13 marlin_is_k_full,
14 marlin_make_empty_g_idx,
15 marlin_make_workspace_new,
16 + marlin_pad_dim,
17 + marlin_pad_qweight,
18 + marlin_pad_scales,
19 + marlin_padded_nk,
20 marlin_permute_bias,
21 marlin_permute_scales,
22 marlin_sort_g_idx,
@@ -54,12 +58,29 @@ class MarlinLinearKernel(MPLinearKernel):
58 f"{MARLIN_SUPPORTED_GROUP_SIZES}",
59 )
60
57 - return check_marlin_supports_shape(
58 - c.partition_weight_shape[1], # out_features
59 - c.partition_weight_shape[0], # in_features
60 - c.full_weight_shape[0], # in_features
61 - c.group_size,
62 - )
61 + if c.has_g_idx:
62 + # Act-order couples K to the full-model group layout, so tile
63 + # padding is not supported; keep the strict shape check.
64 + return check_marlin_supports_shape(
65 + c.partition_weight_shape[1], # out_features
66 + c.partition_weight_shape[0], # in_features
67 + c.full_weight_shape[0], # in_features
68 + c.group_size,
69 + )
70 +
71 + # A group straddling TP ranks cannot be fixed by padding.
72 + if (
73 + c.group_size != -1
74 + and c.group_size < c.full_weight_shape[0]
75 + and c.partition_weight_shape[0] % c.group_size != 0
76 + ):
77 + return False, (
78 + f"in_features per partition {c.partition_weight_shape[0]} is "
79 + f"not divisible by group_size = {c.group_size}."
80 + )
81 +
82 + # Tile misalignment is fixed by zero-padding at weight prep.
83 + return True, None
84
85 # note assumes that
86 # `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
@@ -83,6 +104,13 @@ class MarlinLinearKernel(MPLinearKernel):
104 row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0]
105 self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel)
106
107 + size_k, size_n = c.partition_weight_shape
108 + if c.has_g_idx:
109 + # Act-order shapes were strictly validated in can_implement.
110 + padded_n, padded_k = size_n, size_k
111 + else:
112 + padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
113 +
114 # Allocate marlin workspace.
115 self.workspace = marlin_make_workspace_new(device)
116
@@ -97,10 +125,12 @@ class MarlinLinearKernel(MPLinearKernel):
125 assert isinstance(x, BasevLLMParameter)
126 permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
127 x.data = ops.gptq_marlin_repack(
100 - x.data.contiguous(),
128 + marlin_pad_qweight(
129 + x.data.contiguous(), size_n, size_k, padded_n, padded_k
130 + ),
131 perm=layer.g_idx_sort_indices,
102 - size_k=c.partition_weight_shape[0],
103 - size_n=c.partition_weight_shape[1],
132 + size_k=padded_k,
133 + size_n=padded_n,
134 num_bits=c.weight_type.size_bits,
135 is_a_8bit=is_a_8bit,
136 )
@@ -110,9 +140,16 @@ class MarlinLinearKernel(MPLinearKernel):
140 assert isinstance(x, BasevLLMParameter)
141 permute_param_layout_(x, input_dim=0, output_dim=1)
142 x.data = marlin_permute_scales(
113 - x.data.contiguous(),
114 - size_k=c.partition_weight_shape[0],
115 - size_n=c.partition_weight_shape[1],
143 + marlin_pad_scales(
144 + x.data.contiguous(),
145 + size_n,
146 + size_k,
147 + padded_n,
148 + padded_k,
149 + c.group_size,
150 + ),
151 + size_k=padded_k,
152 + size_n=padded_n,
153 group_size=c.group_size,
154 is_a_8bit=is_a_8bit,
155 )
@@ -143,21 +180,27 @@ class MarlinLinearKernel(MPLinearKernel):
180 layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
181
182 if c.zero_points:
146 - grouped_k = (
147 - c.partition_weight_shape[0] // c.group_size if c.group_size != -1 else 1
148 - )
183 + grouped_k = size_k // c.group_size if c.group_size != -1 else 1
184 + padded_grouped_k = padded_k // c.group_size if c.group_size != -1 else 1
185 self._transform_param(
186 layer,
187 self.w_zp_name,
188 lambda x: marlin_zero_points(
153 - unpack_cols(
154 - x.t(),
155 - c.weight_type.size_bits,
156 - grouped_k,
157 - c.partition_weight_shape[1],
189 + marlin_pad_scales(
190 + unpack_cols(
191 + x.t(),
192 + c.weight_type.size_bits,
193 + grouped_k,
194 + size_n,
195 + ),
196 + size_n,
197 + size_k,
198 + padded_n,
199 + padded_k,
200 + c.group_size,
201 ),
159 - size_k=grouped_k,
160 - size_n=c.partition_weight_shape[1],
202 + size_k=padded_grouped_k,
203 + size_n=padded_n,
204 num_bits=c.weight_type.size_bits,
205 is_a_8bit=is_a_8bit,
206 ),
@@ -168,7 +211,9 @@ class MarlinLinearKernel(MPLinearKernel):
211 self._transform_param(layer, self.w_s_name, transform_w_s)
212
213 if hasattr(layer, "bias") and layer.bias is not None:
171 - layer.bias.data = marlin_permute_bias(layer.bias)
214 + layer.bias.data = marlin_permute_bias(
215 + marlin_pad_dim(layer.bias, size_n, padded_n)
216 + )
217
218 def apply_weights(
219 self,
vllm/model_executor/layers/quantization/awq_marlin.py
+5 -2
@@ -289,8 +289,11 @@ class AWQMarlinConfig(QuantizationConfig):
289 skip_with_substr=True,
290 ):
291 return UnquantizedLinearMethod()
292 - # Check if the layer is supported by AWQMarlin.
293 - if not check_marlin_supports_layer(layer, self.group_size):
292 + # Check if the layer is supported by AWQMarlin; tile-misaligned
293 + # shapes are fixed by padding at weight prep.
294 + if not check_marlin_supports_layer(
295 + layer, self.group_size, allow_tile_padding=True
296 + ):
297 logger.warning_once(
298 "Layer '%s' is not supported by AWQMarlin. Falling back to unoptimized AWQ kernels.", # noqa: E501
299 prefix,
vllm/model_executor/layers/quantization/modelopt.py
+1
@@ -468,6 +468,7 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
468 layer.logical_widths = output_partition_sizes
469 layer.input_size_per_partition = input_size_per_partition
470 layer.output_size_per_partition = output_size_per_partition
471 + layer.orig_dtype = params_dtype
472 weight_dtype = (
473 torch.float8_e4m3fn
474 if self.quant_config.is_checkpoint_fp8_serialized
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+117 -9
@@ -2,6 +2,8 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 +import math
6 +
7 import numpy
8 import torch
9
@@ -17,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
19 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
20 from vllm.platforms import current_platform
21 from vllm.scalar_type import ScalarType, scalar_types
22 +from vllm.utils.math_utils import round_up
23 from vllm.utils.platform_utils import num_compute_units
24
25 from .quant_utils import pack_cols, unpack_cols
@@ -214,7 +217,93 @@ def check_marlin_supports_shape(
217 return True, None
218
219
217 -def check_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool:
220 +def marlin_padded_nk(size_n: int, size_k: int, group_size: int = -1) -> tuple[int, int]:
221 + """Minimal (padded_n, padded_k) satisfying a Marlin thread-tile family.
222 +
223 + Marlin GEMM and repack require (n % 64, k % 128) or (n % 128, k % 64);
224 + shapes satisfying neither are zero-padded up to the cheaper family. K
225 + stays divisible by group_size so padded scales keep an integral group
226 + count. Padded weight regions contribute nothing to the GEMM output:
227 + quantized value 0 decodes to 0.0 (FP4/FP8) or is cancelled by the
228 + zero-padded scales/zero-points (INT).
229 + """
230 + group = group_size if group_size > 0 else 1
231 + candidates = (
232 + (round_up(size_n, 64), round_up(size_k, math.lcm(128, group))),
233 + (round_up(size_n, 128), round_up(size_k, math.lcm(64, group))),
234 + )
235 + padded_nk = min(candidates, key=lambda nk: (nk[0] * nk[1], nk[0] + nk[1]))
236 + if padded_nk != (size_n, size_k):
237 + logger.warning_once(
238 + "Marlin requires thread-tile padding for some weight shapes in "
239 + "this model. Activations and/or outputs of the padded layers are "
240 + "padded/sliced on every forward; performance may be degraded."
241 + )
242 + return padded_nk
243 +
244 +
245 +def marlin_repacked_nk(qweight: torch.Tensor, num_bits: int) -> tuple[int, int]:
246 + """Recover the (size_n, size_k) a Marlin weight was repacked with
247 + (including any tile padding) from its packed shape."""
248 + pack_factor = 32 // num_bits
249 + size_k = qweight.size(0) * GPTQ_MARLIN_TILE
250 + size_n = qweight.size(1) * pack_factor // GPTQ_MARLIN_TILE
251 + return size_n, size_k
252 +
253 +
254 +def marlin_pad_qweight(
255 + qweight: torch.Tensor, size_n: int, size_k: int, padded_n: int, padded_k: int
256 +) -> torch.Tensor:
257 + """Zero-pad a GPTQ-layout packed weight (size_k / pack, size_n) for
258 + gptq_marlin_repack."""
259 + if (padded_n, padded_k) == (size_n, size_k):
260 + return qweight
261 + pack_factor = size_k // qweight.size(0)
262 + return torch.nn.functional.pad(
263 + qweight, (0, padded_n - size_n, 0, (padded_k - size_k) // pack_factor)
264 + )
265 +
266 +
267 +def marlin_pad_scales(
268 + scales: torch.Tensor,
269 + size_n: int,
270 + size_k: int,
271 + padded_n: int,
272 + padded_k: int,
273 + group_size: int,
274 +) -> torch.Tensor:
275 + """Zero-pad weight scales (num_groups, size_n); call before
276 + marlin_permute_scales and pass the padded extents to it."""
277 + if (padded_n, padded_k) == (size_n, size_k):
278 + return scales
279 + pad_rows = padded_k // group_size - scales.size(0) if group_size > 0 else 0
280 + assert pad_rows >= 0
281 + return torch.nn.functional.pad(scales, (0, padded_n - size_n, 0, pad_rows))
282 +
283 +
284 +def marlin_pad_dim(x: torch.Tensor, size: int, padded: int) -> torch.Tensor:
285 + """Zero-pad the last dim from size to padded (activations K, bias N)."""
286 + if padded == size:
287 + return x
288 + return torch.nn.functional.pad(x, (0, padded - size))
289 +
290 +
291 +def marlin_unpad_output(
292 + output: torch.Tensor, size_n: int, padded_n: int
293 +) -> torch.Tensor:
294 + """Strip padded output columns back to the logical N.
295 +
296 + TODO: marlin_gemm could instead write the un-padded columns directly
297 + into a caller-provided `c` buffer so this slice copy disappears.
298 + """
299 + if padded_n == size_n:
300 + return output
301 + return output[..., :size_n].contiguous()
302 +
303 +
304 +def check_marlin_supports_layer(
305 + layer: LinearBase, group_size: int, allow_tile_padding: bool = False
306 +) -> bool:
307 output_size_per_partition = (
308 getattr(layer, "output_size_per_partition", None) or layer.output_size
309 )
@@ -222,6 +311,17 @@ def check_marlin_supports_layer(layer: LinearBase, group_size: int) -> bool:
311 getattr(layer, "input_size_per_partition", None) or layer.input_size
312 )
313
314 + if allow_tile_padding:
315 + # Thread-tile misalignment is fixed by zero-padding at weight prep
316 + # (see marlin_padded_nk); only a quantization group straddling the
317 + # TP shard remains unsupported. Dense layers only - MoE prep does
318 + # not pad yet.
319 + return (
320 + group_size == -1
321 + or group_size >= layer.input_size
322 + or input_size_per_partition % group_size == 0
323 + )
324 +
325 return check_marlin_supports_shape(
326 output_size_per_partition=output_size_per_partition,
327 input_size_per_partition=input_size_per_partition,
@@ -556,10 +656,13 @@ def apply_gptq_marlin_linear(
656 reshaped_x = input.reshape(-1, input.shape[-1])
657 out_shape = input.shape[:-1] + (output_size_per_partition,)
658
659 + padded_n, padded_k = marlin_repacked_nk(weight, wtype.size_bits)
660 + reshaped_x = marlin_pad_dim(reshaped_x, input_size_per_partition, padded_k)
661 +
662 use_atomic_add = should_use_atomic_add_reduce(
663 m=reshaped_x.size(0),
561 - n=output_size_per_partition,
562 - k=reshaped_x.size(1),
664 + n=padded_n,
665 + k=padded_k,
666 device=input.device,
667 dtype=input.dtype,
668 )
@@ -592,14 +695,15 @@ def apply_gptq_marlin_linear(
695 workspace,
696 wtype,
697 size_m=reshaped_x.shape[0],
595 - size_n=output_size_per_partition,
596 - size_k=input_size_per_partition,
698 + size_n=padded_n,
699 + size_k=padded_k,
700 is_k_full=is_k_full,
701 use_atomic_add=use_atomic_add,
702 use_fp32_reduce=use_fp32_reduce,
703 is_zp_float=False,
704 )
705
706 + output = marlin_unpad_output(output, output_size_per_partition, padded_n)
707 return output.reshape(out_shape)
708
709
@@ -622,10 +726,13 @@ def apply_awq_marlin_linear(
726 reshaped_x = input.reshape(-1, input.shape[-1])
727 out_shape = input.shape[:-1] + (output_size_per_partition,)
728
729 + padded_n, padded_k = marlin_repacked_nk(weight, quant_type.size_bits)
730 + reshaped_x = marlin_pad_dim(reshaped_x, input_size_per_partition, padded_k)
731 +
732 use_atomic_add = should_use_atomic_add_reduce(
733 m=reshaped_x.size(0),
627 - n=output_size_per_partition,
628 - k=reshaped_x.size(1),
734 + n=padded_n,
735 + k=padded_k,
736 device=input.device,
737 dtype=input.dtype,
738 )
@@ -657,11 +764,12 @@ def apply_awq_marlin_linear(
764 workspace,
765 quant_type,
766 size_m=reshaped_x.shape[0],
660 - size_n=output_size_per_partition,
661 - size_k=input_size_per_partition,
767 + size_n=padded_n,
768 + size_k=padded_k,
769 use_atomic_add=use_atomic_add,
770 use_fp32_reduce=use_fp32_reduce,
771 is_zp_float=False,
772 )
773
774 + output = marlin_unpad_output(output, output_size_per_partition, padded_n)
775 return output.reshape(out_shape)
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+64 -14
@@ -11,13 +11,20 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
11 USE_FP32_REDUCE_DEFAULT,
12 get_marlin_input_dtype,
13 marlin_make_workspace_new,
14 + marlin_pad_dim,
15 + marlin_pad_qweight,
16 + marlin_pad_scales,
17 + marlin_padded_nk,
18 marlin_permute_bias,
19 marlin_permute_scales,
20 marlin_quant_input,
21 + marlin_repacked_nk,
22 + marlin_unpad_output,
23 should_use_atomic_add_reduce,
24 )
25 from vllm.platforms import current_platform
26 from vllm.scalar_type import scalar_types
27 +from vllm.utils.math_utils import round_up
28
29 FP4_MARLIN_SUPPORTED_GROUP_SIZES = [16]
30
@@ -165,8 +172,15 @@ def apply_fp4_marlin_linear(
172 reshaped_x = input.reshape(-1, input.shape[-1])
173 out_shape = input.shape[:-1] + (size_n,)
174
175 + padded_n, padded_k = marlin_repacked_nk(weight, num_bits=4)
176 + reshaped_x = marlin_pad_dim(reshaped_x, size_k, padded_k)
177 +
178 use_atomic_add = should_use_atomic_add_reduce(
169 - m=reshaped_x.size(0), n=size_n, k=size_k, device=input.device, dtype=input.dtype
179 + m=reshaped_x.size(0),
180 + n=padded_n,
181 + k=padded_k,
182 + device=input.device,
183 + dtype=input.dtype,
184 )
185
186 inputs = reshaped_x
@@ -194,12 +208,13 @@ def apply_fp4_marlin_linear(
208 workspace=workspace,
209 b_q_type=scalar_types.float4_e2m1f,
210 size_m=reshaped_x.size(0),
197 - size_n=size_n,
198 - size_k=size_k,
211 + size_n=padded_n,
212 + size_k=padded_k,
213 use_atomic_add=use_atomic_add,
214 use_fp32_reduce=use_fp32_reduce,
215 )
216
217 + output = marlin_unpad_output(output, size_n, padded_n)
218 return output.reshape(out_shape)
219
220
@@ -217,6 +232,7 @@ def prepare_fp4_layer_for_marlin(
232
233 part_size_n = layer.output_size_per_partition
234 part_size_k = layer.input_size_per_partition
235 + padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size)
236 param_dtype = layer.params_dtype
237
238 assert layer.weight.shape == (part_size_n, part_size_k // 2)
@@ -230,13 +246,14 @@ def prepare_fp4_layer_for_marlin(
246 # Repack weights to marlin format
247 perm = torch.empty(0, dtype=torch.int, device=device)
248 qweight = layer.weight.view(torch.int32).T.contiguous()
249 + qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k)
250
251 is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
252 marlin_qweight = ops.gptq_marlin_repack(
253 b_q_weight=qweight,
254 perm=perm,
238 - size_k=part_size_k,
239 - size_n=part_size_n,
255 + size_k=padded_k,
256 + size_n=padded_n,
257 num_bits=4,
258 is_a_8bit=is_a_8bit,
259 )
@@ -250,10 +267,13 @@ def prepare_fp4_layer_for_marlin(
267 weight_scale = weight_scale.view(torch.float8_e8m0fnu)
268
269 weight_scale = weight_scale.to(param_dtype)
270 + weight_scale = marlin_pad_scales(
271 + weight_scale, part_size_n, part_size_k, padded_n, padded_k, group_size
272 + )
273 weight_scale = marlin_permute_scales(
274 s=weight_scale,
255 - size_k=part_size_k,
256 - size_n=part_size_n,
275 + size_k=padded_k,
276 + size_n=padded_n,
277 group_size=group_size,
278 is_a_8bit=is_a_8bit,
279 )
@@ -280,7 +300,7 @@ def prepare_fp4_layer_for_marlin(
300
301 if hasattr(layer, "bias") and layer.bias is not None:
302 assert layer.bias.shape == (part_size_n,)
283 - bias = marlin_permute_bias(layer.bias)
303 + bias = marlin_permute_bias(marlin_pad_dim(layer.bias, part_size_n, padded_n))
304 layer.bias = torch.nn.Parameter(bias, requires_grad=False)
305
306 return
@@ -313,6 +333,32 @@ def prepare_nvfp4_moe_layer_for_marlin(
333 E = layer.num_experts
334 K = layer.hidden_size
335 N = layer.intermediate_size_per_partition
336 + num_shards = 2 if is_act_and_mul else 1
337 +
338 + # Pad the rank-local intermediate size to satisfy Marlin thread tiles:
339 + # N is an output extent of w13 (per gate/up shard) and the input extent
340 + # of w2, so the padded region never reaches the MoE output.
341 + if K % 128 == 0:
342 + padded_N = round_up(N, 64)
343 + else:
344 + assert K % 64 == 0, f"hidden_size = {K} unsupported by Marlin tiles"
345 + padded_N = round_up(N, 128)
346 +
347 + def pad_w13(x: torch.Tensor) -> torch.Tensor:
348 + """Zero-pad each gate/up shard of a (E, num_shards * N, cols)
349 + tensor to padded_N rows."""
350 + if padded_N == N:
351 + return x
352 + x = x.view(E, num_shards, N, x.size(-1))
353 + x = torch.nn.functional.pad(x, (0, 0, 0, padded_N - N))
354 + return x.reshape(E, num_shards * padded_N, -1)
355 +
356 + def pad_w2(x: torch.Tensor, packing: int) -> torch.Tensor:
357 + """Zero-pad the packed N (last) dim of a (E, K, N / packing)
358 + tensor."""
359 + if padded_N == N:
360 + return x
361 + return torch.nn.functional.pad(x, (0, (padded_N - N) // packing))
362
363 device = w13.device
364 param_dtype = layer.params_dtype
@@ -326,13 +372,16 @@ def prepare_nvfp4_moe_layer_for_marlin(
372 # Repack weights to marlin format
373 def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
374 tensor_list = []
329 - num_shards = 2 if is_act_and_mul else 1
375 if "w13" in name:
376 size_n, size_k = N * num_shards, K
377 + assert weight.shape == (E, size_n, size_k // 2)
378 + weight = pad_w13(weight)
379 + size_n = padded_N * num_shards
380 else:
381 size_n, size_k = K, N
334 -
335 - assert weight.shape == (E, size_n, size_k // 2)
382 + assert weight.shape == (E, size_n, size_k // 2)
383 + weight = pad_w2(weight, packing=2)
384 + size_k = padded_N
385
386 for i in range(E):
387 qweight = weight[i].view(torch.int32).T.contiguous()
@@ -360,11 +409,12 @@ def prepare_nvfp4_moe_layer_for_marlin(
409 scales = scales.to(param_dtype)
410
411 tensor_list = []
363 - num_shards = 2 if is_act_and_mul else 1
412 if "w13" in name:
365 - size_n, size_k = N * num_shards, K
413 + scales = pad_w13(scales)
414 + size_n, size_k = padded_N * num_shards, K
415 else:
367 - size_n, size_k = K, N
416 + scales = pad_w2(scales, packing=GROUP_SIZE)
417 + size_n, size_k = K, padded_N
418
419 # All experts share one global_scale, so compute the max
420 # scale_factor across all experts first, then apply uniformly.
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+45 -18
@@ -10,8 +10,14 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
10 USE_FP32_REDUCE_DEFAULT,
11 get_marlin_input_dtype,
12 marlin_make_workspace_new,
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_repacked_nk,
20 + marlin_unpad_output,
21 should_use_atomic_add_reduce,
22 )
23 from vllm.model_executor.utils import replace_parameter
@@ -56,8 +62,15 @@ def apply_fp8_marlin_linear(
62 reshaped_x = input.reshape(-1, input.shape[-1])
63 out_shape = input.shape[:-1] + (size_n,)
64
65 + padded_n, padded_k = marlin_repacked_nk(weight, num_bits=8)
66 + reshaped_x = marlin_pad_dim(reshaped_x, size_k, padded_k)
67 +
68 use_atomic_add = should_use_atomic_add_reduce(
60 - m=reshaped_x.size(0), n=size_n, k=size_k, device=input.device, dtype=input.dtype
69 + m=reshaped_x.size(0),
70 + n=padded_n,
71 + k=padded_k,
72 + device=input.device,
73 + dtype=input.dtype,
74 )
75
76 inputs = reshaped_x
@@ -80,12 +93,13 @@ def apply_fp8_marlin_linear(
93 workspace=workspace,
94 b_q_type=scalar_types.float8_e4m3fn,
95 size_m=reshaped_x.size(0),
83 - size_n=size_n,
84 - size_k=size_k,
96 + size_n=padded_n,
97 + size_k=padded_k,
98 use_atomic_add=use_atomic_add,
99 use_fp32_reduce=use_fp32_reduce,
100 )
101
102 + output = marlin_unpad_output(output, size_n, padded_n)
103 return output.reshape(out_shape)
104
105
@@ -106,6 +120,8 @@ def prepare_fp8_layer_for_marlin(
120 part_size_n = layer.output_size_per_partition
121 part_size_k = layer.input_size_per_partition
122 weight_block_size = getattr(layer, "weight_block_size", None)
123 + group_size = -1 if weight_block_size is None else weight_block_size[1]
124 + padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size)
125
126 if size_k_first:
127 assert layer.weight.shape == (part_size_k, part_size_n)
@@ -123,12 +139,13 @@ def prepare_fp8_layer_for_marlin(
139 qweight = pack_fp8_to_int32(layer.weight, size_k_first)
140 if not size_k_first:
141 qweight = qweight.T.contiguous()
142 + qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k)
143
144 marlin_qweight = ops.gptq_marlin_repack(
145 b_q_weight=qweight,
146 perm=perm,
130 - size_k=part_size_k,
131 - size_n=part_size_n,
147 + size_k=padded_k,
148 + size_n=padded_n,
149 num_bits=8,
150 )
151 replace_parameter(layer, "weight", marlin_qweight)
@@ -140,8 +157,6 @@ def prepare_fp8_layer_for_marlin(
157 elif "weight_scale_inv" in dir(layer):
158 scales = layer.weight_scale_inv.to(layer.orig_dtype)
159
143 - group_size = -1 if weight_block_size is None else weight_block_size[1]
144 -
160 # marlin kernel only support channel-wise and group-wise quantization
161 # we need to convert the scales
162 if weight_block_size is None:
@@ -182,8 +197,11 @@ def prepare_fp8_layer_for_marlin(
197 # size_n may not divisible by block_size[0]
198 scales = scales[:, :part_size_n]
199
200 + scales = marlin_pad_scales(
201 + scales, part_size_n, part_size_k, padded_n, padded_k, group_size
202 + )
203 marlin_scales = marlin_permute_scales(
186 - s=scales, size_k=part_size_k, size_n=part_size_n, group_size=group_size
204 + s=scales, size_k=padded_k, size_n=padded_n, group_size=group_size
205 )
206 if input_dtype != torch.float8_e4m3fn:
207 marlin_scales = fp8_fused_exponent_bias_into_scales(marlin_scales)
@@ -194,7 +212,7 @@ def prepare_fp8_layer_for_marlin(
212
213 if hasattr(layer, "bias") and layer.bias is not None:
214 assert layer.bias.shape == (part_size_n,)
197 - bias = marlin_permute_bias(layer.bias)
215 + bias = marlin_permute_bias(marlin_pad_dim(layer.bias, part_size_n, padded_n))
216 replace_parameter(layer, "bias", bias)
217
218
@@ -359,10 +377,13 @@ def apply_mxfp8_marlin_linear(
377 reshaped_x = input.reshape(-1, input.shape[-1])
378 out_shape = input.shape[:-1] + (size_n,)
379
380 + padded_n, padded_k = marlin_repacked_nk(weight, num_bits=8)
381 + reshaped_x = marlin_pad_dim(reshaped_x, size_k, padded_k)
382 +
383 use_atomic_add = should_use_atomic_add_reduce(
384 m=reshaped_x.size(0),
364 - n=size_n,
365 - k=size_k,
385 + n=padded_n,
386 + k=padded_k,
387 device=input.device,
388 dtype=input.dtype,
389 )
@@ -381,12 +402,13 @@ def apply_mxfp8_marlin_linear(
402 workspace=workspace,
403 b_q_type=scalar_types.float8_e4m3fn,
404 size_m=reshaped_x.size(0),
384 - size_n=size_n,
385 - size_k=size_k,
405 + size_n=padded_n,
406 + size_k=padded_k,
407 use_atomic_add=use_atomic_add,
408 use_fp32_reduce=use_fp32_reduce,
409 )
410
411 + output = marlin_unpad_output(output, size_n, padded_n)
412 return output.reshape(out_shape)
413
414
@@ -401,6 +423,7 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
423 part_size_n = layer.output_size_per_partition
424 part_size_k = layer.input_size_per_partition
425 group_size = 32 # MX standard block size
426 + padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size)
427
428 device = layer.weight.device
429
@@ -411,12 +434,13 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
434 perm = torch.empty(0, dtype=torch.int, device=device)
435 qweight = pack_fp8_to_int32(layer.weight, size_k_first=False)
436 qweight = qweight.T.contiguous()
437 + qweight = marlin_pad_qweight(qweight, part_size_n, part_size_k, padded_n, padded_k)
438
439 marlin_qweight = ops.gptq_marlin_repack(
440 b_q_weight=qweight,
441 perm=perm,
418 - size_k=part_size_k,
419 - size_n=part_size_n,
442 + size_k=padded_k,
443 + size_n=padded_n,
444 num_bits=8,
445 )
446 replace_parameter(layer, "weight", marlin_qweight)
@@ -429,12 +453,15 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
453 scales = scales.contiguous()
454 scales = scales.view(torch.float8_e8m0fnu).to(param_dtype)
455 scales = scales.T.contiguous()
456 + scales = marlin_pad_scales(
457 + scales, part_size_n, part_size_k, padded_n, padded_k, group_size
458 + )
459
460 # Permute scales to Marlin layout
461 marlin_scales = marlin_permute_scales(
462 s=scales,
436 - size_k=part_size_k,
437 - size_n=part_size_n,
463 + size_k=padded_k,
464 + size_n=padded_n,
465 group_size=group_size,
466 )
467
@@ -445,7 +472,7 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:
472 # BIAS
473 if hasattr(layer, "bias") and layer.bias is not None:
474 assert layer.bias.shape == (part_size_n,)
448 - bias = marlin_permute_bias(layer.bias)
475 + bias = marlin_permute_bias(marlin_pad_dim(layer.bias, part_size_n, padded_n))
476 replace_parameter(layer, "bias", bias)
477
478