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