203 lines
7.0 KiB
Python
203 lines
7.0 KiB
Python
#!/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()
|