# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# See LICENSE for license information.

"""
Utils/Helper classes and methods for attention
"""

import math
import os
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import warnings
import logging
import functools

from dataclasses import dataclass, fields
import numpy as np
from packaging.version import Version as PkgVersion

import torch
import torch.distributed as dist
import torch.nn.functional as F
import transformer_engine_torch as tex
import transformer_engine as te
from transformer_engine.pytorch.cpp_extensions.fused_attn import (
    QKVLayout,
    AttnBiasType,
    AttnMaskType,
    SoftmaxType,
    FusedAttnBackend,
    META_QKV,
    META_DQKV,
    META_O,
    META_DO,
    META_S,
    META_DP,
)
from transformer_engine.pytorch.attention.inference import InferenceParams
from transformer_engine.pytorch.cpu_offload import is_cpu_offload_enabled
from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage
from transformer_engine.pytorch.tensor.float8_tensor import (
    Float8Tensor,
    Float8Quantizer,
    Float8CurrentScalingQuantizer,
)
from transformer_engine.pytorch.tensor.float8_tensor import Float8TensorStorage
from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage

from transformer_engine.pytorch.quantization import get_fp8_te_dtype
from transformer_engine.pytorch.constants import TE_DType, MXFP8_BLOCK_SCALING_SIZE


from transformer_engine.pytorch.utils import (
    get_device_compute_capability,
    get_cudnn_version,
    SplitAlongDim,
    combine_tensors,
)
from transformer_engine.pytorch.export import is_in_onnx_export_mode

from transformer_engine.pytorch.jit import jit_fuser

# NVTE_DEBUG = 0/1 # disables/enables debug mode, default = 0
_NVTE_DEBUG = int(os.getenv("NVTE_DEBUG", "0"))
# NVTE_DEBUG_LEVEL = 0/1/2 # enables more and more verbose debug mode, default = 0
_NVTE_DEBUG_LEVEL = int(os.getenv("NVTE_DEBUG_LEVEL", "0"))
_NVTE_FLASH_ATTN = int(os.getenv("NVTE_FLASH_ATTN", "1"))
# print quantizer info for a particular layer on a particular rank
_print_layer = int(os.getenv("NVTE_PRINT_LAYER_NUMBER", "1"))
_print_rank = int(os.getenv("NVTE_PRINT_RANK", "0"))

_cu_seqlens_cache = {}


class AttentionLogging:
    """
    Manage logging for attention module
    """

    _log_level = _NVTE_DEBUG * _NVTE_DEBUG_LEVEL
    _formatter = logging.Formatter("[%(levelname)-8s | %(name)-19s]: %(message)s")
    _stream_handler = logging.StreamHandler()
    fa_logger = logging.getLogger(__name__)
    _is_logging_setup = False

    @staticmethod
    def setup_logging():
        """
        Set up log levels, logger and handlers
        """
        _log_levels = {0: logging.WARNING, 1: logging.INFO, 2: logging.DEBUG}
        AttentionLogging._log_level = _log_levels[
            AttentionLogging._log_level if AttentionLogging._log_level in [0, 1, 2] else 2
        ]
        AttentionLogging._stream_handler.setFormatter(AttentionLogging._formatter)
        AttentionLogging.fa_logger.setLevel(AttentionLogging._log_level)
        if not AttentionLogging.fa_logger.hasHandlers():
            AttentionLogging.fa_logger.addHandler(AttentionLogging._stream_handler)
        AttentionLogging._is_logging_setup = True


@functools.lru_cache(maxsize=None)
def _get_supported_versions(version_min, version_max):
    """
    Calculate version info based on min and max numbers
    """
    return ">= " + str(version_min) + ", " + "<= " + str(version_max)


def maybe_contiguous(tensor: torch.Tensor) -> torch.Tensor:
    """Make tensor contiguous if final stride is not 1."""
    return tensor.contiguous() if tensor.stride(-1) != 1 else tensor


class FlashAttentionUtils:
    """
    Manage Flash Attention versioning information
    """

    is_installed = False
    version = PkgVersion("0")
    version_required = PkgVersion("2.1.1")
    version_required_blackwell = PkgVersion("2.7.3")
    max_version = PkgVersion("2.8.3")
    v2_plus = False
    v2_1_plus = False
    v2_3_plus = False
    v2_4_plus = False
    v2_4_1_plus = False
    v2_5_plus = False
    v2_5_7_plus = False
    v2_6_0_plus = False
    v2_7_0_plus = False
    warning_printed = False

    v3_is_installed = False
    fa3_version = PkgVersion("0")
    v3_0_0_beta = False
    use_v3 = False
    # FA3 from FA 2.7.3+/hopper has different APIs than FA3 from 2.7.2/hopper
    # Please follow these instructions to install FA3
    v3_installation_steps = """\
(1) git clone https://github.com/Dao-AILab/flash-attention.git
(2) cd flash-attention/hopper && python setup.py install"""
    v3_warning_printed = False

    v4_is_installed = False
    fa4_version = PkgVersion("0")
    use_v4 = False
    v4_installation_steps = """\
pip install flash-attn-4==4.0.0b11 nvidia-cutlass-dsl[cu13]"""
    v4_warning_printed = False
    # Set by backends.py if FA4 is installed; calls flash_attn.cute.interface._validate_head_dims
    # which raises AssertionError for unsupported (head_dim, head_dim_v) combinations.
    v4_validate_head_dims: Callable = None

    @staticmethod
    def set_flash_attention_version():
        """
        Setup version info for FA v2.x
        """
        FlashAttentionUtils.is_installed = True
        FlashAttentionUtils.v2_plus = FlashAttentionUtils.version >= PkgVersion("2")
        FlashAttentionUtils.v2_1_plus = FlashAttentionUtils.version >= PkgVersion("2.1")
        FlashAttentionUtils.v2_3_plus = FlashAttentionUtils.version >= PkgVersion("2.3")
        FlashAttentionUtils.v2_4_plus = FlashAttentionUtils.version >= PkgVersion("2.4")
        FlashAttentionUtils.v2_4_1_plus = FlashAttentionUtils.version >= PkgVersion("2.4.1")
        FlashAttentionUtils.v2_5_plus = FlashAttentionUtils.version >= PkgVersion("2.5.0")
        FlashAttentionUtils.v2_5_7_plus = FlashAttentionUtils.version >= PkgVersion("2.5.7")
        FlashAttentionUtils.v2_6_0_plus = FlashAttentionUtils.version >= PkgVersion("2.6.0")
        FlashAttentionUtils.v2_7_0_plus = FlashAttentionUtils.version >= PkgVersion("2.7.0")

    @staticmethod
    def set_flash_attention_3_params():
        """
        Setup version info for FA v3.x
        """
        FlashAttentionUtils.v3_is_installed = True
        FlashAttentionUtils.v3_0_0_beta = (
            PkgVersion("3.0.0b") < FlashAttentionUtils.fa3_version < PkgVersion("3.0.0")
        )

    @staticmethod
    def set_flash_attention_4_params():
        """
        Setup version info for FA v4.x
        """
        FlashAttentionUtils.v4_is_installed = True


@dataclass(eq=True)
class AttentionParams:
    """
    Attention parameters used to determine which backend to be used.

    Parameters
    ----------
    qkv_type : Union[torch.Tensor, Float8Tensor], default = torch.Tensor
        Type of query/key/value tensors, {`torch.Tensor`, `Float8Tensor`}.
    qkv_dtype : torch.dtype, default = torch.bfloat16
        Data type of query/key/value tensors.
    qkv_layout : str, default = "sbh3d"
        Query/key/value tensor memory layout.
    batch_size : int, default = 1
        Batch size.
    num_heads : int, default = 16
        Number of attention heads in the query tensor.
    num_gqa_groups : int, default = 16
        Number of attention heads in key and value tensors.
    max_seqlen_q : int, default = 128
        Maximum sequence length of the query tensor.
    max_seqlen_kv : int, default = 128
        Maximum sequence length of the key and value tensors.
    head_dim_qk : int, default = 64
        The size of each attention head in query and key tensors.
    head_dim_v : int, default = 64
        The size of each attention head in the value tensor.
    attn_mask_type : str, default = no_mask
        Attention mask type, {`no_mask`, `padding`, `causal`, `padding_causal`,
        `causal_bottom_right`, `padding_causal_bottom_right`, `arbitrary`}
    window_size : Tuple[int, int], default = None
        Sliding window attention size.
    bottom_right_diagonal: bool, default = `None`
        Whether to align sliding window and ALiBi diagonal to the bottom right corner
        of the softmax matrix.
    alibi_slopes_shape : Optional[Union[torch.Size, List]], default = None
        Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`.
    core_attention_bias_type : str, default = no_bias
        Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}.
    core_attention_bias_shape : str, default = 1hss
        Attention bias shape, {`1hss`, `b1ss`, `bhss`}.
    core_attention_bias_requires_grad : bool, default = True
        Whether attention bias requires gradient.
    pad_between_seqs : bool, default = False
        Whether there is padding between sequences in a batch.
        This only applies to `qkv_format=thd`.
    attention_dropout : float, default = 0.0
        Attention dropout.
    context_parallel : bool, default = False
        Whether context parallelism is used or not.
    cp_comm_type : str, default = "p2p"
        The communication type of context parallelism.
    cp_size : int, default = 1
        The group size of context parallelism.
    deterministic : bool, default = False
        Whether to run `DotProductAttention` with determinism or not.
    is_training : bool, default = True
        Whether in training mode (`True`) or inference mode (`False`)
    fp8 : bool, default = False
        Whether `DotProductAttention` is in an `autocast` region.
    fp8_meta : Optional[Dict[str Any]], default = None
        The FP8 metadata tensor of `DotProductAttention`.
    inference_params : Optional[InferenceParams], default = None
        Inference-related parameters. See InferenceParams for details.
    softmax_type : str, default = "vanilla"
        The type of softmax operation. See DotProductAttention for details.
    return_max_logit : bool, default = False
        Whether to output max_logit.
    cuda_graph : bool, default = `False`
        Whether support for cuda graph capture is needed or not.
    num_splits : int, default = 1
        The number of kernels to split attention to.
    fp8_output : bool, default = False
        Whether output is requested in FP8.
    checkpoint_core_attention : bool, default = False
        Whether core attention is recomputed during backward.
    has_score_mod : bool, default = False
        Whether a score_mod callback was provided.
    has_score_mod_bprop : bool, default = False
        Whether a score_mod bprop callback was provided.
    """

    qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor
    qkv_dtype: torch.dtype = torch.bfloat16
    qkv_layout: str = "sbh3d"
    batch_size: int = 1
    num_heads: int = 16
    num_gqa_groups: int = 16
    max_seqlen_q: int = 128
    max_seqlen_kv: int = 128
    head_dim_qk: int = 64
    head_dim_v: int = 64
    attn_mask_type: str = "no_mask"
    window_size: Union[Tuple[int, int], None] = None
    bottom_right_diagonal: bool = True
    alibi_slopes_shape: Union[torch.Size, List, None] = None
    core_attention_bias_type: str = "no_bias"
    core_attention_bias_shape: str = "1hss"
    core_attention_bias_requires_grad: bool = True
    pad_between_seqs: bool = False
    attention_dropout: float = 0.0
    context_parallel: bool = False
    cp_comm_type: str = "p2p"
    cp_size: int = 1
    deterministic: bool = False
    is_training: bool = True
    fp8: bool = False
    fp8_meta: Union[Dict[str, Any], None] = None
    inference_params: Optional[InferenceParams] = None
    softmax_type: str = "vanilla"
    return_max_logit: bool = False
    cuda_graph: bool = False
    num_splits: int = 1
    fp8_output: bool = False
    checkpoint_core_attention: bool = False
    has_score_mod: bool = False
    has_score_mod_bprop: bool = False

    def __eq__(self, other):
        """
        Overwrite dataclass.__eq__ so that only fp8_meta["recipe"] is compared,
        since all other entries of fp8_meta are unused in get_attention_backend.
        """
        if not isinstance(other, self.__class__):
            return NotImplemented
        for field in fields(self):
            fname = field.name
            sf = getattr(self, fname)
            of = getattr(other, fname)
            if fname != "fp8_meta":
                if sf != of:
                    return False
            elif (sf or {}).get("recipe", None) != (of or {}).get("recipe", None):
                return False
        return True


class _NoOpLogger:
    """
    Stand-in for the "DotProductAttention" logger used when get_attention_backend
    is traced by torch.compile. logging.Logger methods are not traceable by dynamo
    (they cause graph breaks), while this class's no-op methods are inlined away.
    """

    def debug(self, *args, **kwargs):
        """No-op."""

    def info(self, *args, **kwargs):
        """No-op."""

    def warning(self, *args, **kwargs):
        """No-op."""

    def error(self, *args, **kwargs):
        """No-op."""


_no_op_logger = _NoOpLogger()


@torch.compiler.assume_constant_result
def _get_fused_attn_backend(
    is_training,
    q_type,
    kv_type,
    qkv_layout,
    bias_type,
    attn_mask_type,
    softmax_type,
    *args,
):
    """Constant-foldable tex.get_fused_attn_backend: the result depends only on
    the attention config, and the python-side enum keeps it traceable by
    torch.compile (see the FusedAttnBackend docstring). Layout/bias/mask/softmax
    are taken as their string keys and resolved to the pybind enums here, so
    that every argument is a python literal or a python enum."""
    return FusedAttnBackend.cast(
        tex.get_fused_attn_backend(
            is_training,
            q_type,
            kv_type,
            QKVLayout[qkv_layout],
            AttnBiasType[bias_type],
            AttnMaskType[attn_mask_type],
            SoftmaxType[softmax_type],
            *args,
        )
    )


