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 
 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
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:
        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)

_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 
 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
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 
156
157
158
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 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
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 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
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
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.")