from enum import Enum
from typing import Callable, Dict, Optional, Tuple

import torch
import torch.distributed
import torch.distributed as dist
import torch.nn.functional as F
from einops import rearrange

from colossalai.kernel.kernel_loader import (
    FlashAttentionForFloatAndCustomMaskLoader,
    FlashAttentionLoader,
    FlashAttentionWithCustomMaskLoader,
    KernelLoader,
)
from colossalai.logging import get_dist_logger

from .utils import RingComm, get_half_index, split_varlen_zigzag

__all__ = [
    "AttnMaskType",
    "ColoAttention",
]

_flash_attn_forward = _flash_attn_backward = None
_unpad_input = _pad_input = None


class AttnMaskType(Enum):
    CUSTOM = 0
    PADDED = 1
    CAUSAL = 2
    PADDED_CAUSAL = 3


def invert_mask(mask: torch.Tensor) -> torch.Tensor:
    """Invert the mask tensor.

    Args:
        mask (torch.Tensor): Mask tensor. Shape should be [B, 1, Sq, Skv]

    Returns:
        torch.Tensor: Inverted mask tensor.
    """
    inverted_mask = 1.0 - mask
    return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(mask.dtype).min)


# adapted from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py
def get_pad_info(
    padding_mask: torch.Tensor, invert: Optional[bool] = False, return_indices: Optional[bool] = True
) -> Tuple[int, torch.Tensor, torch.Tensor]:
    """Get padding information from padding mask.

    Args:
        padding_mask (torch.Tensor): Padding mask tensor. Shape should be [B, Skv]
        invert (Optional[bool], optional): Whether to reverse the padding mask.
        return_indices (Optional[bool], optional): Whether to return the indices of non-masked tokens.

    Returns:
        max_seqlen_in_batch (int): Maximum sequence length in the batch.
        cu_seqlens (torch.Tensor): Shape [B+1]. Cumulative sequence lengths of the sequences in the batch.
        indices (torch.Tensor): Shape [total_nonzero]. The indices of non-masked tokens from the flattened input sequence.
    """
    if invert:
        padding_mask = padding_mask.logical_not()
    seqlens_in_batch = padding_mask.sum(dim=-1, dtype=torch.int32)
    if return_indices:
        indices = torch.nonzero(padding_mask.flatten(), as_tuple=False).flatten()

    max_seqlen_in_batch = seqlens_in_batch.max().item()
    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
    if return_indices:
        return max_seqlen_in_batch, cu_seqlens, indices
    return max_seqlen_in_batch, cu_seqlens


