#!/usr/bin/env python3
"""Run Frontier's linear profiler against the vLLM 0.20 RoPE API.

Frontier passes the pre-v0.20 ``get_rope`` arguments separately, while vLLM
0.20 carries the same values in ``rope_parameters``.  Keep both repositories
unchanged and adapt only that call boundary for the profiling experiment.
"""

from __future__ import annotations

import inspect
import os
from typing import Any

import torch
from vllm import _custom_ops as vllm_ops
from vllm.model_executor.layers.layernorm import (
    GemmaRMSNorm as VllmGemmaRMSNorm,
    fused_add_rms_norm as vllm_fused_add_rms_norm,
)
from vllm.model_executor.layers.rotary_embedding import get_rope as vllm_get_rope

from frontier.profiling.common.layers import layernorm as frontier_layernorm
from frontier.profiling.common.layers import rotary_embedding as frontier_rope


def _vllm020_get_rope_adapter(
    *,
    head_size: int,
    rotary_dim: int,
    max_position: int,
    base: int | float,
    is_neox_style: bool,
    rope_scaling: dict[str, Any] | None,
    dtype: torch.dtype | None = None,
) -> Any:
    rope_parameters = dict(rope_scaling or {})
    rope_parameters["rope_theta"] = base
    rope_parameters["rope_dim"] = rotary_dim
    return vllm_get_rope(
        head_size=head_size,
        max_position=max_position,
        is_neox_style=is_neox_style,
        rope_parameters=rope_parameters,
        dtype=dtype,
    )


def _vllm020_rms_norm_adapter(
    x: torch.Tensor,
    weight: torch.Tensor,
    variance_epsilon: float,
) -> torch.Tensor:
    output = torch.empty_like(x)
    vllm_ops.rms_norm(output, x, weight, variance_epsilon)
    return output


def main() -> None:
    parameters = inspect.signature(vllm_get_rope).parameters
    if "rope_parameters" not in parameters or "rotary_dim" in parameters:
        raise RuntimeError(
            "Expected the vLLM 0.20 get_rope API with rope_parameters; "
            f"found {inspect.signature(vllm_get_rope)}"
        )

    frontier_rope._VLLM_GET_ROPE = _vllm020_get_rope_adapter
    frontier_rope._VLLM_GET_ROPE_IMPORT_ERROR = None
    frontier_layernorm.HAS_VLLM_RMSNORM = True
    frontier_layernorm.VllmGemmaRMSNorm = VllmGemmaRMSNorm
    frontier_layernorm.vllm_rms_norm = _vllm020_rms_norm_adapter
    frontier_layernorm.vllm_fused_add_rms_norm = vllm_fused_add_rms_norm

    from frontier.profiling.linear_op.main import main as frontier_main
    from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config

    model_root = os.environ.get("MODEL_ROOT")
    if not model_root:
        raise RuntimeError("MODEL_ROOT must point to the profiled Qwen checkpoint")
    model_config = ModelConfig(
        model=model_root,
        dtype="bfloat16",
        max_model_len=8192,
        skip_tokenizer_init=True,
        generation_config="vllm",
    )
    with set_current_vllm_config(VllmConfig(model_config=model_config)):
        frontier_main()


if __name__ == "__main__":
    main()