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 5, 06:00 UTC
fix(moe): fix RoutedExpertsCapturer assertion failure with DP>1 and MK path (#37879)
Le Yang committed
Apr 12, 2026 at 22:28 UTC
21fab0a3db091173444112be8caf6b8248ed5043
2 files changed
+103
-5
tests/model_executor/test_routed_experts_capture.py
+82
@@ -1,16 +1,40 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import types
4
+from types import SimpleNamespace
5
+from unittest.mock import patch
6
7
import pytest
8
import torch
9
10
from vllm.distributed.eplb.eplb_state import EplbLayerState
11
from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
12
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
13
+ RoutedExpertsCapturer,
14
+)
15
from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
16
17
pytestmark = pytest.mark.cpu_test
18
19
+_REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"
20
+
21
+
22
+def _capturer_with_buffer(
23
+ *,
24
+ max_tokens: int = 8,
25
+ num_layers: int = 4,
26
+ num_experts_per_tok: int = 2,
27
+ dp_rank: int = 0,
28
+) -> RoutedExpertsCapturer:
29
+ c = RoutedExpertsCapturer()
30
+ c.dp_rank = dp_rank
31
+ c._device_buffer = torch.full(
32
+ (max_tokens, num_layers, num_experts_per_tok),
33
+ -1,
34
+ dtype=torch.int32,
35
+ )
36
+ return c
37
+
38
39
class DummyRouter(BaseRouter):
40
@property
@@ -159,3 +183,61 @@ def test_gpu_model_runner_binding_stage(monkeypatch):
183
assert callable(dummy_module.router.capture_fn)
184
dummy_module.router.capture_fn(torch.tensor([[9, 10]]))
185
assert len(capturer.calls) == 1
186
+
187
+
188
+def test_routed_experts_capturer_single_dp_no_metadata():
189
+ """dp_metadata is None: capture writes the full topk_ids rows."""
190
+ capturer = _capturer_with_buffer(dp_rank=0)
191
+ topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
192
+ ctx = SimpleNamespace(dp_metadata=None)
193
+ with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
194
+ capturer.capture(layer_id=0, topk_ids=topk)
195
+ assert torch.equal(capturer._device_buffer[:3, 0, :], topk)
196
+ assert capturer._device_buffer[3, 0, 0].item() == -1
197
+
198
+
199
+def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
200
+ """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
201
+ capturer = _capturer_with_buffer(dp_rank=1)
202
+ num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
203
+ ctx = SimpleNamespace(
204
+ dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
205
+ )
206
+ # Concatenated order: rank0 rows then rank1 rows.
207
+ topk = torch.tensor(
208
+ [[0, 1], [2, 3], [10, 11], [12, 13], [14, 15]], dtype=torch.int32
209
+ )
210
+ with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
211
+ capturer.capture(layer_id=0, topk_ids=topk)
212
+ want = topk[2:5]
213
+ assert torch.equal(capturer._device_buffer[:3, 0, :], want)
214
+
215
+
216
+def test_routed_experts_capturer_dp_modular_local_tokens():
217
+ """n == token_num_per_dp: topk is already local to this DP rank."""
218
+ capturer = _capturer_with_buffer(dp_rank=1)
219
+ num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
220
+ ctx = SimpleNamespace(
221
+ dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
222
+ )
223
+ topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)
224
+ with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
225
+ capturer.capture(layer_id=0, topk_ids=topk)
226
+ assert torch.equal(capturer._device_buffer[:3, 0, :], topk)
227
+
228
+
229
+def test_routed_experts_capturer_dp_unexpected_batch_raises():
230
+ """Mismatch between topk batch dim and DP layout: fail fast."""
231
+ capturer = _capturer_with_buffer(dp_rank=0)
232
+ num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
233
+ ctx = SimpleNamespace(
234
+ dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
235
+ )
236
+ # total=5, local=2: n=1 matches neither naive (5) nor modular (2).
237
+ topk = torch.tensor([[1, 2]], dtype=torch.int32)
238
+ with (
239
+ patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
240
+ pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),
241
+ ):
242
+ capturer.capture(layer_id=0, topk_ids=topk)
243
+ assert capturer._device_buffer[0, 0, 0].item() == -1
vllm/model_executor/layers/fused_moe/routed_experts_capturer.py
+21
-5
@@ -176,11 +176,27 @@ class RoutedExpertsCapturer:
176
end_loc = topk_ids.shape[0]
177
token_num_per_dp = topk_ids.shape[0]
178
else: # multi dp
179
- token_num_per_dp = ctx.dp_metadata.num_tokens_across_dp_cpu[self.dp_rank]
180
- cumsum = torch.cumsum(ctx.dp_metadata.num_tokens_across_dp_cpu, dim=0)
181
- assert cumsum[-1] == topk_ids.shape[0]
182
- end_loc = cumsum[self.dp_rank]
183
- start_loc = end_loc - token_num_per_dp
179
+ num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu
180
+ token_num_per_dp = int(num_tokens_dp[self.dp_rank].item())
181
+ total = int(num_tokens_dp.sum().item())
182
+ n = topk_ids.shape[0]
183
+
184
+ if n == total:
185
+ # Naive dispatch: all DP ranks' tokens concatenated before routing.
186
+ cumsum = torch.cumsum(num_tokens_dp, dim=0)
187
+ end_loc = int(cumsum[self.dp_rank].item())
188
+ start_loc = end_loc - token_num_per_dp
189
+ elif n == token_num_per_dp:
190
+ # Modular-kernel path: DP combine happens inside quant_method.apply;
191
+ # select_experts only sees this rank's tokens.
192
+ start_loc = 0
193
+ end_loc = token_num_per_dp
194
+ else:
195
+ raise AssertionError(
196
+ "RoutedExpertsCapturer: unexpected topk_ids batch dim "
197
+ f"{n} (expected {total} or {token_num_per_dp} "
198
+ f"for dp_rank={self.dp_rank})"
199
+ )
200
201
if layer_id >= self._device_buffer.shape[1]:
202
return