Add vLLM 0.20 per-TP FlashAttention profile smoke
This commit is contained in:
@@ -0,0 +1,190 @@
|
||||
#!/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 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,
|
||||
),
|
||||
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) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user