[LoRA][MOE] Fix EP correctness in MoE LoRA slicing and virtual-experts kernels (#24171)
This commit is contained in:
Yanbin Jiang
@@ -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)
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user