class ColoAttention:
    _kernel_dispatch_map: Optional[Dict[torch.dtype, Dict[Optional[AttnMaskType], Callable]]] = None

    @staticmethod
    def _init_kernels_dispatch():
        if ColoAttention._kernel_dispatch_map is None:
            # fp16/bf16
            half_dispatch_map = {
                None: FlashAttentionLoader(),
                AttnMaskType.CUSTOM: FlashAttentionWithCustomMaskLoader(),
                AttnMaskType.PADDED: FlashAttentionLoader(),
                AttnMaskType.CAUSAL: FlashAttentionLoader(),
                AttnMaskType.PADDED_CAUSAL: FlashAttentionLoader(),
            }
            # fp32
            float_dispatch_map = {
                None: FlashAttentionForFloatAndCustomMaskLoader(),
                AttnMaskType.CUSTOM: FlashAttentionForFloatAndCustomMaskLoader(),
                AttnMaskType.PADDED: FlashAttentionForFloatAndCustomMaskLoader(),
                AttnMaskType.CAUSAL: FlashAttentionForFloatAndCustomMaskLoader(),
                AttnMaskType.PADDED_CAUSAL: FlashAttentionForFloatAndCustomMaskLoader(),
            }
            ColoAttention._kernel_dispatch_map = {
                torch.float16: half_dispatch_map,
                torch.bfloat16: half_dispatch_map,
                torch.float32: float_dispatch_map,
            }

    @staticmethod
    def _dispatch_kernel(dtype: torch.dtype, mask_type: Optional[AttnMaskType]) -> Callable:
        ColoAttention._init_kernels_dispatch()
        if (
            dtype not in ColoAttention._kernel_dispatch_map
            or mask_type not in ColoAttention._kernel_dispatch_map[dtype]
        ):
            raise ValueError(
                "FlashAttention kernel is not available for dtype {} and mask_type {}".format(dtype, mask_type)
            )
        # lazy load
        if isinstance(ColoAttention._kernel_dispatch_map[dtype][mask_type], KernelLoader):
            ColoAttention._kernel_dispatch_map[dtype][mask_type] = ColoAttention._kernel_dispatch_map[dtype][
                mask_type
            ].load()
        return ColoAttention._kernel_dispatch_map[dtype][mask_type]

    @staticmethod
    def prepare_attn_kwargs(
        shape_4d: Tuple[int],
        dtype: torch.dtype,
        device: torch.device,
        q_padding_mask: Optional[torch.Tensor] = None,
        kv_padding_mask: Optional[torch.Tensor] = None,
        is_causal: bool = False,
        invert: bool = True,
    ) -> Dict[str, torch.Tensor]:
        """Return a dictionary of keyword arguments for attention function. It supports 4 mask type.
        1. custom mask: no padding mask and is_causal=False, return {}, users should handle attention mask by themselves.
        2. padded mask: recv padding mask and is_causal=False, return {attention_mask, attention_mask_type, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, q_indices, kv_indices}.
        3. causal mask: no padding mask and is_causal=True, return {attention_mask, attention_mask_type}.
        4. padded causal mask: recv padding mask and is_causal=True, return {attention_mask, attention_mask_type, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, q_indices, kv_indices}.

        Args:
            shape_4d (Tuple[int]): Should be (B, 1, Sq, Skv)
            dtype (torch.dtype): Dtype of attention mask, generally should be ``hidden_states.dtype``
            device (torch.device): Device of attention mask, generally should be ``hidden_states.device``
            q_padding_mask (Optional[torch.Tensor], optional): Padding mask of query. It should be a long tensor or int tensor.
                The shape should be [B, Sq]. ``1`` means valid token, and ``0`` means padding token. Defaults to None.
            kv_padding_mask (Optional[torch.Tensor], optional): Padding mask of key and value. It should be a long tensor or int tensor.
                The shape should be [B, Skv]. ``1`` means valid token, and ``0`` means padding token.
                If it's None and ``q_padding_mask`` is not None, it will be set to ``q_padding_mask``. Defaults to None.
            is_causal (bool, optional): Whether to use causal attention mask. Defaults to False.
            invert_mask (bool, optional): Whether to invert the mask. Defaults to True.
        Returns:
            Dict[str, torch.Tensor]: Dictionary of keyword arguments for attention function.
        """
        if q_padding_mask is None and not is_causal:
            return {}
        assert len(shape_4d) == 4 and shape_4d[1] == 1
        b, _, s_q, s_kv = shape_4d
        outputs = {}
        if (q_padding_mask is None or q_padding_mask.bool().all()) and (
            kv_padding_mask is None or kv_padding_mask.bool().all()
        ):
            # no padding
            assert is_causal
            outputs["attention_mask_type"] = AttnMaskType.CAUSAL
            attention_mask = torch.ones(s_q, s_kv, dtype=dtype, device=device)
            if s_q != 1:
                attention_mask = attention_mask.tril(diagonal=0)
            attention_mask = attention_mask.expand(b, s_q, s_kv)
        else:
            max_seqlen_q, cu_seqlens_q, q_indices = get_pad_info(q_padding_mask)
            if kv_padding_mask is None:
                # self attention
                kv_padding_mask = q_padding_mask
                max_seqlen_kv, cu_seqlens_kv, kv_indices = max_seqlen_q, cu_seqlens_q, q_indices
            else:
                max_seqlen_kv, cu_seqlens_kv, kv_indices = get_pad_info(kv_padding_mask)
            assert kv_padding_mask.shape == (
                b,
                s_kv,
            ), f"Padding mask shape {kv_padding_mask.shape} should align with shape 4d ({b}, {s_kv})"
            attention_mask = kv_padding_mask[:, None, :].expand(b, s_q, s_kv).to(dtype=dtype, device=device)
            outputs.update(
                {
                    "cu_seqlens_q": cu_seqlens_q,
                    "cu_seqlens_kv": cu_seqlens_kv,
                    "max_seqlen_q": max_seqlen_q,
                    "max_seqlen_kv": max_seqlen_kv,
                    "q_indices": q_indices,
                    "kv_indices": kv_indices,
                }
            )
            if is_causal:
                outputs["attention_mask_type"] = AttnMaskType.PADDED_CAUSAL
                if s_q != 1:
                    attention_mask = attention_mask * attention_mask.new_ones(s_q, s_kv).tril(diagonal=0)
            else:
                outputs["attention_mask_type"] = AttnMaskType.PADDED
        if invert:
            attention_mask = invert_mask(attention_mask).unsqueeze(1)
        outputs["attention_mask"] = attention_mask
        return outputs

    @staticmethod
    def attention(
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        attention_mask_type: AttnMaskType = AttnMaskType.CUSTOM,
        cu_seqlens_q: Optional[torch.Tensor] = None,
        cu_seqlens_kv: Optional[torch.Tensor] = None,
        max_seqlen_q: Optional[int] = None,
        max_seqlen_kv: Optional[int] = None,
        q_indices: Optional[torch.Tensor] = None,
        kv_indices: Optional[torch.Tensor] = None,
        dropout_p: float = 0.0,
        scale: Optional[float] = None,
        **kwargs,
    ) -> torch.Tensor:
        """Flash Attention function. It supports 4 mask type.
        1. custom mask: recv attention_mask
        2. padded mask: recv attention_mask, attention_mask_type, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, indices
        3. causal mask: recv attention_mask, attention_mask_type
        4. padded causal mask: recv attention_mask, attention_mask_type, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, indices

        Args:
            q (torch.Tensor): Query tensor. Shape should be [B, nHeads, Sq, D]
            k (torch.Tensor): Key tensor. Shape should be [B, nHeads, Skv, D]
            v (torch.Tensor): Value tensor. Shape should be [B, nHeads, Skv, D]
            attention_mask (Optional[torch.Tensor], optional): Attention mask tensor. Shape should be [B, 1, Sq, Skv]. Defaults to None.
            attention_mask_type (AttnMaskType, optional): Attention mask type. Defaults to AttnMaskType.CUSTOM.
            cu_seqlens_q (Optional[torch.Tensor], optional): The cumulative sequence lengths
                of the sequences in the batch, used to index into q.
                Shape should be [B+1]. Defaults to None.
            cu_seqlens_kv (Optional[torch.Tensor], optional): The cumulative sequence lengths
                of the sequences in the batch, used to index into kv.
                Shape should be [B+1]. Defaults to None.
            max_seqlen_q (Optional[int], optional): Maximum query sequence length in the batch. Defaults to None.
            max_seqlen_kv (Optional[int], optional): Maximum key/value sequence length in the batch. Defaults to None.
            indices (Optional[torch.Tensor], optional): The indices of non-masked tokens from the flattened input sequence.
                Shape should be [NUM_TOKENS]. Defaults to None.
            dropout_p (float, optional): Dropout probability. Defaults to 0.0.
            scale (Optional[float], optional): Scaling factor applied prior to softmax. Defaults to None.

        Returns:
            torch.Tensor: Output tensor. Shape should be [B, nHeads, Sq, D]
        """
        # known issue: sdpa does not support attention mask which contains whole row of masked tokens, which leads to nan
        # this case is usaul when padding mask is used and self attention is performed
        # thus, we don't use sdpa when padding mask is used
        # sanity check
        if attention_mask is not None:
            assert torch.is_floating_point(attention_mask), "attention_mask should be a floating point tensor."
            if attention_mask_type in (AttnMaskType.CUSTOM, AttnMaskType.CAUSAL):
                assert (
                    cu_seqlens_q is None
                    and cu_seqlens_kv is None
                    and max_seqlen_q is None
                    and max_seqlen_kv is None
                    and q_indices is None
                    and kv_indices is None
                )
                if attention_mask_type == AttnMaskType.CUSTOM:
                    assert not torch.all(attention_mask != 0, dim=-1).any()
            elif attention_mask_type in (
                AttnMaskType.PADDED,
                AttnMaskType.PADDED_CAUSAL,
            ):
                assert (
                    cu_seqlens_q is not None
                    and cu_seqlens_kv is not None
                    and max_seqlen_q is not None
                    and max_seqlen_kv is not None
                    and q_indices is not None
                    and kv_indices is not None
                )
        else:
            # if attention_mask is None, attention_mask_type should be the default value
            assert attention_mask_type == AttnMaskType.CUSTOM

        # kernel dispatch
        mask_type = attention_mask_type if attention_mask is not None else None
        attn_func = ColoAttention._dispatch_kernel(q.dtype, mask_type)
        is_causal = attention_mask is not None and attention_mask_type in (
            AttnMaskType.CAUSAL,
            AttnMaskType.PADDED_CAUSAL,
        )
        return attn_func(
            q,
            k,
            v,
            dropout_p=dropout_p,
            scale=scale,
            attention_mask=attention_mask,
            is_causal=is_causal,
            cu_seqlens_q=cu_seqlens_q,
            cu_seqlens_kv=cu_seqlens_kv,
            max_seqlen_q=max_seqlen_q,
            max_seqlen_kv=max_seqlen_kv,
            q_indices=q_indices,
            kv_indices=kv_indices,
        )


