# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the Transformers modeling backend's linear fusers."""

import inspect
from types import MethodType, SimpleNamespace

import pytest
import torch
import torch.nn as nn
import torch.nn.functional as F

from vllm.model_executor.models.transformers.fuser import get_fuser
from vllm.model_executor.models.transformers.fusers import GLUFuser, QKVFuser


class SiluAndMulStub(nn.Module):
    """Stand-in for vLLM's `SiluAndMul` (no vLLM config required)."""

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        d = x.shape[-1] // 2
        return F.silu(x[..., :d]) * x[..., d:]


class NoDownGLU(nn.Module):
    """`act(gate(x)) * up(x)` with no output projection -> `down_name` is None."""

    def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
        super().__init__()
        self.gate_proj = nn.Linear(hidden, inter, bias=bias)
        self.up_proj = nn.Linear(hidden, inter, bias=bias)
        self.act_fn = nn.SiLU()

    def forward(self, x):
        return self.act_fn(self.gate_proj(x)) * self.up_proj(x)


class GLUMLP(NoDownGLU):
    """`down(act(gate(x)) * up(x))` — the canonical HF GLU MLP."""

    def __init__(self, hidden: int = 16, inter: int = 32, bias: bool = False):
        super().__init__(hidden, inter, bias)
        self.down_proj = nn.Linear(inter, hidden, bias=bias)

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class ReversedGLUMLP(GLUMLP):
    """`up(x) * act(gate(x))` — operands swapped (multiply is commutative)."""

    def forward(self, x):
        return self.down_proj(self.up_proj(x) * self.act_fn(self.gate_proj(x)))


class NotAnMLP(nn.Module):
    """Two linears but no activation*linear multiply -> must not match."""

    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(8, 8)
        self.fc2 = nn.Linear(8, 8)

    def forward(self, x):
        return self.fc2(self.fc1(x))


class NotAnActGLUMLP(GLUMLP):
    """GLU-shaped, but the "activation" is not a known activation module."""

    def __init__(self):
        super().__init__()
        self.act_fn = nn.Dropout()


class UntraceableMLP(GLUMLP):
    """Data-dependent control flow *before* the GLU -> no match."""

    def forward(self, x):
        if x.sum() > 0:  # noqa: SIM108 - intentionally untraceable
            return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return x


class UntraceableTailGLUMLP(GLUMLP):
    """Data-dependent control flow *after* the GLU -> still fusable."""

    def forward(self, x):
        y = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        if y.sum() > torch.inf:  # intentionally untraceable
            y = y * 0
        return y


class FakeAttention(nn.Module):
    """HF v5-style attention: shape unpacking, dead KV branch, kwargs interface."""

    is_causal = True

    def __init__(
        self,
        hidden: int = 32,
        head_dim: int = 8,
        heads: int = 4,
        kv_heads: int = 4,
        bias: bool = False,
        layer_idx: int = 0,
    ):
        super().__init__()
        self.config = SimpleNamespace(_attn_implementation="vllm")
        self.layer_idx = layer_idx
        self.head_dim = head_dim
        self.scaling = head_dim**-0.5
        self.q_proj = nn.Linear(hidden, heads * head_dim, bias=bias)
        self.k_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
        self.v_proj = nn.Linear(hidden, kv_heads * head_dim, bias=bias)
        self.o_proj = nn.Linear(heads * head_dim, hidden, bias=bias)

    def forward(
        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
    ):
        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)
        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        if past_key_values is not None:
            k, v = past_key_values.update(k, v, self.layer_idx)
        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, None
        )
        attn_output, attn_weights = attention_interface(
            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
        )
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        return self.o_proj(attn_output), attn_weights


class ReversedFakeAttention(FakeAttention):
    """Projections computed in (v, k, q) order — q must still be identified."""

    def forward(
        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
    ):
        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS

        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)
        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, None
        )
        attn_output, _ = attention_interface(
            self, q, k, v, attention_mask, scaling=self.scaling, **kwargs
        )
        return self.o_proj(attn_output.reshape(*input_shape, -1)), None


class ExtraProjAttention(FakeAttention):
    """A second non-qkv linear of a different width -> `o_proj` still found."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.sink_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)


class QKNormAttention(FakeAttention):
    """OLMoE-style: a full-dim norm applied to the whole q/k projection output."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.q_norm = nn.RMSNorm(self.q_proj.out_features)
        self.k_norm = nn.RMSNorm(self.k_proj.out_features)

    def forward(
        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
    ):
        q = self.q_norm(self.q_proj(hidden_states))
        k = self.k_norm(self.k_proj(hidden_states))
        v = self.v_proj(hidden_states)
        return self.o_proj(q + k + v), None


