@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 37a682d392

[Kernel] Extend Marlin thread-tile padding to MoE (WNA16 + FP8/MXFP8) (#45703)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

Michael Goin committed Jun 23, 2026 at 12:45 UTC 37a682d392330da665690ee0a77c9d0a875f315f
7 files changed +573 -29
tests/kernels/quantization/test_marlin_tile_padding.py
+389
@@ -5,6 +5,8 @@
5 Run `pytest tests/kernels/quantization/test_marlin_tile_padding.py`.
6 """
7
8 +from types import SimpleNamespace
9 +
10 import pytest
11 import torch
12
@@ -14,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
16 apply_gptq_marlin_linear,
17 marlin_make_empty_g_idx,
18 marlin_make_workspace_new,
19 + marlin_moe_padded_intermediate,
20 marlin_pad_qweight,
21 marlin_pad_scales,
22 marlin_padded_nk,
@@ -115,6 +118,74 @@ def test_marlin_pad_helpers_shapes():
118 assert padded.shape == (1, padded_n)
119
120
121 +# Rank-local MoE intermediate sizes. group<=0 / 32 with a non-multiple-of-64
122 +# size is where tile padding triggers; 64/128 are already tile-aligned.
123 +MOE_INTERMEDIATE_SIZES = [64, 96, 100, 176, 192, 256, 2816]
124 +
125 +
126 +@pytest.mark.parametrize("intermediate", MOE_INTERMEDIATE_SIZES)
127 +@pytest.mark.parametrize("group_size", [-1, 32, 64, 128])
128 +def test_marlin_moe_padded_intermediate(intermediate, group_size):
129 + # The MoE gate only admits shapes where the group does not straddle the
130 + # boundary, i.e. group divides the intermediate size.
131 + if group_size > 0 and intermediate % group_size != 0:
132 + pytest.skip("group straddles the boundary; rejected by the MoE gate")
133 +
134 + padded = marlin_moe_padded_intermediate(intermediate, group_size)
135 + assert padded >= intermediate
136 + # Valid MoE thread tile: gate-up n = 2*intermediate % 128, down k % 64.
137 + assert (2 * padded) % 128 == 0
138 + assert padded % 64 == 0
139 + if group_size > 0:
140 + assert padded % group_size == 0
141 +
142 + # Minimal: no smaller valid intermediate exists.
143 + for cand in range(intermediate, padded):
144 + if (
145 + (2 * cand) % 128 == 0
146 + and cand % 64 == 0
147 + and (group_size <= 0 or cand % group_size == 0)
148 + ):
149 + pytest.fail(f"{cand} beats {padded}")
150 +
151 + # Already-tile-aligned sizes pass through unchanged (zero hot-path cost).
152 + if intermediate % 64 == 0:
153 + assert padded == intermediate
154 +
155 +
156 +def test_marlin_moe_pad_helpers_shapes():
157 + from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
158 + _pad_rows,
159 + _pad_w13_bias,
160 + _pad_w13_shard_cols,
161 + )
162 +
163 + E, rows, N, padded_N = 2, 8, 96, 128
164 +
165 + # w13 stores the two gate/up shards along the last dim; padding each shard
166 + # must preserve the loaded values and zero the padded columns.
167 + w13 = torch.arange(E * rows * 2 * N).reshape(E, rows, 2 * N).float()
168 + padded = _pad_w13_shard_cols(w13, N, padded_N)
169 + assert padded.shape == (E, rows, 2 * padded_N)
170 + shards = padded.view(E, rows, 2, padded_N)
171 + orig = w13.view(E, rows, 2, N)
172 + assert torch.equal(shards[..., :N], orig)
173 + assert shards[..., N:].abs().sum() == 0
174 +
175 + # w2 stores the intermediate dim in the rows.
176 + w2 = torch.ones(E, N // 32, 16)
177 + padded = _pad_rows(w2, padded_N // 32)
178 + assert padded.shape == (E, padded_N // 32, 16)
179 + assert padded[:, N // 32 :, :].abs().sum() == 0
180 +
181 + bias = torch.arange(E * 2 * N).reshape(E, 2 * N).float()
182 + padded = _pad_w13_bias(bias, N, padded_N)
183 + assert padded.shape == (E, 2 * padded_N)
184 + bias_shards = padded.view(E, 2, padded_N)
185 + assert torch.equal(bias_shards[..., :N], bias.view(E, 2, N))
186 + assert bias_shards[..., N:].abs().sum() == 0
187 +
188 +
189 def _gpu_marlin_unsupported() -> bool:
190 return not (
191 current_platform.is_cuda() and current_platform.has_device_capability(80)
@@ -468,3 +539,321 @@ def test_check_marlin_supports_layer_allow_tile_padding():
539 # A group straddling the TP shard cannot be fixed by padding
540 layer = _FakeLinear(4608, 4672, input_size=18688)
541 assert not check_marlin_supports_layer(layer, 128, allow_tile_padding=True)
542 +
543 +
544 +@pytest.mark.skipif(
545 + _gpu_marlin_unsupported(),
546 + reason="Marlin is not supported on this GPU type.",
547 +)
548 +@pytest.mark.parametrize("group_size", [-1, 32])
549 +@pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
550 +def test_gptq_marlin_moe_padded_round_trip(shape, group_size):
551 + """Pad a tile-misaligned MoE intermediate the way the WNA16 Marlin MoE prep
552 + does, run the real repack + fused_marlin_moe, and check against the
553 + dequantized reference. Symmetric int4's quantized zero decodes to -8, so the
554 + padded region only stays out of the output via the zero-padded scales.
555 + """
556 + from tests.kernels.utils import torch_experts
557 + from vllm.config import VllmConfig, set_current_vllm_config
558 + from vllm.model_executor.layers.fused_moe import fused_topk
559 + from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
560 + fused_marlin_moe,
561 + )
562 + from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
563 + _pad_rows,
564 + _pad_w13_shard_cols,
565 + )
566 + from vllm.model_executor.layers.quantization.utils.marlin_utils import (
567 + marlin_moe_padded_intermediate,
568 + marlin_moe_permute_scales,
569 + )
570 +
571 + n, k, e = shape
572 + topk, m = 2, 33
573 + padded_n = marlin_moe_padded_intermediate(n, group_size)
574 + assert padded_n != n, "test should exercise padding"
575 +
576 + dtype = torch.float16
577 + device = torch.device("cuda")
578 + quant_type = scalar_types.uint4b8
579 + bits = quant_type.size_bits
580 + pack = 32 // bits
581 +
582 + a = torch.randn((m, k), device=device, dtype=dtype) / 10
583 + w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
584 + w2 = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
585 +
586 + def quant(w, size_k, size_n):
587 + # w is (size_n, size_k); gptq expects (size_k, size_n).
588 + ref, q_w, s, _, _ = gptq_quantize_weights(
589 + w.T, quant_type, group_size, act_order=False
590 + )
591 + return ref, gptq_pack(q_w, bits, size_k, size_n), s
592 +
593 + w13_qw, w13_s, w13_ref = [], [], []
594 + w2_qw, w2_s, w2_ref = [], [], []
595 + for i in range(e):
596 + ref, qw, s = quant(w1[i], k, 2 * n)
597 + w13_ref.append(ref.T) # (2n, k)
598 + w13_qw.append(qw)
599 + w13_s.append(s)
600 + ref, qw, s = quant(w2[i], n, k)
601 + w2_ref.append(ref.T) # (k, n)
602 + w2_qw.append(qw)
603 + w2_s.append(s)
604 +
605 + w13_qweight = torch.stack(w13_qw)
606 + w2_qweight = torch.stack(w2_qw)
607 + w13_scales = torch.stack(w13_s)
608 + w2_scales = torch.stack(w2_s)
609 + w1_ref = torch.stack(w13_ref) # (e, 2n, k)
610 + w2_ref = torch.stack(w2_ref) # (e, k, n)
611 +
612 + # Pad the intermediate via the production helpers.
613 + w13_qweight = _pad_w13_shard_cols(w13_qweight, n, padded_n)
614 + w2_qweight = _pad_rows(w2_qweight, padded_n // pack)
615 + w13_scales = _pad_w13_shard_cols(w13_scales, n, padded_n)
616 + if group_size > 0:
617 + w2_scales = _pad_rows(w2_scales, padded_n // group_size)
618 +
619 + sort_idx = torch.empty((e, 0), dtype=torch.int32, device=device)
620 + marlin_w13 = ops.gptq_marlin_moe_repack(
621 + w13_qweight, sort_idx, w13_qweight.shape[1] * pack, w13_qweight.shape[2], bits
622 + )
623 + marlin_w2 = ops.gptq_marlin_moe_repack(
624 + w2_qweight, sort_idx, w2_qweight.shape[1] * pack, w2_qweight.shape[2], bits
625 + )
626 + group_or_pack = group_size if group_size != -1 else pack
627 + marlin_w13_s = marlin_moe_permute_scales(
628 + s=w13_scales, size_k=n, size_n=w13_scales.shape[2], group_size=group_size
629 + )
630 + marlin_w2_s = marlin_moe_permute_scales(
631 + s=w2_scales,
632 + size_k=w2_scales.shape[1] * group_or_pack,
633 + size_n=w2_scales.shape[2],
634 + group_size=group_size,
635 + )
636 +
637 + score = torch.randn((m, e), device=device, dtype=dtype)
638 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
639 +
640 + marlin_out = fused_marlin_moe(
641 + a,
642 + marlin_w13,
643 + marlin_w2,
644 + None,
645 + None,
646 + marlin_w13_s,
647 + marlin_w2_s,
648 + topk_weights,
649 + topk_ids,
650 + quant_type_id=quant_type.id,
651 + global_num_experts=e,
652 + is_k_full=True,
653 + )
654 + with set_current_vllm_config(VllmConfig()):
655 + ref = torch_experts(
656 + a,
657 + w1_ref,
658 + w2_ref,
659 + topk_weight=topk_weights,
660 + topk_ids=topk_ids,
661 + global_num_experts=e,
662 + )
663 +
664 + torch.testing.assert_close(marlin_out, ref, atol=5e-2, rtol=0)
665 +
666 +
667 +def test_check_moe_marlin_supports_layer_padding():
668 + from vllm.model_executor.layers.quantization.utils.marlin_utils import (
669 + check_moe_marlin_supports_layer,
670 + )
671 +
672 + def make_layer(hidden, intermediate):
673 + layer = SimpleNamespace()
674 + layer.hidden_size = hidden
675 + layer.apply_router_weight_on_input = False
676 + layer.moe_config = SimpleNamespace(
677 + intermediate_size_per_partition_unpadded=intermediate
678 + )
679 + return layer
680 +
681 + # group=32 with intermediate % 64 != 0: rejected strictly, accepted w/ padding
682 + layer = make_layer(4096, 96)
683 + assert not check_moe_marlin_supports_layer(layer, 32)
684 + assert check_moe_marlin_supports_layer(layer, 32, allow_tile_padding=True)
685 + # channelwise misaligned intermediate is paddable
686 + assert check_moe_marlin_supports_layer(layer, -1, allow_tile_padding=True)
687 +
688 + # A group straddling the boundary cannot be fixed by padding
689 + layer = make_layer(4096, 176)
690 + assert not check_moe_marlin_supports_layer(layer, 128, allow_tile_padding=True)
691 +
692 + # hidden_size is the MoE I/O extent and is never padded
693 + layer = make_layer(4090, 128)
694 + assert not check_moe_marlin_supports_layer(layer, 64, allow_tile_padding=True)
695 +
696 +
697 +@pytest.mark.skipif(
698 + _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
699 + reason="FP8 Marlin is not supported on this GPU type.",
700 +)
701 +@pytest.mark.parametrize("quant", ["channel", "tensor"])
702 +@pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
703 +def test_fp8_marlin_moe_padded_round_trip(shape, quant):
704 + """FP8 weight-only MoE: pad a tile-misaligned intermediate and check the
705 + real prepare + fused_marlin_moe against the dequantized reference."""
706 + from tests.kernels.utils import torch_experts
707 + from vllm.config import VllmConfig, set_current_vllm_config
708 + from vllm.model_executor.layers.fused_moe import fused_topk
709 + from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
710 + fused_marlin_moe,
711 + )
712 + from vllm.model_executor.layers.quantization.utils.marlin_utils import (
713 + marlin_moe_intermediate_size,
714 + marlin_moe_padded_intermediate,
715 + )
716 + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
717 + prepare_fp8_moe_layer_for_marlin,
718 + )
719 +
720 + n, k, e = shape
721 + topk, m = 2, 33
722 + fp8 = torch.float8_e4m3fn
723 + dtype = torch.bfloat16
724 + device = torch.device("cuda")
725 + padded_n = marlin_moe_padded_intermediate(n, -1)
726 + assert padded_n != n
727 +
728 + def q(w): # (out, in) -> fp8 weight, scale, dequant reference
729 + dim = None if quant == "tensor" else 1
730 + s = (w.abs().amax(dim, keepdim=dim is not None) / 448.0).clamp(min=1e-8)
731 + wq = (w / s).clamp(-448, 448).to(fp8)
732 + ref = wq.to(dtype) * s.to(dtype)
733 + s = s.reshape(1) if quant == "tensor" else s.squeeze(1)
734 + return wq, s, ref
735 +
736 + a = torch.randn((m, k), device=device, dtype=dtype) / 10
737 + w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
738 + w2 = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
739 + w13_q, w13_s, w1_ref = zip(*(q(w1[i]) for i in range(e)))
740 + w2_q, w2_s, w2_ref = zip(*(q(w2[i]) for i in range(e)))
741 +
742 + w13_weight, w2_weight = torch.stack(w13_q), torch.stack(w2_q)
743 + layer = SimpleNamespace(
744 + num_experts=e,
745 + hidden_size=k,
746 + intermediate_size_per_partition=n,
747 + orig_dtype=dtype,
748 + w13_weight=w13_weight,
749 + )
750 + pw13, pw2, ps13, ps2 = prepare_fp8_moe_layer_for_marlin(
751 + layer, w13_weight, w2_weight, torch.stack(w13_s), torch.stack(w2_s)
752 + )
753 + assert marlin_moe_intermediate_size(pw13, pw2) == padded_n
754 +
755 + score = torch.randn((m, e), device=device, dtype=dtype)
756 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
757 + out = fused_marlin_moe(
758 + a,
759 + pw13,
760 + pw2,
761 + None,
762 + None,
763 + ps13,
764 + ps2,
765 + topk_weights,
766 + topk_ids,
767 + quant_type_id=scalar_types.float8_e4m3fn.id,
768 + global_num_experts=e,
769 + is_k_full=True,
770 + workspace=layer.workspace,
771 + )
772 + with set_current_vllm_config(VllmConfig()):
773 + ref = torch_experts(
774 + a,
775 + torch.stack(w1_ref),
776 + torch.stack(w2_ref),
777 + topk_weight=topk_weights,
778 + topk_ids=topk_ids,
779 + global_num_experts=e,
780 + )
781 + torch.testing.assert_close(out, ref, atol=8e-2, rtol=0)
782 +
783 +
784 +@pytest.mark.skipif(
785 + _gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
786 + reason="FP8 Marlin is not supported on this GPU type.",
787 +)
788 +@pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
789 +def test_mxfp8_marlin_moe_padded_round_trip(shape):
790 + """MXFP8 weight-only MoE round-trip at a tile-misaligned intermediate, with
791 + unit e8m0 scales so the reference is the exact fp8 dequant."""
792 + from tests.kernels.utils import torch_experts
793 + from vllm.config import VllmConfig, set_current_vllm_config
794 + from vllm.model_executor.layers.fused_moe import fused_topk
795 + from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
796 + fused_marlin_moe,
797 + )
798 + from vllm.model_executor.layers.quantization.utils.marlin_utils import (
799 + marlin_moe_intermediate_size,
800 + marlin_moe_padded_intermediate,
801 + )
802 + from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
803 + prepare_mxfp8_moe_layer_for_marlin,
804 + )
805 +
806 + n, k, e = shape
807 + topk, m, gs, e8m0_one = 2, 33, 32, 127
808 + fp8 = torch.float8_e4m3fn
809 + dtype = torch.bfloat16
810 + device = torch.device("cuda")
811 + padded_n = marlin_moe_padded_intermediate(n, gs)
812 + assert padded_n != n
813 +
814 + a = torch.randn((m, k), device=device, dtype=dtype) / 10
815 + w13_weight = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
816 + w2_weight = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
817 + w13_weight = w13_weight.clamp(-448, 448).to(fp8)
818 + w2_weight = w2_weight.clamp(-448, 448).to(fp8)
819 + w13_scale = torch.full(
820 + (e, 2 * n, k // gs), e8m0_one, dtype=torch.uint8, device=device
821 + )
822 + w2_scale = torch.full((e, k, n // gs), e8m0_one, dtype=torch.uint8, device=device)
823 +
824 + layer = SimpleNamespace(
825 + num_experts=e, hidden_size=k, intermediate_size_per_partition=n
826 + )
827 + with set_current_vllm_config(VllmConfig()):
828 + pw13, pw2, ps13, ps2 = prepare_mxfp8_moe_layer_for_marlin(
829 + layer, w13_weight, w2_weight, w13_scale, w2_scale
830 + )
831 + assert marlin_moe_intermediate_size(pw13, pw2) == padded_n
832 +
833 + score = torch.randn((m, e), device=device, dtype=dtype)
834 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
835 + out = fused_marlin_moe(
836 + a,
837 + pw13,
838 + pw2,
839 + None,
840 + None,
841 + ps13,
842 + ps2,
843 + topk_weights,
844 + topk_ids,
845 + quant_type_id=scalar_types.float8_e4m3fn.id,
846 + global_num_experts=e,
847 + is_k_full=True,
848 + workspace=layer.workspace,
849 + )
850 + with set_current_vllm_config(VllmConfig()):
851 + ref = torch_experts(
852 + a,
853 + w13_weight.to(dtype),
854 + w2_weight.to(dtype),
855 + topk_weight=topk_weights,
856 + topk_ids=topk_ids,
857 + global_num_experts=e,
858 + )
859 + torch.testing.assert_close(out, ref, atol=8e-2, rtol=0)
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+71
@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
29 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
30 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
31 marlin_act_int8_process_scales,
32 + marlin_moe_padded_intermediate,
33 marlin_moe_permute_scales,
34 marlin_permute_bias,
35 moe_awq_to_marlin_zero_points,
@@ -361,6 +362,34 @@ def _process_weights_flashinfer(
362 )
363
364
365 +def _pad_w13_shard_cols(x: torch.Tensor, unit: int, padded_unit: int) -> torch.Tensor:
366 + """Zero-pad each of the two gate/up shards of a ``(E, rows, 2 * unit)``
367 + tensor along its last dim, from ``unit`` to ``padded_unit`` columns."""
368 + if padded_unit == unit:
369 + return x
370 + e, rows, _ = x.shape
371 + x = x.view(e, rows, 2, unit)
372 + x = torch.nn.functional.pad(x, (0, padded_unit - unit))
373 + return x.reshape(e, rows, 2 * padded_unit).contiguous()
374 +
375 +
376 +def _pad_rows(x: torch.Tensor, padded_rows: int) -> torch.Tensor:
377 + """Zero-pad a ``(E, rows, cols)`` tensor to ``padded_rows`` rows."""
378 + if padded_rows == x.size(1):
379 + return x
380 + return torch.nn.functional.pad(x, (0, 0, 0, padded_rows - x.size(1)))
381 +
382 +
383 +def _pad_w13_bias(bias: torch.Tensor, n: int, padded_n: int) -> torch.Tensor:
384 + """Zero-pad each gate/up shard of a ``(E, 2 * n)`` bias to ``padded_n``."""
385 + if padded_n == n:
386 + return bias
387 + e = bias.size(0)
388 + bias = bias.view(e, 2, n)
389 + bias = torch.nn.functional.pad(bias, (0, padded_n - n))
390 + return bias.reshape(e, 2 * padded_n).contiguous()
391 +
392 +
393 def _process_weights_marlin(
394 layer: torch.nn.Module,
395 input_dtype: torch.dtype | None,
@@ -431,6 +460,29 @@ def _process_weights_marlin(
460 marlin_w13_scales = w13_scales
461 marlin_w2_scales = w2_scales
462
463 + # --- Pad the intermediate size to a valid Marlin thread tile ---
464 + # GPTQ packs along K: w13's N is in the (shard) columns, w2's N in the rows.
465 + # Act-order keeps the strict shape and is never padded.
466 + N = layer.intermediate_size_per_partition
467 + padded_N = marlin_moe_padded_intermediate(N, group_size)
468 + if padded_N != N:
469 + assert actorder != "group", (
470 + "Marlin MoE thread-tile padding is unsupported with act-order"
471 + )
472 + marlin_w13_qweight = _pad_w13_shard_cols(marlin_w13_qweight, N, padded_N)
473 + marlin_w2_qweight = _pad_rows(marlin_w2_qweight, padded_N // pack_factor)
474 + marlin_w13_scales = _pad_w13_shard_cols(marlin_w13_scales, N, padded_N)
475 + if group_size > 0:
476 + marlin_w2_scales = _pad_rows(marlin_w2_scales, padded_N // group_size)
477 + if w13_qzeros is not None:
478 + w13_qzeros = _pad_w13_shard_cols(
479 + w13_qzeros, N // pack_factor, padded_N // pack_factor
480 + )
481 + if w2_qzeros is not None and group_size > 0:
482 + w2_qzeros = _pad_rows(w2_qzeros, padded_N // group_size)
483 + if w13_bias is not None:
484 + w13_bias = _pad_w13_bias(w13_bias, N, padded_N)
485 +
486 # --- Process act_order (g_idx) ---
487 if actorder == "group":
488 num_experts = w13_g_idx.shape[0]
@@ -608,6 +660,25 @@ def _process_awq_weights_marlin(
660 w13_scales = w13_scales.data * 512
661 w2_scales = w2_scales.data * 512
662
663 + # --- Pad the intermediate size to a valid Marlin thread tile ---
664 + # AWQ packs along N: w13's N is in the (shard) columns, w2's N in the rows.
665 + N = layer.intermediate_size_per_partition
666 + padded_N = marlin_moe_padded_intermediate(N, group_size)
667 + if padded_N != N:
668 + w13_qweight = _pad_w13_shard_cols(
669 + w13_qweight, N // pack_factor, padded_N // pack_factor
670 + )
671 + w2_qweight = _pad_rows(w2_qweight, padded_N)
672 + w13_scales = _pad_w13_shard_cols(w13_scales, N, padded_N)
673 + w13_qzeros = _pad_w13_shard_cols(
674 + w13_qzeros, N // pack_factor, padded_N // pack_factor
675 + )
676 + if group_size > 0:
677 + w2_scales = _pad_rows(w2_scales, padded_N // group_size)
678 + w2_qzeros = _pad_rows(w2_qzeros, padded_N // group_size)
679 + if w13_bias is not None:
680 + w13_bias = _pad_w13_bias(w13_bias, N, padded_N)
681 +
682 w13_g_idx_sort_indices = torch.nn.Parameter(
683 torch.empty((num_experts, 0), dtype=torch.int32, device=device),
684 requires_grad=False,
vllm/model_executor/layers/quantization/auto_awq.py
+3 -1
@@ -339,7 +339,9 @@ class AutoAWQConfig(QuantizationConfig):
339 ):
340 return UnquantizedFusedMoEMethod(layer.moe_config)
341
342 - if not check_moe_marlin_supports_layer(layer, self.group_size):
342 + if not check_moe_marlin_supports_layer(
343 + layer, self.group_size, allow_tile_padding=True
344 + ):
345 logger.warning_once(
346 f"Layer '{prefix}' is not supported by AutoAWQMoEMarlin. "
347 "Falling back to Moe WNA16 kernels."
vllm/model_executor/layers/quantization/auto_gptq.py
+3 -1
@@ -243,7 +243,9 @@ class AutoGPTQConfig(QuantizationConfig):
243 if isinstance(layer, RoutedExperts):
244 from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
245
246 - if not check_moe_marlin_supports_layer(layer, self.group_size):
246 + if not check_moe_marlin_supports_layer(
247 + layer, self.group_size, allow_tile_padding=not self.desc_act
248 + ):
249 logger.warning_once(
250 f"Layer '{prefix}' is not supported by GPTQMoeMarlin. "
251 "Falling back to Moe WNA16 kernels."
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py
+9 -6
@@ -96,15 +96,18 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
96 )
97
98 # Prefer to use the MarlinMoE kernel when it is supported.
99 + is_actorder = (
100 + weight_quant.strategy == QuantizationStrategy.GROUP
101 + and weight_quant.actorder
102 + in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
103 + )
104 if (
100 - not check_moe_marlin_supports_layer(layer, group_size)
105 + not check_moe_marlin_supports_layer(
106 + layer, group_size, allow_tile_padding=not is_actorder
107 + )
108 or current_platform.is_rocm()
109 ):
103 - if (
104 - weight_quant.strategy == QuantizationStrategy.GROUP
105 - and weight_quant.actorder
106 - in (ActivationOrdering.GROUP, ActivationOrdering.DYNAMIC)
107 - ):
110 + if is_actorder:
111 raise ValueError(
112 "WNA16MoE is not supported with actorder=group/dynamic."
113 )
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+43 -10
@@ -330,12 +330,43 @@ def check_marlin_supports_layer(
330 )[0]
331
332
333 -def check_moe_marlin_supports_layer(layer: RoutedExperts, group_size: int) -> bool:
333 +def marlin_moe_padded_intermediate(intermediate_size: int, group_size: int = -1) -> int:
334 + """Smallest MoE intermediate size satisfying the Marlin MoE thread tiles.
335 +
336 + The kernel needs gate-up ``2 * intermediate % 128 == 0`` and down
337 + ``intermediate % 64 == 0``, i.e. ``intermediate % 64 == 0``. A misaligned
338 + size is zero-padded to the next valid tile at weight prep, kept a multiple
339 + of ``group_size`` so the group count stays integral. The padded region never
340 + reaches the MoE output: w13's padded output channels are zeroed by the
341 + zero-padded scales, so the padded inputs to w2 are zero.
342 + """
343 + group = group_size if group_size > 0 else 1
344 + padded = round_up(intermediate_size, math.lcm(64, group))
345 + if padded != intermediate_size:
346 + logger.warning_once(
347 + "Marlin requires thread-tile padding for the MoE intermediate size "
348 + "of some layers in this model. Padded experts pad/slice activations "
349 + "on every forward; performance may be degraded."
350 + )
351 + return padded
352 +
353 +
354 +def check_moe_marlin_supports_layer(
355 + layer: RoutedExperts, group_size: int, allow_tile_padding: bool = False
356 +) -> bool:
357 + """Whether the fused MoE Marlin kernel supports ``layer``.
358 +
359 + Callers without act-order may pass ``allow_tile_padding=True``: a
360 + tile-misaligned intermediate size is then zero-padded to a valid thread
361 + tile at weight prep (see marlin_moe_padded_intermediate), so only a group
362 + straddling the padded boundary stays unsupported. hidden_size is the MoE
363 + I/O extent and is never padded. Act-order keeps the strict shape.
364 + """
365 if current_platform.is_rocm():
366 return False
367 hidden_size = layer.hidden_size
337 - # Note: The layer has not performed rounding on intermediate_size's at this
338 - # point. Use the unpadded size which won't change.
368 + # The layer has not rounded intermediate_size yet; use the stable unpadded
369 + # size. gate-up needs n=2*intermediate % 128, down needs k=intermediate % 64.
370 intermediate_size_per_partition = (
371 layer.moe_config.intermediate_size_per_partition_unpadded
372 )
@@ -343,13 +374,15 @@ def check_moe_marlin_supports_layer(layer: RoutedExperts, group_size: int) -> bo
374 # apply_router_weight_on_input is not supported for moe marlin
375 supports_router_weight = not layer.apply_router_weight_on_input
376
346 - # gate-up: (n, k) = (intermediate_size_per_partition * 2, hidden_size)
347 - # down: (n, k) = (hidden_size, intermediate_size_per_partition)
348 - # moe marlin requires n % 128 == 0 and k % 64 == 0
349 - supports_shape = (
350 - hidden_size % 128 == 0
351 - and intermediate_size_per_partition % max(64, group_size) == 0
352 - )
377 + if allow_tile_padding:
378 + supports_shape = hidden_size % 128 == 0 and (
379 + group_size <= 0 or intermediate_size_per_partition % group_size == 0
380 + )
381 + else:
382 + supports_shape = (
383 + hidden_size % 128 == 0
384 + and intermediate_size_per_partition % max(64, group_size) == 0
385 + )
386 supports_group_size = group_size in [-1, 32, 64, 128]
387 return supports_shape and supports_group_size and supports_router_weight
388
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+55 -11
@@ -10,6 +10,7 @@ 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_moe_padded_intermediate,
14 marlin_pad_dim,
15 marlin_pad_qweight,
16 marlin_pad_scales,
@@ -216,6 +217,25 @@ def prepare_fp8_layer_for_marlin(
217 replace_parameter(layer, "bias", bias)
218
219
220 +def _moe_pad_shard_rows(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor:
221 + """Zero-pad each gate/up shard of a ``(E, 2 * n, ...)`` tensor to padded_n
222 + rows. FP8 zero decodes to 0.0, so the padded rows contribute nothing."""
223 + if padded_n == n:
224 + return x
225 + e = x.size(0)
226 + rest = x.shape[2:]
227 + x = x.view(e, 2, n, *rest)
228 + x = torch.nn.functional.pad(x, (0, 0) * len(rest) + (0, padded_n - n))
229 + return x.reshape(e, 2 * padded_n, *rest)
230 +
231 +
232 +def _moe_pad_last(x: torch.Tensor, n: int, padded_n: int) -> torch.Tensor:
233 + """Zero-pad the last dim of a ``(E, ..., n)`` tensor to padded_n."""
234 + if padded_n == n:
235 + return x
236 + return torch.nn.functional.pad(x, (0, padded_n - n))
237 +
238 +
239 def prepare_fp8_moe_layer_for_marlin(
240 layer: torch.nn.Module,
241 w13_weight: torch.Tensor,
@@ -246,6 +266,15 @@ def prepare_fp8_moe_layer_for_marlin(
266 n = layer.intermediate_size_per_partition
267 w13_n = w13_weight.size(1)
268 weight_block_size = getattr(layer, "weight_block_size", None)
269 + group_size = -1 if weight_block_size is None else weight_block_size[1]
270 +
271 + # Pad a tile-misaligned intermediate size to a valid Marlin thread tile.
272 + # FP8 zero decodes to 0.0, so padded weights drop out; the converted scales
273 + # are padded to match below (the padded values are irrelevant).
274 + padded_n = marlin_moe_padded_intermediate(n, group_size)
275 + if padded_n != n:
276 + w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n)
277 + w2_weight = _moe_pad_last(w2_weight, n, padded_n)
278
279 # WORKSPACE
280 device = layer.w13_weight.device
@@ -258,13 +287,7 @@ def prepare_fp8_moe_layer_for_marlin(
287 # Repack weights to marlin format
288 def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor:
289 tensor_list = []
261 - if "w13" in name:
262 - size_n, size_k = w13_n, k
263 - else:
264 - size_n, size_k = k, n
265 -
266 - assert weight.shape == (e, size_n, size_k)
267 -
290 + size_n, size_k = weight.size(1), weight.size(2)
291 for i in range(e):
292 qweight = pack_fp8_to_int32(weight[i], size_k_first=False)
293 qweight = qweight.T.contiguous()
@@ -280,9 +303,7 @@ def prepare_fp8_moe_layer_for_marlin(
303 w2_weight = repack_weight("w2", w2_weight)
304
305 # WEIGHT SCALES
283 - # Permute scales
284 - group_size = -1 if weight_block_size is None else weight_block_size[1]
285 -
306 + # Permute scales (convert at the original size, then pad to the tile).
307 def permute_scales(scales: torch.Tensor, name: str) -> torch.Tensor:
308 scales = scales.to(layer.orig_dtype)
309 tensor_list = []
@@ -320,6 +341,20 @@ def prepare_fp8_moe_layer_for_marlin(
341 # size_n may not divisible by block_size[0]
342 scales = scales[..., :size_n].contiguous()
343
344 + # Pad the converted (E, G, size_n) scales to the padded thread tile.
345 + if padded_n != n:
346 + if "w13" in name:
347 + g = scales.size(1)
348 + scales = scales.view(e, g, 2, n)
349 + scales = torch.nn.functional.pad(scales, (0, padded_n - n))
350 + scales = scales.reshape(e, g, 2 * padded_n)
351 + size_n = 2 * padded_n
352 + else:
353 + if group_size > 0:
354 + pad_groups = (padded_n - n) // group_size
355 + scales = torch.nn.functional.pad(scales, (0, 0, 0, pad_groups))
356 + size_k = padded_n
357 +
358 for i in range(e):
359 marlin_scales = marlin_permute_scales(
360 s=scales[i], size_k=size_k, size_n=size_n, group_size=group_size
@@ -497,10 +532,19 @@ def prepare_mxfp8_moe_layer_for_marlin(
532 """
533 group_size = 32
534 e = w13.shape[0]
500 - w13_n = w13.shape[1]
535 k = w13.shape[2]
536 n = w2.shape[2]
537
538 + # Pad a tile-misaligned intermediate size to a valid Marlin thread tile.
539 + padded_n = marlin_moe_padded_intermediate(n, group_size)
540 + if padded_n != n:
541 + w13 = _moe_pad_shard_rows(w13, n, padded_n)
542 + w13_scale = _moe_pad_shard_rows(w13_scale, n, padded_n)
543 + w2 = _moe_pad_last(w2, n, padded_n)
544 + w2_scale = _moe_pad_last(w2_scale, n // group_size, padded_n // group_size)
545 + n = padded_n
546 + w13_n = w13.shape[1]
547 +
548 device = w13.device
549 param_dtype = torch.get_default_dtype()
550 perm = torch.empty(0, dtype=torch.int, device=device)