[Perf][MoE] Index expert mapping lookups in RoutedExperts.load_weights (#58720)
Signed-off-by: Willian <willian@willian.email> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Willian Z committed
Sep 26, 2026 at 09:48 UTC
ad6817b68d9b1d138eca2033cf90233bd0100bf3
2 files changed
+224
-2
tests/kernels/moe/test_moe_weight_loading_padded.py
+192
@@ -544,6 +544,198 @@ class TestLoadWeightsExpertBias:
544
list(RoutedExperts.load_weights(experts, weights))
545
546
547
+class TestLoadWeightsExpertMapping:
548
+ @staticmethod
549
+ def _mapping(projs=("gate_proj", "down_proj", "up_proj"), **kwargs):
550
+ return RoutedExperts.build_expert_params_mapping(
551
+ *projs,
552
+ num_experts=16,
553
+ routed_experts_prefix="",
554
+ include_fused=True,
555
+ **kwargs,
556
+ )
557
+
558
+ @staticmethod
559
+ def _load(
560
+ mapping,
561
+ weights,
562
+ suffix="weight",
563
+ *,
564
+ transposed=False,
565
+ quant_method="tensor",
566
+ layer_name="model.layers.0.mlp.experts",
567
+ ):
568
+ experts = torch.nn.Module()
569
+ experts.layer_name = layer_name
570
+ experts.get_expert_mapping = lambda **_: mapping
571
+ experts.is_fused_checkpoint_transposed = transposed
572
+ experts._orient_fused_weight = RoutedExperts._orient_fused_weight
573
+ calls = []
574
+
575
+ def weight_loader(param, loaded_weight, shard_id, expert_id, **_):
576
+ calls.append((param.test_name, shard_id, expert_id, loaded_weight.clone()))
577
+ return True
578
+
579
+ for prefix in ("w13", "w2"):
580
+ param_name = f"{prefix}_{suffix}"
581
+ param = torch.nn.Parameter(torch.empty(0), requires_grad=False)
582
+ param.test_name = param_name
583
+ param.quant_method = quant_method
584
+ param.weight_loader = weight_loader
585
+ setattr(experts, param_name, param)
586
+ loaded = list(RoutedExperts.load_weights(experts, iter(weights)))
587
+ assert loaded == [call[0] for call in calls]
588
+ return calls
589
+
590
+ @pytest.mark.parametrize(
591
+ "projs,lora_prefix",
592
+ [
593
+ (("gate_proj", "down_proj", "up_proj"), ""),
594
+ (("w1", "w2", "w3"), ""),
595
+ (("up_proj", "down_proj", "up_proj"), ""),
596
+ (("up_proj", "down_proj", ""), ""),
597
+ (("gate_proj", "down_proj", "up_proj"), "base_layer."),
598
+ ],
599
+ )
600
+ @pytest.mark.parametrize(
601
+ "suffix,shape,dtype",
602
+ [
603
+ ("weight", (2, 3), torch.bfloat16),
604
+ ("input_scale", (), torch.float32),
605
+ ("g_idx", (2,), torch.int32),
606
+ ],
607
+ )
608
+ def test_per_expert_weights_keep_all_physical_destinations(
609
+ self, projs, lora_prefix, suffix, shape, dtype
610
+ ):
611
+ mapping = self._mapping(
612
+ projs, num_redundant_experts=4, lora_base_layer_prefix=lora_prefix
613
+ )
614
+ for logical_id, physical_ids in [(1, (1, 17)), (10, (10,))]:
615
+ for proj in dict.fromkeys(projs):
616
+ if not proj:
617
+ continue
618
+ weight = torch.full(shape, logical_id + 1, dtype=dtype)
619
+ name = f"{logical_id}.{proj}.{lora_prefix}{suffix}"
620
+ calls = self._load(mapping, [(name, weight)], suffix)
621
+ shards = [s for p, s in zip(projs, ("w1", "w2", "w3")) if p == proj]
622
+ expected = [
623
+ (f"{'w2' if shard == 'w2' else 'w13'}_{suffix}", shard, physical)
624
+ for physical in physical_ids
625
+ for shard in shards
626
+ ]
627
+ assert [call[:3] for call in calls] == expected
628
+ for call in calls:
629
+ torch.testing.assert_close(call[3], weight, rtol=0, atol=0)
630
+
631
+ @pytest.mark.parametrize(
632
+ "projs,fused_name,lora_prefix",
633
+ [
634
+ (("gate_proj", "down_proj", "up_proj"), "gate_up_proj", ""),
635
+ (("w1", "w2", "w3"), "w13", "base_layer."),
636
+ ],
637
+ )
638
+ def test_per_expert_fused_gate_up_keeps_both_halves(
639
+ self, projs, fused_name, lora_prefix
640
+ ):
641
+ mapping = self._mapping(
642
+ projs, num_redundant_experts=4, lora_base_layer_prefix=lora_prefix
643
+ )
644
+ weight = torch.arange(12).reshape(4, 3)
645
+ calls = self._load(mapping, [(f"1.{fused_name}.{lora_prefix}weight", weight)])
646
+ assert [call[:3] for call in calls] == [
647
+ ("w13_weight", shard, expert)
648
+ for expert in (1, 17)
649
+ for shard in ("w1", "w3")
650
+ ]
651
+ for call, expected in zip(calls, weight.chunk(2) * 2):
652
+ torch.testing.assert_close(call[3], expected, rtol=0, atol=0)
653
+
654
+ @pytest.mark.parametrize(
655
+ "suffix,quant_method,transposed",
656
+ [
657
+ ("weight", "tensor", False),
658
+ ("weight", "tensor", True),
659
+ ("weight_scale", "channel", True),
660
+ ("weight_scale", "block", True),
661
+ ],
662
+ )
663
+ def test_mixed_weights_reset_mapping_and_keep_fused_layout(
664
+ self, transposed, suffix, quant_method
665
+ ):
666
+ weight = torch.arange(192).reshape(16, 4, 3)
667
+ checkpoint = (
668
+ weight.transpose(-1, -2)
669
+ if (transposed and (suffix == "weight" or quant_method == "block"))
670
+ else weight
671
+ )
672
+ calls = self._load(
673
+ self._mapping(),
674
+ [
675
+ (f"1.gate_proj.{suffix}", weight[1, :2]),
676
+ ("gate_up_proj" + suffix[6:], checkpoint),
677
+ (f"10.down_proj.{suffix}", weight[10, :2]),
678
+ ],
679
+ suffix,
680
+ transposed=transposed,
681
+ quant_method=quant_method,
682
+ )
683
+ fused_calls = [
684
+ (f"w13_{suffix}", shard, expert)
685
+ for shard in ("w1", "w3")
686
+ for expert in range(16)
687
+ ]
688
+ assert [call[:3] for call in calls] == [
689
+ (f"w13_{suffix}", "w1", 1),
690
+ *fused_calls,
691
+ (f"w2_{suffix}", "w2", 10),
692
+ ]
693
+ expected = [weight[1, :2]]
694
+ expected += [expert for half in weight.chunk(2, dim=1) for expert in half]
695
+ expected += [weight[10, :2]]
696
+ for call, tensor in zip(calls, expected):
697
+ torch.testing.assert_close(call[3], tensor, rtol=0, atol=0)
698
+
699
+ def test_fused_weights_stop_after_first_matching_run(self):
700
+ mapping = self._mapping()
701
+ calls = self._load(
702
+ [mapping[0], mapping[2], mapping[1]],
703
+ [("gate_up_proj", torch.ones(2, 4, 3))],
704
+ )
705
+ assert [call[:3] for call in calls] == [
706
+ ("w13_weight", "w1", 0),
707
+ ("w13_weight", "w1", 1),
708
+ ]
709
+
710
+ def test_repeated_experts_prefix_does_not_hide_later_matches(self):
711
+ calls = self._load(
712
+ self._mapping(),
713
+ [("1.gate_proj.weight", torch.ones(2, 3))],
714
+ layer_name="model.experts.0.child.experts",
715
+ )
716
+ assert [call[:3] for call in calls] == [("w13_weight", "w1", 1)]
717
+
718
+ @pytest.mark.parametrize(
719
+ "name",
720
+ [
721
+ "01.gate_proj.weight",
722
+ "16.gate_proj.weight",
723
+ "gate.weight",
724
+ "non_expert.weight",
725
+ ],
726
+ )
727
+ def test_unmatched_names_do_not_load_experts(self, name):
728
+ assert self._load(self._mapping(), [(name, torch.ones(2, 3))]) == []
729
+
730
+ @pytest.mark.parametrize(
731
+ "weight_name", ["1.gate_proj.", "experts.", "experts.1", ""]
732
+ )
733
+ def test_unusual_matching_entries_are_not_silently_skipped(self, weight_name):
734
+ mapping = self._mapping() + [("missing", weight_name, 1, "w3")]
735
+ with pytest.raises(AttributeError, match="has no parameter"):
736
+ self._load(mapping, [("1.gate_proj.weight", torch.ones(2, 3))])
737
+
738
+
739
class TestPerTensorScaleCoercion:
740
"""Regression test for shape-(1,) per-tensor scales (issue #43297).
741
vllm/model_executor/layers/fused_moe/routed_experts.py
+32
-2
@@ -36,6 +36,28 @@ if TYPE_CHECKING:
36
logger = init_logger(__name__)
37
38
39
+def _index_expert_mapping(
40
+ mapping: list[tuple[str, str, int, str]],
41
+) -> dict[str, list[tuple[str, str, int, str]]]:
42
+ """Index by logical checkpoint ID, preserving entry order and EPLB replicas.
43
+
44
+ An empty index keeps the full scan for unsupported mapping names.
45
+ """
46
+ mapping_by_expert: dict[str, list[tuple[str, str, int, str]]] = {}
47
+ for entry in mapping:
48
+ prefix, _, suffix = entry[1].partition(".")
49
+ expert_key, sep, _ = suffix.partition(".")
50
+ if (
51
+ prefix != "experts"
52
+ or not expert_key
53
+ or (expert_key.isdecimal() and not sep)
54
+ ):
55
+ return {}
56
+ if expert_key.isdecimal():
57
+ mapping_by_expert.setdefault(expert_key, []).append(entry)
58
+ return mapping_by_expert
59
+
60
+
61
class FusedMoeWeightScaleSupported(Enum):
62
TENSOR = "tensor"
63
CHANNEL = "channel"
@@ -889,14 +911,22 @@ class RoutedExperts(PluggableLayer):
911
self, weights: Iterable[tuple[str, torch.Tensor]]
912
) -> Iterable[str]:
913
expert_mapping = self.get_expert_mapping(include_fused=True)
914
+ mapping_by_expert = _index_expert_mapping(expert_mapping)
915
+
916
for expert_name, loaded_weight in weights:
917
qual_name = f"{self.layer_name}.{expert_name}"
894
- # Fused expert weights can be identified by their 3D tensors
918
is_fused = loaded_weight.dim() == 3
919
+ # Fused tensors and ambiguous names keep the full mapping.
920
+ candidates = expert_mapping
921
+ if not is_fused and qual_name.count("experts.") == 1:
922
+ expert_key = qual_name.partition("experts.")[2].partition(".")[0]
923
+ candidates = mapping_by_expert.get(expert_key, expert_mapping)
924
+
925
matched = False
897
- for param_name, weight_name, expert_id, shard_id in expert_mapping:
926
+ for param_name, weight_name, expert_id, shard_id in candidates:
927
if weight_name not in qual_name:
928
if matched and is_fused:
929
+ # Fused tensors use the first contiguous run of matches.
930
break
931
continue
932
matched = True