From 9cd9f8c20bc34e1a3f57f8a258f5cc8751e76919 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 4 Aug 2026 22:18:50 -0400 Subject: [PATCH] Prune redundant tests points in `correctness_e2e/[test_sequence_parallel,test_async_tp]` (#51068) Signed-off-by: mgoin --- .buildkite/test_areas/compile.yaml | 36 +-- .../compile/correctness_e2e/test_async_tp.py | 6 +- .../correctness_e2e/test_sequence_parallel.py | 244 ++++++++---------- .../passes/distributed/test_async_tp.py | 23 +- tests/models/registry.py | 1 + 5 files changed, 112 insertions(+), 198 deletions(-) diff --git a/.buildkite/test_areas/compile.yaml b/.buildkite/test_areas/compile.yaml index 3a0ea308a9..3ce4d7709e 100644 --- a/.buildkite/test_areas/compile.yaml +++ b/.buildkite/test_areas/compile.yaml @@ -2,44 +2,18 @@ group: Compile depends_on: - image-build steps: -- label: Sequence Parallel Correctness Tests (2 GPUs) - key: sequence-parallel-correctness-tests-2-gpus - timeout_in_minutes: 80 +- label: Sequence Parallel Correctness Tests (2xB200) + key: sequence-parallel-correctness-tests-2xb200 + timeout_in_minutes: 45 working_dir: "/vllm-workspace/" - num_devices: 2 - source_file_dependencies: - - vllm/model_executor/layers/ - - vllm/compilation/ - - vllm/v1/worker/ - - vllm/v1/cudagraph_dispatcher.py - - tests/compile/correctness_e2e/test_sequence_parallel.py - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py - -- label: Sequence Parallel Correctness Tests (2xH100) - key: sequence-parallel-correctness-tests-2xh100 - timeout_in_minutes: 75 - working_dir: "/vllm-workspace/" - device: h100 + device: b200-k8s optional: true num_devices: 2 commands: - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - pytest -v -s tests/compile/correctness_e2e/test_sequence_parallel.py -- label: AsyncTP Correctness Tests (2xH100) - key: asynctp-correctness-tests-2xh100 - timeout_in_minutes: 30 - working_dir: "/vllm-workspace/" - device: h100 - optional: true - num_devices: 2 - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY=1 - - pytest -v -s tests/compile/correctness_e2e/test_async_tp.py - -- label: AsyncTP Correctness Tests (B200) +- label: AsyncTP Correctness Tests (2xB200) key: asynctp-correctness-tests-b200 timeout_in_minutes: 30 working_dir: "/vllm-workspace/" diff --git a/tests/compile/correctness_e2e/test_async_tp.py b/tests/compile/correctness_e2e/test_async_tp.py index e2d597bc7a..97248d153f 100644 --- a/tests/compile/correctness_e2e/test_async_tp.py +++ b/tests/compile/correctness_e2e/test_async_tp.py @@ -29,18 +29,16 @@ NVFP4_HF_OVERRIDES = { @create_new_process_for_each_test() @pytest.mark.parametrize( "model_id", - ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8"], + ["meta-llama/Llama-3.2-1B-Instruct", "RedHatAI/Llama-3.2-1B-Instruct-FP8"], ) @pytest.mark.parametrize("tp_size", [2]) @pytest.mark.parametrize("async_tp_enabled", [True]) @pytest.mark.parametrize("distributed_backend", ["mp"]) -@pytest.mark.parametrize("eager_mode", [False, True]) def test_async_tp_pass_correctness( model_id: str, tp_size: int, async_tp_enabled: bool, distributed_backend: str, - eager_mode: bool, num_gpus_available: int, monkeypatch, ): @@ -64,8 +62,6 @@ def test_async_tp_pass_correctness( "--max-num-seqs", "8", ] - if eager_mode: - common_args.append("--enforce-eager") compilation_config = { "mode": CompilationMode.VLLM_COMPILE, diff --git a/tests/compile/correctness_e2e/test_sequence_parallel.py b/tests/compile/correctness_e2e/test_sequence_parallel.py index 9c9c596198..eed5efb7b3 100644 --- a/tests/compile/correctness_e2e/test_sequence_parallel.py +++ b/tests/compile/correctness_e2e/test_sequence_parallel.py @@ -22,7 +22,7 @@ from vllm.platforms import current_platform from vllm.utils.torch_utils import is_torch_equal_or_newer from ...models.registry import HF_EXAMPLE_MODELS, _HfExamplesInfo -from ...utils import compare_two_settings, create_new_process_for_each_test +from ...utils import compare_all_settings, create_new_process_for_each_test logger = init_logger("test_sequence_parallel") @@ -37,7 +37,6 @@ class ParallelSetup(NamedTuple): fuse_norm_quant: bool fuse_act_quant: bool eager_mode: bool - chunked_prefill: bool class SPTestOptions(NamedTuple): @@ -53,38 +52,6 @@ class SPTestSettings: runner: RunnerOption test_options: SPTestOptions - @staticmethod - def detailed( - *, - tp_base: int = 2, - pp_base: int = 1, - multi_node_only: bool = False, - runner: RunnerOption = "auto", - load_format: str | None = None, - ): - parallel_setups = [] - for eager_mode_val in [False, True]: - for pp_multiplier in [1, 2]: - for chunked_prefill_val in [False, True]: - parallel_setups.append( - ParallelSetup( - tp_size=tp_base, - pp_size=pp_multiplier * pp_base, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=eager_mode_val, - chunked_prefill=chunked_prefill_val, - ) - ) - return SPTestSettings( - parallel_setups=parallel_setups, - distributed_backends=["mp", "ray"], - runner=runner, - test_options=SPTestOptions( - multi_node_only=multi_node_only, load_format=load_format - ), - ) - @staticmethod def fast( *, @@ -94,20 +61,16 @@ class SPTestSettings: multi_node_only: bool = False, load_format: str | None = None, ): - parallel_setups = [] - for eager_mode_val in [False, True]: - for pp_multiplier in [1, 2]: - for chunked_prefill_val in [False, True]: - parallel_setups.append( - ParallelSetup( - tp_size=tp_base, - pp_size=pp_multiplier * pp_base, - fuse_norm_quant=False, - fuse_act_quant=False, - eager_mode=eager_mode_val, - chunked_prefill=chunked_prefill_val, - ) - ) + parallel_setups = [ + ParallelSetup( + tp_size=tp_base, + pp_size=pp_multiplier * pp_base, + fuse_norm_quant=False, + fuse_act_quant=False, + eager_mode=False, + ) + for pp_multiplier in [1, 2] + ] return SPTestSettings( parallel_setups=parallel_setups, distributed_backends=["mp", "ray"], @@ -135,7 +98,6 @@ class SPTestSettings: fuse_norm_quant=fusion_val, fuse_act_quant=fusion_val, eager_mode=True, - chunked_prefill=False, ) ) return SPTestSettings( @@ -161,7 +123,7 @@ class SPTestSettings: ) -def _compare_sp( +def _build_sp_args( model_id: str, parallel_setup: ParallelSetup, distributed_backend: str, @@ -172,17 +134,14 @@ def _compare_sp( fuse_gemm_comms: bool, enable_prompt_embeds: bool, *, - method: Literal["generate", "encode"], is_multimodal: bool, - dtype: str = "float16", -): +) -> tuple[list[str], list[str]] | None: ( tp_size, pp_size, fuse_norm_quant, fuse_act_quant, eager_mode, - chunked_prefill, ) = parallel_setup multi_node_only = test_options.multi_node_only @@ -214,26 +173,20 @@ def _compare_sp( model_info.check_available_online(on_fail="skip") if num_gpus_available < tp_size * pp_size: - pytest.skip(f"Need at least {tp_size} x {pp_size} GPUs") + return None if VLLM_MULTI_NODE and distributed_backend == "mp": - pytest.skip( - "Skipping multi-node pipeline parallel test for " - "multiprocessing distributed backend" - ) + return None if multi_node_only and not VLLM_MULTI_NODE: - pytest.skip("Not in multi-node setting") + return None common_args = [ - # use half precision for speed and memory savings in CI environment "--dtype", - dtype, + "bfloat16", "--max-model-len", "2048", "--max-num-seqs", "8", ] - if chunked_prefill: - common_args.append("--enable-chunked-prefill") if eager_mode: common_args.append("-cc.cudagraph_mode=none") if runner != "auto": @@ -294,85 +247,85 @@ def _compare_sp( "mp", ] - compare_two_settings( - model_id, - tp_sp_args, - tp_args, - method=method, - force_v1_runner=True, - ) + return tp_args, tp_sp_args + + +def _compare_sp_settings( + model_id: str, + settings: list[tuple[list[str], list[str]]], + *, + method: Literal["generate", "encode"], +) -> None: + if not settings: + pytest.skip("No supported sequence-parallel configurations") + + settings_by_tp_args: dict[tuple[str, ...], list[list[str]]] = {} + for tp_args, tp_sp_args in settings: + settings_by_tp_args.setdefault(tuple(tp_args), []).append(tp_sp_args) + + for tp_args_key, grouped_tp_sp_args in settings_by_tp_args.items(): + all_args = [list(tp_args_key), *grouped_tp_sp_args] + compare_all_settings( + model_id, + all_args, + [None] * len(all_args), + method=method, + force_v1_runner=True, + ) SP_TEXT_GENERATION_MODELS = { # [Decoder-only] "hmellor/tiny-random-LlamaForCausalLM": SPTestSettings.fast(), - "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8": SPTestSettings.fp8_quant(), + "RedHatAI/Llama-3.2-1B-Instruct-FP8": SPTestSettings.fp8_quant(), } -SP_TEST_MODELS = [ - # TODO support other models - # [LANGUAGE GENERATION] - "hmellor/tiny-random-LlamaForCausalLM", - "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8", -] - @pytest.mark.parametrize( - ( - "model_id", - "parallel_setup", - "distributed_backend", - "runner", - "test_options", - ), + ("model_id", "settings"), [ - params + pytest.param(model_id, settings, id=model_id) for model_id, settings in SP_TEXT_GENERATION_MODELS.items() - for params in settings.iter_params(model_id) - if model_id in SP_TEST_MODELS ], ) -@pytest.mark.parametrize("use_inductor_graph_partition", [True, False]) -@pytest.mark.parametrize("fuse_gemm_comms", [False]) # TODO: enable async TP @create_new_process_for_each_test() def test_tp_sp_generation( model_id: str, - parallel_setup: ParallelSetup, - distributed_backend: str, - runner: RunnerOption, - test_options: SPTestOptions, + settings: SPTestSettings, num_gpus_available, - use_inductor_graph_partition: bool, - fuse_gemm_comms: bool, ): - if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): - pytest.skip("inductor graph partition is only available in PyTorch 2.9+") - # Skip FP8 SP-only test on sm89 (compute capability 8.9). # An unknown capability (None) is left to fail in the test body rather than # being silently skipped here. capability = current_platform.get_device_capability() - if ( - "fp8" in model_id.lower() - and capability is not None - and capability < (9, 0) - and (not fuse_gemm_comms) - ): + if "fp8" in model_id.lower() and capability is not None and capability < (9, 0): pytest.skip("FP8 reduction support begins with sm90 capable devices.") - _compare_sp( - model_id, - parallel_setup, - distributed_backend, - runner, - test_options, - num_gpus_available, - use_inductor_graph_partition, - fuse_gemm_comms=fuse_gemm_comms, - enable_prompt_embeds=False, - method="generate", - is_multimodal=False, - ) + graph_partition_options = [False] + if is_torch_equal_or_newer("2.9.0.dev"): + graph_partition_options.append(True) + + comparisons = [] + for _, parallel_setup, backend, runner, test_options in settings.iter_params( + model_id + ): + for use_inductor_graph_partition in graph_partition_options: + comparison = _build_sp_args( + model_id, + parallel_setup, + backend, + runner, + test_options, + num_gpus_available, + use_inductor_graph_partition, + fuse_gemm_comms=False, # TODO: enable async TP + enable_prompt_embeds=False, + is_multimodal=False, + ) + if comparison is not None: + comparisons.append(comparison) + + _compare_sp_settings(model_id, comparisons, method="generate") # Focused regression test for the SP + prompt_embeds graph-rewrite path. @@ -385,36 +338,39 @@ SP_PROMPT_EMBEDS_PARALLEL_SETUPS = [ fuse_norm_quant=False, fuse_act_quant=False, eager_mode=False, - chunked_prefill=False, ) for pp_size in [1, 2] ] -@pytest.mark.parametrize("parallel_setup", SP_PROMPT_EMBEDS_PARALLEL_SETUPS) -@pytest.mark.parametrize("use_inductor_graph_partition", [True, False]) @create_new_process_for_each_test() def test_tp_sp_generation_prompt_embeds( - parallel_setup: ParallelSetup, num_gpus_available, - use_inductor_graph_partition: bool, ): - if use_inductor_graph_partition and not is_torch_equal_or_newer("2.9.0.dev"): - pytest.skip("inductor graph partition is only available in PyTorch 2.9+") + model_id = "hmellor/tiny-random-LlamaForCausalLM" + graph_partition_options = [False] + if is_torch_equal_or_newer("2.9.0.dev"): + graph_partition_options.append(True) - _compare_sp( - "hmellor/tiny-random-LlamaForCausalLM", - parallel_setup, - distributed_backend="mp", - runner="auto", - test_options=SPTestOptions(multi_node_only=False, load_format=None), - num_gpus_available=num_gpus_available, - use_inductor_graph_partition=use_inductor_graph_partition, - fuse_gemm_comms=False, - enable_prompt_embeds=True, - method="generate", - is_multimodal=False, - ) + comparisons = [] + for parallel_setup in SP_PROMPT_EMBEDS_PARALLEL_SETUPS: + for use_inductor_graph_partition in graph_partition_options: + comparison = _build_sp_args( + model_id, + parallel_setup, + distributed_backend="mp", + runner="auto", + test_options=SPTestOptions(multi_node_only=False, load_format=None), + num_gpus_available=num_gpus_available, + use_inductor_graph_partition=use_inductor_graph_partition, + fuse_gemm_comms=False, + enable_prompt_embeds=True, + is_multimodal=False, + ) + if comparison is not None: + comparisons.append(comparison) + + _compare_sp_settings(model_id, comparisons, method="generate") @create_new_process_for_each_test() @@ -425,7 +381,7 @@ def test_tp_sp_nvfp4_generation(num_gpus_available: int): ): pytest.skip("NVFP4 requires Blackwell") - _compare_sp( + comparison = _build_sp_args( NVFP4_MODEL_ID, ParallelSetup( tp_size=2, @@ -433,7 +389,6 @@ def test_tp_sp_nvfp4_generation(num_gpus_available: int): fuse_norm_quant=True, fuse_act_quant=True, eager_mode=True, - chunked_prefill=False, ), "mp", "auto", @@ -446,7 +401,10 @@ def test_tp_sp_nvfp4_generation(num_gpus_available: int): use_inductor_graph_partition=False, fuse_gemm_comms=False, enable_prompt_embeds=False, - method="generate", is_multimodal=False, - dtype="bfloat16", + ) + _compare_sp_settings( + NVFP4_MODEL_ID, + [] if comparison is None else [comparison], + method="generate", ) diff --git a/tests/compile/passes/distributed/test_async_tp.py b/tests/compile/passes/distributed/test_async_tp.py index 7d1320bcca..861464c519 100644 --- a/tests/compile/passes/distributed/test_async_tp.py +++ b/tests/compile/passes/distributed/test_async_tp.py @@ -45,7 +45,7 @@ prompts = [ class TestMMRSModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.float16): + def __init__(self, hidden_size=16, dtype=torch.bfloat16): super().__init__() self.hidden_size = hidden_size self.dtype = dtype @@ -77,7 +77,7 @@ class TestMMRSModel(torch.nn.Module): class TestAGMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.float16): + def __init__(self, hidden_size=16, dtype=torch.bfloat16): super().__init__() self.hidden_size = hidden_size self.dtype = dtype @@ -106,7 +106,7 @@ class TestAGMMModel(torch.nn.Module): class _BaseScaledMMModel(torch.nn.Module): - def __init__(self, hidden_size=16, dtype=torch.float16): + def __init__(self, hidden_size=16, dtype=torch.bfloat16): super().__init__() self.hidden_size = hidden_size self.dtype = dtype @@ -254,7 +254,7 @@ class TestAGCutlassScaledMMModel(_BaseScaledMMModel): @pytest.mark.parametrize("batch_size", [8]) @pytest.mark.parametrize("seq_len", [16]) @pytest.mark.parametrize("hidden_size", [16]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("dynamic", [True, False]) @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") def test_async_tp_pass_replace( @@ -265,21 +265,6 @@ def test_async_tp_pass_replace( dtype: torch.dtype, dynamic: bool, ): - if ( - test_model - in ( - TestScaledMMRSModel, - TestAGScaledMMModel, - TestCutlassScaledMMRSModel, - TestAGCutlassScaledMMModel, - ) - and dtype == torch.float16 - ): - pytest.skip( - "Only bf16 high precision output types are supported for " - "per-token (row-wise) scaling" - ) - num_processes = 2 master_port = str(get_open_port()) diff --git a/tests/models/registry.py b/tests/models/registry.py index 4ea3120715..dd36833258 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -376,6 +376,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { "guard": "meta-llama/Llama-Guard-3-1B", "hermes": "NousResearch/Hermes-3-Llama-3.1-8B", "fp8": "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8", + "fp8_1b": "RedHatAI/Llama-3.2-1B-Instruct-FP8", "tiny": "hmellor/tiny-random-LlamaForCausalLM", }, ),