def get_attention_backend(
    attention_params: AttentionParams = None,
):
    """
    Select the appropriate attention backend/sub-backend based on user input and runtime environment.

    Parameters
    ----------
    See `AttentionParams`.

    Returns
    ----------
    use_flash_attention : bool
        Whether the `FlashAttention` backend has been selected.
    use_fused_attention : bool
        Whether the `FusedAttention` backend has been selected.
    fused_attention_backend : FusedAttnBackend
        If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`.
    use_unfused_attention : bool
        Whether the `UnfusedDotProductAttention` backend has been selected.
    available_backends : List[bool]
        All available backends that could support the provided input. A list of Booleans
        in the form of [use_flash_attention, use_fused_attention, use_unfused_attention].
    """
    # NOTE: As part of refactoring attention.py, populating the _attention_backends cache in attention
    # is no longer performed at the end of get_attention_backend(), but the responsibility of doing so
    # is shifted over to the caller of this function
    qkv_type = attention_params.qkv_type
    qkv_dtype = attention_params.qkv_dtype
    qkv_layout = attention_params.qkv_layout
    batch_size = attention_params.batch_size
    num_heads = attention_params.num_heads
    num_gqa_groups = attention_params.num_gqa_groups
    max_seqlen_q = attention_params.max_seqlen_q
    max_seqlen_kv = attention_params.max_seqlen_kv
    head_dim_qk = attention_params.head_dim_qk
    head_dim_v = attention_params.head_dim_v
    attn_mask_type = attention_params.attn_mask_type
    window_size = attention_params.window_size
    bottom_right_diagonal = attention_params.bottom_right_diagonal
    alibi_slopes_shape = attention_params.alibi_slopes_shape
    core_attention_bias_type = attention_params.core_attention_bias_type
    core_attention_bias_shape = attention_params.core_attention_bias_shape
    core_attention_bias_requires_grad = attention_params.core_attention_bias_requires_grad
    pad_between_seqs = attention_params.pad_between_seqs
    attention_dropout = attention_params.attention_dropout
    context_parallel = attention_params.context_parallel
    cp_comm_type = attention_params.cp_comm_type
    cp_size = attention_params.cp_size  # pylint: disable=unused-variable
    deterministic = attention_params.deterministic
    is_training = attention_params.is_training
    fp8 = attention_params.fp8
    fp8_meta = attention_params.fp8_meta
    inference_params = attention_params.inference_params
    softmax_type = attention_params.softmax_type
    return_max_logit = attention_params.return_max_logit
    cuda_graph = attention_params.cuda_graph
    num_splits = attention_params.num_splits
    fp8_output = attention_params.fp8_output
    checkpoint_core_attention = attention_params.checkpoint_core_attention
    has_score_mod = attention_params.has_score_mod
    has_score_mod_bprop = attention_params.has_score_mod_bprop

    # NOTE: environment variables in this function are read with
    # os.environ.get, NOT os.getenv, on purpose: dynamo installs guards on
    # os.environ reads (so changing an NVTE_* variable triggers recompilation
    # under torch.compile), while os.getenv reads are unguarded and would bake
    # stale values into compiled graphs. New code must follow suit.

    # Run config
    if torch.compiler.is_compiling():
        # logging.Logger methods graph-break under torch.compile; backend
        # selection logs are only emitted in eager mode.
        logger = _no_op_logger
    else:
        logger = logging.getLogger("DotProductAttention")
        logger.setLevel(AttentionLogging._log_level)
        if not logger.hasHandlers():
            logger.addHandler(AttentionLogging._stream_handler)
    device_compute_capability = get_device_compute_capability()
    cudnn_version = get_cudnn_version()
    run_config = {
        "transformer_engine_version": te.__version__,
        "compute_capability": (
            "sm" + str(10 * device_compute_capability[0] + device_compute_capability[1])
        ),
        "cuda_version": torch.version.cuda,
        "flash_attn_version": (
            str(FlashAttentionUtils.version)
            if FlashAttentionUtils.is_installed
            else "not installed"
        ),
        "flash_attn_3_version": (
            str(FlashAttentionUtils.fa3_version)
            if FlashAttentionUtils.v3_is_installed
            else "not installed"
        ),
        "flash_attn_4_version": (
            str(FlashAttentionUtils.fa4_version)
            if FlashAttentionUtils.v4_is_installed
            else "not installed"
        ),
        "cudnn_version": ".".join([str(i) for i in cudnn_version]),
    }
    attention_params_dict = {
        field.name: getattr(attention_params, field.name) for field in fields(attention_params)
    }
    run_config.update(attention_params_dict)
    # Add FP8 environment variables to config
    if fp8:
        # all FP8 recipes: 1: (FP8 fwd, FP8 bwd), 0: (FP8 fwd, F16 bwd)
        run_config["NVTE_FP8_DPA_BWD"] = int(os.environ.get("NVTE_FP8_DPA_BWD", "1"))
        # Float8CurrentScaling: 1: use F16 O in bwd, 0: use FP8 O in bwd
        run_config["NVTE_DPA_FP8CS_O_in_F16"] = int(os.environ.get("NVTE_DPA_FP8CS_O_in_F16", "1"))
        # switch recipe to "F16", "DelayedScaling", or "Float8CurrentScaling"
        _dpa_fp8_recipe = os.environ.get("NVTE_DPA_FP8_RECIPE", "")
        run_config["NVTE_DPA_FP8_RECIPE"] = _dpa_fp8_recipe
        if _dpa_fp8_recipe != "":
            # config new recipe if switched
            run_config["NVTE_DPA_FP8_FORMAT"] = os.environ.get("NVTE_DPA_FP8_FORMAT", "HYBRID")
            run_config["NVTE_DPA_FP8DS_AMAX_ALGO"] = os.environ.get(
                "NVTE_DPA_FP8DS_AMAX_ALGO", "most_recent"
            )
            run_config["NVTE_DPA_FP8DS_AMAX_HISTLEN"] = int(
                os.environ.get("NVTE_DPA_FP8DS_AMAX_HISTLEN", "1")
            )
            run_config["NVTE_DPA_FP8DS_REDUCE_AMAX"] = int(
                os.environ.get("NVTE_DPA_FP8DS_REDUCE_AMAX", "1")
            )
        # UnfusedDotProductAttention: 1: allow FP8 emulation, 0: do not allow
        run_config["NVTE_UnfusedDPA_Emulate_FP8"] = int(
            os.environ.get("NVTE_UnfusedDPA_Emulate_FP8", "0")
        )
    logger.debug("Running with config=%s", run_config)

    # The following sections check if `FlashAttention` supports the provided attention params,
    # regardless of whether FA2 or FA3 is installed. If FA2 or FA3 is not installed but is
    # necessary for performance/functionality, a warning will be issued to prompt users to
    # install an appropriate FA version.
    qkv_format, q_format, kv_format = get_qkv_format(qkv_layout, inference_params)

    # Filter: Environment variables
    use_flash_attention = int(os.environ.get("NVTE_FLASH_ATTN", "1"))
    use_flash_attention_2 = use_flash_attention and int(os.environ.get("NVTE_FLASH_ATTN_V2", "1"))
    use_flash_attention_3 = use_flash_attention and int(os.environ.get("NVTE_FLASH_ATTN_V3", "1"))
    use_flash_attention_4 = use_flash_attention and int(os.environ.get("NVTE_FLASH_ATTN_V4", "1"))
    flash_attention_backend = None
    use_fused_attention = int(os.environ.get("NVTE_FUSED_ATTN", "1"))
    use_unfused_attention = int(os.environ.get("NVTE_UNFUSED_ATTN", "1"))
    if not use_flash_attention_2 and FlashAttentionUtils.is_installed:
        logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V2=0")
    if not use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
        logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V3=0")
    if not use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
        logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0 or NVTE_FLASH_ATTN_V4=0")
    if not use_fused_attention:
        logger.debug("Disabling FusedAttention due to NVTE_FUSED_ATTN=0")
    if not use_unfused_attention:
        logger.debug("Disabling UnfusedDotProductAttention due to NVTE_UNFUSED_ATTN=0")

    def _any_flash_attention_enabled() -> bool:
        return bool(
            use_flash_attention
            or use_flash_attention_2
            or use_flash_attention_3
            or use_flash_attention_4
        )

    def _disable_all_flash_attention() -> None:
        nonlocal use_flash_attention
        use_flash_attention = False

    # Filter: Compute capability
    if device_compute_capability < (8, 0):
        if use_flash_attention_2 and FlashAttentionUtils.is_installed:
            logger.debug("Disabling FlashAttention 2 for compute capability < sm80")
        use_flash_attention_2 = False
        if use_fused_attention:
            logger.debug("Disabling FusedAttention for compute capability < sm80")
            use_fused_attention = False
    if device_compute_capability != (9, 0):
        if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
            logger.debug("Disabling FlashAttention 3 for compute capability != sm90")
        use_flash_attention_3 = False
    # FA4 supports SM80, SM90, SM100, SM120
    if device_compute_capability < (8, 0):
        if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
            logger.debug("Disabling FlashAttention 4 for compute capability < sm80")
        use_flash_attention_4 = False
    # On SM90, prefer FA3 over FA4 when FA3 is available.
    # FA3 is more mature on Hopper; FA4's SM90 backward has limitations
    # (MLA, non-standard head dims, SplitKV).
    if (
        device_compute_capability == (9, 0)
        and use_flash_attention_3
        and FlashAttentionUtils.v3_is_installed
        and use_flash_attention_4
        and FlashAttentionUtils.v4_is_installed
    ):
        logger.debug("Disabling FlashAttention 4 to prefer FlashAttention 3 on SM90")
        use_flash_attention_4 = False

    # Filter: Data type
    if qkv_dtype not in [torch.bfloat16, torch.float16]:
        if use_flash_attention_2 and FlashAttentionUtils.is_installed:
            logger.debug(
                "Disabling FlashAttention 2 for unsupported qkv_dtype = %s. "
                "Supported: qkv_dtype = {torch.bfloat16, torch.float16}. ",
                qkv_dtype,
            )
        use_flash_attention_2 = False
        if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
            logger.debug(
                "Disabling FlashAttention 4 for unsupported qkv_dtype = %s. "
                "Supported: qkv_dtype = {torch.bfloat16, torch.float16}. ",
                qkv_dtype,
            )
        use_flash_attention_4 = False
    if qkv_dtype not in [torch.bfloat16, torch.float16, torch.float8_e4m3fn] or qkv_type not in [
        torch.Tensor,
        Float8Tensor,
        Float8TensorStorage,
    ]:
        if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
            logger.debug(
                "Disabling FlashAttention 3 for unsupported qkv_dtype = %s, qkv_type = %s."
                " Supported: qkv_dtype = {torch.bfloat16, torch.float16, torch.float8_e4m3fn},"
                " qkv_type = {torch.Tensor, Float8Tensor, Float8TensorStorage}. ",
                qkv_dtype,
                qkv_type,
            )
        use_flash_attention_3 = False
    if qkv_dtype not in [torch.bfloat16, torch.float16, torch.float8_e4m3fn] or qkv_type not in (
        torch.Tensor,
        Float8Tensor,
        Float8TensorStorage,
        MXFP8Tensor,
        MXFP8TensorStorage,
    ):
        if use_fused_attention:
            logger.debug(
                "Disabling FusedAttention for unsupported qkv_dtype = %s, qkv_type = %s. Supported:"
                " qkv_dtype = {torch.bfloat16, torch.float16, torch.float8_e4m3fn}, qkv_type ="
                " {torch.Tensor, Float8Tensor, Float8TensorStorage, MXFP8Tensor,"
                " MXFP8TensorStorage}. ",
                qkv_dtype,
                qkv_type,
            )
            use_fused_attention = False

    # Filter: Execution type
    if fp8 and fp8_meta["recipe"].fp8_dpa:
        fp8_recipe = fp8_meta["recipe"]
        if fp8_meta.get("local_recipes", None) is not None:
            fp8_recipe = fp8_meta["local_recipes"][0]
        if use_flash_attention_2 and FlashAttentionUtils.is_installed:
            logger.debug("Disabling FlashAttention 2 for FP8 attention")
            use_flash_attention_2 = False
        if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
            logger.debug("Disabling FlashAttention 4 for FP8 attention")
            use_flash_attention_4 = False
        if use_flash_attention_3 and is_training:
            if FlashAttentionUtils.v3_is_installed:
                logger.debug("Disabling FlashAttention 3 for FP8 training")
            use_flash_attention_3 = False
        if use_flash_attention_3 and not (
            fp8_recipe.delayed() or fp8_recipe.float8_current_scaling()
        ):
            if FlashAttentionUtils.v3_is_installed:
                logger.debug("Disabling FlashAttention 3 for %s", fp8_recipe.__class__.__name__)
            use_flash_attention_3 = False
        if use_unfused_attention:
            allow_emulation = (
                os.environ.get("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1"
                or is_in_onnx_export_mode()
            )
            if not allow_emulation:
                logger.debug("Disabling UnfusedDotProductAttention for FP8 attention")
                use_unfused_attention = False
        if use_fused_attention and fp8_recipe.delayed():
            if (
                device_compute_capability >= (10, 0)
                and deterministic
                and cudnn_version < (9, 18, 0)
            ):
                logger.debug(
                    "Disabling FusedAttention for FP8 delayed scaling on arch >= sm100 with"
                    " determinism for cuDNN < 9.18.0"
                )
                use_fused_attention = False
        if use_fused_attention and fp8_recipe.float8_current_scaling():
            if device_compute_capability < (10, 0):
                logger.debug("Disabling FusedAttention for FP8 current scaling on arch < sm100")
                use_fused_attention = False
            else:
                if cudnn_version < (9, 14, 0):
                    logger.debug(
                        "Disabling FusedAttention for FP8 current scaling with cuDNN < 9.14.0"
                    )
                    use_fused_attention = False
                else:
                    if deterministic and cudnn_version < (9, 18, 0):
                        logger.debug(
                            "Disabling FusedAttention for FP8 current scaling with determinism"
                            " for cuDNN < 9.18.0"
                        )
                        use_fused_attention = False
        if use_fused_attention and fp8_recipe.mxfp8():
            if device_compute_capability < (10, 0):
                logger.debug("Disabling FusedAttention for MXFP8 on arch < sm100")
                use_fused_attention = False
            elif fp8_recipe.fp8_mha:
                logger.debug("Disabling FusedAttention for MXFP8 with fp8_mha=True")
                use_fused_attention = False
            else:
                if cudnn_version < (9, 21, 0):
                    logger.debug("Disabling FusedAttention for MXFP8 with cuDNN < 9.21.0")
                    use_fused_attention = False
                elif cudnn_version in ((9, 23, 0), (9, 23, 1)):
                    # 9.23.0/9.23.1: known bugs with MXFP8 SDPA
                    logger.debug("Disabling FusedAttention for MXFP8 with cuDNN 9.23.0/9.23.1")
                    use_fused_attention = False
                elif qkv_format == "thd":
                    logger.debug("Disabling FusedAttention for MXFP8 with qkv_format = thd")
                    use_fused_attention = False
        if use_fused_attention and (fp8_recipe.float8_block_scaling() or fp8_recipe.nvfp4()):
            logger.debug("Disabling FusedAttention for %s", fp8_recipe.__class__.__name__)
            use_fused_attention = False

        if device_compute_capability == (12, 0):
            if use_flash_attention:
                logger.debug(
                    "Disabling FlashAttention as FP8 is not supported"
                    " for compute capability = sm120"
                )
            if use_fused_attention:
                logger.debug(
                    "Disabling FusedAttention as FP8 is not supported"
                    " for compute capability = sm120"
                )
            use_flash_attention = False
            use_fused_attention = False

    # Filter: num_splits
    if num_splits != 1:
        if use_flash_attention_2 and FlashAttentionUtils.is_installed:
            logger.debug("Disabling FlashAttention 2 for num_splits")
            use_flash_attention_2 = False
        # FA4 SplitKV is only supported on SM100+
        if use_flash_attention_4 and device_compute_capability < (10, 0):
            if FlashAttentionUtils.v4_is_installed:
                logger.debug("Disabling FlashAttention 4 for num_splits on SM < 100")
            use_flash_attention_4 = False
        if use_fused_attention:
            logger.debug("Disabling FusedAttention for num_splits")
            use_fused_attention = False
        if use_unfused_attention:
            logger.debug("Disabling UnfusedDotProductAttention for num_splits")
            use_unfused_attention = False

    # Filter: Return max_logit
    if return_max_logit:
        if use_flash_attention:
            use_flash_attention = False
            logger.debug("Disabling FlashAttention for max_logit")
        if fp8 and fp8_meta["recipe"].fp8_dpa:
            use_flash_attention = False
            use_fused_attention = False
            use_unfused_attention = False
            logger.debug("Disabling all backends for max_logit with FP8 attention")

    # Filter: score_mod
    if has_score_mod_bprop and not has_score_mod:
        logger.debug("Disabling all backends because score_mod_bprop requires score_mod")
        _disable_all_flash_attention()
        use_fused_attention = False
        use_unfused_attention = False
    if has_score_mod:
        if _any_flash_attention_enabled():
            logger.debug("Disabling FlashAttention for score_mod")
            _disable_all_flash_attention()
        if use_unfused_attention:
            logger.debug("Disabling UnfusedDotProductAttention for score_mod")
            use_unfused_attention = False

        score_mod_unsupported_reasons = []
        if qkv_dtype not in [torch.float16, torch.bfloat16]:
            score_mod_unsupported_reasons.append(
                f"unsupported qkv_dtype = {qkv_dtype}; supported: torch.float16, torch.bfloat16"
            )
        if qkv_type is not torch.Tensor:
            score_mod_unsupported_reasons.append(
                f"unsupported qkv_type = {qkv_type}; supported: torch.Tensor"
            )
        if fp8:
            score_mod_unsupported_reasons.append("FP8 DotProductAttention is enabled")
        if fp8_output:
            score_mod_unsupported_reasons.append("fp8_output is enabled")
        if inference_params is not None:
            score_mod_unsupported_reasons.append("KV caching is enabled")
        if context_parallel:
            score_mod_unsupported_reasons.append("context parallelism is enabled")
        if "thd" in (qkv_format, q_format, kv_format):
            score_mod_unsupported_reasons.append(
                f"unsupported QKV format: q_format = {q_format}, kv_format = {kv_format}"
            )
        if pad_between_seqs:
            score_mod_unsupported_reasons.append("pad_between_seqs is enabled")
        if attn_mask_type != "no_mask":
            score_mod_unsupported_reasons.append(f"attn_mask_type = {attn_mask_type}")
        if window_size is not None and window_size != (-1, -1):
            score_mod_unsupported_reasons.append(f"window_size = {window_size}")
        if core_attention_bias_type != "no_bias":
            score_mod_unsupported_reasons.append(
                f"core_attention_bias_type = {core_attention_bias_type}"
            )
        if alibi_slopes_shape is not None:
            score_mod_unsupported_reasons.append("ALiBi slopes were provided")
        if softmax_type != "vanilla":
            score_mod_unsupported_reasons.append(f"softmax_type = {softmax_type}")
        if attention_dropout != 0.0:
            score_mod_unsupported_reasons.append(f"attention_dropout = {attention_dropout}")
        if return_max_logit:
            score_mod_unsupported_reasons.append("return_max_logit is enabled")
        if checkpoint_core_attention:
            score_mod_unsupported_reasons.append("checkpoint_core_attention is enabled")
        if cuda_graph:
            score_mod_unsupported_reasons.append("CUDA graph capture is enabled")
        if num_splits != 1:
            score_mod_unsupported_reasons.append(f"num_splits = {num_splits}")
        if score_mod_unsupported_reasons and use_fused_attention:
            logger.debug(
                "Disabling FusedAttention for score_mod because %s",
                "; ".join(score_mod_unsupported_reasons),
            )
            use_fused_attention = False

    # Filter: KV cache
    # backend  | precision      |    KV cache     | architecture | qkv_format    | page_size
    # ---------------------------------------------------------------------------------------
    # Fused    | FP16/BF16      | non-paged/paged | sm80+        | bshd,sbhd,thd | >= 1
    # Flash v2 | FP16/BF16      | non-paged/paged | sm80+        | bshd,sbhd,thd | % 256 == 0
    # Flash v3 | FP16/BF16      | non-paged/paged | sm90         | bshd,sbhd,thd | >= 1
    #          | FP8            | non-paged/paged | sm90         | thd           | >= 1
    # Flash v4 | FP16/BF16      | TODO            | sm80+        | bshd,sbhd,thd | TODO
    # Unfused  | FP32/FP16/BF16 | non-paged/paged | all          | bshd,sbhd,thd | >= 1
    if inference_params is not None:
        # Temporarily disabling fused attention for kv caching for sm89/sm120 irrespective of
        # cuDNN version until the cuDNN bug is resolved.
        if device_compute_capability in ((8, 9), (12, 0)):
            logger.debug("Disabling FusedAttention for KV caching for sm89/sm120")
            use_fused_attention = False
        # Temporarily disable FlashAttention for KV caching on sm120
        if device_compute_capability == (12, 0):
            logger.debug("Disabling FlashAttention for KV caching for sm120")
            use_flash_attention = False
        if context_parallel:
            logger.debug("Disabling all backends for KV caching with context parallelism")
            use_flash_attention = False
            use_fused_attention = False
            use_unfused_attention = False
        if fp8 and fp8_meta["recipe"].fp8_dpa:
            if fp8_meta["recipe"].fp8_mha:
                logger.debug("Disabling all backends for KV caching with FP8 MHA")
                use_flash_attention = False
                use_fused_attention = False
                use_unfused_attention = False
            if use_flash_attention_3 and q_format != "thd":
                if FlashAttentionUtils.v3_is_installed:
                    logger.debug("Disabling FlashAttention 3 for FP8 KV caching and non-THD")
                use_flash_attention_3 = False
            if use_fused_attention:
                logger.debug("Disabling FusedAttention for FP8 KV caching")
                use_fused_attention = False
        else:
            if q_format == "thd" and pad_between_seqs:
                logger.debug("Disabling all backends for pad_between_seqs = True and KV caching")
                use_flash_attention = False
                use_fused_attention = False
                use_unfused_attention = False
        if inference_params.is_paged:
            if use_flash_attention_2 and inference_params.page_size % 256 != 0:
                if FlashAttentionUtils.is_installed:
                    logger.debug("Disabling FlashAttention 2 for page size not divisible by 256")
                use_flash_attention_2 = False
            if use_flash_attention_2:
                if not FlashAttentionUtils.is_installed:
                    FlashAttentionUtils.version_required = PkgVersion("2.5")
                elif not FlashAttentionUtils.v2_5_plus:
                    logger.debug(
                        "Disabling FlashAttention 2 as paged attention requires flash-attn 2.5+"
                    )
                    use_flash_attention_2 = False
        else:
            # Non-paged KV cache still passes a block_table to FA2 for thd_2bshd support,
            # and FA2 enforces page_size % 256 == 0 on the effective page size (max_seqlen_kv).
            if use_flash_attention_2 and max_seqlen_kv % 256 != 0:
                if FlashAttentionUtils.is_installed:
                    logger.debug(
                        "Disabling FlashAttention 2 for non-paged KV cache"
                        " with max_seqlen_kv not divisible by 256"
                    )
                use_flash_attention_2 = False
        if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
            logger.debug("Disabling FlashAttention 4 as it does not support KV cache.")
            use_flash_attention_4 = False

    # Filter: Head dimension
    if head_dim_qk != head_dim_v:
        qkv_layout_group = qkv_layout.replace("b", "").replace("s", "").replace("t", "")
        if use_fused_attention and qkv_layout_group != "hd_hd_hd":
            logger.debug(
                "Disabling FusedAttention as MLA is not supported with qkv_layout = %s.",
                qkv_layout,
            )
            use_fused_attention = False
        if (
            device_compute_capability == (12, 0)
            and (head_dim_qk > 128 or head_dim_qk % 8 != 0)
            and is_training
        ):
            if use_fused_attention:
                logger.debug(
                    "Disabling FusedAttention as MLA for backward pass is not supported for compute"
                    " capability = sm120 for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:"
                    " head_dim_qk = %s",
                    head_dim_qk,
                )
            use_fused_attention = False

    fa2_padded_head_dim = max(head_dim_qk, head_dim_v)
    if (  # pylint: disable=too-many-boolean-expressions
        use_flash_attention_2
        and FlashAttentionUtils.is_installed
        and (
            fa2_padded_head_dim > 256
            or fa2_padded_head_dim % 8 != 0
            or (
                fa2_padded_head_dim > 192
                and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0))
            )
        )
    ):
        logger.debug(
            "Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. "
            "Supported after padding: padded head_dim %%8 = 0, padded head_dim <= 256 "
            "(>192 requires sm80/90/100+). "
            "Found: head_dim_qk = %s, head_dim_v = %s, padded head_dim = %s, on sm%s.",
            head_dim_qk,
            head_dim_v,
            fa2_padded_head_dim,
            ".".join([str(i) for i in device_compute_capability]),
        )
        use_flash_attention_2 = False

    if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:

        def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dtype):
            if head_dim_qk > 256 or num_heads % num_gqa_groups != 0:
                return False
            if head_dim_qk != head_dim_v:
                cond1 = 128 < head_dim_qk <= 192
                cond2 = 96 < head_dim_v <= 128
                cond3 = head_dim_qk <= 64 and head_dim_v <= 512
                if not ((cond1 and cond2) or cond3):
                    return False
                if head_dim_v > 256 and qkv_dtype not in (torch.bfloat16, torch.float16):
                    return False
            return True

        if not _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dtype):
            logger.debug(
                "Disabling FlashAttention 3 due to unsupported num_heads, num_gqa_groups, "
                "head_dim_qk, head_dim_v or qkv_dtype. "
                "Supported: head_dim_qk <= 256, and num_heads %% num_gqa_groups = 0, and "
                "if head_dim_qk is different from head_dim_v, then "
                "(head_dim_qk must in (128, 192] and head_dim_v in (96, 128]) or "
                "(head_dim_qk <= 64 and head_dim_v <= 512), and "
                "if head_dim_qk is different from head_dim_v and head_dim_v > 256, then "
                "qkv_dtype requires fp16 and bf16 data type. "
                "Found: num_heads = %s, num_gqa_groups = %s, "
                "head_dim_qk = %s, head_dim_v = %s and qkv_dtype = %s.",
                num_heads,
                num_gqa_groups,
                head_dim_qk,
                head_dim_v,
                qkv_dtype,
            )
            use_flash_attention_3 = False

    if (
        use_flash_attention_4
        and FlashAttentionUtils.v4_is_installed
        and FlashAttentionUtils.v4_validate_head_dims is not None
    ):
        # Defer to FA4's own _validate_head_dims to keep TE in sync with FA4 supported shapes
        # (e.g., (256, 256) on SM100, (192, 128) DeepSeek, (64, 512) MLA-absorbed).
        # The function asserts on unsupported combinations; SM80/SM120 have no validation branch
        # in FA4 so the call passes through silently for those archs.
        _fa4_alignment = 16 // torch.empty(0, dtype=qkv_dtype).element_size()
        try:
            # pylint: disable-next=not-callable
            FlashAttentionUtils.v4_validate_head_dims(
                head_dim_qk,
                head_dim_v,
                device_compute_capability[0],
                _fa4_alignment,
            )
        except AssertionError:
            logger.debug(
                "Disabling FlashAttention 4 due to unsupported head dimensions. "
                "Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.",
                head_dim_qk,
                head_dim_v,
                device_compute_capability[0] * 10 + device_compute_capability[1],
            )
            use_flash_attention_4 = False
        # flash-attn-4 4.0.0b11 validates (256, 256) on SM100, but its dedicated
        # hd256 kernel diverges from the reference for cross-attention/decode-like
        # shapes such as sq=1, skv=2048. Keep FA4 enabled for the self-attention
        # hd256 path covered by the dedicated test, and fall back for cross-attn.
        if (
            use_flash_attention_4
            and (10, 0) <= device_compute_capability < (12, 0)
            and head_dim_qk == head_dim_v == 256
            and max_seqlen_q != max_seqlen_kv
        ):
            logger.debug(
                "Disabling FlashAttention 4 for SM100 head_dim=256 cross-attention. "
                "Found: max_seqlen_q = %s, max_seqlen_kv = %s.",
                max_seqlen_q,
                max_seqlen_kv,
            )
            use_flash_attention_4 = False
        # Workaround: SM100 backward kernel bug when MLA + 2CTA (head_dim_qk >= 128) for
        # the standard (non-dedicated) kernel path. FlashAttentionBackwardSm100 computes
        # dK_reduce_ncol = gcd(32, tile_hdim // 2) based on Q/K head_dim but reuses it for
        # dV TMEM load atoms. When (tile_hdimv // 2) % dK_reduce_ncol != 0, dV reads are
        # misaligned (e.g. dqk=128, dv=96 gives 48 % 32 != 0). The dedicated (256, 256)
        # kernel uses its own tmem layout and is not affected.
        # See: flash_attn/cute/flash_bwd_sm100.py ~L262 and ~L3890. Still present in
        # flash-attn-4 4.0.0b11.
        if (
            use_flash_attention_4
            and is_training
            and head_dim_qk != head_dim_v
            and head_dim_qk >= 128
            and (10, 0) <= device_compute_capability < (12, 0)
        ):
            _tile_hdim = math.ceil(head_dim_qk / 16) * 16
            _tile_hdimv = math.ceil(head_dim_v / 16) * 16
            _dk_reduce_ncol = math.gcd(32, _tile_hdim // 2)
            if (_tile_hdimv // 2) % _dk_reduce_ncol != 0:
                logger.debug(
                    "Disabling FlashAttention 4 for training due to SM100 backward kernel "
                    "bug with MLA head dimensions (dK_reduce_ncol misalignment for dV). "
                    "Found: head_dim_qk = %s, head_dim_v = %s.",
                    head_dim_qk,
                    head_dim_v,
                )
                use_flash_attention_4 = False

    # Filter: QKV layout
    if qkv_format == "thd":
        if pad_between_seqs:
            if (  # pylint: disable=too-many-boolean-expressions
                use_flash_attention_2 and FlashAttentionUtils.is_installed
            ) or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed):
                logger.debug(
                    "Disabling FlashAttention 2 and 4 for qkv_format = thd when there is "
                    "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]"
                )
            use_flash_attention_2 = False
            use_flash_attention_4 = False
            # FA3 supports pad_between_seqs via seqused_q/seqused_k
            if use_unfused_attention:
                logger.debug("Disabling UnfusedDotProductAttention for pad_between_seqs = True")
            use_unfused_attention = False
        if device_compute_capability == (12, 0):
            if cudnn_version < (9, 18, 1):
                if use_fused_attention:
                    logger.debug(
                        "Disabling FusedAttention as qkv_format = thd is"
                        " not supported for compute capability = sm120 and cuDNN version < 9.18.1"
                    )
                use_fused_attention = False
            elif qkv_layout in {"t3hd", "th3d"}:
                if use_fused_attention:
                    logger.debug(
                        "Disabling FusedAttention as qkv_layout = %s is not supported for"
                        " compute capability = sm120",
                        qkv_layout,
                    )
                use_fused_attention = False

    # Filter: Dropout
    if attention_dropout != 0.0:
        if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
            logger.debug("Disabling FlashAttention 3 for dropout")
            use_flash_attention_3 = False
        if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
            logger.debug("Disabling FlashAttention 4 for dropout")
            use_flash_attention_4 = False

    # Filter: Softmax type
    # context_parallel | softmax_type | supported backends
    # ----------------------------------------------------------------------------------------------------
    # no               | vanilla      | All
    # no               | off-by-one   | FusedAttention, UnfusedDotProductAttention
    # no               | learnable    | FusedAttention, UnfusedDotProductAttention
    # yes              | vanilla      | FusedAttention, FlashAttention
    # yes              | off-by-one   | FusedAttention
    # yes              | learnable    | FusedAttention
    if softmax_type != "vanilla":
        logger.debug("Disabling FlashAttention for softmax_type = %s", softmax_type)
        use_flash_attention = False
        if fp8 and fp8_meta["recipe"].fp8_dpa:
            if use_fused_attention and (
                device_compute_capability < (10, 0) or cudnn_version < (9, 21, 0)
            ):
                logger.debug(
                    "Disabling FusedAttention for softmax_type = %s in FP8 on sm < 100 with cuDNN"
                    " version < 9.21",
                    softmax_type,
                )
                use_fused_attention = False
            if use_unfused_attention:
                logger.debug(
                    "Disabling UnfusedDotProductAttention for softmax_type = %s in FP8",
                    softmax_type,
                )
                use_unfused_attention = False
        if qkv_format == "thd" and cudnn_version < (9, 18, 0):
            logger.debug(
                "Disabling FusedAttention for softmax_type = %s, qkv_format = thd and cuDNN"
                " version < 9.18",
                softmax_type,
            )
            use_fused_attention = False
        if context_parallel and cp_comm_type != "a2a":
            logger.debug(
                "Disabling FusedAttention for context parallelism with softmax_type = %s and"
                " cp_comm_type = %s",
                softmax_type,
                cp_comm_type,
            )
            use_fused_attention = False

    # Filter: Context parallelism
    # qkv_format | attn_mask_type              | attn_bias_type           | supported backends
    # ----------------------------------------------------------------------------------------------------
    # bshd, sbhd | self-attention:             | no_bias, post_scale_bias | FlashAttention, FusedAttention
    #            |     no_mask, causal         |                          |
    #            | cross-attention:            |                          |
    #            |     no_mask                 |                          |
    # thd        | self-attention:             | no_bias                  | FlashAttention, FusedAttention
    #            |     padding, padding_causal |                          | if no padding between sequences,
    #            | cross-attention:            |                          | FusedAttention
    #            |     padding                 |                          | if there is padding between sequences
    # Note: context parallelism requires seq_len % (cp_size * 2) == 0 for each sequence in q, k, v.
    if context_parallel and use_unfused_attention:
        logger.debug(
            "Disabling UnfusedDotProductAttention as it does not support context parallelism"
        )
        use_unfused_attention = False
    if context_parallel and use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
        logger.debug("Disabling FlashAttention 4 as it does not support context parallelism yet")
        use_flash_attention_4 = False
    if context_parallel and (
        use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4
    ):
        if (
            FlashAttentionUtils.is_installed
            or FlashAttentionUtils.v3_is_installed
            or FlashAttentionUtils.v4_is_installed
        ):
            if fp8 and fp8_meta["recipe"].fp8_dpa:
                logger.debug(
                    "Disabling FlashAttention as it does not support context parallelism with FP8"
                )
                use_flash_attention = False
            if "bottom_right" in attn_mask_type:
                logger.debug(
                    "Disabling FlashAttention as it does not support context parallelism with"
                    " causal_bottom_right masking"
                )
                use_flash_attention = False
            elif "causal" in attn_mask_type and max_seqlen_q != max_seqlen_kv:
                logger.debug(
                    "Disabling FlashAttention as it does not support context parallelism with"
                    " causal masking for cross-attention"
                )
                use_flash_attention = False
            elif core_attention_bias_type not in ["no_bias", "post_scale_bias"]:
                logger.debug(
                    "Disabling FlashAttention as it does not support context parallelism with bias"
                    " type of %s",
                    core_attention_bias_type,
                )
                use_flash_attention = False
            elif qkv_format == "thd" and core_attention_bias_type != "no_bias":
                logger.debug(
                    "Disabling FlashAttention as it does not support context parallelism with"
                    " attention bias for THD format"
                )
                use_flash_attention = False

    if context_parallel and use_fused_attention:
        if "bottom_right" in attn_mask_type:
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with"
                " causal_bottom_right masking"
            )
            use_fused_attention = False
        elif "causal" in attn_mask_type and max_seqlen_q != max_seqlen_kv:
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with causal"
                " masking for cross-attention"
            )
            use_fused_attention = False
        elif core_attention_bias_type not in ["no_bias", "post_scale_bias"]:
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with bias type"
                " of %s",
                core_attention_bias_type,
            )
            use_fused_attention = False
        elif qkv_format == "thd" and core_attention_bias_type != "no_bias":
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with attention"
                " bias for THD format"
            )
            use_fused_attention = False
        elif fp8 and fp8_meta["recipe"].fp8_dpa and qkv_format == "thd":
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with FP8"
                " attention and THD format"
            )
            use_fused_attention = False
        elif fp8 and fp8_meta["recipe"].fp8_dpa and core_attention_bias_type != "no_bias":
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with FP8"
                " attention and bias"
            )
            use_fused_attention = False
        elif core_attention_bias_type != "no_bias" and cp_comm_type != "p2p":
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with bias"
                " and cp_comm_type = %s",
                cp_comm_type,
            )
            use_fused_attention = False
        elif (
            window_size is not None
            and (window_size[0] != -1 or window_size[1] not in [-1, 0])
            and cp_comm_type in ["p2p", "a2a+p2p"]
        ):
            logger.debug(
                "Disabling FusedAttention as it does not support context parallelism with sliding"
                " window attention and cp_comm_type = %s",
                cp_comm_type,
            )
            use_fused_attention = False
        elif cp_comm_type in ["a2a", "a2a+p2p"] and (num_heads % 2 != 0 or num_gqa_groups % 2 != 0):
            logger.debug(
                "Disabling FusedAttention as cp_comm_type = %s requires num_heads and"
                " num_gqa_groups divisible by 2 (got num_heads = %s, num_gqa_groups = %s)",
                cp_comm_type,
                num_heads,
                num_gqa_groups,
            )
            use_fused_attention = False

    # Filter: Attention mask
    # attn_mask_type              | attention_mask                       | supported backends
    # ----------------------------------------------------------------------------------------
    # no_mask                     | None                                 | All
    # padding                     |                                      | All
    #     self-attention          | One tensor of shape [b, 1, 1, sq]    |
    #     cross-attention         | Tuple of two tensors of shapes       |
    #                             | [b, 1, 1, sq] and [b, 1, 1, skv]     |
    # causal                      | None                                 |
    #     self-attention          |                                      | All
    #     cross-attention         |                                      | FusedAttention, UnfusedDotProductAttention
    # padding_causal              | Same as "padding"                    |
    #     self-attention          |                                      | All
    #     cross-attention         |                                      | FusedAttention, UnfusedDotProductAttention
    # causal_bottom_right         | None                                 | All
    # padding_causal_bottom_right | Same as "padding"                    | All
    # arbitrary                   | One tensor of shape broadcastable to | UnfusedDotProductAttention
    #                             | [b, h, sq, skv]                      |
    if attn_mask_type == "arbitrary":
        if (  # pylint: disable=too-many-boolean-expressions
            (use_flash_attention_2 and FlashAttentionUtils.is_installed)
            or (use_flash_attention_3 and FlashAttentionUtils.v3_is_installed)
            or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed)
        ):
            logger.debug("Disabling FlashAttention for arbitrary mask")
        use_flash_attention = False
        if use_fused_attention:
            logger.debug("Disabling FusedAttention for arbitrary mask")
        use_fused_attention = False
    if (
        (use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4)
        and attn_mask_type in ["causal", "padding_causal"]
        and max_seqlen_q != max_seqlen_kv
    ):
        logger.warning(
            "Disabling FlashAttention as it only supports bottom-right-diagonal "
            "causal mask since flash-attn 2.1 (our minimum supported version). See "
            "https://github.com/Dao-AILab/flash-attention#21-change-behavior-of-causal-flag"
        )
        use_flash_attention = False

    # Filter: Sliding window attention
    #    backend                 |      window_size       | diagonal alignment
    # ---------------------------------------------------------------------------------
    # FlashAttention             | (-1, -1) or (>=0, >=0) | bottom right
    # FusedAttention             | (-1,  0) or (>=0, >=0) | top left, bottom right
    # UnfusedDotProductAttention | (-1, -1) or (>=0, >=0) | top left, bottom right
    #                            |                        | converts window_size to an 'arbitrary' mask
    if window_size is None:
        window_size = check_set_window_size(attn_mask_type, window_size)
    if (
        use_flash_attention_4
        and (10, 0) <= device_compute_capability < (12, 0)
        and head_dim_qk == head_dim_v == 256
        and (window_size[0] != -1 or window_size[1] not in [-1, 0])
    ):
        logger.debug(
            "Disabling FlashAttention 4 as SM100 head_dim=256 does not support "
            "sliding-window/local attention yet. Found: window_size = %s.",
            window_size,
        )
        use_flash_attention_4 = False
    if use_fused_attention and (window_size[0] != -1 or window_size[1] not in [-1, 0]):
        if (
            fp8
            and (fp8_meta["recipe"].fp8_dpa or fp8_meta["recipe"].fp8_mha)
            and (device_compute_capability < (10, 0) or cudnn_version < (9, 21, 0))
        ):
            logger.debug(
                "Disabling FusedAttention as it does not support sliding window attention for FP8"
                " on sm < 100 with cuDNN version < 9.21"
            )
            use_fused_attention = False
        elif attention_dropout != 0.0:
            logger.debug(
                "Disabling FusedAttention as it only supports sliding window attention "
                "without dropout"
            )
            use_fused_attention = False
        elif max_seqlen_q > max_seqlen_kv:
            logger.debug(
                "Disabling FusedAttention as it does not support sliding window attention "
                "with s_q > s_kv for cross-attention"
            )
            use_fused_attention = False
    if use_flash_attention_2 and (window_size[0] != -1 or window_size[1] not in [-1, 0]):
        if not FlashAttentionUtils.is_installed:
            FlashAttentionUtils.version_required = PkgVersion("2.3")
        elif not FlashAttentionUtils.v2_3_plus:
            logger.debug(
                "Disabling FlashAttention as sliding window attention requires flash-attn 2.3+"
            )
            use_flash_attention_2 = False
        elif not bottom_right_diagonal and max_seqlen_q != max_seqlen_kv:
            logger.debug(
                "Disabling FlashAttention as it only supports sliding window with bottom right"
                " diagonal alignment for cross-attention"
            )
            use_flash_attention = False

    # Filter: Attention bias
    #    backend                 |      bias types              | ALiBi diagonal alignment
    # ---------------------------------------------------------------------------------
    # FlashAttention             | no_bias, alibi/alibi_slopes  | bottom right
    # FusedAttention             | no_bias, post_scale_bias     |
    #                            | alibi/alibi_slopes           | top left,
    #                            |                              | bottom_right (converts to a 'post_scale_bias' bias)
    # UnfusedDotProductAttention | no_bias, pre/post_scale_bias |
    #                            | alibi/alibi_slopes           | both; converts to a 'post_scale_bias' bias
    if core_attention_bias_type == "alibi":
        if use_flash_attention_3:
            if FlashAttentionUtils.v3_is_installed:
                logger.debug("Disabling FlashAttention 3 for ALiBi")
            use_flash_attention_3 = False
        if use_flash_attention_2:
            if not FlashAttentionUtils.is_installed:
                FlashAttentionUtils.version_required = PkgVersion("2.4")
            elif not FlashAttentionUtils.v2_4_plus:
                logger.debug("Disabling FlashAttention as ALiBi requires flash-attn 2.4+")
                use_flash_attention_2 = False
            elif not bottom_right_diagonal and max_seqlen_q != max_seqlen_kv:
                logger.debug(
                    "Disabling FlashAttention as it only supports ALiBi with bottom right diagonal"
                    " alignment for cross-attention"
                )
                use_flash_attention = False
        if use_flash_attention_4:
            if FlashAttentionUtils.v4_is_installed:
                logger.debug("Disabling FlashAttention 4 for ALiBi")
            use_flash_attention_4 = False

    if (
        core_attention_bias_type not in ["no_bias", "alibi"]
        or core_attention_bias_shape is not None
    ):
        if (  # pylint: disable=too-many-boolean-expressions
            (use_flash_attention_2 and FlashAttentionUtils.is_installed)
            or (use_flash_attention_3 and FlashAttentionUtils.v3_is_installed)
            or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed)
        ):
            logger.debug("Disabling FlashAttention for pre/post_scale_bias")
        use_flash_attention = False

    fu_core_attention_bias_type = core_attention_bias_type
    fu_core_attention_bias_shape = core_attention_bias_shape
    fu_core_attention_bias_requires_grad = core_attention_bias_requires_grad
    if (
        use_fused_attention
        and core_attention_bias_type == "alibi"
        and (alibi_slopes_shape is not None)
    ):
        fu_core_attention_bias_type = "post_scale_bias"
        fu_core_attention_bias_requires_grad = False

        if len(alibi_slopes_shape) == 1 and alibi_slopes_shape[0] == num_heads:
            fu_core_attention_bias_shape = "1hss"
        elif (
            len(alibi_slopes_shape) == 2
            and alibi_slopes_shape[0] == batch_size
            and alibi_slopes_shape[1] == num_heads
        ):
            fu_core_attention_bias_shape = "bhss"

    if (
        use_fused_attention
        and fu_core_attention_bias_type == "post_scale_bias"
        and fu_core_attention_bias_shape != "1hss"
    ):
        # dbias calculation is not supported for 111s as of cuDNN 9.18. So, use fused attention backend only if bias does not require grad.
        if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape == "111s":
            logger.warning(
                "Disabling FusedAttention as dbias calculation is not supported for 111s"
            )
            use_fused_attention = False
    # Filter: cuDNN support
    fused_attention_backend = None
    if use_fused_attention:
        # ``DType`` is implicitly convertible to ``transformer_engine::DType``
        # on the C++ side, so pass it straight to the pybind function.
        q_type = TE_DType[qkv_dtype]
        kv_type = q_type
        if fp8 and fp8_meta["recipe"].fp8_dpa:
            q_type = get_fp8_te_dtype(fp8_meta["recipe"], fprop_tensor=True)
            kv_type = q_type
        # NOTE: under torch.compile the numeric args below must not be symbolic
        # (assume_constant_result requires concrete values); ints/floats made
        # dynamic by automatic dynamic currently graph break here.
        fused_attention_backend = _get_fused_attn_backend(
            is_training,
            q_type,
            kv_type,
            qkv_layout,
            fu_core_attention_bias_type,
            attn_mask_type,
            softmax_type,
            attention_dropout,
            num_heads,
            num_gqa_groups,
            max_seqlen_q,
            max_seqlen_kv,
            head_dim_qk,
            head_dim_v,
            window_size[0],
            window_size[1],
            return_max_logit,
            cuda_graph,
            deterministic,
        )
        if fused_attention_backend == FusedAttnBackend["No_Backend"]:
            logger.debug("Disabling FusedAttention as no backend supports the provided input")
            use_fused_attention = False
            fused_attention_backend = None
        elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]:
            logger.debug(
                "Disabling FusedAttention for score_mod because sub-backend %s is not "
                "F16/BF16 arbitrary-seqlen",
                int(fused_attention_backend),
            )
            use_fused_attention = False
            fused_attention_backend = None
    # Filter: Determinism
    # backend                      | deterministic
    # ---------------------------------------------
    # FlashAttention               |
    #     flash-attn >=2.0, <2.4.1 | no
    #     flash-attn >=2.4.1       | yes
    # FusedAttention               |
    #     sub-backend 0            | yes
    #     sub-backend 1            | workspace optimization path and sm90: yes;
    #                              | otherwise: no
    #     sub-backend 2            | no
    # UnfusedDotProductAttention   | yes
    if use_flash_attention_2 and deterministic:
        if not FlashAttentionUtils.is_installed:
            FlashAttentionUtils.version_required = PkgVersion("2.4.1")
        elif not FlashAttentionUtils.v2_4_1_plus:
            logger.warning(
                "Disabling FlashAttention as version <2.4.1 does not support deterministic "
                "execution. To use FlashAttention with deterministic behavior, "
                "please install flash-attn >= 2.4.1."
            )
            use_flash_attention_2 = False
    if use_flash_attention_3 and deterministic and FlashAttentionUtils.v3_is_installed:
        if is_training and max(head_dim_qk, head_dim_v) >= 256:
            logger.debug(
                "Disabling FlashAttention 3 for deterministic backward with"
                " max(head_dim_qk, head_dim_v) >= 256. Found: head_dim_qk = %s, head_dim_v = %s.",
                head_dim_qk,
                head_dim_v,
            )
            use_flash_attention_3 = False
    if use_fused_attention and deterministic:
        if softmax_type != "vanilla":
            logger.debug(
                "Disabling FusedAttention for determinism reasons with softmax_type = %s. "
                "Sink attention (off-by-one and learnable softmax) requires "
                "NVTE_ALLOW_NONDETERMINISTIC_ALGO=1",
                softmax_type,
            )
            use_fused_attention = False
            fused_attention_backend = None
        if (
            fused_attention_backend == FusedAttnBackend["FP8"]
            and is_training
            and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0))
        ):
            logger.debug(
                "Disabling FusedAttention for determinism reasons with FP8 on arch < sm90 or cuDNN"
                " < 9.19.0"
            )
            use_fused_attention = False
            fused_attention_backend = None
        if (
            fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
            and is_training
            and (
                device_compute_capability < (9, 0)
                or core_attention_bias_requires_grad
                or cudnn_version < (8, 9, 5)
            )
        ):
            logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias")
            use_fused_attention = False
            fused_attention_backend = None

    # use_flash_attention may have been set above
    use_flash_attention_2 = use_flash_attention and use_flash_attention_2
    use_flash_attention_3 = use_flash_attention and use_flash_attention_3
    use_flash_attention_4 = use_flash_attention and use_flash_attention_4

    # `FusedAttention` and `FlashAttention` are faster backends than `UnfusedDotProductAttention`.
    # When `FusedAttention` does not support the provided attention params, and `FlashAttention`
    # does, we recommend users to install flash-attn if not installed already.
    if not use_fused_attention and _NVTE_FLASH_ATTN:
        if (
            use_flash_attention_4
            and not FlashAttentionUtils.v4_is_installed
            and not FlashAttentionUtils.v4_warning_printed
            and torch.cuda.current_device() == 0
        ):
            logger.warning(
                "flash-attn v4 may provide important feature support or performance improvement."
                " Please install flash-attn v4 by \n%s",
                FlashAttentionUtils.v4_installation_steps,
            )
            FlashAttentionUtils.v4_warning_printed = True
        elif (
            use_flash_attention_3
            and not FlashAttentionUtils.v3_is_installed
            and not FlashAttentionUtils.v3_warning_printed
            and torch.cuda.current_device() == 0
        ):
            logger.warning(
                "flash-attn v3 may provide important feature support or performance improvement."
                " Please install flash-attn v3 by \n%s",
                FlashAttentionUtils.v3_installation_steps,
            )
            FlashAttentionUtils.v3_warning_printed = True
        elif (
            use_flash_attention_2
            and not FlashAttentionUtils.is_installed
            and not FlashAttentionUtils.warning_printed
            and torch.cuda.current_device() == 0
        ):
            logger.warning(
                "flash-attn may provide important feature support or performance improvement."
                " Please install flash-attn %s by pip3 install flash-attn==<version>.",
                _get_supported_versions(
                    FlashAttentionUtils.version_required,
                    FlashAttentionUtils.max_version,
                ),
            )
            FlashAttentionUtils.warning_printed = True
    # All available backends
    if use_flash_attention_2 and not FlashAttentionUtils.is_installed:
        use_flash_attention_2 = False
    if use_flash_attention_3 and not FlashAttentionUtils.v3_is_installed:
        use_flash_attention_3 = False
    if use_flash_attention_4 and not FlashAttentionUtils.v4_is_installed:
        use_flash_attention_4 = False
    use_flash_attention = use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4
    available_backends = [use_flash_attention, use_fused_attention, use_unfused_attention]
    if use_flash_attention_2:
        flash_attention_backend = FlashAttentionUtils.version
    if use_flash_attention_3:
        flash_attention_backend = FlashAttentionUtils.fa3_version
    if use_flash_attention_4:
        flash_attention_backend = FlashAttentionUtils.fa4_version

    logger.debug(
        "Available backends = {FlashAttention=%s%s, FusedAttention=%s%s,"
        " UnfusedDotProductAttention=%s}",
        bool(available_backends[0]),
        (f" ({str(flash_attention_backend)})" if flash_attention_backend is not None else ""),
        bool(available_backends[1]),
        (
            f" (sub-backend {int(fused_attention_backend)})"
            if fused_attention_backend is not None
            else ""
        ),
        bool(available_backends[2]),
    )

    # Select FusedAttention for performance
    if use_flash_attention and use_fused_attention and device_compute_capability >= (9, 0):
        logger.debug(
            "Disabling FlashAttention to give FusedAttention preference on Hopper+ "
            "for performance reasons"
        )
        use_flash_attention = False

    # Selected backend
    if use_flash_attention:
        use_fused_attention = False
        use_unfused_attention = False
    elif use_fused_attention:
        use_unfused_attention = False
    selected_backend = "NoBackend"
    if use_flash_attention:
        selected_backend = f"FlashAttention ({str(flash_attention_backend)})"
    elif use_fused_attention:
        selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})"
    elif use_unfused_attention:
        selected_backend = "UnfusedDotProductAttention"
    logger.debug("Selected backend = %s.", selected_backend)

    return (
        use_flash_attention,
        flash_attention_backend,
        use_fused_attention,
        fused_attention_backend,
        use_unfused_attention,
        available_backends,
    )


