[Quantization] Bound peak memory when repacking FP4 MoE weights for Marlin (#47851)
Signed-off-by: Joe Rowell <joerowell4@gmail.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Joe Rowell committed
Jul 11, 2026 at 03:46 UTC
1bf3997eaeb41b1ef5cf25d87faa97b7f0be4788
1 file changed
+34
-43
vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py
+34
-43
@@ -306,6 +306,37 @@ def prepare_fp4_layer_for_marlin(
306
return
307
308
309
+def _repack_marlin_experts(
310
+ weight: torch.Tensor,
311
+ size_n: int,
312
+ size_k: int,
313
+ perm: torch.Tensor,
314
+ is_a_8bit: bool,
315
+) -> torch.Tensor:
316
+ """Repack each expert to marlin format into a preallocated output."""
317
+ num_experts = weight.shape[0]
318
+ out: torch.Tensor | None = None
319
+ for i in range(num_experts):
320
+ qweight = weight[i].view(torch.int32).T.contiguous()
321
+ marlin_qweight = ops.gptq_marlin_repack(
322
+ b_q_weight=qweight,
323
+ perm=perm,
324
+ size_k=size_k,
325
+ size_n=size_n,
326
+ num_bits=4,
327
+ is_a_8bit=is_a_8bit,
328
+ )
329
+ if out is None:
330
+ out = torch.empty(
331
+ (num_experts, *marlin_qweight.shape),
332
+ dtype=marlin_qweight.dtype,
333
+ device=marlin_qweight.device,
334
+ )
335
+ out[i] = marlin_qweight
336
+ assert out is not None
337
+ return out
338
+
339
+
340
def prepare_nvfp4_moe_layer_for_marlin(
341
layer: RoutedExperts,
342
w13: torch.Tensor,
@@ -371,7 +402,6 @@ def prepare_nvfp4_moe_layer_for_marlin(
402
# WEIGHT
403
# Repack weights to marlin format
404
def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
374
- tensor_list = []
405
if "w13" in name:
406
size_n, size_k = N * num_shards, K
407
assert weight.shape == (E, size_n, size_k // 2)
@@ -383,20 +413,7 @@ def prepare_nvfp4_moe_layer_for_marlin(
413
weight = pad_w2(weight, packing=2)
414
size_k = padded_N
415
386
- for i in range(E):
387
- qweight = weight[i].view(torch.int32).T.contiguous()
388
-
389
- marlin_qweight = ops.gptq_marlin_repack(
390
- b_q_weight=qweight,
391
- perm=perm,
392
- size_k=size_k,
393
- size_n=size_n,
394
- num_bits=4,
395
- is_a_8bit=is_a_8bit,
396
- )
397
- tensor_list.append(marlin_qweight)
398
-
399
- return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
416
+ return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit)
417
418
w13 = repack_weight(w13, "w13")
419
w2 = repack_weight(w2, "w2")
@@ -472,7 +489,6 @@ def prepare_moe_fp4_layer_for_marlin(
489
# Repack weights to marlin format
490
for name in ["w13_weight", "w2_weight"]:
491
weight = getattr(layer, name)
475
- tensor_list = []
492
if "w13" in name:
493
size_n, size_k = n * 2, k
494
else:
@@ -480,20 +496,7 @@ def prepare_moe_fp4_layer_for_marlin(
496
497
assert weight.shape == (e, size_n, size_k // 2)
498
483
- for i in range(e):
484
- qweight = weight[i].view(torch.int32).T.contiguous()
485
-
486
- marlin_qweight = ops.gptq_marlin_repack(
487
- b_q_weight=qweight,
488
- perm=perm,
489
- size_k=size_k,
490
- size_n=size_n,
491
- num_bits=4,
492
- is_a_8bit=is_a_8bit,
493
- )
494
- tensor_list.append(marlin_qweight)
495
-
496
- weight = torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
499
+ weight = _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit)
500
weight = torch.nn.Parameter(weight, requires_grad=False)
501
502
setattr(layer, name, weight)
@@ -615,7 +618,6 @@ def prepare_moe_mxfp4_layer_for_marlin(
618
619
# WEIGHT: Repack weights to marlin format
620
def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:
618
- tensor_list = []
621
if "w13" in name:
622
size_n, size_k = n * 2, k
623
else:
@@ -623,18 +625,7 @@ def prepare_moe_mxfp4_layer_for_marlin(
625
626
assert weight.shape == (e, size_n, size_k // 2)
627
626
- for i in range(e):
627
- qweight = weight[i].view(torch.int32).T.contiguous()
628
- marlin_qweight = ops.gptq_marlin_repack(
629
- b_q_weight=qweight,
630
- perm=perm,
631
- size_k=size_k,
632
- size_n=size_n,
633
- num_bits=4,
634
- is_a_8bit=is_a_8bit,
635
- )
636
- tensor_list.append(marlin_qweight)
637
- return torch.cat([x.unsqueeze(0) for x in tensor_list], 0)
628
+ return _repack_marlin_experts(weight, size_n, size_k, perm, is_a_8bit)
629
630
w13 = repack_weight(w13, "w13")
631
w2 = repack_weight(w2, "w2")