aituner/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_flashattn.py

#!/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()