@torch.no_grad()
def get_padding_mask(
    batch_size: int,
    cu_seqlens_q: torch.Tensor = None,
    cu_seqlens_kv: torch.Tensor = None,
    max_seqlen_q: int = None,
    max_seqlen_kv: int = None,
    attention_type: str = "self",
):
    """Convert cu_seqlens to attention_mask.

    Built with device-side ops only: reading the sequence lengths on the host
    would synchronize the device once per sequence.
    """
    assert (
        cu_seqlens_q is not None and max_seqlen_q is not None
    ), "cu_seqlens_q and max_seqlen_q are required for self-attention and cross-attention"

    def _mask(cu_seqlens: torch.Tensor, max_seqlen: int) -> torch.Tensor:
        # cu_seqlens may be longer than batch_size + 1 -- inference allocates it
        # for the maximum batch size -- so only its first batch_size + 1 entries
        # describe the current batch.
        seqlens = cu_seqlens[1 : batch_size + 1] - cu_seqlens[:batch_size]
        positions = torch.arange(max_seqlen, device=cu_seqlens.device)
        # True marks a padding token, i.e. one beyond the sequence length. The
        # mask is applied to the attention scores, so it goes on the device
        # those live on -- a no-op unless cu_seqlens is a CPU tensor.
        return (
            (positions.unsqueeze(0) >= seqlens.unsqueeze(1))
            .view(batch_size, 1, 1, max_seqlen)
            .to(device="cuda")
        )

    attention_mask_q = _mask(cu_seqlens_q, max_seqlen_q)
    if attention_type == "self":
        return attention_mask_q
    assert (
        cu_seqlens_kv is not None and max_seqlen_kv is not None
    ), "cu_seqlens_kv and max_seqlen_kv are required for cross-attention"
    return attention_mask_q, _mask(cu_seqlens_kv, max_seqlen_kv)