class PerHeadQKNormAttention(FakeAttention):
    """Qwen3-style: a per-head norm (`head_dim`) applied after the head reshape."""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.q_norm = nn.RMSNorm(self.head_dim)
        self.k_norm = nn.RMSNorm(self.head_dim)

    def forward(
        self, hidden_states, attention_mask=None, past_key_values=None, **kwargs
    ):
        shape = (*hidden_states.shape[:-1], -1, self.head_dim)
        q = self.q_norm(self.q_proj(hidden_states).view(shape))
        k = self.k_norm(self.k_proj(hidden_states).view(shape))
        v = self.v_proj(hidden_states).view(shape)
        return self.o_proj((q + k + v).flatten(-2)), None


class FakeSelfAttn(nn.Module):
    """Stand-in for the vLLM `Attention` looked up in `attention_instances`."""

    def __init__(self):
        super().__init__()
        self.impl = SimpleNamespace(scale=None)

    def forward(self, q, k, v):
        # MHA-shaped stub: any deterministic combination of q/k/v will do
        return q + 2 * k + 3 * v


@pytest.fixture(autouse=True)
def _clear_fuser_cache():
    get_fuser.cache_clear()
    yield
    get_fuser.cache_clear()


def _apply_glu_fuser_with_stubs(module: nn.Module, fuser: GLUFuser):
    """Apply a fuser using plain stand-ins (merged `nn.Linear` + silu AndMul)."""
    gate = module.get_submodule(fuser.gate_name)
    up = module.get_submodule(fuser.up_name)
    merged = nn.Linear(
        gate.in_features,
        gate.out_features + up.out_features,
        bias=gate.bias is not None,
    )
    with torch.no_grad():
        merged.weight.copy_(torch.cat([gate.weight, up.weight], dim=0))
        if gate.bias is not None:
            merged.bias.copy_(torch.cat([gate.bias, up.bias], dim=0))
    setattr(module, fuser.merged_name, merged)
    setattr(module, fuser.act_name, SiluAndMulStub())
    delattr(module, fuser.gate_name)
    delattr(module, fuser.up_name)
    module.forward = MethodType(fuser.fused_forward, module)
    return module


def _apply_qkv_fuser_with_stubs(module: nn.Module, fuser: QKVFuser):
    """Apply a fuser using a plain merged `nn.Linear` (no TP sharding)."""
    q, k, v = (
        module.get_submodule(name)
        for name in (fuser.q_name, fuser.k_name, fuser.v_name)
    )
    merged = nn.Linear(
        q.in_features,
        q.out_features + k.out_features + v.out_features,
        bias=q.bias is not None,
    )
    with torch.no_grad():
        merged.weight.copy_(torch.cat([q.weight, k.weight, v.weight], dim=0))
        if q.bias is not None:
            merged.bias.copy_(torch.cat([q.bias, k.bias, v.bias], dim=0))
    merged.output_sizes = [q.out_features, k.out_features, v.out_features]
    merged.tp_size = 1
    setattr(module, fuser.merged_name, merged)
    for name in (fuser.q_name, fuser.k_name, fuser.v_name):
        delattr(module, name)
    module.forward = MethodType(fuser.fused_forward, module)
    return module


@pytest.mark.parametrize("mlp_cls", [GLUMLP, ReversedGLUMLP])
@pytest.mark.parametrize("bias", [False, True])
def test_detects_and_rewrites_glu(mlp_cls, bias):
    with torch.device("meta"):
        meta = mlp_cls(bias=bias)
    fuser = get_fuser(meta)
    assert isinstance(fuser, GLUFuser)
    assert (
        fuser.gate_name,
        fuser.up_name,
        fuser.act_name,
        fuser.down_name,
    ) == ("gate_proj", "up_proj", "act_fn", "down_proj")

    # The rewritten forward references the merged projection instead of the
    # sources; the rest of the forward is untouched.
    names = fuser.fused_forward.__code__.co_names
    assert "gate_up_proj" in names and "act_fn" in names and "down_proj" in names
    assert not {"gate_proj", "up_proj"} & set(names)

    # Numerics: the fused forward must match the original on a real instance.
    real = mlp_cls(bias=bias)
    for p in real.parameters():
        nn.init.normal_(p, std=0.05)
    x = torch.randn(4, 16)
    expected = real(x)
    fused = _apply_glu_fuser_with_stubs(real, fuser)

    # Fusion is in place: the module keeps its class and other attributes
    assert fused is real and type(fused) is mlp_cls
    torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)


