Blame for kernel/ksa_attn/interface.py

import math
import os
import torch
from typing import Tuple, Optional, Dict
from functools import partial

# Lazy import: CuTe kernel only available on Hopper (sm90+)
_CUTE_AVAILABLE = None


def _check_cute_available():
    global _CUTE_AVAILABLE
    if _CUTE_AVAILABLE is not None:
        return _CUTE_AVAILABLE
    try:
        from flash_attn.cute.interface import _flash_attn_fwd
        from flash_attn.cute.compute_block_sparsity import compute_block_sparsity
        cap = torch.cuda.get_device_capability()
        if cap[0] >= 9:  # sm90+ (Hopper)
            _CUTE_AVAILABLE = True
        else:
            _CUTE_AVAILABLE = False
    except (ImportError, ModuleNotFoundError):
        _CUTE_AVAILABLE = False
    return _CUTE_AVAILABLE


SUMMARY_ATTN_CACHE: Dict[Tuple, Tuple] = {}
HIERARCHICAL_SUMMARY_ATTN_CACHE: Dict[Tuple, Tuple] = {}


# ---------------------------------------------------------------------------
# CuTe kernel path (Hopper only)
# ---------------------------------------------------------------------------

class SummaryAttnFunc(torch.autograd.Function):
    @staticmethod
    def _get_cached_value(
        seqlen_q, seqlen_k, chunk_size, summary_num, sliding_chunk_num,
        skip_old_summary=False, is_fwd=True, m_block_size=128, n_block_size=128,
    ):
        from flash_attn.cute.compute_block_sparsity import compute_block_sparsity
        from .mask import (
            summary_attn_mask_func, summary_attn_contiguous_mask_func,
            sliding_chunk_mask_func, sliding_chunk_contiguous_mask_func,
        )

        prefix_summary_length = seqlen_k - seqlen_q
        cache_key = (seqlen_q, chunk_size, summary_num, sliding_chunk_num,
                     prefix_summary_length, skip_old_summary, is_fwd)
        if cache_key not in SUMMARY_ATTN_CACHE:
            if skip_old_summary:
                func = sliding_chunk_contiguous_mask_func if prefix_summary_length > 0 else sliding_chunk_mask_func
            else:
                func = summary_attn_contiguous_mask_func if prefix_summary_length > 0 else summary_attn_mask_func
            mask_mod = partial(func, chunk_size=chunk_size, summary_num=summary_num,
                               sliding_chunk_num=sliding_chunk_num,
                               prefix_summary_length=prefix_summary_length, swap_indices=False)
            mask_swap_mod = mask_mod
            if not is_fwd:
                mask_swap_mod = partial(func, chunk_size=chunk_size, summary_num=summary_num,
                                        sliding_chunk_num=sliding_chunk_num,
                                        prefix_summary_length=prefix_summary_length, swap_indices=True)
                m_block_size, n_block_size = n_block_size, m_block_size
                seqlen_q, seqlen_k = seqlen_k, seqlen_q
            _, block_sparse = compute_block_sparsity(
                m_block_size, n_block_size, 1, 1, seqlen_q, seqlen_k,
                mask_swap_mod, None, "cuda",
            )
            SUMMARY_ATTN_CACHE[cache_key] = (block_sparse, mask_mod)
        return SUMMARY_ATTN_CACHE[cache_key]

    @staticmethod
    def forward(ctx, q, k, v, chunk_size, summary_num, sliding_chunk_num,
                softmax_scale=None, causal=False, pack_gqa=None,
                deterministic=False, skip_old_summary=False,
                force_return_lse=False):
        from flash_attn.cute.interface import _flash_attn_fwd
        block_sparse, mask_mod = SummaryAttnFunc._get_cached_value(
            q.shape[1], k.shape[1], chunk_size, summary_num, sliding_chunk_num,
            skip_old_summary=skip_old_summary, is_fwd=True, m_block_size=128, n_block_size=64,
        )
        out, lse = _flash_attn_fwd(q, k, v, softmax_scale=softmax_scale, causal=causal,
                                     pack_gqa=pack_gqa, mask_mod=mask_mod,
                                     block_sparse_tensors=block_sparse,
                                     m_block_size=128, n_block_size=64,
                                     return_lse=force_return_lse)
        return out, lse



# ---------------------------------------------------------------------------
# Compressed-KSA CuTe path (Hopper only)
# ---------------------------------------------------------------------------

