feat(attention): fixed budget block selector and token layout

The selector mean-pools query and key tiles, scores the tile pairs and
keeps the highest scoring key tiles per query tile within a budget
expressed as a fraction of the sparsifiable candidates. No scale and no
softmax, since top-k is invariant under both. Scoring runs on query
heads so grouped attention needs no key expansion, the diagonal is
always kept so no query row is left empty, and a budget that covers
every candidate reports dense instead of building a full mask.

The layout says what a packed sequence holds. Only video and image
tokens are sparsifiable; text, conditioning, audio and anything
unrecognized pin their rows and columns dense, padding is dropped, and a
tile straddling a boundary pins. Layouts come from the *_indices tensors
a pipeline passes its transformer by name, from ordered segments where a
transformer packs the sequence itself, or from a leading prefix as a
fallback.

The flex consumer builds a BlockMask with every selected tile in the
full slots, so mask_mod is never invoked and no dense mask is
materialized, and calls flex_attention compiled: called eagerly it reads
mask_mod rather than the block lists, so a block only mask attends
densely and silently. test/test-attention-sparse.py covers this with a
row that fails if the selection stops changing the output, alongside
tile equivalence against sdpa fed the same tiles, measured against the
flex kernel floor rather than an absolute tolerance.

modules/attention/sparse/__init__.py

@@ -0,0 +1,8 @@
+"""Block-sparse attention: the selector, the token layout it respects, and the consumers that apply it."""
+from modules.attention.sparse.selector import BlockSelection, SparseSpec, block_count, radial_blocks, schedule, select_blocks
+from modules.attention.sparse.layout import Span, TokenLayout, block_pins, layout_from_index_kwargs, layout_from_prefix, layout_from_segments
+
+__all__ = [
+    'BlockSelection', 'SparseSpec', 'block_count', 'radial_blocks', 'schedule', 'select_blocks',
+    'Span', 'TokenLayout', 'block_pins', 'layout_from_index_kwargs', 'layout_from_prefix', 'layout_from_segments',
+]

modules/attention/sparse/flex.py

@@ -0,0 +1,37 @@
+"""Turn a BlockSelection into the BlockMask FlexAttention consumes, and call it so the mask is honored."""
+import torch
+from torch.nn.attention.flex_attention import BlockMask, flex_attention, _dense_to_ordered
+from modules.attention.sparse.selector import BlockSelection
+
+
+compiled_flex = None
+
+
+def to_block_mask(selection: BlockSelection, device=None) -> BlockMask:
+    """All selected tiles go in the full slots, so mask_mod is never invoked and no dense S squared mask is built."""
+    keep = selection.keep
+    if device is not None and keep.device != device:
+        keep = keep.to(device)
+    if keep.dim() != 4:
+        raise ValueError(f'block selection must be 4d, got {tuple(keep.shape)}')
+    empty_num, empty_indices = _dense_to_ordered(torch.zeros_like(keep))
+    full_num, full_indices = _dense_to_ordered(keep)
+    return BlockMask.from_kv_blocks(
+        empty_num, empty_indices,
+        full_kv_num_blocks=full_num, full_kv_indices=full_indices,
+        BLOCK_SIZE=(selection.block_q, selection.block_kv),
+        seq_lengths=(selection.seq_q, selection.seq_kv), # exact lengths, so a ragged tail is handled rather than rounded up
+        compute_q_blocks=False, # backward only metadata, and inference never reads it
+    )
+
+
+def flex_call():
+    """flex_attention reads the block lists only when compiled; called eagerly it evaluates mask_mod instead and a block-only mask is silently dense."""
+    global compiled_flex # pylint: disable=global-statement
+    if compiled_flex is None:
+        compiled_flex = torch.compile(flex_attention, dynamic=False)
+    return compiled_flex
+
+
+def attend(query, key, value, selection: BlockSelection, scale=None, enable_gqa=False):
+    return flex_call()(query, key, value, block_mask=to_block_mask(selection, device=query.device), scale=scale, enable_gqa=enable_gqa)

modules/attention/sparse/layout.py