@torch.no_grad()
def get_full_mask(
    max_seqlen_q: int,
    max_seqlen_kv: int,
    attn_mask_type: str = "no_mask",
    attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]] = None,
    window_size: Tuple[int, int] = None,
    attention_type: str = "self",
    bottom_right_alignment: bool = True,
) -> torch.Tensor:
    """
    Get full attention mask in [..., max_seqlen_q, max_seqlen_kv] shape, based on `attn_mask_type`,
    `attention_mask`, and `window_size`. For sliding window attention, the diagonal alignment depends
    on both `attn_mask_type` and `bottom_right_alignment`, as detailed below.::

       attn_mask_type              output shape                                 diagonal alignment
       --------------------------------------------------------------------------------------------
       no_mask                     [1, 1, max_seqlen_q, max_seqlen_kv]          follow bottom_right_alignment
       causal                      [1, 1, max_seqlen_q, max_seqlen_kv]          always top left
       causal_bottom_right         [1, 1, max_seqlen_q, max_seqlen_kv]          always bottom right
       padding                     [batch_size, 1, max_seqlen_q, max_seqlen_kv] follow bottom_right_alignment
       padding_causal              [batch_size, 1, max_seqlen_q, max_seqlen_kv] always top left
       padding_causal_bottom_right [batch_size, 1, max_seqlen_q, max_seqlen_kv] always bottom right
       arbitrary                   same as attention_mask                       follow bottom_right_alignment

    .. note::

    For "padding_bottom_right" mask, or "padding" mask with `bottom_right_alignment` = True, the bottom right
    diagonal comes from the bottom right corner of the [actual_seqlens_q[i], actual_seqlens_kv[i]] matrix,
    i = 0,...,batch_size-1, not the [max_seqlen_q, max_seqlen_kv] matrix. For example, with max_seqlen_q = 4,
    max_seqlen_kv = 4, attn_mask_type = "padding", attention_type = "cross", and attention_mask = (
    [[False, False,  True, True], [False, False, False, False]],
    [[False, False, False, True], [False,  True,  True,  True]]), the returned full attention mask has [2, 4, 4]
    shape and is,::

      [[[False, False, False, True],
        [False, False, False, True],
        [ True,  True,  True, True],
        [ True,  True,  True, True]],
       [[False,  True,  True, True],
        [False,  True,  True, True],
        [False,  True,  True, True],
        [False,  True,  True, True]]]

    Parameters
    ----------
    max_seqlen_q : int
        Maximum sequence length for queries.
    max_seqlen_kv : int
        Maximum sequence length for keys and values.
    attn_mask_type : str, default = no_mask
        Attention mask type, {``"no_mask"``, ``"padding"``, ``"causal"``, ``"padding_causal"``,
        ``"causal_bottom_right"``, ``"padding_causal_bottom_right"``, ``"arbitrary"``}
    attention_mask : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        default = None
        Boolean tensor(s) used to mask out attention softmax input. Please see DotProductAttention
        for the requirements of `attention_mask` for different `attn_mask_type`s.
    window_size : Tuple[int, int], default = None
        Sliding window size for local attention, where query at position i attends to keys
        in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q
        + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding
        window and causal mask specifically. Both `causal` and `causal_bottom_right` masks
        map to `window_size = (-1, 0)` and Transformer Engine distinguishes them based on
        `attn_mask_type`.
    attention_type : str, default = "self"
        Attention type, {"self", "cross"}
    bottom_right_alignment : bool, default = True
        Whether to align the diagonal of the sliding window attention to the bottom right (`True`)
        or top left (`False`) corner of the softmax matrix. Ignored if `attn_mask_type` explicitly
        specifies "causal" or "causal_bottom_right".

    Returns
    ----------
    attn_mask_type : str
        For sliding window attention (>=0, >0), "arbitrary"; otherwise, the same as input `attn_mask_type`
    attention_mask : torch.Tensor
        The full attention mask based on `attn_mask_type`, `attention_mask` and `window_size`
    actual_seqlens_q : torch.Tensor
        For padding masks, the actual sequence lengths for queries, of shape [batch_size].
        For other masks, `None`.
    actual_seqlens_kv : Optional[torch.Tensor], default = None
        For padding masks, the actual sequence lengths for keys and values, of shape [batch_size].
        For other masks, `None`.
    """
    # perform basic checks
    change_type = window_size is not None and (
        window_size[0] != -1 or window_size[1] not in [-1, 0]
    )
    if window_size is None:
        window_size = (-1, -1)
    if "causal" in attn_mask_type:
        window_size = (window_size[0], 0)
    window_size = (
        max_seqlen_kv if window_size[0] == -1 else window_size[0],
        max_seqlen_q if window_size[1] == -1 else window_size[1],
    )

    # apply padding mask
    actual_seqlens_q = None
    actual_seqlens_kv = None
    if "padding" in attn_mask_type:
        if attention_type == "self":
            attention_mask = torch.logical_or(
                attention_mask.squeeze(1).unsqueeze(3), attention_mask
            )
        else:
            attention_mask = torch.logical_or(
                attention_mask[0].squeeze(1).unsqueeze(3), attention_mask[1]
            )
        m = attention_mask.logical_not()
        actual_seqlens_q = m[:, 0, :, 0].sum(dim=1)
        actual_seqlens_kv = m[:, 0, 0, :].sum(dim=1)

    # apply SWA mask
    mask = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view(
        1, 1, max_seqlen_q, 1
    ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view(1, 1, 1, max_seqlen_kv)
    swa_left = None
    swa_right = None
    if attn_mask_type == "causal_bottom_right" or (
        attn_mask_type in ["no_mask", "arbitrary"] and bottom_right_alignment
    ):
        swa_left = mask + max_seqlen_kv - max_seqlen_q - window_size[0]
        swa_right = mask + max_seqlen_kv - max_seqlen_q + window_size[1]
    elif attn_mask_type in ["causal", "padding_causal"] or (
        attn_mask_type in ["no_mask", "padding", "arbitrary"] and not bottom_right_alignment
    ):
        swa_left = mask - window_size[0]
        swa_right = mask + window_size[1]
    elif attn_mask_type == "padding_causal_bottom_right" or (
        attn_mask_type == "padding" and bottom_right_alignment
    ):
        batch_size = attention_mask.shape[0]
        swa_left = mask.expand(batch_size, 1, max_seqlen_q, max_seqlen_kv) + (
            actual_seqlens_kv - actual_seqlens_q - window_size[0]
        ).view(batch_size, 1, 1, 1)
        swa_right = mask.expand(batch_size, 1, max_seqlen_q, max_seqlen_kv) + (
            actual_seqlens_kv - actual_seqlens_q + window_size[1]
        ).view(batch_size, 1, 1, 1)
    swa_mask = torch.logical_not((swa_left <= 0) & ~(swa_right < 0))
    if attention_mask is not None:
        attention_mask = torch.logical_or(swa_mask, attention_mask)
    else:
        attention_mask = swa_mask

    # change mask type
    if change_type:
        attn_mask_type = "arbitrary"

    return attn_mask_type, attention_mask, actual_seqlens_q, actual_seqlens_kv


@torch.no_grad()
def get_alibi(
    _alibi_cache: Dict[str, Any],
    num_heads: int,
    max_seqlen_q: int,
    max_seqlen_kv: int,
    actual_seqlens_q: Optional[torch.Tensor] = None,
    actual_seqlens_kv: Optional[torch.Tensor] = None,
    alibi_slopes: Optional[torch.Tensor] = None,
    bias_dtype: Optional[torch.dtype] = None,
    bottom_right_alignment: bool = True,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Parameters
    ----------
    num_heads : int
        Number of heads.
    max_seqlen_q : int
        Maximum sequence length for queries.
    max_seqlen_kv : int
        Maximum sequence length for keys and values.
    actual_seqlens_q : Optional[torch.Tensor], default = None
        Actual sequence lengths for queries, of shape [batch_size].
    actual_seqlens_kv : Optional[torch.Tensor], default = None
        Actual sequence lengths for keys and values, of shape [batch_size].
    alibi_slopes : Optional[torch.Tensor], default = None
        Custom ALiBi slopes, FP32, CUDA tensor, of shape [num_heads] or [batch_size, num_heads].
    bias_dtype : Optional[torch.dtype], default = None
        Dtype of the generated ALiBi bias. If None, use torch.float32.
    bottom_right_alignment : bool, default = True
        Whether to align the diagonal of the ALiBi bias to the bottom right corner of
        the matrix (`True`) or top left (`False`).

    Returns
    ----------
    alibi_slopes : torch.Tensor
        ALiBi slopes in FP32 and shape [num_heads] or [batch_size, num_heads].
    alibi_bias : torch.Tensor
        ALiBi bias in FP32 or `bias_dtype`. Its shape is
        (1) [1, num_heads, max_seqlen_q, max_seqlen_kv] if `alibi_slopes` is in [num_heads] shape,
        and `actual_seqlens_q` and `actual_seqlens_kv` are `None`; or
        (2) [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] if `alibi_slopes` is in
        [batch_size, num_heads] shape, or, if `alibi_slopes` is in [num_heads] shape and
        `actual_seqlens_q` and `actual_seqlens_kv` are not `None`.
    """
    # NOTE: As part of refactoring attention.py, get_alibi() now receives the alibi cache from the caller
    # as an additional input arg
    if _alibi_cache["_alibi_slopes_require_update"]:
        if alibi_slopes is not None:
            _alibi_cache["_alibi_slopes"] = alibi_slopes
        else:
            n = 2 ** math.floor(math.log2(num_heads))
            m_0 = 2.0 ** (-8.0 / n)
            m = torch.pow(m_0, torch.arange(1, 1 + n))

            if n < num_heads:
                m_hat_0 = 2.0 ** (-4.0 / n)
                m_hat = torch.pow(m_hat_0, torch.arange(1, 1 + 2 * (num_heads - n), 2))
                m = torch.cat([m, m_hat])

            _alibi_cache["_alibi_slopes"] = m.to(dtype=torch.float32, device="cuda")
        _alibi_cache["_num_heads"] = num_heads
        _alibi_cache["_alibi_slopes_require_update"] = False

    if _alibi_cache["_alibi_bias_require_update"]:
        assert _alibi_cache["_alibi_slopes"] is not None, "ALiBi slopes can not be None!"
        if _alibi_cache["_alibi_slopes"].dim() == 1:
            slopes_shape = torch.Size([1, _alibi_cache["_alibi_slopes"].shape[0], 1, 1])
        elif _alibi_cache["_alibi_slopes"].dim() == 2:
            slopes_shape = torch.Size([*_alibi_cache["_alibi_slopes"].shape[:], 1, 1])
        else:
            raise ValueError("ALiBi slopes cannot exceed 2 dimensions.")

        bias = torch.arange(max_seqlen_q, dtype=torch.int32, device="cuda").view(
            1, 1, max_seqlen_q, 1
        ) - torch.arange(max_seqlen_kv, dtype=torch.int32, device="cuda").view(
            1, 1, 1, max_seqlen_kv
        )
        if actual_seqlens_q is None and actual_seqlens_kv is None:
            if bottom_right_alignment:
                bias = bias + max_seqlen_kv - max_seqlen_q
        elif actual_seqlens_q is not None and actual_seqlens_kv is not None:
            batch_size = actual_seqlens_q.shape[0]
            bias = bias.expand(batch_size, 1, max_seqlen_q, max_seqlen_kv)
            if bottom_right_alignment:
                bias = bias + (actual_seqlens_kv - actual_seqlens_q).view(batch_size, 1, 1, 1)
        else:
            assert (
                False
            ), "actual_seqlens_q and actual_seqlens_kv need to be both None or torch.Tensors!"
        bias = bias.abs().mul(-1)
        bias = bias * _alibi_cache["_alibi_slopes"].view(slopes_shape)
        _alibi_cache["_max_seqlen_q"], _alibi_cache["_max_seqlen_kv"] = max_seqlen_q, max_seqlen_kv
        _alibi_cache["_bottom_right_alignment"] = bottom_right_alignment
        bias_dtype = torch.float32 if bias_dtype is None else bias_dtype
        _alibi_cache["_alibi_bias"] = bias.contiguous().to(dtype=bias_dtype, device="cuda")
        _alibi_cache["_alibi_bias_require_update"] = False

    return _alibi_cache["_alibi_slopes"], _alibi_cache["_alibi_bias"]


def get_cu_seqlens(mask: torch.Tensor) -> torch.Tensor:
    """
    Given a padding mask of shape [batch_size, 1, 1, max_seqlen], returns an int32
    tensor of shape [batch_size + 1] containing the cumulative sequence lengths of
    the samples in a batch.
    """
    mask = mask.squeeze(1).squeeze(1)
    reduced_mask = mask.logical_not().sum(dim=1)
    cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
    zero = torch.zeros(1, dtype=torch.int32, device="cuda")
    cu_seqlens = torch.cat((zero, cu_seqlens))

    return cu_seqlens


def get_cu_seqlens_and_indices(mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Given a padding mask of shape [batch_size, 1, 1, max_seqlen], returns an int32
    tensor of shape [batch_size + 1] containing the cumulative sequence lengths of
    the samples in a batch, and another int32 tensor of shape [batch_size * max_seqlen, 1, 1]
    containing the indices for the valid tokens.
    """
    mask = mask.squeeze(1).squeeze(1)
    bs, seqlen = mask.shape

    reduced_mask = mask.logical_not().sum(dim=1)
    cu_seqlens = reduced_mask.cumsum(dim=0).to(torch.int32)
    zero = torch.zeros(1, dtype=torch.int32, device="cuda")
    cu_seqlens = torch.cat((zero, cu_seqlens))

    mask = mask.reshape(-1)
    indices = mask.logical_not().nonzero()
    indices = indices.unsqueeze(-1)

    num_nonzeros = indices.shape[0]
    pad_amount = bs * seqlen - num_nonzeros
    indices = F.pad(
        input=indices, pad=(0, 0, 0, 0, 0, pad_amount), mode="constant", value=float(bs * seqlen)
    )

    return cu_seqlens, indices


def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor:
    """
    Given max_seqlen and cu_seqlens of shape [batch_size + 1], returns an int32
    tensor of shape [batch_size * max_seqlen, 1, 1] containing the indices for
    the valid tokens in a batch.
    """
    bs = len(cu_seqlens) - 1
    seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
    indices = [i * max_seqlen + ii for i, j in enumerate(seqlens) for ii in range(j)]
    indices = torch.Tensor(indices).unsqueeze(1).unsqueeze(1).to(dtype=torch.int64, device="cuda")

    num_nonzeros = indices.shape[0]
    pad_amount = bs * max_seqlen - num_nonzeros
    indices = F.pad(
        input=indices,
        pad=(0, 0, 0, 0, 0, pad_amount),
        mode="constant",
        value=float(bs * max_seqlen),
    )

    return indices


def get_full_cu_seqlens(
    batch_size: int,
    max_seqlen: int,
    device: torch.device,
) -> torch.Tensor:
    """Cumulative sequence lengths in full data batch

    All sequences in batch have the maximum sequence length.

    """
    global _cu_seqlens_cache

    def _get_cu_seqlens(batch_size, max_seqlen, device):
        return torch.arange(
            0,
            (batch_size + 1) * max_seqlen,
            step=max_seqlen,
            dtype=torch.int32,
            device=device,
        )

    if is_in_onnx_export_mode():
        return _get_cu_seqlens(batch_size, max_seqlen, device)

    is_inference = torch.is_inference_mode_enabled()
    cu_seqlens_cache_key = (batch_size, max_seqlen, device, is_inference)
    if cu_seqlens_cache_key not in _cu_seqlens_cache:
        _cu_seqlens_cache[cu_seqlens_cache_key] = _get_cu_seqlens(batch_size, max_seqlen, device)
    return _cu_seqlens_cache[cu_seqlens_cache_key]


@jit_fuser
def _pack_tensor(
    indices: torch.Tensor,
    tensor: torch.Tensor,
) -> torch.Tensor:
    """
    Packs the given tensor using the `indices`.
    """
    dtype = tensor.dtype if not isinstance(tensor, Float8Tensor) else torch.uint8
    padding_indice = torch.zeros(
        1, tensor.shape[1], tensor.shape[2], dtype=dtype, device=tensor.device
    )
    indices = indices.repeat(1, tensor.shape[1], tensor.shape[2])
    if isinstance(tensor, Float8Tensor):
        tensor_data = torch.cat((tensor._data, padding_indice), dim=0)
        gathered_data = torch.gather(tensor_data, 0, indices)

        packed = Float8Tensor.make_like(tensor, data=gathered_data, shape=gathered_data.shape)
    else:
        tensor = torch.cat((tensor, padding_indice), dim=0)

        packed = torch.gather(tensor, 0, indices)
    return packed


@jit_fuser
def _pack_2_tensors(
    indices: torch.Tensor,
    t1: torch.Tensor,
    t2: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Packs the given 2 tensors using the `indices`.
    """
    t1_packed = _pack_tensor(indices, t1)
    t2_packed = _pack_tensor(indices, t2)
    return t1_packed, t2_packed


@jit_fuser
def _pack_3_tensors(
    indices: torch.Tensor,
    t1: torch.Tensor,
    t2: torch.Tensor,
    t3: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Packs the given 3 tensors using the `indices`.
    """
    t1_packed = _pack_tensor(indices, t1)
    t2_packed = _pack_tensor(indices, t2)
    t3_packed = _pack_tensor(indices, t3)
    return t1_packed, t2_packed, t3_packed


@jit_fuser
def _unpack_tensor(
    indices: torch.Tensor,
    dim0: int,
    tensor: torch.Tensor,
) -> torch.Tensor:
    """
    Inverse of `_pack_tensor`.
    """
    indices = indices.repeat(1, tensor.shape[1], tensor.shape[2])
    dtype = tensor.dtype if not isinstance(tensor, Float8Tensor) else torch.uint8
    unpacked = torch.zeros(
        dim0 + 1, tensor.shape[1], tensor.shape[2], dtype=dtype, device=tensor.device
    )
    if isinstance(tensor, Float8Tensor):
        unpacked.scatter_(0, indices, tensor._data)
        unpacked_data = unpacked[0:-1, :, :]
        unpacked = Float8Tensor.make_like(tensor, data=unpacked_data, shape=unpacked_data.shape)
    else:
        unpacked.scatter_(0, indices, tensor)
        unpacked = unpacked[0:-1, :, :]
    return unpacked


@jit_fuser
def _unpack_2_tensors(
    indices: torch.Tensor,
    dim0: int,
    t1: torch.Tensor,
    t2: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Inverse of `_pack_2_tensors`.
    """
    t1_unpacked = _unpack_tensor(indices, dim0, t1)
    t2_unpacked = _unpack_tensor(indices, dim0, t2)
    return t1_unpacked, t2_unpacked


@jit_fuser
def _unpack_3_tensors(
    indices: torch.Tensor,
    dim0: int,
    t1: torch.Tensor,
    t2: torch.Tensor,
    t3: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Inverse of `_pack_3_tensors`.
    """
    t1_unpacked = _unpack_tensor(indices, dim0, t1)
    t2_unpacked = _unpack_tensor(indices, dim0, t2)
    t3_unpacked = _unpack_tensor(indices, dim0, t3)
    return t1_unpacked, t2_unpacked, t3_unpacked


class PackTensors(torch.autograd.Function):
    """
    Autograd function to pack a tensor.
    """

    @staticmethod
    def forward(
        ctx, indices: torch.Tensor, *tensors: Tuple[torch.Tensor, ...]
    ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]:
        # pylint: disable=missing-function-docstring
        assert 1 <= len(tensors) <= 3, f"Packing {len(tensors)} tensors not supported."
        ctx.save_for_backward(indices)
        ctx.dim0 = tensors[0].shape[0]
        if len(tensors) == 1:
            return _pack_tensor(indices, *tensors)
        if len(tensors) == 2:
            return _pack_2_tensors(indices, *tensors)
        return _pack_3_tensors(indices, *tensors)

    @staticmethod
    def backward(ctx, *grad_outputs: Tuple[torch.Tensor, ...]):
        # pylint: disable=missing-function-docstring
        (indices,) = ctx.saved_tensors
        if len(grad_outputs) == 1:
            return None, _unpack_tensor(indices, ctx.dim0, *grad_outputs)
        if len(grad_outputs) == 2:
            return None, *_unpack_2_tensors(indices, ctx.dim0, *grad_outputs)
        return None, *_unpack_3_tensors(indices, ctx.dim0, *grad_outputs)


class UnpackTensor(torch.autograd.Function):
    """
    Autograd function to unpack a tensor.
    """

    @staticmethod
    def forward(
        ctx,
        indices: torch.Tensor,
        dim0: int,
        tensor: torch.Tensor,
    ) -> torch.Tensor:
        # pylint: disable=missing-function-docstring
        ctx.save_for_backward(indices)
        return _unpack_tensor(indices, dim0, tensor)

    @staticmethod
    def backward(ctx, grad_output):
        # pylint: disable=missing-function-docstring
        (indices,) = ctx.saved_tensors
        return None, None, _pack_tensor(indices, grad_output)


# ---------------------------------------------------------------------------
# THD <-> BSHD conversions exposed as `torch.library` custom ops so that
# `torch.compile` can trace them. Backward of each direction is the other
# direction, wired through `register_autograd` + `setup_context`, mirroring
# the pattern in `transformer_engine/pytorch/permutation.py`.
# ---------------------------------------------------------------------------


@torch.library.custom_op("te_attention::convert_thd_to_bshd", mutates_args=())
def _convert_thd_to_bshd_op(
    thd_tensor: torch.Tensor,
    cu_seqlens: torch.Tensor,
    batch_size: int,
    max_seqlen: int,
) -> torch.Tensor:
    """Forward pass for THD->BSHD conversion."""
    if not thd_tensor.is_contiguous():
        thd_tensor = thd_tensor.contiguous()
    return tex.convert_thd_to_bshd(thd_tensor, cu_seqlens, batch_size, max_seqlen)


@_convert_thd_to_bshd_op.register_fake
def _convert_thd_to_bshd_fake(
    thd_tensor: torch.Tensor,
    cu_seqlens: torch.Tensor,
    batch_size: int,
    max_seqlen: int,
) -> torch.Tensor:
    del cu_seqlens
    h, d = thd_tensor.shape[1], thd_tensor.shape[2]
    return torch.empty(
        (batch_size, max_seqlen, h, d), dtype=thd_tensor.dtype, device=thd_tensor.device
    )


@torch.library.custom_op("te_attention::convert_bshd_to_thd", mutates_args=())
def _convert_bshd_to_thd_op(
    bshd_tensor: torch.Tensor,
    cu_seqlens: torch.Tensor,
    num_tokens: int,
) -> torch.Tensor:
    """Forward pass for BSHD->THD conversion."""
    if not bshd_tensor.is_contiguous():
        bshd_tensor = bshd_tensor.contiguous()
    return tex.convert_bshd_to_thd(bshd_tensor, cu_seqlens, num_tokens)


@_convert_bshd_to_thd_op.register_fake
def _convert_bshd_to_thd_fake(
    bshd_tensor: torch.Tensor,
    cu_seqlens: torch.Tensor,
    num_tokens: int,
) -> torch.Tensor:
    del cu_seqlens
    h, d = bshd_tensor.shape[2], bshd_tensor.shape[3]
    return torch.empty((num_tokens, h, d), dtype=bshd_tensor.dtype, device=bshd_tensor.device)


def _convert_thd_to_bshd_setup_context(ctx, inputs, output):
    del output
    thd_tensor, cu_seqlens, _batch_size, _max_seqlen = inputs
    ctx.save_for_backward(cu_seqlens)
    ctx.num_tokens = thd_tensor.size(0)


def _convert_thd_to_bshd_backward_wrapper(ctx, grad_bshd):
    (cu_seqlens,) = ctx.saved_tensors
    grad_thd = torch.ops.te_attention.convert_bshd_to_thd(grad_bshd, cu_seqlens, ctx.num_tokens)
    return grad_thd, None, None, None


_convert_thd_to_bshd_op.register_autograd(
    _convert_thd_to_bshd_backward_wrapper,
    setup_context=_convert_thd_to_bshd_setup_context,
)


def _convert_bshd_to_thd_setup_context(ctx, inputs, output):
    del output
    bshd_tensor, cu_seqlens, _num_tokens = inputs
    ctx.save_for_backward(cu_seqlens)
    ctx.batch_size = bshd_tensor.size(0)
    ctx.max_seqlen = bshd_tensor.size(1)


def _convert_bshd_to_thd_backward_wrapper(ctx, grad_thd):
    (cu_seqlens,) = ctx.saved_tensors
    grad_bshd = torch.ops.te_attention.convert_thd_to_bshd(
        grad_thd, cu_seqlens, ctx.batch_size, ctx.max_seqlen
    )
    return grad_bshd, None, None


_convert_bshd_to_thd_op.register_autograd(
    _convert_bshd_to_thd_backward_wrapper,
    setup_context=_convert_bshd_to_thd_setup_context,
)


class ConvertTHDtoBSHD:
    """
    Convert a tensor from qkv_format = thd to qkv_format = bshd.

    Thin wrapper around the ``te_attention::convert_thd_to_bshd`` custom op,
    exposing an ``.apply(thd_tensor, cu_seqlens, max_seqlen)`` staticmethod so
    callsites keep the ``autograd.Function``-style ``.apply(...)`` invocation.
    """

    @staticmethod
    def apply(thd_tensor, cu_seqlens, max_seqlen):
        # pylint: disable=missing-function-docstring
        batch_size = cu_seqlens.shape[0] - 1
        return torch.ops.te_attention.convert_thd_to_bshd(
            thd_tensor, cu_seqlens, batch_size, max_seqlen
        )


class ConvertBSHDtoTHD:
    """
    Convert a tensor from qkv_format = bshd to qkv_format = thd.

    Thin wrapper around the ``te_attention::convert_bshd_to_thd`` custom op,
    exposing an ``.apply(bshd_tensor, cu_seqlens, num_tokens)`` staticmethod so
    callsites keep the ``autograd.Function``-style ``.apply(...)`` invocation.
    """

    @staticmethod
    def apply(bshd_tensor, cu_seqlens, num_tokens):
        # pylint: disable=missing-function-docstring
        return torch.ops.te_attention.convert_bshd_to_thd(bshd_tensor, cu_seqlens, num_tokens)


def get_qkv_format(
    qkv_layout: str = "bshd_bshd_bshd",
    inference_params: InferenceParams = None,
) -> str:
    """Get qkv format.

    Parameters
    ----------
    qkv_layout : str
       Memory layout of `q`, `k` and `v`. See get_qkv_layout() for more details.
    inference_params : InferenceParams, default = None
        InferenceParams related to KV caching.

    Returns
    ----------
    qkv_format : str, default = sbhd
        Dimension format for `q`, `k` and `v`, {`sbhd`, `bshd`, `thd`}.
    q_format : str
        Format of the `q` tensor, {`bshd`, `sbhd`, `thd`}.
    kv_format : str
        Format of the `k` and `v` tensors, {`bshd`, `sbhd`, `thd`}.
    """
    splited = qkv_layout.replace("paged_kv_", "").split("_")
    if inference_params is not None:
        q_format = "".join([i for i in splited[0] if i.isalpha()])
        kv_format = "".join([i for i in splited[1] if i.isalpha()])
        qkv_format = q_format + "_2" + kv_format if q_format != kv_format else q_format
    else:
        qkv_format = "".join([i for i in splited[0] if i.isalpha()])
        q_format = qkv_format
        kv_format = qkv_format
    return qkv_format, q_format, kv_format


def get_qkv_layout(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    qkv_format: str = "sbhd",
    inference_params: InferenceParams = None,
) -> str:
    """Get qkv layout.

    Parameters
    ----------
    q : torch.Tensor
        Query tensor.
    k : torch.Tensor
        Key tensor.
    v : torch.Tensor
        Value tensor.
    qkv_format : str, default = sbhd
        Dimension format for `q`, `k` and `v`, {`sbhd`, `bshd`, `thd`}. `s` stands for
        the sequence length dimension, `b` batch size, `h` the number of attention heads,
        `d` head size, and `t` the total number of tokens in a batch, i.e.
        `t = sum(s_i) for i = 0...b-1`.
    inference_params : InferenceParams, default = None
        InferenceParams related to KV caching.

    Returns
    ----------
    qkv_layout : str
       Memory layout of `q`, `k` and `v`. Each `qkv_layout` maps to a pair of `q_format` and
       `kv_format` in {`bshd`, `sbhd`, `thd`}. The `paged_kv_` prefix is used to indicate that
       paged KV caching is in play. A few examples of the layouts are as follows.

       (1) `sb3hd` means `q`, `k`, `v` are created as one chunk of memory and that they are
       interleaved in the `2`nd dimension. (2) `sbhd_sbh2d` means `q` and `kv` are created in
       two chunks and that `q` itself is contiguous and `k`, `v` are interleaved with each other
       in the `3`rd dimension, `k = kv[:,:,:,0,:]` and `v = kv[:,:,:,1,:]`. `q_format` and
       `kv_format` in this case are still both `sbhd`. (3) `paged_kv_thd_bshd_bshd` means `q` is
       created in `thd` and `k`, `v` are in `sbhd`. This is likely due to the cache format in
       paged KV caching.

       Mapping:
       `sbhd`: {`sb3hd`, `sbh3d`, `sbhd_sb2hd`, `sbhd_sbh2d`, `sbhd_sbhd_sbhd`, `paged_kv_sbhd_sbhd_sbhd`}
       `bshd`: {`bs3hd`, `bsh3d`, `bshd_bs2hd`, `bshd_bsh2d`, `bshd_bshd_bshd`, `paged_kv_bshd_bshd_bshd`}
       `thd` : {`t3hd`, `th3d`, `thd_t2hd`, `thd_th2d`, `thd_thd_thd`}
       `sbhd_2bshd`: {`sbhd_bshd_bshd`, `paged_kv_sbhd_bshd_bshd`}
       `bshd_2sbhd`: {`bshd_sbhd_sbhd`, `paged_kv_bshd_sbhd_sbhd`}
       `thd_2bshd`: {`thd_bshd_bshd`, `paged_kv_thd_bshd_bshd`}
       `thd_2sbhd`: {`thd_sbhd_sbhd`, `paged_kv_thd_sbhd_sbhd`}

    q : torch.Tensor
        Query tensor. It may be different from input `q` as we try to fit tensors to
        a supported layout.
    k : torch.Tensor
        Key tensor. It may be different from input `k` as we try to fit tensors to
        a supported layout.
    v : torch.Tensor
        Value tensor. It may be different from input `v` as we try to fit tensors to
        a supported layout.
    q_format : str
        Format of the query tensor, {`bshd`, `sbhd`, `thd`}.
    kv_format : str
        Format of the key and value tensors, {`bshd`, `sbhd`, `thd`}.
    """

    check_last_dim_contiguous = all(x.stride(-1) == 1 for x in [q, k, v])
    assert check_last_dim_contiguous, "q, k and v must have stride 1 in their last dimension!"
    if "_2" in qkv_format:
        q_format, kv_format = qkv_format.split("_2")
        is_same_q_kv_format = False
    else:
        q_format = qkv_format
        kv_format = qkv_format
        is_same_q_kv_format = True

    def run_iteratively(q, k, v):
        # check data pointers
        if is_in_onnx_export_mode():
            check_ptrs_qkv = False
            check_ptrs_qk = False
            check_ptrs_kv = False
        else:
            data_ptr = q.untyped_storage().data_ptr()
            check_ptrs_qkv = all(x.untyped_storage().data_ptr() == data_ptr for x in [q, k, v])
            check_ptrs_qk = all(x.untyped_storage().data_ptr() == data_ptr for x in [q, k])
            data_ptr = k.untyped_storage().data_ptr()
            check_ptrs_kv = all(x.untyped_storage().data_ptr() == data_ptr for x in [k, v])

        # check tensor shapes
        shape = q.shape
        check_shapes_qkv = all(shape == x.shape for x in [q, k, v])
        shape = k.shape
        check_shapes_kv = shape[:-1] == v.shape[:-1]

        # check tensor strides
        stride = q.stride()
        check_strides_qkv = all(stride == x.stride() for x in [q, k, v])
        check_strides_kv = tuple(sk / k.shape[-1] for sk in k.stride()[:-1]) == tuple(
            sv / v.shape[-1] for sv in v.stride()[:-1]
        )

        # check tensor offsets for h3d and 3hd layouts
        prod_h_d = q.shape[-1] * q.shape[-2]
        check_3hd_offsets = all(x.storage_offset() == i * prod_h_d for i, x in enumerate([q, k, v]))
        check_h3d_offsets = all(
            x.storage_offset() == i * q.shape[-1] for i, x in enumerate([q, k, v])
        )

        # check tensor offsets for hd_h2d and hd_2hd layouts
        prod_all_dims = [np.prod(x.shape) for x in [q, k]]
        offset = prod_all_dims[0] if check_ptrs_qkv else 0
        prod_h_d = k.shape[-1] * k.shape[-2]
        check_2hd_offsets = all(
            x.storage_offset() == (offset + i * prod_h_d) for i, x in enumerate([k, v])
        )
        check_h2d_offsets = all(
            x.storage_offset() == (offset + i * k.shape[-1]) for i, x in enumerate([k, v])
        )

        # check tensor offsets for hd_hd_hd layouts
        check_hd_offsets_qkv = (
            all(x.storage_offset() == sum(prod_all_dims[:i]) for i, x in enumerate([q, k, v]))
            if check_ptrs_qkv
            else all(x.storage_offset() == 0 for i, x in enumerate([q, k, v]))
        )
        check_hd_offsets_qk = (
            all(x.storage_offset() == sum(prod_all_dims[:i]) for i, x in enumerate([q, k]))
            if not check_ptrs_qkv and check_ptrs_qk
            else all(x.storage_offset() == 0 for i, x in enumerate([q, k]))
        )
        check_hd_offsets_kv = (
            all(x.storage_offset() == sum(prod_all_dims[1 : i + 1]) for i, x in enumerate([k, v]))
            if not check_ptrs_qkv and check_ptrs_kv
            else all(x.storage_offset() == 0 for i, x in enumerate([k, v]))
        )

        if check_ptrs_qkv and check_strides_qkv and check_shapes_qkv and check_3hd_offsets:
            # sb3hd, bs3hd, t3hd
            # one chunk of memory, qkv, with q, k, v interleaved at dim=-3 in qkv
            qkv_layout = qkv_format[:-2] + "3" + qkv_format[-2:]
        elif check_ptrs_qkv and check_strides_qkv and check_shapes_qkv and check_h3d_offsets:
            # sbh3d, bsh3d, th3d
            # one chunk of memory, qkv, with q, k, v interleaved at dim=-2 in qkv
            qkv_layout = qkv_format[:-1] + "3" + qkv_format[-1:]
        elif check_ptrs_kv and check_strides_kv and check_shapes_kv and check_2hd_offsets:
            # sbhd_sb2hd, bshd_bs2hd, thd_t2hd
            # two chunks of memory, q and kv, with k, v interleaved at dim=-3 in kv
            # q and kv may be disjoint or consecutive in memory, and when consecutive, they may
            # have the same data pointer, i.e. check_ptrs_qkv=True
            qkv_layout = qkv_format + "_" + qkv_format[:-2] + "2" + qkv_format[-2:]
        elif check_ptrs_kv and check_strides_kv and check_shapes_kv and check_h2d_offsets:
            # sbhd_sbh2d, bshd_bsh2d, thd_th2d
            # two chunks of memory, q and kv, with k, v interleaved at dim=-2 in kv
            # q and kv may be disjoint or consecutive in memory, and when consecutive, they may
            # have the same data pointer, i.e. check_ptrs_qkv=True
            qkv_layout = qkv_format + "_" + qkv_format[:-1] + "2" + qkv_format[-1:]
        elif (
            check_strides_kv
            and check_shapes_kv
            and (check_hd_offsets_qkv or check_hd_offsets_kv or check_hd_offsets_qk)
        ):
            # sbhd_sbhd_sbhd, bshd_bshd_bshd, thd_thd_thd
            # three chunks of memory, q, k and v, which may be disjoint or consecutive, and
            # when consecutive, they may have the same data pointer, i.e. check_ptrs_qkv=True or
            # check_ptrs_qk=True or check_ptrs_kv=True
            if is_same_q_kv_format:
                qkv_layout = "_".join(list([qkv_format]) * 3)
            else:
                qkv_layout = q_format + "_" + kv_format + "_" + kv_format
        else:
            qkv_layout = "not_supported"

        return qkv_layout

    if not is_in_onnx_export_mode():
        qkv_layout = run_iteratively(q, k, v)
    else:
        qkv_layout = "not_supported"
    if qkv_layout == "not_supported":
        # force q,k,v to be contiguous and run get_layout again
        q, k, v = [x.contiguous() for x in [q, k, v]]
        qkv_layout = run_iteratively(q, k, v)
    if qkv_layout == "not_supported":
        raise RuntimeError("The provided qkv memory layout is not supported!")

    if len(qkv_layout.split("_")) < 3:
        # q/k/v were recognized as views of a packed buffer only by inspecting
        # their data pointers, strides and storage offsets. Skip the nudge while
        # CPU offloading is enabled: offloading forces MultiheadAttention onto
        # its sliced-views fallback, so packed views reaching detection are
        # expected there and the caller has no migration option.
        if not is_cpu_offload_enabled():
            warnings.warn(
                "Relying on pointer-based detection of packed q/k/v layouts"
                f" (detected {qkv_layout!r}) is deprecated: pass the packed buffer"
                " explicitly via qkv_layer/kv_layer (with qkv_interleave_dim) to"
                " DotProductAttention instead.",
                DeprecationWarning,
                stacklevel=2,
            )

    if inference_params is not None and inference_params.is_paged:
        qkv_layout = "paged_kv_" + qkv_layout

    return qkv_layout, q, k, v, q_format, kv_format


def check_set_window_size(
    attn_mask_type: str,
    window_size: Tuple[int, int] = None,
):
    """Check if sliding window size is compliant with attention mask type.
    If not, set it to the appropriate size.

         attn_mask_type                              |   window_size
    -------------------------------------------------------------------------
    no_mask, padding, arbitrary                      | (-1, -1) or (>=0, >=0)
    causal, padding_causal                           | (-1,  0) or (>=0, 0)
    causal_bottom_right, padding_causal_bottom_right | (-1,  0) or (>=0, 0)
    """
    orig_window_size = window_size
    if "causal" in attn_mask_type:
        if orig_window_size is None:
            window_size = (-1, 0)
        elif orig_window_size == (-1, -1) or (
            orig_window_size[0] >= 0 and orig_window_size[1] != 0
        ):
            window_size = (orig_window_size[0], 0)
            warnings.warn(
                "window_size should be (-1, 0) or (>=0, 0) for attn_mask_type=" + attn_mask_type
            )
        elif orig_window_size != (-1, 0) and (orig_window_size[0] < 0 or orig_window_size[1] != 0):
            assert False, (
                "window_size should be (-1, 0) or (>=0, 0) for attn_mask_type=" + attn_mask_type
            )
    elif attn_mask_type in ["no_mask", "padding", "arbitrary"]:
        if orig_window_size is None:
            window_size = (-1, -1)
        elif orig_window_size == (-1, 0):
            window_size = (-1, -1)
            warnings.warn(
                "window_size should be (-1, -1) or (>=0, >=0) for attn_mask_type=" + attn_mask_type
            )
        elif orig_window_size != (-1, -1) and (orig_window_size[0] < 0 or orig_window_size[1] < 0):
            assert False, (
                "window_size should be (-1, -1) or (>=0, >=0) for attn_mask_type=" + attn_mask_type
            )
    else:
        assert False, "Invalid attn_mask_type: " + attn_mask_type
    return window_size


def get_attention_quantizers(fp8, quantizers):
    """Get the list of quantizers used in attention from the quantizers list."""
    if not fp8:
        return [None] * 6

    QKV_quantizer = quantizers["scaling_fwd"][META_QKV]
    QKV_quantizer.internal = False
    QKV_quantizer.set_usage(rowwise=True, columnwise=False)

    S_quantizer = quantizers["scaling_fwd"][META_S]
    S_quantizer.internal = True
    S_quantizer.set_usage(rowwise=True, columnwise=False)

    O_quantizer = quantizers["scaling_fwd"][META_O]
    O_quantizer.internal = False
    O_quantizer.set_usage(rowwise=True, columnwise=False)

    dO_quantizer = quantizers["scaling_bwd"][META_DO]
    dO_quantizer.internal = False
    dO_quantizer.set_usage(rowwise=True, columnwise=False)

    dP_quantizer = quantizers["scaling_bwd"][META_DP]
    dP_quantizer.internal = True
    dP_quantizer.set_usage(rowwise=True, columnwise=False)

    dQKV_quantizer = quantizers["scaling_bwd"][META_DQKV]
    dQKV_quantizer.internal = False
    dQKV_quantizer.set_usage(rowwise=True, columnwise=False)

    # MXFP8 attention: detect from the QKV quantizer instance rather than the
    # recipe predicate so that CustomRecipe (whose `mxfp8()` predicate returns
    # False) gets the same treatment as the built-in MXFP8 recipe. The kernel
    # handles S/dP internally for MXFP8, hence S/dP are nulled out.
    if isinstance(QKV_quantizer, MXFP8Quantizer):
        QKV_quantizer.columnwise_usage = True
        QKV_quantizer.optimize_for_gemm = True
        S_quantizer = None
        O_quantizer.columnwise_usage = True

        dO_quantizer.columnwise_usage = True
        dO_quantizer.optimize_for_gemm = True
        dP_quantizer = None
        dQKV_quantizer.columnwise_usage = True

    _fp8_types = (Float8Quantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer)
    # S/dP are intentionally None under MXFP8 attention; skip the type check
    # for those slots in that case.
    _allow_none = {"S", "dP"} if isinstance(QKV_quantizer, MXFP8Quantizer) else set()
    for _name, _q in [
        ("QKV", QKV_quantizer),
        ("O", O_quantizer),
        ("S", S_quantizer),
        ("dQKV", dQKV_quantizer),
        ("dO", dO_quantizer),
        ("dP", dP_quantizer),
    ]:
        if _q is None and _name in _allow_none:
            continue
        if not isinstance(_q, _fp8_types):
            raise TypeError(
                "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, "
                f"but {_name} quantizer is {type(_q).__name__}. "
                "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an "
                "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or "
                "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles "
                "(boundary slots like O output and dQKV grad-input)."
            )

    return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer


def print_quantizers(
    label,
    layer_number,
    QKV_quantizer,
    O_quantizer,
    S_quantizer,
    dQKV_quantizer,
    dO_quantizer,
    dP_quantizer,
):
    """Print the type and scale/amax of attention quantizers"""
    _to_print = _NVTE_DEBUG * _NVTE_DEBUG_LEVEL == 2
    if (
        _to_print
        and _print_layer == layer_number
        and (
            not dist.is_initialized() or (dist.is_initialized() and dist.get_rank() == _print_rank)
        )
    ):
        names = [
            "QKV_quantizer",
            "S_quantizer",
            "O_quantizer",
            "dO_quantizer",
            "dP_quantizer",
            "dQKV_quantizer",
        ]
        quantizers = [
            QKV_quantizer,
            S_quantizer,
            O_quantizer,
            dO_quantizer,
            dP_quantizer,
            dQKV_quantizer,
        ]
        if "forward" in label:
            names = names[:3]
            quantizers = quantizers[:3]
        if "backward" in label:
            names = names[3:]
            quantizers = quantizers[3:]
        for i, q in enumerate(quantizers):
            type_str = ""
            if q is None:
                type_str = "None"
            elif isinstance(q, Float8Quantizer):
                type_str = "DS"
            elif isinstance(q, Float8CurrentScalingQuantizer):
                type_str = "CS"
            elif isinstance(q, MXFP8Quantizer):
                type_str = "MXFP8"
            if type_str == "DS":
                print(
                    f"{label} >> {names[i]:14s}: {type_str}, {q.scale.item():.4e} x"
                    f" {q.amax.item():.4e} = {q.scale.item()*q.amax.item():.4e}"
                )
            else:
                print(f"{label} >> {names[i]:14s}: {type_str}")


def transpose_to_bhsd_htd_pytorch(tensor, src_format):
    """Permute to BHSD or HTD format using native PyTorch operations."""
    if src_format in ("bhsd", "htd"):
        return tensor
    dim_s = src_format.find("s") if "s" in src_format else src_format.find("t")
    dim_others = [i for i in range(tensor.ndim) if i != dim_s]
    new_dims = [*dim_others[:-1], dim_s, dim_others[-1]]
    return tensor.permute(*new_dims).contiguous()


def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format):
    """MXFP8 attention requires quantization along S and D dimensions. This fast path
    quantizes tensors without swizzle, and pads, permutes and swizzles the scale_invs
    to achieve faster speed due to the smaller sizes of scale_invs compare to the data.
    The output tensors have _rowwise_data and _columnwise_data in src_format, and
    _rowwise_scale_inv and _columnwise_scale_inv in BHSD format.

    Parameters
    ----------
    tensor_quantizer_pairs : list of (torch.Tensor, MXFP8Quantizer)
        Each pair is a tensor and its quantizer (with the desired
        rowwise_usage / columnwise_usage already set).
    src_format : str
        Layout of input tensors: ``"bshd"`` or ``"sbhd"``.
        All tensors in the list must have the same src_format.
    Returns
    -------
    fp8_tensors : list of MXFP8Tensors
        Data in ``src_format``, scale_inv in BHSD format.
    scale_inv_format : str
        Always ``"bhsd"``.
    """
    if not tensor_quantizer_pairs:
        return [], src_format

    fp8_tensors = mxfp8_quantize_only(tensor_quantizer_pairs, src_format)
    mxfp8_transpose_swizzle(fp8_tensors, src_format)
    return fp8_tensors, "bhsd"


def mxfp8_quantize_only(tensor_quantizer_pairs, src_format):
    """Phase 1 of mxfp8_quantize_fast_path: quantize only, no BHSD transpose or GEMM swizzle.

    Returns MXFP8Tensors with data and scale_invs reshaped to src_format layout.
    Call mxfp8_transpose_swizzle to complete the BHSD permute + swizzle when ready
    (e.g. after pre-quantized tensors from fused kernels are also available).

    Parameters
    ----------
    tensor_quantizer_pairs : list of (torch.Tensor, MXFP8Quantizer)
        Same contract as mxfp8_quantize_fast_path.
    src_format : str
        ``"bshd"`` or ``"sbhd"``.

    Returns
    -------
    fp8_tensors : list of MXFP8Tensor
        Data and scale_invs in src_format layout; NOT yet BHSD-permuted or swizzled.
    """
    if not tensor_quantizer_pairs:
        return []
    assert src_format in (
        "bshd",
        "sbhd",
    ), f"mxfp8_quantize_only only supports bshd/sbhd, got {src_format!r}."
    _s_dim = {"bshd": 1, "sbhd": 0}
    _d_dim = {"bshd": 3, "sbhd": 3}

    fp8_tensors = []
    for tensor, quantizer in tensor_quantizer_pairs:
        original_shape = tensor.shape
        rs_shape = list(original_shape)
        rs_shape[_d_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE
        cs_shape = list(original_shape)
        cs_shape[_s_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE
        if src_format == "bshd":
            t2d = tensor.view(*tensor.shape[:2], -1)
        else:
            t2d = tensor.view(tensor.shape[0], -1)
        orig_optimize = quantizer.optimize_for_gemm
        quantizer.optimize_for_gemm = False
        fp8_2d = quantizer(t2d)
        quantizer.optimize_for_gemm = orig_optimize
        # Re-wrap with the original 4D SBHD/BSHD shape so that shape[-1] equals the per-head
        # dimension (matching Q's wrapper shape) and fused_attn_bwd produces 4D dkv that
        # matches key/value's expected gradient shape in _KFQuantizeKVForAttn.backward.
        fp8_t = MXFP8Tensor(
            shape=original_shape,
            dtype=tensor.dtype,
            rowwise_data=(
                fp8_2d._rowwise_data.view(original_shape)
                if fp8_2d._rowwise_data is not None
                else None
            ),
            rowwise_scale_inv=(
                fp8_2d._rowwise_scale_inv.view(rs_shape)
                if fp8_2d._rowwise_scale_inv is not None
                else None
            ),
            columnwise_data=(
                fp8_2d._columnwise_data.view(original_shape)
                if fp8_2d._columnwise_data is not None
                else None
            ),
            columnwise_scale_inv=(
                fp8_2d._columnwise_scale_inv.view(cs_shape)
                if fp8_2d._columnwise_scale_inv is not None
                else None
            ),
            quantizer=quantizer,
            requires_grad=False,
            fp8_dtype=fp8_2d._fp8_dtype,
            with_gemm_swizzled_scales=False,
        )
        fp8_tensors.append(fp8_t)
    return fp8_tensors


def mxfp8_transpose_swizzle(fp8_tensors, src_format):
    """Phase 2 of mxfp8_quantize_fast_path: batched BHSD-transpose + GEMM-swizzle.

    For tensors whose data is already quantized (e.g. from a fused GEMM+quant kernel
    or from mxfp8_quantize_only), permutes each tensor's scale_invs from src_format to
    BHSD and applies the GEMM swizzle in-place.  Complements mxfp8_quantize_only to
    allow pre-quantized tensors (like a fused-kernel Q) to be processed in the same
    batched operation as freshly quantized K/V.

    Parameters
    ----------
    fp8_tensors : list of MXFP8Tensor
        Tensors with _rowwise_scale_inv / _columnwise_scale_inv in src_format layout.
        Modified in-place: scale_invs are replaced with BHSD-permuted, swizzled versions.
    src_format : str
        ``"bshd"`` or ``"sbhd"``.
    """
    if not fp8_tensors:
        return

    assert src_format in (
        "bshd",
        "sbhd",
    ), f"mxfp8_transpose_swizzle only supports bshd/sbhd, got {src_format!r}."

    rs_list = [t._rowwise_scale_inv for t in fp8_tensors]
    cs_list = [t._columnwise_scale_inv for t in fp8_tensors]

    def _align_up(x, a):
        return ((x + a - 1) // a) * a

    def _bhsd_shape(src_4d, d_pad):
        if src_format == "sbhd":
            S, B, H, _ = src_4d.shape
        else:
            B, S, H, _ = src_4d.shape
        return (B, H, S, d_pad)

    def _build_outputs(scale_list, alignment):
        entries = []
        total = 0
        for s in scale_list:
            if s is None:
                entries.append(None)
                continue
            d_pad = _align_up(s.shape[-1], alignment)
            shape = _bhsd_shape(s, d_pad)
            numel = 1
            for dim in shape:
                numel *= dim
            entries.append((total, numel, shape))
            total += numel
        if total == 0:
            return [None] * len(scale_list)
        device = next(s for s in scale_list if s is not None).device
        buf = torch.empty(total, dtype=torch.uint8, device=device)
        return [buf[e[0] : e[0] + e[1]].view(e[2]) if e is not None else None for e in entries]

    rs_outs = _build_outputs(rs_list, 4)
    cs_outs = _build_outputs(cs_list, 128)

    rs_permuted = tex.multi_tensor_transpose_to_bhsd(
        rs_list, original_format=src_format, outputs=rs_outs
    )
    cs_permuted = tex.multi_tensor_transpose_to_bhsd(
        cs_list, original_format=src_format, outputs=cs_outs
    )

    for t, rp, cp in zip(fp8_tensors, rs_permuted, cs_permuted):
        t._rowwise_scale_inv = rp.view(-1, rp.shape[-1]) if rp is not None else None
        t._columnwise_scale_inv = cp.view(-1, cp.shape[-1]) if cp is not None else None

    tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, True, False)
    tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, False, True)
    for t in fp8_tensors:
        t._with_gemm_swizzled_scales = True


def combine_and_quantize(
    qkv_layout,
    q,
    k,
    v,
    qkv_quantizer,
    used_in_forward=True,
    used_in_backward=False,
    keep_same_data_and_scale_inv_format=False,
    combined_qkv: Optional[torch.Tensor] = None,
    combined_kv: Optional[torch.Tensor] = None,
):
    """Combine Q, K, V tensors based on qkv_layout and quantize them together.

    When ``combined_qkv`` (for ``qkv_group=1`` layouts such as ``bs3hd``) or
    ``combined_kv`` (for ``qkv_group=2`` layouts such as ``bshd_bs2hd``) is provided, it must be the
    caller's original packed buffer that q/k/v are views of. It is then quantized
    directly instead of re-deriving the packed buffer from the q/k/v views via
    ``combine_tensors`` (which rebuilds it with a raw ``set_`` under a silent
    adjacency/interleave assumption). Ignored for MXFP8 quantization.
    """
    if isinstance(qkv_quantizer, MXFP8Quantizer):
        qkv_format, q_format, kv_format = get_qkv_format(qkv_layout)
        assert qkv_format in ("bshd", "sbhd"), (
            "combine_and_quantize only supports bshd/sbhd for MXFP8 quantization, got"
            f" {qkv_format!r}."
        )

        _s_dim = {"sbhd": 0, "bshd": 1}
        _d_dim = {"sbhd": 3, "bshd": 3}
        d_qk = q.shape[_d_dim[qkv_format]]
        d_v = v.shape[_d_dim[qkv_format]]
        s_q = q.shape[_s_dim[q_format]]
        s_kv = v.shape[_s_dim[kv_format]]
        assert s_q % 128 == 0 and s_kv % 128 == 0 and d_qk % 32 == 0 and d_v % 32 == 0, (
            "MXFP8 quantization requires s_q % 128 == 0, s_kv % 128 == 0, d_qk % 32 == 0, d_v % 32"
            f" == 0. Found {s_q=}, {s_kv=}, {d_qk=}, {d_v=}."
        )

        if qkv_layout not in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"):
            keep_same_data_and_scale_inv_format = True

        # ---- Fast path: quantize in original layout, permute scale_inv to BHSD, then swizzle ----
        if not keep_same_data_and_scale_inv_format:
            q_quantizer, k_quantizer, v_quantizer = [qkv_quantizer.copy() for _ in range(3)]
            if used_in_forward and not used_in_backward:
                q_quantizer.rowwise_usage = True
                q_quantizer.columnwise_usage = False
                k_quantizer.rowwise_usage = True
                k_quantizer.columnwise_usage = False
                v_quantizer.rowwise_usage = False
                v_quantizer.columnwise_usage = True
            elif (not used_in_forward) and used_in_backward:
                q_quantizer.rowwise_usage = True
                q_quantizer.columnwise_usage = True
                k_quantizer.rowwise_usage = True
                k_quantizer.columnwise_usage = True
                v_quantizer.rowwise_usage = True
                v_quantizer.columnwise_usage = False
            (q_fp8, k_fp8, v_fp8), qkv_scale_inv_format = mxfp8_quantize_fast_path(
                [(q, q_quantizer), (k, k_quantizer), (v, v_quantizer)], qkv_format
            )
            return q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format

        # ---- Slow path: permute data to BHSD, then quantize with swizzle ----
        if qkv_layout in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"):
            q, k, v = tex.multi_tensor_transpose_to_bhsd(
                [q, k, v],
                original_format=qkv_format,
            )
        else:
            q = transpose_to_bhsd_htd_pytorch(q, q_format)
            k = transpose_to_bhsd_htd_pytorch(k, kv_format)
            v = transpose_to_bhsd_htd_pytorch(v, kv_format)
        qkv_layout = "bhsd_bhsd_bhsd"
        qkv_scale_inv_format = "bhsd"

        original_shapes = [x.shape for x in [q, k, v]]
        q, k, v = [x.view(-1, x.shape[-1]) for x in [q, k, v]]

        q_quantizer, k_quantizer, v_quantizer = [qkv_quantizer.copy() for _ in range(3)]
        if used_in_forward and not used_in_backward:
            q_quantizer.rowwise_usage = True
            q_quantizer.columnwise_usage = False
            k_quantizer.rowwise_usage = True
            k_quantizer.columnwise_usage = False
            v_quantizer.rowwise_usage = False
            v_quantizer.columnwise_usage = True
        elif (not used_in_forward) and used_in_backward:
            q_quantizer.rowwise_usage = True
            q_quantizer.columnwise_usage = True
            k_quantizer.rowwise_usage = True
            k_quantizer.columnwise_usage = True
            v_quantizer.rowwise_usage = True
            v_quantizer.columnwise_usage = False
        q_fp8, k_fp8, v_fp8 = [
            quant(x) for quant, x in zip([q_quantizer, k_quantizer, v_quantizer], [q, k, v])
        ]

        for fp8_tensor, shape in zip([q_fp8, k_fp8, v_fp8], original_shapes):
            fp8_tensor._rowwise_data = (
                fp8_tensor._rowwise_data.view(shape)
                if fp8_tensor._rowwise_data is not None
                else None
            )
            fp8_tensor._columnwise_data = (
                fp8_tensor._columnwise_data.view(shape)
                if fp8_tensor._columnwise_data is not None
                else None
            )

        return q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format

    qkv_layout = qkv_layout.replace("paged_kv_", "")
    qkv_group = len(qkv_layout.split("_"))
    src_nominal_dtype = q.dtype
    # 1: qkv packed, 2: kv packed, 3: qkv separate
    match qkv_group:
        case 1:
            dim = qkv_layout.find("3")
            if combined_qkv is not None:
                assert combined_qkv.shape[dim] == 3, (
                    f"combined_qkv does not match qkv_layout {qkv_layout}: expected"
                    f" size 3 at dim {dim}, got shape {tuple(combined_qkv.shape)}."
                )
                qkv = combined_qkv
            else:
                qkv = combine_tensors([q, k, v], dim)
            qkv_fp8 = qkv_quantizer(qkv)
            q_data, k_data, v_data = SplitAlongDim.apply(qkv_fp8._data, dim, [1, 1, 1], True)
        case 2:
            dim = qkv_layout.split("_")[1].find("2")
            if combined_kv is not None:
                assert combined_kv.shape[dim] == 2, (
                    f"combined_kv does not match qkv_layout {qkv_layout}: expected"
                    f" size 2 at dim {dim}, got shape {tuple(combined_kv.shape)}."
                )
                kv = combined_kv
            else:
                kv = combine_tensors([k, v], dim)
            tensors = [q, kv]
            num_tensors = len(tensors)
            shapes = [x.shape for x in tensors]
            numels = [x.numel() for x in tensors]
            numels = [sum(numels[:i]) for i in range(num_tensors + 1)]
            qkv = torch.cat([x.view(-1) for x in tensors], dim=0)
            qkv_fp8 = qkv_quantizer(qkv)
            q_data, kv_data = [
                qkv_fp8._data[numels[i] : numels[i + 1]].view(shapes[i]) for i in range(num_tensors)
            ]
            k_data, v_data = SplitAlongDim.apply(kv_data, dim, [1, 1], True)
        case 3:
            tensors = [q, k, v]
            num_tensors = len(tensors)
            shapes = [x.shape for x in tensors]
            numels = [x.numel() for x in tensors]
            numels = [sum(numels[:i]) for i in range(num_tensors + 1)]
            qkv = torch.cat([x.view(-1) for x in tensors], dim=0)
            qkv_fp8 = qkv_quantizer(qkv)
            q_data, k_data, v_data = [
                qkv_fp8._data[numels[i] : numels[i + 1]].view(shapes[i]) for i in range(num_tensors)
            ]
        case _:
            raise RuntimeError("Invalid qkv_layout " + qkv_layout)

    q_fp8, k_fp8, v_fp8 = [
        Float8Tensor.make_like(qkv_fp8, data=x, dtype=src_nominal_dtype)
        for x in [q_data, k_data, v_data]
    ]

    return q_fp8, k_fp8, v_fp8, qkv_layout, None


def combine_and_dequantize(
    qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=None, des_nominal_dtype=None
):
    """Combine q,k,v based on qkv_layout and dequantize them together"""
    if all(isinstance(x, QuantizedTensorStorage) for x in [q_fp8, k_fp8, v_fp8]):
        src_nominal_dtype = q_fp8.dtype
    else:
        assert src_nominal_dtype is not None, "The nominal dtype of input tensors is required!"
    if des_nominal_dtype is None:
        des_nominal_dtype = src_nominal_dtype

    if all(isinstance(x, (MXFP8Tensor, MXFP8TensorStorage)) for x in [q_fp8, k_fp8, v_fp8]):
        q, k, v = [x.dequantize(dtype=des_nominal_dtype) for x in [q_fp8, k_fp8, v_fp8]]
        return q, k, v

    qkv_layout = qkv_layout.replace("paged_kv_", "")
    qkv_group = len(qkv_layout.split("_"))
    q_data, k_data, v_data = [x._data for x in [q_fp8, k_fp8, v_fp8]]
    # 1: qkv packed, 2: kv packed, 3: qkv separate
    match qkv_group:
        case 1:
            dim = qkv_layout.find("3")
            qkv_data = combine_tensors([q_data, k_data, v_data], dim)
            qkv_fp8 = Float8Tensor.make_like(q_fp8, data=qkv_data)
            qkv = qkv_fp8.dequantize(dtype=des_nominal_dtype)
            q, k, v = SplitAlongDim.apply(qkv, dim, [1, 1, 1], True)
        case 2:
            dim = qkv_layout.split("_")[1].find("2")
            kv_data = combine_tensors([k_data, v_data], dim)
            tensors = [q_data, kv_data]
            num_tensors = len(tensors)
            shapes = [x.shape for x in tensors]
            numels = [x.numel() for x in tensors]
            numels = [sum(numels[:i]) for i in range(num_tensors + 1)]
            qkv_data = torch.cat([x.reshape(-1) for x in tensors], dim=0)
            qkv_fp8 = Float8Tensor.make_like(q_fp8, data=qkv_data, dtype=src_nominal_dtype)
            qkv = qkv_fp8.dequantize(dtype=des_nominal_dtype)
            q, kv = [qkv[numels[i] : numels[i + 1]].view(shapes[i]) for i in range(num_tensors)]
            k, v = SplitAlongDim.apply(kv, dim, [1, 1], True)
        case 3:
            tensors = [q_data, k_data, v_data]
            num_tensors = len(tensors)
            shapes = [x.shape for x in tensors]
            numels = [x.numel() for x in tensors]
            numels = [sum(numels[:i]) for i in range(num_tensors + 1)]
            qkv_data = torch.cat([x.contiguous().reshape(-1) for x in tensors], dim=0)
            qkv_fp8 = Float8Tensor.make_like(q_fp8, data=qkv_data, dtype=src_nominal_dtype)
            qkv = qkv_fp8.dequantize(dtype=des_nominal_dtype)
            q, k, v = [qkv[numels[i] : numels[i + 1]].view(shapes[i]) for i in range(num_tensors)]
        case _:
            raise RuntimeError("Invalid qkv_layout " + qkv_layout)
    return q, k, v
