# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Attention layer with Dual chunk flash attention and sparse attention.
"""
import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type

import torch
import torch.distributed
import torch.nn.functional as F

from vllm import _custom_ops as ops
from vllm.attention.backends.abstract import AttentionLayer, AttentionType
from vllm.attention.backends.flash_attn import (FlashAttentionBackend,
                                                FlashAttentionImpl,
                                                FlashAttentionMetadata,
                                                FlashAttentionMetadataBuilder)
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
from vllm.logger import init_logger
from vllm.utils import async_tensor_h2d
from vllm.platforms import current_platform
if not current_platform.is_rocm():
    from vllm.vllm_flash_attn import (flash_attn_varlen_func,
                                    flash_attn_with_kvcache, sparse_attn_func)
else:
    from flash_attn import (flash_attn_varlen_func,
                        flash_attn_with_kvcache, sparse_attn_func)

if TYPE_CHECKING:
    from vllm.worker.model_runner import ModelInputForGPUBuilder

logger = init_logger(__name__)


class DualChunkFlashAttentionBackend(FlashAttentionBackend):

    accept_output_buffer: bool = False

    @staticmethod
    def get_name() -> str:
        return "DUAL_CHUNK_FLASH_ATTN"

    @staticmethod
    def get_impl_cls() -> Type["DualChunkFlashAttentionImpl"]:
        return DualChunkFlashAttentionImpl

    @staticmethod
    def get_metadata_cls() -> Type["DualChunkFlashAttentionMetadata"]:
        return DualChunkFlashAttentionMetadata

    @staticmethod
    def get_builder_cls() -> Type["DualChunkFlashAttentionMetadataBuilder"]:
        return DualChunkFlashAttentionMetadataBuilder


@dataclass
class DualChunkFlashAttentionMetadata(FlashAttentionMetadata):
    # Block size of the paged kv cache.
    block_size: int = 16

    # Original max position embeddings.
    original_max_position_embeddings: int = 0

    # Chunk size
    chunk_size: int = 8192

    # Local size
    local_size: int = 1024

    # (batch_size,). The orig sequence length per sequence.
    orig_seq_lens: Optional[List[int]] = None

    # orig_seq_lens stored as a tensor.
    orig_seq_lens_tensor: Optional[torch.Tensor] = None

    # Length scaling factor
    scaling_factor: Optional[torch.Tensor] = None

    # (batch_size,). Sequence lengths for intra attention.
    seq_lens_intra: Optional[torch.Tensor] = None

    # Max sequence length for intra attention.
    max_seq_len_intra: Optional[int] = None

    # (batch_size, num_blocks). Block table for intra attention.
    block_tables_intra: Optional[torch.Tensor] = None

    # (batch_size,). Sequence lengths for succ attention.
    seq_lens_succ: Optional[torch.Tensor] = None

    # Max sequence length for succ attention.
    max_seq_len_succ: Optional[int] = None

    # (batch_size, num_blocks). Block table for succ attention.
    block_tables_succ: Optional[torch.Tensor] = None

    # (batch_size,). Sequence lengths for inter attention.
    seq_lens_inter: Optional[torch.Tensor] = None

    # Max sequence length for inter attention.
    max_seq_len_inter: Optional[int] = None

    _cached_prefill_metadata: Optional[
        "DualChunkFlashAttentionMetadata"] = None
    _cached_decode_metadata: Optional["DualChunkFlashAttentionMetadata"] = None

    @property
    def prefill_metadata(self) -> Optional["DualChunkFlashAttentionMetadata"]:
        if self.num_prefills == 0:
            return None

        if self._cached_prefill_metadata is not None:
            return self._cached_prefill_metadata

        prefill_metadata = super().prefill_metadata
        if prefill_metadata is None:
            return None

        prefill_metadata = DualChunkFlashAttentionMetadata(
            **prefill_metadata.asdict_zerocopy())

        prefill_metadata.orig_seq_lens = (
            None if self.orig_seq_lens is None else
            self.orig_seq_lens[:self.num_prefills])
        prefill_metadata.orig_seq_lens_tensor = (
            None if self.orig_seq_lens_tensor is None else
            self.orig_seq_lens_tensor[:self.num_prefills])

        if self.original_max_position_embeddings > 0:
            assert prefill_metadata.orig_seq_lens_tensor is not None
            prefill_metadata.scaling_factor = (
                0.1 * torch.log(prefill_metadata.orig_seq_lens_tensor /
                                self.original_max_position_embeddings) +
                1.0).clip(min=1)

        self._cached_prefill_metadata = prefill_metadata
        return prefill_metadata

    @property
    def decode_metadata(self) -> Optional["DualChunkFlashAttentionMetadata"]:
        if self.num_decode_tokens == 0:
            return None

        if self._cached_decode_metadata is not None:
            return self._cached_decode_metadata

        decode_metadata = super().decode_metadata
        if decode_metadata is None:
            return None

        decode_metadata = DualChunkFlashAttentionMetadata(
            **decode_metadata.asdict_zerocopy())

        decode_metadata.orig_seq_lens_tensor = (
            None if self.orig_seq_lens_tensor is None else
            self.orig_seq_lens_tensor[self.num_prefills:])

        assert decode_metadata.orig_seq_lens_tensor is not None
        assert decode_metadata.block_tables is not None

        cache_seq_lens = decode_metadata.orig_seq_lens_tensor
        chunk_len = self.chunk_size - self.local_size
        chunk_num_curr = (cache_seq_lens - 1) // chunk_len
        batch_size = decode_metadata.num_decode_tokens

        if self.original_max_position_embeddings > 0:
            decode_metadata.scaling_factor = (0.1 * torch.log(
                cache_seq_lens / self.original_max_position_embeddings) +
                                              1.0).clip(min=1)

        seq_lens_intra = cache_seq_lens - chunk_num_curr * chunk_len
        max_seq_len_intra = seq_lens_intra.max().item()
        decode_metadata.seq_lens_intra = seq_lens_intra
        decode_metadata.max_seq_len_intra = max_seq_len_intra

        block_tables_intra = torch.zeros(
            batch_size,
            (max_seq_len_intra - 1) // self.block_size + 1,
            dtype=decode_metadata.block_tables.dtype,
            device=decode_metadata.block_tables.device,
        )
        for i in range(batch_size):
            st = chunk_num_curr[i] * chunk_len // self.block_size
            ed = min(
                st + (max_seq_len_intra - 1) // self.block_size + 1,
                (cache_seq_lens[i] - 1) // self.block_size + 1,
            )
            block_tables_intra[i, :ed -
                               st] = decode_metadata.block_tables[i, st:ed]
        decode_metadata.block_tables_intra = block_tables_intra

        seq_lens_succ = (chunk_num_curr -
                         (chunk_num_curr - 1).clip(min=0)) * chunk_len
        max_seq_len_succ = seq_lens_succ.max().item()
        decode_metadata.seq_lens_succ = seq_lens_succ
        decode_metadata.max_seq_len_succ = max_seq_len_succ
        if max_seq_len_succ:
            block_tables_succ = torch.zeros(
                batch_size,
                (max_seq_len_succ - 1) // self.block_size + 1,
                dtype=decode_metadata.block_tables.dtype,
                device=decode_metadata.block_tables.device,
            )
            for i in range(batch_size):
                start = ((chunk_num_curr[i] - 1).clip(min=0) * chunk_len //
                         self.block_size)
                end = min(
                    start + (max_seq_len_succ - 1) // self.block_size + 1,
                    (cache_seq_lens[i] - 1) // self.block_size + 1,
                )
                block_tables_succ[
                    i, :end - start] = decode_metadata.block_tables[i,
                                                                    start:end]
            decode_metadata.block_tables_succ = block_tables_succ

        seq_lens_inter = (chunk_num_curr - 1).clip(min=0) * chunk_len
        max_seq_len_inter = seq_lens_inter.max().item()
        decode_metadata.seq_lens_inter = seq_lens_inter
        decode_metadata.max_seq_len_inter = max_seq_len_inter

        self._cached_decode_metadata = decode_metadata
        return decode_metadata


class DualChunkFlashAttentionMetadataBuilder(FlashAttentionMetadataBuilder):

    def prepare(self):
        super().prepare()
        self.orig_seq_lens: List[int] = []

    def _add_seq_group(
            self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",
            chunked_prefill_enabled: bool, prefix_cache_hit: bool):
        super()._add_seq_group(inter_data, chunked_prefill_enabled,
                               prefix_cache_hit)
        for prompt_len, seq_len in zip(inter_data.prompt_lens,
                                       inter_data.seq_lens):
            self.orig_seq_lens.append(max(prompt_len, seq_len))

    def build(self, seq_lens: List[int], query_lens: List[int],
              cuda_graph_pad_size: int, batch_size: int):
        attn_metadata = super().build(seq_lens, query_lens,
                                      cuda_graph_pad_size, batch_size)
        attn_metadata = DualChunkFlashAttentionMetadata(
            **attn_metadata.asdict_zerocopy())

        device = self.runner.device
        attn_metadata.orig_seq_lens = self.orig_seq_lens
        attn_metadata.orig_seq_lens_tensor = async_tensor_h2d(
            self.orig_seq_lens, torch.int, device, self.runner.pin_memory)

        attn_metadata.block_size = self.runner.block_size
        dual_chunk_attn_config = getattr(self.runner.model_config.hf_config,
                                         "dual_chunk_attention_config", {})
        attn_metadata.original_max_position_embeddings = \
            dual_chunk_attn_config.get("original_max_position_embeddings", 0)
        attn_metadata.chunk_size = dual_chunk_attn_config.get(
            "chunk_size", 8192)
        attn_metadata.local_size = dual_chunk_attn_config.get(
            "local_size", 1024)

        return attn_metadata


class DualChunkFlashAttentionImpl(FlashAttentionImpl):
    """
    If the input tensors contain prompt tokens, the layout is as follows:
    |<--------------- num_prefill_tokens ----------------->|
    |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->|
    Otherwise, the layout is as follows:
    |<----------------- num_decode_tokens ------------------>|
    |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->|
    Generation tokens can contain padding when cuda-graph is used.
    Currently, prompt tokens don't contain any padding.
    The prompts might have different lengths, while the generation tokens
    always have length 1.
    If chunked prefill is enabled, prefill tokens and decode tokens can be
    batched together in a flattened 1D query.
    |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|
    |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|
    Currently, cuda graph is disabled for chunked prefill, meaning there's no
    padding between prefill and decode tokens.
    """

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: int,
        alibi_slopes: Optional[List[float]],
        sliding_window: Optional[int],
        kv_cache_dtype: str,
        blocksparse_params: Optional[Dict[str, Any]] = None,
        logits_soft_cap: Optional[float] = None,
        attn_type: str = AttentionType.DECODER,
        kv_sharing_target_layer_name: Optional[str] = None,
        layer_idx: int = -1,
        dual_chunk_attention_config: Optional[Dict[str, Any]] = None,
    ) -> None:
        if kv_sharing_target_layer_name is not None:
            raise NotImplementedError("KV sharing is not supported in V0.")
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = float(scale)
        self.num_kv_heads = num_kv_heads
        if alibi_slopes is not None:
            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
        self.alibi_slopes = alibi_slopes
        self.sliding_window = ((sliding_window, sliding_window)
                               if sliding_window is not None else (-1, -1))
        self.kv_cache_dtype = kv_cache_dtype

        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
        if sliding_window is not None:
            # NOTE(woosuk): flash-attn's sliding window does not work with
            # paged KV cache.
            raise ValueError(
                "Sliding window is not supported in FlashAttention.")

        support_head_sizes = (
            DualChunkFlashAttentionBackend.get_supported_head_sizes())

        if head_size not in support_head_sizes:
            raise ValueError(
                f"Head size {head_size} is not supported by FlashAttention. "
                f"Supported head sizes are: {support_head_sizes}.")

        assert dual_chunk_attention_config is not None
        self.chunk_size = dual_chunk_attention_config.get("chunk_size", 8192)
        self.local_size = dual_chunk_attention_config.get("local_size", 1024)
        self.original_max_position_embeddings = dual_chunk_attention_config.get(
            "original_max_position_embeddings", 0)
        self.sparse_attention_config = dual_chunk_attention_config.get(
            "sparse_attention_config", None)
        if not self.sparse_attention_config:
            logger.warning_once("Sparse attention will not be enabled as "
                                "sparse attention config is not provided.")
        self.sparse_attention_enabled = dual_chunk_attention_config.get(
            "sparse_attention_enabled", self.sparse_attention_config
            is not None)
        self.sparse_attention_threshold = dual_chunk_attention_config.get(
            "sparse_attention_threshold", 32768)
        self.sparse_attention_last_q = dual_chunk_attention_config.get(
            "sparse_attention_last_q", 64)
        self.layer_idx = layer_idx
        self.dual_chunk_attention_config = dual_chunk_attention_config

        if self.sparse_attention_config:
            self.sparse_attention_config = {
                int(i): j
                for i, j in self.sparse_attention_config[
                    self.layer_idx].items()
            }
            start_head = self.num_heads * get_tensor_model_parallel_rank()
            end_head = start_head + self.num_heads
            self.sparse_attention_config = [
                self.sparse_attention_config[i]
                for i in range(start_head, end_head)
            ]

        if self.sparse_attention_enabled:
            self.arange = torch.arange(self.sparse_attention_last_q,
                                       device="cuda")
            self.last_q_mask = (self.arange[None, None, :, None]
                                >= self.arange[None, None, None, :])

    def forward(  # type: ignore
        self,
        layer: AttentionLayer,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        kv_cache: torch.Tensor,
        attn_metadata: DualChunkFlashAttentionMetadata,
        output: Optional[torch.Tensor] = None,
        output_scale: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Forward pass with DualChunkFlashAttention.
        Args:
            query: shape = [num_tokens, num_heads * head_size]
            query_succ: shape = [num_tokens, num_heads * head_size]
            query_inter: 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 = [2, num_blocks, block_size, num_kv_heads * head_size]
            attn_metadata: Metadata for attention.
        Returns:
            shape = [num_tokens, num_heads * head_size]
        """
        assert output is None, "Output tensor not supported for DualChunk"

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

        (
            query,
            query_succ,
            query_inter,
            query_succ_critical,
            query_inter_critical,
        ) = torch.split(query, query.shape[-1] // 5, dim=-1)

        assert (
            query_succ is not None and query_inter is not None
        ), "query_succ and query_inter are required in Dual Chunk Attention."

        num_tokens, hidden_size = query.shape

        # Reshape the query, key, and value tensors.
        query = query.view(-1, self.num_heads, self.head_size)
        query_succ = query_succ.view(-1, self.num_heads, self.head_size)
        query_inter = query_inter.view(-1, self.num_heads, self.head_size)
        query_succ_critical = query_succ_critical.view(-1, self.num_heads,
                                                       self.head_size)
        query_inter_critical = query_inter_critical.view(
            -1, self.num_heads, self.head_size)
        key = key.view(-1, self.num_kv_heads, self.head_size)
        value = value.view(-1, self.num_kv_heads, self.head_size)

        if self.original_max_position_embeddings > 0:
            if prefill_meta := attn_metadata.prefill_metadata:
                assert prefill_meta.scaling_factor is not None
                assert prefill_meta.query_start_loc is not None
                assert prefill_meta.orig_seq_lens is not None
                current_start = 0
                query_start_loc_cpu = prefill_meta.query_start_loc.cpu()
                for i in range(len(prefill_meta.orig_seq_lens)):
                    current_end = (current_start +
                                   (query_start_loc_cpu[i + 1] -
                                    query_start_loc_cpu[i]).item())
                    key[current_start:current_end].mul_(
                        prefill_meta.scaling_factor[i])
                    current_start = current_end
                assert current_end <= attn_metadata.num_prefill_tokens
            if decode_meta := attn_metadata.decode_metadata:
                assert decode_meta.scaling_factor is not None
                scaling_factor = decode_meta.scaling_factor
                key[attn_metadata.num_prefill_tokens:].mul_(
                    scaling_factor.unsqueeze(-1).unsqueeze(-1))

        if kv_cache is not None and kv_cache.numel() > 0:
            key_cache = kv_cache[0]
            value_cache = kv_cache[1]

            # Reshape the input keys and values and store them in the cache.
            # If kv_cache is not provided, the new key and value tensors are
            # not cached. This happens during the initial memory profiling run.
            ops.reshape_and_cache_flash(
                key,
                value,
                key_cache,
                value_cache,
                attn_metadata.slot_mapping.flatten(),
                self.kv_cache_dtype,
                layer._k_scale,
                layer._v_scale,
            )

        num_prefill_tokens = attn_metadata.num_prefill_tokens
        num_decode_tokens = attn_metadata.num_decode_tokens
        assert key.shape[0] == num_prefill_tokens + num_decode_tokens
        assert value.shape[0] == num_prefill_tokens + num_decode_tokens
        output = torch.empty_like(query)

        # Query for decode. KV is not needed because it is already cached.
        decode_query = query[num_prefill_tokens:]
        decode_query_succ = query_succ[num_prefill_tokens:]
        decode_query_inter = query_inter[num_prefill_tokens:]

        # QKV for prefill.
        query = query[:num_prefill_tokens]
        query_succ = query_succ[:num_prefill_tokens]
        query_inter = query_inter[:num_prefill_tokens]
        query_succ_critical = query_succ_critical[:num_prefill_tokens]
        query_inter_critical = query_inter_critical[:num_prefill_tokens]
        key = key[:num_prefill_tokens]
        value = value[:num_prefill_tokens]
        assert query.shape[0] == num_prefill_tokens
        assert decode_query.shape[0] == num_decode_tokens

        if prefill_meta := attn_metadata.prefill_metadata:
            # Prompt run.
            if (kv_cache is None or prefill_meta.block_tables is None
                    or prefill_meta.block_tables.numel() == 0):
                # normal attention, called during the profiling run.
                out = flash_attn_varlen_func(
                    q=query,
                    k=key,
                    v=value,
                    cu_seqlens_q=prefill_meta.seq_start_loc,
                    cu_seqlens_k=prefill_meta.seq_start_loc,
                    max_seqlen_q=prefill_meta.max_prefill_seq_len,
                    max_seqlen_k=prefill_meta.max_prefill_seq_len,
                    softmax_scale=self.scale,
                    causal=True,
                    window_size=self.sliding_window,
                    alibi_slopes=self.alibi_slopes,
                )
                assert output[:num_prefill_tokens].shape == out.shape
                output[:num_prefill_tokens] = out
            else:
                # prefix-enabled attention
                assert prefill_meta.seq_lens is not None
                assert prefill_meta.orig_seq_lens is not None
                output[:num_prefill_tokens] = (
                    self._dual_chunk_flash_attn_prefill(
                        q=query,
                        q_succ=query_succ,
                        q_inter=query_inter,
                        q_succ_critical=query_succ_critical,
                        q_inter_critical=query_inter_critical,
                        k=key_cache,
                        v=value_cache,
                        cu_seqlens_q=prefill_meta.query_start_loc,
                        cu_seqlens_k=prefill_meta.seq_start_loc,
                        orig_seq_lens=prefill_meta.orig_seq_lens,
                        scaling_factor=prefill_meta.scaling_factor,
                        softmax_scale=self.scale,
                        causal=True,
                        window_size=(-1, -1),
                        alibi_slopes=self.alibi_slopes,
                        block_table=prefill_meta.block_tables,
                        chunk_size=self.chunk_size,
                        local_size=self.local_size,
                    ))

        if decode_meta := attn_metadata.decode_metadata:
            # Decoding run.
            output[num_prefill_tokens:] = (
                self._dual_chunk_flash_attn_decoding(
                    decode_query.unsqueeze(1),
                    decode_query_succ.unsqueeze(1),
                    decode_query_inter.unsqueeze(1),
                    key_cache,
                    value_cache,
                    block_table=decode_meta.block_tables,
                    cache_seqlens=decode_meta.seq_lens_tensor,
                    softmax_scale=self.scale,
                    causal=True,
                    alibi_slopes=self.alibi_slopes,
                    chunk_size=self.chunk_size,
                    local_size=self.local_size,
                    original_max_position_embeddings=self.
                    original_max_position_embeddings,
                    decode_meta=decode_meta,
                ).squeeze(1))
        # Reshape the output tensor.
        return output.view(num_tokens, hidden_size)

    def _dual_chunk_flash_attn_prefill(
        self,
        q,
        q_succ,
        q_inter,
        q_succ_critical,
        q_inter_critical,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        orig_seq_lens: List[int],
        scaling_factor: torch.Tensor,
        softmax_scale: float,
        causal: Optional[bool] = True,
        window_size: Tuple[int, int] = (-1, -1),
        alibi_slopes: Optional[torch.Tensor] = None,
        block_table: Optional[torch.Tensor] = None,
        chunk_size: int = 8192,
        local_size: int = 1024,
    ):
        if alibi_slopes is not None:
            raise ValueError(
                "Dual Chunk Attention does not support alibi_slopes")
        if not causal:
            raise ValueError(
                "Dual Chunk Attention does not support causal=False")
        if window_size != (-1, -1):
            raise ValueError(
                "Dual Chunk Attention does not support window_size")

        cu_seqlens_q_cpu = cu_seqlens_q.cpu().tolist()
        cu_seqlens_k_cpu = cu_seqlens_k.cpu().tolist()
        all_outputs = []

        for i in range(0, len(cu_seqlens_q_cpu) - 1):
            qs = cu_seqlens_q_cpu[i]
            qe = cu_seqlens_q_cpu[i:i + 2][-1]
            ks = cu_seqlens_k_cpu[i]
            ke = cu_seqlens_k_cpu[i:i + 2][-1]

            current_q = q[qs:qe]
            current_q_succ = q_succ[qs:qe]
            current_q_inter = q_inter[qs:qe]
            current_q_succ_critical = q_succ_critical[qs:qe]
            current_q_inter_critical = q_inter_critical[qs:qe]

            if block_table is None:
                current_k = k[ks:ke]
                current_v = v[ks:ke]
                current_block_table = None
                current_orig_seq_len = orig_seq_lens[i]
            else:
                current_block_table = block_table[i]
                current_orig_seq_len = orig_seq_lens[i]
                current_k = k
                current_v = v
            sparse_attn_enabled = (self.sparse_attention_enabled
                                   and current_orig_seq_len
                                   > self.sparse_attention_threshold)

            if current_q.shape[0] == 0:
                continue

            if current_k.shape[0] == 0:
                all_outputs.append(
                    torch.zeros(
                        (current_q.shape[0], current_q.shape[1], v.shape[2]),
                        device=q.device,
                        dtype=q.dtype,
                    ))
                continue

            current_output = torch.empty_like(current_q)
            group_size = int(current_q.size(-2) / current_k.size(-2))

            if sparse_attn_enabled:
                num_device_q_heads = current_q.size(-2)
                heads_vertical_size = torch.empty(size=(num_device_q_heads, ),
                                                  dtype=torch.int32)
                heads_slash_size = torch.empty(size=(num_device_q_heads, ),
                                               dtype=torch.int32)
                for head_id in range(current_q.size(-2)):
                    (
                        ty,
                        vertical_size,
                        slash_size,
                        _,
                    ) = self.sparse_attention_config[head_id]
                    assert ty == "vertical_and_slash", "only support slash mode"

                    if vertical_size == 30:
                        vertical_size += 100
                    heads_vertical_size[head_id] = vertical_size
                    heads_slash_size[head_id] = slash_size

                current_output = self._dual_chunk_flash_attn_prefill_func(
                    current_q,  # allheads
                    current_q_succ,
                    current_q_inter,
                    current_q_succ_critical,
                    current_q_inter_critical,
                    current_k,
                    current_v,
                    current_block_table,
                    softmax_scale,
                    chunk_size,
                    local_size,
                    scaling_factor[i].item(),
                    ke - ks,
                    sparse_attn_enabled=sparse_attn_enabled,
                    heads_vertical_size=heads_vertical_size,
                    heads_slash_size=heads_slash_size,
                    group_size=group_size)
            else:
                for head_id in range(current_q.size(-2)):
                    # (seq_len, num_heads, head_size)
                    current_q_head = current_q[:, head_id, :].unsqueeze(1)
                    current_q_succ_head = \
                        current_q_succ[:, head_id, :].unsqueeze(1)
                    current_q_inter_head = \
                        current_q_inter[:, head_id, :].unsqueeze(1)
                    current_q_succ_head_critical = \
                        current_q_succ_critical[:, head_id, :].unsqueeze(1)
                    current_q_inter_head_critical = \
                        current_q_inter_critical[:, head_id, :].unsqueeze(1)
                    if block_table is not None:
                        current_k_head = current_k[..., head_id //
                                                   group_size, :].unsqueeze(2)
                        current_v_head = current_v[..., head_id //
                                                   group_size, :].unsqueeze(2)

                    else:
                        current_k_head = current_k[:, head_id, :].unsqueeze(1)
                        current_v_head = current_v[:, head_id, :].unsqueeze(1)

                    current_out = self._dual_chunk_flash_attn_prefill_func(
                        current_q_head,
                        current_q_succ_head,
                        current_q_inter_head,
                        current_q_succ_head_critical,
                        current_q_inter_head_critical,
                        current_k_head,
                        current_v_head,
                        current_block_table,
                        softmax_scale,
                        chunk_size,
                        local_size,
                        scaling_factor[i].item(),
                        ke - ks,
                        sparse_attn_enabled=sparse_attn_enabled,
                    )
                    current_output[:, head_id:head_id + 1, :] = current_out
            all_outputs.append(current_output)
        return torch.cat(all_outputs, dim=0)

    def _dual_chunk_flash_attn_prefill_func(
        self,
        q,
        q_succ,
        q_inter,
        q_succ_critical,
        q_inter_critical,
        k,
        v,
        block_table,
        softmax_scale: float,
        chunk_size: int,
        local_size: int,
        scaling_factor: float,
        k_length: int,
        sparse_attn_enabled: Optional[bool] = True,
        heads_vertical_size=None,
        heads_slash_size=None,
        group_size=None,
    ):
        flash_results = []
        chunk_len = chunk_size - local_size

        if block_table is not None:
            block_size = v.shape[1]
            if chunk_len % block_size != 0:
                raise ValueError("chunk_len must be divisible by block_size.")
        else:
            block_size = 1

        if self.original_max_position_embeddings > 0:
            softmax_scale = softmax_scale * scaling_factor

        begin = k_length - q.shape[0]
        while begin < k_length:
            flash_per_chunk = []

            prev_chunk_end_pos = (begin // chunk_len) * chunk_len
            next_chunk_end_pos = prev_chunk_end_pos + chunk_len
            end = min(next_chunk_end_pos, k_length)
            qbegin = begin - (k_length - q.shape[0])
            qend = end - (k_length - q.shape[0])

            qk_chunks = []
            q_states_intra = q[qbegin:qend]
            # choose critical token
            if block_table is not None:
                block_tables_intra = _get_block(block_table, block_size,
                                                prev_chunk_end_pos, end)
                k_states_intra = k[block_tables_intra].view(
                    -1, *k.shape[-2:])[:(end - prev_chunk_end_pos)]
                v_states_intra = v[block_tables_intra].view(
                    -1, *v.shape[-2:])[:(end - prev_chunk_end_pos)]
            else:
                block_tables_intra = None
                k_states_intra = k[prev_chunk_end_pos:end]
                v_states_intra = v[prev_chunk_end_pos:end]

            if sparse_attn_enabled:
                last_q_size = min(qend - qbegin, self.sparse_attention_last_q)
                _, num_device_k_heads, head_dim = k_states_intra.shape
                k_states_intra = (k_states_intra.unsqueeze(2).repeat(
                    1, 1, group_size,
                    1).reshape(-1, num_device_k_heads * group_size, head_dim))
                v_states_intra = (v_states_intra.unsqueeze(2).repeat(
                    1, 1, group_size,
                    1).reshape(-1, num_device_k_heads * group_size, head_dim))
                qk_chunks.append(
                    (q_states_intra.transpose(0, 1)[:, -last_q_size:] *
                     softmax_scale) @ k_states_intra.permute(1, 2, 0))

            if prev_chunk_end_pos - chunk_len >= 0:
                q_states_succ = q_succ[qbegin:qend]
                q_states_succ_critical = q_succ_critical[qbegin:qend]
                if block_table is not None:
                    block_tables_succ = _get_block(
                        block_table, block_size,
                        prev_chunk_end_pos - chunk_len, prev_chunk_end_pos)
                    k_states_succ = k[block_tables_succ].view(
                        -1, *k.shape[-2:])[:chunk_len]
                    v_states_succ = v[block_tables_succ].view(
                        -1, *v.shape[-2:])[:chunk_len]
                else:
                    k_states_succ = k[prev_chunk_end_pos -
                                      chunk_len:prev_chunk_end_pos]
                    v_states_succ = v[prev_chunk_end_pos -
                                      chunk_len:prev_chunk_end_pos]

                if sparse_attn_enabled:
                    k_states_succ = (k_states_succ.unsqueeze(2).repeat(
                        1, 1, group_size,
                        1).reshape(-1, num_device_k_heads * group_size,
                                   head_dim))
                    v_states_succ = (v_states_succ.unsqueeze(2).repeat(
                        1, 1, group_size,
                        1).reshape(-1, num_device_k_heads * group_size,
                                   head_dim))
                    qk_chunks.append((q_states_succ_critical.transpose(
                        0, 1)[:, -last_q_size:] * softmax_scale)
                                     @ k_states_succ.permute(1, 2, 0))

            if prev_chunk_end_pos - chunk_len * 2 >= 0:
                q_states_inter = q_inter[qbegin:qend]
                q_states_inter_critical = q_inter_critical[qbegin:qend]
                if block_table is not None:
                    block_tables_inter = _get_block(
                        block_table, block_size, 0,
                        prev_chunk_end_pos - chunk_len)
                    k_states_inter = k[block_tables_inter].view(
                        -1, *k.shape[-2:])[:(prev_chunk_end_pos - chunk_len)]
                    v_states_inter = v[block_tables_inter].view(
                        -1, *v.shape[-2:])[:(prev_chunk_end_pos - chunk_len)]
                else:
                    k_states_inter = k[:prev_chunk_end_pos - chunk_len]
                    v_states_inter = v[:prev_chunk_end_pos - chunk_len]

                if sparse_attn_enabled:
                    k_states_inter = (k_states_inter.unsqueeze(2).repeat(
                        1, 1, group_size,
                        1).reshape(-1, num_device_k_heads * group_size,
                                   head_dim))
                    v_states_inter = (v_states_inter.unsqueeze(2).repeat(
                        1, 1, group_size,
                        1).reshape(-1, num_device_k_heads * group_size,
                                   head_dim))
                    qk_chunks.append((q_states_inter_critical.transpose(
                        0, 1)[:, -last_q_size:] * softmax_scale)
                                     @ k_states_inter.permute(1, 2, 0))

            if sparse_attn_enabled:
                reversed_qk = qk_chunks[::-1]
                qk = torch.cat(reversed_qk, dim=-1)

                qk[:, :, -last_q_size:] = torch.where(
                    self.last_q_mask[..., -last_q_size:,
                                     -last_q_size:].to(qk.device),
                    qk[:, :, -last_q_size:], -torch.inf)
                qk = F.softmax(qk, dim=-1, dtype=torch.float32)

                vertical = qk.sum(-2, keepdim=True)
                vertical[..., :30] = torch.inf

                # Avoid sorting by using the min/max ints to fill the indexer
                # buffers.
                int32_max = torch.iinfo(torch.int32).max
                int32_min = torch.iinfo(torch.int32).min
                n_heads = qk.size()[0]
                max_slash_topk = torch.max(heads_slash_size).item()
                max_vertical_topk = torch.max(heads_vertical_size).item()
                # store each head's slash topk, vertical topk
                vertical = vertical.reshape((n_heads, -1))
                # prevent out of range when prompt size < max_vertical_topk
                max_vertical_topk = min(vertical.shape[-1], max_vertical_topk)
                vertical_topk_buffer = torch.topk(vertical, max_vertical_topk,
                                                  -1).indices
                slash_topk_buffer = torch.empty(size=(n_heads, max_slash_topk),
                                                dtype=torch.int64,
                                                device=qk.device)
                for head_i in range(n_heads):
                    #  (nqheads=1, lastq, k_len)
                    head_score = qk[head_i:head_i + 1, :, :]
                    slash_scores = _sum_all_diagonal_matrix(head_score)
                    if head_score.size(1) != 1:
                        # drop right up corner
                        slash_scores = slash_scores[..., :-last_q_size + 1]
                    slash_scores[..., -100:] = torch.inf

                    head_slash_size = heads_slash_size[head_i]
                    head_slash_size = min(head_slash_size, vertical.size(-1))
                    slash_topk = torch.topk(slash_scores, head_slash_size,
                                            -1).indices
                    #（nheads, max_topk）
                    slash_topk_buffer[head_i, :head_slash_size] = slash_topk

                    # reset heads topk
                    heads_slash_size[head_i] = head_slash_size
                    heads_vertical_size[head_i] = min(
                        heads_vertical_size[head_i], max_vertical_topk)

                # store
                vertical_buffer = torch.full((n_heads, max_vertical_topk),
                                             int32_max,
                                             dtype=torch.int64,
                                             device=q.device)
                slash_buffer = torch.full((n_heads, max_slash_topk),
                                          int32_min,
                                          dtype=torch.int64,
                                          device=q.device)
                succ_vertical_buffer = torch.full((n_heads, max_vertical_topk),
                                                  int32_max,
                                                  dtype=torch.int64,
                                                  device=q.device)
                succ_slash_buffer = torch.full((n_heads, max_slash_topk),
                                               int32_min,
                                               dtype=torch.int64,
                                               device=q.device)
                inter_vertical_buffer = torch.full(
                    (n_heads, max_vertical_topk),
                    int32_max,
                    dtype=torch.int64,
                    device=q.device)
                inter_slash_buffer = torch.full((n_heads, max_slash_topk),
                                                int32_min,
                                                dtype=torch.int64,
                                                device=q.device)

                vertical_size_buffer = torch.empty(size=(n_heads, ),
                                                   dtype=torch.int32,
                                                   device=q.device)
                slash_sizes_buffer = torch.empty(size=(n_heads, ),
                                                 dtype=torch.int32,
                                                 device=q.device)
                succ_vertical_size_buffer = torch.empty(size=(n_heads, ),
                                                        dtype=torch.int32,
                                                        device=q.device)
                succ_slash_sizes_buffer = torch.empty(size=(n_heads, ),
                                                      dtype=torch.int32,
                                                      device=q.device)
                inter_vertical_size_buffer = torch.empty(size=(n_heads, ),
                                                         dtype=torch.int32,
                                                         device=q.device)
                inter_slash_sizes_buffer = torch.empty(size=(n_heads, ),
                                                       dtype=torch.int32,
                                                       device=q.device)

                for head_i in range(n_heads):
                    vertical_topk = vertical_topk_buffer[
                        head_i, :heads_vertical_size[head_i]]
                    # intra
                    intra_vertical_indices = vertical_topk[
                        vertical_topk >=
                        prev_chunk_end_pos] - prev_chunk_end_pos
                    if intra_vertical_indices.nelement() == 0:
                        intra_vertical_indices = torch.cat([
                            intra_vertical_indices,
                            torch.arange(0,
                                         k_states_intra.size(0),
                                         max(1,
                                             k_states_intra.size(0) / 5),
                                         dtype=torch.int32,
                                         device=intra_vertical_indices.device)
                        ])
                    slash_topk = slash_topk_buffer[
                        head_i, :heads_slash_size[head_i]]
                    intra_slash_indices = (
                        (qk.size(-1) - 1) -
                        slash_topk[slash_topk >= prev_chunk_end_pos])
                    # fill buffer
                    v_count = intra_vertical_indices.nelement()
                    s_count = intra_slash_indices.nelement()
                    vertical_size_buffer[head_i] = v_count
                    slash_sizes_buffer[head_i] = s_count
                    vertical_buffer[head_i, :v_count].copy_(
                        intra_vertical_indices)
                    slash_buffer[head_i, :s_count].copy_(intra_slash_indices)
                    # succ
                    if prev_chunk_end_pos - chunk_len >= 0:
                        succ_vertical_indices = vertical_topk[
                            (vertical_topk < prev_chunk_end_pos)
                            & (vertical_topk >= prev_chunk_end_pos -
                               chunk_len)] - (prev_chunk_end_pos - chunk_len)
                        # TODO: support no vertical
                        if succ_vertical_indices.nelement() == 0:
                            succ_vertical_indices = torch.cat([
                                succ_vertical_indices,
                                torch.arange(
                                    0,
                                    k_states_succ.size(0),
                                    max(1,
                                        k_states_succ.size(0) / 5),
                                    dtype=torch.int32,
                                    device=intra_vertical_indices.device)
                            ])
                        succ_slash_indices = (
                            (prev_chunk_end_pos + (qend - qbegin) - 1) -
                            slash_topk[((slash_topk >=
                                         (prev_chunk_end_pos - chunk_len)) &
                                        (slash_topk < (prev_chunk_end_pos +
                                                       (qend - qbegin))))])
                        if succ_slash_indices.nelement() == 0:
                            succ_slash_indices = torch.cat([
                                succ_slash_indices,
                                torch.arange(
                                    0,
                                    k_states_succ.size(0),
                                    max(1,
                                        k_states_succ.size(0) / 5),
                                    dtype=torch.int32,
                                    device=intra_vertical_indices.device)
                            ])
                        # fill buffer
                        v_count = succ_vertical_indices.nelement()
                        s_count = succ_slash_indices.nelement()
                        succ_vertical_size_buffer[head_i] = v_count
                        succ_slash_sizes_buffer[head_i] = s_count
                        succ_vertical_buffer[head_i, :v_count].copy_(
                            succ_vertical_indices)
                        succ_slash_buffer[head_i, :s_count].copy_(
                            succ_slash_indices)

                    if prev_chunk_end_pos - 2 * chunk_len >= 0:
                        inter_vertical_indices = vertical_topk[
                            vertical_topk < prev_chunk_end_pos - chunk_len]

                        if inter_vertical_indices.nelement() == 0:
                            inter_vertical_indices = torch.cat([
                                inter_vertical_indices,
                                torch.arange(
                                    0,
                                    k_states_inter.size(0),
                                    max(1,
                                        k_states_inter.size(0) / 5),
                                    dtype=torch.int32,
                                    device=intra_vertical_indices.device)
                            ])
                        inter_slash_indices = (
                            (prev_chunk_end_pos - chunk_len +
                             (qend - qbegin) - 1) -
                            slash_topk[slash_topk < (prev_chunk_end_pos -
                                                     chunk_len +
                                                     (qend - qbegin))])
                        if inter_slash_indices.nelement() == 0:
                            inter_slash_indices = torch.cat([
                                inter_slash_indices,
                                torch.arange(
                                    0,
                                    k_states_inter.size(0),
                                    max(1,
                                        k_states_inter.size(0) / 5),
                                    dtype=torch.int32,
                                    device=intra_vertical_indices.device)
                            ])
                        # fill buffer
                        v_count = inter_vertical_indices.nelement()
                        s_count = inter_slash_indices.nelement()
                        inter_vertical_size_buffer[head_i] = v_count
                        inter_slash_sizes_buffer[head_i] = s_count
                        inter_vertical_buffer[head_i, :v_count].copy_(
                            inter_vertical_indices)
                        inter_slash_buffer[head_i, :s_count].copy_(
                            inter_slash_indices)
            else:
                intra_vertical_indices, intra_slash_indices = None, None
                succ_vertical_indices, succ_slash_indices = None, None
                inter_vertical_indices, inter_slash_indices = None, None

            if sparse_attn_enabled:
                flash_result = self._do_flash_attn(
                    q_states_intra,
                    k_states_intra,
                    v_states_intra,
                    softmax_scale=softmax_scale,
                    causal=True,
                    block_table=block_table,
                    stage="intra",
                    vertical_indices=vertical_buffer,
                    slash_indices=slash_buffer,
                    vertical_indices_count=vertical_size_buffer,
                    slash_indices_count=slash_sizes_buffer,
                    mergehead_softmax_scale=softmax_scale,
                    sparse_attn_enabled=sparse_attn_enabled)
            else:
                flash_result = self._do_flash_attn(
                    q_states_intra,
                    k_states_intra,
                    v_states_intra,
                    softmax_scale=softmax_scale,
                    causal=True,
                    block_table=block_table,
                    stage="intra",
                    vertical_indices=intra_vertical_indices,
                    slash_indices=intra_slash_indices,
                    sparse_attn_enabled=sparse_attn_enabled)
            flash_per_chunk.append(flash_result)

            if prev_chunk_end_pos - chunk_len >= 0:
                if sparse_attn_enabled:
                    flash_result = self._do_flash_attn(
                        q_states_succ,
                        k_states_succ,
                        v_states_succ,
                        softmax_scale=softmax_scale,
                        causal=False,
                        block_table=block_table,
                        stage="succ",
                        vertical_indices=succ_vertical_buffer,
                        slash_indices=succ_slash_buffer,
                        vertical_indices_count=succ_vertical_size_buffer,
                        slash_indices_count=succ_slash_sizes_buffer,
                        mergehead_softmax_scale=softmax_scale,
                        sparse_attn_enabled=sparse_attn_enabled)
                else:
                    flash_result = self._do_flash_attn(
                        q_states_succ,
                        k_states_succ,
                        v_states_succ,
                        softmax_scale=softmax_scale,
                        causal=False,
                        block_table=block_table,
                        stage="succ",
                        vertical_indices=succ_vertical_indices,
                        slash_indices=succ_slash_indices,
                        sparse_attn_enabled=sparse_attn_enabled)
                flash_per_chunk.append(flash_result)

            if prev_chunk_end_pos - chunk_len * 2 >= 0:
                if sparse_attn_enabled:
                    flash_result = self._do_flash_attn(
                        q_states_inter,
                        k_states_inter,
                        v_states_inter,
                        softmax_scale=softmax_scale,
                        causal=False,
                        block_table=block_table,
                        stage="inter",
                        vertical_indices=inter_vertical_buffer,
                        slash_indices=inter_slash_buffer,
                        vertical_indices_count=inter_vertical_size_buffer,
                        slash_indices_count=inter_slash_sizes_buffer,
                        mergehead_softmax_scale=softmax_scale,
                        sparse_attn_enabled=sparse_attn_enabled)
                else:
                    flash_result = self._do_flash_attn(
                        q_states_inter,
                        k_states_inter,
                        v_states_inter,
                        softmax_scale=softmax_scale,
                        causal=False,
                        block_table=block_table,
                        stage="inter",
                        vertical_indices=inter_vertical_indices,
                        slash_indices=inter_slash_indices,
                        sparse_attn_enabled=sparse_attn_enabled)
                flash_per_chunk.append(flash_result)

            flash_results.append(flash_per_chunk)
            begin = end

        attn_output = self._merge_attn_outputs(flash_results)
        del flash_results
        return attn_output

    def _do_flash_attn(
        self,
        query_states: torch.Tensor,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        softmax_scale: float,
        causal: bool = True,
        block_table: torch.Tensor = None,
        max_seqlen_k: Optional[int] = None,
        stage: str = "intra",
        vertical_indices: Optional[torch.Tensor] = None,
        slash_indices: Optional[torch.Tensor] = None,
        vertical_indices_count: Optional[torch.Tensor] = None,
        slash_indices_count: Optional[torch.Tensor] = None,
        mergehead_softmax_scale: Optional[float] = None,
        sparse_attn_enabled: Optional[bool] = False,
    ):
        if max_seqlen_k is None:
            max_seqlen_k = key_states.shape[0]

        q_len = query_states.shape[0]
        q_heads = query_states.shape[1]
        h_dim = query_states.shape[-1]

        if sparse_attn_enabled:
            assert slash_indices is not None
            if stage == "intra":
                assert causal
            else:
                assert not causal

            query_states = query_states.unsqueeze(0).transpose(1, 2)
            key_states = key_states.unsqueeze(0).transpose(1, 2)
            value_states = value_states.unsqueeze(0).transpose(1, 2)

            q = query_states
            k = key_states
            v = value_states

            if (vertical_indices_count is not None and \
                    slash_indices_count is not None):
                assert mergehead_softmax_scale is not None

                res, s_lse = _vertical_slash_sparse_attention(
                    q,
                    k,
                    v,
                    vertical_indices,
                    slash_indices,
                    mergehead_softmax_scale,
                    causal=causal,
                    stage=stage,
                    vertical_indices_count=vertical_indices_count,
                    slash_indices_count=slash_indices_count)
                res = res.view(q_heads, q_len,
                               h_dim).transpose(0, 1)  # (qlen,nhead,h_dim)
                s_lse = s_lse.view(
                    q_heads, q_len,
                    1).squeeze(-1).unsqueeze(0).float()  # (1, nhead,qlen)
            else:
                res, s_lse = _vertical_slash_sparse_attention(q,
                                                              k,
                                                              v,
                                                              vertical_indices,
                                                              slash_indices,
                                                              softmax_scale,
                                                              causal=causal,
                                                              stage=stage)
                res = res.view(q_len, q_heads, h_dim)
                s_lse = s_lse.view(q_len, q_heads, 1).transpose(0, 2).float()
            return res, s_lse

        if not current_platform.is_rocm():
            output, softmax_lse = flash_attn_varlen_func(
                q=query_states,
                k=key_states,
                v=value_states,
                softmax_scale=softmax_scale,
                cu_seqlens_q=torch.tensor([0, query_states.shape[0]],
                                        dtype=torch.int32,
                                        device=query_states.device),
                max_seqlen_q=query_states.shape[0],
                cu_seqlens_k=torch.tensor([0, max_seqlen_k],
                                        dtype=torch.int32,
                                        device=query_states.device),
                max_seqlen_k=max_seqlen_k,
                causal=causal,
                block_table=block_table.unsqueeze(0),
                return_softmax_lse=True,
            )
        else:
            output, softmax_lse = flash_attn_varlen_func(
                q=query_states,
                k=key_states,
                v=value_states,
                softmax_scale=softmax_scale,
                cu_seqlens_q=torch.tensor([0, query_states.shape[0]],
                                        dtype=torch.int32,
                                        device=query_states.device),
                max_seqlen_q=query_states.shape[0],
                cu_seqlens_k=torch.tensor([0, max_seqlen_k],
                                        dtype=torch.int32,
                                        device=query_states.device),
                max_seqlen_k=max_seqlen_k,
                causal=causal,
                block_table=block_table.unsqueeze(0),
                return_attn_probs=True,
            )
        softmax_lse = softmax_lse.view(q_len, q_heads, 1).transpose(0,
                                                                    2).float()
        return output, softmax_lse

    def _merge_attn_outputs(
        self,
        flash_results: List[List[Tuple[torch.Tensor, torch.Tensor]]],
        return_lse: Optional[bool] = False,
    ) -> torch.Tensor:
        attn_outputs_all = []
        logits_all = []

        for flash_per_chunk in flash_results:
            if len(flash_per_chunk) == 1:
                attn_outputs_all.append(flash_per_chunk[0][0])
                if return_lse:
                    logits_all.append(flash_per_chunk[0][1])
                continue

            attn_outputs = torch.stack([
                flash_attn_output[0] for flash_attn_output in flash_per_chunk
            ])
            logits = torch.stack([
                flash_attn_output[1] for flash_attn_output in flash_per_chunk
            ])
            logits = logits.to(torch.float32)

            if return_lse:
                max_val = torch.max(logits, dim=0).values
                diff = torch.abs(logits[0] - logits[1])
                log_sum_exp = max_val + torch.log1p(torch.exp(-diff))
                logits_all.append(log_sum_exp)

            max_logits = torch.max(logits, dim=0).values
            stable_logits = logits - max_logits.unsqueeze(0)
            lse_s = torch.exp(stable_logits).detach()
            lse_sum = torch.sum(lse_s, dim=0)
            lse_s /= lse_sum
            attn_outputs *= lse_s.unsqueeze(-1).transpose(2, 3).squeeze(1)
            attn_outputs_all.append(attn_outputs.sum(dim=0))

        if return_lse:
            return (torch.cat(attn_outputs_all,
                              dim=0), torch.cat(logits_all, dim=-1))
        else:
            return torch.cat(attn_outputs_all, dim=0)

    def _dual_chunk_flash_attn_decoding(
        self,
        query: torch.Tensor,
        query_succ: torch.Tensor,
        query_inter: torch.Tensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        block_table: torch.Tensor,
        cache_seqlens: torch.Tensor,
        softmax_scale: float,
        causal: bool,
        alibi_slopes: Optional[torch.Tensor],
        chunk_size: int,
        local_size: int,
        original_max_position_embeddings: int,
        decode_meta: DualChunkFlashAttentionMetadata,
    ):
        if not causal:
            raise ValueError(
                "Dual Chunk Attention does not support causal=False")

        block_size = value_cache.shape[1]
        chunk_len = chunk_size - local_size
        if chunk_len % block_size != 0:
            raise ValueError("chunk_len must be divisible by block_size.")
        if original_max_position_embeddings > 0:
            assert decode_meta.scaling_factor is not None
            scaling_factor = decode_meta.scaling_factor
            query = (query * scaling_factor.view(-1, 1, 1, 1)).to(
                query.dtype
            )  # possible for numerical issue, need to fused in the kernel
            query_succ = (query_succ * scaling_factor.view(-1, 1, 1, 1)).to(
                query.dtype)
            query_inter = (query_inter * scaling_factor.view(-1, 1, 1, 1)).to(
                query.dtype)
        outputs_list = []
        softmax_lses_list = []

        # intra-attention
        intra_output, intra_softmax_lse = (
            self._dual_chunk_flash_attn_decoding_with_exp_sums(
                query,
                key_cache,
                value_cache,
                decode_meta.block_tables_intra,
                decode_meta.seq_lens_intra,
                softmax_scale,
                alibi_slopes,
                causal=False,
            ))
        outputs_list.append(intra_output)
        softmax_lses_list.append(intra_softmax_lse)

        # succ-attention
        if decode_meta.max_seq_len_succ:
            succ_output, succ_softmax_lse = (
                self._dual_chunk_flash_attn_decoding_with_exp_sums(
                    query_succ,
                    key_cache,
                    value_cache,
                    decode_meta.block_tables_succ,
                    decode_meta.seq_lens_succ,
                    softmax_scale,
                    alibi_slopes,
                    causal=False,
                ))
            outputs_list.append(succ_output)
            softmax_lses_list.append(succ_softmax_lse)

        # inter-attention
        if decode_meta.max_seq_len_inter:
            inter_output, inter_softmax_lse = (
                self._dual_chunk_flash_attn_decoding_with_exp_sums(
                    query_inter,
                    key_cache,
                    value_cache,
                    block_table[:, :decode_meta.max_seq_len_inter],
                    decode_meta.seq_lens_inter,
                    softmax_scale,
                    alibi_slopes,
                    causal=False,
                ))
            outputs_list.append(inter_output)
            softmax_lses_list.append(inter_softmax_lse)
        outputs = torch.stack(outputs_list, dim=0)
        del outputs_list
        softmax_lses = torch.stack(softmax_lses_list, dim=0).to(torch.float32)
        del softmax_lses_list
        max_logits = torch.max(softmax_lses, dim=0).values
        stable_logits = softmax_lses - max_logits.unsqueeze(0)
        lse_s = torch.exp(stable_logits).detach()
        lse_sum = torch.sum(lse_s, dim=0)
        lse_s /= lse_sum
        outputs *= lse_s.unsqueeze(-1).transpose(2, 3)
        return outputs.sum(0)

    def _dual_chunk_flash_attn_decoding_with_exp_sums(
        self,
        query: torch.Tensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        block_table: torch.Tensor,
        cache_seqlens: torch.Tensor,
        softmax_scale: float,
        alibi_slopes: Optional[torch.Tensor],
        causal: bool,
    ):
        out, softmax_lse = flash_attn_with_kvcache(
            q=query,
            k_cache=key_cache,
            v_cache=value_cache,
            block_table=block_table,
            cache_seqlens=cache_seqlens,
            softmax_scale=softmax_scale,
            alibi_slopes=alibi_slopes,
            causal=causal,
            return_softmax_lse=True,
        )
        mask = (cache_seqlens == 0)
        out[mask] = 0
        softmax_lse[mask] = -float("inf")
        return out, softmax_lse


def _vertical_slash_sparse_attention(
    query: torch.Tensor,  # [BATCH, N_HEADS, N_CTX, D_HEAD]
    key: torch.Tensor,  # [BATCH, N_HEADS, N_KV_CTX, D_HEAD]
    value: torch.Tensor,  # [BATCH, N_HEADS, N_KV_CTX, D_HEAD]
    v_idx: torch.Tensor,  # [BATCH, N_HEADS, NNZ_V]
    s_idx: torch.Tensor,  # [BATCH, N_HEADS, NNZ_S]
    softmax_scale: float,
    causal: bool = True,
    stage: str = "intra",
    block_size_M: int = 64,
    block_size_N: int = 64,
    vertical_indices_count: torch.Tensor = None,  # [N_HEADS,]
    slash_indices_count: torch.Tensor = None,
):
    if stage == "intra":
        assert causal
    else:
        assert not causal

    batch_size, num_heads, context_size, head_dim = query.shape
    _, _, kv_seq_len, _ = key.shape

    if head_dim not in [16, 32, 64, 128, 256, 512]:
        target_dim = 2**math.ceil(math.log2(head_dim)) - head_dim
        query = F.pad(query, [0, target_dim, 0, 0, 0, 0, 0, 0])
        key = F.pad(key, [0, target_dim, 0, 0, 0, 0, 0, 0])
        value = F.pad(value, [0, target_dim, 0, 0, 0, 0, 0, 0])

    v_idx = v_idx.to(torch.int32).reshape(
        (batch_size, num_heads, -1)).sort(dim=-1, descending=False)[0]
    s_idx = s_idx.to(torch.int32).reshape(
        (batch_size, num_heads, -1)).sort(dim=-1, descending=True)[0]
    q_seqlens = torch.tensor([context_size],
                             dtype=torch.int32,
                             device=query.device)
    kv_seqlens = torch.tensor([kv_seq_len],
                              dtype=torch.int32,
                              device=query.device)

    if vertical_indices_count is not None and slash_indices_count is not None:
        (
            block_count,
            block_offset,
            column_count,
            column_index,
        ) = ops.convert_vertical_slash_indexes_mergehead(
            q_seqlens, kv_seqlens, v_idx, s_idx, vertical_indices_count,
            slash_indices_count, context_size, block_size_M, block_size_N,
            causal)
    else:
        (
            block_count,
            block_offset,
            column_count,
            column_index,
        ) = ops.convert_vertical_slash_indexes(q_seqlens, kv_seqlens, v_idx,
                                               s_idx, context_size,
                                               block_size_M, block_size_N,
                                               causal)

    q = query.transpose(1, 2).contiguous()
    k = key.transpose(1, 2).contiguous()
    v = value.transpose(1, 2).contiguous()
    out, lse = sparse_attn_func(
        q,
        k,
        v,
        block_count,
        block_offset,
        column_count,
        column_index,
        causal=causal,
        softmax_scale=softmax_scale,
        return_softmax_lse=True,
    )
    out = out.transpose(1, 2).contiguous()
    softmax_lse = lse.reshape(*lse.shape, 1)
    return (out[..., :context_size, :head_dim],
            softmax_lse[..., :context_size, :])


def _sum_all_diagonal_matrix(mat: torch.tensor):
    h, n, m = mat.shape
    # Zero matrix used for padding
    zero_mat = torch.zeros((h, n, n), device=mat.device)
    # pads the matrix on left and right
    mat_padded = torch.cat((zero_mat, mat, zero_mat), -1)
    # Change the strides
    mat_strided = mat_padded.as_strided((1, n, n + m),
                                        (n * (2 * n + m), 2 * n + m + 1, 1))
    # Sums the resulting matrix's columns
    sum_diags = torch.sum(mat_strided, 1)
    return sum_diags[:, 1:]  # drop left bottom corner


def _get_block(block_table: torch.Tensor, block_size: int, begin: int,
               end: int):
    begin_block = begin // block_size
    end_block = (end - 1) // block_size + 1
    return block_table[begin_block:end_block]
