#!/usr/bin/env python3 """Profile the exact vLLM 0.20 FlashAttention backend at TP-local shapes. This deliberately uses vLLM's own v0.20.0 attention benchmark runner instead of Frontier's FlashInfer-only attention wrapper. The output is raw evidence; projection into Frontier's split attention CSV schema is a separate step. """ from __future__ import annotations import argparse import json import subprocess import sys import types from pathlib import Path import torch 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("--tp", type=int, choices=(1, 2, 4), nargs="+", default=[1, 2, 4]) parser.add_argument( "--batch-specs", nargs="+", default=["q128", "4q1s128", "q128_4q1s128"], ) parser.add_argument("--warmup-iters", type=int, default=3) parser.add_argument("--repeats", type=int, default=5) parser.add_argument("--device", default="cuda:0") 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 json_default(value: object) -> object: if isinstance(value, (Path, torch.dtype, torch.device)): return str(value) item = getattr(value, "item", None) if callable(item): return item() raise TypeError(f"cannot serialize {type(value).__name__}") 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 not args.model.joinpath("config.json").is_file(): raise SystemExit(f"missing model config: {args.model / 'config.json'}") bench_dir = args.vllm_source / "benchmarks" / "attention_benchmarks" sys.path.insert(0, str(bench_dir)) import runner # type: ignore[import-not-found] # noqa: PLC0415 from common import BenchmarkConfig # type: ignore[import-not-found] # noqa: PLC0415 from vllm.config import ( # noqa: PLC0415 CacheConfig, CompilationConfig, DeviceConfig, LoadConfig, ModelConfig, ParallelConfig, SchedulerConfig, VllmConfig, ) from vllm.v1.worker.workspace import init_workspace_manager # noqa: PLC0415 def create_vllm_config(config: BenchmarkConfig, max_num_blocks: int) -> VllmConfig: model_config = ModelConfig( model=str(args.model), tokenizer=str(args.model), trust_remote_code=False, dtype="bfloat16", seed=0, max_model_len=40960, ) cache_config = CacheConfig(block_size=config.block_size, cache_dtype="auto") cache_config.num_gpu_blocks = max_num_blocks cache_config.num_cpu_blocks = 0 parallel_config = ParallelConfig(tensor_parallel_size=1) scheduler_config = SchedulerConfig( max_num_seqs=256, max_num_batched_tokens=8192, max_model_len=40960, is_encoder_decoder=False, enable_chunked_prefill=True, ) model_config.get_num_layers = types.MethodType( lambda self: config.num_layers, model_config ) model_config.get_sliding_window_for_layer = types.MethodType( lambda self, i: None, model_config ) model_config.get_logits_soft_cap_for_layer = types.MethodType( lambda self, i: 0.0, model_config ) model_config.get_sm_scale_for_layer = types.MethodType( lambda self, i: 1.0 / config.head_dim**0.5, model_config ) model_config.get_num_attention_heads = types.MethodType( lambda self, parallel_config=None: config.num_q_heads, model_config ) model_config.get_num_kv_heads = types.MethodType( lambda self, parallel_config=None: config.num_kv_heads, model_config ) model_config.get_head_size = types.MethodType( lambda self: config.head_dim, model_config ) model_config.get_sliding_window = types.MethodType( lambda self: None, model_config ) return VllmConfig( model_config=model_config, cache_config=cache_config, parallel_config=parallel_config, scheduler_config=scheduler_config, device_config=DeviceConfig(), load_config=LoadConfig(), compilation_config=CompilationConfig(), ) runner._create_vllm_config = create_vllm_config init_workspace_manager(args.device) rows: list[dict[str, object]] = [] for tp in args.tp: for batch_spec in args.batch_specs: config = BenchmarkConfig( backend="FLASH_ATTN", batch_spec=batch_spec, num_layers=1, head_dim=128, num_q_heads=32 // tp, num_kv_heads=4 // tp, block_size=16, device=args.device, dtype=torch.bfloat16, repeats=args.repeats, warmup_iters=args.warmup_iters, profile_memory=True, kv_cache_dtype="auto", use_cuda_graphs=False, ) result = runner.run_attention_benchmark(config) row = result.to_dict() row["tensor_parallel_size"] = tp rows.append(row) print( json.dumps( { "tp": tp, "batch_spec": batch_spec, "mean_time_s": result.mean_time, "error": result.error, }, sort_keys=True, default=json_default, ), flush=True, ) if not result.success: raise SystemExit(f"attention profile failed: {row}") payload = { "schema_version": "qwen30_vllm020_flashattn_raw.v1", "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(torch.device(args.device)), "model": str(args.model), "dtype": "bfloat16", "attention_backend": "FLASH_ATTN", "block_size": 16, }, "rows": rows, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text( json.dumps(payload, indent=2, sort_keys=True, default=json_default) + "\n" ) if __name__ == "__main__": main()