# Source code for tensorplay.nn.attention.omni_attention

Source: https://www.tensorplay.cn/docs/_modules/tensorplay/nn/attention/omni_attention.html

```
# mypy: allow-untyped-defs
"""User-facing API for omni attention: block-sparse mask metadata and
attention with an arbitrary attention-score modification function."""

from __future__ import annotations

import functools
import inspect
import itertools
import math
import operator
import types
import typing
import warnings
from collections.abc import Callable
from enum import Enum
from typing import Any, cast, Literal, NamedTuple, overload, TypeAlias, TypeVar
from typing_extensions import deprecated, Never, NotRequired, Self, TypedDict

import tensorplay
from tensorplay import Tensor
from tensorplay.graph import GraphCaptureError
from tensorplay._higher_order_ops.omni_attention import (
    omni_attention as omni_attention_hop,
)
from tensorplay._higher_order_ops.utils import setup_compilation_env
from tensorplay.nn.attention._utils import _validate_sdpa_input
from tensorplay.utils._pytree import (
    GetAttrKey,
    tree_flatten,
    tree_map_only,
    tree_unflatten,
    TreeSpec,
)

if typing.TYPE_CHECKING:
    from tensorplay.graph.node import BaseArgumentTypes

    DeviceLikeType = str | tensorplay.device

# Private debug flag to disable internal compilation wrapping for debugging purposes.
# WARNING: This is intended ONLY for debugging score_mod and mask_mod functions.
# When enabled, this bypasses the required internal compilation that ensures correctness
# and performance. Only use this temporarily when you need to set breakpoints
# in your score_mod/mask_mod functions during development.
#
# This flag only affects the internal compilation when omni_attention is called directly.
# If you have already wrapped omni_attention in compile(), this flag has no effect
# and the user's compilation will still occur.
#
# Usage:
#   import tensorplay.nn.attention.omni_attention as oa
#   oa._OMNI_ATTENTION_DISABLE_COMPILE_DEBUG = True
#   # Now you can set breakpoints in your score_mod/mask_mod
#   output = oa.omni_attention(q, k, v, score_mod=my_score_mod)
#
_OMNI_ATTENTION_DISABLE_COMPILE_DEBUG = False

_BLOCK_MASK_TOO_SMALL_ERROR = (
    "block_mask was created for a smaller length than you're "
    "using it for, you likely need to create a new block mask."
)
_BLOCK_MASK_TOO_LARGE_ERROR = (
    "block_mask was created for a larger length than you're "
    "using it for, you can either 1. create a new block mask with "
    "the correct length, or 2. 'adjust' the existing block mask to "
    "the correct length by calling block_mask._adjust(q_len, kv_len). "
    "This essentially 'crops' the block mask to the upper left corner, "
    "which does not work for all mask_mods!"
)

_WARNINGS_SHOWN: set[str] = set()

def _warn_once(
    warning_id: str, message: str, category: type[Warning] = UserWarning
) -> None:
    """Helper to ensure each warning is shown only once per process."""
    if warning_id not in _WARNINGS_SHOWN:
        if not tensorplay.compiler.is_compiling():
            warnings.warn(message, category, stacklevel=2)
        _WARNINGS_SHOWN.add(warning_id)

def _validate_block_mask_shape(
    query: Tensor,
    key: Tensor,
    q_len: int | tensorplay.SymInt,
    kv_len: int | tensorplay.SymInt,
    block_mask_q_len: int | tensorplay.SymInt,
    block_mask_kv_len: int | tensorplay.SymInt,
) -> None:
    """Preserve unbacked checks without specializing regular dynamic lengths."""
    # Imported lazily: the symbolic-shapes package pulls in the graph passes,
    # which import back into this package.
    from tensorplay.graph.experimental.symbolic_shapes import has_free_unbacked_symbols

    has_unbacked_input_lengths = has_free_unbacked_symbols(
        query
    ) or has_free_unbacked_symbols(key)
    if not tensorplay.compiler.is_compiling():
        lengths = (q_len, kv_len, block_mask_q_len, block_mask_kv_len)
        has_unbacked_input_lengths = (
            has_unbacked_input_lengths or has_free_unbacked_symbols(lengths)
        )

    # Case 1: unbacked lengths need symbolic checks that can become runtime assertions.
    if has_unbacked_input_lengths:
        tensorplay._check(q_len <= block_mask_q_len, _BLOCK_MASK_TOO_SMALL_ERROR)
        tensorplay._check(kv_len <= block_mask_kv_len, _BLOCK_MASK_TOO_SMALL_ERROR)
        tensorplay._check(q_len >= block_mask_q_len, _BLOCK_MASK_TOO_LARGE_ERROR)
        tensorplay._check(kv_len >= block_mask_kv_len, _BLOCK_MASK_TOO_LARGE_ERROR)
        return

    # Case 2: backed lengths use regular validation so the tracer can generalize equality.
    if q_len > block_mask_q_len or kv_len > block_mask_kv_len:
        raise RuntimeError(_BLOCK_MASK_TOO_SMALL_ERROR)
    if q_len < block_mask_q_len or kv_len < block_mask_kv_len:
        raise RuntimeError(_BLOCK_MASK_TOO_LARGE_ERROR)

    if q_len != block_mask_q_len:
        raise AssertionError(
            f"query.size(-2) ({q_len}) != block_mask_q_len ({block_mask_q_len})"
        )
    if kv_len != block_mask_kv_len:
        raise AssertionError(
            f"key.size(-2) ({kv_len}) != block_mask_kv_len ({block_mask_kv_len})"
        )

__all__ = [
    "BlockMask",
    "omni_attention",
    "AuxOutput",
    "AuxRequest",
    "OmniKernelOptions",
    "create_block_mask",
    "create_mask",
    "or_masks",
    "and_masks",
    "noop_mask",
]
_score_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor, Tensor], Tensor]
_mask_mod_signature = Callable[[Tensor, Tensor, Tensor, Tensor], Tensor]
_Backend: TypeAlias = Literal["AUTO", "TRITON", "FLASH", "TRITON_DECODE"]
_R = TypeVar("_R")

class OmniKernelOptions(TypedDict, total=False):
    """Options for controlling the behavior of omni attention kernels.

    These options are passed to the underlying kernels to control performance
    and numerical behavior. Most users will not need to specify these options as the
    default autotuning provides good performance.

    The options can be prefixed with ``fwd_`` or ``bwd_`` to apply only to forward or
    backward pass respectively. For example: ``fwd_BLOCK_M`` and ``bwd_BLOCK_M1``.

    Note:
      We currently do not provide any backward compatibility guarantees for these options.
      That being said most of these have remained pretty stable since their introduction. But
      We do not consider this part of the public API just yet. We think that some documentation
      Is better than secret hidden flags, but we may change these options in the future.

    Example Usage:
        .. code-block:: python

            # Using dictionary (backward compatible)
            kernel_opts = {"BLOCK_M": 64, "BLOCK_N": 64, "PRESCALE_QK": True}
            output = omni_attention(q, k, v, kernel_options=kernel_opts)

            # Using TypedDict (recommended for type safety)
            from tensorplay.nn.attention.omni_attention import OmniKernelOptions

            kernel_opts: OmniKernelOptions = {
                "BLOCK_M": 64,
                "BLOCK_N": 64,
                "PRESCALE_QK": True,
            }
            output = omni_attention(q, k, v, kernel_options=kernel_opts)

            # Forward/backward specific options
            kernel_opts: OmniKernelOptions = {
                "fwd_BLOCK_M": 64,
                "bwd_BLOCK_M1": 32,
                "PRESCALE_QK": False,
            }
            output = omni_attention(q, k, v, kernel_options=kernel_opts)
    """

    # Performance tuning options

    num_warps: NotRequired[int]
    """Number of warps to use in the kernel. Higher values may improve performance
    but increase register pressure. Default is determined by autotuning."""

    num_stages: NotRequired[int]
    """Number of pipeline stages in the kernel. Higher values may improve performance
    but increase shared memory usage. Default is determined by autotuning."""

    BLOCK_M: NotRequired[int]
    """Thread block size for the sequence length dimension of Q in forward pass.
    Must be a power of 2. Common values: 16, 32, 64, 128. Default is determined by autotuning."""

    BLOCK_N: NotRequired[int]
    """Thread block size for the sequence length dimension of K/V in forward pass.
    Must be a power of 2. Common values: 16, 32, 64, 128. Default is determined by autotuning."""

    # Backward-specific block sizes (when prefixed with 'bwd_')

    BLOCK_M1: NotRequired[int]
    """Thread block size for Q dimension in backward pass. Use as 'bwd_BLOCK_M1'.
    Default is determined by autotuning."""

    BLOCK_N1: NotRequired[int]
    """Thread block size for K/V dimension in backward pass. Use as 'bwd_BLOCK_N1'.
    Default is determined by autotuning."""

    BLOCK_M2: NotRequired[int]
    """Thread block size for second Q dimension in backward pass. Use as 'bwd_BLOCK_M2'.
    Default is determined by autotuning."""

    BLOCK_N2: NotRequired[int]
    """Thread block size for second K/V dimension in backward pass. Use as 'bwd_BLOCK_N2'.
    Default is determined by autotuning."""

    PRESCALE_QK: NotRequired[bool]
    """Whether to pre-scale QK by 1/sqrt(d) and change of base. This is slightly faster but
    may have more numerical error. Default: False."""

    ROWS_GUARANTEED_SAFE: NotRequired[bool]
    """If True, guarantees that at least one value in each row is not masked out.
    Allows skipping safety checks for better performance. Only set this if you are certain
    your mask guarantees this property. For example, causal attention is guaranteed safe
    because each query has at least 1 key-value to attend to. Default: False."""

    BLOCKS_ARE_CONTIGUOUS: NotRequired[bool]
    """If True, guarantees that all blocks in the mask are contiguous.
    Allows optimizing block traversal. For example, causal masks would satisfy this,
    but prefix_lm + sliding window would not. Default: False."""

    WRITE_DQ: NotRequired[bool]
    """Controls whether gradient scatters are done in the DQ iteration loop of the backward pass.
    Setting this to False will force this to happen in the DK loop which depending on your
    specific score_mod and mask_mod might be faster. Default: True."""

    FORCE_USE_OMNI_ATTENTION: NotRequired[bool]
    """If True, forces the use of the omni attention kernel instead of potentially using
    the more optimized decoding kernel for short sequences. This can be a helpful
    option for debugging. Default: False."""

    USE_TMA: NotRequired[bool]
    """Whether to use Tensor Memory Accelerator (TMA) on supported hardware.
    This is experimental and may not work on all hardware, currently specific
    to NVIDIA GPUs Hopper+. Default: False."""

    # ROCm-specific options

    kpack: NotRequired[int]
    """ROCm-specific kernel packing parameter."""

    matrix_instr_nonkdim: NotRequired[int]
    """ROCm-specific matrix instruction non-K dimension."""

    waves_per_eu: NotRequired[int]
    """ROCm-specific waves per execution unit."""

    BACKEND: NotRequired[_Backend]
    """Selects a specific kernel backend.

    Options:
        - "AUTO": Use current heuristics (typically tiling-based kernels with
          automatic selection between omni_attention and decoding)
        - "TRITON": Standard omni_attention kernel
        - "TRITON_DECODE": Decoding kernel, only available for short sequence lengths with specific configurations
        - "FLASH": Experimental: Flash Attention kernel (cute-dsl), user needs to have flash installed

    This option cannot be combined with legacy knobs such as ``FORCE_USE_OMNI_ATTENTION``.
    Raises an error if the requested backend cannot be used. Default: "AUTO"
    """

class _KernelOptionsWithInternals(OmniKernelOptions, total=False):
    OUTPUT_LOGSUMEXP: bool
    OUTPUT_MAX: bool

class AuxRequest(NamedTuple):
    """Request which auxiliary outputs to compute from omni_attention.

    Each field is a boolean indicating whether that auxiliary output should be computed.
    """

    lse: bool = False
    max_scores: bool = False

class AuxOutput(NamedTuple):
    """Auxiliary outputs from omni_attention operation.

    Fields will be None if not requested, or contain the tensor if requested.
    """

    lse: Tensor | None = None
    max_scores: Tensor | None = None

class _ModificationType(Enum):
    """Enum for the type of modification function.
    - SCORE_MOD: score_mod function which accepts a score as the first argument
    - mask_mod: mask function which does not accept a score and is only used for generating
    block mask
    """

    SCORE_MOD = 1
    MASK_MOD = 2
    UNKNOWN = 3

def _get_mod_type(
    fn: _score_mod_signature | _mask_mod_signature | Callable[..., Any],
) -> _ModificationType:
    """Get the type of modification function.
    This function inspects the number of positional arguments of the function to determine
    the type of modification function. If the function has 5 positional arguments, it is
    considered as a score_mod function. If the function has 4 positional arguments, it is
    considered as a mask function.
    """
    if hasattr(fn, "__code__"):
        code = fn.__code__
        num_positional_total = code.co_argcount
        defaults = ()
        if hasattr(fn, "__defaults__"):
            defaults = fn.__defaults__ or ()
        num_defaults = len(defaults)
        num_positional_args = num_positional_total - num_defaults
    else:
        num_positional_args = sum(
            1
            for param in inspect.signature(fn).parameters.values()
            if param.default is inspect.Parameter.empty
        )
    if num_positional_args != 5 and num_positional_args != 4:
        raise AssertionError(
            f"Expected 4 or 5 positional args, got {num_positional_args}"
        )
    if num_positional_args == 5:
        return _ModificationType.SCORE_MOD
    elif num_positional_args == 4:
        return _ModificationType.MASK_MOD
    else:
        return _ModificationType.UNKNOWN

# Need to define it here so that the tracer doesn't skip it
def _vmap_for_bhqkv(
    fn: Callable[..., _R],
    prefix: tuple[int | None, ...],
    suffix: tuple[int | None, ...] = (),
    out_dims: int | list[int | None] = 0,
    group_dim: bool = False,
) -> Callable[..., _R]:
    """Used to vmap both score_mods and mask_mods over 4-dimensional/5-dimension inputs.
    Mapping over the [b, hq, q_idx, kv_idx] or [b, hkv, g, q_idx, kv_idx] dimensions.

    Args:
        fn (callable): The function to vmap.
        prefix (tuple): The prefix of the vmap. For score mod functions,
                        this should be set to (0,). For mask_mods = ()
        suffix (tuple): We need to add (0,) if gradOut is being mapped over,
                        and (None,) * len(other_buffers).
        out_dims (tuple): For forward cases, keep this as the default 0 since
                          we are only returning 1 output. For backwards, the joint
                          graph returns grads for B, H, Q_idx, KV_idx and other_buffers,
                          so we set this to (0, None, None, None, None) + (None,) * len(other_buffers).

    Returns:
        callable: The vmapped function.
    """
    # We vmap a function 4 times, broadcasting the [b, h, q_idx, kv_idx] dimensions
    dimensions: list[tuple[int | None, int | None, int | None, int | None]] = []
    dimensions = [
        (None, None, None, 0),
        (None, None, 0, None),
        (None, 0, None, None),
    ]

    if group_dim:
        dimensions += [
            (None, 0, None, None),
        ]

    dimensions += [
        (0, None, None, None),
    ]

    for dims in dimensions:
        fn = tensorplay.func.vmap(fn, in_dims=prefix + dims + suffix, out_dims=out_dims)  # type: ignore[arg-type]
    return fn

def _identity(
    score: Tensor,
    batch: Tensor,
    head: Tensor,
    token_q: Tensor,
    token_kv: Tensor,
) -> Tensor:
    return score

[docs]
def noop_mask(
    batch: Tensor,
    head: Tensor,
    token_q: Tensor,
    token_kv: Tensor,
) -> Tensor:
    """Returns a noop mask_mod"""
    return batch.new_ones(size=(), dtype=tensorplay.bool, device=batch.device)

def _sliced_mask_mod_error(
    batch: Tensor,
    head: Tensor,
    token_q: Tensor,
    token_kv: Tensor,
) -> Never:
    """
    Raises helpful error when using mask_mod from a sliced BlockMask.

    After slicing a BlockMask, the mask_mod is reset and cannot be used directly.
    Users must reassign mask_mod from the original (unsliced) BlockMask.
    """
    raise RuntimeError(
        "Cannot use mask_mod from a sliced BlockMask. "
        "When you slice a BlockMask using [], the mask_mod attribute is reset. "
        "You must set it from the original BlockMask's mask_mod."
        "\n\nIncorrect usage:"
        "\n  base_mask = create_block_mask(my_mask_fn, ...)"
        "\n  sliced_mask = base_mask[:, :, block_idx]"
        "\n  sliced_mask.mask_mod = apply_offset(sliced_mask.mask_mod, offset)  # WRONG!"
        "\n\nCorrect usage:"
        "\n  base_mask = create_block_mask(my_mask_fn, ...)"
        "\n  sliced_mask = base_mask[:, :, block_idx]"
        "\n  sliced_mask.mask_mod = apply_offset(base_mask.mask_mod, offset)  # Use base_mask!"
    )

_DEFAULT_SPARSE_BLOCK_SIZE = 128
_LARGE_SPARSE_BLOCK_SIZE = 1 << 30

def _cdiv(a: int, b: int) -> int:
    return (a + b - 1) // b

def _ordered_to_dense(
    num_blocks_in_row: Tensor, col_indices: Tensor, num_cols: int
) -> Tensor:
    num_rows = col_indices.shape[-2]
    max_entries = col_indices.shape[-1]
    batch_dims = num_blocks_in_row.shape[:-1]
    device = num_blocks_in_row.device

    def create_dense_one(kv_num_blocks, kv_indices):
        dense_mask = kv_indices.new_zeros(num_rows, num_cols + 1, dtype=tensorplay.int32)

        row_indices = tensorplay.arange(num_rows, dtype=tensorplay.int64, device=device).unsqueeze(
            -1
        )
        col_range = tensorplay.arange(max_entries, dtype=tensorplay.int64, device=device)
        index_mask = col_range < kv_num_blocks.unsqueeze(-1)

        # We write to one spot "out of bounds"
        valid_indices = tensorplay.where(index_mask, kv_indices, num_cols)

        # set the values in 'a' to 1 where the indices are valid
        dense_mask[row_indices, valid_indices] = dense_mask.new_ones(())
        return dense_mask[:, :num_cols].contiguous()

    create_dense_batched = create_dense_one
    for _ in range(len(batch_dims)):
        create_dense_batched = tensorplay.func.vmap(create_dense_batched, in_dims=(0, 0))

    out = create_dense_batched(num_blocks_in_row, col_indices)
    return out

def _dense_to_ordered(dense_mask: Tensor) -> tuple[Tensor, Tensor]:
    dense_mask = dense_mask.to(dtype=tensorplay.int32)
    num_blocks_in_row = dense_mask.sum(dim=-1)
    with tensorplay.graph.traceback.annotate({"fallback_to_eager": True}):
        col_indices = tensorplay.argsort(dense_mask, dim=-1, descending=True, stable=True)
    return (
        num_blocks_in_row.to(tensorplay.int32, memory_format=tensorplay.contiguous_format),
        col_indices.to(tensorplay.int32, memory_format=tensorplay.contiguous_format),
    )

def _transpose_ordered(
    num_blocks_in_row: Tensor, col_indices: Tensor, num_cols: int | None = None
) -> tuple[Tensor, Tensor]:
    dense = _ordered_to_dense(
        num_blocks_in_row,
        col_indices,
        col_indices.shape[-1] if num_cols is None else num_cols,
    )
    return _dense_to_ordered(dense.transpose(-2, -1))

def _adjust_num_blocks_and_indices(
    num_blocks: Tensor,
    indices: Tensor,
    new_num_rows: int,
    new_num_cols: int,
) -> tuple[Tensor, Tensor]:
    """Crop an ordered block list while ignoring undefined entries past num_blocks."""
    indices = indices[:, :, :new_num_rows, :new_num_cols]
    num_blocks = num_blocks[:, :, :new_num_rows]
    valid_entries = (
        tensorplay.arange(indices.shape[-1], dtype=num_blocks.dtype, device=indices.device)
        < num_blocks[..., None]
    )
    num_blocks = tensorplay.sum(valid_entries & (indices < new_num_cols), dim=-1).to(
        tensorplay.int32
    )
    return num_blocks, indices

# TreeSpec for an empty tuple — used as the sentinel when there are no closure leaves.
_EMPTY_CLOSURE_SPEC = tree_flatten(())[1]

class _ExtractedLeaf:
    """Sentinel in _StrippedClosure.leaf_entries marking a position that is
    filled from the extracted pytree leaves list during reconstruction."""

    __slots__ = ()

    def __repr__(self) -> str:
        return "_EXTRACTED_LEAF"

_EXTRACTED_LEAF = _ExtractedLeaf()

class _CallableLeaf(typing.NamedTuple):
    """Entry in stripped callable metadata for a recursively processed
    nested callable. Stores enough information to reconstruct it from the
    extracted leaves during unflattening."""

    stripped: _CallableMetadata | Callable[..., Any]
    callable_spec: TreeSpec
    n_extracted: int  # number of extracted leaves this function contributes

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, _CallableLeaf):
            return False
        return (
            _callable_entry_eq(self.stripped, other.stripped)
            and self.callable_spec == other.callable_spec
            and self.n_extracted == other.n_extracted
        )

    def __hash__(self) -> int:
        return hash(
            (
                _callable_entry_hash(self.stripped),
                self.callable_spec,
                self.n_extracted,
            )
        )

class _StrippedClosure(typing.NamedTuple):
    """Data container holding the parts of a function needed for reconstruction.

    Created by _extract_callable_pytree when closure tensors are lifted into
    pytree leaves.  Unlike a FunctionType with None-filled cells, this is not
    callable — it is pure data stored in the pytree context.
    """

    code: types.CodeType
    globals_dict: dict[str, Any]
    name: str
    qualname: str
    defaults: tuple[Any, ...] | None
    kwdefaults: dict[str, Any] | None
    extra_dict: dict[str, Any]
    # Per-position info for the closure's flattened leaves.
    # _EXTRACTED_LEAF → position filled from the extracted leaves list.
    # _CallableLeaf   → recursively processed nested callable.
    leaf_entries: tuple[_ExtractedLeaf | _CallableLeaf, ...]

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, _StrippedClosure):
            return False
        return self.code == other.code and self.leaf_entries == other.leaf_entries

    def __hash__(self) -> int:
        return hash((self.code, self.leaf_entries))

class _StrippedPartial(typing.NamedTuple):
    """Data container holding the parts of a functools.partial needed for reconstruction."""

    leaf_entries: tuple[_ExtractedLeaf | _CallableLeaf, ...]

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, _StrippedPartial):
            return False
        return self.leaf_entries == other.leaf_entries

    def __hash__(self) -> int:
        return hash(self.leaf_entries)

class _PlainFunction(typing.NamedTuple):
    fn: Callable[..., Any]

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, _PlainFunction):
            return False
        return self.fn.__code__ == other.fn.__code__

    def __hash__(self) -> int:
        return hash(self.fn.__code__)

_CallableMetadata: TypeAlias = _StrippedClosure | _StrippedPartial | _PlainFunction

def _callable_entry_eq(
    lhs: _CallableMetadata | types.FunctionType,
    rhs: _CallableMetadata | types.FunctionType,
) -> bool:
    if inspect.isfunction(lhs) and inspect.isfunction(rhs):
        return lhs.__code__ == rhs.__code__
    return lhs == rhs

def _callable_entry_hash(value: _CallableMetadata | types.FunctionType) -> int:
    if inspect.isfunction(value):
        return hash(value.__code__)
    return hash(value)

def _extract_callable_leaves(
    leaves: list[Any], _seen: set[int]
) -> tuple[tuple[BaseArgumentTypes, ...], tuple[_ExtractedLeaf | _CallableLeaf, ...]]:
    extracted: list[BaseArgumentTypes] = []
    leaf_entries: list[_ExtractedLeaf | _CallableLeaf] = []
    for leaf in leaves:
        if inspect.isfunction(leaf) or isinstance(leaf, functools.partial):
            child_extracted, child_spec, child_stripped = _extract_callable_pytree(
                leaf, _seen
            )
            if not isinstance(
                child_stripped, (_StrippedClosure, _StrippedPartial, _PlainFunction)
            ) and not inspect.isfunction(child_stripped):
                raise AssertionError(
                    "expected nested callable extraction to produce callable metadata"
                )
            extracted.extend(child_extracted)
            leaf_entries.append(
                _CallableLeaf(child_stripped, child_spec, len(child_extracted))
            )
        else:
            extracted.append(leaf)
            leaf_entries.append(_EXTRACTED_LEAF)
    return tuple(extracted), tuple(leaf_entries)

def _extract_callable_pytree(
    fn, _seen: set[int] | None = None
) -> tuple[
    tuple[BaseArgumentTypes, ...],
    TreeSpec,
    _CallableMetadata | Callable[..., Any],
]:
    """Extract closure contents as a flattened sub-pytree.

    Returns (extracted_leaves, callable_spec, fn_or_stripped) where:
    - extracted_leaves: flattened non-function contents from the closure,
      functools.partial payload, plus any tensors/scalars recursively extracted
      from nested function closures
    - callable_spec: TreeSpec describing how to reconstruct the callable state
    - fn_or_stripped: either the original fn (skipped extraction) or a
      stripped callable carrying the parts needed for reconstruction

    Functions found among the closure or partial leaves are recursively
    processed: their own closure tensors are extracted into the leaves list,
    and their skeleton is stored in the stripped metadata. All other values
    (tensors, scalars, None, etc.) remain as extracted leaves.

    If fn is not a plain function or functools.partial, returns the original
    function unchanged with no closure leaves. Plain functions without a
    closure are wrapped as _PlainFunction so structural equality/hash can use
    their __code__.

    Skipped under graph tracing (tensorplay.compiler.is_compiling) because the
    tracer can't trace through closure cell introspection and handles freevars
    via its own lifting mechanism.
    """
    if tensorplay.compiler.is_compiling() or not (
        inspect.isfunction(fn) or isinstance(fn, functools.partial)
    ):
        return (), _EMPTY_CLOSURE_SPEC, fn

    # Cycle detection for self-referencing closures.
    if _seen is None:
        _seen = set()
    if id(fn) in _seen:
        return (), _EMPTY_CLOSURE_SPEC, fn
    _seen.add(id(fn))

    if isinstance(fn, functools.partial):
        partial_leaves, partial_spec = tree_flatten((fn.func, fn.args, fn.keywords))
        extracted, leaf_entries = _extract_callable_leaves(partial_leaves, _seen)
        return tuple(extracted), partial_spec, _StrippedPartial(tuple(leaf_entries))

    closure = fn.__closure__
    if not closure:
        return (), _EMPTY_CLOSURE_SPEC, _PlainFunction(fn)

    try:
        contents = tuple(cell.cell_contents for cell in closure)
    except ValueError:
        # Empty cell (created but not yet assigned) — can't extract
        return (), _EMPTY_CLOSURE_SPEC, _PlainFunction(fn)

    closure_leaves, callable_spec = tree_flatten(contents)

    extracted, leaf_entries = _extract_callable_leaves(closure_leaves, _seen)

    stripped = _StrippedClosure(
        code=fn.__code__,
        globals_dict=fn.__globals__,
        name=fn.__name__,
        qualname=fn.__qualname__,
        defaults=fn.__defaults__,
        kwdefaults=fn.__kwdefaults__,
        extra_dict=dict(fn.__dict__) if fn.__dict__ else {},
        leaf_entries=tuple(leaf_entries),
    )

    return tuple(extracted), callable_spec, stripped

def _reconstruct_closure_fn(stripped, extracted_leaves, callable_spec):
    """Rebuild a stripped callable from flattened extracted leaves."""
    if isinstance(stripped, _PlainFunction):
        return stripped.fn
    if not isinstance(stripped, (_StrippedClosure, _StrippedPartial)):
        return stripped

    all_leaves: list[BaseArgumentTypes | Callable[..., Any]] = []
    idx = 0
    for entry in stripped.leaf_entries:
        if isinstance(entry, _CallableLeaf):
            child_fn = _reconstruct_closure_fn(
                entry.stripped,
                extracted_leaves[idx : idx + entry.n_extracted],
                entry.callable_spec,
            )
            all_leaves.append(child_fn)
            idx += entry.n_extracted
        else:
            # _EXTRACTED_LEAF — take from extracted leaves
            all_leaves.append(extracted_leaves[idx])
            idx += 1

    if isinstance(stripped, _StrippedPartial):
        fn, args, keywords = tree_unflatten(all_leaves, callable_spec)
        return functools.partial(fn, *args, **keywords)

    contents = tree_unflatten(all_leaves, callable_spec)
    new_cells = tuple(types.CellType(v) for v in contents)

    restored = types.FunctionType(
        stripped.code,
        stripped.globals_dict,
        stripped.name,
        stripped.defaults,
        new_cells,
    )
    restored.__qualname__ = stripped.qualname
    if stripped.kwdefaults:
        restored.__kwdefaults__ = stripped.kwdefaults
    if stripped.extra_dict:
        restored.__dict__.update(stripped.extra_dict)

    return restored

class _MaskModWrapper:
    """Wraps a mask_mod or stripped callable metadata with value-based equality.

    BlockMask stores an arbitrary callable (mask_mod) in its pytree context.
    The default __eq__ for functions uses identity comparison, which is too
    strict when the same closure is recreated (e.g., defined inside forward()).

    When callable state has been extracted (by _extract_callable_pytree), fn is
    stripped metadata (pure data, not callable). Equality compares the stripped
    callable structure + callable_spec without triggering tensor dispatch.

    When extraction is skipped (e.g., under tracing), fn is the original
    callable and equality compares code objects for plain functions or
    delegates to __eq__ for callable objects.
    """

    __slots__ = ("fn", "callable_spec")

    def __init__(self, fn, callable_spec=None) -> None:
        self.fn = fn
        self.callable_spec = callable_spec

    def __call__(self, b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
        if isinstance(self.fn, (_StrippedClosure, _StrippedPartial, _PlainFunction)):
            raise RuntimeError(
                "_MaskModWrapper with stripped callable is not callable — "
                "use _reconstruct_closure_fn to rebuild the function first"
            )
        return self.fn(b, h, q_idx, kv_idx)

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, _MaskModWrapper):
            return False
        if self.fn is other.fn and self.callable_spec is other.callable_spec:
            return True
        if isinstance(
            self.fn, (_StrippedClosure, _StrippedPartial, _PlainFunction)
        ) and isinstance(
            other.fn, (_StrippedClosure, _StrippedPartial, _PlainFunction)
        ):
            return self.fn == other.fn and self.callable_spec == other.callable_spec
        # Non-extracted plain functions: compare code objects
        if inspect.isfunction(self.fn) and inspect.isfunction(other.fn):
            return self.fn.__code__ == other.fn.__code__
        # Callable objects: delegate to their __eq__
        if not isinstance(
            self.fn, (_StrippedClosure, _StrippedPartial, _PlainFunction)
        ) and not isinstance(
            other.fn, (_StrippedClosure, _StrippedPartial, _PlainFunction)
        ):
            return self.fn == other.fn
        return False

    def __hash__(self) -> int:
        if isinstance(self.fn, (_StrippedClosure, _StrippedPartial, _PlainFunction)):
            return hash((self.fn, self.callable_spec))
        if inspect.isfunction(self.fn):
            return hash(self.fn.__code__)
        return hash(self.fn)

    def __repr__(self) -> str:
        return f"_MaskModWrapper({self.fn})"

[docs]
class BlockMask:
    r"""
    BlockMask is our format for representing a block-sparse attention mask.
    It is somewhat of a cross in-between BCSR and a non-sparse format.

    **Basics**

    A block-sparse mask means that instead of representing the sparsity of
    individual elements in the mask, a KV_BLOCK_SIZE x Q_BLOCK_SIZE block is
    considered sparse only if every element within that block is sparse.
    This aligns well with hardware, which generally expects to perform
    contiguous loads and computation.

    This format is primarily optimized for 1. simplicity, and 2. kernel
    efficiency. Notably, it is *not* optimized for size, as this mask is always
    reduced by a factor of KV_BLOCK_SIZE * Q_BLOCK_SIZE. If the size is a
    concern, the tensors can be reduced in size by increasing the block size.

    The essentials of our format are:

    num_blocks_in_row: Tensor[ROWS]:
    Describes the number of blocks present in each row.

    col_indices: Tensor[ROWS, MAX_BLOCKS_IN_COL]:
    `col_indices[i]` is the sequence of block positions for row i. The values of
    this row after `col_indices[i][num_blocks_in_row[i]]` are undefined.

    For example, to reconstruct the original tensor from this format:

    .. code-block:: python

        dense_mask = tensorplay.zeros(ROWS, COLS)
        for row in range(ROWS):
            for block_idx in range(num_blocks_in_row[row]):
                dense_mask[row, col_indices[row, block_idx]] = 1

    Notably, this format makes it easier to implement a reduction along the
    *rows* of the mask.

    **Details**

    The basics of our format require only kv_num_blocks and kv_indices. The
    primary block-sparse layout is represented by up to 4 tensor pairs:

    1. (kv_num_blocks, kv_indices): Used for the forwards pass of attention, as
    we reduce along the KV dimension.

    2. [OPTIONAL] (full_kv_num_blocks, full_kv_indices): This is optional and
    purely an optimization. As it turns out, applying masking to every block
    is quite expensive! If we specifically know which blocks are "full" and
    don't require masking at all, then we can skip applying mask_mod to these
    blocks. This requires the user to split out a separate mask_mod from the
    score_mod. For causal masks, this is about a 15% speedup.

    3. [GENERATED] (q_num_blocks, q_indices): Required for the backwards pass,
    as computing dKV requires iterating along the mask along the Q dimension. These are autogenerated from 1.

    4. [GENERATED] (full_q_num_blocks, full_q_indices): Same as above, but for
    the backwards pass. These are autogenerated from 2.

    Additional optional tensors may carry deterministic dQ metadata for
    block-sparse backward:

    5. [OPTIONAL] dq_write_order: Write-order metadata for partial blocks. This
    is produced by create_block_mask when compute_dq_write_order=True, or passed
    directly to BlockMask.from_kv_blocks by callers that precompute it.

    6. [OPTIONAL] dq_write_order_full: Write-order metadata for full blocks,
    produced or passed the same way as dq_write_order.

    7. [OPTIONAL] dq_kv_order: Explicit KV scheduler order used to produce the
    write-order metadata. create_block_mask currently accepts a boolean
    dq_kv_order; BlockMask.from_kv_blocks also accepts a tensor for callers that
    provide precomputed write-order metadata directly.
    """

    seq_lengths: tuple[int, int]
    kv_num_blocks: Tensor
    kv_indices: Tensor
    full_kv_num_blocks: Tensor | None
    full_kv_indices: Tensor | None
    q_num_blocks: Tensor | None
    q_indices: Tensor | None
    full_q_num_blocks: Tensor | None
    full_q_indices: Tensor | None
    dq_write_order: Tensor | None
    dq_write_order_full: Tensor | None
    dq_kv_order: Tensor | None
    dq_kv_order_spt: bool | None
    BLOCK_SIZE: tuple[int, int]
    mask_mod: _mask_mod_signature

    # Attribute lists for pytree flatten/unflatten
    _TENSOR_ATTRS = [
        "kv_num_blocks",
        "kv_indices",
        "full_kv_num_blocks",
        "full_kv_indices",
        "q_num_blocks",
        "q_indices",
        "full_q_num_blocks",
        "full_q_indices",
        "dq_write_order",
        "dq_write_order_full",
        "dq_kv_order",
    ]

    _CONTEXT_ATTRS = [
        "seq_lengths",
        "BLOCK_SIZE",
        "mask_mod",
        "dq_kv_order_spt",
    ]

    def __init__(
        self,
        seq_lengths: tuple[int, int],
        kv_num_blocks: Tensor,
        kv_indices: Tensor,
        full_kv_num_blocks: Tensor | None,
        full_kv_indices: Tensor | None,
        q_num_blocks: Tensor | None,
        q_indices: Tensor | None,
        full_q_num_blocks: Tensor | None,
        full_q_indices: Tensor | None,
        BLOCK_SIZE: tuple[int, int] = (
            _DEFAULT_SPARSE_BLOCK_SIZE,
            _DEFAULT_SPARSE_BLOCK_SIZE,
        ),
        mask_mod: _mask_mod_signature = noop_mask,
        *,
        dq_write_order: Tensor | None = None,
        dq_write_order_full: Tensor | None = None,
        dq_kv_order: Tensor | None = None,
        dq_kv_order_spt: bool | None = None,
    ) -> None:
        if kv_indices.dim() < 2:
            raise RuntimeError("BlockMask must have at least 2 dimensions")
        if kv_num_blocks is None:
            raise AssertionError("kv_num_blocks must be provided")
        if kv_indices is None:
            raise AssertionError("kv_indices must be provided")
        if (full_kv_num_blocks is None) != (full_kv_indices is None):
            raise AssertionError(
                "full_kv_num_blocks and full_kv_indices must be both provided or omitted"
            )
        if (full_q_num_blocks is None) != (full_q_indices is None):
            raise AssertionError(
                "full_q_num_blocks and full_q_indices must be both provided or omitted"
            )

        self.seq_lengths = seq_lengths
        self.kv_num_blocks = kv_num_blocks
        self.kv_indices = kv_indices
        self.full_kv_num_blocks = full_kv_num_blocks
        self.full_kv_indices = full_kv_indices
        self.q_num_blocks = q_num_blocks
        self.q_indices = q_indices
        self.full_q_num_blocks = full_q_num_blocks
        self.full_q_indices = full_q_indices
        if dq_write_order_full is not None and dq_write_order is None:
            raise ValueError("dq_write_order_full requires dq_write_order")

        self.dq_write_order = dq_write_order
        self.dq_write_order_full = dq_write_order_full
        self.dq_kv_order = dq_kv_order
        self.dq_kv_order_spt = dq_kv_order_spt
        self.BLOCK_SIZE = BLOCK_SIZE
        self.mask_mod = mask_mod

    def _dq_kv_order(self) -> Tensor | bool | None:
        if self.dq_kv_order is not None:
            return self.dq_kv_order
        return self.dq_kv_order_spt

    @staticmethod
    def _slice_dq_kv_order(
        dq_kv_order: Tensor | None,
        index: tuple[int | slice | Tensor, ...],
    ) -> Tensor | None:
        if dq_kv_order is None:
            return None
        return dq_kv_order[index[:2] + (slice(None),)]

    def _query_length_for_sliced_blocks(
        self, q_index: int | slice | Tensor, selected_q_blocks: int
    ) -> int:
        """Maps a Q-block selection to its packed logical token length."""
        q_block_size = self.BLOCK_SIZE[0]
        if not isinstance(q_index, slice):
            return selected_q_blocks * q_block_size

        selected_blocks = range(self.kv_num_blocks.shape[-1])[q_index]
        if len(selected_blocks) == 0:
            return 0
        if selected_blocks.step == 1:
            return min(selected_blocks.stop * q_block_size, self.seq_lengths[0]) - min(
                selected_blocks.start * q_block_size, self.seq_lengths[0]
            )

        q_length = len(selected_blocks) * q_block_size
        trailing_q_tokens = self.seq_lengths[0] % q_block_size
        if trailing_q_tokens and self.kv_num_blocks.shape[-1] - 1 in selected_blocks:
            q_length -= q_block_size - trailing_q_tokens
        return q_length

[docs]
    @classmethod
    def from_kv_blocks(
        cls,
        kv_num_blocks: Tensor,
        kv_indices: Tensor,
        full_kv_num_blocks: Tensor | None = None,
        full_kv_indices: Tensor | None = None,
        BLOCK_SIZE: int | tuple[int, int] = _DEFAULT_SPARSE_BLOCK_SIZE,
        mask_mod: _mask_mod_signature | None = None,
        seq_lengths: tuple[int, int] | None = None,
        compute_q_blocks: bool = True,
        *,
        dq_write_order: Tensor | None = None,
        dq_write_order_full: Tensor | None = None,
        dq_kv_order: Tensor | bool | None = None,
    ) -> Self:
        """
        Creates a BlockMask instance from key-value block information.

        Args:
            kv_num_blocks (Tensor): Number of kv_blocks in each Q_BLOCK_SIZE row tile.
            kv_indices (Tensor): Indices of key-value blocks in each Q_BLOCK_SIZE row tile.
            full_kv_num_blocks (Optional[Tensor]): Number of full kv_blocks in each Q_BLOCK_SIZE row tile.
            full_kv_indices (Optional[Tensor]): Indices of full key-value blocks in each Q_BLOCK_SIZE row tile.
            BLOCK_SIZE (Union[int, tuple[int, int]]): Size of KV_BLOCK_SIZE x Q_BLOCK_SIZE tiles.
            mask_mod (Optional[Callable]): Function to modify the mask.
            dq_write_order (Optional[Tensor]): Precomputed deterministic dQ write-order metadata.
            dq_write_order_full (Optional[Tensor]): Precomputed deterministic dQ write-order metadata for full blocks.
            dq_kv_order (Optional[Union[Tensor, bool]]): KV-column scheduler order used to produce dq_write_order. A bool selects a built-in order; a tensor gives an explicit scheduler-rank to n-block permutation.

        Returns:
            BlockMask: Instance with full Q information generated via _transposed_ordered

        Raises:
            RuntimeError: If kv_indices has < 2 dimensions.
            AssertionError: If only one of full_kv_* args is provided.
        """
        if kv_indices.dim() < 2:
            raise RuntimeError("BlockMask must have at least 2 dimensions")

        if (full_kv_num_blocks is None) != (full_kv_indices is None):
            raise AssertionError(
                "full_kv_num_blocks and full_kv_indices must be both provided or omitted"
            )

        if isinstance(BLOCK_SIZE, int):
            BLOCK_SIZE = (BLOCK_SIZE, BLOCK_SIZE)

        mask_mod = mask_mod if mask_mod is not None else noop_mask
        if seq_lengths is None:
            q_length = kv_indices.shape[-2] * BLOCK_SIZE[0]
            kv_length = kv_indices.shape[-1] * BLOCK_SIZE[1]
            seq_lengths = (q_length, kv_length)
        kv_num_cols = (seq_lengths[1] + BLOCK_SIZE[1] - 1) // BLOCK_SIZE[1]

        # Generate q_num_blocks and q_indices
        if compute_q_blocks:
            q_num_blocks, q_indices = _transpose_ordered(
                kv_num_blocks, kv_indices, kv_num_cols
            )
            if full_kv_num_blocks is not None:
                if full_kv_indices is None:
                    raise AssertionError("full_kv_indices must not be None")
                full_q_num_blocks, full_q_indices = _transpose_ordered(
                    full_kv_num_blocks, full_kv_indices, kv_num_cols
                )
            else:
                full_q_num_blocks, full_q_indices = None, None
        else:
            q_num_blocks, q_indices = None, None
            full_q_num_blocks, full_q_indices = None, None

        if dq_kv_order is not None and not isinstance(dq_kv_order, (bool, Tensor)):
            raise ValueError("dq_kv_order must be a bool, Tensor, or None")
        dq_kv_order_tensor = dq_kv_order if isinstance(dq_kv_order, Tensor) else None
        dq_kv_order_spt = dq_kv_order if isinstance(dq_kv_order, bool) else None

        return cls(
            seq_lengths=seq_lengths,
            kv_num_blocks=kv_num_blocks,
            kv_indices=kv_indices,
            full_kv_num_blocks=full_kv_num_blocks,
            full_kv_indices=full_kv_indices,
            q_num_blocks=q_num_blocks,
            q_indices=q_indices,
            full_q_num_blocks=full_q_num_blocks,
            full_q_indices=full_q_indices,
            BLOCK_SIZE=BLOCK_SIZE,
            mask_mod=mask_mod,
            dq_write_order=dq_write_order,
            dq_write_order_full=dq_write_order_full,
            dq_kv_order=dq_kv_order_tensor,
            dq_kv_order_spt=dq_kv_order_spt,
        )

    @overload
    def as_tuple(
        self, flatten: Literal[True] = ...
    ) -> tuple[
        int,
        int,
        Tensor,
        Tensor,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        bool | None,
        int,
        int,
        _mask_mod_signature,
    ]: ...

    @overload
    def as_tuple(
        self, flatten: Literal[False]
    ) -> tuple[
        tuple[int, int],
        Tensor,
        Tensor,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        Tensor | None,
        bool | None,
        tuple[int, int],
        _mask_mod_signature,
    ]: ...

[docs]
    def as_tuple(self, flatten: bool = True) -> tuple[Any, ...]:
        """
        Returns a tuple of the attributes of the BlockMask.

        Args:
            flatten (bool): If True, it will flatten the tuple of (KV_BLOCK_SIZE, Q_BLOCK_SIZE)
        """
        if flatten:
            block_size = (self.BLOCK_SIZE[0], self.BLOCK_SIZE[1])  # type: ignore[assignment]
            seq_lengths = (self.seq_lengths[0], self.seq_lengths[1])  # type: ignore[assignment]
        else:
            block_size = (self.BLOCK_SIZE,)  # type: ignore[assignment]
            seq_lengths = (self.seq_lengths,)  # type: ignore[assignment]

        # pyrefly: ignore [not-iterable]
        return (
            *seq_lengths,
            self.kv_num_blocks,
            self.kv_indices,
            self.full_kv_num_blocks,
            self.full_kv_indices,
            self.q_num_blocks,
            self.q_indices,
            self.full_q_num_blocks,
            self.full_q_indices,
            self.dq_write_order,
            self.dq_write_order_full,
            self.dq_kv_order,
            self.dq_kv_order_spt,
            *block_size,
            self.mask_mod,
        )

    @property
    def shape(self) -> tuple[int, ...]:
        *batch_dims, _, _ = self.kv_indices.shape
        return tuple(batch_dims) + self.seq_lengths

    def __str__(self) -> str:
        s = f"BlockMask(shape={self.shape}, sparsity={self.sparsity():.2f}%, \n"
        mask_str = self.to_string().strip()
        s += mask_str
        s += "\n)"
        return s

    def __getitem__(
        self, index: int | slice | Tensor | tuple[int | slice | Tensor, ...]
    ) -> Self:
        """
        Returns a new BlockMask by selecting batch, head, and Q-block rows.

        BlockMask indexing accepts up to three indices. The first two select batch
        and head metadata. The third selects rows in the block-sparse Q grid, not
        individual query tokens. For example, ``block_mask[:, :, 1]`` selects one
        Q-block row, covering the second chunk of ``Q_BLOCK_SIZE`` query tokens,
        not query token 1. Integer indices are normalized to length-one slices so
        the corresponding BlockMask dimension is preserved.

        The returned BlockMask is packed: selecting Q block row ``i`` produces
        local Q block row 0 in the result. Its ``shape[-2]`` is the packed logical
        query-token length of the selected Q-block rows, clipped for partial final
        blocks when the selection is a Python integer or slice. Tensor indexing is
        handled without reading index values, so its logical query length is the
        number of selected Q-block rows times ``Q_BLOCK_SIZE``. Use a Python
        integer or slice when selecting a partial final block and an exact logical
        length is required.

        Callers must slice or gather the query tensor separately and install a new
        ``mask_mod`` if the mask depends on absolute query positions. KV block
        columns are not sliced by this API.

        Args:
            index: Batch, head, and Q-block-row index.

        Returns:
            A BlockMask containing metadata for the selected batch, head, and
            packed Q-block rows.

        Raises:
            IndexError: If more than three indices are provided.
            NotImplementedError: If tensor dq_kv_order metadata is present.

        Example:
            .. code-block:: python

                def causal_mask(b, h, q_idx, kv_idx):
                    return q_idx >= kv_idx

                block_mask = create_block_mask(
                    causal_mask, 4, 2, 512, 512, device="cuda"
                )
                assert block_mask.kv_num_blocks.shape == (4, 2, 4)
                assert block_mask.kv_indices.shape == (4, 2, 4, 4)

                batch_slice = block_mask[0]
                assert batch_slice.kv_num_blocks.shape == (1, 2, 4)
                assert batch_slice.kv_indices.shape == (1, 2, 4, 4)

                head_slice = block_mask[0, 1]
                assert head_slice.kv_num_blocks.shape == (1, 1, 4)
                assert head_slice.kv_indices.shape == (1, 1, 4, 4)

                query_block_slice = block_mask[:, :, :1]
                assert query_block_slice.shape == (4, 2, 128, 512)
                assert query_block_slice.kv_num_blocks.shape == (4, 2, 1)
                assert query_block_slice.kv_indices.shape == (4, 2, 1, 4)

                same_block_slice = block_mask[:, :, 0]
                assert same_block_slice.shape == (4, 2, 128, 512)

                q = tensorplay.randn(4, 2, 512, 64, device="cuda")
                q_chunk = q[:, :, :128, :]
                assert q_chunk.shape[-2] == query_block_slice.shape[-2]
        """
        if self.dq_kv_order is not None:
            raise NotImplementedError(
                "Slicing BlockMask with tensor dq_kv_order is not supported yet. "
                "Construct a new BlockMask with matching precomputed dq_write_order metadata."
            )

        index = (index,) if not isinstance(index, tuple) else index
        if len(index) > 3:
            raise IndexError(
                "BlockMask indexing supports batch, head, and Q-block dimensions"
            )
        padded = (*index, slice(None), slice(None), slice(None))[:3]
        sizes = self.kv_num_blocks.shape[:3]
        index = tuple(
            (slice(i + n, i + n + 1) if -n <= i < 0 else slice(i, i + 1))
            if isinstance(i, int)
            else i
            for i, n in zip(padded, sizes, strict=True)
        )
        new_kv_num_blocks = self.kv_num_blocks[index]
        new_kv_indices = self.kv_indices[index]
        if self.full_kv_num_blocks is not None:
            if self.full_kv_indices is None:
                raise AssertionError("full_kv_indices must not be None")
            new_full_kv_num_blocks = self.full_kv_num_blocks[index]
            new_full_kv_indices = self.full_kv_indices[index]
        else:
            new_full_kv_num_blocks = None
            new_full_kv_indices = None
        new_block_mask = BlockMask.from_kv_blocks(
            new_kv_num_blocks,
            new_kv_indices,
            new_full_kv_num_blocks,
            new_full_kv_indices,
            BLOCK_SIZE=self.BLOCK_SIZE,
            mask_mod=_sliced_mask_mod_error,
            seq_lengths=(
                self._query_length_for_sliced_blocks(
                    index[2], new_kv_indices.shape[-2]
                ),
                self.seq_lengths[1],
            ),
            compute_q_blocks=self.q_indices is not None,
        )
        new_block_mask.dq_kv_order = self._slice_dq_kv_order(self.dq_kv_order, index)
        new_block_mask.dq_kv_order_spt = self.dq_kv_order_spt
        if self.dq_write_order is not None and new_block_mask.dq_kv_order is None:
            dq_wo, dq_wo_full = _compute_dq_write_order_from_block_mask(
                new_block_mask, dq_kv_order=new_block_mask.dq_kv_order_spt
            )
            new_block_mask.dq_write_order = dq_wo
            new_block_mask.dq_write_order_full = dq_wo_full
        return new_block_mask

    def __repr__(self) -> str:
        def shape_or_none(x: Tensor | None):
            return x.shape if x is not None else None

        return (
            f"BlockMask(\n"
            f"    kv_num_blocks={self.kv_num_blocks.shape},\n"
            f"    kv_indices={self.kv_indices.shape},\n"
            f"    full_kv_num_blocks={shape_or_none(self.full_kv_num_blocks)},\n"
            f"    full_kv_indices={shape_or_none(self.full_kv_indices)},\n"
            f"    q_num_blocks={shape_or_none(self.q_num_blocks)},\n"
            f"    q_indices={shape_or_none(self.q_indices)},\n"
            f"    full_q_num_blocks={shape_or_none(self.full_q_num_blocks)},\n"
            f"    full_q_indices={shape_or_none(self.full_q_indices)},\n"
            f"    BLOCK_SIZE={self.BLOCK_SIZE},\n"
            f"    shape={self.shape},\n"
            f"    sparsity={self.sparsity():.2f}%,\n"
            f"    mask_mod={self.mask_mod.__name__ if hasattr(self.mask_mod, '__name__') else self.mask_mod}\n"
            f")"
        )

    def _adjust(self, new_q_len: int, new_kv_len: int) -> Self:
        if self.dq_kv_order is not None:
            raise NotImplementedError(
                "Adjusting BlockMask with tensor dq_kv_order is not supported yet. "
                "Construct a new BlockMask with matching precomputed dq_write_order metadata."
            )

        new_num_rows = (new_q_len + self.BLOCK_SIZE[0] - 1) // self.BLOCK_SIZE[0]
        new_num_cols = (new_kv_len + self.BLOCK_SIZE[1] - 1) // self.BLOCK_SIZE[1]
        new_kv_num_blocks, new_kv_indices = _adjust_num_blocks_and_indices(
            self.kv_num_blocks, self.kv_indices, new_num_rows, new_num_cols
        )
        if self.full_kv_num_blocks is not None:
            if self.full_kv_indices is None:
                raise AssertionError("full_kv_indices must not be None")
            (
                new_full_kv_num_blocks,
                new_full_kv_indices,
            ) = _adjust_num_blocks_and_indices(
                self.full_kv_num_blocks,
                self.full_kv_indices,
                new_num_rows,
                new_num_cols,
            )
        else:
            new_full_kv_num_blocks = None
            new_full_kv_indices = None
        new_block_mask = self.from_kv_blocks(
            new_kv_num_blocks,
            new_kv_indices,
            new_full_kv_num_blocks,
            new_full_kv_indices,
            BLOCK_SIZE=self.BLOCK_SIZE,
            mask_mod=self.mask_mod,
            seq_lengths=(new_q_len, new_kv_len),
            compute_q_blocks=self.q_indices is not None,
        )
        new_block_mask.dq_kv_order = self.dq_kv_order
        new_block_mask.dq_kv_order_spt = self.dq_kv_order_spt
        if self.dq_write_order is not None and new_block_mask.dq_kv_order is None:
            dq_wo, dq_wo_full = _compute_dq_write_order_from_block_mask(
                new_block_mask, dq_kv_order=new_block_mask.dq_kv_order_spt
            )
            new_block_mask.dq_write_order = dq_wo
            new_block_mask.dq_write_order_full = dq_wo_full
        return new_block_mask

[docs]
    def numel(self) -> int:
        """Returns the number of elements (not accounting for sparsity) in the mask."""
        shape = self.shape

        def _prod(xs):
            return functools.reduce(operator.mul, xs, 1)

        return _prod(shape)

[docs]
    def sparsity(self) -> float:
        """Computes the percentage of blocks that are sparse (i.e. not computed)"""
        total_blocks = math.prod(self.kv_num_blocks.shape[:-1])
        for seq_len, block_size in zip(self.seq_lengths, self.BLOCK_SIZE, strict=True):
            total_blocks *= _cdiv(seq_len, block_size)

        computed_blocks = self.kv_num_blocks.sum()
        if self.full_kv_num_blocks is not None:
            computed_blocks += self.full_kv_num_blocks.sum()

        dense_ratio = computed_blocks.item() / total_blocks
        return 100 * (1 - dense_ratio)

[docs]
    def to_dense(self) -> Tensor:
        """Returns a dense block that is equivalent to the block mask."""
        partial_dense = _ordered_to_dense(
            self.kv_num_blocks, self.kv_indices, self.kv_indices.shape[-1]
        )
        if self.full_kv_num_blocks is not None:
            if self.full_kv_indices is None:
                raise AssertionError("full_kv_indices must not be None")
            # pyrefly: ignore [bad-return]
            return partial_dense | _ordered_to_dense(
                self.full_kv_num_blocks,
                self.full_kv_indices,
                self.full_kv_indices.shape[-1],
            )
        return partial_dense

[docs]
    def to_string(
        self, grid_size: int | tuple[int, int] = (20, 20), limit: int = 4
    ) -> str:
        """Returns a string representation of the block mask. Quite nifty.

        If grid_size is -1, prints out an uncompressed version. Warning, it can be quite big!
        """
        dense_mask = self.to_dense()
        *batch_dims, num_rows, num_cols = dense_mask.shape
        if isinstance(grid_size, int):
            max_rows = grid_size
            max_cols = grid_size
        elif grid_size == -1:
            max_rows = num_rows
            max_cols = num_cols
        else:
            max_rows, max_cols = grid_size

        def create_block_vis(*batch_idx):
            descriptors = []

            descriptors.append(f"{batch_idx}")

            vis = ", ".join(reversed(descriptors)) + "\n"

            def summarize_section(section) -> str:
                percentage = section.float().mean().item()
                if percentage == 1:
                    return "█"
                elif percentage == 0:
                    return " "
                else:
                    return "░"

            row_step = max(1, _cdiv(num_rows, max_rows))
            col_step = max(1, _cdiv(num_cols, max_cols))

            for r in range(0, num_rows, row_step):
                for c in range(0, num_cols, col_step):
                    cur_mask = dense_mask
                    for idx in batch_idx:
                        cur_mask = cur_mask[idx]
                    char = summarize_section(
                        cur_mask[r : r + row_step, c : c + col_step]
                    )
                    vis += char * 2
                vis += "\n"
            return vis

        total_vis = []
        for idx, batch_idx in enumerate(
            itertools.product(*[range(i) for i in batch_dims])
        ):
            if idx == limit:
                total_vis.append("...")
                total_vis.append("To print out more, set BlockMask.to_string(limit=N)")
                total_vis.append(
                    "You can also index (BlockMask[batch, head]) to choose a specific batch or head"
                )
                break
            block_vis = create_block_vis(*batch_idx)
            total_vis.append(block_vis)

        return "\n".join(total_vis)

[docs]
    def to(self, device: tensorplay.device | str) -> BlockMask:
        """Moves the BlockMask to the specified device.

        Args:
            device (tensorplay.device or str): The target device to move the BlockMask to.
                Can be a device object or a string (e.g., 'cpu', 'cuda:0').

        Returns:
            BlockMask: A new BlockMask instance with all tensor components moved
            to the specified device.

        Note:
            This method does not modify the original BlockMask in-place.
            Instead, it returns a new BlockMask instance where individual tensor attributes
            may or may not be moved to the specified device, depending on their
            current device placement.
        """
        mapped_tensors = tree_map_only(
            Tensor,
            lambda x: x.to(device),
            tuple(getattr(self, attr) for attr in self._TENSOR_ATTRS),
        )
        return BlockMask(
            seq_lengths=self.seq_lengths,
            **dict(zip(self._TENSOR_ATTRS, mapped_tensors, strict=True)),
            BLOCK_SIZE=self.BLOCK_SIZE,
            mask_mod=self.mask_mod,
            dq_kv_order_spt=self.dq_kv_order_spt,
        )

    @staticmethod
    def _wrap_context_value(attr: str, value: Any) -> Any:
        if attr == "mask_mod":
            return _MaskModWrapper(value)
        return value

    @staticmethod
    def _unwrap_context_value(attr: str, value: Any) -> Any:
        if attr == "mask_mod":
            if not isinstance(value, _MaskModWrapper):
                raise AssertionError(f"Expected _MaskModWrapper, got {type(value)}")
            return value.fn
        return value

    def _flatten(
        self,
    ) -> tuple[tuple[BaseArgumentTypes | None, ...], tuple[Any, ...]]:
        """Flatten BlockMask into a list of tensors and context.

        Closure tensors from mask_mod are extracted into the leaves via
        _extract_callable_pytree so they are visible to the tracing
        infrastructure (instead of being hidden in the pytree context).
        """
        optional_tensor_attrs = tuple(
            attr for attr in self._TENSOR_ATTRS if getattr(self, attr) is None
        )
        tensors = tuple(
            getattr(self, attr)
            for attr in self._TENSOR_ATTRS
            if attr not in optional_tensor_attrs
        )
        closure_leaves, callable_spec, stripped = _extract_callable_pytree(
            self.mask_mod
        )
        all_leaves = tensors + closure_leaves
        context = (
            *(
                self._wrap_context_value(attr, getattr(self, attr))
                if attr != "mask_mod"
                else _MaskModWrapper(stripped, callable_spec)
                for attr in self._CONTEXT_ATTRS
            ),
            optional_tensor_attrs,
        )
        return all_leaves, context

    @classmethod
    def _unflatten(
        cls,
        leaves: tuple[Any, ...],
        context: tuple[Any, ...],
    ) -> Self:
        """Unflatten leaves and context back into a BlockMask."""
        optional_tensor_attrs = context[-1]
        tensor_attrs = tuple(
            attr for attr in cls._TENSOR_ATTRS if attr not in optional_tensor_attrs
        )
        n_regular = len(tensor_attrs)
        regular_leaves = leaves[:n_regular]
        closure_leaves = leaves[n_regular:]
        tensor_values = dict.fromkeys(optional_tensor_attrs)
        tensor_values.update(zip(tensor_attrs, regular_leaves))
        context_values = {}
        for attr, val in zip(cls._CONTEXT_ATTRS, context[:-1]):
            if attr == "mask_mod" and isinstance(val, _MaskModWrapper):
                context_values[attr] = _reconstruct_closure_fn(
                    val.fn, closure_leaves, val.callable_spec
                )
            else:
                context_values[attr] = cls._unwrap_context_value(attr, val)
        return cls(
            seq_lengths=cast(tuple[int, int], context_values["seq_lengths"]),
            kv_num_blocks=cast(Tensor, tensor_values["kv_num_blocks"]),
            kv_indices=cast(Tensor, tensor_values["kv_indices"]),
            full_kv_num_blocks=cast(Tensor | None, tensor_values["full_kv_num_blocks"]),
            full_kv_indices=cast(Tensor | None, tensor_values["full_kv_indices"]),
            q_num_blocks=cast(Tensor | None, tensor_values["q_num_blocks"]),
            q_indices=cast(Tensor | None, tensor_values["q_indices"]),
            full_q_num_blocks=cast(Tensor | None, tensor_values["full_q_num_blocks"]),
            full_q_indices=cast(Tensor | None, tensor_values["full_q_indices"]),
            BLOCK_SIZE=cast(tuple[int, int], context_values["BLOCK_SIZE"]),
            mask_mod=cast(_mask_mod_signature, context_values["mask_mod"]),
            dq_write_order=cast(Tensor | None, tensor_values["dq_write_order"]),
            dq_write_order_full=cast(
                Tensor | None, tensor_values["dq_write_order_full"]
            ),
            dq_kv_order=cast(Tensor | None, tensor_values["dq_kv_order"]),
            dq_kv_order_spt=cast(bool | None, context_values["dq_kv_order_spt"]),
        )

    def _flatten_with_keys(
        self,
    ) -> tuple[tuple[tuple[GetAttrKey, Any], ...], tuple[Any, ...]]:
        """Flatten BlockMask with keys for better tracing.

        Closure tensors from mask_mod are extracted into the leaves via
        _extract_callable_pytree so they are visible to the tracing
        infrastructure (instead of being hidden in the pytree context).
        """
        optional_tensor_attrs = tuple(
            attr for attr in self._TENSOR_ATTRS if getattr(self, attr) is None
        )
        tensors = tuple(
            (GetAttrKey(attr), getattr(self, attr))
            for attr in self._TENSOR_ATTRS
            if attr not in optional_tensor_attrs
        )
        closure_leaves, callable_spec, stripped = _extract_callable_pytree(
            self.mask_mod
        )
        closure_with_keys = tuple(
            (GetAttrKey(f"_closure_{i}"), leaf) for i, leaf in enumerate(closure_leaves)
        )
        all_leaves = tensors + closure_with_keys
        context = (
            *(
                self._wrap_context_value(attr, getattr(self, attr))
                if attr != "mask_mod"
                else _MaskModWrapper(stripped, callable_spec)
                for attr in self._CONTEXT_ATTRS
            ),
            optional_tensor_attrs,
        )
        return all_leaves, context

def _broadcast_to_dim(x: Tensor, dim: int) -> Tensor:
    while x.dim() < dim:
        x = x.unsqueeze(0)
    return x

def _round_up_to_multiple(x: int, multiple: int) -> int:
    return (x + multiple - 1) // multiple * multiple

def _convert_mask_to_block_mask(
    mask: Tensor,
    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
    separate_full_blocks: bool = False,
) -> tuple[Tensor, Tensor | None]:
    if mask.dtype != tensorplay.bool:
        raise AssertionError(f"mask.dtype must be bool, got {mask.dtype}")
    mask = _broadcast_to_dim(mask, 4)

    def padding_needed_for_multiple(x, multiple):
        return _round_up_to_multiple(x, multiple) - x

    mask = tensorplay.nn.functional.pad(
        mask,
        (
            0,
            padding_needed_for_multiple(mask.shape[-1], KV_BLOCK_SIZE),
            0,
            padding_needed_for_multiple(mask.shape[-2], Q_BLOCK_SIZE),
        ),
    )
    B, H, Q, KV = mask.shape
    if Q % Q_BLOCK_SIZE != 0:
        raise AssertionError(
            f"Q ({Q}) must be divisible by Q_BLOCK_SIZE ({Q_BLOCK_SIZE})"
        )
    if KV % KV_BLOCK_SIZE != 0:
        raise AssertionError(
            f"KV ({KV}) must be divisible by KV_BLOCK_SIZE ({KV_BLOCK_SIZE})"
        )
    mask = mask.view(
        B, H, Q // Q_BLOCK_SIZE, Q_BLOCK_SIZE, KV // KV_BLOCK_SIZE, KV_BLOCK_SIZE
    )  # [B, H, Q//Q_BLOCK_SIZE, Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE, KV_BLOCK_SIZE]
    mask = mask.permute(
        0, 1, 2, 4, 3, 5
    )  # [B, H, Q//Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE, Q_BLOCK_SIZE, KV_BLOCK_SIZE]
    mask_block_sum = mask.sum(
        dim=[-2, -1]
    )  # [B, H, Q//Q_BLOCK_SIZE, KV//KV_BLOCK_SIZE]
    if separate_full_blocks:
        full_block_sum = Q_BLOCK_SIZE * KV_BLOCK_SIZE
        full_blocks = mask_block_sum == full_block_sum
        partial_blocks = (mask_block_sum > 0) & (mask_block_sum < full_block_sum)
        partial_blocks = partial_blocks.to(dtype=tensorplay.int8)
        full_blocks = full_blocks.to(dtype=tensorplay.int8)
        return partial_blocks, full_blocks
    else:
        partial_blocks = mask_block_sum > 0
        partial_blocks = partial_blocks.to(dtype=tensorplay.int8)
        return partial_blocks, None

[docs]
def or_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:
    """Returns a mask_mod that's the union of provided mask_mods"""
    if not all(callable(arg) for arg in mask_mods):
        raise RuntimeError(f"All inputs should be callable mask_mods: {mask_mods}")

    def or_mask(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
        result = b.new_zeros((), dtype=tensorplay.bool)
        for mask in mask_mods:
            result = result | mask(b, h, q_idx, kv_idx)
        return result

    return or_mask

[docs]
def and_masks(*mask_mods: _mask_mod_signature) -> _mask_mod_signature:
    """Returns a mask_mod that's the intersection of provided mask_mods"""
    if not all(callable(arg) for arg in mask_mods):
        raise RuntimeError(f"All inputs should be callable mask_mods: {mask_mods}")

    def and_mask(b: Tensor, h: Tensor, q_idx: Tensor, kv_idx: Tensor) -> Tensor:
        result = b.new_ones((), dtype=tensorplay.bool)
        for mask in mask_mods:
            result = result & mask(b, h, q_idx, kv_idx)
        return result

    return and_mask

def _convert_block_mask_to_mask(
    block_mask: Tensor,
    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
) -> Tensor:
    if block_mask.dim() != 4:
        raise AssertionError(f"block_mask.dim() must be 4, got {block_mask.dim()}")
    B, H, Q, KV = block_mask.shape
    block_mask = block_mask.expand(Q_BLOCK_SIZE, KV_BLOCK_SIZE, *block_mask.shape)
    block_mask = block_mask.permute(2, 3, 4, 0, 5, 1).reshape(
        B, H, Q * Q_BLOCK_SIZE, KV * KV_BLOCK_SIZE
    )
    return block_mask

def _create_sparse_block_from_block_mask(
    block_mask: tuple[Tensor, Tensor | None],
    mask_mod: _mask_mod_signature | None,
    seq_lengths: tuple[int, int],
    Q_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
    KV_BLOCK_SIZE: int = _DEFAULT_SPARSE_BLOCK_SIZE,
) -> BlockMask:
    partial_blocks, full_blocks = block_mask

    partial_bm = _dense_to_ordered(partial_blocks)
    if full_blocks is not None:
        full_bm: tuple[Tensor | None, Tensor | None] = _dense_to_ordered(full_blocks)
    else:
        full_bm = (None, None)

    return BlockMask.from_kv_blocks(
        partial_bm[0],
        partial_bm[1],
        full_bm[0],
        full_bm[1],
        BLOCK_SIZE=(Q_BLOCK_SIZE, KV_BLOCK_SIZE),
        mask_mod=mask_mod,
        seq_lengths=seq_lengths,
    )

[docs]
def create_mask(
    mod_fn: _score_mod_signature | _mask_mod_signature,
    B: int | None,
    H: int | None,
    Q_LEN: int,
    KV_LEN: int,
    device: DeviceLikeType | None = None,
) -> Tensor:
    r"""This function creates a mask tensor from a mod_fn function.

    Args:
        mod_fn (Union[_score_mod_signature, _mask_mod_signature]): Function to modify attention scores.
        B (int): Batch size.
        H (int): Number of query heads.
        Q_LEN (int): Sequence length of query.
        KV_LEN (int): Sequence length of key/value.
        device (str): Device to run the mask creation on.

    Returns:
        mask (Tensor): A mask tensor with shape (B, H, M, N).
    """
    if device is None:
        device = tensorplay.accelerator.current_accelerator() or "cpu"
    if B is None:
        B = 1
    if H is None:
        H = 1
    b = tensorplay.arange(0, B, device=device)
    h = tensorplay.arange(0, H, device=device)
    m = tensorplay.arange(0, Q_LEN, device=device)
    n = tensorplay.arange(0, KV_LEN, device=device)
    mod_type = _get_mod_type(mod_fn)

    from tensorplay._higher_order_ops.utils import TransformGetItemToIndex

    with TransformGetItemToIndex():
        if mod_type == _ModificationType.SCORE_MOD:
            score_mod = mod_fn
            score_mod = _vmap_for_bhqkv(score_mod, prefix=(0,))  # first input is score
            out = score_mod(tensorplay.zeros(B, H, Q_LEN, KV_LEN, device=device), b, h, m, n)
            mask = tensorplay.where(tensorplay.isneginf(out), False, True)
            return mask
        elif mod_type == _ModificationType.MASK_MOD:
            mask_mod = mod_fn
            mask_mod = _vmap_for_bhqkv(mask_mod, prefix=())
            mask = mask_mod(b, h, m, n)
            return mask
        else:
            raise AssertionError

def _compute_dq_write_order_from_block_mask(
    block_mask: BlockMask,
    dq_kv_order: bool | None = False,
) -> tuple[Tensor, Tensor | None]:
    """Compute dQ write-order metadata for deterministic block-sparse backward.

    Algorithm (per batch/head):
    1. Convert the partial and full `kv` block lists into a dense `m_block -> n_block`
       contributor mask and merge them.
    2. Build a rank table in the requested scheduler order. ``dq_kv_order=False``
       means ascending n-block order and ``True`` means descending/SPT order.
    3. Gather ranks using backward `q_indices` order to produce semaphore values
       used by deterministic dQ accumulation.
    """
    kv_num_blocks = block_mask.kv_num_blocks
    kv_indices = block_mask.kv_indices
    full_kv_num_blocks = block_mask.full_kv_num_blocks
    full_kv_indices = block_mask.full_kv_indices
    q_num_blocks = block_mask.q_num_blocks
    q_indices = block_mask.q_indices
    full_q_num_blocks = block_mask.full_q_num_blocks
    full_q_indices = block_mask.full_q_indices

    if q_num_blocks is None or q_indices is None:
        raise ValueError(
            "BlockMask must have q_num_blocks and q_indices to compute dq_write_order"
        )

    device = kv_indices.device

    def _expand_bh(tensor: Tensor, target_b: int, target_h: int) -> Tensor:
        return tensor.expand(target_b, target_h, *tensor.shape[2:])

    broadcast_shapes = [
        kv_indices.shape[:2],
        q_indices.shape[:2],
    ]
    if full_kv_indices is not None:
        broadcast_shapes.append(full_kv_indices.shape[:2])
    if full_q_indices is not None:
        broadcast_shapes.append(full_q_indices.shape[:2])
    B, H = tensorplay.broadcast_shapes(*broadcast_shapes)
    kv_num_blocks = _expand_bh(kv_num_blocks, B, H)
    kv_indices = _expand_bh(kv_indices, B, H)
    if full_kv_num_blocks is not None:
        full_kv_num_blocks = _expand_bh(full_kv_num_blocks, B, H)
    if full_kv_indices is not None:
        full_kv_indices = _expand_bh(full_kv_indices, B, H)

    _, _, num_m, _ = kv_indices.shape
    _, _, num_n, _ = q_indices.shape

    has_full = full_kv_num_blocks is not None and full_kv_indices is not None

    dense_partial = _ordered_to_dense(kv_num_blocks, kv_indices, num_n)
    if has_full:
        if full_kv_num_blocks is None or full_kv_indices is None:
            raise AssertionError(
                "full_kv_num_blocks and full_kv_indices must not be None"
            )
        dense_full = _ordered_to_dense(full_kv_num_blocks, full_kv_indices, num_n)
        dense = (dense_partial + dense_full).clamp(max=1)
    else:
        dense = dense_partial

    cumsum = dense.cumsum(dim=-1)
    rank_table = (cumsum - dense).to(tensorplay.int32)

    if dq_kv_order:
        total_per_m = cumsum[:, :, :, -1:]
        rank_table = (total_per_m - 1 - rank_table).to(tensorplay.int32)

    def _gather_write_order(bwd_idx: Tensor) -> Tensor:
        b_i = tensorplay.arange(B, device=device)[:, None, None, None]
        h_i = tensorplay.arange(H, device=device)[None, :, None, None]
        n_i = tensorplay.arange(bwd_idx.shape[2], device=device)[None, None, :, None]
        m_vals = bwd_idx.long().clamp(0, num_m - 1)
        return rank_table[b_i, h_i, m_vals, n_i].to(tensorplay.int32)

    dq_write_order = _gather_write_order(q_indices)

    dq_write_order_full = None
    if has_full and full_q_num_blocks is not None and full_q_indices is not None:
        dq_write_order_full = _gather_write_order(full_q_indices)

    return dq_write_order, dq_write_order_full

[docs]
def create_block_mask(
    mask_mod: _mask_mod_signature,
    B: int | None,
    H: int | None,
    Q_LEN: int,
    KV_LEN: int,
    device: DeviceLikeType | None = None,
    BLOCK_SIZE: int | tuple[int, int] = _DEFAULT_SPARSE_BLOCK_SIZE,
    _compile=False,
    separate_full_blocks: bool = True,
    compute_dq_write_order: bool = False,
    dq_kv_order: bool = True,
) -> BlockMask:
    r"""This function creates a block mask tuple from a mask_mod function.

    Args:
        mask_mod (Callable): mask_mod function. This is a callable that defines the
            masking pattern for the attention mechanism. It takes four arguments:
            b (batch size), h (number of heads), q_idx (query index), and kv_idx (key/value index).
            It should return a boolean tensor indicating which attention connections are allowed (True)
            or masked out (False).
        B (int): Batch size.
        H (int): Number of query heads.
        Q_LEN (int): Sequence length of query.
        KV_LEN (int): Sequence length of key/value.
        device (str): Device to run the mask creation on.
        BLOCK_SIZE (int or tuple[int, int]): Block size for the block mask. If a single int is provided it is used for both query and key/value.
        separate_full_blocks (bool): If True, fully unmasked blocks are stored
            separately so kernels can skip mask_mod on those blocks. If False,
            all non-empty blocks are stored as partial blocks and mask_mod is
            applied to every block.
        compute_dq_write_order (bool): If True, precompute dQ write-order
            metadata needed by deterministic block-sparse backward.
        dq_kv_order (bool): KV-column scheduler order used for deterministic
            dQ accumulation when compute_dq_write_order is True. False means
            ascending n-block order and True means descending/SPT order.
            Explicit tensor schedules are not supported by create_block_mask
            yet; they are supported by BlockMask.from_kv_blocks for callers
            that provide precomputed write-order metadata directly.

    Returns:
        BlockMask:  A BlockMask object that contains the block mask information.

    Example Usage:
        .. code-block:: python

            def causal_mask(b, h, q_idx, kv_idx):
                return q_idx >= kv_idx

            block_mask = create_block_mask(causal_mask, 1, 1, 8192, 8192, device="cuda")
            query = tensorplay.randn(1, 1, 8192, 64, device="cuda", dtype=tensorplay.float16)
            key = tensorplay.randn(1, 1, 8192, 64, device="cuda", dtype=tensorplay.float16)
            value = tensorplay.randn(1, 1, 8192, 64, device="cuda", dtype=tensorplay.float16)
            output = omni_attention(query, key, value, block_mask=block_mask)
    """
    if device is None:
        device = tensorplay.accelerator.current_accelerator() or "cpu"
    mod_type = _get_mod_type(mask_mod)
    if mod_type != _ModificationType.MASK_MOD:
        raise AssertionError(
            f"create-block_mask requires a mask_mod function! Got {mask_mod}"
        )
    if B is None:
        B = 1
    if H is None:
        H = 1
    if isinstance(BLOCK_SIZE, int):
        Q_BLOCK_SIZE = BLOCK_SIZE
        KV_BLOCK_SIZE = BLOCK_SIZE
    else:
        Q_BLOCK_SIZE, KV_BLOCK_SIZE = BLOCK_SIZE

    if _compile:
        warnings.warn(
            "_compile flag on create_block_mask was originally added to work around a compile limitation. That limitation has since been addressed. So, to compile create_block_mask, we suggest doing compile(create_block_mask). This still works for now, but will be removed in the future.",
            DeprecationWarning,
            stacklevel=2,
        )
        return tensorplay.compile(create_block_mask)(
            mask_mod,
            B,
            H,
            Q_LEN,
            KV_LEN,
            device,
            BLOCK_SIZE,
            False,
            separate_full_blocks,
            compute_dq_write_order,
            dq_kv_order,
        )

    if not isinstance(dq_kv_order, bool):
        raise ValueError("dq_kv_order must be a bool when using create_block_mask")

    mask_tensor = create_mask(mask_mod, B, H, Q_LEN, KV_LEN, device)
    partial_block_mask, full_block_mask = _convert_mask_to_block_mask(
        mask_tensor,
        Q_BLOCK_SIZE=Q_BLOCK_SIZE,
        KV_BLOCK_SIZE=KV_BLOCK_SIZE,
        separate_full_blocks=separate_full_blocks,
    )
    block_mask = _create_sparse_block_from_block_mask(
        (partial_block_mask, full_block_mask),
        mask_mod,
        (Q_LEN, KV_LEN),
        Q_BLOCK_SIZE,
        KV_BLOCK_SIZE,
    )

    if compute_dq_write_order:
        dq_wo, dq_wo_full = _compute_dq_write_order_from_block_mask(
            block_mask, dq_kv_order=dq_kv_order
        )
        block_mask.dq_write_order = dq_wo
        block_mask.dq_write_order_full = dq_wo_full
        block_mask.dq_kv_order = None
        block_mask.dq_kv_order_spt = dq_kv_order

    return block_mask

def _create_empty_block_mask(query: Tensor, key: Tensor) -> BlockMask:
    r"""Default block mask for omni attention.
    If users don't specify any block sparse mask info, we create this
    empty block sparse mask. Which creates a BlockMask with 1 block that is the full length
    of the query and key tensors.
    """
    device = query.device
    return BlockMask.from_kv_blocks(
        kv_num_blocks=tensorplay.ones(1, 1, 1, dtype=tensorplay.int32, device=device),
        kv_indices=tensorplay.zeros(1, 1, 1, 1, dtype=tensorplay.int32, device=device),
        BLOCK_SIZE=_LARGE_SPARSE_BLOCK_SIZE,
        seq_lengths=(1, 1),
    )

def _apply_kernel_options(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    return_lse: bool,
    kernel_options: OmniKernelOptions | None,
    return_aux: AuxRequest | None = None,
) -> _KernelOptionsWithInternals:
    kernel_options = cast(
        _KernelOptionsWithInternals,
        {} if kernel_options is None else dict(kernel_options),
    )

    if "BACKEND" in kernel_options and kernel_options.get(
        "FORCE_USE_OMNI_ATTENTION", False
    ):
        # TODO: remove FORCE_USE_OMNI_ATTENTION once BACKEND is fully adopted.
        raise RuntimeError(
            "BACKEND cannot be combined with legacy FORCE_USE_OMNI_ATTENTION. "
            "BACKEND supersedes the legacy knob; please drop FORCE_USE_OMNI_ATTENTION "
            "and only specify the desired BACKEND."
        )

    if "BACKEND" in kernel_options:
        valid_backends = typing.get_args(_Backend)
        if kernel_options["BACKEND"] not in valid_backends:
            raise ValueError(
                f"Invalid BACKEND value '{kernel_options['BACKEND']}'. "
                f"Must be one of {valid_backends}"
            )

    kernel_options.setdefault("BACKEND", "AUTO")
    kernel_options.setdefault("PRESCALE_QK", False)
    kernel_options.setdefault("ROWS_GUARANTEED_SAFE", False)
    kernel_options.setdefault("BLOCKS_ARE_CONTIGUOUS", False)
    # This forces all biases grad scatters to be done in the DQ iteration loop of the backwards
    kernel_options.setdefault("WRITE_DQ", True)

    any_inputs_on_cpu_device = (
        query.device.type == "cpu"
        or key.device.type == "cpu"
        or value.device.type == "cpu"
    )
    any_inputs_on_mps_device = (
        query.device.type == "mps"
        or key.device.type == "mps"
        or value.device.type == "mps"
    )

    # Determine what auxiliary outputs are needed
    output_lse = return_lse
    output_max = False

    if return_aux is not None:
        # New API takes precedence over legacy parameters
        output_lse = return_aux.lse
        output_max = return_aux.max_scores

    # If forward kernel needs to return logsumexp is decided by this rule internally.
    if "OUTPUT_LOGSUMEXP" in kernel_options:
        raise AssertionError("OUTPUT_LOGSUMEXP must not be in kernel_options")
    kernel_options["OUTPUT_LOGSUMEXP"] = True
    if not output_lse:
        # We used to check if q,k,v required grads but since captured buffers can require grad
        # we always write unless in no_grad
        kernel_options["OUTPUT_LOGSUMEXP"] = tensorplay.is_grad_enabled()
        if any_inputs_on_cpu_device or any_inputs_on_mps_device:
            # CPU/MPS support inference only, no LSE/backward yet.
            # TODO: support CPU for training and return lse
            kernel_options["OUTPUT_LOGSUMEXP"] = False
        if any_inputs_on_mps_device:
            # MPS supports inference only; backward / LSE not yet implemented
            kernel_options["OUTPUT_LOGSUMEXP"] = False

    # If forward kernel needs to return max is decided by this rule internally.
    if "OUTPUT_MAX" in kernel_options:
        raise AssertionError("OUTPUT_MAX must not be in kernel_options")
    if kernel_options["BACKEND"] == "FLASH" and output_max:
        raise NotImplementedError(
            "Returning max scores is not supported with BACKEND='FLASH'. "
            "Use return_aux=AuxRequest(lse=True) or omit max_scores."
        )
    kernel_options["OUTPUT_MAX"] = output_max
    if any_inputs_on_cpu_device and output_max:
        # CPU doesn't support returning max yet
        # TODO: support CPU for returning max
        raise NotImplementedError("Returning max scores is not supported on CPU.")

    return kernel_options

def _validate_embed_dim(query: Tensor, key: Tensor, value: Tensor) -> None:
    # Read through the shape rather than through the method form: while a
    # region is being captured the extents a value stands for are settled --
    # they are part of what the region was captured for -- so they can be
    # compared here, where the method form would leave a value that has no
    # answer until the graph is built, and a comparison cannot wait for that.
    if query.shape[-1] != key.shape[-1]:
        raise ValueError(
            f"Expect query and key/value to have the same embedding dimension "
            f"but got E={query.shape[-1]} and E={key.shape[-1]}."
        )

def _validate_device(query: Tensor, key: Tensor, value: Tensor) -> None:
    """TODO: Remove once non cuda/cpu devices support is added
    We only need to check query since we have already that q,k,v are on the same device
    """
    if query.device.type == "cpu" and (
        query.requires_grad or key.requires_grad or value.requires_grad
    ):
        raise NotImplementedError(
            "OmniAttention does not support backward on CPU. Please set the input requires_grad to False or use another device."
        )
    if query.device.type == "mps" and (
        query.requires_grad or key.requires_grad or value.requires_grad
    ):
        raise NotImplementedError(
            "OmniAttention does not support backward on MPS. Please set the input requires_grad to False or use another device."
        )
    supported_devices = {"cuda", "cpu", "xpu", "hpu", "mps"}
    if query.device.type not in supported_devices:
        raise ValueError(
            "OmniAttention is only supported on CUDA, CPU, HPU, or MPS devices. "
            f"Found input tensors on {query.device.type} device."
        )

def _validate_no_nested_tensors(query: Tensor, key: Tensor, value: Tensor) -> None:
    if query.is_nested() or key.is_nested() or value.is_nested():
        raise NotImplementedError(
            "omni_attention does not support NestedTensor inputs, including "
            "compile(omni_attention) with jagged NestedTensor inputs. "
            "Convert inputs to dense tensors before calling omni_attention."
        )

def _enforce_mem_layouts(
    query: Tensor, key: Tensor, value: Tensor
) -> tuple[Tensor, Tensor, Tensor]:
    """
    Enforce memory layouts for query, key, and value tensors.

    For non-FP8 dtypes, no action is taken.

    For FP8 dtypes, we enforce the following memory layouts:
    - Query tensor must be in row-major memory layout, as it will be the left-operand in the FP8 GEMM `q @ k.T`.
    - Key tensor must be in row-major memory layout, as it will be transposed when used as the right-operand
      in the FP8 GEMM `q @ k.T`, meaning it will correctly be in column-major memory layout for the GEMM.
    - Value tensor must be in column-major memory layout, as it will be the right-operand in the FP8 GEMM `softmax_scores @ v`.

    Returns the query, key, and value tensors with the enforced memory layouts.
    """

    def is_row_major(tensor: Tensor) -> bool:
        return tensor.stride()[-1] == 1

    def is_col_major(tensor: Tensor) -> bool:
        return tensor.stride()[-2] == 1

    # These memory layout constraints are only for FP8 GEMMs on NVIDIA GPU architectures >= SM89 and < SM100.
    # This is because GPU arch < SM89 does not support FP8 GEMMs, and
    # SM100 has support for TN, NT, TT, NN layouts for FP8 GEMMs
    # (i.e., left and right operands can be in row or column major layouts)
    # so this check is only needed for older architectures.
    fp8_dtypes = (
        tensorplay.float8_e4m3fn,
        tensorplay.float8_e5m2,
    )
    gemm_precision = query.dtype

    should_enforce_mem_layout = (
        gemm_precision in fp8_dtypes
        and tensorplay.version.cuda is not None
        and tensorplay.cuda.get_device_capability("cuda") >= (8, 9)
        and tensorplay.cuda.get_device_capability("cuda") < (10, 0)
    )
    if not should_enforce_mem_layout:
        return query, key, value

    # Query must be in row-major memory layout as the left-operand in the FP8 GEMM `q @ k.T`
    if not is_row_major(query):
        query = query.contiguous()

    # Key must be in row-major memory layout as it will be transposed when used as the right-operand
    # in the FP8 GEMM `q @ k.T`, meaning it will correctly be in column-major memory layout for the GEMM.
    if not is_row_major(key):
        key = key.contiguous()

    # Value must be in column-major memory layout as the right-operand in the FP8 GEMM `softmax_scores @ v`
    if not is_col_major(value):
        value = value.transpose(-2, -1).contiguous().transpose(-2, -1)
    return query, key, value

@overload
def omni_attention(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    score_mod: _score_mod_signature | None = ...,
    block_mask: BlockMask | None = ...,
    scale: float | None = ...,
    enable_gqa: bool = ...,
    return_lse: Literal[False] = ...,
    kernel_options: OmniKernelOptions | None = ...,
    *,
    return_aux: None = ...,
) -> Tensor: ...

@overload
@deprecated(
    "return_lse is deprecated and will be removed in a future release. "
    "Use return_aux=AuxRequest(lse=True) instead.",
    category=FutureWarning,
)
def omni_attention(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    score_mod: _score_mod_signature | None = ...,
    block_mask: BlockMask | None = ...,
    scale: float | None = ...,
    enable_gqa: bool = ...,
    return_lse: Literal[True] = ...,
    kernel_options: OmniKernelOptions | None = ...,
    *,
    return_aux: None = ...,
) -> tuple[Tensor, Tensor]: ...

@overload
def omni_attention(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    score_mod: _score_mod_signature | None = ...,
    block_mask: BlockMask | None = ...,
    scale: float | None = ...,
    enable_gqa: bool = ...,
    return_lse: bool = ...,
    kernel_options: OmniKernelOptions | None = ...,
    *,
    return_aux: AuxRequest,
) -> tuple[Tensor, AuxOutput]: ...

@overload
def omni_attention(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    score_mod: _score_mod_signature | None = ...,
    block_mask: BlockMask | None = ...,
    scale: float | None = ...,
    enable_gqa: bool = ...,
    return_lse: Literal[True] = ...,
    kernel_options: OmniKernelOptions | None = ...,
    *,
    return_aux: AuxRequest,
) -> Never: ...

def omni_attention(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    score_mod: _score_mod_signature | None = None,
    block_mask: BlockMask | None = None,
    scale: float | None = None,
    enable_gqa: bool = False,
    return_lse: bool = False,
    kernel_options: OmniKernelOptions | None = None,
    *,
    return_aux: AuxRequest | None = None,
) -> Tensor | tuple[Tensor, Tensor] | tuple[Tensor, AuxOutput]:
    r"""This function implements scaled dot product attention with an arbitrary attention score modification function.

    This function computes the scaled dot product attention between query, key, and value tensors with a user-defined
    attention score modification function. The attention score modification function will be applied after the attention
    scores have been calculated between the query and key tensors. The attention scores are calculated as follows:

    The ``score_mod`` function should have the following signature:

    .. code-block:: python

        def score_mod(
            score: Tensor,
            batch: Tensor,
            head: Tensor,
            q_idx: Tensor,
            k_idx: Tensor
        ) -> Tensor:

    Where:
        - ``score``: A scalar tensor representing the attention score,
          with the same data type and device as the query, key, and value tensors.
        - ``batch``, ``head``, ``q_idx``, ``k_idx``: Scalar tensors indicating
          the batch index, query head index, query index, and key/value index, respectively.
          These should have the ``tensorplay.int64`` data type and be located on the same device as the score tensor.

    Args:
        query (Tensor): Query tensor; shape :math:`(B, Hq, L, E)`. For FP8 dtypes, should be in row-major memory layout for optimal performance.
        key (Tensor): Key tensor; shape :math:`(B, Hkv, S, E)`. For FP8 dtypes, should be in row-major memory layout for optimal performance.
        value (Tensor): Value tensor; shape :math:`(B, Hkv, S, Ev)`. For FP8 dtypes, should be in column-major memory layout for optimal performance.
        score_mod (Optional[Callable]): Function to modify attention scores. By default no score_mod is applied.
        block_mask (Optional[BlockMask]): BlockMask object that controls the blocksparsity pattern of the attention.
        scale (Optional[float]): Scaling factor applied prior to softmax. If none, the default value is set to :math:`\frac{1}{\sqrt{E}}`.
        enable_gqa (bool): If set to True, enables Grouped Query Attention (GQA) and broadcasts key/value heads to query heads.
        return_lse (bool): Whether to return the logsumexp of the attention scores. Default is False. **Deprecated**: Use ``return_aux=AuxRequest(lse=True)`` instead.
        kernel_options (Optional[OmniKernelOptions]):
            Options to control the behavior of the underlying kernels.
            See :class:`OmniKernelOptions` for available options and usage examples.
        return_aux (Optional[AuxRequest]): Specifies which auxiliary outputs to compute and return.
            If None, only the attention output is returned. Use ``AuxRequest(lse=True, max_scores=True)``
            to request both auxiliary outputs.

    Returns:
        output (Tensor): Attention output; shape :math:`(B, Hq, L, Ev)`.

        When ``return_aux`` is not None:
            aux (AuxOutput): Auxiliary outputs with requested fields populated.

        When ``return_aux`` is None (deprecated paths):
            lse (Tensor): Log-sum-exp of attention scores; shape :math:`(B, Hq, L)`. Only returned if ``return_lse=True``.

    Shape legend:
        - :math:`N: \text{Batch size} ... : \text{Any number of other batch dimensions (optional)}`
        - :math:`S: \text{Source sequence length}`
        - :math:`L: \text{Target sequence length}`
        - :math:`E: \text{Embedding dimension of the query and key}`
        - :math:`Ev: \text{Embedding dimension of the value}`

    .. warning::
        `tensorplay.nn.attention.omni_attention` is a prototype feature.
        Please look forward to a more stable implementation in a future version.
    """
    # Some basic input validation
    _validate_no_nested_tensors(query, key, value)
    _validate_sdpa_input(query, key, value, allow_lowp_kv=True)
    _validate_embed_dim(query, key, value)
    _validate_device(query, key, value)
    query, key, value = _enforce_mem_layouts(query, key, value)
    if len(query.shape) != 4 or len(key.shape) != 4 or len(value.shape) != 4:
        raise NotImplementedError("NYI: query, key, and value must be 4D tensors")
    if (not enable_gqa) and query.shape[-3] != key.shape[-3]:
        raise ValueError(
            f"Expect query and key/value to have the same number of heads "
            f"but got Hq={query.shape[-3]} and Hkv={key.shape[-3]}. "
            f"Try setting enable_gqa=True for GQA."
        )
    if enable_gqa:
        Hq = query.shape[1]
        Hkv = key.shape[1]
        if Hq % Hkv != 0:
            raise ValueError(
                f"Expect number of query heads to be a multiple of kv heads for GQA "
                f"but got Hq={Hq} and Hkv={Hkv}."
            )
    if query.shape[0] != key.shape[0]:
        if block_mask is None:
            raise ValueError(
                f"Expect query and key/value to have the same batch size, "
                f"or non-none block_mask, "
                f"but got block_mask=None, Bq={query.shape[0]}, and Bkv={key.shape[0]}."
            )

        if block_mask.kv_num_blocks.size(0) != query.shape[0]:
            raise ValueError(
                f"Expect query and key/value to have the same batch size, "
                f"or block_mask and query to have the same batch size, "
                f"but got Bq={query.shape[0]}, Bkv={key.shape[0]}, B_block_mask={block_mask.kv_num_blocks.size(0)}."
            )

    if score_mod is None:
        score_mod = _identity

    if block_mask is None:
        block_mask = _create_empty_block_mask(query, key)

    # If BlockMask was sliced, its mask_mod is intentionally replaced with an error-raising stub.
    # This guard ensures we surface the intended error message before any shape-based checks.
    if getattr(block_mask, "mask_mod", None) is _sliced_mask_mod_error:
        raise RuntimeError("Cannot use mask_mod from a sliced BlockMask")

    if (
        block_mask.BLOCK_SIZE[0] == _LARGE_SPARSE_BLOCK_SIZE
        and block_mask.BLOCK_SIZE[1] == _LARGE_SPARSE_BLOCK_SIZE
    ):
        # This corresponds to the case where we essentially have a "no-op" block mask.
        pass
    else:
        q_len = query.size(-2)
        kv_len = key.size(-2)
        block_mask_q_len = block_mask.shape[-2]
        block_mask_kv_len = block_mask.shape[-1]

        _validate_block_mask_shape(
            query, key, q_len, kv_len, block_mask_q_len, block_mask_kv_len
        )

    if scale is None:
        scale = 1.0 / math.sqrt(query.size(-1))

    if query.device != block_mask.kv_num_blocks.device:  # type: ignore[union-attr]
        raise RuntimeError(
            f"Expect q/k/v and block_mask to be on the same device "
            f"but got {query.device} and {block_mask.kv_num_blocks.device}."  # type: ignore[union-attr]
        )

    # Handle deprecation warnings for old parameters
    if return_lse and return_aux is not None:
        raise ValueError(
            "Cannot specify both return_lse and return_aux. "
            "return_lse is deprecated, please use return_aux=AuxRequest(lse=True) instead."
        )
    elif return_lse and return_aux is None:
        _warn_once(
            "deprecated_return_lse",
            "return_lse is deprecated and will be removed in a future release. "
            "Please use return_aux=AuxRequest(lse=True) instead.",
            category=FutureWarning,
        )

    kernel_options = _apply_kernel_options(
        query,
        key,
        value,
        return_lse,
        kernel_options,
        return_aux,
    )

    def _finalize_outputs(
        out,
        lse,
        max_scores,
        *,
        return_aux: AuxRequest | None,
        return_lse: bool,
        stats_are_log2: bool,
    ):
        """Normalize stats and build return value (aux-aware, legacy-compatible)."""
        ln2 = math.log(2.0)
        return_lse = return_lse or return_aux is not None and return_aux.lse
        return_max = return_aux is not None and return_aux.max_scores

        lse_scaled = lse if (return_lse and lse.numel() > 0) else None
        max_scaled = max_scores if (return_max and max_scores.numel() > 0) else None

        if stats_are_log2:
            lse_scaled = lse_scaled * ln2 if lse_scaled is not None else None
            max_scaled = max_scaled * ln2 if max_scaled is not None else None

        if return_aux is not None:
            return out, AuxOutput(
                lse=lse_scaled,
                max_scores=max_scaled,
            )

        if return_lse:
            return out, lse_scaled

        return out

    if tensorplay.compiler.is_compiling():
        # mark head_dim and number of heads to be static
        for x in [query, key, value]:
            tensorplay.compiler.mark_static(x, -3)
            tensorplay.compiler.mark_static(x, -1)

        out, lse, max_scores = omni_attention_hop(
            query,
            key,
            value,
            score_mod,
            block_mask.as_tuple(),
            scale,
            kernel_options,  # type: ignore[union-attr]
        )
        return _finalize_outputs(
            out,
            lse,
            max_scores,
            return_aux=return_aux,
            return_lse=return_lse,
            stats_are_log2=kernel_options["BACKEND"] != "FLASH",
        )

    if not _OMNI_ATTENTION_DISABLE_COMPILE_DEBUG:
        _warn_once(
            warning_id="omni_attention_performance",
            message=(
                "omni_attention called without compile() - this will use an unfused implementation that materializes the full scores matrix instead of generating a fused kernel.\n\n"
                "SOLUTION: Use compile(omni_attention)(...)\n\n"
                "If you want to debug your score_mod/mask_mod, you can set:\n"
                "tensorplay.nn.attention.omni_attention._OMNI_ATTENTION_DISABLE_COMPILE_DEBUG = True\n\n"
                "This will allow you to use print statements or breakpoints. Note: This doesn't work with the backwards pass and may produce incorrect results."
            ),
        )

    if not hasattr(tensorplay.compiler, "compile"):
        raise RuntimeError("omni_attention requires graph tracer support")

    block_mask_tensors = tuple(
        getattr(block_mask, name) for name in BlockMask._TENSOR_ATTRS
    )
    seq_q, seq_kv = block_mask.seq_lengths
    dq_kv_order_spt = block_mask.dq_kv_order_spt
    q_block_size, kv_block_size = block_mask.BLOCK_SIZE
    mask_mod = block_mask.mask_mod

    def _omni_attention_hop_wrapper(
        query,
        key,
        value,
        kv_num_blocks,
        kv_indices,
        full_kv_num_blocks,
        full_kv_indices,
        q_num_blocks,
        q_indices,
        full_q_num_blocks,
        full_q_indices,
        dq_write_order,
        dq_write_order_full,
        dq_kv_order,
    ):
        traced_block_mask = (
            seq_q,
            seq_kv,
            kv_num_blocks,
            kv_indices,
            full_kv_num_blocks,
            full_kv_indices,
            q_num_blocks,
            q_indices,
            full_q_num_blocks,
            full_q_indices,
            dq_write_order,
            dq_write_order_full,
            dq_kv_order,
            dq_kv_order_spt,
            q_block_size,
            kv_block_size,
            mask_mod,
        )
        return omni_attention_hop(
            query,
            key,
            value,
            score_mod,
            traced_block_mask,
            scale,
            kernel_options,
        )

    with setup_compilation_env() as backend:
        if _OMNI_ATTENTION_DISABLE_COMPILE_DEBUG:
            omni_fn = _omni_attention_hop_wrapper
        else:
            omni_fn = tensorplay.compile(
                _omni_attention_hop_wrapper, backend=backend, fullgraph=True
            )

        out, lse, max_scores = omni_fn(
            query,
            key,
            value,
            *block_mask_tensors,
        )
    return _finalize_outputs(
        out,
        lse,
        max_scores,
        return_aux=return_aux,
        return_lse=return_lse,
        stats_are_log2=_OMNI_ATTENTION_DISABLE_COMPILE_DEBUG
        or kernel_options["BACKEND"] != "FLASH",
    )
```