diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index 5416fdc8ea6..896a214a4a6 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -7,9 +7,16 @@ from unittest.mock import Mock, patch import pytest import torch -from vllm.config import VllmConfig +from vllm.config import CUDAGraphMode, VllmConfig from vllm.config.compilation import CompilationMode from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import ( + BatchDescriptor, + DPMetadata, + ForwardContext, + get_forward_context, + set_forward_context, +) from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import CPUExpertsInt4 from vllm.model_executor.layers.fused_moe.modular_kernel import ( @@ -36,6 +43,12 @@ pytestmark = pytest.mark.cpu_test _REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer" +def _forward_context(**kwargs): + return ForwardContext( + no_compile_layers={}, attn_metadata={}, slot_mapping={}, **kwargs + ) + + def _capturer_with_buffer( *, max_tokens: int = 8, @@ -319,13 +332,120 @@ def test_routed_experts_capturer_single_dp_no_metadata(): """dp_metadata is None: capture writes the full topk_ids rows.""" capturer = _capturer_with_buffer(dp_rank=0) topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32) - ctx = SimpleNamespace(dp_metadata=None) + ctx = _forward_context() with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): capturer.capture(layer_id=0, topk_ids=topk) assert torch.equal(capturer.device_buffer[:3, 0, :], topk) assert capturer.device_buffer[3, 0, 0].item() == -1 +@pytest.mark.parametrize("tp_size, num_tokens", [(2, 6), (2, 3), (2, 1), (4, 2)]) +@pytest.mark.parametrize( + "cg_mode", [CUDAGraphMode.NONE, CUDAGraphMode.FULL, CUDAGraphMode.PIECEWISE] +) +def test_single_dp_ep_capture_restores_tp_shards(tp_size, num_tokens, cg_mode): + """Capture all current rows after model-level SP, including buffer reuse.""" + capturer = _capturer_with_buffer(tp_size=tp_size) + config = SimpleNamespace( + parallel_config=SimpleNamespace( + data_parallel_size=1, + tensor_parallel_size=tp_size, + enable_expert_parallel=True, + use_sequence_parallel_moe=False, + is_moe_model=True, + ), + compilation_config=SimpleNamespace( + fast_moe_cold_start=False, static_forward_context={} + ), + ) + tp_group = Mock() + shard_n = (num_tokens + tp_size - 1) // tp_size + with ( + set_forward_context( + None, + config, + num_tokens=num_tokens, + cudagraph_runtime_mode=cg_mode, + ), + patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group), + ): + assert get_forward_context().dp_metadata is None + for offset in (20, 1): + expected = torch.arange(num_tokens * 2, dtype=torch.int32).reshape(-1, 2) + expected += offset + padded = torch.full((shard_n * tp_size, 2), 99, dtype=torch.int32) + padded[:num_tokens] = expected + tp_group.all_gather.return_value = padded + capturer.capture(0, padded[:shard_n]) + assert torch.equal( + capturer.snapshot_routing_data(num_tokens)[:, 0].int(), expected + ) + assert (capturer.device_buffer[num_tokens:, 0] == -1).all() + if num_tokens == 1: + tp_group.all_gather.assert_not_called() + else: + assert tp_group.all_gather.call_count == 2 + + +@pytest.mark.parametrize("num_tokens, num_rows", [(6, 6), (3, 4), (0, 0)]) +def test_single_dp_tp_capture_does_not_regather_full_batch(num_tokens, num_rows): + capturer = _capturer_with_buffer(tp_size=2) + topk = torch.arange(num_rows * 2, dtype=torch.int32).reshape(-1, 2) + ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens)) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group, + ): + capturer.capture(0, topk) + get_tp_group.assert_not_called() + assert torch.equal(capturer.device_buffer[:num_tokens, 0], topk[:num_tokens]) + assert (capturer.device_buffer[num_tokens:, 0] == -1).all() + + +@pytest.mark.parametrize("num_tokens, num_rows", [(6, 2), (0, 1)]) +def test_single_dp_tp_capture_rejects_unexpected_batch_without_writing( + num_tokens, num_rows +): + capturer = _capturer_with_buffer(tp_size=2) + topk = torch.ones(num_rows, 2, dtype=torch.int32) + ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=num_tokens)) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + pytest.raises(AssertionError, match="unexpected topk_ids batch dim"), + ): + capturer.capture(0, topk) + assert (capturer.device_buffer == -1).all() + + +@pytest.mark.parametrize("tp_rank", [0, 1]) +def test_capture_uses_pre_sp_size_after_model_shards_padding_mask(tp_rank): + """Model-level SP shortens the mask; capture still needs all three rows.""" + from vllm.models.common.ops import sequence_parallel as sp + + capturer = _capturer_with_buffer(tp_size=2) + expected = torch.arange(8, dtype=torch.int32).reshape(4, 2) + ctx = _forward_context( + is_padding=torch.zeros(3, dtype=torch.bool), + batch_descriptor=BatchDescriptor(num_tokens=3), + ) + with ( + patch.object(sp, "get_tensor_model_parallel_world_size", return_value=2), + patch.object(sp, "get_tensor_model_parallel_rank", return_value=tp_rank), + ): + ctx.is_padding = sp.sp_padding_mask(ctx.is_padding, torch.empty(3, 1)) + assert len(ctx.is_padding) == 2 + tp_group = Mock() + tp_group.all_gather.return_value = expected + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group", return_value=tp_group), + ): + capturer.capture(0, expected[2 * tp_rank : 2 * (tp_rank + 1)]) + tp_group.all_gather.assert_called_once() + torch.testing.assert_close(capturer.device_buffer[:3, 0], expected[:3]) + assert (capturer.device_buffer[3:, 0] == -1).all() + + @pytest.mark.parametrize( ("num_experts", "dtype_name", "torch_dtype"), [(256, "uint8", torch.uint8), (257, "uint16", torch.uint16)], @@ -359,7 +479,7 @@ def test_routed_experts_capturer_narrows_snapshot(output_dtype): capturer = _capturer_with_buffer(dtype=torch.int32) capturer.output_dtype = output_dtype topk = torch.tensor([[1, 2], [254, 255]], dtype=torch.int64) - ctx = SimpleNamespace(dp_metadata=None) + ctx = _forward_context() with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): capturer.capture(layer_id=0, topk_ids=topk) @@ -375,8 +495,8 @@ def test_routed_experts_capturer_dp_naive_concatenated_all_ranks(): """N == sum(num_tokens_dp): slice this rank's segment from concatenated topk.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) ) # Concatenated order: rank0 rows then rank1 rows. topk = torch.tensor( @@ -392,8 +512,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens(): """N == token_num_per_dp: topk is already local to this DP rank.""" capturer = _capturer_with_buffer(dp_rank=1) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) ) topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32) with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx): @@ -404,8 +524,8 @@ def test_routed_experts_capturer_dp_modular_local_tokens(): def test_routed_experts_capturer_dp_ep_gathered_shards(): capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) num_tokens_dp = torch.tensor([3, 2], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace( + ctx = _forward_context( + dp_metadata=DPMetadata( num_tokens_across_dp_cpu=num_tokens_dp, local_sizes=[2, 2, 1, 1], ) @@ -421,13 +541,17 @@ def test_routed_experts_capturer_dp_ep_gathered_shards(): assert torch.equal(capturer.device_buffer[:2, 0], topk[4:]) -def test_routed_experts_capturer_sp_modular_gathers_tp_shards(): - capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) +@pytest.mark.parametrize("single_dp", [False, True]) +def test_routed_experts_capturer_sp_modular_gathers_tp_shards(single_dp): + capturer = _capturer_with_buffer(dp_rank=0 if single_dp else 1, tp_size=2) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=( + None if single_dp else DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) + ), + batch_descriptor=BatchDescriptor(num_tokens=3), ) - local_shard = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32) + local_shard = torch.tensor([[10, 12], [11, 13]], dtype=torch.int32).t() gathered = torch.tensor([[10, 11], [12, 13], [14, 15], [-1, -1]], dtype=torch.int32) tp_group = Mock() tp_group.all_gather.return_value = gathered @@ -438,16 +562,83 @@ def test_routed_experts_capturer_sp_modular_gathers_tp_shards(): ): capturer.capture(layer_id=0, topk_ids=local_shard) - tp_group.all_gather.assert_called_once_with(local_shard, dim=0) + tp_group.all_gather.assert_called_once() + gather_input = tp_group.all_gather.call_args.args[0] + assert gather_input.is_contiguous() + torch.testing.assert_close(gather_input, local_shard) + assert tp_group.all_gather.call_args.kwargs == {"dim": 0} assert torch.equal(capturer.device_buffer[:3, 0], gathered[:3]) +def test_routed_experts_capturer_rejects_incomplete_dp_shard_groups(): + """Invalid gather metadata must not silently select another rank's rows.""" + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) + ctx = _forward_context( + dp_metadata=DPMetadata( + num_tokens_across_dp_cpu=torch.tensor([3, 2]), + local_sizes=[2, 2, 2], + ) + ) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + pytest.raises(ValueError, match="not divisible by DP rank count"), + ): + capturer.capture(0, torch.ones(6, 2, dtype=torch.int32)) + assert (capturer.device_buffer == -1).all() + + +@pytest.mark.parametrize("layer_id", [-1, 4]) +def test_routed_experts_capturer_rejects_layer_before_restoring_rows(layer_id): + capturer = _capturer_with_buffer(tp_size=2) + ctx = _forward_context(batch_descriptor=BatchDescriptor(num_tokens=3)) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group, + pytest.raises(IndexError, match="exceeds capture buffer"), + ): + capturer.capture(layer_id, torch.ones(2, 2, dtype=torch.int32)) + get_tp_group.assert_not_called() + assert (capturer.device_buffer == -1).all() + + +@pytest.mark.parametrize( + ("counts", "local_sizes", "start", "stop"), + [ + ([3, 0], [2, 2, 0, 0], 3, 3), + ([0, 3], None, 0, 3), + ([0, 0], [0, 0, 0, 0], 0, 0), + ([2, 2], [1, 1, 1, 1], 2, 4), + ], + ids=["idle-rank", "only-active-rank", "all-idle", "unpadded-gather-tie"], +) +def test_routed_experts_capturer_dp_preserves_empty_and_unpadded_spans( + counts, local_sizes, start, stop +): + capturer = _capturer_with_buffer(dp_rank=1, tp_size=2) + topk = torch.arange(sum(counts) * 2, dtype=torch.int32).reshape(-1, 2) + ctx = _forward_context( + dp_metadata=DPMetadata( + num_tokens_across_dp_cpu=torch.tensor(counts), local_sizes=local_sizes + ) + ) + with ( + patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx), + patch(f"{_REC_MODULE}.get_tp_group") as get_tp_group, + ): + capturer.capture(0, topk) + get_tp_group.assert_not_called() + torch.testing.assert_close( + capturer.device_buffer[: stop - start, 0], topk[start:stop] + ) + assert (capturer.device_buffer[stop - start :, 0] == -1).all() + + def test_routed_experts_capturer_dp_unexpected_batch_raises(): """Mismatch between topk batch dim and DP layout: fail fast.""" capturer = _capturer_with_buffer(dp_rank=0) num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32) - ctx = SimpleNamespace( - dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp) + ctx = _forward_context( + dp_metadata=DPMetadata(num_tokens_across_dp_cpu=num_tokens_dp) ) # total=5, local=2: n=1 matches neither naive (5) nor modular (2). topk = torch.tensor([[1, 2]], dtype=torch.int32) diff --git a/tests/models/test_deepseek_v41_replay_batch.py b/tests/models/test_deepseek_v41_replay_batch.py index 42d6d4d7256..d34f9a2002f 100644 --- a/tests/models/test_deepseek_v41_replay_batch.py +++ b/tests/models/test_deepseek_v41_replay_batch.py @@ -4,7 +4,7 @@ from dataclasses import replace from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import numpy as np import pytest @@ -16,6 +16,9 @@ from vllm.forward_context import ( ForwardContext, override_forward_context, ) +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, +) from vllm.models.deepseek_v41.decoder_replay_layers import DecoderReplayLayers from vllm.models.deepseek_v41.nvidia.model_state import DeepseekV41ModelState from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor @@ -173,6 +176,31 @@ def test_replay_batch_keeps_each_request_window(state): assert context.dp_metadata is None and context.is_padding is None +def test_replay_context_restores_tp_routing_rows(state): + """Capture uses the compact replay batch's size, not the outer batch's.""" + batch = _input_batch(QUERY_LENS, SEQ_LENS, PREFILLING) + replay, _ = _prepare(state, batch, CUDAGraphMode.NONE) + context = replay.forward_context + assert context.is_padding is None and context.dp_metadata is None + count = len(REPLAY_ROWS) + assert context.batch_descriptor.num_tokens == count + capturer = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer) + capturer.dp_rank, capturer.tp_size = 0, 2 + capturer.device_buffer = torch.full((401, 1, 2), -1, device=DEVICE) + expected = torch.tensor(REPLAY_ROWS, device=DEVICE)[:, None].expand(-1, 2) + gathered = torch.cat((expected, torch.full((1, 2), -1, device=DEVICE))) + with ( + override_forward_context(context), + patch( + "vllm.model_executor.layers.fused_moe.routed_experts_capturer.get_tp_group" + ) as group, + ): + group.return_value.all_gather.return_value = gathered + capturer.capture(0, gathered[: (count + 1) // 2]) + torch.testing.assert_close(capturer.device_buffer[:count, 0], expected) + assert (capturer.device_buffer[count:] == -1).all() + + def test_replay_batch_keeps_device_decode_boundaries(state): """Adaptive verification resizes decodes on the GPU (CPU [2, 2], device [1, 3]): their rows stay and the boundaries shift only by trimmed rows.""" diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 9ddc8e9255e..95bb83dee8c 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -327,10 +327,8 @@ def set_forward_context( vllm_config.parallel_config, num_tokens or 0, num_tokens_across_dp ) - # Convenience: if cudagraph is used and num_tokens is given, we can just - # create a batch descriptor here if not given (there's no harm since if it - # doesn't match in the wrapper it'll fall through). - if cudagraph_runtime_mode != CUDAGraphMode.NONE and num_tokens is not None: + # Preserve the batch size before model-level sequence parallelism. + if num_tokens is not None: batch_descriptor = batch_descriptor or BatchDescriptor(num_tokens=num_tokens) additional_kwargs = current_platform.set_additional_forward_context( diff --git a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py index 5d111fa635f..8ece4da679e 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts_capturer.py @@ -14,7 +14,7 @@ import torch from vllm.config import VllmConfig from vllm.distributed.parallel_state import get_tp_group -from vllm.forward_context import get_forward_context +from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig from vllm.platforms import current_platform @@ -85,103 +85,74 @@ class RoutedExpertsCapturer: self.tp_size = vllm_config.parallel_config.tensor_parallel_size def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None: - """Capture expert routing decisions for a specific layer. - - Under data parallelism, ``topk_ids`` may have four different batch - layouts depending on where the DP combine happens and whether - Expert Parallelism (EP) or Sequence Parallelism (SP) is active for the - MoE layer: - - ``n == total`` (naive dispatch): all DP ranks' tokens are - concatenated before routing; we slice out this rank's span - using the cumulative per-rank counts. - - ``n == token_num_per_dp`` (modular-kernel path): DP combine - happens inside ``quant_method.apply``; ``select_experts`` only - ever sees this rank's tokens, so we take the whole tensor. - - ``n == sum(dp_metadata.local_sizes)`` (naive DP+EP dispatch): - sequence-parallel shards from every DP rank are gathered through - the flattened EP group. The shard sizes include CUDA-graph / SP - padding, so we use them to locate this DP rank's unpadded rows. - - ``n == ceil(token_num_per_dp / tp_size)`` (SP + modular-kernel - path): tokens were split along dim=0 across the TP group by - ``_sequence_parallel_context`` - (``moe_runner_base.py:_sequence_parallel_context``), so each - TP rank only sees its shard. We all-gather along dim=0 to - reconstruct this DP rank's full routing tensor. SP pads with - ceil-div (see ``_compute_sp_num_tokens`` in - ``forward_context.py``), so the gathered tensor may contain a - few trailing padding rows which are trimmed by the downstream - ``[:token_num_per_dp]`` slice. - - Args: - layer_id: The layer index. - topk_ids: Tensor of shape (batch_size, num_routed_experts). - - """ - ctx = get_forward_context() - if ctx.dp_metadata is None: - local_topk_ids = topk_ids - else: - num_tokens_dp = ctx.dp_metadata.num_tokens_across_dp_cpu - num_local_tokens = int(num_tokens_dp[self.dp_rank].item()) - total = int(num_tokens_dp.sum().item()) - n = topk_ids.shape[0] - shard_sizes = getattr(ctx.dp_metadata, "local_sizes", None) - gathered_size = sum(shard_sizes) if shard_sizes is not None else -1 - - if n == total: - end = int(num_tokens_dp[: self.dp_rank + 1].sum().item()) - local_topk_ids = topk_ids[end - num_local_tokens : end] - elif n == num_local_tokens: - local_topk_ids = topk_ids - elif shard_sizes is not None and n == gathered_size: - # Naive DP+EP dispatch gathers the sequence-parallel shards in - # flattened DP-then-TP order. ``local_sizes`` is the exact - # all-gatherv layout, including padding. Select this DP rank's - # contiguous shard group, then trim its trailing padding. - num_dp_ranks = len(num_tokens_dp) - assert len(shard_sizes) % num_dp_ranks == 0 - shards_per_dp_rank = len(shard_sizes) // num_dp_ranks - first_shard = self.dp_rank * shards_per_dp_rank - start = sum(shard_sizes[:first_shard]) - local_topk_ids = topk_ids[start : start + num_local_tokens] - elif ( - self.tp_size > 1 - and n == (num_local_tokens + self.tp_size - 1) // self.tp_size - ): - # SP + modular-kernel path. All-gather across the TP - # group along dim=0 to reconstruct the full per-DP-rank - # tensor; keep only the first ``token_num_per_dp`` rows - # (trailing rows are SP ceil-div padding). The TP group - # is always initialized on real rollout workers, and - # every rank in the group reaches this branch in - # lockstep (bind is per-FusedMoEFactory layer, SP is a global - # condition), so a bare all_gather here will not - # deadlock -- let it raise if the precondition is - # violated rather than skip silently. - # - # ``topk_ids`` is already whatever the router produced - # (typically int32/int64, both supported by NCCL); the - # downstream buffer assignment narrows to the capture dtype. - local_topk_ids = get_tp_group().all_gather(topk_ids, dim=0)[ - :num_local_tokens - ] - else: - sp_expected = (num_local_tokens + self.tp_size - 1) // self.tp_size - raise AssertionError( - "RoutedExpertsCapturer: unexpected topk_ids batch " - f"dim {n} (expected {total}, {num_local_tokens}, " - f"{gathered_size}, or {sp_expected} for " - f"dp_rank={self.dp_rank}, tp_size={self.tp_size})" - ) - if not 0 <= layer_id < self.device_buffer.shape[1]: raise IndexError( f"routed-experts layer {layer_id} exceeds capture buffer " f"layer count {self.device_buffer.shape[1]}" ) + local_topk_ids = self._restore_local_topk_ids(topk_ids, get_forward_context()) self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids + def _restore_local_topk_ids( + self, topk_ids: torch.Tensor, ctx: ForwardContext + ) -> torch.Tensor: + """Restore this DP rank's token rows, trimming trailing SP padding. + + Layout priority: cross-DP batch, local batch, gathered DP/TP shards, + then a local TP shard. Gathered rows are DP-major/TP-minor; preserve + this priority when layout sizes coincide. PCP is unsupported. + + Without DP metadata, use the pre-SP batch descriptor. Unknown counts + pass through; known counts must match a full, sharded or padded batch. + A full local batch avoids gathering, including zero/one-token batches + where only the publishing TP0 rank needs the real rows. + """ + n = topk_ids.shape[0] + dp_metadata = ctx.dp_metadata + if dp_metadata is None: + if ctx.batch_descriptor is None: + return topk_ids + num_local_tokens = ctx.batch_descriptor.num_tokens + else: + num_tokens_dp = dp_metadata.num_tokens_across_dp_cpu + num_local_tokens = int(num_tokens_dp[self.dp_rank].item()) + if n == int(num_tokens_dp.sum().item()): + end = int(num_tokens_dp[: self.dp_rank + 1].sum().item()) + return topk_ids[end - num_local_tokens : end] + + if n == num_local_tokens: + return topk_ids + + if dp_metadata is not None: + shard_sizes = dp_metadata.local_sizes + if shard_sizes is not None and n == sum(shard_sizes): + num_dp_ranks = len(num_tokens_dp) + if len(shard_sizes) % num_dp_ranks != 0: + raise ValueError( + "RoutedExpertsCapturer: gathered shard count " + f"{len(shard_sizes)} is not divisible by DP rank count " + f"{num_dp_ranks}" + ) + shards_per_dp_rank = len(shard_sizes) // num_dp_ranks + first_shard = self.dp_rank * shards_per_dp_rank + start = sum(shard_sizes[:first_shard]) + return topk_ids[start : start + num_local_tokens] + + shard_n = (num_local_tokens + self.tp_size - 1) // self.tp_size + if n == shard_n: + return get_tp_group().all_gather(topk_ids.contiguous(), dim=0)[ + :num_local_tokens + ] + if dp_metadata is None and n == shard_n * self.tp_size: + return topk_ids[:num_local_tokens] + + raise AssertionError( + f"RoutedExpertsCapturer: unexpected topk_ids batch dim {n} " + f"(local_tokens={num_local_tokens}, " + f"dp_rank={self.dp_rank}, tp_size={self.tp_size})" + ) + def snapshot_routing_data(self, num_tokens: int) -> torch.Tensor: """Return a stable snapshot of the current routing data.""" return self.device_buffer[:num_tokens].to(self.output_dtype)