class HierarchicalSummaryAttnFunc(torch.autograd.Function):
    @staticmethod
    def _get_cached_value(
        device_index, seqlen_q, seqlen_k, chunk_size, sliding_chunk_num,
        group_size, meta_length, *, is_fwd,
        m_block_size, n_block_size,
    ):
        from flash_attn.cute.compute_block_sparsity import compute_block_sparsity
        from .mask import hierarchical_summary_attn_mask_func

        cache_key = (
            device_index, seqlen_q, seqlen_k, chunk_size, sliding_chunk_num,
            group_size, meta_length, is_fwd,
        )
        cached = HIERARCHICAL_SUMMARY_ATTN_CACHE.get(cache_key)
        if cached is not None:
            return cached

        mask_mod = partial(
            hierarchical_summary_attn_mask_func,
            chunk_size=chunk_size,
            sliding_chunk_num=sliding_chunk_num,
            group_size=group_size,
            meta_length=meta_length,
            swap_indices=False,
        )
        sparsity_mask_mod = mask_mod
        sparsity_seqlen_q, sparsity_seqlen_k = seqlen_q, seqlen_k
        if not is_fwd:
            sparsity_mask_mod = partial(
                hierarchical_summary_attn_mask_func,
                chunk_size=chunk_size,
                sliding_chunk_num=sliding_chunk_num,
                group_size=group_size,
                meta_length=meta_length,
                swap_indices=True,
            )
            m_block_size, n_block_size = n_block_size, m_block_size
            sparsity_seqlen_q, sparsity_seqlen_k = seqlen_k, seqlen_q

        _, block_sparse = compute_block_sparsity(
            m_block_size,
            n_block_size,
            1,
            1,
            sparsity_seqlen_q,
            sparsity_seqlen_k,
            sparsity_mask_mod,
            None,
            "cuda",
        )
        cached = (block_sparse, mask_mod)
        HIERARCHICAL_SUMMARY_ATTN_CACHE[cache_key] = cached
        return cached

    @staticmethod
    def forward(
        ctx, q, k, v, meta_k, meta_v,
        chunk_size, sliding_chunk_num, group_size,
        softmax_scale=None, deterministic=False, mass_correction=True,
        force_return_lse=False,
    ):
        from flash_attn.cute.interface import _flash_attn_fwd

        from .mask import (
            hierarchical_summary_attn_score_mod,
        )

        # Keep one physical entry per meta K/V and apply log(G) directly to
        # its score. This preserves the mass correction without undoing the
        # sequence compression inside the attention kernel.
        needs_backward = any(
            tensor.requires_grad for tensor in (q, k, v, meta_k, meta_v)
        )
        meta_length = meta_k.shape[1]
        full_k = torch.cat((meta_k, k), dim=1).contiguous()
        full_v = torch.cat((meta_v, v), dim=1).contiguous()
        meta_logit_bias = math.log(group_size)
        score_mod = None
        if mass_correction:
            score_mod = partial(
                hierarchical_summary_attn_score_mod,
                meta_length=meta_length,
                meta_logit_bias=meta_logit_bias,
            )
        block_sparse, mask_mod = HierarchicalSummaryAttnFunc._get_cached_value(
            q.device.index,
            q.shape[1],
            full_k.shape[1],
            chunk_size,
            sliding_chunk_num,
            group_size,
            meta_length,
            is_fwd=True,
            m_block_size=128,
            n_block_size=64,
        )
        out, lse = _flash_attn_fwd(
            q,
            full_k,
            full_v,
            softmax_scale=softmax_scale,
            causal=False,
            pack_gqa=None,
            score_mod=score_mod,
            mask_mod=mask_mod,
            block_sparse_tensors=block_sparse,
            m_block_size=128,
            n_block_size=64,
            # torch.autograd.Function.forward runs with grad recording
            # disabled.  Therefore full_k/full_v created above do not carry
            # requires_grad even when meta_k/meta_v are the only trainable
            # inputs (the meta-pool-only warm-up case).  Force LSE retention
            # based on the original inputs so custom backward has the state
            # it needs without adding inference-only storage.
            return_lse=needs_backward or force_return_lse,
        )
        return out, lse



# ---------------------------------------------------------------------------
# Per-query attention-reuse S-KSA CuTe path (Hopper only)
# ---------------------------------------------------------------------------





# ---------------------------------------------------------------------------
# Summary-indexed, 8-token block sparse Full attention (Hopper only)
# ---------------------------------------------------------------------------

