Skip to content

vllm.models.minimax_m3.amd.specialized_decode.runtime.executor

FFN-launch executor: vLLM runs each sparse layer's input norm and attention, then one fused launch runs the rest of the layer.

Classes:

  • FfnExecutor –

    Collective construction: every TP rank builds it at the same point, and

FfnExecutor

Collective construction: every TP rank builds it at the same point, and every phase that can fail on one rank ends in a TP agreement.

Methods:

  • abort –

    Mark this rank failed on every rank, so peers' waits end.

  • begin_step –

    This step's slot mapping (-1: padding), read by every launch on the

  • output_step –

    Every rank, every sampled step (eager or replayed): the epoch reset

  • reset_epoch –

    Between steps, on every rank together: no launch in flight, every

Source code in vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.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
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
class FfnExecutor:
    """Collective construction: every TP rank builds it at the same point, and
    every phase that can fail on one rank ends in a TP agreement."""

    uses_caches = False

    def __init__(
        self,
        layers,
        weights: dict[int, MonoLayerWeights],
        widths: tuple[int, ...],
        tp: int,
        reset_launches: int = RESET_LAUNCHES,
    ) -> None:
        from vllm.models.deepseek_v41.amd.mono.common.peer_memory import PeerBuffer

        group = get_tp_group()
        self.rank = group.rank_in_group
        device = torch.device("cuda", torch.accelerator.current_device_index())
        self._layers = {layer.layer_id: layer for layer in layers}
        self._weights = weights
        self.peer = PeerBuffer(peer_bytes(tp), group.cpu_group, self.rank, tp, device)
        failure: Exception | None = None
        try:
            self.peer.bytes.zero_()
            self.runner = FfnRunner(
                tp, self.rank, self.peer.local, self.peer.addresses, device, widths
            )
            cus = current_platform.num_compute_units(device.index)
            self.runner.warm(next(iter(weights.values())), compute_units=cus)
            torch.accelerator.synchronize()
        except Exception as err:
            failure = err
        if not tp_agree(failure is None):
            self.peer.close()
            raise RuntimeError(
                "MiniMax-M3 FFN launch setup failed on "
                f"{'this' if failure else 'another'} TP rank"
            ) from failure
        logger.info_once(
            "MiniMax-M3 FFN launch: rows %s, resident CTAs per CU %s",
            widths,
            tuple(sorted(set(self.runner.residency.values()))),
        )
        self.fault_word = self.peer.bytes[: 4 * ERROR_WORDS].view(torch.int32)
        self.launches = dict.fromkeys(self._layers, 0)
        self.steps = 0
        self.outputs = 0
        self.resets = 0
        self._since_reset = 0
        self._reset_launches = reset_launches
        self._epoch_host = torch.zeros(1, dtype=torch.int32, pin_memory=True)
        self._epoch_event: torch.cuda.Event | None = None
        self._slots: torch.Tensor | None = None
        self._kernels = {
            layer_id: self._layer_kernel(layer_id) for layer_id in self._layers
        }

    def _layer_kernel(self, layer_id: int):
        layer = self._layers[layer_id]
        weights = self._weights[layer_id]
        runner = self.runner
        launches = self.launches

        final = layer_id == max(self._layers)

        def run(positions, hidden_states, residual):
            hidden_states, residual = layer.input_layernorm(hidden_states, residual)
            part = layer.self_attn(positions=positions, hidden_states=hidden_states)
            launches[layer_id] += 1
            assert self._slots is not None
            return runner(part, residual, weights, self._slots, final)

        return run

    def begin_step(self, slots: torch.Tensor) -> None:
        """This step's slot mapping (-1: padding), read by every launch on the
        device; under a FULL graph, the runner's persistent buffer."""
        self._slots = slots

    def layer(self, layer_id: int):
        return self._kernels.get(layer_id)

    def bind(self, layer_id, binding) -> None:
        raise NotImplementedError("the FFN launch reads no KV cache")

    def health(self) -> MonoHealth:
        words = self.fault_word.tolist()
        return MonoHealth(bool(words[0]), *words[1:])

    def abort(self) -> None:
        """Mark this rank failed on every rank, so peers' waits end."""
        try:
            words = torch.tensor(
                [1, self.rank, -1, 0], dtype=torch.int32, device=self.peer.bytes.device
            )
            for base in self.peer.addresses.tolist():
                view = _device_words(base, ERROR_BYTES // 4, self.peer.bytes.device)
                view[:ERROR_WORDS].copy_(words)
            torch.accelerator.synchronize()
        except Exception:
            logger.exception("MiniMax-M3 FFN launch: marking peers failed")

    def after_step(self) -> None:
        self.steps += 1

    def output_step(self) -> None:
        """Every rank, every sampled step (eager or replayed): the epoch reset
        when the launches since the last one could near the tag limit, and
        launch coverage now and then, read without a host sync."""
        self.outputs += 1
        self._since_reset += 1
        if self._since_reset * len(self._layers) >= self._reset_launches:
            self.reset_epoch()
        if self.outputs % COVERAGE_LOG_EVERY:
            return
        done = self._epoch_event is not None and self._epoch_event.query()
        if done and self.rank == 0:
            logger.info(
                "MiniMax-M3 FFN launch: %d output steps, %d launches since "
                "reset %d (%d sparse layers each)",
                self.outputs,
                int(self._epoch_host[0]),
                self.resets,
                len(self._layers),
            )
        self._epoch_host.copy_(self.runner.epoch[:1], non_blocking=True)
        self._epoch_event = torch.cuda.Event()
        self._epoch_event.record()

    def reset_epoch(self) -> None:
        """Between steps, on every rank together: no launch in flight, every
        mailbox region and the epoch cleared, so tags restart from one. The
        sticky error word is kept. Every rank reaches each phase's agreement;
        a phase that failed anywhere fails the reset on every rank (the
        failing rank marks all ranks' error words), so no launch resumes on a
        partly cleared state."""
        self._reset_phase("drain", torch.accelerator.synchronize)
        self._reset_phase("clear", self._clear)
        self._since_reset = 0
        self.resets += 1
        logger.info("MiniMax-M3 FFN launch: epoch reset %d", self.resets)

    def _clear(self) -> None:
        self.peer.bytes[4 * ERROR_WORDS :].zero_()
        self.runner.reset()
        torch.accelerator.synchronize()

    def _reset_phase(self, phase: str, run) -> None:
        failure: Exception | None = None
        try:
            run()
        except Exception as err:
            failure = err
        try:
            agreed = tp_agree(failure is None)
        except Exception as err:
            agreed, failure = False, failure or err
        if agreed:
            return
        if failure is not None:
            self.abort()
        raise RuntimeError(
            f"MiniMax-M3 FFN launch epoch reset ({phase}) failed on "
            f"{'this' if failure else 'another'} TP rank"
        ) from failure

    def close(self) -> None:
        self.peer.close()

abort()

Mark this rank failed on every rank, so peers' waits end.

Source code in vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py 
134
135
136
137
138
139
140
141
142
143
144
145
def abort(self) -> None:
    """Mark this rank failed on every rank, so peers' waits end."""
    try:
        words = torch.tensor(
            [1, self.rank, -1, 0], dtype=torch.int32, device=self.peer.bytes.device
        )
        for base in self.peer.addresses.tolist():
            view = _device_words(base, ERROR_BYTES // 4, self.peer.bytes.device)
            view[:ERROR_WORDS].copy_(words)
        torch.accelerator.synchronize()
    except Exception:
        logger.exception("MiniMax-M3 FFN launch: marking peers failed")

begin_step(slots)

This step's slot mapping (-1: padding), read by every launch on the device; under a FULL graph, the runner's persistent buffer.

Source code in vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py 
119
120
121
122
def begin_step(self, slots: torch.Tensor) -> None:
    """This step's slot mapping (-1: padding), read by every launch on the
    device; under a FULL graph, the runner's persistent buffer."""
    self._slots = slots

output_step()

Every rank, every sampled step (eager or replayed): the epoch reset when the launches since the last one could near the tag limit, and launch coverage now and then, read without a host sync.

Source code in vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def output_step(self) -> None:
    """Every rank, every sampled step (eager or replayed): the epoch reset
    when the launches since the last one could near the tag limit, and
    launch coverage now and then, read without a host sync."""
    self.outputs += 1
    self._since_reset += 1
    if self._since_reset * len(self._layers) >= self._reset_launches:
        self.reset_epoch()
    if self.outputs % COVERAGE_LOG_EVERY:
        return
    done = self._epoch_event is not None and self._epoch_event.query()
    if done and self.rank == 0:
        logger.info(
            "MiniMax-M3 FFN launch: %d output steps, %d launches since "
            "reset %d (%d sparse layers each)",
            self.outputs,
            int(self._epoch_host[0]),
            self.resets,
            len(self._layers),
        )
    self._epoch_host.copy_(self.runner.epoch[:1], non_blocking=True)
    self._epoch_event = torch.cuda.Event()
    self._epoch_event.record()

reset_epoch()

Between steps, on every rank together: no launch in flight, every mailbox region and the epoch cleared, so tags restart from one. The sticky error word is kept. Every rank reaches each phase's agreement; a phase that failed anywhere fails the reset on every rank (the failing rank marks all ranks' error words), so no launch resumes on a partly cleared state.

Source code in vllm/models/minimax_m3/amd/specialized_decode/runtime/executor.py 
174
175
176
177
178
179
180
181
182
183
184
185
def reset_epoch(self) -> None:
    """Between steps, on every rank together: no launch in flight, every
    mailbox region and the epoch cleared, so tags restart from one. The
    sticky error word is kept. Every rank reaches each phase's agreement;
    a phase that failed anywhere fails the reset on every rank (the
    failing rank marks all ranks' error words), so no launch resumes on a
    partly cleared state."""
    self._reset_phase("drain", torch.accelerator.synchronize)
    self._reset_phase("clear", self._clear)
    self._since_reset = 0
    self.resets += 1
    logger.info("MiniMax-M3 FFN launch: epoch reset %d", self.resets)