Skip to content

vllm.model_executor.layers.fused_moe.routed_experts_capturer

Classes:

Functions:

RoutedExpertsCapturer

Worker-side capturer for routed experts, lives on GPU.

Layer-level hooks call :meth:capture inside the forward pass. Routing rows owned by this DP rank are written into a preallocated device buffer.

The device buffer uses int32. Stable snapshots use the narrowest dtype that can represent every logical expert ID.

Invariants
  • One instance per worker; shape is fixed at init and covers the worst-case step (max_num_batched_tokens tokens).
  • Every routed layer overwrites the current step's token rows.

Methods:

  • capture –

    Capture expert routing decisions for a specific layer.

  • snapshot_routing_data –

    Return a stable snapshot of the current routing data.

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
class RoutedExpertsCapturer:
    """Worker-side capturer for routed experts, lives on GPU.

    Layer-level hooks call :meth:`capture` inside the forward pass. Routing
    rows owned by this DP rank are written into a preallocated device buffer.

    The device buffer uses ``int32``. Stable snapshots use the narrowest dtype
    that can represent every logical expert ID.

    Invariants:
        - One instance per worker; shape is fixed at init and covers the
          worst-case step (``max_num_batched_tokens`` tokens).
        - Every routed layer overwrites the current step's token rows.
    """

    def __init__(
        self,
        max_num_batched_tokens: int,
        vllm_config: VllmConfig,
    ) -> None:
        num_layers, num_experts, num_experts_per_tok = _get_routed_experts_shape(
            vllm_config
        )
        self.shape_per_token = (num_layers, num_experts_per_tok)
        self.output_dtype_name = "uint8" if num_experts <= 256 else "uint16"
        self.output_dtype = getattr(torch, self.output_dtype_name)
        logger.info(
            "RoutedExpertsCapturer: allocating buffer with "
            "max_tokens=%d, num_layers=%d, num_experts_per_tok=%d "
            "(hf_config.model_type=%s)",
            max_num_batched_tokens,
            num_layers,
            num_experts_per_tok,
            vllm_config.model_config.hf_text_config.model_type,
        )
        self.device_buffer = torch.zeros(
            (max_num_batched_tokens, *self.shape_per_token),
            dtype=torch.int32,
            device=current_platform.device_type,
        )
        self.dp_rank = vllm_config.parallel_config.data_parallel_rank
        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]}"
            )

        self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids

    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)

capture(layer_id, topk_ids)

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.

Parameters:

  • layer_id

    (int) –

    The layer index.

  • topk_ids

    (Tensor) –

    Tensor of shape (batch_size, num_routed_experts).

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
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]}"
        )

    self.device_buffer[: len(local_topk_ids), layer_id] = local_topk_ids

snapshot_routing_data(num_tokens)

Return a stable snapshot of the current routing data.

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
185
186
187
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)

RoutedExpertsSink

Layer-owned buffer and callback for the expert ids a monolithic kernel routes to; it outlives kernel rebuilds on weight reload.

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
class RoutedExpertsSink:
    """Layer-owned buffer and callback for the expert ids a monolithic kernel
    routes to; it outlives kernel rebuilds on weight reload."""

    def __init__(
        self, moe_config: FusedMoEConfig, capture_fn: Callable[[torch.Tensor], None]
    ) -> None:
        # Sized for per-rank batches gathered across the DP or EP group.
        group_size = moe_config.ep_size if moe_config.use_ep else moe_config.dp_size
        self.buffer = torch.empty(
            (moe_config.max_num_tokens * group_size, moe_config.experts_per_token),
            dtype=torch.int16,
            device=moe_config.device,
        )
        self.capture_fn = capture_fn

bind_routed_experts_capturer(model, capturer)

Attach capture callbacks to the target model's MoE routers.

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
def bind_routed_experts_capturer(
    model: torch.nn.Module,
    capturer: RoutedExpertsCapturer,
) -> None:
    """Attach capture callbacks to the target model's MoE routers."""
    from vllm.model_executor.layers.fused_moe.layer import MoERunner
    from vllm.model_executor.layers.fused_moe.modular_kernel import (
        FusedMoEExpertsMonolithic,
    )
    from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter

    num_bound = 0
    for module in model.modules():
        if isinstance(module, RoutedExpertsCaptureSource):
            module.capture_fn = partial(capturer.capture, module.layer_id)
            num_bound += 1
            continue
        if not isinstance(module, MoERunner):
            continue
        capture_fn = partial(capturer.capture, module.layer_id)
        quant_method = module._quant_method
        if quant_method.is_monolithic:
            moe_kernel = getattr(quant_method, "moe_kernel", None)
            impl = getattr(moe_kernel, "impl", None)
            fused_experts = getattr(impl, "fused_experts", None)
            if not (
                isinstance(fused_experts, FusedMoEExpertsMonolithic)
                and fused_experts.supports_routing_replay_capture()
            ):
                raise ValueError(
                    "Routed-experts capture is not supported with monolithic "
                    f"MoE kernel {type(fused_experts).__name__}."
                )
            module.routed_experts.routing_sink = RoutedExpertsSink(
                fused_experts.moe_config, capture_fn
            )
            num_bound += 1
        elif isinstance(module.router, BaseRouter):
            module.router.set_capture_fn(capture_fn)
            num_bound += 1
        else:
            raise ValueError(
                "Routed-experts capture is not supported with router "
                f"{type(module.router).__name__}."
            )

    if num_bound == 0:
        raise ValueError("No supported MoE router found for routed-experts capture.")