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)])

test/test-attention-sparse.py

@@ -0,0 +1,389 @@
+#!/usr/bin/env python
+"""
+Offline unit tests for block-sparse attention in modules.attention.sparse.
+
+Covers:
+
+- block pooling, including the ragged tail, against a per-block reference
+- the diagonal invariant: every query tile keeps the key tiles its tokens overlap
+- budget semantics: density tracks the budget over the candidates, pins survive, drops never do
+- the dense short circuit, and the force flag that suppresses it for tests
+- determinism of the selection for identical inputs
+- layout reading: the *_indices form a pipeline passes by name, with a non-final video run
+  relabelled as conditioning, and the segment form a transformer knows at its packing site
+- pins and drops derived from a layout: pinned columns, dropped padding, pinned boundary tiles
+- the flex consumer: a full-keep selection through flex_attention reproduces dense sdpa, and a
+  selection with dropped tiles reproduces sdpa given the same tiles masked out
+- the density matched radial control and the step schedule
+
+The flex rows need a cuda device and compile the flex kernel; they skip on cpu.
+
+Usage:
+    python test/test-attention-sparse.py
+"""
+
+import os
+import sys
+
+import torch
+
+script_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, script_dir)
+os.chdir(script_dir)
+
+os.environ['SD_INSTALL_QUIET'] = '1'
+
+# Bootstrap cmd_args before any module that pulls in shared.py.
+import modules.cmd_args  # pylint: disable=wrong-import-position
+import installer  # pylint: disable=wrong-import-position
+orig_argv = sys.argv
+sys.argv = [sys.argv[0]]
+try:
+    modules.cmd_args.parse_args()
+finally:
+    sys.argv = orig_argv
+installer.add_args(modules.cmd_args.parser)
+modules.cmd_args.parsed, _ = modules.cmd_args.parser.parse_known_args([])
+
+stock_sdpa = torch.nn.functional.scaled_dot_product_attention # captured before shared installs the configured hijacks
+
+from modules.errors import log                                # pylint: disable=wrong-import-position
+from modules import shared                                    # pylint: disable=wrong-import-position,unused-import
+from modules.attention import sparse                          # pylint: disable=wrong-import-position
+from modules.attention.sparse import flex as sparse_flex      # pylint: disable=wrong-import-position
+
+
+results: dict[str, dict] = {}
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+
+def category(name: str):
+    if name not in results:
+        results[name] = {'passed': 0, 'failed': 0, 'skipped': 0, 'tests': []}
+    return name
+
+
+def record(cat: str, passed, name: str, detail: str = ''):
+    status = 'SKIP' if passed is None else ('PASS' if passed else 'FAIL')
+    key = {'SKIP': 'skipped', 'PASS': 'passed', 'FAIL': 'failed'}[status]
+    results[cat][key] += 1
+    results[cat]['tests'].append((status, name))
+    msg = f'  {status}: {name}'
+    if detail:
+        msg += f' ({detail})'
+    (log.info if status != 'FAIL' else log.error)(msg)
+
+
+def run_test(cat: str, fn):
+    name = fn.__name__
+    try:
+        outcome = fn()
+        record(cat, None if outcome is None else bool(outcome), name)
+    except AssertionError as e:
+        record(cat, False, name, str(e))
+    except Exception as e:  # pylint: disable=broad-except
+        record(cat, False, name, f'exception: {e}')
+        import traceback
+        traceback.print_exc()
+
+
+generator = torch.Generator(device=device).manual_seed(1234)
+
+
+def randn(*shape, dtype=torch.float32):
+    return torch.randn(*shape, generator=generator, device=device, dtype=dtype)
+
+
+def qkv(heads=4, seq=1024, dim=64):
+    return randn(1, heads, seq, dim), randn(1, heads, seq, dim), randn(1, heads, seq, dim)
+
+
+# ============================================================
+# Selector
+# ============================================================
+
+def test_pooling_matches_a_per_block_reference():
+    x = randn(1, 2, 300, 8)
+    pooled = sparse.selector.pool_blocks(x, 128)
+    assert pooled.shape == (1, 2, 3, 8), pooled.shape
+    for index, (start, end) in enumerate([(0, 128), (128, 256), (256, 300)]):
+        expected = x[..., start:end, :].to(torch.float32).mean(dim=-2)
+        assert torch.allclose(pooled[..., index, :], expected, atol=1e-5), index
+    return True
+
+
+def test_diagonal_covers_every_overlapping_tile():
+    nq, nk, bq, bk = 4, 8, 128, 64
+    diagonal = sparse.selector.diagonal_blocks(nq, nk, bq, bk, device)
+    for i in range(nq):
+        for j in range(nk):
+            overlaps = (i * bq < (j + 1) * bk) and (j * bk < (i + 1) * bq)
+            assert bool(diagonal[i, j]) == overlaps, (i, j)
+    assert int(diagonal.sum().item()) == nq * (bq // bk) # two kv tiles per query tile at 128 over 64
+    return True
+
+
+def test_budget_sets_density_over_the_candidates():
+    q, k = randn(1, 4, 1024, 32), randn(1, 4, 1024, 32)
+    for budget in (0.15, 0.30, 0.50):
+        spec = sparse.SparseSpec(budget=budget)
+        selection = sparse.select_blocks(q, k, spec)
+        assert selection is not None, budget
+        keep = selection.keep
+        diagonal = sparse.selector.diagonal_blocks(keep.shape[-2], keep.shape[-1], spec.block_q, spec.block_kv, device)
+        candidates = int((~diagonal).sum().item())
+        chosen = int((keep.bool() & ~diagonal).sum().item()) / keep.shape[1]
+        expected = candidates * budget
+        assert abs(chosen - expected) <= keep.shape[-2], f'budget={budget} chose {chosen} of {candidates}, expected about {expected}'
+        assert bool((keep.bool() | ~diagonal).all()), 'a diagonal tile was dropped'
+    return True
+
+
+def test_pins_survive_and_drops_never_appear():
+    q, k = randn(1, 2, 512, 32), randn(1, 2, 512, 32)
+    spec = sparse.SparseSpec(budget=0.10)
+    nq = sparse.block_count(512, spec.block_q)
+    nk = sparse.block_count(512, spec.block_kv)
+    pins = torch.zeros(1, 1, nq, nk, dtype=torch.bool, device=device)
+    drops = torch.zeros_like(pins)
+    pins[..., 0] = True # a pinned column, as a text prefix produces
+    drops[..., -1] = True # a padding column
+    selection = sparse.select_blocks(q, k, spec, pins=pins, drops=drops)
+    assert selection is not None
+    assert bool(selection.keep[..., 0].all()), 'pinned column not kept'
+    assert not bool(selection.keep[..., -1].any()), 'dropped column kept'
+    return True
+
+
+def test_dense_short_circuit_and_force():
+    q, k = randn(1, 2, 512, 32), randn(1, 2, 512, 32)
+    assert sparse.select_blocks(q, k, sparse.SparseSpec(budget=1.0)) is None, 'full budget must report dense'
+    forced = sparse.select_blocks(q, k, sparse.SparseSpec(budget=1.0, force=True))
+    assert forced is not None and bool(forced.keep.all()), 'forced full budget must keep every tile'
+    return True
+
+
+def test_selection_is_deterministic():
+    q, k = randn(1, 4, 1024, 32), randn(1, 4, 1024, 32)
+    spec = sparse.SparseSpec(budget=0.25)
+    first = sparse.select_blocks(q, k, spec)
+    second = sparse.select_blocks(q, k, spec)
+    assert torch.equal(first.keep, second.keep)
+    return True
+
+
+def test_head_shared_collapses_the_head_dimension():
+    q, k = randn(1, 8, 1024, 32), randn(1, 8, 1024, 32)
+    selection = sparse.select_blocks(q, k, sparse.SparseSpec(budget=0.25, head_shared=True))
+    assert selection.keep.shape[1] == 1, selection.keep.shape
+    return True
+
+
+def test_gqa_scores_on_query_heads():
+    q, k = randn(1, 8, 1024, 32), randn(1, 2, 1024, 32)
+    selection = sparse.select_blocks(q, k, sparse.SparseSpec(budget=0.25))
+    assert selection.keep.shape[1] == 8, selection.keep.shape # both consumers need the mask head dim to be Hq or 1
+    return True
+
+
+# ============================================================
+# Layout
+# ============================================================
+
+def test_layout_from_index_kwargs_relabels_the_conditioning_video_run():
+    kwargs = { # the shape MiniMax H3 passes its transformer: text, a keyframe video run, audio, then the generated video
+        'text_indices': torch.arange(0, 8, device=device),
+        'video_indices': torch.cat([torch.arange(8, 12, device=device), torch.arange(20, 40, device=device)]),
+        'audio_indices': torch.arange(12, 20, device=device),
+        'hidden_states': torch.zeros(1, device=device), # not an index tensor, must be ignored
+    }
+    layout = sparse.layout_from_index_kwargs(kwargs, length=40)
+    kinds = [(s.kind, s.start, s.end) for s in layout.spans]
+    assert kinds == [('text', 0, 8), ('cond', 8, 12), ('audio', 12, 20), ('video', 20, 40)], kinds
+    assert layout.sparsifiable_tokens() == 20
+    return True
+
+
+def test_layout_from_index_kwargs_returns_none_without_indices():
+    assert sparse.layout_from_index_kwargs({'hidden_states': torch.zeros(4, device=device)}, length=4) is None
+    return True
+
+
+def test_layout_from_segments_and_prefix():
+    layout = sparse.layout_from_segments([('text', 128), ('image', 4096), ('pad', 128)])
+    assert layout.length == 4352 and layout.sparsifiable_tokens() == 4096
+    prefix = sparse.layout_from_prefix(1024, 64)
+    assert prefix.sparsifiable_tokens() == 960 and prefix.source == 'prefix'
+    return True
+
+
+def test_block_pins_pin_conditioning_and_drop_padding():
+    block_q, block_kv = 128, 64
+    layout = sparse.layout_from_segments([('text', 128), ('video', 1024), ('pad', 128)])
+    pins, drops = sparse.block_pins(layout, 1280, 1280, block_q, block_kv, device)
+    assert pins.shape == (1, 1, 10, 20) and drops.shape == pins.shape, (pins.shape, drops.shape)
+    assert bool(pins[0, 0, :, 0:2].all()), 'the text columns must be pinned'
+    assert bool(drops[0, 0, :, 18:20].all()), 'the padding columns must be dropped'
+    assert not bool(drops[0, 0, :, 0:18].any()), 'only padding may be dropped'
+    assert bool(pins[0, 0, 0, 0:18].all()), 'the query tile holding text must stay dense over every column that is not padding'
+    assert not bool(pins[0, 0, :, 18:20].any()), 'a dropped column is skipped, never pinned'
+    assert not bool(pins[0, 0, 1:9, 2:18].any()), 'video against video must remain sparsifiable'
+    return True
+
+
+def test_block_pins_pin_a_boundary_tile():
+    layout = sparse.layout_from_segments([('text', 100), ('video', 1180)]) # the boundary falls inside the first tile
+    pins, drops = sparse.block_pins(layout, 1280, 1280, 128, 64, device)
+    assert not bool(drops.any()), 'nothing is padding here'
+    assert bool(pins[0, 0, 0, :].all()), 'a query tile straddling a boundary must stay dense'
+    assert bool(pins[0, 0, :, 0:2].all()), 'a key tile straddling a boundary must stay dense'
+    return True
+
+
+def test_block_pins_are_cached_per_geometry():
+    layout = sparse.layout_from_segments([('text', 128), ('video', 1024)])
+    first = sparse.block_pins(layout, 1152, 1152, 128, 64, device)
+    second = sparse.block_pins(layout, 1152, 1152, 128, 64, device)
+    assert first[0] is second[0] and first[1] is second[1], 'identical geometry should hit the cache'
+    return True
+
+
+# ============================================================
+# Consumers and controls
+# ============================================================
+
+def test_radial_control_matches_the_requested_density():
+    spec = sparse.SparseSpec()
+    for density in (0.15, 0.30):
+        control = sparse.radial_blocks(4096, 4096, density, spec, device)
+        assert abs(control.density - density) < 0.05, f'requested {density}, got {control.density}'
+    return True
+
+
+def test_schedule_has_at_most_two_budgets():
+    flat = sparse.schedule(20, 0.3)
+    assert set(flat) == {0.3} and len(flat) == 20
+    bumped = sparse.schedule(20, 0.3, bump=0.3, bump_steps=2)
+    assert len(set(bumped)) == 2, set(bumped)
+    assert bumped[0] == bumped[1] == 0.6 and bumped[-1] == bumped[-2] == 0.6 and bumped[10] == 0.3
+    return True
+
+
+def flex_available():
+    return device.type == 'cuda'
+
+
+def kernel_floor(q, k, v):
+    """How far the flex kernel sits from sdpa on the same dense problem, which bounds what any sparse row can prove."""
+    full = sparse.select_blocks(q, k, sparse.SparseSpec(budget=1.0, force=True))
+    return (sparse_flex.attend(q, k, v, full) - stock_sdpa(q, k, v)).abs().max().item()
+
+
+def test_flex_full_selection_reproduces_dense_sdpa():
+    if not flex_available():
+        return None
+    q, k, v = qkv()
+    floor = kernel_floor(q, k, v)
+    assert floor < 5e-3, f'a full selection should reproduce dense sdpa, differs by {floor}'
+    log.info(f'    flex kernel floor vs sdpa: {floor:.6f}')
+    return True
+
+
+def test_flex_sparse_selection_matches_the_same_tiles_under_sdpa():
+    if not flex_available():
+        return None
+    q, k, v = qkv()
+    spec = sparse.SparseSpec(budget=0.25)
+    selection = sparse.select_blocks(q, k, spec)
+    got = sparse_flex.attend(q, k, v, selection)
+    # expand the tile selection to tokens and hand sdpa the same thing
+    token_mask = selection.keep.bool().repeat_interleave(spec.block_q, dim=-2).repeat_interleave(spec.block_kv, dim=-1)
+    expected = stock_sdpa(q, k, v, attn_mask=token_mask[..., :q.shape[-2], :k.shape[-2]])
+    delta = (got - expected).abs().max().item()
+    floor = kernel_floor(q, k, v)
+    assert delta <= max(4 * floor, 2e-3), f'sparse selection differs from the same tiles under sdpa by {delta}, floor {floor}'
+    return True
+
+
+def test_flex_applies_the_selection_at_all():
+    if not flex_available():
+        return None
+    # flex reads the block lists only when compiled; eager evaluates mask_mod instead, so a
+    # block only mask silently attends densely. this row fails if the consumer stops compiling.
+    q, k, v = qkv()
+    selection = sparse.select_blocks(q, k, sparse.SparseSpec(budget=0.25))
+    delta = (sparse_flex.attend(q, k, v, selection) - stock_sdpa(q, k, v)).abs().max().item()
+    floor = kernel_floor(q, k, v)
+    assert delta > 20 * max(floor, 1e-6), f'a 25 percent selection changed the output by only {delta}, floor {floor}: the mask is not being applied'
+    return True
+
+
+def test_flex_handles_a_ragged_tail():
+    if not flex_available():
+        return None
+    seq = 1000 # neither block size divides this
+    q, k, v = qkv(heads=2, seq=seq)
+    selection = sparse.select_blocks(q, k, sparse.SparseSpec(budget=1.0, force=True))
+    delta = (sparse_flex.attend(q, k, v, selection) - stock_sdpa(q, k, v)).abs().max().item()
+    assert delta < 5e-3, f'ragged tail differs by {delta}'
+    return True
+
+
+def run_all():
+    log.warning(f'=== selector (device={device}) ===')
+    cat = category('selector')
+    for fn in [
+        test_pooling_matches_a_per_block_reference,
+        test_diagonal_covers_every_overlapping_tile,
+        test_budget_sets_density_over_the_candidates,
+        test_pins_survive_and_drops_never_appear,
+        test_dense_short_circuit_and_force,
+        test_selection_is_deterministic,
+        test_head_shared_collapses_the_head_dimension,
+        test_gqa_scores_on_query_heads,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== layout ===')
+    cat = category('layout')
+    for fn in [
+        test_layout_from_index_kwargs_relabels_the_conditioning_video_run,
+        test_layout_from_index_kwargs_returns_none_without_indices,
+        test_layout_from_segments_and_prefix,
+        test_block_pins_pin_conditioning_and_drop_padding,
+        test_block_pins_pin_a_boundary_tile,
+        test_block_pins_are_cached_per_geometry,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== consumers ===')
+    cat = category('consumers')
+    for fn in [
+        test_radial_control_matches_the_requested_density,
+        test_schedule_has_at_most_two_budgets,
+        test_flex_full_selection_reproduces_dense_sdpa,
+        test_flex_sparse_selection_matches_the_same_tiles_under_sdpa,
+        test_flex_applies_the_selection_at_all,
+        test_flex_handles_a_ragged_tail,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== Results ===')
+    total_passed = total_failed = total_skipped = 0
+    for cat_name, info in results.items():
+        ok = info['failed'] == 0
+        log.info(f"  {cat_name}: {info['passed']} passed, {info['failed']} failed, {info['skipped']} skipped [{'PASS' if ok else 'FAIL'}]")
+        total_passed += info['passed']
+        total_failed += info['failed']
+        total_skipped += info['skipped']
+    log.warning(f'Total: {total_passed} passed, {total_failed} failed, {total_skipped} skipped')
+    return total_failed == 0
+
+
+if __name__ == '__main__':
+    import time
+    t0 = time.time()
+    ok = run_all()
+    torch.nn.functional.scaled_dot_product_attention = stock_sdpa
+    log.warning(f'Total time: {time.time() - t0:.2f}s')
+    sys.exit(0 if ok else 1)