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:

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
 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
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
        self.restore_rows: Callable[[torch.Tensor, ForwardContext], torch.Tensor] = (
            self._restore_local_topk_ids
        )

    def capture(self, layer_id: int, topk_ids: torch.Tensor) -> None:
        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]}"
            )

        ctx = get_forward_context()
        local_topk_ids = self.restore_rows(topk_ids, ctx)
        offset = dbo_current_ubatch_token_offset()
        self.device_buffer[offset : offset + 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)

_restore_local_topk_ids(topk_ids, ctx)

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.

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

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 
179
180
181
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)

RoutedExpertsRowRestorer

Bases: Protocol

Methods:

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
@runtime_checkable
class RoutedExpertsRowRestorer(Protocol):
    def restore_routed_experts_rows(
        self, topk_ids: torch.Tensor, ctx: ForwardContext
    ) -> torch.Tensor:
        """Restore local token-ordered logical IDs for all captured MoE layers.

        Binding selects the first provider. Read current metadata from ``ctx``;
        do not depend on the provider's layer-specific state or modify input.
        Return a 2-D tensor with the input device, dtype and top-k width.
        Collectives must be rank-consistent and graph-replay safe.
        """
        ...

restore_routed_experts_rows(topk_ids, ctx)

Restore local token-ordered logical IDs for all captured MoE layers.

Binding selects the first provider. Read current metadata from ctx; do not depend on the provider's layer-specific state or modify input. Return a 2-D tensor with the input device, dtype and top-k width. Collectives must be rank-consistent and graph-replay safe.

Source code in vllm/model_executor/layers/fused_moe/routed_experts_capturer.py 
33
34
35
36
37
38
39
40
41
42
43
def restore_routed_experts_rows(
    self, topk_ids: torch.Tensor, ctx: ForwardContext
) -> torch.Tensor:
    """Restore local token-ordered logical IDs for all captured MoE layers.

    Binding selects the first provider. Read current metadata from ``ctx``;
    do not depend on the provider's layer-specific state or modify input.
    Return a 2-D tensor with the input device, dtype and top-k width.
    Collectives must be rank-consistent and graph-replay safe.
    """
    ...

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 
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
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 
201
202
203
204
205
206
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
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
    restorer_bound = False
    for module in model.modules():
        if not isinstance(module, (MoERunner, RoutedExpertsCaptureSource)):
            continue
        if not restorer_bound and isinstance(module, RoutedExpertsRowRestorer):
            capturer.restore_rows = module.restore_routed_experts_rows
            restorer_bound = True
        if isinstance(module, RoutedExpertsCaptureSource):
            module.capture_fn = partial(capturer.capture, module.layer_id)
            num_bound += 1
            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.")