@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9cd9f8c20b

Prune redundant tests points in `correctness_e2e/[test_sequence_parallel,test_async_tp]` (#51068)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Aug 4, 2026 at 22:18 UTC 9cd9f8c20bc34e1a3f57f8a258f5cc8751e76919
5 files changed +113 -199
.buildkite/test_areas/compile.yaml
+5 -31
@@ -2,44 +2,18 @@ group: Compile
2 depends_on:
3 - image-build
4 steps:
5 -- label: Sequence Parallel Correctness Tests (2 GPUs)
6 - key: sequence-parallel-correctness-tests-2-gpus
7 - timeout_in_minutes: 80
8 - working_dir: "/vllm-workspace/"
9 - num_devices: 2
10 - source_file_dependencies:
11 - - vllm/model_executor/layers/
12 - - vllm/compilation/
13 - - vllm/v1/worker/
14 - - vllm/v1/cudagraph_dispatcher.py
15 - - tests/compile/correctness_e2e/test_sequence_parallel.py
16 - commands:
17 - - export VLLM_TEST_CLEAN_GPU_MEMORY=1
18 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
19 -
20 -- label: Sequence Parallel Correctness Tests (2xH100)
21 - key: sequence-parallel-correctness-tests-2xh100
22 - timeout_in_minutes: 75
5 +- label: Sequence Parallel Correctness Tests (2xB200)
6 + key: sequence-parallel-correctness-tests-2xb200
7 + timeout_in_minutes: 45
8 working_dir: "/vllm-workspace/"
24 - device: h100
9 + device: b200-k8s
10 optional: true
11 num_devices: 2
12 commands:
13 - export VLLM_TEST_CLEAN_GPU_MEMORY=1
14 - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py
15
31 -- label: AsyncTP Correctness Tests (2xH100)
32 - key: asynctp-correctness-tests-2xh100
33 - timeout_in_minutes: 30
34 - working_dir: "/vllm-workspace/"
35 - device: h100
36 - optional: true
37 - num_devices: 2
38 - commands:
39 - - export VLLM_TEST_CLEAN_GPU_MEMORY=1
40 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py
41 -
42 -- label: AsyncTP Correctness Tests (B200)
16 +- label: AsyncTP Correctness Tests (2xB200)
17 key: asynctp-correctness-tests-b200
18 timeout_in_minutes: 30
19 working_dir: "/vllm-workspace/"
tests/compile/correctness_e2e/test_async_tp.py
+1 -5
@@ -29,18 +29,16 @@ NVFP4_HF_OVERRIDES = {
29 @create_new_process_for_each_test()
30 @pytest.mark.parametrize(
31 "model_id",
32 - ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"],
32 + ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Llama-3.2-1B-Instruct-FP8"],
33 )
34 @pytest.mark.parametrize("tp_size", [2])
35 @pytest.mark.parametrize("async_tp_enabled", [True])
36 @pytest.mark.parametrize("distributed_backend", ["mp"])
37 -@pytest.mark.parametrize("eager_mode", [False, True])
37 def test_async_tp_pass_correctness(
38 model_id: str,
39 tp_size: int,
40 async_tp_enabled: bool,
41 distributed_backend: str,
43 - eager_mode: bool,
42 num_gpus_available: int,
43 monkeypatch,
44 ):
@@ -64,8 +62,6 @@ def test_async_tp_pass_correctness(
62 "--max-num-seqs",
63 "8",
64 ]
67 - if eager_mode:
68 - common_args.append("--enforce-eager")
65
66 compilation_config = {
67 "mode": CompilationMode.VLLM_COMPILE,
tests/compile/correctness_e2e/test_sequence_parallel.py
+102 -144
@@ -22,7 +22,7 @@ from vllm.platforms import current_platform
22 from vllm.utils.torch_utils import is_torch_equal_or_newer
23
24 from ...models.registry import HF_EXAMPLE_MODELS, _HfExamplesInfo
25 -from ...utils import compare_two_settings, create_new_process_for_each_test
25 +from ...utils import compare_all_settings, create_new_process_for_each_test
26
27 logger = init_logger("test_sequence_parallel")
28
@@ -37,7 +37,6 @@ class ParallelSetup(NamedTuple):
37 fuse_norm_quant: bool
38 fuse_act_quant: bool
39 eager_mode: bool
40 - chunked_prefill: bool
40
41
42 class SPTestOptions(NamedTuple):
@@ -53,38 +52,6 @@ class SPTestSettings:
52 runner: RunnerOption
53 test_options: SPTestOptions
54
56 - @staticmethod
57 - def detailed(
58 - *,
59 - tp_base: int = 2,
60 - pp_base: int = 1,
61 - multi_node_only: bool = False,
62 - runner: RunnerOption = "auto",
63 - load_format: str | None = None,
64 - ):
65 - parallel_setups = []
66 - for eager_mode_val in [False, True]:
67 - for pp_multiplier in [1, 2]:
68 - for chunked_prefill_val in [False, True]:
69 - parallel_setups.append(
70 - ParallelSetup(
71 - tp_size=tp_base,
72 - pp_size=pp_multiplier * pp_base,
73 - fuse_norm_quant=False,
74 - fuse_act_quant=False,
75 - eager_mode=eager_mode_val,
76 - chunked_prefill=chunked_prefill_val,
77 - )
78 - )
79 - return SPTestSettings(
80 - parallel_setups=parallel_setups,
81 - distributed_backends=["mp", "ray"],
82 - runner=runner,
83 - test_options=SPTestOptions(
84 - multi_node_only=multi_node_only, load_format=load_format
85 - ),
86 - )
87 -
55 @staticmethod
56 def fast(
57 *,
@@ -94,20 +61,16 @@ class SPTestSettings:
61 multi_node_only: bool = False,
62 load_format: str | None = None,
63 ):
97 - parallel_setups = []
98 - for eager_mode_val in [False, True]:
99 - for pp_multiplier in [1, 2]:
100 - for chunked_prefill_val in [False, True]:
101 - parallel_setups.append(
102 - ParallelSetup(
103 - tp_size=tp_base,
104 - pp_size=pp_multiplier * pp_base,
105 - fuse_norm_quant=False,
106 - fuse_act_quant=False,
107 - eager_mode=eager_mode_val,
108 - chunked_prefill=chunked_prefill_val,
109 - )
110 - )
64 + parallel_setups = [
65 + ParallelSetup(
66 + tp_size=tp_base,
67 + pp_size=pp_multiplier * pp_base,
68 + fuse_norm_quant=False,
69 + fuse_act_quant=False,
70 + eager_mode=False,
71 + )
72 + for pp_multiplier in [1, 2]
73 + ]
74 return SPTestSettings(
75 parallel_setups=parallel_setups,
76 distributed_backends=["mp", "ray"],
@@ -135,7 +98,6 @@ class SPTestSettings:
98 fuse_norm_quant=fusion_val,
99 fuse_act_quant=fusion_val,
100 eager_mode=True,
138 - chunked_prefill=False,
101 )
102 )
103 return SPTestSettings(
@@ -161,7 +123,7 @@ class SPTestSettings:
123 )
124
125
164 -def _compare_sp(
126 +def _build_sp_args(
127 model_id: str,
128 parallel_setup: ParallelSetup,
129 distributed_backend: str,
@@ -172,17 +134,14 @@ def _compare_sp(
134 fuse_gemm_comms: bool,
135 enable_prompt_embeds: bool,
136 *,
175 - method: Literal["generate", "encode"],
137 is_multimodal: bool,
177 - dtype: str = "float16",
178 -):
138 +) -> tuple[list[str], list[str]] | None:
139 (
140 tp_size,
141 pp_size,
142 fuse_norm_quant,
143 fuse_act_quant,
144 eager_mode,
185 - chunked_prefill,
145 ) = parallel_setup
146
147 multi_node_only = test_options.multi_node_only
@@ -214,26 +173,20 @@ def _compare_sp(
173 model_info.check_available_online(on_fail="skip")
174
175 if num_gpus_available < tp_size * pp_size:
217 - pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs")
176 + return None
177 if VLLM_MULTI_NODE and distributed_backend == "mp":
219 - pytest.skip(
220 - "Skipping multi-node pipeline parallel test for "
221 - "multiprocessing distributed backend"
222 - )
178 + return None
179 if multi_node_only and not VLLM_MULTI_NODE:
224 - pytest.skip("Not in multi-node setting")
180 + return None
181
182 common_args = [
227 - # use half precision for speed and memory savings in CI environment
183 "--dtype",
229 - dtype,
184 + "bfloat16",
185 "--max-model-len",
186 "2048",
187 "--max-num-seqs",
188 "8",
189 ]
235 - if chunked_prefill:
236 - common_args.append("--enable-chunked-prefill")
190 if eager_mode:
191 common_args.append("-cc.cudagraph_mode=none")
192 if runner != "auto":
@@ -294,85 +247,85 @@ def _compare_sp(
247 "mp",
248 ]
249
297 - compare_two_settings(
298 - model_id,
299 - tp_sp_args,
300 - tp_args,
301 - method=method,
302 - force_v1_runner=True,
303 - )
250 + return tp_args, tp_sp_args
251 +
252 +
253 +def _compare_sp_settings(
254 + model_id: str,
255 + settings: list[tuple[list[str], list[str]]],
256 + *,
257 + method: Literal["generate", "encode"],
258 +) -> None:
259 + if not settings:
260 + pytest.skip("No supported sequence-parallel configurations")
261 +
262 + settings_by_tp_args: dict[tuple[str, ...], list[list[str]]] = {}
263 + for tp_args, tp_sp_args in settings:
264 + settings_by_tp_args.setdefault(tuple(tp_args), []).append(tp_sp_args)
265 +
266 + for tp_args_key, grouped_tp_sp_args in settings_by_tp_args.items():
267 + all_args = [list(tp_args_key), *grouped_tp_sp_args]
268 + compare_all_settings(
269 + model_id,
270 + all_args,
271 + [None] * len(all_args),
272 + method=method,
273 + force_v1_runner=True,
274 + )
275
276
277 SP_TEXT_GENERATION_MODELS = {
278 # [Decoder-only]
279 "hmellor/tiny-random-LlamaForCausalLM": SPTestSettings.fast(),
309 - "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8": SPTestSettings.fp8_quant(),
280 + "RedHatAI/Llama-3.2-1B-Instruct-FP8": SPTestSettings.fp8_quant(),
281 }
282
312 -SP_TEST_MODELS = [
313 - # TODO support other models
314 - # [LANGUAGE GENERATION]
315 - "hmellor/tiny-random-LlamaForCausalLM",
316 - "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
317 -]
318 -
283
284 @pytest.mark.parametrize(
321 - (
322 - "model_id",
323 - "parallel_setup",
324 - "distributed_backend",
325 - "runner",
326 - "test_options",
327 - ),
285 + ("model_id", "settings"),
286 [
329 - params
287 + pytest.param(model_id, settings, id=model_id)
288 for model_id, settings in SP_TEXT_GENERATION_MODELS.items()
331 - for params in settings.iter_params(model_id)
332 - if model_id in SP_TEST_MODELS
289 ],
290 )
335 -@pytest.mark.parametrize("use_inductor_graph_partition", [True, False])
336 -@pytest.mark.parametrize("fuse_gemm_comms", [False]) # TODO: enable async TP
291 @create_new_process_for_each_test()
292 def test_tp_sp_generation(
293 model_id: str,
340 - parallel_setup: ParallelSetup,
341 - distributed_backend: str,
342 - runner: RunnerOption,
343 - test_options: SPTestOptions,
294 + settings: SPTestSettings,
295 num_gpus_available,
345 - use_inductor_graph_partition: bool,
346 - fuse_gemm_comms: bool,
296 ):
348 - if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
349 - pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
350 -
297 # Skip FP8 SP-only test on sm89 (compute capability 8.9).
298 # An unknown capability (None) is left to fail in the test body rather than
299 # being silently skipped here.
300 capability = current_platform.get_device_capability()
355 - if (
356 - "fp8" in model_id.lower()
357 - and capability is not None
358 - and capability < (9, 0)
359 - and (not fuse_gemm_comms)
360 - ):
301 + if "fp8" in model_id.lower() and capability is not None and capability < (9, 0):
302 pytest.skip("FP8 reduction support begins with sm90 capable devices.")
303
363 - _compare_sp(
364 - model_id,
365 - parallel_setup,
366 - distributed_backend,
367 - runner,
368 - test_options,
369 - num_gpus_available,
370 - use_inductor_graph_partition,
371 - fuse_gemm_comms=fuse_gemm_comms,
372 - enable_prompt_embeds=False,
373 - method="generate",
374 - is_multimodal=False,
375 - )
304 + graph_partition_options = [False]
305 + if is_torch_equal_or_newer("2.9.0.dev"):
306 + graph_partition_options.append(True)
307 +
308 + comparisons = []
309 + for _, parallel_setup, backend, runner, test_options in settings.iter_params(
310 + model_id
311 + ):
312 + for use_inductor_graph_partition in graph_partition_options:
313 + comparison = _build_sp_args(
314 + model_id,
315 + parallel_setup,
316 + backend,
317 + runner,
318 + test_options,
319 + num_gpus_available,
320 + use_inductor_graph_partition,
321 + fuse_gemm_comms=False, # TODO: enable async TP
322 + enable_prompt_embeds=False,
323 + is_multimodal=False,
324 + )
325 + if comparison is not None:
326 + comparisons.append(comparison)
327 +
328 + _compare_sp_settings(model_id, comparisons, method="generate")
329
330
331 # Focused regression test for the SP + prompt_embeds graph-rewrite path.
@@ -385,36 +338,39 @@ SP_PROMPT_EMBEDS_PARALLEL_SETUPS = [
338 fuse_norm_quant=False,
339 fuse_act_quant=False,
340 eager_mode=False,
388 - chunked_prefill=False,
341 )
342 for pp_size in [1, 2]
343 ]
344
345
394 -@pytest.mark.parametrize("parallel_setup", SP_PROMPT_EMBEDS_PARALLEL_SETUPS)
395 -@pytest.mark.parametrize("use_inductor_graph_partition", [True, False])
346 @create_new_process_for_each_test()
347 def test_tp_sp_generation_prompt_embeds(
398 - parallel_setup: ParallelSetup,
348 num_gpus_available,
400 - use_inductor_graph_partition: bool,
349 ):
402 - if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"):
403 - pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
404 -
405 - _compare_sp(
406 - "hmellor/tiny-random-LlamaForCausalLM",
407 - parallel_setup,
408 - distributed_backend="mp",
409 - runner="auto",
410 - test_options=SPTestOptions(multi_node_only=False, load_format=None),
411 - num_gpus_available=num_gpus_available,
412 - use_inductor_graph_partition=use_inductor_graph_partition,
413 - fuse_gemm_comms=False,
414 - enable_prompt_embeds=True,
415 - method="generate",
416 - is_multimodal=False,
417 - )
350 + model_id = "hmellor/tiny-random-LlamaForCausalLM"
351 + graph_partition_options = [False]
352 + if is_torch_equal_or_newer("2.9.0.dev"):
353 + graph_partition_options.append(True)
354 +
355 + comparisons = []
356 + for parallel_setup in SP_PROMPT_EMBEDS_PARALLEL_SETUPS:
357 + for use_inductor_graph_partition in graph_partition_options:
358 + comparison = _build_sp_args(
359 + model_id,
360 + parallel_setup,
361 + distributed_backend="mp",
362 + runner="auto",
363 + test_options=SPTestOptions(multi_node_only=False, load_format=None),
364 + num_gpus_available=num_gpus_available,
365 + use_inductor_graph_partition=use_inductor_graph_partition,
366 + fuse_gemm_comms=False,
367 + enable_prompt_embeds=True,
368 + is_multimodal=False,
369 + )
370 + if comparison is not None:
371 + comparisons.append(comparison)
372 +
373 + _compare_sp_settings(model_id, comparisons, method="generate")
374
375
376 @create_new_process_for_each_test()
@@ -425,7 +381,7 @@ def test_tp_sp_nvfp4_generation(num_gpus_available: int):
381 ):
382 pytest.skip("NVFP4 requires Blackwell")
383
428 - _compare_sp(
384 + comparison = _build_sp_args(
385 NVFP4_MODEL_ID,
386 ParallelSetup(
387 tp_size=2,
@@ -433,7 +389,6 @@ def test_tp_sp_nvfp4_generation(num_gpus_available: int):
389 fuse_norm_quant=True,
390 fuse_act_quant=True,
391 eager_mode=True,
436 - chunked_prefill=False,
392 ),
393 "mp",
394 "auto",
@@ -446,7 +401,10 @@ def test_tp_sp_nvfp4_generation(num_gpus_available: int):
401 use_inductor_graph_partition=False,
402 fuse_gemm_comms=False,
403 enable_prompt_embeds=False,
449 - method="generate",
404 is_multimodal=False,
451 - dtype="bfloat16",
405 + )
406 + _compare_sp_settings(
407 + NVFP4_MODEL_ID,
408 + [] if comparison is None else [comparison],
409 + method="generate",
410 )
tests/compile/passes/distributed/test_async_tp.py
+4 -19
@@ -45,7 +45,7 @@ prompts = [
45
46
47 class TestMMRSModel(torch.nn.Module):
48 - def __init__(self, hidden_size=16, dtype=torch.float16):
48 + def __init__(self, hidden_size=16, dtype=torch.bfloat16):
49 super().__init__()
50 self.hidden_size = hidden_size
51 self.dtype = dtype
@@ -77,7 +77,7 @@ class TestMMRSModel(torch.nn.Module):
77
78
79 class TestAGMMModel(torch.nn.Module):
80 - def __init__(self, hidden_size=16, dtype=torch.float16):
80 + def __init__(self, hidden_size=16, dtype=torch.bfloat16):
81 super().__init__()
82 self.hidden_size = hidden_size
83 self.dtype = dtype
@@ -106,7 +106,7 @@ class TestAGMMModel(torch.nn.Module):
106
107
108 class _BaseScaledMMModel(torch.nn.Module):
109 - def __init__(self, hidden_size=16, dtype=torch.float16):
109 + def __init__(self, hidden_size=16, dtype=torch.bfloat16):
110 super().__init__()
111 self.hidden_size = hidden_size
112 self.dtype = dtype
@@ -254,7 +254,7 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel):
254 @pytest.mark.parametrize("batch_size", [8])
255 @pytest.mark.parametrize("seq_len", [16])
256 @pytest.mark.parametrize("hidden_size", [16])
257 -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
257 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
258 @pytest.mark.parametrize("dynamic", [True, False])
259 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA")
260 def test_async_tp_pass_replace(
@@ -265,21 +265,6 @@ def test_async_tp_pass_replace(
265 dtype: torch.dtype,
266 dynamic: bool,
267 ):
268 - if (
269 - test_model
270 - in (
271 - TestScaledMMRSModel,
272 - TestAGScaledMMModel,
273 - TestCutlassScaledMMRSModel,
274 - TestAGCutlassScaledMMModel,
275 - )
276 - and dtype == torch.float16
277 - ):
278 - pytest.skip(
279 - "Only bf16 high precision output types are supported for "
280 - "per-token (row-wise) scaling"
281 - )
282 -
268 num_processes = 2
269 master_port = str(get_open_port())
270
tests/models/registry.py
+1
@@ -376,6 +376,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
376 "guard": "meta-llama/Llama-Guard-3-1B",
377 "hermes": "NousResearch/Hermes-3-Llama-3.1-8B",
378 "fp8": "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
379 + "fp8_1b": "RedHatAI/Llama-3.2-1B-Instruct-FP8",
380 "tiny": "hmellor/tiny-random-LlamaForCausalLM",
381 },
382 ),