| 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 |
| 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 |
| 24 | |
| 25 | |
| 26 | class _FakeMoERunner: |
| 27 | """Minimal MoERunner stand-in for token-count discovery tests.""" |
| 28 | |
| 29 | moe_config: Any |
| 30 | |
| 31 | |
| 32 | def _make_moe(*, max_deferred_tokens: int = 128, enabled: bool = True): |
| 33 | """Create a fake MoE layer with a deferred-finalize token limit.""" |
| 34 | moe = _FakeMoERunner() |
| 35 | moe.moe_config = SimpleNamespace( |
| 36 | use_deferred_moe_finalize=enabled, |
| 37 | defer_moe_finalize_max_num_tokens=max_deferred_tokens, |
| 38 | ) |
| 39 | return moe |
| 40 | |
| 41 | |
| 42 | def _make_runner(modules, *, max_tokens: int = 8192, linear_backend: str = "auto"): |
| 43 | """Create a runner carrying only the state used by FlashInfer warmup.""" |
| 44 | return SimpleNamespace( |
| 45 | scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens), |
| 46 | vllm_config=SimpleNamespace( |
| 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)) |
| 52 | ), |
| 53 | _dummy_run=Mock(), |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | def test_flashinfer_autotune_token_counts_include_deferred_moe_limits(): |
| 58 | runner = _make_runner( |
| 59 | [ |
| 60 | object(), |
| 61 | _make_moe(max_deferred_tokens=128), |
| 62 | _make_moe(max_deferred_tokens=128), |
| 63 | _make_moe(max_deferred_tokens=64, enabled=False), |
| 64 | _make_moe(max_deferred_tokens=-1), |
| 65 | ], |
| 66 | linear_backend="flashinfer_cutedsl", |
| 67 | ) |
| 68 | |
| 69 | with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): |
| 70 | token_counts = _flashinfer_autotune_token_counts(runner) |
| 71 | |
| 72 | assert token_counts == (8192, 32, 128) |
| 73 | |
| 74 | |
| 75 | def test_flashinfer_autotune_token_counts_are_bounded_and_deduplicated(): |
| 76 | runner = _make_runner( |
| 77 | [_make_moe(max_deferred_tokens=4096)], |
| 78 | max_tokens=32, |
| 79 | linear_backend="flashinfer_cutedsl", |
| 80 | ) |
| 81 | |
| 82 | with patch("vllm.model_executor.layers.fused_moe.MoERunner", _FakeMoERunner): |
| 83 | token_counts = _flashinfer_autotune_token_counts(runner) |
| 84 | |
| 85 | assert token_counts == (32,) |
| 86 | |
| 87 | |
| 88 | @pytest.mark.parametrize("skip_attn", [False, True]) |
| 89 | def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): |
| 90 | from vllm.v1.worker.gpu.model_runner import GPUModelRunner as V2Runner |
| 91 | from vllm.v1.worker.gpu_model_runner import GPUModelRunner as V1Runner |
| 92 | |
| 93 | runner = _make_runner([]) |
| 94 | dummy_run_signature = signature((V2Runner if skip_attn else V1Runner)._dummy_run) |
| 95 | runner._dummy_run.side_effect = lambda **kwargs: dummy_run_signature.bind( |
| 96 | None, **kwargs |
| 97 | ) |
| 98 | max_buckets = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 8192) |
| 99 | deferred_buckets = (1, 2, 4, 8, 16, 32, 64, 128) |
| 100 | |
| 101 | with ( |
| 102 | patch( |
| 103 | "vllm.model_executor.warmup.kernel_warmup." |
| 104 | "_flashinfer_autotune_token_counts", |
| 105 | return_value=(8192, 128), |
| 106 | ), |
| 107 | patch( |
| 108 | "vllm.utils.flashinfer.flashinfer_get_hybrid_num_tokens_buckets", |
| 109 | side_effect=(max_buckets, deferred_buckets), |
| 110 | ) as get_buckets, |
| 111 | patch("vllm.utils.flashinfer.autotune") as autotune, |
| 112 | ): |
| 113 | _run_flashinfer_autotune_dummy_runs(runner, skip_attn=skip_attn) |
| 114 | |
| 115 | assert get_buckets.call_args_list == [call(8192), call(128)] |
| 116 | assert autotune.call_args_list == [ |
| 117 | call(tuning_buckets=max_buckets), |
| 118 | call(tuning_buckets=deferred_buckets), |
| 119 | ] |
| 120 | assert runner._dummy_run.call_args_list == [ |
| 121 | call( |
| 122 | num_tokens=8192, |
| 123 | skip_eplb=True, |
| 124 | is_profile=True, |
| 125 | randomize_inputs=True, |
| 126 | **({"skip_attn": True} if skip_attn else {}), |
| 127 | ), |
| 128 | call( |
| 129 | num_tokens=128, |
| 130 | skip_eplb=True, |
| 131 | is_profile=True, |
| 132 | randomize_inputs=True, |
| 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()) |