def _validate_summary_block_sparse_inputs(
    q,
    k,
    v,
    selected_block_indices,
    block_size_tokens,
    local_window_tokens,
    local_alignment_tokens=None,
):
    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:
        raise ValueError("q, k, and v must have [batch, text, heads, head_dim]")
    if q.shape[0] != k.shape[0] or k.shape != v.shape:
        raise ValueError("q, k, and v batch/KV shapes must agree")
    if q.shape[1] != k.shape[1]:
        raise ValueError(
            "summary_block_sparse_attn_func currently requires equal Q/K text lengths"
        )
    if selected_block_indices.ndim != 3:
        raise ValueError("selected_block_indices must have [batch, query, topk]")
    if tuple(selected_block_indices.shape[:2]) != tuple(q.shape[:2]):
        raise ValueError("one selected-block row is required per text query")
    if selected_block_indices.dtype not in (torch.int32, torch.int64):
        raise ValueError("selected_block_indices must be int32 or int64")
    if not isinstance(block_size_tokens, int) or block_size_tokens <= 0:
        raise ValueError("block_size_tokens must be a positive integer")
    if block_size_tokens not in (1, 2, 4, 8):
        raise ValueError("route span must be 1, 2, 4, or 8 tokens")
    if not isinstance(local_window_tokens, int) or local_window_tokens <= 0:
        raise ValueError("local_window_tokens must be a positive integer")
    if local_alignment_tokens is None:
        local_alignment_tokens = block_size_tokens
    if (
        not isinstance(local_alignment_tokens, int)
        or local_alignment_tokens <= 0
        or local_alignment_tokens % block_size_tokens
    ):
        raise ValueError(
            "local_alignment_tokens must be a positive multiple of the "
            "route span"
        )
    if q.shape[-1] != k.shape[-1]:
        raise ValueError("Q and K head dimensions must agree")
    if q.shape[2] % k.shape[2] != 0:
        raise ValueError("the number of Q heads must be divisible by KV heads")
    if not all(t.is_cuda for t in (q, k, v, selected_block_indices)):
        raise RuntimeError("summary_block_sparse_attn_func requires CUDA tensors")
    if not (q.device == k.device == v.device == selected_block_indices.device):
        raise ValueError("q, k, v, and selected_block_indices must share a device")

    if os.environ.get("SKSA_VALIDATE_BLOCK_INDICES", "0") == "1":
        raw_block_count = math.ceil(k.shape[1] / block_size_tokens)
        # This diagnostic intentionally synchronises.  The production path
        # relies on the indexer's TopK contract and remains asynchronous.
        invalid = (selected_block_indices < -1) | (
            selected_block_indices >= raw_block_count
        )
        if bool(invalid.any().item()):
            raise ValueError(
                f"selected block IDs must be -1 or in [0, {raw_block_count})"
            )


def _pack_summary_block_indices(
    selected_block_indices,
    *,
    raw_block_count,
    max_workspace_elements=16 * 1024 * 1024,
):
    """Pack ``[B,Q,K]`` unique block IDs into an 8x smaller bit route.

    Top-k is expected to produce unique non-negative IDs per row.  Repeated
    IDs would make scatter-add carry into a neighbouring bit and are rejected
    by the public wrapper before this helper is called.
    """
    batch, query_count, _ = selected_block_indices.shape
    # CuTe's DLPack bridge cannot infer a unique leading dimension for a
    # rank-3 auxiliary tensor whose final extent is 1 (both trailing strides
    # then equal one).  One unused padding byte keeps short smoke/decode-like
    # sequences on the same code path without changing any route bit.
    route_byte_count = max(2, math.ceil(raw_block_count / 8))
    packed = torch.zeros(
        (batch, query_count, route_byte_count),
        dtype=torch.uint8,
        device=selected_block_indices.device,
    )
    query_chunk = max(
        1,
        max_workspace_elements // max(batch * route_byte_count, 1),
    )
    for query_start in range(0, query_count, query_chunk):
        query_end = min(query_start + query_chunk, query_count)
        indices = selected_block_indices[:, query_start:query_end]
        valid = indices >= 0
        safe_indices = indices.clamp_min(0).to(torch.long)
        byte_indices = safe_indices // 8
        bit_values = (
            torch.ones_like(safe_indices, dtype=torch.int32)
            << (safe_indices % 8).to(torch.int32)
        )
        bit_values = bit_values * valid.to(torch.int32)
        workspace = torch.zeros(
            (batch, query_end - query_start, route_byte_count),
            dtype=torch.int32,
            device=indices.device,
        )
        workspace.scatter_add_(2, byte_indices, bit_values)
        packed[:, query_start:query_end] = workspace.to(torch.uint8)
    return packed.contiguous()


def _check_unique_summary_block_indices(selected_block_indices):
    """Reject duplicate valid IDs because the packed route uses scatter-add."""
    if selected_block_indices.shape[-1] <= 1:
        return
    sorted_indices = selected_block_indices.sort(dim=-1).values
    duplicate = (
        (sorted_indices[..., 1:] == sorted_indices[..., :-1])
        & (sorted_indices[..., 1:] >= 0)
    )
    if bool(duplicate.any().item()):
        raise ValueError("valid selected block IDs must be unique within each query")


