#!/usr/bin/env python3 """Profile the vLLM 0.20 TP all-reduce path under a frozen runtime contract.""" from __future__ import annotations import argparse import json import os import statistics import subprocess from pathlib import Path from typing import Any import torch import torch.distributed as dist import vllm VLLM_VERSION = "0.20.0" VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--vllm-source", type=Path, required=True) parser.add_argument("--model", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--num-tokens", type=int, nargs="+", default=[8]) parser.add_argument("--hidden-dim", type=int, default=2048) parser.add_argument( "--collective-contract", choices=("flashinfer-trtllm", "qwen235-serving-projected"), default="flashinfer-trtllm", ) parser.add_argument("--warmup-iters", type=int, default=3) parser.add_argument("--repeats", type=int, default=10) parser.add_argument("--trials", type=int, default=1) return parser.parse_args() def git_head(repo: Path) -> str: return subprocess.check_output( ["git", "-C", str(repo), "rev-parse", "HEAD"], text=True ).strip() def stats_ms(samples: list[float]) -> dict[str, float]: return { "min": min(samples), "max": max(samples), "mean": statistics.fmean(samples), "median": statistics.median(samples), "std": statistics.pstdev(samples), } def main() -> None: args = parse_args() if vllm.__version__ != VLLM_VERSION: raise SystemExit(f"expected vLLM {VLLM_VERSION}, got {vllm.__version__}") source_head = git_head(args.vllm_source) if source_head != VLLM_COMMIT: raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}") if args.repeats <= 0 or args.trials <= 0 or args.warmup_iters < 0: raise SystemExit("warmup/repeats/trials must be non-negative/positive") if os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER") != "1": raise SystemExit("VLLM_ALLREDUCE_USE_FLASHINFER must equal 1") if os.getenv("VLLM_FLASHINFER_ALLREDUCE_BACKEND") != "trtllm": raise SystemExit("VLLM_FLASHINFER_ALLREDUCE_BACKEND must equal trtllm") expected_symm_mem = ( "1" if args.collective_contract == "qwen235-serving-projected" else "0" ) if os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM") != expected_symm_mem: raise SystemExit( f"VLLM_ALLREDUCE_USE_SYMM_MEM must equal {expected_symm_mem}" ) if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ: raise SystemExit("launch with torchrun") from vllm.distributed import ( destroy_distributed_environment, destroy_model_parallel, init_distributed_environment, initialize_model_parallel, set_custom_all_reduce, tensor_model_parallel_all_reduce, ) from vllm.distributed.parallel_state import get_tp_group from vllm.config import ( ModelConfig, ParallelConfig, VllmConfig, set_current_vllm_config, ) rank = int(os.environ["RANK"]) local_rank = int(os.environ["LOCAL_RANK"]) world_size = int(os.environ["WORLD_SIZE"]) if world_size not in (2, 4, 8): raise SystemExit(f"expected TP world size 2, 4, or 8, got {world_size}") device = torch.device(f"cuda:{local_rank}") torch.accelerator.set_device_index(device) torch.set_default_device(device) disable_custom_all_reduce = ( args.collective_contract == "qwen235-serving-projected" ) set_custom_all_reduce(not disable_custom_all_reduce) init_distributed_environment() model_config = ModelConfig( model=str(args.model), dtype="bfloat16", max_model_len=8192, skip_tokenizer_init=True, generation_config="vllm", ) vllm_config = VllmConfig( model_config=model_config, parallel_config=ParallelConfig( tensor_parallel_size=world_size, disable_custom_all_reduce=disable_custom_all_reduce, ) ) with set_current_vllm_config(vllm_config): initialize_model_parallel(tensor_model_parallel_size=world_size) from vllm.distributed.device_communicators.all_reduce_utils import ( should_nccl_symm_mem_allreduce, ) def resolve_profile_operation(input_tensor: torch.Tensor): communicator = get_tp_group().device_communicator if args.collective_contract == "qwen235-serving-projected": from vllm.compilation.passes.fusion.allreduce_rms_fusion import ( FI_ALLREDUCE_FUSION_MAX_SIZE_MB, ) major, minor = torch.cuda.get_device_capability(device) capability = major * 10 + minor try: fusion_limit_bytes = int( FI_ALLREDUCE_FUSION_MAX_SIZE_MB[capability][world_size] * 1024 * 1024 ) except KeyError as error: raise SystemExit( f"no FlashInfer fusion limit for SM{capability}, TP{world_size}" ) from error if input_tensor.nbytes <= fusion_limit_bytes: from vllm.distributed.device_communicators.flashinfer_all_reduce import ( get_fi_ar_workspace, ) fi_comm = communicator.fi_ar_comm if fi_comm is None or fi_comm.disabled: raise SystemExit("FlashInfer all-reduce communicator is unavailable") max_token_num = fusion_limit_bytes // ( args.hidden_dim * input_tensor.element_size() ) workspace = get_fi_ar_workspace( world_size=world_size, rank=rank, max_token_num=max_token_num, hidden_dim=args.hidden_dim, dtype=input_tensor.dtype, group=get_tp_group().device_group, ) if workspace is None: raise SystemExit( "FlashInfer rejected the serving-matched integer workspace" ) if not fi_comm.should_use_fi_ar(input_tensor): raise SystemExit("FlashInfer rejected a fusion-eligible payload") return ( "flashinfer_trtllm_fused_projection", fi_comm.all_reduce, fusion_limit_bytes, ) if ( communicator.pynccl_comm is not None and should_nccl_symm_mem_allreduce(world_size, input_tensor) ): return ( "pynccl_symmetric_with_copy", torch.ops.vllm.all_reduce_symmetric_with_copy, fusion_limit_bytes, ) if ( communicator.symm_mem_comm is not None and communicator.symm_mem_comm.should_use_symm_mem(input_tensor) ): return ( "torch_symmetric_memory", communicator.symm_mem_comm.all_reduce, fusion_limit_bytes, ) if communicator.pynccl_comm is not None: return "pynccl", communicator.pynccl_comm.all_reduce, fusion_limit_bytes return "torch_distributed", tensor_model_parallel_all_reduce, fusion_limit_bytes fi_comm = communicator.fi_ar_comm selected_backend = ( "flashinfer_trtllm" if fi_comm is not None and not fi_comm.disabled and fi_comm.should_use_fi_ar(input_tensor) else "non_flashinfer_fallback" ) return selected_backend, tensor_model_parallel_all_reduce, None rows: list[dict[str, Any]] = [] expected_sum = world_size * (world_size + 1) / 2 try: for num_tokens in args.num_tokens: input_tensor = torch.full( (num_tokens, args.hidden_dim), float(rank + 1), dtype=torch.bfloat16, device=device, ) selected_backend, profile_operation, fusion_limit_bytes = ( resolve_profile_operation(input_tensor) ) for _ in range(args.warmup_iters): output = profile_operation(input_tensor) torch.accelerator.synchronize() torch.testing.assert_close( output, torch.full_like(output, expected_sum), atol=0.0, rtol=0.0, ) if args.collective_contract == "flashinfer-trtllm": if selected_backend != "flashinfer_trtllm": raise SystemExit( f"expected FlashInfer TRTLLM, got {selected_backend} at " f"TP={world_size}, tokens={num_tokens}" ) elif selected_backend not in { "flashinfer_trtllm_fused_projection", "pynccl_symmetric_with_copy", "torch_symmetric_memory", "pynccl", }: raise SystemExit( f"unexpected Qwen235 serving backend: {selected_backend}" ) all_trial_rank_samples: list[list[list[float]]] = [] critical_path_samples: list[float] = [] for _ in range(args.trials): samples: list[float] = [] for _ in range(args.repeats): dist.barrier() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() output = profile_operation(input_tensor) end.record() torch.accelerator.synchronize() samples.append(float(start.elapsed_time(end))) gathered: list[list[float] | None] = [None] * world_size dist.all_gather_object(gathered, samples) if rank == 0: rank_samples = [item for item in gathered if item is not None] all_trial_rank_samples.append(rank_samples) critical_path_samples.extend( max(per_rank[index] for per_rank in rank_samples) for index in range(args.repeats) ) if rank == 0: flattened_by_rank = [ [ sample for trial in all_trial_rank_samples for sample in trial[rank_index] ] for rank_index in range(world_size) ] per_rank = [stats_ms(samples) for samples in flattened_by_rank] row = { "tensor_parallel_size": world_size, "num_tokens": num_tokens, "hidden_dim": args.hidden_dim, "payload_bytes": num_tokens * args.hidden_dim * torch.tensor([], dtype=torch.bfloat16).element_size(), "dtype": "bfloat16", "communicator": "vllm.tensor_model_parallel_all_reduce", "collective_contract": args.collective_contract, "selected_backend": selected_backend, "real_fusion_limit_bytes": fusion_limit_bytes, "trials": args.trials, "repeats_per_trial": args.repeats, "per_trial_rank_samples_ms": all_trial_rank_samples, "per_rank_time_ms": per_rank, "critical_path_time_ms": stats_ms(critical_path_samples), "critical_path_median_ms": statistics.median( critical_path_samples ), } rows.append(row) print(json.dumps(row, sort_keys=True), flush=True) finally: destroy_model_parallel() destroy_distributed_environment() if rank == 0: payload = { "schema_version": ( "qwen30_vllm020_allreduce_raw.v1" if args.collective_contract == "flashinfer-trtllm" else "vllm020_allreduce_raw.v2" ), "environment": { "vllm_version": vllm.__version__, "vllm_source_commit": source_head, "torch_version": torch.__version__, "torch_cuda": torch.version.cuda, "gpu": torch.cuda.get_device_name(device), "model": str(args.model), "backend_env": { "VLLM_ALLREDUCE_USE_FLASHINFER": os.getenv( "VLLM_ALLREDUCE_USE_FLASHINFER", "0" ), "VLLM_FLASHINFER_ALLREDUCE_BACKEND": os.getenv( "VLLM_FLASHINFER_ALLREDUCE_BACKEND" ), "VLLM_ALLREDUCE_USE_SYMM_MEM": os.getenv( "VLLM_ALLREDUCE_USE_SYMM_MEM", "1" ), }, "collective_contract": args.collective_contract, "disable_custom_all_reduce": disable_custom_all_reduce, }, "rows": rows, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") if __name__ == "__main__": main()