Hugging Face's logo

OpenOneRec
/
KSA-1.9B-base

Text Generation
Transformers
Safetensors
English
Chinese
qwen3
ksa
ksa-v1
custom_code
long-context
conversational
text-generation-inference
Model card Files
xet
 Community
KSA-1.9B-base / kernel /ksa_attn /interface.py
OpenOneRec's picture
OpenOneRec
Release KSA-1.9B-base weights and inference runtime
6c91dd5 verified
Raw History Blame Contribute Delete
40.4 kB
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()
)