def _summary_block_sparse_tensors(
    selected_block_indices,
    *,
    text_length,
    block_size_tokens,
    local_window_tokens,
    local_alignment_tokens=None,
    query_tile_size=128,
    forward_kv_tile_size=64,
    backward_kv_tile_size=64,
    packed_route=None,
):
    """Build conservative CuTe schedules from compact block IDs.

    Forward uses a narrower 64-token KV execution tile.  This retains the
    cross-query reuse of a 128-token Q tile while reducing over-computation
    when only some of the eight 8-token raw blocks in a 64-token KV tile are
    selected.  The model-local FlashAttention wheel supports a 64-token SM90
    backward KV tile, so the same schedule can be transposed directly instead
    of conservatively merging adjacent forward tiles.  Passing 128 remains
    available for numerical and performance A/B tests.

    Remote tiles are always marked partial because every query row may have a
    different route.  The exact per-query mask is still applied inside both
    kernels; these schedules only decide which execution tiles are launched.
    """
    from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch

    if query_tile_size != 128:
        raise ValueError("SM90 S-KSA schedules currently require 128-token Q tiles")
    if forward_kv_tile_size != 64:
        raise ValueError("SM90 S-KSA forward currently requires 64-token KV tiles")
    if backward_kv_tile_size not in (64, 128):
        raise ValueError("SM90 S-KSA backward requires a 64- or 128-token KV tile")
    if backward_kv_tile_size % forward_kv_tile_size != 0:
        raise ValueError("backward KV tile must be a multiple of forward KV tile")
    if local_alignment_tokens is None:
        local_alignment_tokens = block_size_tokens

    batch, route_length, _ = selected_block_indices.shape
    if route_length != text_length:
        raise ValueError("one selected-block row is required per text query")
    q_tile_count = math.ceil(text_length / query_tile_size)
    forward_kv_tile_count = math.ceil(text_length / forward_kv_tile_size)
    backward_kv_tile_count = math.ceil(text_length / backward_kv_tile_size)
    padded_query_count = q_tile_count * query_tile_size

    if packed_route is not None:
        # Reuse the already-built bitmap instead of traversing
        # [B,Q,TopK] again.  An 8-token route uses one byte per 64-token KV
        # tile; a token route uses eight bytes per tile.
        if forward_kv_tile_size % (block_size_tokens * 8):
            raise ValueError(
                "a forward KV tile must contain a whole number of route bytes"
            )
        route_bytes_per_kv_tile = (
            forward_kv_tile_size // (block_size_tokens * 8)
        )
        required_route_bytes = (
            forward_kv_tile_count * route_bytes_per_kv_tile
        )
        route_byte_count = math.ceil(math.ceil(text_length / block_size_tokens) / 8)
        if (
            packed_route.ndim != 3
            or packed_route.shape[0] != batch
            or packed_route.shape[1] != text_length
            or packed_route.shape[2] < route_byte_count
            or packed_route.dtype != torch.uint8
            or packed_route.device != selected_block_indices.device
        ):
            raise ValueError(
                "packed S-KSA route must be uint8 [batch, query, route_byte]"
            )
        if padded_query_count != text_length:
            packed_padding = packed_route.new_zeros(
                (
                    batch,
                    padded_query_count - text_length,
                    packed_route.shape[2],
                )
            )
            packed_padded = torch.cat((packed_route, packed_padding), dim=1)
        else:
            packed_padded = packed_route
        packed_padded = packed_padded[..., :required_route_bytes]
        if packed_padded.shape[-1] < required_route_bytes:
            packed_padded = torch.nn.functional.pad(
                packed_padded, (0, required_route_bytes - packed_padded.shape[-1])
            )
        allowed = packed_padded[..., :required_route_bytes].view(
            batch,
            q_tile_count,
            query_tile_size,
            forward_kv_tile_count,
            route_bytes_per_kv_tile,
        ).any(dim=(2, 4)).to(torch.bool).contiguous()
    else:
        # Compatibility/oracle path for callers that only have selected IDs.
        if padded_query_count != text_length:
            padding = selected_block_indices.new_full(
                (
                    batch,
                    padded_query_count - text_length,
                    selected_block_indices.shape[-1],
                ),
                -1,
            )
            selected_padded = torch.cat(
                (selected_block_indices, padding), dim=1
            )
        else:
            selected_padded = selected_block_indices
        allowed_with_sentinel = torch.zeros(
            (batch, q_tile_count, forward_kv_tile_count + 1),
            dtype=torch.bool,
            device=selected_block_indices.device,
        )
        # Chunk the int64 scatter indices so 128K x large-topk routes do not
        # create another sequence-sized temporary tensor.
        tile_chunk = 8
        for q_tile_start in range(0, q_tile_count, tile_chunk):
            q_tile_end = min(q_tile_start + tile_chunk, q_tile_count)
            query_start = q_tile_start * query_tile_size
            query_end = q_tile_end * query_tile_size
            logical = selected_padded[:, query_start:query_end]
            valid = logical >= 0
            physical = torch.where(
                valid,
                (logical * block_size_tokens) // forward_kv_tile_size,
                logical.new_full((), forward_kv_tile_count),
            )
            physical = physical.reshape(
                batch, q_tile_end - q_tile_start, -1
            )
            allowed_with_sentinel[:, q_tile_start:q_tile_end].scatter_(
                2,
                physical.to(torch.long),
                True,
            )
        allowed = allowed_with_sentinel[
            ..., :forward_kv_tile_count
        ].contiguous()

    q_tiles = torch.arange(q_tile_count, device=allowed.device)
    kv_tiles = torch.arange(forward_kv_tile_count, device=allowed.device)
    q_first_token = q_tiles * query_tile_size
    q_last_token = ((q_tiles + 1) * query_tile_size - 1).clamp_max(
        text_length - 1
    )
    kv_first_token = kv_tiles * forward_kv_tile_size
    kv_last_token = ((kv_tiles + 1) * forward_kv_tile_size - 1).clamp_max(
        text_length - 1
    )

    # Schedule every tile with at least one causal-local (q, k) pair.  Each
    # query's local start is clamped at zero and aligned down to a raw summary
    # block boundary.  The token mask still clips the two boundary tiles
    # exactly.
    q_first_local_start = (
        (q_first_token - local_window_tokens + 1).clamp_min(0)
        // local_alignment_tokens
    ) * local_alignment_tokens
    local_overlap = (
        (kv_first_token[None, :] <= q_last_token[:, None])
        & (kv_last_token[None, :] >= q_first_local_start[:, None])
    )
    causal_overlap = kv_first_token[None, :] <= q_last_token[:, None]
    allowed |= local_overlap.unsqueeze(0)
    allowed &= causal_overlap.unsqueeze(0)

    # A tile can skip mask_mod only when every actual q/k pair is causal and
    # inside the local window.  Selected remote tiles always stay partial.
    q_last_local_start = (
        (q_last_token - local_window_tokens + 1).clamp_min(0)
        // local_alignment_tokens
    ) * local_alignment_tokens
    local_full = (
        (kv_last_token[None, :] <= q_first_token[:, None])
        & (kv_first_token[None, :] >= q_last_local_start[:, None])
    )
    full_allowed = allowed & local_full.unsqueeze(0)
    partial_allowed = allowed & ~full_allowed

    def compact(partial, full, *, block_size):
        counts = partial.sum(dim=-1).to(torch.int32).unsqueeze(1)
        indices = torch.argsort(
            partial.to(torch.int8), dim=-1, descending=True, stable=True
        ).to(torch.int32).unsqueeze(1)
        full_counts = full.sum(dim=-1).to(torch.int32).unsqueeze(1)
        full_indices = torch.argsort(
            full.to(torch.int8), dim=-1, descending=True, stable=True
        ).to(torch.int32).unsqueeze(1)
        return BlockSparseTensorsTorch(
            mask_block_cnt=counts,
            mask_block_idx=indices,
            full_block_cnt=full_counts,
            full_block_idx=full_indices,
            block_size=block_size,
        )

    forward_sparse = compact(
        partial_allowed,
        full_allowed,
        block_size=(query_tile_size, forward_kv_tile_size),
    )

    # Backward consumes the KV-major transpose of the forward Q-major schedule.
    # With the default 64-token tile this is a direct transpose.  The generic
    # merge below also keeps the old 128-token schedule available for A/B:
    # a merged tile is full only if both halves are full.
    return forward_sparse, None