def _load_varlen_helpers():
    """Helper to load functions for padding and unpadding packed sequences.
    Use only when flash attn is installed
    """
    global _pad_input, _unpad_input
    # Flash attn claims this is more efficient than torch's bool indexing due to avoiding
    # broadcast
    if _pad_input is None or _unpad_input is None:
        try:
            from flash_attn.bert_padding import index_first_axis, pad_input

            def unpad_input(hidden_states: torch.Tensor, indices: torch.Tensor):
                return index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices)

            _pad_input = pad_input
            _unpad_input = unpad_input
        except ImportError as e:
            raise RuntimeError(
                f"Flash Attention is not installed. You can install it via 'pip install flash-attn --no-build-isolation'"
            ) from e


def _load_flash_attn():
    """A light-weight loader to check whether flash-attn is installed.
    Can't use ColoAttention._dispatch_kernel because we mutate the backward pass
    """
    global _flash_attn_forward, _flash_attn_backward
    if _flash_attn_forward is None or _flash_attn_backward is None:
        try:
            from flash_attn.flash_attn_interface import _flash_attn_varlen_backward as _flash_attn_backward
            from flash_attn.flash_attn_interface import _flash_attn_varlen_forward as _flash_attn_forward
        except ImportError as e:
            raise RuntimeError(
                f"Flash Attention is not installed. You can install it via 'pip install flash-attn --no-build-isolation'"
            ) from e

    _load_varlen_helpers()


# NOTE: This can cause spawned processes to hang on exit
# with python 3.9
@torch.compile()
def _rescale_out_lse(out, block_out, lse, block_lse):
    """
    Compute the new attention denominator:
        exp(lse) + exp(block_lse) = exp(max_scale) * (exp(min_scale - max_scale) + 1)
    Args:
        out: (T, H, D)
        block_out: (T, H, D)
        lse: (H, T, 1)
        block_lse: (H, T, 1)
    """

    # min_scale = torch.min(lse, block_lse)
    # max_scale = torch.max(lse, block_lse)
    # new_lse = max_scale + torch.log(1 + torch.exp(min_scale - max_scale))

    # NOTE: directly assigning to .data here is buggy
    # probably due to casting dtypes/strides
    new_lse = lse + torch.log(1 + torch.exp(block_lse - lse))

    new_block_lse = torch.exp(block_lse - new_lse)
    out = (torch.exp(lse - new_lse) * out + new_block_lse * block_out).to(out)
    lse = new_lse

    # Equivalent to the above
    # See https://github.com/zhuzilin/ring-flash-attention/pull/34#issuecomment-2076126795
    # out = (out - F.sigmoid(block_lse - lse) * (out - block_out))
    # lse = (lse - F.logsigmoid(lse - block_lse))
    return out, lse


