445 lines
18 KiB
Python
445 lines
18 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 statistics
|
|
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")
|
|
parser.add_argument("--profile-kv-update", action="store_true")
|
|
parser.add_argument(
|
|
"--profile-method",
|
|
choices=("cuda_event", "record_function"),
|
|
default="cuda_event",
|
|
)
|
|
parser.add_argument("--frontier-source", type=Path)
|
|
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'}")
|
|
if args.profile_method == "record_function" and args.frontier_source is None:
|
|
raise SystemExit("--frontier-source is required for --profile-method record_function")
|
|
|
|
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 batch_spec import parse_batch_spec # type: ignore[import-not-found] # noqa: PLC0415
|
|
from common import ( # type: ignore[import-not-found] # noqa: PLC0415
|
|
BenchmarkConfig,
|
|
BenchmarkResult,
|
|
)
|
|
from vllm.config import ( # noqa: PLC0415
|
|
CacheConfig,
|
|
CompilationConfig,
|
|
DeviceConfig,
|
|
LoadConfig,
|
|
ModelConfig,
|
|
ParallelConfig,
|
|
SchedulerConfig,
|
|
VllmConfig,
|
|
set_current_vllm_config,
|
|
)
|
|
from vllm.v1.attention.backends.utils import ( # noqa: PLC0415
|
|
get_kv_cache_layout,
|
|
set_kv_cache_layout,
|
|
)
|
|
from vllm.v1.kv_cache_interface import FullAttentionSpec # noqa: PLC0415
|
|
from vllm.v1.worker.workspace import init_workspace_manager # noqa: PLC0415
|
|
|
|
record_function_tracer = None
|
|
if args.profile_method == "record_function":
|
|
sys.path.insert(0, str(args.frontier_source.resolve()))
|
|
from frontier.profiling.utils.record_function_tracer import RecordFunctionTracer
|
|
|
|
record_function_tracer = RecordFunctionTracer
|
|
|
|
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)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
if args.profile_method == "record_function":
|
|
(args.output.parent / "profiler_traces").mkdir(exist_ok=True)
|
|
|
|
def profile_kv_cache_update(config: BenchmarkConfig) -> dict[str, float]:
|
|
device = torch.device(config.device)
|
|
requests = parse_batch_spec(config.batch_spec)
|
|
total_q = sum(request.q_len for request in requests)
|
|
max_kv = max(request.kv_len for request in requests)
|
|
max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
|
|
max_num_blocks = len(requests) * max_blocks_per_request
|
|
vllm_config = create_vllm_config(config, max_num_blocks)
|
|
dtype = vllm_config.model_config.dtype
|
|
with set_current_vllm_config(vllm_config):
|
|
backend_config = runner._get_backend_config(config.backend)
|
|
backend_class, impl, layer = runner._create_backend_impl(
|
|
backend_config, config, device, dtype
|
|
)
|
|
required_layout = backend_class.get_required_kv_cache_layout()
|
|
if required_layout is not None:
|
|
set_kv_cache_layout(required_layout)
|
|
get_kv_cache_layout.cache_clear()
|
|
common_metadata = runner._build_common_attn_metadata(
|
|
[request.q_len for request in requests],
|
|
[request.kv_len for request in requests],
|
|
config.block_size,
|
|
device,
|
|
)
|
|
kv_cache_spec = FullAttentionSpec(
|
|
block_size=config.block_size,
|
|
num_kv_heads=config.num_kv_heads,
|
|
head_size=config.head_dim,
|
|
dtype=dtype,
|
|
)
|
|
layer._kv_cache_spec = kv_cache_spec
|
|
_, key_list, value_list = runner._create_input_tensors(
|
|
config, total_q, device, dtype
|
|
)
|
|
cache_list = runner._create_kv_cache(
|
|
config, max_num_blocks, backend_class, device, dtype
|
|
)
|
|
for _ in range(config.warmup_iters):
|
|
for layer_index in range(config.num_layers):
|
|
impl.do_kv_cache_update(
|
|
layer,
|
|
key_list[layer_index],
|
|
value_list[layer_index],
|
|
cache_list[layer_index],
|
|
common_metadata.slot_mapping,
|
|
)
|
|
torch.accelerator.synchronize()
|
|
samples: list[float] = []
|
|
for _ in range(config.repeats):
|
|
start = torch.cuda.Event(enable_timing=True)
|
|
end = torch.cuda.Event(enable_timing=True)
|
|
start.record()
|
|
for layer_index in range(config.num_layers):
|
|
impl.do_kv_cache_update(
|
|
layer,
|
|
key_list[layer_index],
|
|
value_list[layer_index],
|
|
cache_list[layer_index],
|
|
common_metadata.slot_mapping,
|
|
)
|
|
end.record()
|
|
torch.accelerator.synchronize()
|
|
samples.append(float(start.elapsed_time(end)) / config.num_layers)
|
|
return {
|
|
"min_ms": min(samples),
|
|
"max_ms": max(samples),
|
|
"mean_ms": statistics.fmean(samples),
|
|
"median_ms": statistics.median(samples),
|
|
"std_ms": statistics.pstdev(samples),
|
|
}
|
|
|
|
def profile_kernel_only(
|
|
config: BenchmarkConfig,
|
|
) -> tuple[BenchmarkResult, dict[str, float] | None]:
|
|
"""Trace exactly one vLLM FA3 forward/KV-update per annotation.
|
|
|
|
`RecordFunctionTracer` is Frontier's actual KERNEL_ONLY collector: it
|
|
follows CUDA launch correlations and sums kernels under the annotation.
|
|
The profiling loop therefore contains no CUDA-event value relabeling.
|
|
"""
|
|
device = torch.device(config.device)
|
|
torch.accelerator.set_device_index(device)
|
|
backend_config = runner._get_backend_config(config.backend)
|
|
requests = parse_batch_spec(config.batch_spec)
|
|
q_lens = [request.q_len for request in requests]
|
|
kv_lens = [request.kv_len for request in requests]
|
|
total_q = sum(q_lens)
|
|
max_kv = max(kv_lens)
|
|
max_blocks_per_request = (max_kv + config.block_size - 1) // config.block_size
|
|
max_num_blocks = len(requests) * max_blocks_per_request
|
|
|
|
with runner.log_warnings_and_errors_only():
|
|
vllm_config = create_vllm_config(config, max_num_blocks)
|
|
dtype = vllm_config.model_config.dtype
|
|
with set_current_vllm_config(vllm_config):
|
|
backend_class, impl, layer = runner._create_backend_impl(
|
|
backend_config, config, device, dtype
|
|
)
|
|
required_layout = backend_class.get_required_kv_cache_layout()
|
|
if required_layout is not None:
|
|
set_kv_cache_layout(required_layout)
|
|
get_kv_cache_layout.cache_clear()
|
|
common_metadata = runner._build_common_attn_metadata(
|
|
q_lens, kv_lens, config.block_size, device
|
|
)
|
|
kv_cache_spec = FullAttentionSpec(
|
|
block_size=config.block_size,
|
|
num_kv_heads=config.num_kv_heads,
|
|
head_size=config.head_dim,
|
|
dtype=dtype,
|
|
)
|
|
builder = runner._create_metadata_builder(
|
|
backend_class, kv_cache_spec, vllm_config, device, config.backend
|
|
)
|
|
attn_metadata = builder.build(
|
|
common_prefix_len=0, common_attn_metadata=common_metadata
|
|
)
|
|
quantize_query = config.kv_cache_dtype.startswith("fp8") and getattr(
|
|
impl, "supports_quant_query_input", False
|
|
)
|
|
q_list, k_list, v_list = runner._create_input_tensors(
|
|
config, total_q, device, dtype, quantize_query=quantize_query
|
|
)
|
|
cache_list = runner._create_kv_cache(
|
|
config, max_num_blocks, backend_class, device, dtype
|
|
)
|
|
output = torch.empty(
|
|
total_q,
|
|
config.num_q_heads,
|
|
config.head_dim,
|
|
device=device,
|
|
dtype=dtype,
|
|
)
|
|
|
|
def run_core() -> None:
|
|
for layer_index in range(config.num_layers):
|
|
impl.forward(
|
|
layer,
|
|
q_list[layer_index],
|
|
k_list[layer_index],
|
|
v_list[layer_index],
|
|
cache_list[layer_index],
|
|
attn_metadata,
|
|
output=output,
|
|
)
|
|
|
|
for _ in range(config.warmup_iters):
|
|
run_core()
|
|
torch.accelerator.synchronize()
|
|
core_tracer = record_function_tracer(str(args.output.parent))
|
|
with core_tracer:
|
|
for _ in range(config.repeats):
|
|
with torch.profiler.record_function("vidur_attention_core"):
|
|
run_core()
|
|
core_stats = core_tracer.get_operation_time_stats()
|
|
if "attention_core" not in core_stats:
|
|
raise RuntimeError("missing KERNEL_ONLY FlashAttention core stats")
|
|
core = {
|
|
name: float(value) / config.num_layers
|
|
for name, value in core_stats["attention_core"].items()
|
|
}
|
|
|
|
kv_stats = None
|
|
if args.profile_kv_update:
|
|
def run_kv_cache_update() -> None:
|
|
for layer_index in range(config.num_layers):
|
|
impl.do_kv_cache_update(
|
|
layer,
|
|
k_list[layer_index],
|
|
v_list[layer_index],
|
|
cache_list[layer_index],
|
|
common_metadata.slot_mapping,
|
|
)
|
|
|
|
for _ in range(config.warmup_iters):
|
|
run_kv_cache_update()
|
|
torch.accelerator.synchronize()
|
|
kv_tracer = record_function_tracer(str(args.output.parent))
|
|
with kv_tracer:
|
|
for _ in range(config.repeats):
|
|
with torch.profiler.record_function("vidur_attn_kv_cache_save"):
|
|
run_kv_cache_update()
|
|
kv_time_stats = kv_tracer.get_operation_time_stats()
|
|
if "attn_kv_cache_save" not in kv_time_stats:
|
|
raise RuntimeError("missing KERNEL_ONLY KV-update stats")
|
|
kv_stats = {
|
|
f"{name}_ms": float(value) / config.num_layers
|
|
for name, value in kv_time_stats["attn_kv_cache_save"].items()
|
|
}
|
|
|
|
result = BenchmarkResult(
|
|
config=config,
|
|
mean_time=core["mean"] / 1000.0,
|
|
std_time=core["std"] / 1000.0,
|
|
min_time=core["min"] / 1000.0,
|
|
max_time=core["max"] / 1000.0,
|
|
throughput_tokens_per_sec=(
|
|
total_q / (core["mean"] / 1000.0) if core["mean"] > 0 else 0.0
|
|
),
|
|
)
|
|
return result, kv_stats
|
|
|
|
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,
|
|
)
|
|
if args.profile_method == "record_function":
|
|
result, kv_stats = profile_kernel_only(config)
|
|
else:
|
|
result = runner.run_attention_benchmark(config)
|
|
kv_stats = (
|
|
profile_kv_cache_update(config) if args.profile_kv_update else None
|
|
)
|
|
row = result.to_dict()
|
|
row["tensor_parallel_size"] = tp
|
|
row["attention_core_excludes_kv_cache_update"] = True
|
|
if kv_stats is not None:
|
|
row["kv_cache_update_time"] = kv_stats
|
|
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,
|
|
"profile_kv_update": args.profile_kv_update,
|
|
"profile_method": args.profile_method,
|
|
},
|
|
"rows": rows,
|
|
}
|
|
args.output.write_text(
|
|
json.dumps(payload, indent=2, sort_keys=True, default=json_default) + "\n"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|