# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

import itertools
import logging
from dataclasses import dataclass
from typing import Any

import torch
from torch.nn.attention import (
    activate_flash_attention_impl,
    current_flash_attention_impl,
)
from torch.nn.attention.varlen import AuxRequest
from torchtitan.distributed.batch_invariant import is_in_batch_invariant_mode
from torchtitan.models.common.attention import InnerAttention
from torchtitan.observability.logging import warn_once
from torchtitan.protocols.module import Module
from torchtitan.tools.utils import get_cuda_flash_attention_impl
from vllm.model_executor.layers.attention import Attention
from vllm.model_executor.layers.attention.attention import get_attention_context
from vllm.v1.attention.backend import AttentionCGSupport, AttentionType
from vllm.v1.attention.backends.flash_attn import (
    FlashAttentionBackend,
    FlashAttentionImpl,
    FlashAttentionMetadata,
    FlashAttentionMetadataBuilder,
)
from vllm.v1.attention.backends.flash_attn_diffkv import (
    FlashAttentionDiffKVBackend,
    FlashAttentionDiffKVImpl,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum, register_backend

logger = logging.getLogger(__name__)


@register_backend(AttentionBackendEnum.CUSTOM)
class TorchTitanVarlenInnerAttentionBackend(FlashAttentionBackend):
    """Custom vLLM attention backend using PyTorch's native FlashAttention kernel.

    This class is not directly referenced in user code. It is registered into
    vLLM's attention backend registry via the ``@register_backend`` decorator
    and selected at runtime when the vLLM engine is configured to use a CUSTOM
    attention backend.

    Inheriting from ``FlashAttentionBackend`` is not strictly required for all
    backends, but it is convenient here to reuse metadata construction logic.
    """

    @staticmethod
    def get_name():
        # vLLM requires any custom attention backend to return "CUSTOM" as its
        # name so the backend registry can look it up correctly.
        return "CUSTOM"

    @staticmethod
    def get_impl_cls():
        return TorchTitanVarlenInnerAttentionImpl

    @staticmethod
    def get_builder_cls():
        class TorchTitanVarlenInnerAttentionMetadataBuilder(
            FlashAttentionMetadataBuilder
        ):
            _cudagraph_support = AttentionCGSupport.ALWAYS

        return TorchTitanVarlenInnerAttentionMetadataBuilder


class TorchTitanVarlenInnerAttentionDiffKVBackend(FlashAttentionDiffKVBackend):
    """PyTorch varlen backend using vLLM's unequal K/V cache support."""

    @staticmethod
    def get_name():
        return "CUSTOM"

    @staticmethod
    def get_impl_cls():
        return TorchTitanVarlenInnerAttentionDiffKVImpl

    @staticmethod
    def get_builder_cls():
        return TorchTitanVarlenInnerAttentionBackend.get_builder_cls()


class TorchTitanVarlenInnerAttentionImpl(FlashAttentionImpl):
    """
    Custom vLLM attention backend impl using PyTorch's native FlashAttention varlen API.
    Instead of using vLLM's FlashAttention kernel, this implementation takes the kernel
    dependency from torch directly while supporting the same interface.
    """

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)

        # optional post-attention epilogue transform (out, lse) -> out
        # set by VLLMModelWrapper in vllm_wrapper.py
        self.out_transform = None

        self.enable_gqa = self.num_heads > self.num_kv_heads

        flash_attention_impl = get_cuda_flash_attention_impl()
        if flash_attention_impl is not None:
            # activate_flash_attention_impl() will restore internal global state
            # and re-run register function, so we want to only call it once.
            if current_flash_attention_impl() != flash_attention_impl:
                try:
                    activate_flash_attention_impl(flash_attention_impl)
                except (ImportError, RuntimeError, ValueError) as error:
                    capability = torch.cuda.get_device_capability()
                    raise RuntimeError(
                        f"{flash_attention_impl} is required on detected SM "
                        f"{capability[0]}.{capability[1]}, but activation failed."
                    ) from error
        else:
            warn_once(
                logger,
                "FA3/FA4 not available on this CUDA architecture, falling back to FA2. ",
            )

    # Based on vLLM's FlashAttentionImpl.forward():
    # https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/flash_attn.py
    def forward(
        self,
        layer: torch.nn.Module,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: FlashAttentionMetadata,
        output: torch.Tensor | None = None,
        output_scale: torch.Tensor | None = None,
        output_block_scale: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Forward pass with FlashAttention.

        Args:
            query: shape = [num_tokens, num_heads, head_size]
            key: shape = [num_tokens, num_kv_heads, head_size]
            value: shape = [num_tokens, num_kv_heads, head_size]
            kv_cache: shape is ``[num_blocks, num_kv_heads, block_size,
                2 * head_size]``, with K and V packed in the last dimension.
            attn_metadata: Metadata for attention.
        Returns:
            shape = [num_tokens, num_heads * head_size]
        """
        assert output is not None, "Output tensor must be provided."
        assert (
            self.vllm_flash_attn_version is not None
        ), "FlashAttention version not detected."

        if output_scale is not None or output_block_scale is not None:
            raise NotImplementedError(
                "fused output quantization is not yet supported for FlashAttentionImpl"
            )

        # Breakable cuda_graph: under VLLM_USE_BREAKABLE_CUDAGRAPH the recorded
        # forward closure pins capture-time args (attn_metadata None, varlen metadata
        # absent). Re-read live per-layer metadata + kv_cache from the forward context
        # (vLLM refreshes them before each replay) BEFORE the None check, else replay
        # short-circuits to output.fill_(0) (zeroed attention). No-op outside capture.
        attn_metadata, _, kv_cache, _ = get_attention_context(layer.layer_name)

        if attn_metadata is None:
            # Profiling / CUDA graph dummy-capture run (no real metadata yet).
            return output.fill_(0)

        attn_type = self.attn_type

        # IMPORTANT!
        # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
        # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
        # in this method. For example, `view` and `slice` (or `[:n]`) operations
        # are surprisingly slow even in the case they do not invoke any GPU ops.
        # Minimize the PyTorch ops in this method as much as possible.
        # Whenever making a change in this method, please benchmark the
        # performance to make sure it does not introduce any overhead.

        num_actual_tokens = attn_metadata.num_actual_tokens

        assert attn_type not in (
            AttentionType.ENCODER_ONLY,
            AttentionType.ENCODER,
        ), "Encoder-only attention not supported yet."

        # vLLM #44455 packs K and V into the content dimension for full-attention
        # layers, including those in hybrid models. Restore the layout expected by
        # varlen attention before splitting them.
        key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)

        assert not self.kv_cache_dtype.startswith(
            "fp8"
        ), "FP8 KV cache not supported yet."

        assert not attn_metadata.use_cascade, "Cascade not supported yet."

        cu_seqlens_q = attn_metadata.query_start_loc
        seqused_k = attn_metadata.seq_lens
        max_seqlen_q = attn_metadata.max_query_len
        max_seqlen_k = attn_metadata.max_seq_len
        block_table = attn_metadata.block_table

        assert self.dcp_world_size == 1, "DCP not supported yet."

        if not attn_metadata.causal:
            raise RuntimeError("Non-causal attention not supported yet.")

        # vLLM assigns sliding_window_size = None to (-1, -1) w/ optional causal flag
        # but varlen only encode with sliding window. so we need to convert vllm (-1, -1) to (-1, 0)
        # for proper full causal attention instead of bidirectional
        if self.sliding_window == (-1, -1):
            sliding_window_size = (-1, 0)
        else:
            # by default vLLM sets attention type = DECODER, which will set (W-1, 0)
            sliding_window_size = self.sliding_window

        assert self.alibi_slopes is None, "Alibi slopes not supported yet."

        # FA3 and FA4 infer key lengths from block_table + seqused_k for paged
        # KV. FA4 explicitly rejects page_table together with cu_seqlens_k.
        # FA2 requires cu_seqlens_k to be explicitly set.
        fa_impl = current_flash_attention_impl()
        if fa_impl in ("FA3", "FA4"):
            cu_seqlens_k = None
        else:
            num_seqs = seqused_k.shape[0]
            cu_seqlens_k = torch.zeros(
                num_seqs + 1, dtype=torch.int32, device=query.device
            )
            cu_seqlens_k[1:] = torch.cumsum(seqused_k, dim=0)
        extra_kwargs: dict[str, Any] = {}

        # TODO(pytorch/pytorch#179760): FA2's auto num_splits heuristic
        # produces NaN intermittently with paged KV (block_table). Force
        # num_splits=1 as a workaround until the root cause is fixed
        # upstream. current_flash_attention_impl() returns None when FA2
        # is the implicit default (SM < 9.0). For FA3, only force
        # num_splits=1 in batch-invariant mode (determinism).
        if fa_impl in (None, "FA2") or is_in_batch_invariant_mode():
            extra_kwargs["num_splits"] = 1

        if self.enable_gqa:
            extra_kwargs["enable_gqa"] = True

        if self.out_transform is not None:
            extra_kwargs["return_aux"] = AuxRequest(lse=True)

        result = torch.nn.attention.varlen.varlen_attn_out(
            output[:num_actual_tokens],
            query[:num_actual_tokens],
            key_cache,
            value_cache,
            cu_seqlens_q,
            cu_seqlens_k,
            max_seqlen_q,
            max_seqlen_k,
            scale=self.scale,
            window_size=sliding_window_size,
            block_table=block_table,
            seqused_k=seqused_k,
            **extra_kwargs,
        )
        if self.out_transform is None:
            return result

        out, lse = result
        out = self.out_transform(out, lse.transpose(0, 1))
        output[:num_actual_tokens].copy_(out)
        return output[:num_actual_tokens]


class TorchTitanVarlenInnerAttentionDiffKVImpl(TorchTitanVarlenInnerAttentionImpl):
    """PyTorch varlen implementation using vLLM's unequal K/V cache writer."""

    def do_kv_cache_update(
        self,
        layer: torch.nn.Module,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        slot_mapping: torch.Tensor,
    ) -> None:
        return FlashAttentionDiffKVImpl.do_kv_cache_update(
            self, layer, key, value, kv_cache, slot_mapping
        )


class VLLMAttentionWrapper(Module):
    """Adapter from TorchTitan tensor layout to ``vllm.Attention``.

    vLLM's ``Attention`` layer manages KV-cache and paged attention internally,
    and expects ``(num_tokens, num_heads, head_dim)`` inputs.

    Used as ``inner_attention`` in GQAttention via Config-based construction.
    """

    # vLLM requires a unique prefix per Attention layer for
    # static_forward_context registration.
    # TODO: Pass layer_id through the build chain instead of using a
    # global counter. The counter breaks with pipeline parallelism
    # where layers are built on different ranks.
    _layer_counter: itertools.count = itertools.count()
    _module_protocol_exempt_children = frozenset({"vllm_attn"})

    @dataclass(kw_only=True, slots=True)
    class Config(Module.Config):
        attention_metadata_key: type[InnerAttention]
        hidden_size: int
        num_heads: int
        num_kv_heads: int
        head_dim: int
        value_head_dim: int | None = None
        scale: float | None = None
        sliding_window_size: int | None = None
        """Causal sliding-window size (``None`` => full attention)."""

    def __init__(self, config: Config) -> None:
        super().__init__()
        self.attention_metadata_key = config.attention_metadata_key

        from vllm.config import get_current_vllm_config

        vllm_config = get_current_vllm_config()
        tp_degree = vllm_config.parallel_config.tensor_parallel_size

        num_heads = config.num_heads
        num_kv_heads = config.num_kv_heads

        if num_kv_heads < tp_degree:
            raise ValueError(
                f"num_kv_heads ({num_kv_heads}) must be >= "
                f"tensor_parallel_size ({tp_degree})"
            )
        if num_kv_heads % tp_degree != 0:
            raise ValueError(
                f"num_kv_heads ({num_kv_heads}) must be divisible by "
                f"tensor_parallel_size ({tp_degree})"
            )
        if num_heads % tp_degree != 0:
            raise ValueError(
                f"num_heads ({num_heads}) must be divisible by "
                f"tensor_parallel_size ({tp_degree})"
            )

        num_heads = num_heads // tp_degree
        num_kv_heads = num_kv_heads // tp_degree
        head_dim = config.head_dim
        value_head_dim = (
            config.value_head_dim if config.value_head_dim is not None else head_dim
        )
        scale = config.scale if config.scale is not None else head_dim**-0.5

        self.hidden_size = config.hidden_size
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.head_dim = head_dim
        self.value_head_dim = value_head_dim
        self.scale = scale

        cache_config = (
            vllm_config.cache_config if hasattr(vllm_config, "cache_config") else None
        )

        # TODO: This need to be compatible with Pipeline Parallelism
        layer_id = next(VLLMAttentionWrapper._layer_counter)
        diff_kv_kwargs: dict[str, Any] = {}
        if value_head_dim != head_dim:
            FlashAttentionDiffKVBackend.set_head_size_v(value_head_dim)
            diff_kv_kwargs = {
                "head_size_v": value_head_dim,
                "attn_backend": TorchTitanVarlenInnerAttentionDiffKVBackend,
            }
        self.vllm_attn = Attention(
            num_heads=num_heads,
            head_size=head_dim,
            scale=scale,
            num_kv_heads=num_kv_heads,
            cache_config=cache_config,
            quant_config=None,
            per_layer_sliding_window=config.sliding_window_size,
            prefix=f"model.layers.{layer_id}.attention.inner_attention",
            **diff_kv_kwargs,
        )

    def forward(
        self,
        q_THK: torch.Tensor,
        k_THK: torch.Tensor,
        v_THV: torch.Tensor,
        *,
        attention_metadata: None = None,
        **kwargs,
    ) -> torch.Tensor:
        """Run vLLM paged attention on local (non-DTensor) tensors.

        Args:
            q_THK: ``(num_tokens, num_heads, query/key_head_dim)``
            k_THK: ``(num_tokens, num_kv_heads, query/key_head_dim)``
            v_THV: ``(num_tokens, num_kv_heads, value_head_dim)``

        Returns:
            ``(num_tokens, num_heads, value_head_dim)``.
        """
        if attention_metadata is not None:
            raise ValueError(
                "VLLMAttentionWrapper does not support attention_metadata; vLLM "
                "manages causal masking and the KV-cache internally."
            )

        if self.value_head_dim != self.head_dim:
            # V may be a strided view (e.g. MLA splits it from a fused per-head
            # KV projection), but vLLM's cache writer assumes V heads are packed
            # with stride value_head_dim, so repack before updating the cache.
            v_THV = v_THV.contiguous()
        out_TD = self.vllm_attn(q_THK, k_THK, v_THV)

        # vLLM's flash attention backend may pad the token count (e.g.
        # round up to an even number), which introduces a new symbolic
        # shape under torch.compile.  Narrow to trim this padding.
        num_tokens = q_THK.shape[0]
        out_TD = out_TD.narrow(0, 0, num_tokens)
        return out_TD.view(num_tokens, -1, self.value_head_dim)


def get_attention_dimensions(
    attention_config, model_dim: int
) -> tuple[int, int, int, int]:
    """Return query heads, KV heads, Q/K head dim, and value head dim."""
    num_heads = attention_config.n_heads
    num_kv_heads = getattr(attention_config, "n_kv_heads", None) or num_heads
    if hasattr(attention_config, "qk_nope_head_dim"):
        head_dim = attention_config.qk_nope_head_dim + attention_config.qk_rope_head_dim
        value_head_dim = attention_config.v_head_dim
    else:
        head_dim = getattr(attention_config, "head_dim", None)
        head_dim = head_dim if head_dim is not None else model_dim // num_heads
        value_head_dim = head_dim
    return num_heads, num_kv_heads, head_dim, value_head_dim