class SummaryBlockSparseAttentionFunc(torch.autograd.Function):
    """Exact per-query Full attention over selected raw-token route units."""

    @staticmethod
    def forward(
        ctx,
        q,
        k,
        v,
        packed_route,
        selected_block_indices,
        block_size_tokens,
        local_window_tokens,
        local_alignment_tokens,
        softmax_scale=None,
        deterministic=False,
        cached_output=None,
        cached_lse=None,
    ):
        from flash_attn.cute.interface import _flash_attn_fwd
        from .mask import summary_block_sparse_attn_mask_func

        fwd_sparse, bwd_sparse = _summary_block_sparse_tensors(
            selected_block_indices,
            text_length=q.shape[1],
            block_size_tokens=block_size_tokens,
            local_window_tokens=local_window_tokens,
            local_alignment_tokens=local_alignment_tokens,
            packed_route=packed_route,
        )
        mask_mod = partial(
            summary_block_sparse_attn_mask_func,
            block_size_tokens=block_size_tokens,
            local_window_tokens=local_window_tokens,
            local_alignment_tokens=local_alignment_tokens,
        )
        if (cached_output is None) != (cached_lse is None):
            raise ValueError("cached sparse output and LSE must be provided together")
        if cached_output is None:
            out, lse = _flash_attn_fwd(
                q,
                k,
                v,
                softmax_scale=softmax_scale,
                causal=False,
                pack_gqa=None,
                mask_mod=mask_mod,
                block_sparse_tensors=fwd_sparse,
                aux_tensors=[packed_route],
                m_block_size=128,
                n_block_size=64,
                return_lse=any(t.requires_grad for t in (q, k, v)),
            )
        else:
            expected_lse_shape = (q.shape[0], q.shape[2], q.shape[1])
            if (
                cached_output.shape != q.shape
                or cached_output.dtype != q.dtype
                or cached_output.device != q.device
                or cached_lse.shape != expected_lse_shape
                or cached_lse.dtype != torch.float32
                or cached_lse.device != q.device
            ):
                raise ValueError("cached sparse attention tensor metadata mismatch")
            # The enclosing non-reentrant layer checkpoint only needs to
            # rebuild tensors saved for this Function's backward. Reuse the
            # detached original-forward values instead of executing the same
            # sparse forward kernel a second time.
            out, lse = cached_output, cached_lse
        if lse is not None:
            ctx.mark_non_differentiable(lse)
        return out, lse





