diff --git a/modules/attention/sparse/__init__.py b/modules/attention/sparse/__init__.py
new file mode 100644
index 000000000..5d8efc1f1
--- /dev/null
+++ b/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',
+]
diff --git a/modules/attention/sparse/flex.py b/modules/attention/sparse/flex.py
new file mode 100644
index 000000000..42ef1dad3
--- /dev/null
+++ b/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)
diff --git a/modules/attention/sparse/layout.py b/modules/attention/sparse/layout.py
new file mode 100644
index 000000000..772e49b65
--- /dev/null
+++ b/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
diff --git a/modules/attention/sparse/selector.py b/modules/attention/sparse/selector.py
new file mode 100644
index 000000000..89539e8cb
--- /dev/null
+++ b/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)])
diff --git a/test/test-attention-sparse.py b/test/test-attention-sparse.py
new file mode 100644
index 000000000..b4bd2f5e3
--- /dev/null
+++ b/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)