From 8975479f87dc544c01160891f55a6fe0da3240c1 Mon Sep 17 00:00:00 2001 From: Yanbin Jiang Date: Thu, 30 Apr 2026 22:42:10 -0700 Subject: [PATCH] [LoRA][MOE] Fix EP correctness in MoE LoRA slicing and virtual-experts kernels (#24171) --- python/sglang/srt/lora/lora_manager.py | 6 + python/sglang/srt/lora/mem_pool.py | 28 +- .../srt/lora/triton_ops/virtual_experts.py | 54 +++- .../lora/test_virtual_experts_kernels.py | 259 ++++++++++++++++++ .../unit/lora/test_mem_pool_ep_unit.py | 119 ++++++++ 5 files changed, 452 insertions(+), 14 deletions(-) create mode 100644 test/registered/lora/test_virtual_experts_kernels.py diff --git a/python/sglang/srt/lora/lora_manager.py b/python/sglang/srt/lora/lora_manager.py index 4a70b8123..3ef5ae7f4 100644 --- a/python/sglang/srt/lora/lora_manager.py +++ b/python/sglang/srt/lora/lora_manager.py @@ -810,6 +810,12 @@ class LoRAManager: x in self.target_modules for x in ["gate_up_proj", "down_proj"] ): layer_id = get_layer_id(module_name) + if layer_id is None: + # FusedMoE submodules outside the decoder layer hierarchy + # (e.g. nested helpers under non-".layers." prefixes) have + # no resolvable layer id; skip them so we don't index + # `self.lora_modules` with `None`. + continue lora_module = self.set_lora_module(module_name, module) lora_module.experts_shared_outer_loras = self.experts_shared_outer_loras lora_module.lora_use_virtual_experts = self.lora_use_virtual_experts diff --git a/python/sglang/srt/lora/mem_pool.py b/python/sglang/srt/lora/mem_pool.py index 22a9faed2..1656e03a3 100644 --- a/python/sglang/srt/lora/mem_pool.py +++ b/python/sglang/srt/lora/mem_pool.py @@ -753,6 +753,15 @@ class LoRAMemoryPool: else: temp_B_buffer[target_module] = weights + # Track which buffer keys correspond to a real wrapped module on + # this layer. `temp_A/B_buffer` is seeded with every key in the + # global `A/B_buffer` (union across all layer types), but a + # hybrid-architecture layer (e.g. Qwen3.5 linear-attn vs full-attn, + # or first-k-dense MoE) only owns a subset of those modules. The + # buffer-copy loops below skip non-owned keys to avoid the + # redundant zero-fills on slots no `update_lora_info` ever points + # a forward-time module at. + active_target_modules: Set[str] = set() cur_layer_modules = lora_modules[layer_id] for module_name, module in cur_layer_modules.items(): # TODO (Jonahcb): check if the code can be refactored to avoid the special handling for FusedMoEWithLoRA @@ -760,13 +769,19 @@ class LoRAMemoryPool: from sglang.srt.lora.layers import FusedMoEWithLoRA if isinstance(module, FusedMoEWithLoRA): + # Per-expert MoE weights are sharded along `moe_tp_size` + # (= tp_size // ep_size // dp_size), so the slice index + # must be `moe_tp_rank`. Passing the outer `tp_rank` here + # produces an off-the-end slice when ep_size < tp_size + # (e.g. tp=4 ep=2 → ranks 2,3 slice past intermediate_size). moe_target_modules = ["gate_up_proj_moe", "down_proj_moe"] for target_module in moe_target_modules: + active_target_modules.add(target_module) if temp_A_buffer.get(target_module) is not None: temp_A_buffer[target_module] = ( module.slice_moe_lora_a_weights( temp_A_buffer[target_module], - self.tp_rank, + self.moe_tp_rank, target_module, ) ) @@ -774,7 +789,7 @@ class LoRAMemoryPool: temp_B_buffer[target_module] = ( module.slice_moe_lora_b_weights( temp_B_buffer[target_module], - self.tp_rank, + self.moe_tp_rank, target_module, ) ) @@ -783,6 +798,11 @@ class LoRAMemoryPool: # Handle regular modules target_module = get_target_module_name(module_name, self.target_modules) + # Mark active even if the adapter has no weights for this + # module on this layer — the buffer still needs to be zeroed + # (so a previously-evicted adapter's weights don't leak into + # the new slot) and the wrapped layer module will read it. + active_target_modules.add(target_module) if temp_A_buffer[target_module] is None: # Skip weight slicing if the weight is not present in the adapter @@ -797,6 +817,8 @@ class LoRAMemoryPool: ) for name, weights in temp_A_buffer.items(): + if name not in active_target_modules: + continue c = get_stacked_multiply(name, self.base_model) max_r = self.max_lora_rank target_buffer = self.A_buffer[name][layer_id] @@ -885,6 +907,8 @@ class LoRAMemoryPool: load_lora_weight_tensor(buffer_view, weights) for name, weights in temp_B_buffer.items(): + if name not in active_target_modules: + continue target_buffer = self.B_buffer[name][layer_id] if name in ["gate_up_proj_moe", "down_proj_moe"]: diff --git a/python/sglang/srt/lora/triton_ops/virtual_experts.py b/python/sglang/srt/lora/triton_ops/virtual_experts.py index 3da01ea58..e3467fba5 100644 --- a/python/sglang/srt/lora/triton_ops/virtual_experts.py +++ b/python/sglang/srt/lora/triton_ops/virtual_experts.py @@ -46,7 +46,12 @@ def _fused_virtual_topk_ids_kernel( safe_lora = tl.maximum(lora_id, 0) base = tl.load(topk_ids_ptr + offs, mask=valid, other=0) - result = base + safe_lora * num_experts_for_weight + # Preserve negative sentinel topk_ids (e.g. -1 for non-local experts after + # EP dispatch). Without this, `-1 + safe_lora * num_experts` would land on + # a real virtual-expert slot belonging to another adapter and trigger OOB + # loads in downstream LoRA kernels. + shifted = base + safe_lora * num_experts_for_weight + result = tl.where(base < 0, base, shifted) tl.store(virtual_topk_ids_ptr + offs, result, mask=valid) # Write mask once per row (at first k position) @@ -299,19 +304,40 @@ def _align_block_size_torch( block_size: int, num_experts: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Pure-PyTorch align_block_size for num_experts > 1024, compiled via torch.compile.""" + """Pure-PyTorch align_block_size for num_experts > 1024, compiled via torch.compile. + + Out-of-range topk_ids (negative sentinels left by EP dispatch, or virtual- + expert IDs >= num_experts produced when those sentinels are combined with + a per-adapter offset) are routed into a dedicated sentinel bucket. Without + this, indexing ``padded_offsets[sorted_expert_ids]`` would wrap (-1) or + OOB-read, and the bad expert ids would propagate into the downstream LoRA + GEMM as real expert slots. + """ device = topk_ids.device flat_topk_ids = topk_ids.reshape(-1).to(torch.int64) - num_valid_tokens = flat_topk_ids.numel() + num_total_tokens = flat_topk_ids.numel() + + # Map every invalid id to the sentinel bucket (`num_experts`). The bucket + # itself is allocated below via `bucket_count = num_experts + 1` and is + # excluded from block→expert assignment so its blocks stay marked -1. + sentinel = num_experts + valid_mask = (flat_topk_ids >= 0) & (flat_topk_ids < num_experts) + safe_topk_ids = torch.where( + valid_mask, + flat_topk_ids, + torch.full_like(flat_topk_ids, sentinel), + ) + + bucket_count = num_experts + 1 max_total_padded_tokens = ( - (num_valid_tokens + num_experts * (block_size - 1) + block_size - 1) + (num_total_tokens + bucket_count * (block_size - 1) + block_size - 1) // block_size ) * block_size max_num_blocks = max_total_padded_tokens // block_size sorted_token_ids = torch.full( (max_total_padded_tokens,), - num_valid_tokens, + num_total_tokens, dtype=torch.int32, device=device, ) @@ -322,13 +348,13 @@ def _align_block_size_torch( device=device, ) - if num_valid_tokens == 0: + if num_total_tokens == 0: num_tokens_post_padded = torch.zeros((1,), dtype=torch.int32, device=device) return sorted_token_ids, expert_ids, num_tokens_post_padded - sorted_order = torch.argsort(flat_topk_ids) - sorted_expert_ids = flat_topk_ids[sorted_order] - expert_range = torch.arange(num_experts, device=device, dtype=torch.int64) + sorted_order = torch.argsort(safe_topk_ids) + sorted_expert_ids = safe_topk_ids[sorted_order] + expert_range = torch.arange(bucket_count, device=device, dtype=torch.int64) counts_offsets = torch.searchsorted(sorted_expert_ids, expert_range, right=False) counts_end = torch.searchsorted(sorted_expert_ids, expert_range, right=True) counts = counts_end - counts_offsets @@ -337,7 +363,7 @@ def _align_block_size_torch( padded_offsets = torch.cumsum(padded_counts, dim=0) - padded_counts token_ranks = ( - torch.arange(num_valid_tokens, device=device, dtype=torch.int64) + torch.arange(num_total_tokens, device=device, dtype=torch.int64) - counts_offsets[sorted_expert_ids] ) output_positions = padded_offsets[sorted_expert_ids] + token_ranks @@ -347,13 +373,17 @@ def _align_block_size_torch( sorted_order.to(torch.int32), ) + # Drop the sentinel bucket from the block→expert assignment so its blocks + # remain -1 instead of getting a real expert id from `searchsorted`. block_counts = padded_counts // block_size - actual_num_blocks = block_counts.sum() + real_block_counts = block_counts.clone() + real_block_counts[sentinel] = 0 + actual_num_blocks = real_block_counts.sum() if max_num_blocks <= 0: return sorted_token_ids, expert_ids, total_padded_tokens - block_offsets = torch.cumsum(block_counts, dim=0) + block_offsets = torch.cumsum(real_block_counts, dim=0) all_block_positions = torch.arange(max_num_blocks, device=device, dtype=torch.int64) assigned_experts = torch.searchsorted( block_offsets, all_block_positions, right=True diff --git a/test/registered/lora/test_virtual_experts_kernels.py b/test/registered/lora/test_virtual_experts_kernels.py new file mode 100644 index 000000000..0730e57c3 --- /dev/null +++ b/test/registered/lora/test_virtual_experts_kernels.py @@ -0,0 +1,259 @@ +"""Unit tests for the LoRA virtual-experts kernels under post-EP-dispatch +sentinel `topk_ids` (-1) and out-of-range expert IDs. + +Covers two regression bugs that surface only with `--lora-use-virtual-experts` ++ `ep_size > 1`: + +- `_fused_virtual_topk_ids` must preserve negative sentinel topk_ids. After + EP dispatch, non-local experts arrive as `-1`; the pre-fix kernel mapped + them onto a real virtual-expert slot belonging to another adapter and + triggered OOB loads in downstream LoRA kernels. + +- `_align_block_size_torch` (the `>= 1024`-expert torch.compile fallback) + must route `-1` and `>= num_experts` IDs into a sentinel bucket so they + don't OOB-index `padded_offsets[sorted_expert_ids]` (negative wrap, or + past-end) and don't get assigned to a real expert in the consumer-block + table. + +Both kernels run on CUDA. The torch-compile fallback is gated on +`virtual_num_experts > 1024` in production, but we exercise it directly +here at smaller sizes for cheaper iteration; one test sticks to the >1024 +regime to mirror the production trigger. + +Usage: + python -m pytest test/registered/lora/test_virtual_experts_kernels.py -v +""" + +import unittest + +import torch + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=15, suite="stage-b-test-1-gpu-small") + +from sglang.srt.lora.triton_ops.virtual_experts import ( + _align_block_size_torch, + _fused_virtual_topk_ids, +) + + +class TestFusedVirtualTopkIdsPreservesSentinels(CustomTestCase): + """Item B regression: post-EP-dispatch -1 sentinels must NOT be remapped.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + cls.device = "cuda:0" + + def test_negative_sentinels_preserved(self): + # Mix of valid topk_ids in [0, num_experts), -1 sentinels (typical + # post-EP-dispatch), and a synthetic -2 to ensure the fix doesn't + # depend on the exact -1 value. + topk_ids = torch.tensor( + [ + [3, -1], + [-1, 5], + [0, 7], + [-1, -1], + [2, 9], + [11, -2], + [-1, 4], + [6, -1], + ], + dtype=torch.int32, + device=self.device, + ) + token_lora_mapping = torch.tensor( + [0, 1, 0, 2, -1, 1, 0, 1], dtype=torch.int32, device=self.device + ) + num_experts = 16 + max_loras = 4 + + virtual_ids, _, _ = _fused_virtual_topk_ids( + topk_ids, + token_lora_mapping, + num_experts, + shared_outer=False, + max_loras=max_loras, + ) + + # Every negative input must stay negative (and equal) in the output. + for m in range(topk_ids.shape[0]): + for k in range(topk_ids.shape[1]): + base = topk_ids[m, k].item() + if base < 0: + self.assertEqual( + virtual_ids[m, k].item(), + base, + f"negative sentinel at ({m},{k}) was remapped: " + f"{base} -> {virtual_ids[m, k].item()}", + ) + + def test_positive_topk_remapped_correctly(self): + """Sanity: valid (non-negative) IDs follow the + `base + safe_lora * num_experts` rule.""" + topk_ids = torch.tensor( + [[3, 1], [0, 7], [2, 9]], dtype=torch.int32, device=self.device + ) + token_lora_mapping = torch.tensor( + [0, 1, 2], dtype=torch.int32, device=self.device + ) + num_experts = 16 + max_loras = 4 + + virtual_ids, _, _ = _fused_virtual_topk_ids( + topk_ids, token_lora_mapping, num_experts, False, max_loras + ) + + for m in range(topk_ids.shape[0]): + lora = token_lora_mapping[m].item() + for k in range(topk_ids.shape[1]): + base = topk_ids[m, k].item() + expected = base + max(lora, 0) * num_experts + self.assertEqual(virtual_ids[m, k].item(), expected) + + def test_no_lora_token_does_not_shift_base(self): + """`token_lora_mapping[m] == -1` (no LoRA) keeps `safe_lora=0`, + so positive bases pass through unchanged and the row mask is False.""" + topk_ids = torch.tensor([[3, 5]], dtype=torch.int32, device=self.device) + token_lora_mapping = torch.tensor([-1], dtype=torch.int32, device=self.device) + num_experts = 16 + + virtual_ids, mask, _ = _fused_virtual_topk_ids( + topk_ids, token_lora_mapping, num_experts, False, max_loras=4 + ) + self.assertEqual(virtual_ids[0, 0].item(), 3) + self.assertEqual(virtual_ids[0, 1].item(), 5) + self.assertFalse(bool(mask[0].item())) + + +class TestAlignBlockSizeTorchSentinelBucket(CustomTestCase): + """Item C regression: invalid `topk_ids` must not OOB-index + `padded_offsets[sorted_expert_ids]`, must not be assigned to any real + expert in the consumer-block table, and the function must remain + correct on legitimate (all-valid) inputs.""" + + @classmethod + def setUpClass(cls): + if not torch.cuda.is_available(): + raise unittest.SkipTest("CUDA required") + cls.device = "cuda:0" + + @staticmethod + def _assigned_experts(expert_ids: torch.Tensor) -> list: + """Return the list of real expert ids assigned to blocks (filtering + out -1 sentinels for padding/exclusion).""" + return expert_ids[expert_ids != -1].cpu().tolist() + + def _assert_only_real_or_sentinel(self, expert_ids: torch.Tensor, num_experts: int): + for eid in expert_ids.cpu().tolist(): + self.assertTrue( + eid == -1 or 0 <= eid < num_experts, + f"expert_ids contains invalid value {eid}", + ) + + def test_all_valid_baseline(self): + """Sanity: with no invalid IDs, every present real expert appears + in the assignment, and no junk values leak through.""" + num_experts = 8 + block_size = 16 + topk_ids = torch.tensor( + [[0, 3], [4, 7], [1, 2], [5, 6]], + dtype=torch.int32, + device=self.device, + ) + + _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + + self._assert_only_real_or_sentinel(expert_ids, num_experts) + assigned = set(self._assigned_experts(expert_ids)) + # Every distinct real input expert must be assigned to at least one + # block. + self.assertEqual(assigned, set(range(num_experts))) + + def test_negative_ids_routed_to_sentinel(self): + """`-1` tokens must not appear as real expert assignments and must + not corrupt the assignment of real IDs.""" + num_experts = 8 + block_size = 16 + topk_ids = torch.tensor( + [[0, -1], [-1, 7], [1, -1], [-1, -1]], + dtype=torch.int32, + device=self.device, + ) + + _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + + self._assert_only_real_or_sentinel(expert_ids, num_experts) + assigned = self._assigned_experts(expert_ids) + for valid_eid in (0, 1, 7): + self.assertIn(valid_eid, assigned) + + def test_oor_ids_routed_to_sentinel(self): + """IDs `>= num_experts` (e.g. virtual-experts remap when combined + with non-local sentinels) must not break cumsum/searchsorted and + must not show up as real assignments.""" + num_experts = 8 + block_size = 16 + topk_ids = torch.tensor( + [[0, 100], [50, 7], [1, 200]], + dtype=torch.int32, + device=self.device, + ) + + _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + + self._assert_only_real_or_sentinel(expert_ids, num_experts) + assigned = self._assigned_experts(expert_ids) + for valid_eid in (0, 1, 7, 50): + # 50 is OOR (>= 8), should NOT be in assigned. + if valid_eid >= num_experts: + self.assertNotIn(valid_eid, assigned) + else: + self.assertIn(valid_eid, assigned) + + def test_mixed_invalid_at_production_size(self): + """Mirror the production trigger: `num_experts > 1024` (only path + where `_align_block_size_torch` is invoked instead of the native + align kernel).""" + num_experts = 1500 + block_size = 16 + topk_ids = torch.tensor( + [ + [-1, 500], + [num_experts + 7, 1000], + [num_experts * 2, 100], + [-1, 0], + ], + dtype=torch.int32, + device=self.device, + ) + + _, expert_ids, _ = _align_block_size_torch(topk_ids, block_size, num_experts) + + self._assert_only_real_or_sentinel(expert_ids, num_experts) + assigned = self._assigned_experts(expert_ids) + for valid_eid in (0, 100, 500, 1000): + self.assertIn(valid_eid, assigned) + + def test_empty_topk_ids_does_not_crash(self): + """Edge: empty input. Should return empty/zero outputs without + OOB indexing on the sentinel bucket.""" + num_experts = 8 + block_size = 16 + topk_ids = torch.empty((0, 2), dtype=torch.int32, device=self.device) + + sorted_token_ids, expert_ids, num_post_padded = _align_block_size_torch( + topk_ids, block_size, num_experts + ) + + self.assertEqual(num_post_padded.item(), 0) + # Whatever expert_ids contains, it must be sentinel only. + self.assertEqual(self._assigned_experts(expert_ids), []) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/test/registered/unit/lora/test_mem_pool_ep_unit.py b/test/registered/unit/lora/test_mem_pool_ep_unit.py index 6a374e058..56a68dc98 100644 --- a/test/registered/unit/lora/test_mem_pool_ep_unit.py +++ b/test/registered/unit/lora/test_mem_pool_ep_unit.py @@ -607,5 +607,124 @@ class TestMoeBufferShardsByMoeTp(unittest.TestCase): self.assertEqual(q_b, (2, 48, 8)) +class TestLoadBufferPassesMoeTpRankToSlice(unittest.TestCase): + """Regression: `load_lora_weight_to_buffer` must hand `moe_tp_rank` (not + the outer `tp_rank`) to `slice_moe_lora_{a,b}_weights`. + + Per-expert MoE weights are sharded along + `moe_tp_size = tp_size // ep_size // dp_size`, NOT the outer `tp_size`. + The bug only surfaces when those two values differ — i.e. when + `1 < ep_size < tp_size`. Concrete reproducer (`tp=4 ep=2`): + + moe_tp_size = 2; outer rank 3 has moe_tp_rank=1. + `intermediate_size_per_partition = moe_inter / 2 = 384`. + Slicing with the OUTER rank (3) computes `start = 3 * 384 = 1152`, + which is past the full `moe_inter = 768`, returning a `[r, 0]`-shaped + tensor that fails the shape-match assert in `load_lora_weight_tensor`. + + This test exercises `load_lora_weight_to_buffer` end-to-end with a + minimal mocked `FusedMoEWithLoRA` whose slicer captures-and-raises so + we don't need to satisfy buffer-copy shape constraints. + """ + + class _StopAfterCapture(Exception): + """Sentinel raised from the mocked slicer to short-circuit + execution before the buffer-copy phase (which would need real + shapes the test does not provide).""" + + def test_moe_tp_rank_used_for_slicing_when_ep_lt_tp(self): + from sglang.srt.lora.layers import FusedMoEWithLoRA + + # tp=4 ep=2 → moe_tp_size=2. Pick OUTER rank 3 so moe_tp_rank=1. + # The two values differ; the bug would surface on this exact rank. + pool = LoRAMemoryPool.__new__(LoRAMemoryPool) + pool.tp_size = 4 + pool.tp_rank = 3 + pool.moe_tp_size = 2 + pool.moe_tp_rank = 1 + pool.moe_ep_size = 2 + pool.moe_ep_rank = 1 + pool.moe_use_local_expert_ids = True + pool._num_experts_local = 1 + pool.num_layer = 1 + pool.target_modules = {"gate_up_proj", "down_proj"} + pool.experts_shared_outer_loras = False + pool.strict_loading = False + pool.lora_added_tokens_size = 0 + # Tiny placeholder buffers — the mocked slicer raises before any of + # this is read in the buffer-copy phase. + pool.A_buffer = { + "gate_up_proj_moe": [torch.zeros(1, 1, 1, 1)], + "down_proj_moe": [torch.zeros(1, 1, 1, 1)], + } + pool.B_buffer = { + "gate_up_proj_moe": [torch.zeros(1, 1, 1, 1)], + "down_proj_moe": [torch.zeros(1, 1, 1, 1)], + } + pool.embedding_A_buffer = {} + pool.embedding_B_buffer = {} + pool.lm_head_A_buffer = {} + pool.lm_head_B_buffer = {} + pool.new_embeddings_buffer = {} + + captured_ranks = [] + + moe_mod = mock.MagicMock(spec=FusedMoEWithLoRA) + + def capture_a(weights, tp_rank, target_module): + captured_ranks.append(("A", target_module, tp_rank)) + raise TestLoadBufferPassesMoeTpRankToSlice._StopAfterCapture() + + def capture_b(weights, tp_rank, target_module): + captured_ranks.append(("B", target_module, tp_rank)) + raise TestLoadBufferPassesMoeTpRankToSlice._StopAfterCapture() + + moe_mod.slice_moe_lora_a_weights.side_effect = capture_a + moe_mod.slice_moe_lora_b_weights.side_effect = capture_b + + # Adapter with one per-expert MoE LoRA-A weight. The expert regex + # `experts\.(\d+)\.` must match the key, which routes the weight + # into `temp_A_buffer["gate_up_proj_moe"]` — the dict shape that + # makes `temp_A_buffer.get("gate_up_proj_moe") is not None` true, + # which in turn triggers `slice_moe_lora_a_weights` (and the + # capture). + adapter = mock.MagicMock() + adapter.config.r = 4 + adapter.scaling = 1.0 + adapter.embedding_layers = {} + adapter.added_tokens_embeddings = {} + adapter.layers = [ + types.SimpleNamespace( + weights={ + "model.layers.0.mlp.experts.0.gate_up_proj.lora_A.weight": ( + torch.zeros(8, 4) + ), + }, + ) + ] + + with self.assertRaises(TestLoadBufferPassesMoeTpRankToSlice._StopAfterCapture): + pool.load_lora_weight_to_buffer( + uid="test", + buffer_id=0, + lora_adapter=adapter, + lora_modules=[{"mlp.experts": moe_mod}], + lora_embed_tokens_module=None, + lora_lm_head_module=None, + ) + + self.assertGreater(len(captured_ranks), 0, "slicing was never invoked") + for ab, target_module, rank in captured_ranks: + self.assertEqual( + rank, + pool.moe_tp_rank, + f"slice_moe_lora_{ab.lower()}_weights for {target_module} " + f"received rank={rank}; expected moe_tp_rank=" + f"{pool.moe_tp_rank} (outer tp_rank is {pool.tp_rank}). " + "Passing the outer tp_rank slices past " + "intermediate_size_per_partition when ep_size < tp_size.", + ) + + if __name__ == "__main__": unittest.main()