def test_glu_identifies_down_projection():
    """The row projection consuming `act(gate(x)) * up(x)` is identified.

    It is forced to `RowParallelLinear` in `update_attrs` so its sharded input
    matches the column-parallel merged gate/up; `None` when there is no such
    projection to force (fusion of gate/up still applies)."""
    with torch.device("meta"):
        assert get_fuser(GLUMLP()).down_name == "down_proj"
        assert get_fuser(ReversedGLUMLP()).down_name == "down_proj"
        assert get_fuser(NoDownGLU()).down_name is None


@pytest.mark.parametrize("attn_cls", [FakeAttention, ReversedFakeAttention])
@pytest.mark.parametrize("kv_heads", [4, 2])
def test_detects_and_rewrites_qkv(attn_cls, kv_heads):
    if attn_cls is ReversedFakeAttention and kv_heads == 4:
        pytest.skip("MHA q/k/v assignment is order-based by design")
    with torch.device("meta"):
        meta = attn_cls(kv_heads=kv_heads)
    fuser = get_fuser(meta)
    assert isinstance(fuser, QKVFuser)
    # q (sharded differently under TP) must be identified exactly; k/v may be
    # swapped for non-canonical compute order, which is numerically consistent
    # because the weight mapping and the split indices follow the same
    # assignment.
    assert fuser.q_name == "q_proj"
    assert {fuser.k_name, fuser.v_name} == {"k_proj", "v_proj"}
    assert fuser.o_name == "o_proj"

    # The projections are merged; everything else stays live Python with its
    # original semantics (branches, kwargs, attribute reads)
    code = fuser.fused_forward.__code__
    names = code.co_names
    assert "qkv_proj" in names and "output_sizes" in names and "o_proj" in names
    assert "tp_size" in names
    assert not {"q_proj", "k_proj", "v_proj"} & set(names)
    if attn_cls is FakeAttention:
        assert "update" in names  # the cache branch survives
    assert code.co_flags & inspect.CO_VARKEYWORDS  # **kwargs survives

    # Numerics: the fused forward must match the original on a real instance,
    # with a different layer_idx than the traced instance (kv_heads == heads so
    # the q/k/v stub combination is shape-compatible).
    real = attn_cls(kv_heads=4, layer_idx=3)
    for p in real.parameters():
        nn.init.normal_(p, std=0.05)
    x = torch.randn(1, 5, 32)
    attention_instances = {3: FakeSelfAttn()}
    expected, _ = real(x, attention_instances=attention_instances)
    fused = _apply_qkv_fuser_with_stubs(real, fuser)

    # Fusion is in place: the module keeps its class and other attributes
    assert fused is real and type(fused) is attn_cls
    assert fused.layer_idx == 3 and fused.is_causal and fused.config is not None
    out, _ = fused(x, attention_instances=attention_instances)
    torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5)


def test_qkv_identifies_output_projection():
    with torch.device("meta"):
        assert get_fuser(FakeAttention()).o_name == "o_proj"
        assert get_fuser(ReversedFakeAttention()).o_name == "o_proj"
        assert get_fuser(ExtraProjAttention()).o_name == "o_proj"
        # Norm children (q_norm/k_norm) must not disturb o_proj identification.
        assert get_fuser(QKNormAttention()).o_name == "o_proj"
        assert get_fuser(PerHeadQKNormAttention()).o_name == "o_proj"


def test_fuser_is_cached_per_class():
    with torch.device("meta"):
        fuser_a = get_fuser(GLUMLP())
        fuser_b = get_fuser(GLUMLP())
    assert fuser_a is fuser_b
    assert GLUMLP in get_fuser.cache


@pytest.mark.parametrize("cls", [NotAnMLP, UntraceableMLP])
def test_non_matching_modules_return_none(cls):
    with torch.device("meta"):
        module = cls()
    assert get_fuser(module) is None


