[LoRA][MOE] Fix EP correctness in MoE LoRA slicing and virtual-experts kernels (#24171)

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

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"]:

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

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)

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()