class RingAttention(torch.autograd.Function):
    """Implements the Ring Attention from `Ring Attention with Blockwise Transformers for Near-Infinite Context`
    (https://arxiv.org/abs/2310.01889).
    For load-balancing we adopted the "zigzag" attention scheme from https://github.com/zhuzilin/ring-flash-attention/tree/main
    For portable integration with more models, we don't follow the spirit of "block-wise FNN" in the original paper,
    which requires fusing FFN with the Flash Attention kernel/function (see https://arxiv.org/pdf/2305.19370;
    implemented in Jax and not optimized).
    We adopt the double ring topology from LoongTrain (https://arxiv.org/pdf/2406.18485) to fully utilize available
    NICs on each node, by computing attention within a inner ring first and then sending all KVs to the next
    ring at once.
    """

    # Globle cache to avoid recomputation for same-lengthed sequences
    CU_SEQLENS: torch.Tensor = None  # [B+1]
    TOTAL_SEQLEN: int = None
    HALF_INDICES: Tuple = None
    SUPPORTED_MASK_TYPES = (AttnMaskType.CAUSAL, AttnMaskType.PADDED_CAUSAL)
    ATTN_DONE: torch.cuda.Event = None
    SP_STREAM: torch.cuda.Stream = None
    SP_GROUP: dist.ProcessGroup = None
    # duplicate process group for concurrent NCCL streams
    # while both PyTorch and NCCL warns(https://github.com/pytorch/pytorch/commit/2dbe5cb979f674f0052a8eea1f7b6c3c0ba441d7)
    # against this, in practice it seems to work fine.
    INNER_RING_GROUP: dist.ProcessGroup = None
    INNER_RING_GROUP_COPY: dist.ProcessGroup = None
    INTER_RING_GROUP: dist.ProcessGroup = None
    INTER_RING_GROUP_COPY: dist.ProcessGroup = None

    @staticmethod
    def get_double_ring_groups(sp_group, inner_ring_size=None):
        """
        Get 2D ring groups for the given process group. Generally, to avoid congestion, the inner ring size
        shouldn't be larger than the number of NICs on each node.
        Args:
            sp_group (dist.ProcessGroup): Process group for sequence parallelism
            inner_ring_size (Optional[int], optional): Inner ring size. Defaults to None.
        Returns:
            Tuple[dist.ProcessGroup, dist.ProcessGroup]: Inner-ring process group and inter-ring process group.
        """
        sp_size = dist.get_world_size(sp_group)
        sp_rank = dist.get_rank(sp_group)

        if inner_ring_size is None:
            if torch.cuda.device_count() >= dist.get_world_size():
                # single node, no need to consider NICs
                return sp_group, sp_group
            if sp_size <= 4:
                inner_ring_size = min(2, sp_size)
            else:
                inner_ring_size = min(4, sp_size)
        else:
            assert (
                inner_ring_size <= sp_size and sp_size % inner_ring_size == 0
            ), f"Error: sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"

        if inner_ring_size == sp_size:
            return sp_group, sp_group
        assert (
            sp_size % inner_ring_size == 0
        ), f"sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"

        logger = get_dist_logger()
        logger.info(
            f"Using 2D Ring Attention with inner ring size {inner_ring_size} to maximze NIC util for inter-node comm. Cross your fingers for speed-ups!",
            ranks=[0],
        )
        num_rings = sp_size // inner_ring_size
        inner_ring_group = None
        inter_ring_group = None

        # Create inner ring groups
        for i in range(inner_ring_size):
            ranks = list(range(i * inner_ring_size, (i + 1) * inner_ring_size))
            group = dist.new_group(ranks)
            if sp_rank in ranks:
                inner_ring_group = group

        # Create inter ring groups
        for i in range(num_rings):
            ranks = list(range(i, sp_size, num_rings))
            group = dist.new_group(ranks)
            if sp_rank in ranks:
                inter_ring_group = group

        return inner_ring_group, inter_ring_group

    @staticmethod
    def attention(
        q,  # (B, H, Sq, D)
        k,
        v,
        sp_group,
        attention_mask_type,
        cu_seqlens=None,
        max_seqlen=None,
        valid_indices=None,
        dropout_p=0.0,
        softmax_scale=None,
        deterministic=False,
        return_softmax=False,
        inner_ring_size=None,
        **kwargs,
    ):
        """
        Ring Attention forward pass supporting variable-length sequences. When using varlen mode,
        each sequence in the batch should have length divisible by sp_size * 2.

        Args:
            q (torch.Tensor): Query tensor. Shape should be [B, nHeads, Sq, D]
            k (torch.Tensor): Key tensor. Shape should be [B, nHeads, Sq, Sq, D]
            v (torch.Tensor): Value tensor. Shape should be [B, nHeads, Sq, Sq, D]
            sp_group (Optional[dist.ProcessGroup]): Process group for sequence parallelism
            sp_tream (torch.cuda.Stream): An different stream for output correction.
            cu_seqlens (Optional[torch.Tensor], optional): The cumulative sequence lengths
                of the sequences in the batch, used to index into q.
                Shape should be [B+1].
            max_seqlen (Optional[int], optional): Maximum query sequence length in the batch.
            valid_indices (Optional[torch.Tensor], optional): The indices of non-masked tokens from get_pad_info.
                Shape should be [t].
            dropout_p (float, optional): Dropout probability. Defaults to 0.0.
            softmax_scale (Optional[float], optional): Scaling factor applied prior to softmax.
            deterministic (bool, optional): Whether to force deterministic backward pass. See https://github.com/Dao-AILab/flash-attention/issues/349
            return_softmax (bool, optional): Whether to return the softmax denominator (logsumexp).
            inner_ring_size (Optional[int], optional): Inner ring size of the 2D ring. By default use a heuristic to decide.

        Returns:
            out: Output tensor of shape [B, nHeads, Sq, D] or [T, nHeads, D] if pad_output is False.
            softmax_lse: (if return_softmax is True) Softmax denominator (logsumexp).
                Shape should be [total_q_seqlen, nHeads]
        """
        # Check input args
        _load_flash_attn()
        if RingAttention.ATTN_DONE is None:
            RingAttention.ATTN_DONE = torch.cuda.Event()
        if RingAttention.SP_STREAM is None:
            RingAttention.SP_STREAM = torch.cuda.Stream()

        assert (
            q.shape[2] == k.shape[2]
        ), "Q, K and V having different sequence lengths (inference or cross-attn)\
            is not supported yet in training."
        assert (
            attention_mask_type in RingAttention.SUPPORTED_MASK_TYPES
        ), f"Mask type {attention_mask_type} is not supported yet."

        clone_pg = lambda pg: dist.new_group(dist.get_process_group_ranks(pg))

        if RingAttention.SP_GROUP is not sp_group:
            RingAttention.SP_GROUP = sp_group
            inner_ring_group, inter_ring_group = RingAttention.get_double_ring_groups(sp_group, inner_ring_size)
            RingAttention.INNER_RING_GROUP = inner_ring_group
            RingAttention.INTER_RING_GROUP = inter_ring_group
        else:
            inner_ring_group = RingAttention.INNER_RING_GROUP
            inter_ring_group = RingAttention.INTER_RING_GROUP

        # (B, H, Sq, D) -> (B, Sq, H, D)
        q, k, v = [x.transpose(1, 2).contiguous() for x in (q, k, v)]
        pad_output = q.dim() == 4

        # Get sequence length info for varlen forward
        if attention_mask_type == AttnMaskType.CAUSAL:
            # All sequences share the same length
            b, sq, h, d = q.shape
            max_seqlen = sq
            # Cache to avoid recreation for a single sequence
            if sq * b == RingAttention.TOTAL_SEQLEN:
                cu_seqlens = RingAttention.CU_SEQLENS
            else:
                cu_seqlens = torch.arange(0, b * sq + 1, sq, device=q.device, dtype=torch.int32)
                RingAttention.TOTAL_SEQLEN = b * sq

        # "Packed" mode where sequences of different lengths are packed into [total_q_seqlen, H, D]
        elif attention_mask_type == AttnMaskType.PADDED_CAUSAL:
            assert (
                cu_seqlens is not None and max_seqlen is not None and valid_indices is not None
            ), "Packed mode requires pre-computed cu_seqlens and max_seq_len."
            if pad_output:
                b, sq, h, d = q.shape
                q, k, v = [_unpad_input(x, valid_indices) for x in (q, k, v)]

        out, softmax_lse = RingAttention.apply(
            q,
            k,
            v,
            sp_group,
            RingAttention.SP_STREAM,
            cu_seqlens,
            max_seqlen,
            dropout_p,
            softmax_scale,
            deterministic,
            return_softmax,
            attention_mask_type == AttnMaskType.PADDED_CAUSAL,
            inner_ring_group,
            inter_ring_group,
        )

        if attention_mask_type == AttnMaskType.PADDED_CAUSAL:
            if pad_output:
                out = _pad_input(out, valid_indices, b, sq)  # (T, ...) -> (B, Sq, ...)
                out = out.transpose(1, 2)  # (B, Sq, H, D) -> (B, H, Sq, D)
        else:
            out = out.transpose(1, 2)

        if return_softmax:
            return out, softmax_lse
        return out

    @staticmethod
    def forward(
        ctx,
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        sp_group: dist.ProcessGroup,
        sp_stream: torch.cuda.Stream,
        cu_seqlens: torch.Tensor,
        max_seqlen: int,
        dropout_p: float = 0.0,
        softmax_scale: Optional[float] = None,
        deterministic: Optional[bool] = False,
        return_softmax: Optional[bool] = False,
        is_packed: Optional[bool] = False,
        inner_ring_group: Optional[dist.ProcessGroup] = None,
        inter_ring_group: Optional[dist.ProcessGroup] = None,
    ):

        cu_seqlens_q = cu_seqlens_kv = cu_seqlens
        max_seqlen_q = max_seqlen_kv = max_seqlen
        cu_seqlens_half = cu_seqlens // 2
        max_seqlen_half = max_seqlen // 2

        misc_kwargs = {
            "window_size": (-1, -1),
            "alibi_slopes": None,
            "softmax_scale": q.shape[-1] ** -0.5 if softmax_scale is None else softmax_scale,
            "dropout_p": dropout_p,
            "block_table": None,
            "softcap": 0.0,
            "return_softmax": False,
        }

        if (
            RingAttention.HALF_INDICES is not None
            and cu_seqlens.shape == RingAttention.CU_SEQLENS.shape
            and (cu_seqlens == RingAttention.CU_SEQLENS).all()
        ):
            half_idx_front, half_idx_back = RingAttention.HALF_INDICES
        else:
            half_idx_front = get_half_index(cu_seqlens, front=True)
            half_idx_back = get_half_index(cu_seqlens, front=False)
            RingAttention.HALF_INDICES = (half_idx_front, half_idx_back)
            RingAttention.CU_SEQLENS = cu_seqlens

        if is_packed:
            t, h, d = q.shape
        else:
            b, sq, h, d = q.shape
            t = b * sq
            # Be careful about GQA/MQA in reshape
            q, k, v = [x.view(t, *x.shape[-2:]) for x in (q, k, v)]

        if inner_ring_group is None or inter_ring_group is None:
            # Use one ring if not specified
            inner_ring_group = inter_ring_group = sp_group

        sp_size = dist.get_world_size(sp_group)
        sp_rank = dist.get_rank(sp_group)
        # Attempt to achieve concurrent comm in the two-stream forward
        local_kv_comms = [RingComm(inner_ring_group) for _ in range(2)]
        inter_ring_comm = RingComm(inter_ring_group)
        local_sp_size = dist.get_world_size(inner_ring_group)
        local_sp_rank = dist.get_rank(inner_ring_group)
        inter_ring_rank = dist.get_rank(inter_ring_group) if inter_ring_group is not sp_group else 0
        num_rings = dist.get_world_size(inter_ring_group) if inter_ring_group is not sp_group else 1

        # Non-contiguous indexing copies to a new contiguous tensor,
        # so only do it once
        if sp_rank != sp_size - 1:
            q1 = q[half_idx_back]

        # Pre-allocate double buffer for overlapping and receiving next step's inputs
        kv_buffers = [torch.stack((k, v))]  # (2, B, Sq, H, D)
        kv_buffers.append(torch.empty_like(kv_buffers[0]))

        # outputs
        out = None
        block_out = [None, None]
        softmax_lse = [None, None]
        block_softmax_lse = [None, None]  # log sum exp, the denominator of softmax in attention
        rng_states = [None for _ in range(sp_size)]
        sp_streams = [torch.cuda.current_stream(), sp_stream]

        def _forward(q, k, v, causal):
            (
                _,
                _,
                _,
                _,
                out,
                softmax_lse,
                _,
                rng_state,
            ) = _flash_attn_forward(
                q,
                k,
                v,
                cu_seqlens_q if q.shape[0] == t else cu_seqlens_half,
                cu_seqlens_kv if k.shape[0] == t else cu_seqlens_half,
                max_seqlen_q if q.shape[0] == t else max_seqlen_half,
                max_seqlen_kv if k.shape[0] == t else max_seqlen_half,
                causal=causal,
                **misc_kwargs,
            )
            return out, softmax_lse, rng_state

        def _kv_comm(i):
            # Avoid overwriting attn input when it shares mem with buffer
            if not RingAttention.ATTN_DONE.query():
                kv_buffers[(i + 1) % 2] = torch.empty_like(kv_buffers[i % 2])
            if i < local_sp_size - 1:
                local_kv_comms[i % 2].send_recv(kv_buffers[i % 2], kv_buffers[(i + 1) % 2])

        def _local_ring_forward():
            # (Hopefully) overlap output correction with next flash attn
            for i in range(local_sp_size):
                with torch.cuda.stream(sp_streams[i % 2]):
                    # Wait for current kv from prev rank
                    # NOTE: waiting outside the current stream will NOT correctly synchronize.
                    if i > 0:
                        local_kv_comms[(i + 1) % 2].wait()
                    if i == 0:
                        _kv_comm(i)

                    if i == 0:
                        # Compute with local KV; no mask
                        kv_block = kv_buffers[0]
                        q_block = q
                        (block_out[i % 2], block_softmax_lse[i % 2], rng_states[i]) = _forward(  # (T, H, D)  # (H, T)
                            q_block, kv_block[0], kv_block[1], causal=True
                        )
                    elif i <= local_sp_rank:
                        # Received the "surrounding" kv chunks
                        # Drop the second half of received kv
                        # (2, t // 2, H, D)
                        kv_block = kv_buffers[i % 2][:, half_idx_front]
                        q_block = q
                        (
                            block_out[i % 2],  # (T, H, D)
                            block_softmax_lse[i % 2],  # (H, T)
                            rng_states[i],
                        ) = _forward(q_block, kv_block[0], kv_block[1], causal=False)
                    else:
                        # Received the inner kv chunks
                        # Drop the first half of q
                        kv_block = kv_buffers[i % 2]
                        q_block = q1
                        (
                            block_out[i % 2],  # (T, H, D)
                            block_softmax_lse[i % 2],  # (H, T)
                            rng_states[i],
                        ) = _forward(q_block, kv_block[0], kv_block[1], causal=False)
                    RingAttention.ATTN_DONE.record()
                    # Pipeline the next KV comm with output correction instead of the next flash attn
                    # to minimize idle time when comm takes longer than attn.
                    _kv_comm(i + 1)

                    block_softmax_lse[i % 2] = (
                        block_softmax_lse[i % 2].transpose(0, 1).unsqueeze(-1).contiguous().float()
                    )  # (H, T) -> (T, H, 1)
                    assert block_out[i % 2].shape[:-1] == block_softmax_lse[i % 2].shape[:-1]
                    # Output and log sum exp correction. Ideally overlap this with the next flash attn kernel.
                    # In reality this always finishes before next flash attn; no need for extra sync.
                    if i == 0:
                        out = block_out[0]
                        softmax_lse = block_softmax_lse[0]
                    elif i <= local_sp_rank:
                        out, softmax_lse = _rescale_out_lse(
                            out, block_out[i % 2], softmax_lse, block_softmax_lse[i % 2]
                        )
                    else:
                        out[half_idx_back], softmax_lse[half_idx_back] = _rescale_out_lse(
                            out[half_idx_back], block_out[i % 2], softmax_lse[half_idx_back], block_softmax_lse[i % 2]
                        )

            torch.cuda.current_stream().wait_stream(sp_stream)
            return out, softmax_lse

        def _other_ring_forward(ring_num_idx, out, softmax_lse):
            # Loop through the inner ring after receiving
            # all new KVs from the previous inner ring
            for i in range(local_sp_size):
                with torch.cuda.stream(sp_streams[i % 2]):
                    # Send & recv KV
                    if i > 0:
                        local_kv_comms[(i + 1) % 2].wait()

                    if i == 0:
                        _kv_comm(i)

                    if ring_num_idx > inter_ring_rank:
                        kv_block = kv_buffers[i % 2]
                        (
                            block_out[i % 2],
                            block_softmax_lse[i % 2],
                            rng_states[i + local_sp_size * ring_num_idx],
                        ) = _forward(q1, kv_block[0], kv_block[1], causal=False)
                        RingAttention.ATTN_DONE.record()

                        _kv_comm(i + 1)
                        block_softmax_lse[i % 2] = (
                            block_softmax_lse[i % 2].transpose(0, 1).unsqueeze(-1).contiguous().float()
                        )
                        out[half_idx_back], softmax_lse[half_idx_back] = _rescale_out_lse(
                            out[half_idx_back], block_out[i % 2], softmax_lse[half_idx_back], block_softmax_lse[i % 2]
                        )
                    else:
                        kv_block = kv_buffers[i % 2][:, half_idx_front]
                        (
                            block_out[i % 2],
                            block_softmax_lse[i % 2],
                            rng_states[i + local_sp_size * ring_num_idx],
                        ) = _forward(q, kv_block[0], kv_block[1], causal=False)
                        RingAttention.ATTN_DONE.record()

                        _kv_comm(i + 1)
                        block_softmax_lse[i % 2] = (
                            block_softmax_lse[i % 2].transpose(0, 1).unsqueeze(-1).contiguous().float()
                        )
                        out, softmax_lse = _rescale_out_lse(
                            out, block_out[i % 2], softmax_lse, block_softmax_lse[i % 2]
                        )

            torch.cuda.current_stream().wait_stream(sp_stream)
            return out, softmax_lse

        # Send and recv KV between rings at once to maximize NIC util.
        inter_ring_kv = None
        for ring_num_idx in range(num_rings):
            if ring_num_idx > 0:
                inter_ring_comm.wait()
                # Reset indices
                kv_buffers[0] = inter_ring_kv

            if ring_num_idx < num_rings - 1:
                if ring_num_idx == 0:
                    to_send = kv_buffers[0]
                else:
                    # The last received KV
                    to_send = kv_buffers[(local_sp_size - 1) % 2]
                inter_ring_kv = inter_ring_comm.send_recv(to_send)

            if ring_num_idx == 0:
                out, softmax_lse = _local_ring_forward()
            else:
                out, softmax_lse = _other_ring_forward(ring_num_idx, out, softmax_lse)

        out = out.to(q.dtype)
        if not is_packed:
            out = out.view(b, sq, h, d)
            q, k, v = [x.view(b, sq, *x.shape[-2:]) for x in (q, k, v)]  # (T, H, D) -> (B, Sq, H, D)
        softmax_lse = softmax_lse.squeeze(-1)

        ctx.sp_group = sp_group
        ctx.max_seqlen_q = ctx.max_seqlen_kv = max_seqlen
        misc_kwargs["deterministic"] = deterministic
        del misc_kwargs["return_softmax"]
        ctx.misc_kwargs = misc_kwargs
        ctx.is_packed = is_packed

        ctx.kv_group = inner_ring_group
        ctx.inter_kv_group = inter_ring_group

        ctx.save_for_backward(
            q,
            k,
            v,
            out,
            softmax_lse.transpose(0, 1).contiguous(),  # (T, H) -> (H, T)
            cu_seqlens_q,
            cu_seqlens_kv,
            half_idx_front,
            half_idx_back,
            *rng_states,
        )

        if return_softmax:
            return out, softmax_lse
        return out, None

    def backward(ctx, dout, _):
        """
        During backward, we accumulate q grads on each rank locally, but iterate kv and their grads
        over all ranks for accumulation.
        """
        (q, k, v, out, softmax_lse, cu_seqlens_q, cu_seqlens_kv, half_idx_front, half_idx_back) = ctx.saved_tensors[:9]
        rng_states = ctx.saved_tensors[9:]

        is_packed = ctx.is_packed
        max_seqlen_q = ctx.max_seqlen_q
        max_seqlen_kv = ctx.max_seqlen_kv
        cu_seqlens_half = cu_seqlens_q // 2
        max_seqlen_half = max_seqlen_q // 2
        misc_kwargs = ctx.misc_kwargs
        del misc_kwargs["block_table"]

        assert (
            out.shape == dout.shape == q.shape
        ), f"out {out.shape} and dout {dout.shape} should have the same shape ({q.shape})."

        if is_packed:
            t, h, d = q.shape
        else:
            b, sq, h, d = q.shape
            t = b * sq
        q, k, v, out, dout = [x.view(t, *x.shape[-2:]) for x in (q, k, v, out, dout)]

        # Sequence parallel args
        sp_group = ctx.sp_group
        local_kv_group = ctx.kv_group
        inter_kv_group = ctx.inter_kv_group

        local_sp_rank = dist.get_rank(sp_group)
        sp_size = dist.get_world_size(sp_group)
        # Using separate streams (pg) for concurrent kv and dkv comm may
        # cause NCCL "software caused connection abort" here...
        local_kv_comm = RingComm(local_kv_group)
        local_dkv_comm = RingComm(local_kv_group)
        inter_kv_comm = RingComm(inter_kv_group)
        inter_dkv_comm = RingComm(inter_kv_group)
        local_sp_size = dist.get_world_size(local_kv_group)
        local_sp_rank = dist.get_rank(local_kv_group)

        if dist.get_world_size(inter_kv_group) != sp_size:
            num_rings = dist.get_world_size(inter_kv_group)
            inter_ring_rank = dist.get_rank(inter_kv_group)
        else:
            num_rings = 1
            inter_ring_rank = 0

        if local_sp_rank != sp_size - 1:
            softmax_lse1 = softmax_lse[:, half_idx_back]
        dout = dout.contiguous()

        # Double comm buffers for sending and receiving kv
        kv_buffers = [torch.stack((k, v))]  # (2, T, H, D)
        kv_buffers.append(torch.empty_like(kv_buffers[0]))

        dq = None  # (T, H, D)
        # Intermediate outputs
        dq_block = torch.empty_like(q)  # (T, H, D)
        dk_block = torch.empty_like(k)  # (T, H, D)
        dv_block = torch.empty_like(v)  # (T, H, D)
        dkv_buffers = [torch.empty_like(kv, dtype=torch.float32) for kv in kv_buffers]  # (T, H, D)
        del k, v

        def _backward(dout, q, k, v, out, softmax_lse, dq, dk, dv, rng_state, causal):
            _flash_attn_backward(
                dout,
                q,
                k,
                v,
                out,
                softmax_lse,
                dq,
                dk,
                dv,
                cu_seqlens_q if dq.shape[0] == t else cu_seqlens_half,
                cu_seqlens_kv if dk.shape[0] == t else cu_seqlens_half,
                max_seqlen_q if dq.shape[0] == t else max_seqlen_half,
                max_seqlen_kv if dk.shape[0] == t else max_seqlen_half,
                causal=causal,
                rng_state=rng_state,
                **misc_kwargs,
            )

        # NOTE: We avoid using two streams due to doubled buffers
        # and that backward is more communication intensive.
        def _local_ring_backward():
            for i in range(local_sp_size):
                if i > 0:
                    local_kv_comm.wait()

                if i < local_sp_size - 1:
                    # Send kv to next rank for backward
                    local_kv_comm.send_recv(kv_buffers[i % 2], kv_buffers[(i + 1) % 2])

                if i == 0:
                    # Backward with local kv
                    k_, v_ = kv_buffers[i % 2]
                    q_, dout_, out_ = q, dout, out
                    dq_, dk_, dv_ = dq_block, dk_block, dv_block
                    _backward(dout_, q_, k_, v_, out_, softmax_lse, dq_, dk_, dv_, rng_states[i], causal=True)

                elif i <= local_sp_rank:
                    # Drop the second half of kv
                    # (T, H, D) -> (T // 2, H, D)
                    k_, v_ = [x[half_idx_front] for x in kv_buffers[i % 2]]
                    dk_, dv_ = [x[: t // 2] for x in (dk_block, dv_block)]
                    dq_, q_, out_, dout_ = (dq_block, q, out, dout)
                    _backward(dout_, q_, k_, v_, out_, softmax_lse, dq_, dk_, dv_, rng_states[i], causal=False)

                else:
                    # Drop the first half of q
                    k_, v_ = kv_buffers[i % 2]
                    dk_, dv_ = dk_block, dv_block
                    q_, out_, dout_ = [x[half_idx_back] for x in (q, out, dout)]
                    dq_ = dq_block[: t // 2]
                    _backward(dout_, q_, k_, v_, out_, softmax_lse1, dq_, dk_, dv_, rng_states[i], causal=False)

                # Accumulate grads
                if i == 0:
                    dq = dq_block.float()
                    dkv_buffers[i % 2][0] = dk_block.float()
                    dkv_buffers[i % 2][1] = dv_block.float()
                else:
                    # Accumulate local dq
                    if i <= local_sp_rank:
                        dq += dq_  # (T, H, D)
                    else:
                        dq[half_idx_back] += dq_

                    # Wait for mobile kv grad accumulators
                    local_dkv_comm.wait()

                    if i <= local_sp_rank:
                        # q blocks "surrounded" by kv blocks
                        dkv_buffers[i % 2][0][half_idx_front] += dk_
                        dkv_buffers[i % 2][1][half_idx_front] += dv_
                    else:
                        # q blocks "surrounding" kv blocks
                        dkv_buffers[i % 2][0] += dk_
                        dkv_buffers[i % 2][1] += dv_
                local_dkv_comm.send_recv(send_tensor=dkv_buffers[i % 2], recv_tensor=dkv_buffers[(i + 1) % 2])

            local_dkv_comm.wait()
            dkv_recv = dkv_buffers[local_sp_size % 2]
            dkv_send = dkv_buffers[(local_sp_size - 1) % 2]
            return dq, dkv_recv, dkv_send

        def _other_ring_backward(ring_num_idx, dq):
            if ring_num_idx > inter_ring_rank:
                # Indexing is expensive
                q_, out_, dout_ = [x[half_idx_back] for x in (q, out, dout)]
            else:
                q_, out_, dout_ = (q, out, dout)

            for i in range(local_sp_size):
                if i > 0:
                    local_kv_comm.wait()

                if i < local_sp_size - 1:
                    local_kv_comm.send_recv(kv_buffers[i % 2], kv_buffers[(i + 1) % 2])

                rng_state = rng_states[i + local_sp_size * ring_num_idx]
                if ring_num_idx > inter_ring_rank:
                    k_, v_ = kv_buffers[i % 2]
                    dk_, dv_ = dk_block, dv_block
                    dq_ = dq_block[: t // 2]
                    _backward(dout_, q_, k_, v_, out_, softmax_lse1, dq_, dk_, dv_, rng_state, causal=False)

                    dq[half_idx_back] += dq_
                    if i > 0:
                        local_dkv_comm.wait()
                    else:
                        inter_dkv_comm.wait()

                    dkv_buffers[i % 2][0] += dk_
                    dkv_buffers[i % 2][1] += dv_
                else:
                    k_, v_ = [x[half_idx_front] for x in kv_buffers[i % 2]]
                    dk_, dv_ = [x[: t // 2] for x in (dk_block, dv_block)]
                    dq_ = dq_block
                    _backward(dout_, q_, k_, v_, out_, softmax_lse, dq_, dk_, dv_, rng_state, causal=False)

                    dq += dq_
                    if i > 0:
                        local_dkv_comm.wait()
                    else:
                        inter_dkv_comm.wait()

                    dkv_buffers[i % 2][0][half_idx_front] += dk_
                    dkv_buffers[i % 2][1][half_idx_front] += dv_

                local_dkv_comm.send_recv(send_tensor=dkv_buffers[i % 2], recv_tensor=dkv_buffers[(i + 1) % 2])

            local_dkv_comm.wait()
            dkv_recv = dkv_buffers[local_sp_size % 2]
            dkv_send = dkv_buffers[(local_sp_size - 1) % 2]
            return dq, dkv_recv, dkv_send

        inter_ring_kv = None
        for ring_num_idx in range(num_rings):
            if ring_num_idx > 0:
                inter_kv_comm.wait()
                kv_buffers[0] = inter_ring_kv

            if ring_num_idx < num_rings - 1:
                # Re-allocate a buffer in each inter-ring step
                inter_ring_kv = inter_kv_comm.send_recv(kv_buffers[0])

            if ring_num_idx == 0:
                dq, dkv_recv, dkv_send = _local_ring_backward()
            else:
                dq, dkv_recv, dkv_send = _other_ring_backward(ring_num_idx, dq)

            if num_rings > 1:
                # Reuse the local buffers
                inter_dkv_comm.send_recv(send_tensor=dkv_recv, recv_tensor=dkv_send)
                # Reset indices
                dkv_buffers[0] = dkv_send
                dkv_buffers[1] = dkv_recv
                if ring_num_idx == num_rings - 1:
                    inter_dkv_comm.wait()
                    dkv_recv = dkv_buffers[0]

        dq, dk, dv = [x.to(q.dtype) for x in (dq, *dkv_recv)]
        if not is_packed:
            dq, dk, dv = [x.view(b, sq, *x.shape[-2:]) for x in (dq, dk, dv)]

        return (dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None)

    @staticmethod
    def prepare_varlen_batch(
        attention_mask: torch.Tensor,
        sp_group: dist.ProcessGroup,
        inputs_embeds: torch.Tensor = None,
        position_ids: Optional[torch.Tensor] = None,
        is_label: bool = False,
        is_2d: bool = True,
    ):
        """
        Preprocess a batch of padded sequence by splitting input sequence by sp_size
        sequence-wise and packing them into one sequence. Updates the mask info accordingly.
        Args:
            attention_mask (torch.Tensor): Contains the mask [B, Sq], where True means the token is NOT masked.
            sp_group (dist.ProcessGroup): Process group for sequence parallelism
            inputs_embeds (torch.Tensor): Input embeddings. Shape should be [B, Sq, ...]
            position_ids (Optional[torch.Tensor], optional): Position ids of shape [Sq] or [1, Sq]. Defaults to None.
            is_label (bool, optional): Whether inputs_embeds is instead a label tensor. If True, mask out the first
                token of each sequence.
            is_2d (bool, optional): Whether to return 2D outputs padded to max_seqlen // sp_size or flatten
                the batch dim to a packed 1d sequence. Contingent on model forward shape definitions.

        Returns:
            inputs_embeds: Packed input embeddings of shape [B, Sq // sp_size, ...].
            mask_info: A dictionary of mask info.
            position_ids: Packed position ids of shape [..., Sq // sp_size].

        """
        _load_varlen_helpers()
        sp_size = dist.get_world_size(group=sp_group)
        sp_rank = dist.get_rank(group=sp_group)
        mask_info = {}
        mask_info["max_seqlen"], mask_info["cu_seqlens"] = get_pad_info(attention_mask, return_indices=False)

        # Unpad, split seq-wise, then pad back to (B, max_seqlen // sp_size)
        # Split mask to compute local nonzero position indices
        # (B, Sq) -> (B, max_seqlen // sp_size)
        attention_mask = attention_mask[:, : mask_info["max_seqlen"]]
        if inputs_embeds is not None:
            inputs_embeds = inputs_embeds[:, : mask_info["max_seqlen"]]
            inputs_embeds = split_varlen_zigzag(
                inputs_embeds,
                mask_info["cu_seqlens"],
                sp_group,
                mask_info["max_seqlen"],
                is_2d=is_2d,
                is_label=is_label,
            )
        attention_mask = split_varlen_zigzag(
            attention_mask, mask_info["cu_seqlens"], sp_group, mask_info["max_seqlen"], is_2d=is_2d
        )

        if position_ids is not None:
            indices = torch.tensor([sp_rank, 2 * sp_size - sp_rank - 1], device=inputs_embeds.device)
            position_ids = (
                position_ids[..., : mask_info["max_seqlen"]]  # unpad
                .view(-1, sp_size * 2, mask_info["max_seqlen"] // (sp_size * 2))
                .index_select(-2, indices)
                .view(-1, mask_info["max_seqlen"] // sp_size)
            )

        mask_info["max_seqlen"] //= sp_size
        mask_info["valid_indices"] = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
        mask_info["cu_seqlens"] //= sp_size
        mask_info["attention_mask_type"] = AttnMaskType.PADDED_CAUSAL
        return inputs_embeds, mask_info, position_ids
