[Bugfix][PP][Spec Decode] MiniMax-M3 EAGLE3 aux-state relay at PP > 1 and per-stage FlashInfer autotune (#57197)
Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Venky committed
Sep 30, 2026 at 14:14 UTC
42f0c17ea755a70656f92d278ba086e31fc72086
4 files changed
+382
-15
tests/model_executor/test_flashinfer_autotune_warmup.py
+190
-2
@@ -1,16 +1,23 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import json
5
+import sys
6
+from collections import defaultdict
7
+from contextlib import nullcontext
8
from inspect import signature
5
-from types import SimpleNamespace
9
+from pathlib import Path
10
+from types import ModuleType, SimpleNamespace
11
from typing import Any
12
from unittest.mock import Mock, call, patch
13
14
import pytest
15
16
+from vllm.model_executor.warmup import kernel_warmup as warmup
17
from vllm.model_executor.warmup.kernel_warmup import (
18
_flashinfer_autotune_token_counts,
19
_run_flashinfer_autotune_dummy_runs,
20
+ flashinfer_autotune,
21
)
22
23
pytestmark = pytest.mark.cpu_test
@@ -37,7 +44,8 @@ def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto
44
return SimpleNamespace(
45
scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens),
46
vllm_config=SimpleNamespace(
40
- kernel_config=SimpleNamespace(linear_backend=linear_backend)
47
+ kernel_config=SimpleNamespace(linear_backend=linear_backend),
48
+ attention_config=SimpleNamespace(hisparse_config=None),
49
),
50
get_model=Mock(
51
return_value=SimpleNamespace(modules=Mock(return_value=modules))
@@ -125,3 +133,183 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn):
133
**({"skip_attn": True} if skip_attn else {}),
134
),
135
]
136
+
137
+
138
+class _AutotuneGroup:
139
+ def __init__(self, run, ranks):
140
+ self.run = run
141
+ self.ranks = tuple(ranks)
142
+ self.world_size = len(self.ranks)
143
+ self.rank_in_group = self.ranks.index(run.rank)
144
+ self.cpu_group = self
145
+
146
+ def record(self, operation):
147
+ self.run.collectives[self.ranks][self.run.rank].append(operation)
148
+
149
+ def broadcast_object(self, obj, src=0):
150
+ assert src == 0
151
+ self.record(("broadcast", src))
152
+ if self.rank_in_group == src:
153
+ self.run.broadcasts[self.ranks] = obj
154
+ return self.run.broadcasts[self.ranks]
155
+
156
+ def barrier(self):
157
+ self.record(("barrier",))
158
+
159
+
160
+class _AutotuneTuner:
161
+ def __init__(self, run):
162
+ self.run = run
163
+ self.cache = {}
164
+ self.loaded = None
165
+
166
+ def load_configs(self, path):
167
+ self.loaded = json.loads(Path(path).read_text())
168
+ self.cache.update(self.loaded)
169
+
170
+ def save_configs(self, path):
171
+ self.run.saves.append((self.run.rank, Path(path), dict(self.cache)))
172
+ Path(path).write_text(json.dumps(self.cache))
173
+
174
+ def profile(self, operation):
175
+ if operation in self.cache:
176
+ return
177
+ group = self.run.tuning_group
178
+ self.run.profile_groups[self.run.rank].append(
179
+ None if group is None else group.ranks
180
+ )
181
+ for tactic in range(2):
182
+ if group is not None:
183
+ group.record(("all_reduce", operation, tactic))
184
+ self.cache[operation] = self.run.rank // self.run.tp
185
+
186
+
187
+class _AutotuneRun:
188
+ def __init__(self, pp, tp):
189
+ self.pp, self.tp = pp, tp
190
+ self.rank = 0
191
+ self.tuning_group = None
192
+ self.collectives: dict[tuple[int, ...], dict[int, list[tuple[Any, ...]]]] = (
193
+ defaultdict(lambda: defaultdict(list))
194
+ )
195
+ self.broadcasts = {}
196
+ self.tuners = {}
197
+ self.saves = []
198
+ self.profile_groups = defaultdict(list)
199
+
200
+ def world(self):
201
+ return _AutotuneGroup(self, range(self.pp * self.tp))
202
+
203
+ def tensor_group(self):
204
+ start = self.rank // self.tp * self.tp
205
+ return _AutotuneGroup(self, range(start, start + self.tp))
206
+
207
+ def pipeline_group(self):
208
+ return SimpleNamespace(world_size=self.pp)
209
+
210
+ def set_group(self, group):
211
+ self.tuning_group = group
212
+
213
+ def dummy_runs(self, runner, *, skip_attn=False):
214
+ self.tuners[self.rank].profile("shared_gemm")
215
+ if self.pp > 1 and self.rank // self.tp == 0:
216
+ self.tuners[self.rank].profile("pp0_extra_gemm")
217
+
218
+ def execute(self):
219
+ for rank in range(self.pp * self.tp):
220
+ self.rank = rank
221
+ self.tuners[rank] = _AutotuneTuner(self)
222
+ flashinfer_autotune(_make_runner([]))
223
+ return self
224
+
225
+ def assert_collectives_match(self):
226
+ for ranks, traces in self.collectives.items():
227
+ assert set(traces) == set(ranks)
228
+ expected = traces[ranks[0]]
229
+ assert all(trace == expected for trace in traces.values()), dict(traces)
230
+
231
+
232
+@pytest.fixture
233
+def autotune_run(monkeypatch, tmp_path):
234
+ import vllm.utils.flashinfer as fi_utils
235
+ from vllm.distributed import parallel_state
236
+
237
+ def make_run(*, pp=2, tp=4):
238
+ run = _AutotuneRun(pp, tp)
239
+ autotuner = ModuleType("flashinfer.autotuner")
240
+ monkeypatch.setattr(
241
+ autotuner,
242
+ "AutoTuner",
243
+ SimpleNamespace(get=lambda: run.tuners[run.rank]),
244
+ raising=False,
245
+ )
246
+ monkeypatch.setattr(
247
+ autotuner, "set_autotune_process_group", run.set_group, raising=False
248
+ )
249
+ monkeypatch.setitem(sys.modules, "flashinfer.autotuner", autotuner)
250
+ monkeypatch.setattr(parallel_state, "get_world_group", run.world)
251
+ monkeypatch.setattr(parallel_state, "get_tp_group", run.tensor_group)
252
+ monkeypatch.setattr(parallel_state, "get_pp_group", run.pipeline_group)
253
+ monkeypatch.setattr(fi_utils, "autotune", lambda **kwargs: nullcontext())
254
+ monkeypatch.setattr(
255
+ warmup,
256
+ "resolve_flashinfer_autotune_file",
257
+ lambda runner: tmp_path / "autotune_configs.json",
258
+ )
259
+ monkeypatch.setattr(
260
+ warmup, "_flashinfer_autotune_skip_ops", lambda runner: None
261
+ )
262
+ monkeypatch.setattr(
263
+ warmup, "_run_flashinfer_autotune_dummy_runs", run.dummy_runs
264
+ )
265
+ monkeypatch.setattr(warmup, "replayssm_autotune_warmup", lambda runner: None)
266
+ monkeypatch.setattr(warmup, "_autotune_kimi_k3_kda_qkvg", lambda model: None)
267
+ return run
268
+
269
+ return make_run
270
+
271
+
272
+@pytest.mark.parametrize("tp", [1, 4])
273
+def test_heterogeneous_pp_stages_have_compatible_collectives(autotune_run, tp):
274
+ run = autotune_run(tp=tp).execute()
275
+ run.assert_collectives_match()
276
+ for rank, groups in run.profile_groups.items():
277
+ expected = tuple(range(rank // tp * tp, (rank // tp + 1) * tp))
278
+ assert groups and all(
279
+ group == (expected if tp > 1 else None) for group in groups
280
+ )
281
+
282
+
283
+def test_pp_stage_cache_roundtrip_isolated_and_asymmetric_hits_safe(
284
+ autotune_run, tmp_path
285
+):
286
+ legacy = tmp_path / "autotune_configs.json"
287
+ legacy.write_text('{"legacy_world_cache": 99}')
288
+ cold = autotune_run().execute()
289
+ assert [rank for rank, _, _ in cold.saves] == [0, 4]
290
+ paths = [path for _, path, _ in cold.saves]
291
+ assert len(set(paths)) == 2 and legacy not in paths
292
+ assert json.loads(legacy.read_text()) == {"legacy_world_cache": 99}
293
+ assert cold.saves[0][2] == {"shared_gemm": 0, "pp0_extra_gemm": 0}
294
+ assert cold.saves[1][2] == {"shared_gemm": 1}
295
+ cold.assert_collectives_match()
296
+ warm = autotune_run().execute()
297
+ warm.assert_collectives_match()
298
+ assert not warm.profile_groups
299
+ for rank, tuner in warm.tuners.items():
300
+ assert tuner.loaded == cold.saves[rank // 4][2]
301
+ paths[1].unlink()
302
+ mixed = autotune_run().execute()
303
+ mixed.assert_collectives_match()
304
+ assert set(mixed.profile_groups) == {4, 5, 6, 7}
305
+ assert all(mixed.tuners[rank].loaded is not None for rank in range(4))
306
+ assert all(mixed.tuners[rank].loaded is None for rank in range(4, 8))
307
+
308
+
309
+def test_pp1_retains_world_synchronization_and_existing_cache_name(autotune_run):
310
+ run = autotune_run(pp=1, tp=4).execute()
311
+ run.assert_collectives_match()
312
+ assert [(rank, path.name) for rank, path, _ in run.saves] == [
313
+ (0, "autotune_configs.json")
314
+ ]
315
+ assert all(groups == [(0, 1, 2, 3)] for groups in run.profile_groups.values())
tests/v1/worker/test_eagle3_aux_hidden_states_pp.py
+152
@@ -3,10 +3,15 @@
3
from types import SimpleNamespace
4
5
import pytest
6
+import torch
7
8
+from vllm.distributed import parallel_state
9
from vllm.model_executor.models.interfaces import EagleModelMixin
10
from vllm.model_executor.models.mimo import MiMoModel
11
+from vllm.models.minimax_m3.nvidia import model as minimax_m3
12
+from vllm.sequence import IntermediateTensors
13
from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import (
14
+ reserve_aux_intermediate_tensor_slots,
15
verify_supports_aux_hidden_states_over_pp,
16
)
17
@@ -24,3 +29,150 @@ def test_mimo_does_not_inherit_aux_hidden_state_pp_support():
29
assert not inner.supports_aux_hidden_states_over_pp
30
with pytest.raises(ValueError, match="does not support eagle3"):
31
verify_supports_aux_hidden_states_over_pp(target, "eagle3")
32
+
33
+
34
+class _PPAuxLayer(torch.nn.Module):
35
+ """Aliases and mutates residual in place like MiniMaxM3DecoderLayer."""
36
+
37
+ ffn_all_reduce_deferred = False
38
+
39
+ def __init__(self, index):
40
+ super().__init__()
41
+ self.index = index
42
+
43
+ def forward(self, positions, hidden_states, residual):
44
+ total = hidden_states if residual is None else hidden_states + residual
45
+ residual = hidden_states if residual is None else residual
46
+ residual.copy_(total + self.index + 1)
47
+ return torch.full_like(hidden_states, 10 * (self.index + 1)), residual
48
+
49
+
50
+class _PPAuxNorm(torch.nn.Module):
51
+ """Clobbers its inputs so aux taps that alias them would fail."""
52
+
53
+ def forward(self, hidden_states, residual):
54
+ output = hidden_states + residual
55
+ hidden_states.fill_(-101)
56
+ residual.fill_(-103)
57
+ return output, residual
58
+
59
+
60
+@pytest.fixture
61
+def minimax_pp_stage(monkeypatch):
62
+ """Use the real MiniMax constructor/forward and aux helpers on CPU tensors."""
63
+ group = SimpleNamespace(world_size=2, is_first_rank=True, is_last_rank=False)
64
+ bounds = [0, 3]
65
+ monkeypatch.setattr(minimax_m3, "get_pp_group", lambda: group)
66
+ monkeypatch.setattr(parallel_state, "get_pp_group", lambda: group)
67
+ monkeypatch.setattr(parallel_state, "model_parallel_is_initialized", lambda: True)
68
+ monkeypatch.setattr(
69
+ minimax_m3,
70
+ "VocabParallelEmbedding",
71
+ lambda *args, **kwargs: torch.nn.Identity(),
72
+ )
73
+ monkeypatch.setattr(
74
+ minimax_m3, "MiniMAXGemmaRMSNorm", lambda *args, **kwargs: _PPAuxNorm()
75
+ )
76
+ monkeypatch.setattr(
77
+ minimax_m3,
78
+ "MiniMaxM3DecoderLayer",
79
+ lambda **kwargs: _PPAuxLayer(int(kwargs["prefix"].rsplit(".", 1)[1])),
80
+ )
81
+
82
+ def make_layers(num_layers, build, prefix):
83
+ return (
84
+ *bounds,
85
+ torch.nn.ModuleList(
86
+ build(f"{prefix}.{index}")
87
+ if bounds[0] <= index < bounds[1]
88
+ else torch.nn.Identity()
89
+ for index in range(num_layers)
90
+ ),
91
+ )
92
+
93
+ monkeypatch.setattr(minimax_m3, "make_layers", make_layers)
94
+
95
+ def make_stage(start, end, taps, world_size=2):
96
+ bounds[:] = [start, end]
97
+ group.world_size = world_size
98
+ group.is_first_rank = start == 0
99
+ group.is_last_rank = end == 6
100
+ config = SimpleNamespace(
101
+ vocab_size=16, hidden_size=6144, num_hidden_layers=6, rms_norm_eps=1e-6
102
+ )
103
+ model = minimax_m3.MiniMaxM3Model(
104
+ vllm_config=SimpleNamespace(
105
+ model_config=SimpleNamespace(hf_text_config=config), quant_config=None
106
+ )
107
+ )
108
+ model._set_aux_hidden_state_layers(taps)
109
+ target = SimpleNamespace(
110
+ model=model,
111
+ make_empty_intermediate_tensors=model.make_empty_intermediate_tensors,
112
+ )
113
+ reserve_aux_intermediate_tensor_slots(target)
114
+ return model, target
115
+
116
+ return make_stage
117
+
118
+
119
+# Input 7; _PPAuxLayer i adds 11 * (i + 1) to hidden + residual.
120
+def _expected_aux_value(tap):
121
+ return 7 + 11 * tap * (tap + 1) // 2
122
+
123
+
124
+def _pp_aux_state(value):
125
+ return torch.full((2, 6144), float(value), device="cpu")
126
+
127
+
128
+@pytest.mark.parametrize("taps", [(0, 3, 6), (0, 1, 3), (4, 5, 6)])
129
+def test_minimax_pp2_preserves_ordered_global_aux_taps(minimax_pp_stage, taps):
130
+ first, _ = minimax_pp_stage(0, 3, taps)
131
+ sent = first(None, None, None, _pp_aux_state(7))
132
+ last, target = minimax_pp_stage(3, 6, taps)
133
+ reserved = target.make_empty_intermediate_tensors(2, torch.float32, "cpu")
134
+ assert set(sent.tensors) == set(reserved.tensors)
135
+ result = last(None, None, sent)
136
+ assert isinstance(result, tuple)
137
+ output, aux = result
138
+ torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6)))
139
+ for hidden, tap in zip(aux, taps, strict=True):
140
+ torch.testing.assert_close(hidden, _pp_aux_state(_expected_aux_value(tap)))
141
+
142
+
143
+def test_minimax_missing_upstream_aux_slot_is_not_silently_dropped(minimax_pp_stage):
144
+ model, _ = minimax_pp_stage(3, 6, (0, 3, 6))
145
+ incoming = IntermediateTensors(
146
+ {
147
+ "hidden_states": _pp_aux_state(30),
148
+ "residual": _pp_aux_state(43),
149
+ "aux_hidden_states_0": _pp_aux_state(7),
150
+ }
151
+ )
152
+ with pytest.raises(RuntimeError, match="Missing aux_hidden_states_1"):
153
+ model(None, None, incoming)
154
+
155
+
156
+@pytest.mark.parametrize("taps", [(), (0, 3, 6)])
157
+def test_minimax_pp1_output_and_aux_taps(minimax_pp_stage, taps):
158
+ model, _ = minimax_pp_stage(0, 6, taps, world_size=1)
159
+ output = model(None, None, None, _pp_aux_state(7))
160
+ if taps:
161
+ output, aux = output
162
+ for hidden, tap in zip(aux, taps, strict=True):
163
+ torch.testing.assert_close(hidden, _pp_aux_state(_expected_aux_value(tap)))
164
+ torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6)))
165
+
166
+
167
+def test_minimax_pp2_without_aux_keeps_hidden_and_residual_transport(minimax_pp_stage):
168
+ first, _ = minimax_pp_stage(0, 3, ())
169
+ sent = first(None, None, None, _pp_aux_state(7))
170
+ assert set(sent.tensors) == {"hidden_states", "residual"}
171
+ torch.testing.assert_close(sent["hidden_states"], _pp_aux_state(30))
172
+ torch.testing.assert_close(sent["residual"], _pp_aux_state(43))
173
+ last, target = minimax_pp_stage(3, 6, ())
174
+ assert set(
175
+ target.make_empty_intermediate_tensors(12, torch.bfloat16, "cpu").tensors
176
+ ) == {"hidden_states", "residual"}
177
+ output = last(None, None, sent)
178
+ torch.testing.assert_close(output, _pp_aux_state(_expected_aux_value(6)))
vllm/model_executor/warmup/kernel_warmup.py
+21
-9
@@ -410,17 +410,24 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
410
Without autotuning, FlashInfer will rely on heuristics, which may
411
be significantly slower.
412
413
- Every rank profiles the same tactics. When distributed, per-tactic
414
- timings are averaged over the world CPU group so all ranks select the
415
- same tactic.
413
+ With PP > 1, stages run different layers and may profile different ops,
414
+ so each stage's TP group tunes separately with its own cache file;
415
+ otherwise the world group tunes together. Per-tactic timings are
416
+ averaged over the tuning group so all its ranks select the same tactic.
417
"""
418
from flashinfer.autotuner import AutoTuner, set_autotune_process_group
419
420
import vllm.utils.flashinfer as fi_utils
420
- from vllm.distributed.parallel_state import get_world_group
421
+ from vllm.distributed.parallel_state import (
422
+ get_pp_group,
423
+ get_tp_group,
424
+ get_world_group,
425
+ )
426
427
world = get_world_group()
423
- is_leader = world.rank_in_group == 0
428
+ pp_size = get_pp_group().world_size
429
+ tune_group = get_tp_group() if pp_size > 1 else world
430
+ is_leader = tune_group.rank_in_group == 0
431
tuner = AutoTuner.get()
432
433
autotune_kwargs: dict = {}
@@ -433,6 +440,11 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
440
autotune_kwargs["skip_ops"] = skip_ops
441
442
cache_path = resolve_flashinfer_autotune_file(runner)
443
+ if pp_size > 1:
444
+ ranks = "-".join(str(rank) for rank in tune_group.ranks)
445
+ cache_path = cache_path.with_name(
446
+ f"{cache_path.stem}_tp_{ranks}{cache_path.suffix}"
447
+ )
448
if is_leader:
449
logger.info_once("Using FlashInfer autotune cache file: %s", cache_path)
450
@@ -441,18 +453,18 @@ def flashinfer_autotune(runner: "GPUModelRunner") -> None:
453
# which lead to some EP ranks receiving no tokens and skipping their
454
# MoE kernel entirely, and cause hang due to all-reduce collective
455
# during synchronized autotuning.
444
- # Read cached autotune results and broadcast to all ranks.
456
+ # Read cached autotune results and broadcast within the tuning group.
457
cached_results: bytes | None = None
458
if is_leader and cache_path.exists():
459
with open(cache_path, "rb") as f:
460
cached_results = f.read()
449
- cached_results = world.broadcast_object(cached_results, src=0)
461
+ cached_results = tune_group.broadcast_object(cached_results, src=0)
462
if cached_results is not None:
463
write_flashinfer_autotune_cache(cache_path, cached_results)
452
- world.barrier()
464
+ tune_group.barrier()
465
tuner.load_configs(str(cache_path))
466
455
- group = world.cpu_group if world.world_size > 1 else None
467
+ group = tune_group.cpu_group if tune_group.world_size > 1 else None
468
set_autotune_process_group(group)
469
try:
470
with (
vllm/models/minimax_m3/nvidia/model.py
+19
-4
@@ -794,6 +794,7 @@ class MiniMaxM3DecoderLayer(nn.Module):
794
795
class MiniMaxM3Model(nn.Module, EagleModelMixin):
796
fall_back_to_pt_during_load = False
797
+ supports_aux_hidden_states_over_pp = True
798
799
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
800
super().__init__()
@@ -886,9 +887,18 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
887
hidden_states = intermediate_tensors["hidden_states"]
888
residual = intermediate_tensors["residual"]
889
889
- # EAGLE3 is not yet compatible with pipeline parallel
890
- aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)
891
- for idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]):
890
+ remote_aux = self.collect_remote_aux_hidden_states(intermediate_tensors)
891
+ aux_hidden_states: list[torch.Tensor] = []
892
+ if get_pp_group().is_first_rank:
893
+ self._maybe_add_hidden_state(
894
+ aux_hidden_states, self.start_layer, hidden_states, residual
895
+ )
896
+ # Preserve the entry tap before in-place residual updates.
897
+ if aux_hidden_states:
898
+ aux_hidden_states[0] = aux_hidden_states[0].clone()
899
+ for idx, layer in enumerate(
900
+ self.layers[self.start_layer : self.end_layer], start=self.start_layer
901
+ ):
902
hidden_states, residual = layer(positions, hidden_states, residual)
903
self._maybe_add_hidden_state(
904
aux_hidden_states, idx + 1, hidden_states, residual
@@ -896,7 +906,11 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
906
907
if not get_pp_group().is_last_rank:
908
return IntermediateTensors(
899
- {"hidden_states": hidden_states, "residual": residual}
909
+ {
910
+ "hidden_states": hidden_states,
911
+ "residual": residual,
912
+ **self.pack_local_aux_hidden_states(aux_hidden_states),
913
+ }
914
)
915
916
if self.fuse_final_norm_allreduce:
@@ -906,6 +920,7 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):
920
else:
921
hidden_states, _ = self.norm(hidden_states, residual)
922
923
+ aux_hidden_states = remote_aux + aux_hidden_states
924
if len(aux_hidden_states) > 0:
925
return hidden_states, aux_hidden_states
926
return hidden_states