iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Michael Goin
2026-08-04 19:18:50 -07:00
committed by GitHub
parent 33c50587d2
commit 9cd9f8c20b
5 changed files with 112 additions and 198 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/compile.yaml
+5 -31
View File
@@ -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/"
tests/compile/correctness_e2e/test_async_tp.py
+1 -5
View File
@@ -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,
tests/compile/correctness_e2e/test_sequence_parallel.py
+101 -143
View File
@@ -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",
)
tests/compile/passes/distributed/test_async_tp.py
+4 -19
View File
@@ -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())
tests/models/registry.py
+1
View File
@@ -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",
},
),