# ---------------------------------------------------------------------------
# FlexAttention path (Ampere / any GPU with PyTorch >= 2.5)
# ---------------------------------------------------------------------------

_FLEX_BLOCK_MASK_CACHE: Dict[Tuple, object] = {}


def _create_block_mask_no_materialize(mask_mod, seq_len_q, seq_len_k, device, BLOCK_SIZE=128):
    """Create BlockMask without materializing the full (seq_len_q, seq_len_k) dense mask.

    Processes one q_block row at a time, reducing peak memory from
    O(seq_q * seq_k) to O(BLOCK_SIZE * seq_k).
    """
    from torch.nn.attention.flex_attention import BlockMask

    num_q_blocks = (seq_len_q + BLOCK_SIZE - 1) // BLOCK_SIZE
    num_kv_blocks = (seq_len_k + BLOCK_SIZE - 1) // BLOCK_SIZE

    # Block-level sparsity: (num_q_blocks, num_kv_blocks)
    any_mask = torch.zeros(num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)
    full_mask = torch.zeros(num_q_blocks, num_kv_blocks, dtype=torch.bool, device=device)

    kv_arange = torch.arange(seq_len_k, device=device)
    pad_k = (-seq_len_k) % BLOCK_SIZE
    last_kv_len = seq_len_k - (num_kv_blocks - 1) * BLOCK_SIZE if seq_len_k % BLOCK_SIZE else BLOCK_SIZE

    for q_blk in range(num_q_blocks):
        q_start = q_blk * BLOCK_SIZE
        q_end = min(q_start + BLOCK_SIZE, seq_len_q)
        q_len = q_end - q_start
        q_idx = torch.arange(q_start, q_end, device=device)

        # (q_len, seq_len_k) bool
        mask = mask_mod(0, 0, q_idx[:, None], kv_arange[None, :])

        if pad_k > 0:
            mask = torch.nn.functional.pad(mask, (0, pad_k), value=False)

        # (q_len, num_kv_blocks, BLOCK_SIZE) -> sum per block
        block_sums = mask.reshape(q_len, num_kv_blocks, BLOCK_SIZE).sum(dim=(0, 2))

        any_mask[q_blk] = block_sums > 0
        full_mask[q_blk] = block_sums == (q_len * BLOCK_SIZE)
        # Last kv block may be shorter
        if pad_k > 0:
            full_mask[q_blk, -1] = block_sums[-1] == (q_len * last_kv_len)

    def _bool_mask_to_compact(mask_2d):
        """Convert (rows, cols) bool mask to (num_per_row, indices) format."""
        num_per_row = mask_2d.sum(dim=1).to(torch.int32)
        max_num = max(num_per_row.max().item(), 1)
        sorted_idx = mask_2d.to(torch.int32).argsort(dim=1, descending=True, stable=True)
        indices = sorted_idx[:, :max_num].to(torch.int32)
        return num_per_row, indices

    kv_num_blocks, kv_indices = _bool_mask_to_compact(any_mask)
    full_kv_num_blocks, full_kv_indices = _bool_mask_to_compact(full_mask)
    q_num_blocks, q_indices = _bool_mask_to_compact(any_mask.t().contiguous())
    full_q_num_blocks, full_q_indices = _bool_mask_to_compact(full_mask.t().contiguous())

    return BlockMask(
        kv_num_blocks=kv_num_blocks[None, None],
        kv_indices=kv_indices[None, None],
        full_kv_num_blocks=full_kv_num_blocks[None, None],
        full_kv_indices=full_kv_indices[None, None],
        q_num_blocks=q_num_blocks[None, None],
        q_indices=q_indices[None, None],
        full_q_num_blocks=full_q_num_blocks[None, None],
        full_q_indices=full_q_indices[None, None],
        BLOCK_SIZE=(BLOCK_SIZE, BLOCK_SIZE),
        mask_mod=mask_mod,
        seq_lengths=(seq_len_q, seq_len_k),
    )