@@ -0,0 +1,132 @@
+"""What each token in a packed sequence is, so the selector knows what it may sparsify."""
+from dataclasses import dataclass
+import torch
+
+
+# only the bulk modalities are sparsifiable; everything else is pinned dense, and an unrecognized kind pins too
+SPARSIFIABLE = frozenset({'video', 'image'})
+DROPPED = frozenset({'pad'})
+
+
+@dataclass(frozen=True)
+class Span:
+    kind: str
+    start: int
+    end: int
+
+
+@dataclass(frozen=True)
+class TokenLayout:
+    """Ordered spans covering one packed sequence."""
+    spans: tuple[Span, ...]
+    length: int
+    source: str = 'unknown' # how the layout was obtained, for the log
+
+    def key(self) -> tuple:
+        return (self.length, self.source, tuple((s.kind, s.start, s.end) for s in self.spans))
+
+    def kinds(self) -> tuple[str, ...]:
+        return tuple(dict.fromkeys(s.kind for s in self.spans))
+
+    def sparsifiable_tokens(self) -> int:
+        return sum(s.end - s.start for s in self.spans if s.kind in SPARSIFIABLE)
+
+    def token_flags(self, device) -> tuple[torch.Tensor, torch.Tensor]:
+        """Per token: may this be sparsified, and is it padding."""
+        sparse = torch.zeros(self.length, dtype=torch.bool, device=device)
+        pad = torch.zeros(self.length, dtype=torch.bool, device=device)
+        for span in self.spans:
+            if span.kind in SPARSIFIABLE:
+                sparse[span.start:span.end] = True
+            elif span.kind in DROPPED:
+                pad[span.start:span.end] = True
+        return sparse, pad
+
+
+def runs(indices: torch.Tensor) -> list[tuple[int, int]]:
+    """Contiguous [start, end) runs in a sorted 1d index tensor."""
+    if indices.numel() == 0:
+        return []
+    values = indices.detach().to('cpu', torch.int64).sort().values
+    breaks = (values[1:] - values[:-1] != 1).nonzero().flatten().tolist()
+    bounds = [0, *[b + 1 for b in breaks], values.numel()]
+    return [(int(values[bounds[i]].item()), int(values[bounds[i + 1] - 1].item()) + 1) for i in range(len(bounds) - 1)]
+
+
+def layout_from_index_kwargs(kwargs: dict, length: int) -> TokenLayout | None:
+    """Read a layout off the *_indices tensors a pipeline passes its transformer by name."""
+    spans: list[Span] = []
+    for name, value in kwargs.items():
+        if not name.endswith('_indices') or not torch.is_tensor(value) or value.dim() != 1 or value.is_floating_point():
+            continue
+        kind = name[:-len('_indices')].lower()
+        found = runs(value)
+        for position, (start, end) in enumerate(found):
+            # a video run that is not the last one is keyframe conditioning, which stays dense
+            resolved = 'cond' if (kind == 'video' and position < len(found) - 1) else kind
+            spans.append(Span(kind=resolved, start=start, end=end))
+    if not spans:
+        return None
+    spans.sort(key=lambda s: s.start)
+    return TokenLayout(spans=tuple(spans), length=length, source='indices')
+
+
+def layout_from_segments(segments, length: int | None = None, source: str = 'segments') -> TokenLayout:
+    """Build a layout from ordered (kind, count) pairs, the form a transformer knows at its packing site."""
+    spans: list[Span] = []
+    cursor = 0
+    for kind, count in segments:
+        if count <= 0:
+            continue
+        spans.append(Span(kind=kind, start=cursor, end=cursor + count))
+        cursor += count
+    return TokenLayout(spans=tuple(spans), length=length if length is not None else cursor, source=source)
+
+
+def layout_from_prefix(length: int, prefix: int) -> TokenLayout:
+    """Fallback when nothing published a layout: treat a leading run as conditioning and sparsify the rest."""
+    return layout_from_segments([('text', prefix), ('image', length - prefix)], length=length, source='prefix')
+
+
+def block_flags(flags: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per block: do all tokens carry the flag, does any token carry it."""
+    seq = flags.shape[0]
+    whole = (seq // block) * block
+    parts_all, parts_any = [], []
+    if whole:
+        view = flags[:whole].view(whole // block, block)
+        parts_all.append(view.all(dim=-1))
+        parts_any.append(view.any(dim=-1))
+    if whole < seq:
+        parts_all.append(flags[whole:].all(dim=-1, keepdim=True))
+        parts_any.append(flags[whole:].any(dim=-1, keepdim=True))
+    def join(parts):
+        return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)
+    return join(parts_all), join(parts_any)
+
+
+pin_cache: dict = {}
+
+
+def block_pins(layout: TokenLayout, seq_q: int, seq_kv: int, block_q: int, block_kv: int, device) -> tuple[torch.Tensor, torch.Tensor]:
+    """Tiles that must stay dense and tiles that can be skipped outright, as (1, 1, NQ, NK) masks."""
+    cache_key = (layout.key(), seq_q, seq_kv, block_q, block_kv, str(device))
+    hit = pin_cache.get(cache_key)
+    if hit is not None:
+        return hit
+    sparse_tokens, pad_tokens = layout.token_flags(device)
+    q_sparse = sparse_tokens[:seq_q] if layout.length >= seq_q else torch.nn.functional.pad(sparse_tokens, (0, seq_q - layout.length))
+    kv_sparse = sparse_tokens[:seq_kv] if layout.length >= seq_kv else torch.nn.functional.pad(sparse_tokens, (0, seq_kv - layout.length))
+    kv_pad = pad_tokens[:seq_kv] if layout.length >= seq_kv else torch.nn.functional.pad(pad_tokens, (0, seq_kv - layout.length))
+    q_all_sparse, _ = block_flags(q_sparse, block_q)
+    kv_all_sparse, _ = block_flags(kv_sparse, block_kv)
+    kv_all_pad, _ = block_flags(kv_pad, block_kv)
+    # a tile is pinned when its query tile or its key tile carries anything that is not sparsifiable, boundary tiles included
+    pins = (~q_all_sparse).unsqueeze(-1) | (~kv_all_sparse).unsqueeze(0)
+    drops = kv_all_pad.unsqueeze(0).expand_as(pins)
+    pins = (pins & ~drops).unsqueeze(0).unsqueeze(0).contiguous()
+    drops = drops.unsqueeze(0).unsqueeze(0).contiguous()
+    if len(pin_cache) > 32:
+        pin_cache.clear()
+    pin_cache[cache_key] = (pins, drops)
+    return pins, drops

modules/attention/sparse/selector.py

@@ -0,0 +1,129 @@
+"""Fixed-budget block selection: which KV tiles each query tile attends to."""
+from dataclasses import dataclass
+import math
+import torch
+
+
+@dataclass(frozen=True)
+class SparseSpec:
+    """How much to keep and at what granularity. Budget is a fraction of the sparsifiable candidates, pins are added on top."""
+    budget: float = 0.30
+    block_q: int = 128
+    block_kv: int = 64
+    head_shared: bool = False # score once for all heads, cheaper and coarser
+    force: bool = False # skip the dense short circuit, so tests can exercise the path at budget 1.0
+    score_chunk_bytes: int = 256 << 20
+
+
+@dataclass(frozen=True)
+class BlockSelection:
+    """int8 keep flags per (query tile, kv tile); the geometry every consumer reads."""
+    keep: torch.Tensor # (B, H, NQ, NK), H is the query head count or 1
+    block_q: int
+    block_kv: int
+    budget: float
+    seq_q: int
+    seq_kv: int
+    density: float
+
+    @property
+    def shape(self) -> tuple[int, int, int, int]:
+        return tuple(self.keep.shape)
+
+
+def block_count(length: int, block: int) -> int:
+    return (length + block - 1) // block
+
+
+def pool_blocks(x: torch.Tensor, block: int) -> torch.Tensor:
+    """Mean over each block of tokens, fp32, without materializing a padded copy."""
+    seq = x.shape[-2]
+    whole = (seq // block) * block
+    parts = []
+    if whole:
+        head = x[..., :whole, :]
+        parts.append(head.unflatten(-2, (whole // block, block)).mean(dim=-2, dtype=torch.float32))
+    if whole < seq:
+        parts.append(x[..., whole:, :].mean(dim=-2, dtype=torch.float32, keepdim=True))
+    return parts[0] if len(parts) == 1 else torch.cat(parts, dim=-2)
+
+
+def diagonal_blocks(nq: int, nk: int, block_q: int, block_kv: int, device) -> torch.Tensor:
+    """Tiles whose query and key token ranges overlap; keeping them removes the empty-row case."""
+    q_index = torch.arange(nq, device=device).unsqueeze(-1)
+    k_index = torch.arange(nk, device=device).unsqueeze(0)
+    return (q_index * block_q < (k_index + 1) * block_kv) & (k_index * block_kv < (q_index + 1) * block_q)
+
+
+def score_blocks(query: torch.Tensor, key: torch.Tensor, spec: SparseSpec) -> torch.Tensor:
+    """Mean-pooled query-key affinity per tile pair. No scale and no softmax: top-k is invariant under both."""
+    pooled_q = pool_blocks(query, spec.block_q) # (B, Hq, NQ, D)
+    pooled_k = pool_blocks(key, spec.block_kv) # (B, Hkv, NK, D)
+    heads_q, heads_kv = pooled_q.shape[1], pooled_k.shape[1]
+    if spec.head_shared:
+        pooled_q = pooled_q.mean(dim=1, keepdim=True)
+        pooled_k = pooled_k.mean(dim=1, keepdim=True)
+    elif heads_kv != heads_q: # gqa: score on query heads, the geometry both consumers expect
+        pooled_k = pooled_k.repeat_interleave(heads_q // heads_kv, dim=1)
+    heads = pooled_q.shape[1]
+    per_head = pooled_q.shape[2] * pooled_k.shape[2] * 4
+    chunk = max(1, min(heads, spec.score_chunk_bytes // max(per_head, 1)))
+    if chunk >= heads:
+        return pooled_q @ pooled_k.transpose(-1, -2)
+    return torch.cat([pooled_q[:, i:i + chunk] @ pooled_k[:, i:i + chunk].transpose(-1, -2) for i in range(0, heads, chunk)], dim=1)
+
+
+def select_blocks(query: torch.Tensor, key: torch.Tensor, spec: SparseSpec, pins: torch.Tensor | None = None, drops: torch.Tensor | None = None) -> BlockSelection | None:
+    """Keep the highest scoring KV tiles per query tile within the budget, plus pins and the diagonal. None means attend densely."""
+    seq_q, seq_kv = query.shape[-2], key.shape[-2]
+    nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
+    device = query.device
+    must = diagonal_blocks(nq, nk, spec.block_q, spec.block_kv, device).unsqueeze(0).unsqueeze(0)
+    if pins is not None:
+        must = must | pins
+    forbidden = drops if drops is not None else torch.zeros_like(must)
+    candidates = ~must & ~forbidden
+    per_row = candidates.sum(dim=-1, keepdim=True) # (.., NQ, 1)
+    keep_per_row = torch.ceil(per_row * spec.budget).to(torch.int64)
+    if not spec.force and bool((keep_per_row >= per_row).all()):
+        return None # the budget covers every candidate, so the mask would be dense
+
+    scores = score_blocks(query, key, spec)
+    scores = scores.masked_fill(~candidates.expand_as(scores), float('-inf'))
+    limit = int(keep_per_row.max().item())
+    keep = must.expand(scores.shape).clone()
+    if limit > 0:
+        order = scores.argsort(dim=-1, descending=True, stable=True)
+        rank = torch.empty_like(order)
+        rank.scatter_(-1, order, torch.arange(nk, device=device).expand_as(order))
+        keep |= (rank < keep_per_row) & candidates
+    keep &= ~forbidden
+    keep_int8 = keep.to(torch.int8)
+    density = float(keep_int8.sum().item()) / max(keep_int8.numel(), 1)
+    return BlockSelection(keep=keep_int8, block_q=spec.block_q, block_kv=spec.block_kv, budget=spec.budget, seq_q=seq_q, seq_kv=seq_kv, density=density)
+
+
+def radial_blocks(seq_q: int, seq_kv: int, density: float, spec: SparseSpec, device) -> BlockSelection:
+    """A band around the diagonal at the requested density: the static control the selector has to beat."""
+    nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
+    q_center = (torch.arange(nq, device=device).unsqueeze(-1) + 0.5) * spec.block_q
+    k_center = (torch.arange(nk, device=device).unsqueeze(0) + 0.5) * spec.block_kv
+    distance = (q_center - k_center).abs()
+    low, high = 0.0, float(max(seq_q, seq_kv))
+    for _ in range(40): # bisect the bandwidth, since the band width to density map has no closed form at the edges
+        mid = (low + high) / 2
+        if float((distance <= mid).to(torch.float32).mean().item()) < density:
+            low = mid
+        else:
+            high = mid
+    keep = (distance <= high).unsqueeze(0).unsqueeze(0).to(torch.int8)
+    return BlockSelection(keep=keep, block_q=spec.block_q, block_kv=spec.block_kv, budget=density, seq_q=seq_q, seq_kv=seq_kv, density=float(keep.sum().item()) / max(keep.numel(), 1))
+
+
+def schedule(steps: int, budget: float, bump: float = 0.0, bump_steps: int = 0) -> tuple[float, ...]:
+    """Per-step budgets, precomputed. At most two distinct values, so a compiled consumer sees at most two specializations."""
+    if bump <= 0 or bump_steps <= 0 or steps <= 0:
+        return tuple([budget] * max(steps, 0))
+    raised = min(1.0, budget + bump)
+    edge = min(bump_steps, math.ceil(steps / 2))
+    return tuple([raised if (i < edge or i >= steps - edge) else budget for i in range(steps)])