def test_untraceable_tail_still_fuses():
    with torch.device("meta"):
        meta = UntraceableTailGLUMLP()
    fuser = get_fuser(meta)
    assert isinstance(fuser, GLUFuser)

    # Numerics: the live tail must survive the rewrite
    real = UntraceableTailGLUMLP()
    for p in real.parameters():
        nn.init.normal_(p, std=0.05)
    x = torch.randn(4, 16)
    expected = real(x)
    fused = _apply_glu_fuser_with_stubs(real, fuser)
    torch.testing.assert_close(fused(x), expected, atol=1e-5, rtol=1e-5)


def test_weight_mappings_are_scoped_to_fused_prefixes():
    from vllm.model_executor.models.utils import WeightsMapper

    with torch.device("meta"):
        glu_fuser = get_fuser(GLUMLP())
        qkv_fuser = get_fuser(FakeAttention())

    mapper = WeightsMapper()
    for prefix in ("model.layers.0.mlp", "model.layers.1.mlp"):
        mapper.orig_to_new_stacked.update(glu_fuser.orig_to_new_stacked(prefix))
    mapper.orig_to_new_stacked.update(
        qkv_fuser.orig_to_new_stacked("model.layers.0.self_attn")
    )

    names = [
        "model.layers.0.mlp.gate_proj.weight",
        "model.layers.0.mlp.up_proj.weight",
        "model.layers.1.mlp.gate_proj.weight",
        "model.layers.0.self_attn.q_proj.weight",
        "model.layers.0.self_attn.k_proj.weight",
        "model.layers.0.self_attn.v_proj.weight",
        # Unfused modules at other prefixes must be left untouched.
        "model.layers.2.mlp.experts.0.gate_proj.weight",
        "model.layers.1.self_attn.q_proj.weight",
    ]
    # `apply` rewrites the name and stamps the shard id onto each tensor.
    weights = [(name, torch.empty(0)) for name in names]
    mapped = list(mapper.apply(weights))
    mapped_names = [name for name, _ in mapped]
    shard_ids = [getattr(data, "shard_id", None) for _, data in mapped]

    assert mapped_names == [
        "model.layers.0.mlp.gate_up_proj.weight",
        "model.layers.0.mlp.gate_up_proj.weight",
        "model.layers.1.mlp.gate_up_proj.weight",
        "model.layers.0.self_attn.qkv_proj.weight",
        "model.layers.0.self_attn.qkv_proj.weight",
        "model.layers.0.self_attn.qkv_proj.weight",
        # Only the exact fused layers are remapped; everything else is untouched.
        "model.layers.2.mlp.experts.0.gate_proj.weight",
        "model.layers.1.self_attn.q_proj.weight",
    ]
    assert shard_ids == [0, 1, 0, "q", "k", "v", None, None]

    # The fused layers are exposed to the quantization machinery via their
    # original constituent projection names (what the checkpoint stores).
    assert glu_fuser.packed_modules_mapping == {
        "gate_up_proj": ["gate_proj", "up_proj"],
    }
    assert qkv_fuser.packed_modules_mapping == {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
    }


@pytest.mark.parametrize("cls", [NotAnMLP, NotAnActGLUMLP])
def test_unfusable_modules_are_not_fused(cls, default_vllm_config):
    with torch.device("meta"):
        module = cls()
    fuser = get_fuser(module)
    # Either no pattern matches the class, or this instance fails validation
    # (`recursive_replace` gates fusion and its weight mappings on `validate`)
    model_config = default_vllm_config.model_config
    assert fuser is None or not fuser.validate(module, model_config)


def test_act_and_mul_derived_from_module(default_vllm_config):
    from transformers.activations import GELUTanh, SiLUActivation

    from vllm.model_executor.layers.activation import GeluAndMul, SiluAndMul

    assert isinstance(GLUFuser._get_act_and_mul(nn.SiLU()), SiluAndMul)
    assert isinstance(GLUFuser._get_act_and_mul(SiLUActivation()), SiluAndMul)
    gelu_tanh = GLUFuser._get_act_and_mul(GELUTanh())
    assert isinstance(gelu_tanh, GeluAndMul) and gelu_tanh.approximate == "tanh"
    gelu = GLUFuser._get_act_and_mul(nn.GELU())
    assert isinstance(gelu, GeluAndMul) and gelu.approximate == "none"
    # Not activations at all -> no fusion
    assert GLUFuser._get_act_and_mul_name(nn.Dropout()) is None
    assert GLUFuser._get_act_and_mul_name(nn.LayerNorm(8)) is None
    with pytest.raises(ValueError, match="No AndMul equivalent"):
        GLUFuser._get_act_and_mul(nn.Dropout())