A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 16:15 UTC
[EPLB] Fix replica selection bias in fused_moe router (#40810)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com>
Artem Perevedentsev committed
Apr 25, 2026 at 01:06 UTC
333529deae59cd4100df540f225470c9bc539bee
2 files changed
+70
-15
tests/kernels/moe/test_routing.py
+62
-12
@@ -662,6 +662,52 @@ def test_eplb_map_no_redundancy(
662
assert load.sum().item() == 0
663
664
665
+@pytest.mark.parametrize("top_k,R", [(2, 2), (4, 2), (8, 4), (8, 8)])
666
+def test_eplb_map_hot_expert_replica_balance(top_k, R):
667
+ """Hot logical expert with R replicas must be balanced across replicas
668
+ even when ``top_k`` is a multiple of ``R``. In that regime every top-k
669
+ offset for the hot expert lands on a multiple of ``top_k`` in the flat
670
+ ``topk_ids`` view, so per-replica assignment must not collapse onto a
671
+ single replica.
672
+ """
673
+ num_tokens = 8192
674
+ num_logical = 16
675
+ num_physical = R + (num_logical - 1)
676
+
677
+ l2p = torch.full((num_logical, R), -1, dtype=torch.int64, device="cuda")
678
+ l2p[0] = torch.arange(R, dtype=torch.int64, device="cuda")
679
+ for i in range(1, num_logical):
680
+ l2p[i, 0] = R + i - 1
681
+ rc = torch.tensor([R] + [1] * (num_logical - 1), dtype=torch.int64, device="cuda")
682
+
683
+ torch.manual_seed(0)
684
+ topk_ids = torch.randint(
685
+ 1,
686
+ num_logical,
687
+ (num_tokens, top_k),
688
+ dtype=torch.int32,
689
+ device="cuda",
690
+ )
691
+ topk_ids[:, 0] = 0
692
+
693
+ load = torch.zeros(num_physical, dtype=torch.int32, device="cuda")
694
+ rec = torch.tensor(True, dtype=torch.bool, device="cuda")
695
+
696
+ eplb_map_to_physical_and_record(
697
+ topk_ids=topk_ids,
698
+ expert_load_view=load,
699
+ logical_to_physical_map=l2p,
700
+ logical_replica_count=rc,
701
+ record_enabled=rec,
702
+ )
703
+
704
+ hot_load = load[:R].float()
705
+ max_mean = (hot_load.max() / hot_load.mean()).item()
706
+ assert max_mean < 1.15, (
707
+ f"Hot expert replicas uneven: {hot_load.tolist()}, max/mean={max_mean:.3f}"
708
+ )
709
+
710
+
711
@pytest.mark.parametrize("record_enabled", [True, False])
712
@pytest.mark.parametrize(
713
"l2p_map, replica_count, num_physical, topk_ids, expected_out, expected_load",
@@ -672,10 +718,12 @@ def test_eplb_map_no_redundancy(
718
[2, 2, 1, 1],
719
6,
720
[[0, 1], [2, 3], [0, 2]],
675
- # offs: 0→0%2=0→p0, 1→1%2=1→p5, 2→2%1=0→p2,
676
- # 3→3%1=0→p3, 4→4%2=0→p0, 5→5%1=0→p2
677
- [[0, 5], [2, 3], [0, 2]],
678
- [2, 0, 2, 1, 0, 1],
721
+ # replica = (token_idx * KNUTH) & 0xFFFFFFFF % R.
722
+ # token 0 hash=0x00000000: %2=0, %1=0.
723
+ # token 1 hash=0x9E3779B9: %2=1, %1=0.
724
+ # token 2 hash=0x3C6EF372: %2=0, %1=0.
725
+ [[0, 1], [2, 3], [0, 2]],
726
+ [2, 1, 2, 1, 0, 0],
727
id="partial",
728
),
729
pytest.param(
@@ -684,10 +732,11 @@ def test_eplb_map_no_redundancy(
732
[2, 2, 2, 2],
733
8,
734
[[0, 1], [2, 3], [0, 2]],
687
- # offs: 0→0%2=0→p0, 1→1%2=1→p5, 2→2%2=0→p2,
688
- # 3→3%2=1→p7, 4→4%2=0→p0, 5→5%2=1→p6
689
- [[0, 5], [2, 7], [0, 6]],
690
- [2, 0, 1, 0, 0, 1, 1, 1],
735
+ # token 0 hash=0x00000000: %2=0.
736
+ # token 1 hash=0x9E3779B9: %2=1.
737
+ # token 2 hash=0x3C6EF372: %2=0.
738
+ [[0, 1], [6, 7], [0, 2]],
739
+ [2, 1, 1, 0, 0, 0, 1, 1],
740
id="full",
741
),
742
pytest.param(
@@ -696,10 +745,11 @@ def test_eplb_map_no_redundancy(
745
[4, 2, 2],
746
8,
747
[[0, 1], [2, 0], [1, 2]],
699
- # offs: 0→0%4=0→p0, 1→1%2=1→p4, 2→2%2=0→p2,
700
- # 3→3%4=3→p7, 4→4%2=0→p1, 5→5%2=1→p6
701
- [[0, 4], [2, 7], [1, 6]],
702
- [1, 1, 1, 0, 1, 0, 1, 1],
748
+ # token 0 hash=0x00000000: %4=0, %2=0.
749
+ # token 1 hash=0x9E3779B9: %4=1, %2=1.
750
+ # token 2 hash=0x3C6EF372: %4=2, %2=0.
751
+ [[0, 1], [6, 3], [1, 2]],
752
+ [1, 2, 1, 1, 0, 0, 1, 0],
753
id="uneven",
754
),
755
],
vllm/model_executor/layers/fused_moe/router/base_router.py
+8
-3
@@ -26,6 +26,7 @@ if current_platform.is_cuda_alike():
26
map_slots,
27
out_size,
28
numel,
29
+ num_active_experts,
30
BLOCK_SIZE: tl.constexpr,
31
):
32
pid = tl.program_id(0)
@@ -37,7 +38,6 @@ if current_platform.is_cuda_alike():
38
safe_expert_id = tl.where(valid_expert, expert_id, 0)
39
40
# 1. Convert the logical expert ids to physical expert ids
40
- # Directly select a random replica for each logical expert
41
replica_count = tl.load(
42
logical_replica_count_ptr + safe_expert_id,
43
mask=mask & valid_expert,
@@ -45,8 +45,11 @@ if current_platform.is_cuda_alike():
45
)
46
# Avoid invalid modulo/div by forcing at least 1.
47
replica_count = tl.maximum(replica_count, 1)
48
- # Match torch.compile path: use flattened token position.
49
- replica_idx = offs % replica_count
48
+ # floor(2^32 / phi), classic Knuth multiplicative hash multiplier.
49
+ KNUTH_MULTIPLIER = 2654435769
50
+ token_idx = (offs // num_active_experts).to(tl.int64)
51
+ hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF
52
+ replica_idx = hashed % replica_count
53
54
# 2. Record expert load metrics.
55
@@ -85,6 +88,7 @@ if current_platform.is_cuda_alike():
88
numel = topk_ids_in.numel()
89
if numel == 0:
90
return topk_ids
91
+ num_active_experts = topk_ids_in.shape[-1]
92
out_flat = torch.empty((numel,), device=topk_ids.device, dtype=topk_ids.dtype)
93
grid = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),)
94
assert expert_load_view.is_contiguous()
@@ -99,6 +103,7 @@ if current_platform.is_cuda_alike():
103
logical_to_physical_map.shape[1],
104
expert_load_view.shape[0],
105
numel,
106
+ num_active_experts,
107
BLOCK_SIZE=256,
108
)
109
return out_flat.reshape(topk_ids.shape)