def _get_flex_block_mask(seq_len_q, seq_len_k, chunk_size, summary_num, sliding_chunk_num,
                          skip_old_summary, device):
    cache_key = (seq_len_q, seq_len_k, chunk_size, summary_num, sliding_chunk_num, skip_old_summary)
    if cache_key not in _FLEX_BLOCK_MASK_CACHE:
        block = chunk_size + summary_num
        prefix_summary_length = seq_len_k - seq_len_q

        if prefix_summary_length > 0:
            def mask_mod(b, h, q_idx, kv_idx):
                q_chunk_id = q_idx // block
                kv_chunk_id = (kv_idx - prefix_summary_length) // block
                is_sliding = (q_chunk_id - kv_chunk_id) <= sliding_chunk_num
                is_q_summary = (q_idx % block) >= chunk_size
                is_kv_summary = ((kv_idx - prefix_summary_length) % block) >= chunk_size
                is_causal = q_idx >= (kv_idx - prefix_summary_length)
                valid_kv_q_summary = is_q_summary & (kv_chunk_id == q_chunk_id) & is_causal
                valid_sliding_q_text = is_sliding & (~is_kv_summary) & (kv_idx >= prefix_summary_length)
                if skip_old_summary:
                    valid_kv_q_text = (~is_q_summary) & is_causal & valid_sliding_q_text
                else:
                    valid_summary_q_text = ((kv_idx // summary_num) + sliding_chunk_num) < q_chunk_id
                    valid_kv_q_text = (~is_q_summary) & ((is_causal & valid_sliding_q_text) | valid_summary_q_text)
                return valid_kv_q_summary | valid_kv_q_text
        else:
            def mask_mod(b, h, q_idx, kv_idx):
                q_chunk_id = q_idx // block
                kv_chunk_id = kv_idx // block
                is_sliding = (q_chunk_id - kv_chunk_id) <= sliding_chunk_num
                is_q_summary = (q_idx % block) >= chunk_size
                is_kv_summary = (kv_idx % block) >= chunk_size
                is_causal = q_idx >= kv_idx
                valid_kv_q_summary = is_q_summary & (kv_chunk_id == q_chunk_id) & is_causal
                valid_sliding_q_text = is_sliding & (~is_kv_summary)
                if skip_old_summary:
                    valid_kv_q_text = (~is_q_summary) & is_causal & valid_sliding_q_text
                else:
                    valid_summary_q_text = is_kv_summary & ((q_chunk_id - kv_chunk_id) > sliding_chunk_num)
                    valid_kv_q_text = (~is_q_summary) & is_causal & (valid_sliding_q_text | valid_summary_q_text)
                return valid_kv_q_summary | valid_kv_q_text

        block_mask = _create_block_mask_no_materialize(
            mask_mod, seq_len_q, seq_len_k, device,
        )
        _FLEX_BLOCK_MASK_CACHE[cache_key] = block_mask
    return _FLEX_BLOCK_MASK_CACHE[cache_key]


_compiled_flex_attention = None


def _get_compiled_flex_attention():
    global _compiled_flex_attention
    if _compiled_flex_attention is None:
        from torch.nn.attention.flex_attention import flex_attention
        _compiled_flex_attention = torch.compile(flex_attention)
    return _compiled_flex_attention


def _summary_attn_flex(q, k, v, chunk_size, summary_num, sliding_chunk_num, skip_old_summary):
    """FlexAttention-based implementation. Auto-compiles to efficient Triton block-sparse kernel.

    Works on any GPU with PyTorch >= 2.5 (Ampere, Hopper, etc).
    No KV duplication, no Python loops — single fused kernel call.
    """
    b, seq_len_q, h_q, d = q.shape
    _, seq_len_k, h_kv, _ = k.shape
    enable_gqa = (h_q != h_kv)

    block_mask = _get_flex_block_mask(seq_len_q, seq_len_k, chunk_size, summary_num,
                                       sliding_chunk_num, skip_old_summary, q.device)

    # flex_attention expects [B, H, S, D]
    q_flex = q.transpose(1, 2)  # [B, H, S, D]
    k_flex = k.transpose(1, 2)
    v_flex = v.transpose(1, 2)

    compiled_fn = _get_compiled_flex_attention()
    out = compiled_fn(q_flex, k_flex, v_flex, block_mask=block_mask, enable_gqa=enable_gqa)
    out = out.transpose(1, 2)  # [B, S, H, D]

    return out, None


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def summary_attn_func(
    q, k, v, chunk_size, summary_num, sliding_chunk_num,
    summary_pos: Optional[torch.Tensor] = None,
    skip_old_summary: bool = False,
    force_return_lse: bool = False,
):
    if summary_pos is not None:
        k = torch.cat((k[:, summary_pos, :], k), dim=1)
        v = torch.cat((v[:, summary_pos, :], v), dim=1)

    if _check_cute_available():
        out = SummaryAttnFunc.apply(
            q, k, v, chunk_size, summary_num, sliding_chunk_num,
            None, False, None, False, skip_old_summary, force_return_lse,
        )
    else:
        out = _summary_attn_flex(
            q, k, v, chunk_size, summary_num, sliding_chunk_num, skip_old_summary,
        )
    return out


def hierarchical_summary_attn_func(
    q, k, v, meta_k, meta_v,
    chunk_size, sliding_chunk_num, group_size,
    softmax_scale=None, mass_correction=True, force_return_lse=False,
):
    """Compressed-KSA attention; tensors use [batch, sequence, heads, dim]."""
    tensors = (q, k, v, meta_k, meta_v)
    if not all(t.is_cuda for t in tensors):
        raise RuntimeError("hierarchical_summary_attn_func requires CUDA tensors")
    if any(t.ndim != 4 for t in tensors):
        raise ValueError("q/k/v/meta_k/meta_v must be rank-4 tensors")
    if k.shape != v.shape or meta_k.shape != meta_v.shape:
        raise ValueError("key/value tensor shapes must match")
    if q.shape[0] != k.shape[0] or q.shape[0] != meta_k.shape[0]:
        raise ValueError("batch sizes must match")
    if q.shape[1] != k.shape[1] or k.shape[2:] != meta_k.shape[2:]:
        raise ValueError("base and meta K/V layouts are incompatible")
    if q.shape[2] % k.shape[2]:
        raise ValueError("query heads must be divisible by KV heads")
    if group_size <= 1 or chunk_size <= 0 or sliding_chunk_num <= 0:
        raise ValueError("group_size must exceed one and chunk/window sizes must be positive")
    with torch.cuda.device(q.device):
        if not _check_cute_available():
            raise RuntimeError("hierarchical_summary_attn_func requires CuTe on sm90+")
        return HierarchicalSummaryAttnFunc.apply(
            q, k, v, meta_k, meta_v,
            chunk_size, sliding_chunk_num, group_size,
            softmax_scale, False, mass_correction, force_return_lse,
        )








def _pad_summary_block_sparse_cute_inputs(
    q,
    k,
    v,
    selected_block_indices,
    *,
    query_tile_size=128,
):
    """Pad CuTe's physical Q tile without changing logical attention.

    The block-sparse CuTe consumer may evaluate ``mask_mod`` for every
    physical row in its final Q tile before applying the sequence-length
    predicate to later partial KV blocks.  Supplying real all-invalid route
    rows for those physical lanes keeps the auxiliary lookup in bounds.
    Future zero K/V rows remain invisible to every real query through the
    kernel's causal guard, and the padded query outputs are discarded by the
    public wrapper.
    """
    text_length = q.shape[1]
    padding = (-text_length) % query_tile_size
    if not padding:
        return q, k, v, selected_block_indices, text_length

    def pad_sequence(tensor):
        return torch.cat(
            (
                tensor,
                tensor.new_zeros(
                    tensor.shape[0], padding, *tensor.shape[2:]
                ),
            ),
            dim=1,
        )

    padded_route = torch.cat(
        (
            selected_block_indices,
            selected_block_indices.new_full(
                (
                    selected_block_indices.shape[0],
                    padding,
                    selected_block_indices.shape[2],
                ),
                -1,
            ),
        ),
        dim=1,
    )
    return (
        pad_sequence(q),
        pad_sequence(k),
        pad_sequence(v),
        padded_route,
        text_length,
    )


def summary_block_sparse_attn_func(
    q,
    k,
    v,
    selected_block_indices,
    block_size_tokens,
    local_window_tokens,
    softmax_scale=None,
    deterministic=False,
):
    """Run exact per-query Full attention over selected raw-token blocks.

    ``selected_block_indices[b, q]`` contains the summary/raw-block IDs chosen
    by query ``q``; ``-1`` pads rows with fewer candidates.  A selected ID
    exposes the corresponding raw K/V block, never the summary K/V itself.
    Local causal tokens are added independently.  The caller is responsible
    for reserving block 0 inside the TopK budget when an attention sink is
    desired; this kernel always applies a final causal guard.

    The selected IDs must be unique within each row.  This is naturally true
    for ``torch.topk``.  Set ``SKSA_VALIDATE_BLOCK_INDICES=1`` to perform the
    synchronising bounds/duplicate checks during diagnostics.
    """
    _validate_summary_block_sparse_inputs(
        q,
        k,
        v,
        selected_block_indices,
        block_size_tokens,
        local_window_tokens,
        block_size_tokens,
    )
    if os.environ.get("SKSA_VALIDATE_BLOCK_INDICES", "0") == "1":
        _check_unique_summary_block_indices(selected_block_indices)
    route_indices = selected_block_indices.to(torch.int32).contiguous()
    with torch.cuda.device(q.device):
        if not _check_cute_available():
            raise RuntimeError(
                "summary_block_sparse_attn_func requires CuTe on sm90+"
            )
        # Qwen's ordinary multi-head raw K/V is much wider than DeepSeek MLA's
        # shared latent KV.  Direct per-query gather therefore loses crucial
        # cross-query KV reuse and is bandwidth-bound.  Keep the exact CuTe
        # tiled implementation as production default; ``direct`` is an
        # explicit diagnostic prototype only.
        q, k, v, route_indices, text_length = (
            _pad_summary_block_sparse_cute_inputs(
                q,
                k,
                v,
                route_indices,
            )
        )
        raw_block_count = math.ceil(k.shape[1] / block_size_tokens)
        packed_route = _pack_summary_block_indices(
            route_indices,
            raw_block_count=raw_block_count,
        )
        output, _ = SummaryBlockSparseAttentionFunc.apply(
            q,
            k,
            v,
            packed_route,
            route_indices,
            block_size_tokens,
            local_window_tokens,
            block_size_tokens,
            softmax_scale,
            deterministic,
            None,
            None,
        )
        return (
            output
            if output.shape[1] == text_length
            else output[:, :text_length].contiguous()
        )