Merge branch 'dev' into feat/lora-sdnq-stack

cli/sdnq-attention-benchmark.py

@@ -96,25 +96,75 @@ def save_transcript(path):
             file_console.print(renderable)
     console.print(f"results saved to {path}")
 
+def key_padding_mask(cfg, device, keep=0.75):
+    # boolean key-padding mask over the kv axis, first keep fraction of keys valid
+    kv_tokens = cfg.get("kv_tokens", cfg["tokens"])
+    attn_mask = torch.zeros(cfg["batch"], 1, 1, kv_tokens, device=device, dtype=torch.bool)
+    attn_mask[..., :int(kv_tokens * keep)] = True
+    return attn_mask
+
+
+def krea2_segment_mask(cfg, device):
+    # the transformer's segment_mask: text is padded to a fixed 512 tokens ahead of the
+    # image tokens and the padded tail is masked for queries and keys both, so padding
+    # query rows are fully masked and yield nan under sdpa (the model nan_to_num's them)
+    valid = torch.ones(cfg["batch"], cfg["tokens"], device=device, dtype=torch.bool)
+    valid[:, 128:512] = False
+    return valid.unsqueeze(1).unsqueeze(2) * valid.unsqueeze(1).unsqueeze(3)
+
+
+def build_preset_mask(cfg, device):
+    # dense masks come from the preset's mask_fn; the element guard keeps h3-scale presets
+    # from materializing multi-gigabyte masks, those shapes belong to block-granular masks
+    mask_fn = cfg.get("mask_fn")
+    if mask_fn is None:
+        return None
+    attn_mask = mask_fn(cfg, device)
+    if attn_mask is not None and attn_mask.numel() > 2**31:
+        raise ValueError(f"preset dense mask holds {attn_mask.numel():,} elements; this shape needs a block-granular mask, not a token mask")
+    return attn_mask
+
+
 shape_presets = {
     # geometry from the model transformer and text-encoder configs;
-    # optional keys: kv_tokens (cross-attention), kv_heads (gqa), causal
+    # optional keys: kv_tokens (cross-attention), kv_heads (gqa), causal,
+    # mask_fn (token-granular attn_mask builder), mask_nan_guard (fully-masked query
+    # rows nan under stock sdpa), iters/warmup (per-preset run overrides for very large
+    # shapes), ref_head_chunk (head-sliced fp32 reference to bound peak memory),
+    # sparse (token layout driving the sparse selector rows)
     "sd15": dict(batch=2, heads=8, tokens=4096, head_dim=40, desc="SD 1.5 unet self-attention at 512px, batched cfg, head dim padded 40 to 64"),
     "sdxl": dict(batch=2, heads=10, tokens=4096, head_dim=64, desc="SDXL unet self-attention at 1024px, batched cfg"),
     "sdxl-cross": dict(batch=2, heads=10, tokens=4096, kv_tokens=77, head_dim=64, desc="SDXL unet cross-attention at 1024px, 77 text tokens"),
     "qwen3-te": dict(batch=2, heads=16, kv_heads=8, tokens=512, head_dim=128, causal=True, desc="Qwen3 text encoder (Anima), causal gqa 16:8 heads, 512 token prompt"),
     "anima": dict(batch=1, heads=16, tokens=4096, head_dim=128, desc="Anima 1.0 self-attention at 1024px, one cfg pass"),
     "flux2": dict(batch=1, heads=32, tokens=4608, head_dim=128, desc="FLUX.2 Klein 9B joint attention at 1024px, 4096 image plus 512 text tokens"),
-    "krea2": dict(batch=1, heads=48, tokens=4608, head_dim=128, desc="Krea 2 12B joint attention at 1024px, 4096 image plus 512 text tokens (128 real), kv expanded from gqa 48:12, segment mask"),
+    "krea2": dict(batch=1, heads=48, tokens=4608, head_dim=128, mask_fn=krea2_segment_mask, mask_nan_guard=True, desc="Krea 2 12B joint attention at 1024px, 4096 image plus 512 text tokens (128 real), kv expanded from gqa 48:12, segment mask"),
     "wan22": dict(batch=1, heads=40, tokens=32760, head_dim=128, desc="Wan 2.2 A14B self-attention, 832x480 81 frames, one cfg pass"),
     "wan22-cfg": dict(batch=2, heads=40, tokens=32760, head_dim=128, desc="Wan 2.2 A14B self-attention, 832x480 81 frames, batched cfg"),
     "ltx2": dict(batch=1, heads=32, tokens=13376, head_dim=128, desc="LTX 2.3 self-attention, 1216x704 121 frames, one cfg pass"),
-    "masked": dict(batch=1, heads=32, tokens=4608, head_dim=128, desc="FLUX.2 Klein shape with boolean key-padding mask, 25% of keys masked"),
+    "h3": dict(batch=1, heads=56, tokens=38222, head_dim=128, iters=8, warmup=3, ref_head_chunk=14,
+        sparse=dict(layout=[("text", 0, 512), ("audio", 512, 926), ("video", 926, 38222)]),
+        desc="MiniMax H3 packed self-attention, 1344x768 124 frames (5.2s): 512 text + 414 audio + 37296 video rows, guidance-free"),
+    "h3-long": dict(batch=1, heads=56, tokens=109574, head_dim=128, iters=6, warmup=2, ref_head_chunk=8, config_timeout=1200,
+        sparse=dict(layout=[("text", 0, 512), ("audio", 512, 1718), ("video", 1718, 109574)]),
+        desc="MiniMax H3 packed self-attention, 1344x768 362 frames (15.1s): 512 text + 1206 audio + 107856 video rows"),
+    "masked": dict(batch=1, heads=32, tokens=4608, head_dim=128, mask_fn=key_padding_mask, desc="FLUX.2 Klein shape with boolean key-padding mask, 25% of keys masked"),
 }
-full_run = ["sd15", "sdxl", "sdxl-cross", "qwen3-te", "anima", "flux2", "krea2", "wan22", "ltx2"]
+# sparse crossover probes at fixed h3 geometry; the smallest token count where a sparse row
+# beats dense past the verdict threshold is the measured minimum-sequence gate
+for gate_tokens in (2048, 4096, 8192, 16384, 32768, 65536):
+    shape_presets[f"gate-{gate_tokens // 1024}k"] = dict(
+        batch=1, heads=56, tokens=gate_tokens, head_dim=128,
+        sparse=dict(layout=[("text", 0, 512), ("video", 512, gate_tokens)]),
+        desc=f"sparse crossover probe at h3 geometry, {gate_tokens} tokens",
+        **(dict(iters=8, warmup=3) if gate_tokens >= 32768 else {}),
+    )
+full_run = ["sd15", "sdxl", "sdxl-cross", "qwen3-te", "anima", "flux2", "krea2", "wan22", "ltx2", "h3"]
+sparse_run = ["krea2", "h3", "h3-long"]
+gate_run = [f"gate-{tokens // 1024}k" for tokens in (2048, 4096, 8192, 16384, 32768, 65536)]
 # settings advice comes from a self-attention shape with the full config set; cross-attention
 # and text-encoder shapes measure the hijack's cost there but would mislead as global advice
-recommendation_presets = ["flux2", "krea2", "anima", "sdxl", "wan22", "ltx2", "sd15"]
+recommendation_presets = ["flux2", "krea2", "anima", "sdxl", "wan22", "ltx2", "h3", "sd15"]
 default_shapes = "sdxl,flux2"
 all_sections = ["attention", "dequant", "block"]
 
@@ -122,11 +172,16 @@ all_sections = ["attention", "dequant", "block"]
 # measures a complete configuration of weights dtype x matmul path x attention end to end
 # generic dit-block geometries from the model transformer configs: flux.1 (3072 wide,
 # 24 heads, 4x gelu ff, 4096 image plus 512 text tokens) and krea 2 (6144 wide, 48 heads
-# after gqa expansion, swiglu at 16384, same joint sequence)
+# after gqa expansion, swiglu at 16384, same joint sequence); optional keys: head_dim
+# (attention width when heads*head_dim != hidden) and mlp ("gelu" default or "swiglu")
 block_geometries = {
     "flux1": dict(hidden=3072, heads=24, mlp_dim=12288, tokens=4608),
     "krea2": dict(hidden=6144, heads=48, mlp_dim=16384, tokens=4608),
+    # minimax h3: attention wider than the residual stream (56*128 > 5376), swiglu mlp; the
+    # full 124-frame token count makes the block section long, so it runs only when selected
+    "h3": dict(hidden=5376, heads=56, head_dim=128, mlp_dim=14336, mlp="swiglu", tokens=38222, iters=6, warmup=2, config_timeout=900),
 }
+default_block_geometries = "flux1,krea2"
 block_geometry = block_geometries["flux1"] # active geometry; bench_block_section iterates
 block_attention_specs = {
     "sdpa": None, # stock torch sdpa
@@ -139,6 +194,19 @@ block_attention_specs = {
     "sage": "sage", # external baselines, resolved to the sage wrappers in build_bench_block
     "sage fp16 accum": "sagefp16",
 }
+# attention-table config id measuring the same kernel as each block attention spec, for the
+# cross-instrument compute split; specs without a standalone row map to None
+block_spec_attention_ids = {
+    "sdpa": "base",
+    "atten int8": "int8",
+    "atten int8 smooth": "smooth",
+    "atten int8 hadamard": "hadamard",
+    "atten full": "full",
+    "atten pv accum": "pvaccum",
+    "atten fp16 accum": "fp16full-accum",
+    "sage": "sage",
+    "sage fp16 accum": "sagefp16",
+}
 # id, weights config (None = bf16), use quantized matmul, attention spec; fp8/fp4 rows use the
 # dequant path: quantized matmul auto-selects fp8 for float dtypes, unsupported before sm_89
 block_configs = [
@@ -173,8 +241,19 @@ bench_configs = [
     ("sage", "sageattention", None), # label resolved to the dispatched kernel by sage_kernel_label
     ("sagefp16", "sage int8 qk + fp16 pv, fp16 accum", None), # sm86 only
     ("amdflash", "triton flash (amd)", None),
+    ("flex", "flex attention, dense", None), # compiled: flex reads its block lists only under compile
+    ("flex-sparse100", "flex + selector, budget 100%", None), # the selector runs but keeps everything, so this row is its overhead alone
+    ("flex-sparse50", "flex + selector, budget 50%", None),
+    ("flex-sparse30", "flex + selector, budget 30%", None),
+    ("flex-sparse15", "flex + selector, budget 15%", None),
+    ("flex-radial30", "flex + static radial band, 30%", None), # density matched control with no per-call producer
     ("noquant", "sdnq, quantized matmul off", dict(do_quantize=False)),
     ("int8", "sdnq int8 qk", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
+    ("int8-sparse100", "sdnq int8 qk + selector, budget 100%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), # the selector runs but keeps everything, so this row is its overhead on the quantized kernel
+    ("int8-sparse50", "sdnq int8 qk + selector, budget 50%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
+    ("int8-sparse30", "sdnq int8 qk + selector, budget 30%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
+    ("int8-sparse15", "sdnq int8 qk + selector, budget 15%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
+    ("int8-radial30", "sdnq int8 qk + static radial band, 30%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), # density matched control with no per-call producer
     ("smooth", "sdnq int8 qk + smooth k", dict(matmul_dtype="auto", pv_matmul_dtype="auto", smooth_k=True)),
     ("hadamard", "sdnq int8 qk + hadamard", dict(matmul_dtype="auto", pv_matmul_dtype="auto", use_hadamard=True)),
     ("smooth_hadamard", "sdnq int8 qk + smooth + hadamard", dict(matmul_dtype="auto", pv_matmul_dtype="auto", smooth_k=True, use_hadamard=True)),
@@ -193,7 +272,10 @@ bench_configs = [
 # external baselines are compared against but never starred or recommended as sdnq configs;
 # the unsafe accum mode is measured and displayed under the same rule, since its overflow
 # tail lives outside what mean error can see
-external_config_ids = ("base", "sage", "sagefp16", "amdflash")
+# baselines and non-sdnq rows: reported, never recommended as an sdnq setting, and the sparse
+# rows are lossy by design so a recommendation must not pick one for being fast
+external_config_ids = ("base", "sage", "sagefp16", "amdflash", "flex")
+sparse_config_ids = ("flex-sparse100", "flex-sparse50", "flex-sparse30", "flex-sparse15", "flex-radial30", "int8-sparse100", "int8-sparse50", "int8-sparse30", "int8-sparse15", "int8-radial30")
 unsafe_config_ids = ("pvaccum",)
 # every preset runs the full config list (availability gates still apply per config); only
 # hard technical exclusions live here, never runtime trims. sd15: compiling hadamard with
@@ -281,18 +363,22 @@ def parse_cli():
     parser.add_argument("--dequant-sweeps", type=str, default="all", help=f"comma-separated setting sweeps in the dequant section: {', '.join(all_dequant_sweeps)}; 'all' or 'none' (default: %(default)s)")
     parser.add_argument("--mm-backends", type=str, default="none", help=f"comma-separated quantized-matmul backends to compare in one run: {', '.join(all_mm_backends)}; 'none' benches only the backend this device selects (default: %(default)s)")
     parser.add_argument("--mm-rounds", type=int, default=2, help="alternating rounds per matmul backend, fastest kept, so clock drift cancels instead of favouring one backend (default: %(default)s)")
+    parser.add_argument("--configs", type=str, default="all", help=f"comma-separated attention configs: {', '.join(config_id for config_id, _label, _kwargs in bench_configs)}; 'all' runs every one (default: %(default)s)")
     parser.add_argument("--block-configs", type=str, default="all", help=f"comma-separated combined block configs: {', '.join(config_id for config_id, _w, _mm, _a in block_configs)}; 'all' runs every one (default: %(default)s)")
-    parser.add_argument("--shapes", type=str, default=default_shapes, help=f"comma-separated attention shape presets: {', '.join(shape_presets)}; 'all' runs {', '.join(full_run)} (default: %(default)s)")
+    parser.add_argument("--block-geometries", type=str, default=default_block_geometries, help=f"comma-separated block geometries: {', '.join(block_geometries)}; 'all' runs every one (default: %(default)s)")
+    parser.add_argument("--shapes", type=str, default=default_shapes, help=f"comma-separated attention shape presets: {', '.join(shape_presets)}; 'all' runs {', '.join(full_run)}, 'sparse' and 'gate' run the sparse and crossover lists (default: %(default)s)")
     parser.add_argument("--iters", type=int, default=12, help="minimum timed iterations per config, scaled up for fast kernels (default: %(default)s)")
     parser.add_argument("--warmup", type=int, default=4, help="minimum warmup iterations per config, scaled up for fast kernels (default: %(default)s)")
     parser.add_argument("--skip-checks", action="store_true", help="skip kernel correctness checks")
     parser.add_argument("--skip-bench", action="store_true", help="skip benchmarks, run checks and the fp8 and compile probes only")
     parser.add_argument("--dtype", type=str, default="auto", choices=["auto", "bf16", "fp16"], help="tensor dtype for benchmarks; auto uses the dtype the webui selected for this gpu (default: %(default)s)")
-    parser.add_argument("--config-timeout", type=int, default=300, help="best effort: abort a config whose compile plus first call exceeds this many seconds, 0 disables; cannot interrupt native-level hangs (default: %(default)s)")
+    parser.add_argument("--config-timeout", type=int, default=None, help="best effort: abort a config whose compile plus first call exceeds this many seconds, 0 disables; cannot interrupt native-level hangs (default: 300, or the limit a preset or block geometry declares for itself)")
     parser.add_argument("--save", type=str, default="auto", help="plain-text copy of all tables and notes; 'auto' (default) names it <gpu>-t<torch>-<date>.txt in the output directory, 'none' disables, anything else is used as the path")
     parser.add_argument("--json", type=str, default="auto", help="structured results (environment, probes, per-shape and dequant timings, recommendations); 'auto' (default) names it <gpu>-t<torch>-<date>.json in the output directory, 'none' disables, anything else is used as the path")
     parser.add_argument("--outdir", type=str, default=None, help="directory for auto-named outputs (default: $SDNQ_BENCH_DIR, or benchmarks/ under the sdnext root)")
     args = parser.parse_args()
+    args.timeout_flag = args.config_timeout # None lets a preset or block geometry declare its own limit
+    args.config_timeout = resolve_timeout(args.timeout_flag)
     sys.argv = sys.argv[:1] # sdnext parses argv again on import and rejects unknown arguments
     return args
 
@@ -425,6 +511,16 @@ def atten_supports_fp16_accum():
         return False
 
 
+def atten_supports_block_mask():
+    # the sdnq sparse rows feed the kernel's block_mask kwarg; skip them on builds without it
+    if sdnq_triton_atten is None:
+        return False
+    try:
+        return "block_mask" in inspect.signature(inspect.unwrap(sdnq_triton_atten)).parameters
+    except (TypeError, ValueError):
+        return False
+
+
 def triton_mm_supports_fp16_accum():
     try:
         from sdnq.kernels import triton_mm, triton_scaled_mm
@@ -450,6 +546,56 @@ def triton_mm_fp16_accum():
         triton_mm.USE_FP16_ACCUM, triton_scaled_mm.USE_FP16_ACCUM = saved
 
 
+flex_budgets = {"flex-sparse100": 1.0, "flex-sparse50": 0.50, "flex-sparse30": 0.30, "flex-sparse15": 0.15}
+sdnq_sparse_budgets = {"int8-sparse100": 1.0, "int8-sparse50": 0.50, "int8-sparse30": 0.30, "int8-sparse15": 0.15}
+
+
+def is_sdnq_sparse(config_id):
+    return config_id in sdnq_sparse_budgets or config_id == "int8-radial30"
+
+
+def make_sdnq_sparse_fn(config_id, q, k, v, attn_mask, kwargs, causal, gqa):
+    """The producer the flex rows time, feeding the quantized kernel's block mask input instead."""
+    from modules.attention.sparse import selector as sparse_selector
+
+    def attend(selection):
+        return sdnq_triton_atten(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa, block_mask=selection.keep, block_mask_m=selection.block_q, block_mask_n=selection.block_kv, **kwargs)
+    if config_id == "int8-radial30":
+        static = sparse_selector.radial_blocks(q.shape[-2], k.shape[-2], 0.30, sparse_selector.SparseSpec(), q.device)
+        return lambda: attend(static)
+    spec = sparse_selector.SparseSpec(budget=sdnq_sparse_budgets[config_id], force=True)
+    cache_key = ("bench", config_id, tuple(q.shape), tuple(k.shape))
+    return lambda: attend(sparse_selector.select_blocks(q, k, spec, cache_key=cache_key))
+
+
+def flex_available():
+    try:
+        import modules.attention.sparse.flex # pylint: disable=unused-import
+        return torch.cuda.is_available()
+    except Exception:
+        return False
+
+
+def make_flex_fn(config_id, q, k, v, scale, gqa):
+    """Time the selector inside the attention it accelerates; a producer measured on its own looks free and is not."""
+    from modules.attention.sparse import flex as sparse_flex, selector as sparse_selector
+    call = sparse_flex.flex_call()
+    if config_id == "flex":
+        return lambda: call(q, k, v, scale=scale, enable_gqa=gqa)
+    if config_id == "flex-radial30":
+        # a static pattern is built once by construction, which is exactly the advantage it has to defend
+        static = sparse_flex.to_block_mask(sparse_selector.radial_blocks(q.shape[-2], k.shape[-2], 0.30, sparse_selector.SparseSpec(), q.device))
+        return lambda: call(q, k, v, block_mask=static, scale=scale, enable_gqa=gqa)
+    spec = sparse_selector.SparseSpec(budget=flex_budgets[config_id], force=True)
+
+    cache_key = ("bench", config_id, tuple(q.shape), tuple(k.shape)) # the webui caches the geometry per layout, so measure that path
+
+    def run():
+        selection = sparse_selector.select_blocks(q, k, spec, cache_key=cache_key)
+        return call(q, k, v, block_mask=sparse_flex.to_block_mask(selection), scale=scale, enable_gqa=gqa)
+    return run
+
+
 def sage_attention():
     # mirror the backend selection from modules/attention.py: sm86 needs the cuda backend
     try:
@@ -545,13 +691,23 @@ def make_qkv(batch, heads, tokens, head_dim, structured=True, kv_heads=None, kv_
     return q, k, v
 
 
-def fp32_reference(q, k, v, **kwargs):
+def fp32_reference(q, k, v, head_chunk=0, **kwargs):
     # sdnext enables tf32 globally; a math-backend dispatch fallback would degrade the reference to tf32 precision
     tf32_matmul = torch.backends.cuda.matmul.allow_tf32
     tf32_cudnn = torch.backends.cudnn.allow_tf32
     torch.backends.cuda.matmul.allow_tf32 = False
     torch.backends.cudnn.allow_tf32 = False
     try:
+        if head_chunk and not kwargs.get("enable_gqa") and q.shape[1] > head_chunk:
+            # head-sliced reference: bounds the fp32 peak on very long sequences; gqa shapes
+            # keep the one-shot path since slicing q heads would have to regroup kv heads
+            attn_mask = kwargs.pop("attn_mask", None)
+            outs = []
+            for start in range(0, q.shape[1], head_chunk):
+                heads = slice(start, start + head_chunk)
+                mask_slice = attn_mask[:, heads] if attn_mask is not None and attn_mask.shape[1] > 1 else attn_mask
+                outs.append(torch.nn.functional.scaled_dot_product_attention(q[:, heads].to(torch.float32), k[:, heads].to(torch.float32), v[:, heads].to(torch.float32), attn_mask=mask_slice, **kwargs))
+            return torch.cat(outs, dim=1)
         return torch.nn.functional.scaled_dot_product_attention(q.to(torch.float32), k.to(torch.float32), v.to(torch.float32), **kwargs)
     finally:
         torch.backends.cuda.matmul.allow_tf32 = tf32_matmul
@@ -641,6 +797,13 @@ def live_progress():
     return progress, task
 
 
+def resolve_timeout(flag, declared=None):
+    # the cli flag wins when given; otherwise a preset or block geometry may declare its own limit
+    if flag is not None:
+        return flag
+    return 300 if declared is None else declared
+
+
 @contextmanager
 def time_limit(seconds, label):
     # torch.compile can spin indefinitely in sympy/inductor on pathological graphs
@@ -721,6 +884,9 @@ def run_drift_sigma():
 
 drift_override = None
 verdict_z = 1.28 # one-sided 90%: an on/off verdict is only stated when its margin test clears this
+# split instruments A and B differ systematically (strided views out of the fused projection vs
+# contiguous standalone tensors): 0.1-2.5% over 19 same-kernel rows at h3, far below the 17-25% hadamard gap
+split_instrument_offset = 0.05
 
 
 def sidak_z_for(count):
@@ -853,6 +1019,8 @@ def print_environment(fp8_result, prep_status, prep_detail, weight_dequant_resul
         lines.append(f"compiled weight dequant, float8_e5m2 storage: {e5m2_verdict}")
     if not atten_supports_fp16_accum():
         lines.append("fp16 accumulation kwarg: [yellow]absent in this sdnq build, accum rows skipped[/yellow]")
+    if not atten_supports_block_mask():
+        lines.append("block mask kwarg: [yellow]absent in this sdnq build, sdnq sparse rows skipped[/yellow]")
     overrides = [f"{key}={value}" for key, value in os.environ.items() if key.startswith("SDNQ_TRITON_ATTEN") or key.startswith("SDNQ_TRITON_MM") or key.startswith("SDNQ_ALLOW_FP8") or key.startswith("SDNQ_COMPILE")]
     if overrides:
         lines.append(f"env overrides: {' '.join(overrides)}")
@@ -870,6 +1038,7 @@ def print_environment(fp8_result, prep_status, prep_detail, weight_dequant_resul
         **runtime_versions,
         fp8_attention_matmul=fp8_result["qk"][0] if fp8_result is not None else None,
         atten_fp16_accum=atten_supports_fp16_accum(),
+        atten_block_mask=atten_supports_block_mask(),
         triton_mm_fp16_accum=os.environ.get("SDNQ_TRITON_MM_USE_FP16_ACCUM", None),
         compiled_input_prep=prep_status,
         fp8_compile_gate=fp8_compile_gate_flag(),
@@ -1009,35 +1178,56 @@ def make_source_weight(out_features, in_features, seed=1234):
 
 
 class BenchBlock(torch.nn.Module):
-    # dit-style block: fused qkv self-attention plus a gelu mlp, both with residuals; the
-    # attention_fn attribute is set per benchmark config (stock sdpa or sdnq attention)
-    def __init__(self, hidden, heads, mlp_dim, device=None, dtype=None):
+    # dit-style block: fused qkv self-attention plus a gelu or swiglu mlp, both with residuals;
+    # head_dim decouples attention width from hidden for models whose attention is wider than
+    # the residual stream; the attention_fn attribute is set per benchmark config
+    def __init__(self, hidden, heads, mlp_dim, head_dim=None, mlp="gelu", device=None, dtype=None):
         super().__init__()
         self.heads = heads
+        self.head_dim = head_dim if head_dim is not None else hidden // heads
+        self.mlp = mlp
+        inner = heads * self.head_dim
         self.norm1 = torch.nn.LayerNorm(hidden, elementwise_affine=False, device=device, dtype=dtype)
         self.norm2 = torch.nn.LayerNorm(hidden, elementwise_affine=False, device=device, dtype=dtype)
-        self.qkv = torch.nn.Linear(hidden, hidden * 3, bias=False, device=device, dtype=dtype)
-        self.proj = torch.nn.Linear(hidden, hidden, bias=False, device=device, dtype=dtype)
+        self.qkv = torch.nn.Linear(hidden, inner * 3, bias=False, device=device, dtype=dtype)
+        self.proj = torch.nn.Linear(inner, hidden, bias=False, device=device, dtype=dtype)
         self.up = torch.nn.Linear(hidden, mlp_dim, bias=False, device=device, dtype=dtype)
+        if mlp == "swiglu":
+            self.gate = torch.nn.Linear(hidden, mlp_dim, bias=False, device=device, dtype=dtype)
         self.down = torch.nn.Linear(mlp_dim, hidden, bias=False, device=device, dtype=dtype)
         self.attention_fn = None
 
     def forward(self, x):
-        batch, tokens, channels = x.shape
+        batch, tokens, _channels = x.shape
         h = self.norm1(x)
-        qkv = self.qkv(h).view(batch, tokens, 3, self.heads, channels // self.heads).permute(2, 0, 3, 1, 4)
-        attn = self.attention_fn(qkv[0], qkv[1], qkv[2]).transpose(1, 2).reshape(batch, tokens, channels)
+        qkv = self.qkv(h).view(batch, tokens, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        attn = self.attention_fn(qkv[0], qkv[1], qkv[2]).transpose(1, 2).reshape(batch, tokens, self.heads * self.head_dim)
         x = x + self.proj(attn)
         h = self.norm2(x)
+        if self.mlp == "swiglu":
+            return x + self.down(torch.nn.functional.silu(self.gate(h)) * self.up(h))
         return x + self.down(torch.nn.functional.gelu(self.up(h)))
 
 
+def build_block_module(dtype=None):
+    # construct a block for the active geometry; every construction site goes through here so
+    # geometry keys are read in exactly one place
+    return BenchBlock(
+        block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"],
+        head_dim=block_geometry.get("head_dim"), mlp=block_geometry.get("mlp", "gelu"),
+        device=torch_device, dtype=dtype if dtype is not None else bench_dtype,
+    )
+
+
 def make_block_master():
-    # one master weight set shared by every block config, so all rows quantize identical weights
-    hidden, heads, mlp_dim = block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"]
-    block = BenchBlock(hidden, heads, mlp_dim, device=torch_device, dtype=bench_dtype)
+    # one master weight set shared by every block config, so all rows quantize identical weights;
+    # seed order keeps gelu geometries bitwise stable, swiglu appends its gate after up
+    block = build_block_module()
+    linears = [block.qkv, block.proj, block.up, block.down]
+    if hasattr(block, "gate"):
+        linears.append(block.gate)
     with torch.no_grad():
-        for seed, linear in enumerate((block.qkv, block.proj, block.up, block.down), start=1):
+        for seed, linear in enumerate(linears, start=1):
             linear.weight.copy_(make_source_weight(linear.out_features, linear.in_features, seed=seed))
     return {key: value.clone() for key, value in block.state_dict().items()}
 
@@ -1045,8 +1235,7 @@ def make_block_master():
 def build_bench_block(master_sd, weights_cfg, use_mm, attention_spec):
     from sdnq import SDNQConfig
     from sdnq.quantizer import apply_sdnq_to_module
-    hidden, heads, mlp_dim = block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"]
-    block = BenchBlock(hidden, heads, mlp_dim, device=torch_device, dtype=bench_dtype)
+    block = build_block_module()
     block.load_state_dict(master_sd)
     block.eval()
     for param in block.parameters():
@@ -1335,36 +1524,29 @@ def make_prep_fn(q, k, v, attn_mask, kwargs, is_causal=False, enable_gqa=False):
     return prep
 
 
-def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_result=None):
+def bench_shape(preset, iters, warmup, position=None, config_timeout=None, fp8_result=None, selected=None):
     preset_cfg = shape_presets[preset]
+    config_timeout = resolve_timeout(config_timeout, preset_cfg.get("config_timeout"))
     batch, heads, tokens, head_dim = preset_cfg["batch"], preset_cfg["heads"], preset_cfg["tokens"], preset_cfg["head_dim"]
     kv_tokens = preset_cfg.get("kv_tokens", tokens)
     kv_heads = preset_cfg.get("kv_heads", heads)
     causal = preset_cfg.get("causal", False)
     gqa = kv_heads != heads
     description = preset_cfg["desc"]
+    iters = preset_cfg.get("iters", iters)
+    warmup = preset_cfg.get("warmup", warmup)
+    ref_head_chunk = preset_cfg.get("ref_head_chunk", 0)
     excluded_configs = preset_excluded_configs.get(preset, set())
     if preset == "sd15":
         emit("[yellow]sd15: hadamard configs skipped, compiling hadamard with a non pow2 head dim currently hangs torch inductor[/yellow]")
-    attn_mask = None
-    mask_nan_guard = False
-    if preset == "masked":
-        attn_mask = torch.zeros(batch, 1, 1, tokens, device=torch_device, dtype=torch.bool)
-        attn_mask[..., :int(tokens * 0.75)] = True
-    elif preset == "krea2":
-        # the transformer's segment_mask: text is padded to a fixed 512 tokens ahead of the
-        # image tokens and the padded tail is masked for queries and keys both, so padding
-        # query rows are fully masked and yield nan under sdpa (the model nan_to_num's them)
-        valid = torch.ones(batch, tokens, device=torch_device, dtype=torch.bool)
-        valid[:, 128:512] = False
-        attn_mask = valid.unsqueeze(1).unsqueeze(2) * valid.unsqueeze(1).unsqueeze(3)
-        mask_nan_guard = True
+    attn_mask = build_preset_mask(preset_cfg, torch_device)
+    mask_nan_guard = preset_cfg.get("mask_nan_guard", False)
     sage = sage_attention()
     sage_fp16 = sage_attention_fp16_accum()
     amd_flash = amd_triton_flash()
     selected_configs = []
     for config_id, label, kwargs in bench_configs:
-        if config_id in excluded_configs:
+        if config_id in excluded_configs or (selected is not None and config_id not in selected):
             continue
         if config_id == "sage" and (sage is None or attn_mask is not None or head_dim not in {64, 96, 128} or kv_tokens != tokens or gqa or causal):
             continue
@@ -1372,6 +1554,10 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_re
             continue
         if config_id == "amdflash" and (amd_flash is None or attn_mask is not None or head_dim > 128 or gqa):
             continue
+        if config_id.startswith("flex") and (not flex_available() or attn_mask is not None or causal or kv_tokens != tokens):
+            continue # a block only mask cannot carry a token mask or a causal rule, and cross attention is not sparsified
+        if is_sdnq_sparse(config_id) and (not atten_supports_block_mask() or causal or kv_tokens != tokens):
+            continue # the kernel composes a token mask with the block mask, so only the causal and cross attention rules apply
         if config_id == "fp8qk" and not (fp8_result and fp8_result["qk"][0]):
             continue
         if config_id == "fp8pv" and not (fp8_result and fp8_result["pv"][0]):
@@ -1413,7 +1599,7 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_re
         progress.update(task, description=f"{prefix}{preset}: preparing inputs and fp32 reference")
         q, k, v = make_qkv(batch, heads, tokens, head_dim, kv_heads=kv_heads, kv_tokens=kv_tokens)
         scale = head_dim ** -0.5
-        ref = fp32_reference(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa)
+        ref = fp32_reference(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa, head_chunk=ref_head_chunk)
         if mask_nan_guard:
             ref = torch.nan_to_num(ref)
         anchor_fn = None
@@ -1432,6 +1618,10 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_re
             elif config_id == "amdflash":
                 def fn(sm=scale):
                     return amd_flash(q, k, v, sm, is_causal=causal)
+            elif config_id.startswith("flex"):
+                fn = make_flex_fn(config_id, q, k, v, scale, gqa)
+            elif is_sdnq_sparse(config_id):
+                fn = make_sdnq_sparse_fn(config_id, q, k, v, attn_mask, kwargs, causal, gqa)
             else:
                 def fn(kw=kwargs, mask=attn_mask):
                     return sdnq_triton_atten(q, k, v, attn_mask=mask, is_causal=causal, enable_gqa=gqa, **kw)
@@ -2514,17 +2704,24 @@ def block_label(weights_cfg, use_mm, attention_spec):
     return f"{weights_part} + {attention_spec}"
 
 
-def bench_block_section(iters, warmup, config_timeout=300, selected=None):
+def bench_block_section(iters, warmup, config_timeout=None, selected=None, geometries=None):
     global block_geometry # pylint: disable=global-statement
     all_results = {}
     for family, geometry in block_geometries.items():
+        if geometries is not None and family not in geometries:
+            continue
         block_geometry = geometry
-        results = bench_block_geometry(iters, warmup, config_timeout=config_timeout, selected=selected)
+        geometry_iters = geometry.get("iters", iters)
+        geometry_warmup = geometry.get("warmup", warmup)
+        geometry_timeout = resolve_timeout(config_timeout, geometry.get("config_timeout"))
+        results = bench_block_geometry(geometry_iters, geometry_warmup, config_timeout=geometry_timeout, selected=selected)
         all_results[family] = results
-        report.setdefault("blocks", {})[family] = dict(geometry=dict(geometry), results=results)
-    # the first family also lands at the flat block key, which replays and the buyback
+        report.setdefault("blocks", {})[family] = dict(geometry=dict(geometry), results=results, split=None)
+    if not all_results:
+        return {}
+    # the first family run also lands at the flat block key, which replays and the buyback
     # veto fall back to when no family matches the reference shape
-    primary = next(iter(block_geometries))
+    primary = next(iter(all_results))
     report["block"] = report["blocks"][primary]
     return all_results.get(primary, {})
 
@@ -2556,7 +2753,7 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
     table.add_column("out err", justify="right")
     table.add_column("max tok err", justify="right")
     table.add_column("err x4 blocks", justify="right")
-    panel = Panel(table, title=f"combined block: hidden={hidden} heads={heads} mlp={mlp_dim} tokens={tokens} {dtype_label()}", subtitle="[dim]dit block, fused qkv + gelu mlp with residuals; err vs an fp32 reference block, max tok = worst single token, x4 = four stacked blocks[/dim]", box=ROUNDED_BOX, expand=False)
+    panel = Panel(table, title=f"combined block: hidden={hidden} heads={heads} mlp={mlp_dim} tokens={tokens} {dtype_label()}", subtitle="[dim]dit block, fused qkv attention + mlp with residuals; err vs an fp32 reference block, max tok = worst single token, x4 = four stacked blocks[/dim]", box=ROUNDED_BOX, expand=False)
 
     def run_depth(block, x0, depth):
         h = x0
@@ -2572,7 +2769,7 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
         master = make_block_master()
         generator = torch.Generator(device=torch_device).manual_seed(7)
         x = torch.randn(1, tokens, hidden, device=torch_device, dtype=bench_dtype, generator=generator)
-        ref_block = BenchBlock(hidden, heads, mlp_dim, device=torch_device, dtype=torch.float32)
+        ref_block = build_block_module(dtype=torch.float32)
         ref_block.load_state_dict(master)
         ref_block.eval()
         def ref_attention(q, k, v):
@@ -2619,6 +2816,15 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
                 phase("measuring depth-4 error")
                 with torch.no_grad():
                     entry["err4"] = rel_err(run_depth(block, x, 4), ref_out4)
+                phase("timing identity-attention variant")
+                real_attention_fn = block.attention_fn
+                def identity_attention_fn(q, k, v): # pylint: disable=unused-argument # same shapes and permutes, zero attention flops
+                    return v
+                block.attention_fn = identity_attention_fn
+                try:
+                    entry["identity_ms"], entry["identity_ms_sigma"] = bench_stats(fn, warmup, iters, on_phase=phase)
+                finally:
+                    block.attention_fn = real_attention_fn
                 del block, out
                 if base_ms is None:
                     base_ms = entry["ms"]
@@ -2653,6 +2859,72 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
     return results
 
 
+def emit_block_splits():
+    # instrument B reads the attention tables, so the split renders once both sections are in
+    for family, data in (report.get("blocks") or {}).items():
+        data["split"] = block_split_table(family, data["geometry"], data["results"])
+
+
+def block_split_table(family, geometry, results):
+    # compute split per config from two independent instruments: A subtracts the identity-
+    # attention variant timed inside the block, B reads the standalone attention table at the
+    # same geometry from this run; a speedup ceiling is only stated where the two agree
+    head_dim = geometry.get("head_dim") or geometry["hidden"] // geometry["heads"]
+    expected_geometry = f"batch=1 heads={geometry['heads']} tokens={geometry['tokens']} head_dim={head_dim}"
+    attention_preset = None
+    for preset_name, data in (report.get("attention") or {}).items():
+        if data.get("geometry") == expected_geometry and not shape_presets.get(preset_name, {}).get("mask_fn"):
+            attention_preset = preset_name
+            break
+    attention_results = (report.get("attention") or {}).get(attention_preset, {}).get("results", {}) if attention_preset else {}
+    spec_by_config = {config_id: spec for config_id, _w, _mm, spec in block_configs}
+    budgets = (0.5, 0.3, 0.15)
+
+    table = Table(box=box.SIMPLE_HEAVY)
+    table.add_column("config")
+    table.add_column("block ms", justify="right")
+    table.add_column("rest ms", justify="right")
+    table.add_column("attn A", justify="right")
+    table.add_column("attn B", justify="right")
+    table.add_column("agree", justify="right")
+    for budget in budgets:
+        table.add_column(f"ceil@{int(budget * 100)}%", justify="right")
+
+    split = {}
+    for config_id, entry in results.items():
+        ms, identity_ms = entry.get("ms"), entry.get("identity_ms")
+        if not ms or not identity_ms:
+            continue
+        attn_a = ms - identity_ms
+        if attn_a <= 0:
+            continue
+        attention_id = block_spec_attention_ids.get(spec_by_config.get(config_id))
+        attention_entry = attention_results.get(attention_id) or {}
+        attn_b = attention_entry.get("ms")
+        agree = None
+        if attn_b:
+            # subtraction amplifies the relative sigma of instrument A by ms/attn_a
+            sigma_a = (row_sigma(entry) or 0.0) * (ms / attn_a)
+            sigma_b = row_sigma(attention_entry, "ms") or 0.0
+            threshold = max(verdict_z * math.sqrt(sigma_a * sigma_a + sigma_b * sigma_b), run_drift_sigma(), split_instrument_offset)
+            agree = abs(math.log(attn_a / attn_b)) <= threshold
+        ceilings = {budget: ms / (budget * attn_a + identity_ms) for budget in budgets} if agree else None
+        split[config_id] = dict(ms=ms, rest_ms=identity_ms, attn_a_ms=attn_a, attn_b_ms=attn_b, agree=agree, ceilings=ceilings)
+        agree_cell = "-" if agree is None else ("yes" if agree else "[yellow]no[/yellow]")
+        ceiling_cells = [f"{ceilings[budget]:.2f}x" if ceilings else "-" for budget in budgets]
+        table.add_row(entry["label"], f"{ms:8.3f} ms", f"{identity_ms:8.3f} ms", f"{attn_a:8.3f} ms", f"{attn_b:8.3f} ms" if attn_b else "-", agree_cell, *ceiling_cells)
+
+    if split:
+        subtitle = "[dim]rest = identity-attention variant; A = block minus rest, B = standalone attention table"
+        subtitle += f" ({attention_preset})" if attention_preset else " (no matching attention preset this run)"
+        subtitle += "; ceilings are per-block upper bounds at the given kv budget, generation adds te/vae/projections[/dim]"
+        emit(Panel(table, title=f"compute split: {family}", subtitle=subtitle, box=ROUNDED_BOX, expand=False))
+        disagreements = [config_id for config_id, row in split.items() if row["agree"] is False]
+        if disagreements:
+            emit(f"[yellow]split instruments disagree on {', '.join(disagreements)}; ceilings withheld there, treat the split with suspicion[/yellow]")
+    return split
+
+
 def measured(results, config_id):
     entry = results.get(config_id) or {}
     ms = entry.get("ms")
@@ -2661,7 +2933,7 @@ def measured(results, config_id):
 
 def best_config(results):
     # lowest error among rows within 5% of the fastest sdnq time
-    candidates = [(config_id, entry["ms"], entry["err"]) for config_id, entry in results.items() if entry.get("ms") is not None and config_id not in external_config_ids and config_id not in unsafe_config_ids]
+    candidates = [(config_id, entry["ms"], entry["err"]) for config_id, entry in results.items() if entry.get("ms") is not None and config_id not in external_config_ids and config_id not in sparse_config_ids and config_id not in unsafe_config_ids]
     if not candidates:
         return None
     fastest = min(ms for _config_id, ms, _err in candidates)
@@ -2695,7 +2967,7 @@ def select_attention_config(results):
     pool, capped = [], []
     for config_id, label, kwargs in bench_configs:
         settings = config_settings(kwargs)
-        if settings is None:
+        if settings is None or config_id in sparse_config_ids: # a sparse row shares a settings tuple with its dense row but is a stage over it, not a setting
             continue
         if settings["accum"] and settings["pv"] == "disabled":
             continue # the unsafe accumulation combo is never a candidate; the accum row cites it directly
@@ -3503,11 +3775,24 @@ def main():
             sys.exit(1)
     known_blocks = [config_id for config_id, _w, _mm, _a in block_configs]
     selected_blocks = None if args.block_configs.strip().lower() == "all" else [s.strip() for s in args.block_configs.split(",") if s.strip()]
+    selected_geometries = None if args.block_geometries.strip().lower() == "all" else [s.strip() for s in args.block_geometries.split(",") if s.strip()]
+    if selected_geometries is not None:
+        unknown_geometries = [s for s in selected_geometries if s not in block_geometries]
+        if unknown_geometries:
+            console.print(f"[red]unknown block geometry(ies): {', '.join(unknown_geometries)}; available: {', '.join(block_geometries)}[/red]")
+            sys.exit(1)
     if selected_blocks is not None:
         unknown_blocks = [s for s in selected_blocks if s not in known_blocks]
         if unknown_blocks:
             console.print(f"[red]unknown block config(s): {', '.join(unknown_blocks)}; available: {', '.join(known_blocks)}[/red]")
             sys.exit(1)
+    known_attention = [config_id for config_id, _label, _kwargs in bench_configs]
+    selected_attention = None if args.configs.strip().lower() == "all" else [s.strip() for s in args.configs.split(",") if s.strip()]
+    if selected_attention is not None:
+        unknown_attention = [s for s in selected_attention if s not in known_attention]
+        if unknown_attention:
+            console.print(f"[red]unknown attention config(s): {', '.join(unknown_attention)}; available: {', '.join(known_attention)}[/red]")
+            sys.exit(1)
     known_variants = [variant_id for variant_id, _cfg in dequant_variant_configs]
     variants_arg = args.dequant_variants.strip().lower()
     if variants_arg == "all":
@@ -3550,7 +3835,9 @@ def main():
             bench_dtype = devices.dtype
     else:
         bench_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}[args.dtype]
-    selected = list(full_run) if args.shapes.strip().lower() == "all" else [s.strip() for s in args.shapes.split(",") if s.strip()]
+    shapes_arg = args.shapes.strip().lower()
+    shape_run_aliases = {"all": full_run, "sparse": sparse_run, "gate": gate_run}
+    selected = list(shape_run_aliases[shapes_arg]) if shapes_arg in shape_run_aliases else [s.strip() for s in args.shapes.split(",") if s.strip()]
     unknown = [s for s in selected if s not in shape_presets]
     if unknown:
         console.print(f"[red]unknown shape preset(s): {', '.join(unknown)}; available: {', '.join(shape_presets)}[/red]")
@@ -3615,7 +3902,9 @@ def main():
         if free_vram_gb() < 3.0:
             emit(f"[yellow]skipping block benchmarks: needs about 3 gb free vram, {free_vram_gb():.1f} gb available[/yellow]")
         else:
-            bench_block_section(args.iters, args.warmup, config_timeout=args.config_timeout, selected=selected_blocks)
+            bench_block_section(args.iters, args.warmup, config_timeout=args.timeout_flag, selected=selected_blocks, geometries=selected_geometries)
+            if "attention" not in sections:
+                emit_block_splits()
 
     if "attention" in sections:
         # bench the prep mode the advice points to: compiled, static workaround, or eager
@@ -3637,7 +3926,8 @@ def main():
             if free_vram_gb() < needed:
                 emit(f"[yellow]skipping {preset}: needs about {needed:.0f} gb free vram, {free_vram_gb():.1f} gb available[/yellow]")
                 continue
-            all_results[preset] = bench_shape(preset, args.iters, args.warmup, position=(index, len(selected)), config_timeout=args.config_timeout, fp8_result=fp8_result)
+            all_results[preset] = bench_shape(preset, args.iters, args.warmup, position=(index, len(selected)), config_timeout=args.timeout_flag, fp8_result=fp8_result, selected=selected_attention)
+        emit_block_splits()
         build_recommendations(all_results, fp8_result, prep_status, block_results=(report.get("block") or {}).get("results"), block_variants=report.get("blocks"))
 
     if drift_samples:

modules/attention/__init__.py

@@ -1,12 +1,12 @@
 """Attention backends: one scaled_dot_product_attention router over the registered backends, the per-generation context, and the diffusers-side processor and dispatcher setup."""
 from modules.attention.registry import AttentionBackend, AttentionCall, Constraints, Platform, Registry, registry
 from modules.attention.router import Plan, PlanEntry, build_plan, get_plan, install_router, reapply, reapply_options, report
-from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, hijack_kernels, get_kernel_hijack, get_hf_api_hijack
+from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, list_dispatcher_backends, hijack_kernels, get_kernel_hijack, get_hf_api_hijack
 from modules.attention import backends, context, debug
 
 __all__ = [
     'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry',
     'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'reapply', 'reapply_options', 'report',
-    'set_diffusers_attention', 'set_attention_dispatcher', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack',
+    'set_diffusers_attention', 'set_attention_dispatcher', 'list_dispatcher_backends', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack',
     'backends', 'context', 'debug',
 ]

modules/attention/backends/flex.py

@@ -4,12 +4,18 @@ from modules.attention.registry import AttentionBackend, Constraints, Platform
 
 
 def prepare(platform: Platform, original): # pylint: disable=unused-argument
-    from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+    from torch.nn.attention.flex_attention import create_block_mask
+    from modules.attention.sparse import flex as sparse_flex
 
     def causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument
         return q_idx >= kv_idx
 
-    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=None): # pylint: disable=unused-argument
+        if selection is not None:
+            return sparse_flex.attend(query, key, value, selection, scale=scale, enable_gqa=enable_gqa)
+        # compiled, always: eager flex_attention materializes the whole score matrix, which is
+        # tens of gigabytes at video sequence lengths and fails in the driver rather than cleanly
+        flex_attention = sparse_flex.flex_call()
         score_mod = None
         block_mask = None
         if attn_mask is not None:
@@ -36,4 +42,5 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument
 backend = AttentionBackend(
     name='flex', label='Flex attention', priority=20, prepare=prepare,
     constraints=Constraints(min_ndim=4, same_device=True), # flex_attention takes 4d tensors on one device and compiles on cpu
+    caps=frozenset({'block_mask'}),
 )

modules/attention/backends/sdnq.py

@@ -1,7 +1,16 @@
+import inspect
 from modules.logger import log
 from modules.attention.registry import AttentionBackend, Constraints, Platform
 
 
+def supports_block_mask(entry) -> bool:
+    """Whether the installed sdnq takes a block mask; the chain must not promise what the kernel cannot do."""
+    try:
+        return 'block_mask' in inspect.signature(inspect.unwrap(entry)).parameters
+    except (TypeError, ValueError):
+        return False
+
+
 def prepare(platform: Platform, original): # pylint: disable=unused-argument
     from modules import shared
     from sdnq.kernels.triton_atten import sdnq_triton_atten
@@ -14,11 +23,17 @@ def prepare(platform: Platform, original): # pylint: disable=unused-argument
         'quantize_fp32': shared.opts.sdnq_attention_quantize_fp32,
         'use_fp16_accum': shared.opts.sdnq_attention_use_fp16_accum,
     }
+    block_mask = supports_block_mask(sdnq_triton_atten)
 
-    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
+    def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=None): # pylint: disable=unused-argument
+        if selection is not None:
+            return sdnq_triton_atten(query=query, key=key, value=value, attn_mask=attn_mask, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, block_mask=selection.keep, block_mask_m=selection.block_q, block_mask_n=selection.block_kv, **options)
         return sdnq_triton_atten(query=query, key=key, value=value, attn_mask=attn_mask, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **options)
 
-    log.debug(f'Torch attention: type="SDNQ attention" matmul={options["matmul_dtype"]}:{options["pv_matmul_dtype"]} smooth={options["smooth_k"]} hadamard={options["use_hadamard"]} quantize_fp32={options["quantize_fp32"]} fp16_accum={options["use_fp16_accum"]}')
+    call.caps = backend.caps if block_mask else frozenset()
+    if not block_mask and getattr(shared.opts, 'sparse_attention_enabled', False):
+        log.warning('SDNQ attention: the installed sdnq has no block mask input, sparse attention cannot use it; update the sdnq submodule')
+    log.debug(f'Torch attention: type="SDNQ attention" matmul={options["matmul_dtype"]}:{options["pv_matmul_dtype"]} smooth={options["smooth_k"]} hadamard={options["use_hadamard"]} quantize_fp32={options["quantize_fp32"]} fp16_accum={options["use_fp16_accum"]} block_mask={block_mask}')
     return call
 
 
@@ -26,4 +41,5 @@ backend = AttentionBackend(
     name='sdnq', label='SDNQ attention', priority=60, prepare=prepare,
     constraints=Constraints(min_tokens=32, min_long_side=512, min_heads=2), # sequences of 512 or fewer are text encoders, single-head calls the vae
     options=('sdnq_attention_matmul_type', 'sdnq_attention_pv_matmul_type', 'sdnq_attention_smooth_k', 'sdnq_attention_use_hadamard', 'sdnq_attention_hadamard_group_size', 'sdnq_attention_quantize_fp32', 'sdnq_attention_use_fp16_accum'),
+    caps=frozenset({'block_mask', 'masked_block'}), # the kernel takes attn_mask and block_mask together
 )

modules/attention/context.py

@@ -11,8 +11,9 @@ class GenerationContext:
     step: int = 0 # index of the denoiser forward about to run
     steps: int = 0 # forwards in the current pass
     forwards: int = 0
-    model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, telemetry only
+    model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, for telemetry and the sparse exclusion list
     step_buffer: torch.Tensor | None = None # the step as a device scalar updated in place, so compiled readers keep their graph
+    layout: object | None = None # TokenLayout published by whoever knows the packing, None until something does
 
 
 current = GenerationContext()
@@ -26,11 +27,36 @@ def denoiser_name(pipe) -> str | None:
     return None
 
 
+# every slot a pipeline can enter once per denoising step, from sd_offload_state.group_offload_main. The aux
+# components on that list (decoder, controlnet, prior) are left alone: they pack no attention sequence, and a
+# publication from one would clear the layout the denoiser just set
+DENOISER_SLOTS = ('transformer', 'unet', 'transformer_2', 'transformer_ref', 'unconditional_transformer')
+
+
+def install_layout_hook(pipe) -> None:
+    """Let a classic pipeline's denoiser publish its own packing: the modular path has its own hook, this is the rest."""
+    from modules import shared
+    if pipe is None or not getattr(shared.opts, 'sparse_attention_enabled', False):
+        return
+    from modules.attention.sparse import layout as sparse_layout
+
+    def publish(denoiser, args, kwargs): # pylint: disable=unused-argument
+        set_layout(sparse_layout.layout_from_kwargs(kwargs, denoiser.__class__.__name__))
+
+    for name in DENOISER_SLOTS:
+        module = getattr(pipe, name, None)
+        if module is None or getattr(module, 'sdnext_layout_hook', None) is not None or getattr(module, 'sdnext_state_hook', None) is not None:
+            continue
+        module.sdnext_layout_hook = module.register_forward_pre_hook(publish, with_kwargs=True)
+
+
 def begin(pipe, steps: int = 0) -> None:
     from modules import devices
     current.active = True
     current.role = 'transformer'
+    current.layout = None
     current.model_key = (pipe.__class__.__name__, denoiser_name(pipe)) if pipe is not None else None
+    install_layout_hook(pipe)
     device = devices.device if devices.device is not None else torch.device('cpu')
     if current.step_buffer is None or current.step_buffer.device != device:
         current.step_buffer = torch.zeros((), dtype=torch.int64, device=device)
@@ -56,11 +82,19 @@ def tick(step: int | None = None) -> None:
     current.forwards = current.step + 1
 
 
+def set_layout(layout) -> None:
+    """Publish what the packed sequence holds; callers that know the packing set this per forward."""
+    current.layout = layout
+
+
 def end() -> None:
+    from modules.attention import debug
     current.active = False
     current.role = None
     current.model_key = None
+    current.layout = None
     new_pass(0)
+    debug.end_generation()
 
 
 def set_role(name: str | None) -> None:

modules/attention/debug.py

@@ -6,16 +6,37 @@ from modules.attention import context
 
 enabled = os.environ.get('SD_ATTN_DEBUG', None) is not None
 seen: set[tuple] = set()
+counts: dict[tuple, int] = {}
 
 
 def observe(name: str, query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None) -> None:
-    """Log each distinct route once: backend, component role, step, shapes, dtype and mask presence."""
-    signature = (name, context.current.role, tuple(query.shape), tuple(key.shape), str(query.dtype), attn_mask is not None)
+    """Log each distinct route once: backend, component role, step, shapes, dtype, mask presence and whether the inputs are contiguous; count every call."""
+    contiguous = query.is_contiguous() and key.is_contiguous()
+    signature = (name, context.current.role, tuple(query.shape), tuple(key.shape), str(query.dtype), attn_mask is not None, contiguous)
+    counts[signature] = counts.get(signature, 0) + 1
     if signature in seen:
         return
     seen.add(signature)
-    log.debug(f'Attention route: backend={name} role={context.current.role} step={context.current.step} q={list(query.shape)} k={list(key.shape)} dtype={query.dtype} mask={attn_mask is not None}')
+    log.debug(f'Attention route: backend={name} role={context.current.role} step={context.current.step} q={list(query.shape)} k={list(key.shape)} dtype={query.dtype} mask={attn_mask is not None} contiguous={contiguous}')
+
+
+def summary() -> list[str]:
+    """One line per route with its call count since the last generation, busiest first."""
+    lines = []
+    for signature, count in sorted(counts.items(), key=lambda item: -item[1]):
+        name, role, q_shape, k_shape, dtype, masked, contiguous = signature
+        lines.append(f'backend={name} role={role} q={list(q_shape)} k={list(k_shape)} dtype={dtype} mask={masked} contiguous={contiguous} calls={count}')
+    return lines
+
+
+def end_generation() -> None:
+    """Log the route counts of the generation that just ended and start the next count."""
+    if enabled and counts:
+        for line in summary():
+            log.debug(f'Attention routes: {line}')
+    counts.clear()
 
 
 def reset() -> None:
     seen.clear()
+    counts.clear()

modules/attention/dispatcher.py

@@ -94,3 +94,13 @@ def set_attention_dispatcher(pipe):
         log.warning(f'Attention dispatcher: active={prev[0].value} list={backends} target={attn} not found')
     else:
         log.debug(f'Attention dispatcher: active={prev[0].value} list={backends}')
+
+
+def list_dispatcher_backends() -> list:
+    """The kernels diffusers can dispatch attention to, for anything that offers hf_attention as a choice."""
+    try:
+        from diffusers.models import attention_dispatch as a
+        return sorted(b.value for b in a._AttentionBackendRegistry.list_backends()) # pylint: disable=protected-access
+    except Exception as e:
+        log.error(f'Attention dispatcher: {e}')
+        return []

modules/attention/registry.py

@@ -64,6 +64,7 @@ class AttentionBackend:
     terminal: bool = False # serves every call the entries decline, in place of the original sdpa
     platforms: frozenset[str] | None = None # devices backends the implementation exists for, None for all
     options: tuple[str, ...] = () # settings the prepared call captures; a change to one rebuilds the chain
+    caps: frozenset[str] = frozenset() # what the call can consume beyond plain sdpa arguments: 'block_mask', and 'masked_block' when it composes one with a token mask
 
     def available_on(self, platform: Platform) -> bool:
         return self.platforms is None or platform.backend in self.platforms
@@ -94,5 +95,8 @@ class Registry:
     def options(self) -> list[str]:
         return sorted({name for backend in self.backends.values() for name in backend.options})
 
+    def with_cap(self, cap: str) -> list[AttentionBackend]:
+        return [backend for backend in self.ordered() if cap in backend.caps]
+
 
 registry = Registry()

modules/attention/router.py

@@ -13,6 +13,7 @@ from modules.attention.registry import AttentionBackend, AttentionCall, Platform
 class PlanEntry:
     backend: AttentionBackend
     call: AttentionCall
+    caps: frozenset[str] = frozenset() # the backend's declared caps, narrowed to what prepare verified in the installed implementation
 
 
 @dataclass(frozen=True)
@@ -50,7 +51,7 @@ def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registr
             continue
         if call is None:
             continue
-        entry = PlanEntry(backend=backend, call=call)
+        entry = PlanEntry(backend=backend, call=call, caps=backend.caps & frozenset(getattr(call, 'caps', backend.caps)))
         if backend.terminal:
             terminal = entry
         else:
@@ -59,7 +60,7 @@ def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registr
     return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels))
 
 
-def make_router(plan: Plan, observer: Callable | None = None) -> AttentionCall:
+def make_router(plan: Plan, observer: Callable | None = None, stage: Callable | None = None) -> AttentionCall:
     entries = plan.entries
     terminal = plan.terminal.call if plan.terminal is not None else None
     terminal_name = plan.terminal.backend.name if plan.terminal is not None else 'sdpa'
@@ -69,10 +70,16 @@ def make_router(plan: Plan, observer: Callable | None = None) -> AttentionCall:
     def sdpa_router(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs):
         for entry in entries:
             if entry.backend.constraints.accepts(query, key, value, attn_mask):
+                if stage is not None and 'block_mask' in entry.caps:
+                    selection = stage(query, key, value, attn_mask, is_causal, entry.caps)
+                    if selection is not None:
+                        if observer is not None:
+                            observer(f'{entry.backend.name}+sparse', query, key, attn_mask)
+                        return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=selection)
                 if observer is not None:
                     observer(entry.backend.name, query, key, attn_mask)
                 return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa)
-        if observer is not None:
+        if observer is not None: # pylint: disable=duplicate-code
             observer(terminal_name, query, key, attn_mask)
         if terminal is not None:
             return terminal(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **kwargs)
@@ -89,21 +96,43 @@ def install_router(labels, platform: Platform, original: AttentionCall, reg: Reg
     plan = build_plan(labels, platform, original, reg)
     debug.reset()
     observer = debug.observe if debug.enabled else None
-    torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer) if (plan.entries or plan.terminal is not None) else original
+    stage = build_sparse_stage(plan)
+    torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer, stage) if (plan.entries or plan.terminal is not None) else original
     current_plan = plan
     torch_info.set(attention='>'.join(plan.chain()))
-    log.debug(f'Torch attention: chain={">".join(plan.chain())} overrides={list(labels)} backend={platform.backend}')
+    log.debug(f'Torch attention: chain={">".join(plan.chain())} overrides={list(labels)} backend={platform.backend} sparse={stage is not None}')
     return plan
 
 
+def build_sparse_stage(plan: Plan):
+    """Sparse attention is a stage over the chain rather than a chain member, so it needs a backend in the chain that consumes a block mask."""
+    from modules.attention.sparse import stage as sparse_stage
+    try:
+        options = sparse_stage.read_options()
+    except Exception:
+        return None
+    if not options.enabled:
+        return None
+    capable = [entry.backend.name for entry in plan.entries if 'block_mask' in entry.caps]
+    if not capable:
+        names = [backend.label for backend in default_registry.with_cap('block_mask')]
+        log.warning(f'Sparse attention: enabled but no active backend consumes a block mask, enable one of {names} in sdp overrides; attention stays dense')
+        return None
+    built = sparse_stage.make_stage(options)
+    if built is not None:
+        log.info(f'Sparse attention: backend={capable[0]} budget={options.budget:.0%} gate={options.min_tokens} schedule={options.schedule_steps}x+{options.schedule_bump:.0%}')
+    return built
+
+
 def get_plan() -> Plan | None:
     return current_plan
 
 
 def reapply_options(reg: Registry | None = None) -> list[str]:
-    """Settings whose change rebuilds the chain: the override set, the torch kernel flags, and every option a backend captures."""
+    """Settings whose change rebuilds the chain: the override set, the torch kernel flags, every option a backend captures, and the sparse stage."""
+    from modules.attention.sparse import stage as sparse_stage
     reg = reg if reg is not None else default_registry
-    return ['sdp_options', 'sdp_overrides', *reg.options()]
+    return ['sdp_options', 'sdp_overrides', *reg.options(), *sparse_stage.OPTION_NAMES]
 
 
 def reapply() -> None:
@@ -117,12 +146,22 @@ def reapply() -> None:
 
 
 def report() -> dict:
-    """The active chain and generation context, for the api and the debug log."""
+    """The active chain, sparse stage and generation context, for the api and the debug log."""
+    from modules.attention.sparse import stage as sparse_stage
     plan = current_plan
     state = context.current
+    options = sparse_stage.read_options()
+    layout = state.layout
     return {
         'chain': plan.chain() if plan is not None else ['sdpa'],
         'overrides': list(plan.labels) if plan is not None else [],
+        'sparse': {
+            'enabled': options.enabled,
+            'budget': options.budget,
+            'gate': options.min_tokens,
+            'capable': [entry.backend.name for entry in plan.entries if 'block_mask' in entry.caps] if plan is not None else [],
+            'layout': {'source': layout.source, 'kinds': list(layout.kinds()), 'length': layout.length} if layout is not None else None,
+        },
         'backend': plan.platform.backend if plan is not None else None,
         'context': {'active': state.active, 'role': state.role, 'step': state.step, 'steps': state.steps, 'model': state.model_key},
     }

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_kwargs, layout_from_prefix, layout_from_segments, publish_segments, segments_from_live
+
+__all__ = [
+    'BlockSelection', 'SparseSpec', 'block_count', 'radial_blocks', 'schedule', 'select_blocks',
+    'Span', 'TokenLayout', 'block_pins', 'layout_from_index_kwargs', 'layout_from_kwargs', 'layout_from_prefix', 'layout_from_segments', 'publish_segments', 'segments_from_live',
+]

modules/attention/sparse/flex.py

@@ -0,0 +1,39 @@
+"""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)}')
+    # the partial slots stay empty by construction, so build them directly rather than sorting a mask of zeros
+    empty_num = torch.zeros(keep.shape[:-1], dtype=torch.int32, device=keep.device)
+    empty_indices = torch.zeros(keep.shape, dtype=torch.int32, device=keep.device)
+    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,172 @@
+"""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 | None = None) -> 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 if length is not None else spans[-1].end, source='indices')
+
+
+# how an architecture orders its joint sequence, which the call itself does not reveal. Verified against the
+# diffusers transformers that take txt_ids and img_ids; HiDream packs image first and is deliberately absent, so
+# it falls back rather than being pinned backwards. An unlisted class publishes nothing.
+JOINT_TEXT_FIRST = frozenset({
+    'FluxTransformer2DModel', 'Flux2Transformer2DModel', 'ChromaTransformer2DModel', 'BriaTransformer2DModel',
+    'BriaFiboTransformer2DModel', 'LongCatImageTransformer2DModel', 'OvisImageTransformer2DModel',
+})
+
+
+def layout_from_stream_ids(kwargs: dict, cls_name: str | None) -> TokenLayout | None:
+    """Read the stream lengths off the rotary id tensors a joint transformer is given by name."""
+    if cls_name not in JOINT_TEXT_FIRST:
+        return None
+    text, image = kwargs.get('txt_ids'), kwargs.get('img_ids')
+    if not torch.is_tensor(text) or not torch.is_tensor(image) or text.dim() < 2 or image.dim() < 2:
+        return None
+    return layout_from_segments((('text', text.shape[-2]), ('image', image.shape[-2])), source='stream-ids')
+
+
+def layout_from_kwargs(kwargs: dict, cls_name: str | None = None) -> TokenLayout | None:
+    """Whatever the denoiser says about its own packing, by whichever convention it uses."""
+    return layout_from_index_kwargs(kwargs or {}) or layout_from_stream_ids(kwargs or {}, cls_name)
+
+
+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 segments_from_live(live: torch.Tensor, kind: str, pad_kind: str = 'pad') -> list[tuple[str, int]]:
+    """Run length encode a boolean live mask into ordered (kind, count) pairs, the dead runs labelled as padding."""
+    values = live.detach().to('cpu').bool()
+    if values.numel() == 0:
+        return []
+    changes = (values[1:] != values[:-1]).nonzero().flatten().tolist()
+    bounds = [0, *[c + 1 for c in changes], values.numel()]
+    return [(kind if bool(values[bounds[i]]) else pad_kind, bounds[i + 1] - bounds[i]) for i in range(len(bounds) - 1)]
+
+
+def publish_segments(segments, length: int | None = None, source: str = 'segments') -> None:
+    """Publish a layout from the site that packs the sequence, which is the only place the segment lengths are all known."""
+    from modules.attention import context
+    context.set_layout(layout_from_segments(segments, length=length, 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,147 @@
+"""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
+
+    @property
+    def shape(self) -> tuple[int, int, int, int]:
+        return tuple(self.keep.shape)
+
+    def density(self) -> float:
+        """Fraction of tiles kept. Reads back from the accelerator, so this is for reporting and tests, never the hot path."""
+        return float(self.keep.sum().item()) / max(self.keep.numel(), 1)
+
+
+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)
+
+
+plan_cache: dict = {}
+
+
+def selection_plan(spec: SparseSpec, nq: int, nk: int, pins, drops, device, cache_key=None):
+    """The parts that depend only on geometry and layout, not on the tensors: what must be kept, what may be chosen, and how many."""
+    key = (cache_key, nq, nk, spec.block_q, spec.block_kv, spec.budget, str(device))
+    hit = plan_cache.get(key) if cache_key is not None else None
+    if hit is not None:
+        return hit
+    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)
+    covers_everything = bool((keep_per_row >= per_row).all()) # one readback, amortized over the generation by the cache
+    built = (must, forbidden, candidates, keep_per_row, covers_everything)
+    if cache_key is not None:
+        if len(plan_cache) > 32:
+            plan_cache.clear()
+        plan_cache[key] = built
+    return built
+
+
+def select_blocks(query: torch.Tensor, key: torch.Tensor, spec: SparseSpec, pins: torch.Tensor | None = None, drops: torch.Tensor | None = None, cache_key=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, forbidden, candidates, keep_per_row, covers_everything = selection_plan(spec, nq, nk, pins, drops, device, cache_key)
+    if covers_everything and not spec.force:
+        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'))
+    # rank rather than topk, so the per row budget varies without a host side k
+    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 = must | ((rank < keep_per_row) & candidates)
+    keep &= ~forbidden
+    return BlockSelection(keep=keep.to(torch.int8), block_q=spec.block_q, block_kv=spec.block_kv, budget=spec.budget, seq_q=seq_q, seq_kv=seq_kv)
+
+
+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)
+
+
+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)])

modules/attention/sparse/stage.py

@@ -0,0 +1,173 @@
+"""The router stage that turns settings plus a published layout into a per call block selection."""
+import os
+from dataclasses import dataclass
+import torch
+from modules.logger import log
+from modules.attention import context
+from modules.attention.sparse import layout as layout_mod
+from modules.attention.sparse.selector import BlockSelection, SparseSpec, block_count, radial_blocks, schedule, select_blocks
+
+# SD_SPARSE_PATTERN=radial replaces the content aware selection with a static band around the
+# diagonal at the same density: the control the selector has to beat, and the fallback if it does not
+pattern = os.environ.get('SD_SPARSE_PATTERN', 'adaptive').strip().lower()
+
+
+# measured on a 3090: below roughly this length a 30 percent budget caps under 1.25x per block,
+# so the selector cannot pay for itself; see docs/sparse-attention-tracker.md. The settings registry
+# carries the same number as the option default, this is the fallback when the option is absent
+DEFAULT_MIN_TOKENS = 8192
+
+# settings the stage reads, so a change to any of them rebuilds the chain
+OPTION_NAMES = ('sparse_attention_enabled', 'sparse_attention_budget', 'sparse_attention_min_tokens', 'sparse_attention_schedule_steps', 'sparse_attention_schedule_bump', 'sparse_attention_head_shared', 'sparse_attention_exclude')
+
+
+@dataclass(frozen=True)
+class StageOptions:
+    enabled: bool = False
+    budget: float = 0.30
+    min_tokens: int = DEFAULT_MIN_TOKENS # 0 sparsifies every sequence that reaches the stage
+    schedule_steps: int = 0
+    schedule_bump: float = 0.0
+    head_shared: bool = False
+    exclude: tuple = () # architectures, pipeline classes or denoiser classes that stay dense
+
+
+def read_options() -> StageOptions:
+    from modules import shared
+    opts = shared.opts
+    return StageOptions(
+        enabled=bool(getattr(opts, 'sparse_attention_enabled', False)),
+        budget=float(getattr(opts, 'sparse_attention_budget', 30)) / 100.0,
+        min_tokens=int(getattr(opts, 'sparse_attention_min_tokens', DEFAULT_MIN_TOKENS)),
+        schedule_steps=int(getattr(opts, 'sparse_attention_schedule_steps', 0)),
+        schedule_bump=float(getattr(opts, 'sparse_attention_schedule_bump', 0)) / 100.0,
+        head_shared=bool(getattr(opts, 'sparse_attention_head_shared', False)),
+        exclude=parse_exclusions(getattr(opts, 'sparse_attention_exclude', '')),
+    )
+
+
+def parse_exclusions(raw) -> tuple:
+    """The exclusion list as lowercase entries, each naming an architecture, a pipeline class or a denoiser class."""
+    return tuple(entry.strip().lower() for entry in str(raw or '').split(',') if len(entry.strip()) > 0)
+
+
+def match_exclusion(model_key, arch, exclude: tuple) -> str:
+    """The entry the loaded model matches, empty when none of them do."""
+    names = {str(name).lower() for name in (*(model_key or ()), arch) if name}
+    return next((entry for entry in exclude if entry in names), '')
+
+
+def resolve_layout(seq: int, reported: set) -> layout_mod.TokenLayout:
+    """The published layout when there is one, otherwise sparsify the whole sequence and say so once."""
+    published = context.current.layout
+    if isinstance(published, layout_mod.TokenLayout) and published.length == seq:
+        return published
+    if seq not in reported:
+        reported.add(seq)
+        detail = 'none published' if published is None else f'published length {getattr(published, "length", None)} does not match {seq}'
+        log.info(f'Sparse attention: no token layout ({detail}), sparsifying the whole sequence at tokens={seq}')
+    return layout_mod.layout_from_prefix(seq, 0)
+
+
+def make_stage(options: StageOptions):
+    """Return the per call selector, or None when the feature is off."""
+    if not options.enabled or options.budget >= 1.0:
+        return None
+    reported: set = set()
+    inactive: set = set()
+    notified: set = set()
+    cache: dict = {}
+    static: dict = {}
+    excluded: dict = {}
+
+    def static_selection(query, key, spec, pins, drops, cache_key):
+        """A density matched band, built once per geometry, honoring the same layout pins so the control differs from the selector only in how it chooses video tiles."""
+        static_key = (cache_key, spec.budget, query.shape[-2], key.shape[-2])
+        built = static.get(static_key)
+        if built is None:
+            reference = select_blocks(query, key, spec, pins=pins, drops=drops, cache_key=cache_key)
+            if reference is None:
+                return None
+            target = reference.density()
+            pinned = float(pins.to(torch.float32).mean().item()) if pins is not None else 0.0
+            band = radial_blocks(query.shape[-2], key.shape[-2], max(target - pinned, 0.0), spec, query.device)
+            keep = band.keep.bool()
+            if pins is not None:
+                keep = keep | pins
+            if drops is not None:
+                keep = keep & ~drops
+            built = BlockSelection(keep=keep.to(torch.int8), block_q=spec.block_q, block_kv=spec.block_kv, budget=spec.budget, seq_q=query.shape[-2], seq_kv=key.shape[-2])
+            static.clear()
+            static[static_key] = built
+            log.info(f'Sparse attention: static radial pattern density={built.density():.3f} against selector {target:.3f} at budget={spec.budget:.0%}')
+        return built
+
+    def budget_for_step() -> float:
+        state = context.current
+        if options.schedule_steps <= 0 or options.schedule_bump <= 0 or state.steps <= 0:
+            return options.budget
+        key = (state.steps, options.budget, options.schedule_bump, options.schedule_steps)
+        table = cache.get(key)
+        if table is None:
+            table = schedule(state.steps, options.budget, options.schedule_bump, options.schedule_steps)
+            cache.clear()
+            cache[key] = table
+        return table[min(state.step, len(table) - 1)] if table else options.budget
+
+    def on_exclusion_list(state) -> bool:
+        """Whether the loaded model is excluded, resolved once per model since the answer cannot change within one."""
+        if not options.exclude:
+            return False
+        hit = excluded.get(state.model_key)
+        if hit is None:
+            from modules import shared
+            hit = match_exclusion(state.model_key, getattr(shared, 'sd_model_type', None), options.exclude)
+            excluded[state.model_key] = hit
+            if hit: # an enabled setting that cannot act says so rather than doing nothing quietly
+                log.info(f'Sparse attention: "{hit}" is on the exclusion list; attention stays dense')
+        return len(hit) > 0
+
+    def decline(reason: str):
+        stage.last_skip = reason
+        return None
+
+    def stage(query, key, value, attn_mask, is_causal, caps=frozenset()): # pylint: disable=unused-argument
+        state = context.current
+        if state.role != 'transformer' or not state.active:
+            return decline('not the denoiser')
+        if on_exclusion_list(state):
+            return decline('excluded')
+        if is_causal: # the selection keeps the diagonal but encodes no causality
+            return decline('causal')
+        if attn_mask is not None and 'masked_block' not in caps: # flex would need a mask_mod to combine the two
+            if 'masked' not in notified: # an enabled setting that cannot act says so rather than doing nothing quietly
+                notified.add('masked')
+                log.info('Sparse attention: this model passes an attention mask and the serving backend cannot combine it with a block selection; attention stays dense')
+            return decline('masked')
+        if query.device.type == 'cpu' or query.dim() != 4:
+            return decline('unsupported tensor')
+        seq_q, seq_kv = query.shape[-2], key.shape[-2]
+        if seq_q != seq_kv: # cross attention is short and already cheap
+            return decline('cross attention')
+        if seq_q < options.min_tokens:
+            if seq_q not in inactive: # an enabled setting that cannot act says so rather than doing nothing quietly
+                inactive.add(seq_q)
+                log.info(f'Sparse attention: inactive at tokens={seq_q}, below the minimum sequence of {options.min_tokens}; attention stays dense')
+            return decline('below the minimum sequence')
+        budget = budget_for_step()
+        if budget >= 1.0:
+            return decline('budget covers everything')
+        spec = SparseSpec(budget=budget, head_shared=options.head_shared)
+        token_layout = resolve_layout(seq_q, reported)
+        nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
+        pins, drops = layout_mod.block_pins(token_layout, seq_q, seq_kv, spec.block_q, spec.block_kv, query.device)
+        if pins.shape[-2:] != (nq, nk):
+            return decline('layout geometry mismatch')
+        stage.last_skip = None
+        if pattern == 'radial':
+            return static_selection(query, key, spec, pins, drops, token_layout.key())
+        return select_blocks(query, key, spec, pins=pins, drops=drops, cache_key=token_layout.key())
+
+    stage.options = options
+    stage.last_skip = None
+    return stage

modules/sd_hijack_modular.py

@@ -75,6 +75,15 @@ class InterruptLogFilter(logging.Filter):
         return 'Interrupted...' not in record.msg
 
 
+def publish_layout(kwargs):
+    """Hand the attention router whatever the pipeline says about its packed sequence, keyed on the *_indices tensors rather than the model."""
+    try:
+        from modules.attention.sparse import layout as sparse_layout
+        attention_context.set_layout(sparse_layout.layout_from_index_kwargs(kwargs or {}))
+    except Exception as e:
+        log.debug(f'Pipeline: token layout {e}')
+
+
 def install_state_hook(pipe):
     runner_log = logging.getLogger('diffusers.modular_pipelines.modular_pipeline')
     if not any(isinstance(f, InterruptLogFilter) for f in runner_log.filters):
@@ -91,9 +100,10 @@ def install_state_hook(pipe):
             return True
         return False
 
-    def _pre_transformer_hook(module, args): # pylint: disable=unused-argument
+    def _pre_transformer_hook(module, args, kwargs): # pylint: disable=unused-argument
         new_phase = set_phase('Generate', module)
         attention_context.set_role('transformer')
+        publish_layout(kwargs)
         if new_phase:
             sd_offload.offload_ondemand(pipe, exclude=['transformer', 'transformer_ref'], reason='generate', force=hasattr(pipe, 'sdnext_force_offload'))
         if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0:
@@ -138,7 +148,8 @@ def install_state_hook(pipe):
         if module is not None:
             target = getattr(module, 'model', module) # conditioning calls the inner model directly
             if isinstance(target, torch.nn.Module) and getattr(target, 'sdnext_state_hook', None) is None:
-                target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook)
+                # with_kwargs, because the blocks call the transformer entirely by keyword and the token layout rides in those kwargs
+                target.sdnext_state_hook = target.register_forward_pre_hook(_pre_transformer_hook, with_kwargs=True)
 
     for name in ('text_encoder', 'text_encoder_2'):
         module = getattr(pipe, name, None)

modules/ui_definitions.py

@@ -262,6 +262,15 @@ def create_settings(cmd_opts):
         "sdnq_attention_pv_matmul_type": OptionInfo("disabled", "SDNQ Attention PV MatMul type", gr.Radio, {"choices": sdnq_matmul_modes}),
         "sdnq_attention_hadamard_group_size": OptionInfo(256, "SDNQ Attention Hadamard Group Size", gr.Slider, {"minimum": 4, "maximum": 1024, "step": 1}),
 
+        "sparse_attention_sep": OptionInfo("<h2>Sparse Attention</h2>", "", gr.HTML),
+        "sparse_attention_enabled": OptionInfo(False, "Sparse Attention", gr.Checkbox),
+        "sparse_attention_budget": OptionInfo(30, "Sparse Attention KV budget", gr.Slider, {"minimum": 5, "maximum": 100, "step": 5}),
+        "sparse_attention_min_tokens": OptionInfo(8192, "Sparse Attention minimum sequence", gr.Slider, {"minimum": 0, "maximum": 32768, "step": 1024}),
+        "sparse_attention_schedule_steps": OptionInfo(0, "Sparse Attention dense steps", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
+        "sparse_attention_schedule_bump": OptionInfo(30, "Sparse Attention dense step bonus", gr.Slider, {"minimum": 0, "maximum": 70, "step": 5}),
+        "sparse_attention_head_shared": OptionInfo(False, "Sparse Attention share selection across heads", gr.Checkbox),
+        "sparse_attention_exclude": OptionInfo("CosmosTransformer3DModel", "Sparse Attention excluded models", gr.Textbox),
+
         "hf_attention_sep": OptionInfo("<h2>Attention Dispatcher</h2>", "", gr.HTML),
         "hf_attention": OptionInfo('', "Attention dispatcher kernel", gr.Textbox),
     }))

pipelines/krea2/transformer_krea2.py

@@ -342,6 +342,16 @@ class Krea2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOri
         else:
             mask = None
 
+        from modules.attention.sparse import layout as sparse_layout # delayed, this module is also importable without the webui
+        if attention_mask is not None:
+            # the text stream is padded to a fixed length and the joint sequence to a multiple of 256, so the live
+            # runs are what the layout must report; a wholly padded key block is dropped rather than pinned
+            live = attention_mask.any(dim=0)
+            segments = sparse_layout.segments_from_live(live[:txtlen], "text") + sparse_layout.segments_from_live(live[txtlen:], "image")
+        else:
+            segments = [("text", txtlen), ("image", imglen)]
+        sparse_layout.publish_segments(segments, source="krea2")
+
         freqs = self.posemb(position_ids)
 
         for block in self.blocks:

scripts/xyz/xyz_grid_classes.py

@@ -3,6 +3,12 @@ from scripts.xyz.xyz_grid_shared import ( # pylint: disable=no-name-in-module, u
     apply_task_arg,
     apply_task_args,
     apply_setting,
+    apply_attention,
+    apply_attention_overrides,
+    apply_attention_dispatcher,
+    list_sdp_overrides,
+    save_attention,
+    restore_attention,
     apply_prompt_primary,
     apply_prompt_refine,
     apply_prompt_detailer,
@@ -43,7 +49,7 @@ from scripts.xyz.xyz_grid_shared import ( # pylint: disable=no-name-in-module, u
     format_nothing,
     str_permutations,
  )
-from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet
+from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet, attention
 from modules.control.units import controlnet, t2iadapter
 from modules.control import processor
 
@@ -114,6 +120,7 @@ class SharedSettingsStackHelper():
     disable_apply_metadata = None
     disable_apply_params = None
     sdnq_quant_mode = None
+    attention_settings = None
 
     def __enter__(self):
         # Save overridden settings so they can be restored later
@@ -151,6 +158,7 @@ class SharedSettingsStackHelper():
         self.disable_apply_metadata = shared.opts.disable_apply_metadata
         self.disable_apply_params = shared.opts.disable_apply_params
         self.sdnq_quant_mode = shared.opts.sdnq_quantize_weights_mode
+        self.attention_settings = save_attention()
         shared.opts.data["disable_apply_metadata"] = []
         shared.opts.data["disable_apply_params"] = ''
 
@@ -203,6 +211,7 @@ class SharedSettingsStackHelper():
         if self.sdnq_quant_mode != shared.opts.sdnq_quantize_weights_mode:
             shared.opts.data["sdnq_quantize_weights_mode"] = self.sdnq_quant_mode
             sd_models.reload_model_weights(op='model')
+        restore_attention(self.attention_settings)
 
 
 axis_options = [
@@ -270,6 +279,20 @@ axis_options = [
     AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")),
     AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared_items.sdnq_quant_modes)),
     AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared_items.sdnq_quant_modes)),
+    AxisOption("[Attention] Method", str, apply_setting('cross_attention_optimization'), cost=0.2, choices=shared_items.list_crossattention),
+    AxisOption("[Attention] SDP override", str, apply_attention_overrides, cost=0.2, choices=list_sdp_overrides),
+    AxisOption("[Attention] Dispatcher", str, apply_attention_dispatcher, cost=0.2, choices=lambda: ['None'] + attention.list_dispatcher_backends()),
+    AxisOption("[Attention] SDNQ matmul", str, apply_attention('sdnq_attention_matmul_type'), cost=0.2, choices=lambda: list(shared_items.sdnq_matmul_modes)),
+    AxisOption("[Attention] SDNQ PV matmul", str, apply_attention('sdnq_attention_pv_matmul_type'), cost=0.2, choices=lambda: list(shared_items.sdnq_matmul_modes)),
+    AxisOption("[Attention] SDNQ smooth K", str, apply_attention('sdnq_attention_smooth_k'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Attention] SDNQ hadamard", str, apply_attention('sdnq_attention_use_hadamard'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Attention] SDNQ fp16 accumulation", str, apply_attention('sdnq_attention_use_fp16_accum'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Sparse] Enabled", str, apply_attention('sparse_attention_enabled'), cost=0.2, choices=lambda: ['False', 'True']),
+    AxisOption("[Sparse] KV budget", int, apply_attention('sparse_attention_budget'), cost=0.2),
+    AxisOption("[Sparse] Minimum sequence", int, apply_attention('sparse_attention_min_tokens'), cost=0.2),
+    AxisOption("[Sparse] Dense steps", int, apply_attention('sparse_attention_schedule_steps'), cost=0.2),
+    AxisOption("[Sparse] Dense step bonus", int, apply_attention('sparse_attention_schedule_bump'), cost=0.2),
+    AxisOption("[Sparse] Shared heads", str, apply_attention('sparse_attention_head_shared'), cost=0.2, choices=lambda: ['False', 'True']),
     AxisOption("[HDR] Mode", int, apply_field("hdr_mode")),
     AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")),
     AxisOption("[HDR] Color", float, apply_field("hdr_color")),

scripts/xyz/xyz_grid_shared.py

@@ -81,6 +81,75 @@ def apply_setting(field):
     return fun
 
 
+def attention_options() -> list:
+    """Attention settings an axis can change; the stack helper restores exactly this set."""
+    from modules import attention
+    return ['cross_attention_optimization', 'hf_attention', *attention.reapply_options()]
+
+
+def list_sdp_overrides() -> list:
+    item = shared.opts.data_labels.get('sdp_overrides', None)
+    args = item.component_args if item is not None else None
+    args = args() if callable(args) else args
+    return ['None'] + list((args or {}).get('choices', None) or [])
+
+
+def apply_attention(field):
+    def fun(p, x, xs):
+        from modules import attention
+        apply_setting(field)(p, x, xs)
+        attention.reapply() # backends read their settings when the chain is built, so a write on its own changes nothing
+        owner = next((backend for backend in attention.registry.backends.values() if field in backend.options), None)
+        plan = attention.get_plan()
+        if owner is not None and plan is not None and owner.name not in plan.chain():
+            log.warning(f'XYZ grid apply attention: {field} is read by "{owner.label}" which is not in the active chain={plan.chain()}')
+    return fun
+
+
+def apply_attention_overrides(p, x, xs):
+    from modules import attention
+    labels = [label.strip() for label in str(x).split('+') if len(label.strip()) > 0 and label.strip().lower() != 'none']
+    unknown = [label for label in labels if attention.registry.by_label(label) is None]
+    if len(unknown) > 0:
+        log.warning(f'XYZ grid apply attention: unknown overrides={unknown} available={attention.registry.labels()}')
+    shared.opts.data['sdp_overrides'] = labels
+    attention.reapply()
+    log.debug(f'XYZ grid apply attention: overrides={labels}')
+
+
+def apply_attention_dispatcher(p, x, xs):
+    from modules import attention
+    value = '' if str(x).strip().lower() in ['none', 'default'] else str(x).strip()
+    shared.opts.data['hf_attention'] = value
+    if shared.sd_loaded:
+        attention.set_attention_dispatcher(shared.sd_model)
+    log.debug(f'XYZ grid apply attention: dispatcher="{value}"')
+
+
+def save_attention() -> dict:
+    return {field: shared.opts.data[field] for field in attention_options() if field in shared.opts.data}
+
+
+def restore_attention(saved: dict):
+    """Put back whatever an attention axis changed, keys it introduced included, then rebuild what reads them."""
+    from modules import attention
+    changed = []
+    for field in attention_options():
+        if (field in saved) == (field in shared.opts.data) and saved.get(field, None) == shared.opts.data.get(field, None):
+            continue
+        changed.append(field)
+        if field in saved:
+            shared.opts.data[field] = saved[field]
+        else:
+            shared.opts.data.pop(field, None)
+    if len(changed) == 0:
+        return
+    attention.reapply()
+    if 'hf_attention' in changed and shared.sd_loaded:
+        attention.set_attention_dispatcher(shared.sd_model)
+    log.debug(f'XYZ grid restore attention: {changed}')
+
+
 def apply_seed(p, x, xs):
     p.seed = x
     p.all_seeds = None

test/test-attention-router.py

@@ -284,6 +284,40 @@ def test_prepare_failure_skips_backend():
     return True
 
 
+def test_prepared_call_narrows_caps():
+    # a backend declares what it can consume; prepare may narrow that to what the installed
+    # implementation verified, never widen it, and the router hands a selection only to survivors
+    reg = attention.Registry()
+
+    def add(name, declared, verified=None):
+        def prepare(platform, original): # pylint: disable=unused-argument
+            def call(*args, **kwargs): # pylint: disable=unused-argument
+                return name
+            if verified is not None:
+                call.caps = verified
+            return call
+        reg.register(attention.AttentionBackend(name=name, label=f'{name} attention', priority=10, prepare=prepare, caps=declared))
+
+    add('declared', frozenset({'block_mask'}))
+    add('narrowed', frozenset({'block_mask'}), frozenset())
+    add('widened', frozenset(), frozenset({'block_mask'}))
+    platform = attention.Platform(backend='cuda')
+    plan = attention.build_plan(['declared attention', 'narrowed attention', 'widened attention'], platform, sdpa_stub, reg)
+    caps = {entry.backend.name: entry.caps for entry in plan.entries}
+    assert caps == {'declared': frozenset({'block_mask'}), 'narrowed': frozenset(), 'widened': frozenset()}, caps
+    q = shaped((1, 8, 1024, 64))
+    calls = []
+
+    def stage(*args): # pylint: disable=unused-argument
+        calls.append('stage')
+        return 'selection'
+    router = attention_router.make_router(attention.build_plan(['narrowed attention'], platform, sdpa_stub, reg), stage=stage)
+    assert router(q, q, q) == 'narrowed' and not calls, calls
+    router = attention_router.make_router(attention.build_plan(['declared attention'], platform, sdpa_stub, reg), stage=stage)
+    assert router(q, q, q) == 'declared' and calls == ['stage'], calls
+    return True
+
+
 def test_install_router_keeps_original_for_empty_plan():
     saved = torch.nn.functional.scaled_dot_product_attention
     saved_plan = attention_router.current_plan
@@ -318,12 +352,15 @@ def test_dynamic_backend_pins_pre_dynamic_sdpa():
 def test_context_classic_ticks_follow_the_callback():
     ctx = attention.context
 
+    class Denoiser(torch.nn.Module): # begin installs the layout pre-hook, so the stand-in has to accept one
+        pass
+
     class Pipe:
-        transformer = object()
+        transformer = Denoiser()
 
     ctx.begin(Pipe(), steps=4)
     assert ctx.current.active and ctx.current.role == 'transformer' and ctx.current.step == 0 and ctx.current.steps == 4
-    assert ctx.current.model_key == ('Pipe', 'object'), ctx.current.model_key
+    assert ctx.current.model_key == ('Pipe', 'Denoiser'), ctx.current.model_key
     buffer = ctx.current.step_buffer
     for completed in range(4):
         ctx.tick(completed + 1) # the diffusers callback reports the step just completed
@@ -515,6 +552,7 @@ def run_all():
         test_choices_match_backends,
         test_router_dispatch_prefers_priority_then_terminal_then_original,
         test_prepare_failure_skips_backend,
+        test_prepared_call_narrows_caps,
         test_install_router_keeps_original_for_empty_plan,
         test_dynamic_backend_pins_pre_dynamic_sdpa,
     ]:

test/test-attention-sdnq-sparse.py

@@ -0,0 +1,688 @@
+#!/usr/bin/env python
+"""
+Offline tests for the block mask input of the sdnq Triton attention kernel.
+
+Covers:
+
+- the kernel contract through the raw triton op, which takes the per query block count and
+  ascending index list of kept kv blocks that get_block_mask_input builds from the int8 mask: a
+  block mask is within one ulp of the same kernel fed the token-expanded dense mask, over a ragged
+  tail with sub tiles past the end of the sequence, GQA, batch and head broadcasts, the contiguous padded list contract
+  and an empty block row; bitwise against no mask with every block kept, and bitwise against the
+  token mask when both arms carry one. The token-mask path has its own -inf-safe softmax
+  update, which the compiler rounds differently by up to one ulp depending on tile and dtype; the
+  block path shares the dense arithmetic, which the all-ones row proves
+- the launcher's validation of the block mask
+- the nesting filter: prune_configs keeps only tiles that sit inside one mask block and raises,
+  naming the env vars, when none do
+- the public entry: sdnq_triton_atten(block_mask=...) on the quantized paths, under the same one
+  ulp and all-ones rules, bool and 3d masks normalized, and the backward entry refusing a block
+  mask; these rows skip until the entry accepts block_mask
+- the flex consumer and the kernel fed one BlockSelection both sit within tolerance of fp32 sdpa
+  on the same tiles
+
+Bitwise rows need the autotuner pinned to one config: the two arms have different autotune keys,
+and a different tile changes the accumulation order. A single run pins 64x32 unless the
+SDNQ_TRITON_ATTEN_*_LIST env says otherwise, and the last row asserts the pin left the autotuner
+exactly one config. --tiles runs the file once per tile pair in a subprocess, which is also what
+proves every candidate tile nests the 128x64 block: a tile that did not would fail the kernel's
+static_assert at compile time.
+
+Usage:
+    python test/test-attention-sdnq-sparse.py
+    python test/test-attention-sdnq-sparse.py --tiles
+"""
+
+import os
+import sys
+import inspect
+
+TILES = [(32, 16), (32, 32), (32, 64), (64, 16), (64, 32), (64, 64), (128, 16), (128, 32), (128, 64)]
+PIN = {
+    'SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST': '64',
+    'SDNQ_TRITON_ATTEN_BLOCK_SIZE_N_LIST': '32',
+    'SDNQ_TRITON_ATTEN_NUM_WARPS_LIST': '4',
+    'SDNQ_TRITON_ATTEN_NUM_STAGES_LIST': '1',
+}
+
+
+def run_tiles() -> bool:
+    import subprocess
+    ok = True
+    for block_m, block_n in TILES:
+        env = dict(os.environ, **PIN)
+        env['SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST'] = str(block_m)
+        env['SDNQ_TRITON_ATTEN_BLOCK_SIZE_N_LIST'] = str(block_n)
+        print(f'=== tile {block_m}x{block_n} ===', flush=True)
+        code = subprocess.call([sys.executable, os.path.abspath(__file__)], env=env)
+        print(f'=== tile {block_m}x{block_n}: {"PASS" if code == 0 else "FAIL"} ===', flush=True)
+        ok = ok and code == 0
+    return ok
+
+
+if __name__ == '__main__' and '--tiles' in sys.argv:
+    sys.exit(0 if run_tiles() else 1)
+
+for pin_name, pin_value in PIN.items():
+    os.environ.setdefault(pin_name, pin_value)
+
+
+def env_list(name: str) -> list[int]:
+    return [int(x) for x in os.environ[name].replace(' ', '').split(',')]
+
+
+pinned = {name: env_list(name) for name in PIN}
+if any(len(values) != 1 for values in pinned.values()):
+    print('this test needs the attention autotuner pinned to one config: set each SDNQ_TRITON_ATTEN_*_LIST to a single value', flush=True)
+    sys.exit(2)
+block_size_m = pinned['SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST'][0]
+block_size_n = pinned['SDNQ_TRITON_ATTEN_BLOCK_SIZE_N_LIST'][0]
+
+import torch  # pylint: disable=wrong-import-position
+
+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
+from sdnq.kernels import triton_atten as atten_module         # pylint: disable=wrong-import-position
+from sdnq.kernels import triton_atten_backward as backward_module # pylint: disable=wrong-import-position
+
+
+results: dict[str, dict] = {}
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+kernel_available = device.type == 'cuda'
+int8_ok = block_size_n >= 32 # the int8 paths prune tiles narrower than 32, so a narrower pin cannot exercise them
+BLOCK_M, BLOCK_N = 128, 64
+
+
+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(4321)
+
+
+def randn(*shape, dtype=torch.bfloat16):
+    return torch.randn(*shape, generator=generator, device=device, dtype=torch.float32).to(dtype).contiguous()
+
+
+def qkv(batch=2, heads=4, kv_heads=None, seq=1000, dim=64, dtype=torch.bfloat16):
+    kv_heads = kv_heads or heads
+    return randn(batch, heads, seq, dim, dtype=dtype), randn(batch, kv_heads, seq, dim, dtype=dtype), randn(batch, kv_heads, seq, dim, dtype=dtype)
+
+
+def blocks(seq_q: int, seq_kv: int):
+    return sparse.selector.block_count(seq_q, BLOCK_M), sparse.selector.block_count(seq_kv, BLOCK_N)
+
+
+def random_keep(batch: int, heads: int, seq_q: int, seq_kv: int, density: float = 0.4):
+    """A random block mask with the diagonal kept, so no row is empty unless a test empties it."""
+    nq, nk = blocks(seq_q, seq_kv)
+    keep = torch.rand(batch, heads, nq, nk, generator=generator, device=device) < density
+    keep |= sparse.selector.diagonal_blocks(nq, nk, BLOCK_M, BLOCK_N, device)
+    return keep.to(torch.int8)
+
+
+def expand(keep: torch.Tensor, seq_q: int, seq_kv: int):
+    """The token-granular int8 mask a block mask stands for."""
+    return keep.repeat_interleave(BLOCK_M, dim=-2).repeat_interleave(BLOCK_N, dim=-1)[..., :seq_q, :seq_kv].contiguous()
+
+
+def ordered(keep: torch.Tensor):
+    """The count and ascending index list per query block that the kernel walks, built the way the entry builds them: left-padded to 4d, size-1 batch and head dims kept, padded to the descriptor chunks."""
+    while keep.ndim < 4:
+        keep = keep.unsqueeze(0)
+    return atten_module.get_block_mask_input(keep)
+
+
+def raw(q, k, v, attn_mask=None, block_mask=None, block_count=None, block_index=None):
+    """The triton op without the input prep: nothing is quantized, so the mask path is the only difference between arms."""
+    if block_mask is not None:
+        block_count, block_index = ordered(block_mask)
+    sparse = block_count is not None
+    return atten_module.sdnq_triton_atten_fwd(
+        q, k, v, None, None, None,
+        attn_mask=attn_mask, is_causal=False, sm_scale=q.shape[-1] ** -0.5, use_fp16_accum=False, out_dtype=q.dtype,
+        block_count=block_count, block_index=block_index, block_mask_m=BLOCK_M if sparse else 0, block_mask_n=BLOCK_N if sparse else 0,
+    )
+
+
+def entry_accepts_block_mask() -> bool:
+    return 'block_mask' in inspect.signature(inspect.unwrap(atten_module.sdnq_triton_atten)).parameters
+
+
+def entry(q, k, v, **kwargs):
+    return atten_module.sdnq_triton_atten(q, k, v, **kwargs)
+
+
+def configs_for(sizes_m, sizes_n):
+    import triton
+    return [triton.Config({'BLOCK_SIZE_M': m, 'BLOCK_SIZE_N': n}, num_warps=4, num_stages=1) for m in sizes_m for n in sizes_n]
+
+
+def nests(conf) -> bool:
+    return BLOCK_M % conf.kwargs['BLOCK_SIZE_M'] == 0 and BLOCK_N % conf.kwargs['BLOCK_SIZE_N'] == 0
+
+
+def assert_one_ulp(got: torch.Tensor, want: torch.Tensor, label: str = ''):
+    """Within one unit in the last place at each element, plus an absolute floor for the drift between the two rounding sequences (a straddled rounding of one p before the pv dot, amplified on rows with few keys), two orders below what a mis-skipped tile moves."""
+    if torch.equal(got, want):
+        return
+    eps = torch.finfo(got.dtype).eps
+    atol = want.float().abs().max().item() * 2**-10
+    if not torch.allclose(got.float(), want.float(), rtol=eps, atol=atol):
+        gap = (got.float() - want.float()).abs()
+        raise AssertionError(f'{label} max diff {gap.max().item():.3e} over {int((got != want).sum().item())} elements, bound one ulp plus {atol:.1e}')
+
+
+# ============================================================
+# Kernel contract
+# ============================================================
+
+def test_block_mask_matches_the_expanded_token_mask():
+    if not kernel_available:
+        return None
+    for seq in (1000, 970): # 970 leaves the last block with sub tiles past the end of the sequence under the narrower tiles
+        q, k, v = qkv(seq=seq)
+        keep = random_keep(2, 4, seq, seq)
+        got = raw(q, k, v, block_mask=keep)
+        assert_one_ulp(got, raw(q, k, v, attn_mask=expand(keep, seq, seq)), f'seq {seq}')
+        assert torch.isfinite(got).all()
+        assert not torch.equal(got, raw(q, k, v)), 'a 40 percent block mask left the output identical to dense: the mask is not applied'
+    return True
+
+
+def test_gqa_block_mask_indexes_query_heads():
+    if not kernel_available:
+        return None
+    q, k, v = qkv(heads=8, kv_heads=2)
+    for heads in (8, 1):
+        keep = random_keep(2, heads, 1000, 1000)
+        assert_one_ulp(raw(q, k, v, block_mask=keep), raw(q, k, v, attn_mask=expand(keep, 1000, 1000)), f'mask heads {heads}')
+    return True
+
+
+def test_block_mask_broadcasts_over_batch_and_heads():
+    if not kernel_available:
+        return None
+    q, k, v = qkv(batch=3, heads=4)
+    for batch, heads in ((1, 1), (1, 4), (3, 1)):
+        keep = random_keep(batch, heads, 1000, 1000)
+        assert_one_ulp(raw(q, k, v, block_mask=keep), raw(q, k, v, attn_mask=expand(keep, 1000, 1000)), f'mask shape {tuple(keep.shape)}')
+    return True
+
+
+def test_block_lists_must_be_contiguous_and_padded():
+    """The kernel reads the lists through descriptors from shapes alone, so the launcher takes only what the entry builds."""
+    if not kernel_available:
+        return None
+    q, k, v = qkv(seq=1100) # 9 query blocks and 18 kv blocks, so both lists get padded (to 12 and 32)
+    keep = random_keep(2, 4, 1100, 1100)
+    count, index = ordered(keep)
+    nq, nk = blocks(1100, 1100)
+    assert count.is_contiguous() and index.is_contiguous()
+    assert count.shape[-1] % atten_module.block_count_chunk == 0 and index.shape[-1] % atten_module.block_index_chunk == 0
+    assert count.shape[-1] > nq and index.shape[-2] == nq and index.shape[-1] > nk, 'this geometry should pad both lists'
+    assert torch.equal(count[..., :nq], keep.sum(-1, dtype=torch.int32)), 'padding changed the counts'
+    reference = raw(q, k, v, block_count=count, block_index=index)
+    count_view = count.transpose(-1, -2).contiguous().transpose(-1, -2)
+    index_view = index.transpose(-1, -2).contiguous().transpose(-1, -2)
+    unpadded_count = count[..., :nq].contiguous()
+    unpadded_index = index[..., :nk].contiguous()
+    rejected = (
+        ('a strided count', {'block_count': count_view, 'block_index': index}),
+        ('a strided index', {'block_count': count, 'block_index': index_view}),
+        ('an unpadded count', {'block_count': unpadded_count, 'block_index': index}),
+        ('an unpadded index', {'block_count': count, 'block_index': unpadded_index}),
+    )
+    for label, lists in rejected:
+        try:
+            raw(q, k, v, **lists)
+        except ValueError:
+            continue
+        raise AssertionError(f'{label} was accepted')
+    assert torch.equal(raw(q, k, v, block_mask=keep), reference)
+    shared_keep = keep[:1, :1]
+    shared_count, shared_index = ordered(shared_keep)
+    assert tuple(shared_count.shape[:2]) == (1, 1) and tuple(shared_index.shape[:2]) == (1, 1), 'size-1 batch and head dims stay as given, the kernel broadcasts by shape'
+    assert torch.equal(raw(q, k, v, block_count=shared_count, block_index=shared_index), raw(q, k, v, block_mask=shared_keep.expand(2, 4, -1, -1).contiguous()))
+    return True
+
+
+def test_empty_block_row_gives_zeros_without_nan():
+    if not kernel_available:
+        return None
+    q, k, v = qkv()
+    keep = random_keep(2, 4, 1000, 1000)
+    keep[..., 2, :] = 0
+    got = raw(q, k, v, block_mask=keep)
+    assert_one_ulp(got, raw(q, k, v, attn_mask=expand(keep, 1000, 1000)))
+    assert torch.isfinite(got).all()
+    assert got[..., 256:384, :].abs().max().item() == 0.0, 'an empty block row should give zeros'
+    assert got[..., :256, :].abs().max().item() > 0.0
+    return True
+
+
+def test_block_mask_composes_with_a_token_mask():
+    if not kernel_available:
+        return None
+    q, k, v = qkv()
+    keep = random_keep(2, 4, 1000, 1000)
+    padding = torch.ones(2, 1, 1000, 1000, dtype=torch.int8, device=device)
+    padding[..., -37:] = 0 # the last keys are padding, as a packed sequence with a tail would have
+    got = raw(q, k, v, attn_mask=padding, block_mask=keep)
+    want = raw(q, k, v, attn_mask=(padding.bool() & expand(keep, 1000, 1000).bool()).to(torch.int8))
+    assert torch.equal(got, want), 'both arms carry a token mask, so this one is bitwise' # pylint: disable=line-too-long
+    assert not torch.equal(got, raw(q, k, v, block_mask=keep)), 'the token mask was ignored beside the block mask'
+    return True
+
+
+def same_arithmetic(got: torch.Tensor, want: torch.Tensor, label: str):
+    """Every block kept against no mask: the same tiles in the same order, so any gap is compiler scheduling of a loop with a runtime trip count; report whether it was bitwise and bound it either way."""
+    bitwise = torch.equal(got, want)
+    if not bitwise:
+        gap = (got.float() - want.float()).abs()
+        log.info(f'    {label}: all-ones differs from dense by {gap.max().item():.3e} over {int((got != want).sum().item())} elements (compiler scheduling, not tiles)')
+    assert_one_ulp(got, want, label)
+    return bitwise
+
+
+def test_all_ones_block_mask_matches_no_mask():
+    if not kernel_available:
+        return None
+    q, k, v = qkv()
+    nq, nk = blocks(1000, 1000)
+    ones = torch.ones(1, 1, nq, nk, dtype=torch.int8, device=device)
+    same_arithmetic(raw(q, k, v, block_mask=ones), raw(q, k, v), 'raw bf16')
+    return True
+
+
+def test_sub_tile_past_the_sequence_is_masked():
+    """992 divides every kernel tile but not the 64-wide mask block, so the last block's trailing sub tile sits entirely past the keys.
+
+    Without a tail rule that knows the block geometry it loads zeros, scores them as zero and inflates the softmax denominator.
+    """
+    if not kernel_available:
+        return None
+    seq = 992
+    assert seq % block_size_n == 0 or seq % BLOCK_N != 0, 'this row needs a sequence the tile divides and the mask block does not'
+    q, k, v = qkv(seq=seq)
+    nq, nk = blocks(seq, seq)
+    assert nk * BLOCK_N > seq, 'the last mask block should run past the sequence'
+    ones = torch.ones(1, 1, nq, nk, dtype=torch.int8, device=device)
+    assert_one_ulp(raw(q, k, v, block_mask=ones), raw(q, k, v), f'seq {seq}, every block kept')
+    keep = random_keep(2, 4, seq, seq)
+    assert_one_ulp(raw(q, k, v, block_mask=keep), raw(q, k, v, attn_mask=expand(keep, seq, seq)), f'seq {seq}, 40 percent kept')
+    return True
+
+
+def test_launcher_validates_the_block_lists():
+    if not kernel_available:
+        return None
+    q, k, v = qkv()
+    ones = torch.ones(1, 1, *blocks(1000, 1000), dtype=torch.int8, device=device)
+    count, index = ordered(ones)
+    wide_count, wide_index = ordered(ones.expand(5, 4, -1, -1).contiguous())
+    bad = (
+        ('a count without an index', {'block_count': count}),
+        ('an index without a count', {'block_index': index}),
+        ('an int8 index', {'block_count': count, 'block_index': index.to(torch.int8)}),
+        ('a 3d index', {'block_count': count, 'block_index': index[0]}),
+        ('a short index', {'block_count': count, 'block_index': index[..., :-1].contiguous()}),
+        ('a wrong batch', {'block_count': wide_count, 'block_index': wide_index}),
+    )
+    for label, lists in bad:
+        try:
+            raw(q, k, v, **lists)
+        except ValueError:
+            continue
+        raise AssertionError(f'{label} was accepted')
+    try:
+        atten_module.sdnq_triton_atten_fwd(q, k, v, None, None, None, sm_scale=0.125, out_dtype=q.dtype, block_count=count, block_index=index, block_mask_m=0, block_mask_n=0)
+    except ValueError:
+        return True
+    raise AssertionError('block lists without block sizes were accepted')
+
+
+# ============================================================
+# Autotune nesting
+# ============================================================
+
+def test_prune_configs_applies_the_filter():
+    """The nesting filter keeps exactly the tiles a mask block divides by, leaves a dense launch alone, and raises rather than letting a tile list that cannot nest run."""
+    if not kernel_available:
+        return None
+    q, k, v = qkv()
+    args = {
+        'q_ptr': q, 'k_ptr': k, 'v_ptr': v, 'out_ptr': q,
+        'QN': 1000, 'KN': 1000, 'QHD': 64, 'KHD': 64, 'VHD': 64,
+        'is_causal': 0, 'do_block_mask': 1, 'BLOCK_MASK_M': BLOCK_M, 'BLOCK_MASK_N': BLOCK_N,
+    }
+    with_mask = atten_module.prune_configs(configs_for((32, 64, 128, 256), (16, 32, 64, 128)), args)
+    without = atten_module.prune_configs(configs_for((32, 64, 128, 256), (16, 32, 64, 128)), dict(args, do_block_mask=0))
+    assert with_mask, 'nothing survived'
+    assert [c.kwargs for c in with_mask] == [c.kwargs for c in without if nests(c)], ([c.kwargs for c in with_mask], [c.kwargs for c in without])
+    try: # from_small is the terminal arm, so the fallback to the small configs cannot refill the list before the filter reaches it
+        atten_module.prune_configs(configs_for((256,), (128,)), args, from_small=True)
+    except ValueError as e:
+        assert 'SDNQ_TRITON_ATTEN_BLOCK_SIZE_M_LIST' in str(e), e
+        return True
+    raise AssertionError('a tile list with nothing nesting did not raise')
+
+
+# ============================================================
+# Public entry
+# ============================================================
+
+def entry_rows_ready() -> bool:
+    return kernel_available and entry_accepts_block_mask()
+
+
+def entry_rows(q, k, v, keep, **kwargs) -> bool:
+    """The block arm within one ulp of the token mask, applied at all, and bitwise the dense path with every block kept."""
+    seq_q, seq_kv = q.shape[-2], k.shape[-2]
+    got = entry(q, k, v, block_mask=keep, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, **kwargs)
+    assert_one_ulp(got, entry(q, k, v, attn_mask=expand(keep, seq_q, seq_kv).bool(), **kwargs))
+    dense = entry(q, k, v, **kwargs)
+    assert not torch.equal(got, dense), 'the block mask is not applied'
+    nq, nk = blocks(seq_q, seq_kv)
+    ones = torch.ones(1, 1, nq, nk, dtype=torch.int8, device=device)
+    same_arithmetic(entry(q, k, v, block_mask=ones, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, **kwargs), dense, f'entry {kwargs.get("matmul_dtype", "int8")}')
+    return True
+
+
+def test_entry_int8_block_mask():
+    if not entry_rows_ready() or not int8_ok:
+        return None
+    q, k, v = qkv()
+    return entry_rows(q, k, v, random_keep(2, 4, 1000, 1000), matmul_dtype='int8')
+
+
+def test_entry_int8_pv_block_mask():
+    if not entry_rows_ready() or not int8_ok:
+        return None
+    q, k, v = qkv()
+    return entry_rows(q, k, v, random_keep(2, 4, 1000, 1000), matmul_dtype='int8', pv_matmul_dtype='int8')
+
+
+def test_entry_fp16_accum_block_mask():
+    if not entry_rows_ready():
+        return None
+    q, k, v = qkv(dtype=torch.float16)
+    return entry_rows(q, k, v, random_keep(2, 4, 1000, 1000), matmul_dtype='float16', pv_matmul_dtype='float16', use_fp16_accum=True)
+
+
+def test_entry_normalizes_bool_and_3d_block_masks():
+    if not entry_rows_ready():
+        return None
+    q, k, v = qkv(batch=1)
+    keep = random_keep(1, 4, 1000, 1000)
+    want = entry(q, k, v, block_mask=keep, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, do_quantize=False)
+    assert torch.equal(entry(q, k, v, block_mask=keep.bool(), block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, do_quantize=False), want), 'bool'
+    assert torch.equal(entry(q, k, v, block_mask=keep[0], block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, do_quantize=False), want), '3d'
+    return True
+
+
+# ============================================================
+# Backward
+# ============================================================
+
+def backward_available() -> bool:
+    return kernel_available and 'block_mask' in inspect.signature(backward_module.sdnq_triton_atten_with_backward).parameters
+
+
+def upstream_grad(shape, seed=17):
+    """The same upstream gradient for every arm: two arms fed different ones are not comparable at all."""
+    return torch.randn(shape, generator=torch.Generator(device=device).manual_seed(seed), device=device, dtype=torch.float32)
+
+
+def grads(q, k, v, upstream=None, **kwargs):
+    """dq, dk and dv for one attention call through the kernel's autograd function."""
+    q, k, v = q.detach().clone().requires_grad_(True), k.detach().clone().requires_grad_(True), v.detach().clone().requires_grad_(True)
+    out = backward_module.sdnq_triton_atten_with_backward(q, k, v, **kwargs)
+    if upstream is None:
+        upstream = upstream_grad(out.shape)
+    out.backward(upstream.to(out.dtype))
+    return q.grad, k.grad, v.grad
+
+
+def test_backward_matches_the_expanded_token_mask():
+    """The load-bearing row: the same selection through the block lists and through a token mask has to give the same gradients."""
+    if not backward_available():
+        return None
+    seq = 1000
+    q, k, v = qkv(seq=seq)
+    keep = random_keep(2, 4, seq, seq)
+    block = grads(q, k, v, block_mask=keep, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, do_quantize=False)
+    token = grads(q, k, v, attn_mask=expand(keep, seq, seq).bool(), do_quantize=False)
+    dense = grads(q, k, v, do_quantize=False)
+    for name, got, want, other in zip(('dq', 'dk', 'dv'), block, token, dense):
+        assert got is not None and torch.isfinite(got).all(), f'{name} is not finite'
+        assert_one_ulp(got, want, name)
+        assert not torch.equal(got, other), f'{name} matches the dense gradient, so the mask never reached the backward'
+    return True
+
+
+def test_backward_gqa_and_int8():
+    """GQA sends every query head's selection into one kv head's gradient, and the quantized path is the one that ships."""
+    if not backward_available():
+        return None
+    seq = 1000
+    q, k, v = qkv(heads=4, kv_heads=2, seq=seq)
+    keep = random_keep(2, 4, seq, seq)
+    for label, kwargs in (('bf16', {'do_quantize': False}), ('int8', {'matmul_dtype': 'int8'})):
+        if label == 'int8' and not int8_ok:
+            continue
+        block = grads(q, k, v, block_mask=keep, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, **kwargs)
+        token = grads(q, k, v, attn_mask=expand(keep, seq, seq).bool(), **kwargs)
+        for name, got, want in zip(('dq', 'dk', 'dv'), block, token):
+            assert_one_ulp(got, want, f'{label} {name}')
+    return True
+
+
+def test_backward_zeroes_dropped_blocks():
+    """A query block that keeps nothing gets no gradient, and a kv block no query block kept gets none either."""
+    if not backward_available():
+        return None
+    seq = 1000
+    q, k, v = qkv(seq=seq)
+    keep = random_keep(2, 4, seq, seq)
+    keep[..., 2, :] = 0 # query block 2 attends to nothing
+    keep[..., :, 5] = 0 # kv block 5 is attended by nothing
+    dq, dk, dv = grads(q, k, v, block_mask=keep, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, do_quantize=False)
+    assert torch.isfinite(dq).all() and torch.isfinite(dk).all() and torch.isfinite(dv).all()
+    assert dq[..., 2 * BLOCK_M:3 * BLOCK_M, :].abs().max().item() == 0.0, 'an empty query block still got a gradient'
+    assert dq[..., :BLOCK_M, :].abs().max().item() > 0.0
+    assert dk[..., 5 * BLOCK_N:6 * BLOCK_N, :].abs().max().item() == 0.0, 'a dropped kv block still got a key gradient'
+    assert dv[..., 5 * BLOCK_N:6 * BLOCK_N, :].abs().max().item() == 0.0, 'a dropped kv block still got a value gradient'
+    assert dk[..., :BLOCK_N, :].abs().max().item() > 0.0
+    return True
+
+
+def test_backward_matches_fp32_autograd():
+    """Against truth rather than against the other masked path: torch's own gradients through fp32 sdpa on the same selection.
+
+    Run in fp32, where the kernel sits within 0.1 percent of autograd; in bf16 the dense backward is already 1.8 percent
+    out, so a bf16 arm could only carry a bound too loose to catch a mis-walked block. Head dim 32 keeps fp32 inside the
+    shared memory the widest pinned tile has, which fp32 at head dim 64 exceeds on the dense backward too.
+    """
+    if not backward_available():
+        return None
+    seq, dim = 1000, 32
+    q, k, v = (t.float().contiguous() for t in qkv(batch=1, heads=4, seq=seq, dim=dim))
+    keep = random_keep(1, 4, seq, seq)
+    upstream = upstream_grad((1, 4, seq, dim))
+    reference = [t.detach().clone().requires_grad_(True) for t in (q, k, v)]
+    stock_sdpa(*reference, attn_mask=expand(keep, seq, seq).bool()).backward(upstream)
+    got = grads(q, k, v, upstream=upstream, block_mask=keep, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, do_quantize=False)
+    for name, mine, want in zip(('dq', 'dk', 'dv'), got, reference):
+        scale = want.grad.abs().max().item()
+        gap = (mine.float() - want.grad).abs().max().item()
+        log.info(f'    {name}: {100 * gap / scale:.3f} percent of the gradient scale against fp32 autograd')
+        assert gap < 5e-3 * scale, f'{name} is {gap:.3e} off fp32 autograd, more than 0.5 percent of {scale:.3e}'
+    return True
+
+
+def test_backward_ragged_sub_tile():
+    """The transposed lists have the same ragged tail as the forward's: 992 divides the tile and not the mask block."""
+    if not backward_available():
+        return None
+    seq = 992
+    q, k, v = qkv(seq=seq)
+    keep = random_keep(2, 4, seq, seq)
+    block = grads(q, k, v, block_mask=keep, block_mask_m=BLOCK_M, block_mask_n=BLOCK_N, do_quantize=False)
+    token = grads(q, k, v, attn_mask=expand(keep, seq, seq).bool(), do_quantize=False)
+    for name, got, want in zip(('dq', 'dk', 'dv'), block, token):
+        assert torch.isfinite(got).all(), f'{name} is not finite'
+        assert_one_ulp(got, want, f'seq {seq} {name}')
+    return True
+
+
+# ============================================================
+# Consumers
+# ============================================================
+
+def test_kernel_and_flex_agree_on_one_selection():
+    if not entry_rows_ready():
+        return None
+    seq = 2048
+    q, k, v = qkv(batch=1, heads=4, seq=seq)
+    selection = sparse.select_blocks(q, k, sparse.SparseSpec(budget=0.25))
+    assert selection is not None
+    reference = stock_sdpa(q.float(), k.float(), v.float(), attn_mask=expand(selection.keep, seq, seq).bool())
+    flex_gap = (sparse_flex.attend(q, k, v, selection).float() - reference).abs().max().item()
+    kernel_gap = (entry(q, k, v, block_mask=selection.keep, block_mask_m=selection.block_q, block_mask_n=selection.block_kv, do_quantize=False).float() - reference).abs().max().item()
+    log.info(f'    gaps against fp32 sdpa on the same tiles: flex {flex_gap:.5f} kernel {kernel_gap:.5f}')
+    assert flex_gap < 1e-2 and kernel_gap < 1e-2, (flex_gap, kernel_gap)
+    return True
+
+
+def test_the_pin_left_the_autotuner_one_config():
+    # with one config the autotuner benchmarks nothing and prunes nothing, so both arms of every
+    # bitwise row compiled the same tile; the in-kernel static_assert is the nesting guard there
+    configs = [conf.kwargs for conf in atten_module.autotune_configs]
+    assert configs == [{'BLOCK_SIZE_M': block_size_m, 'BLOCK_SIZE_N': block_size_n}], f'the pin did not take: {configs}'
+    tuner_configs = [conf.kwargs for conf in getattr(atten_module.sdnq_attn_kernel, 'configs', [])]
+    assert tuner_configs == configs, f'the autotuner holds {tuner_configs}'
+    return True
+
+
+def run_all():
+    log.warning(f'Running sdnq block mask tests on {device}, tile pinned to {block_size_m}x{block_size_n}')
+
+    log.warning('=== kernel contract ===')
+    cat = category('contract')
+    for fn in [
+        test_block_mask_matches_the_expanded_token_mask,
+        test_gqa_block_mask_indexes_query_heads,
+        test_block_mask_broadcasts_over_batch_and_heads,
+        test_block_lists_must_be_contiguous_and_padded,
+        test_empty_block_row_gives_zeros_without_nan,
+        test_block_mask_composes_with_a_token_mask,
+        test_all_ones_block_mask_matches_no_mask,
+        test_sub_tile_past_the_sequence_is_masked,
+        test_launcher_validates_the_block_lists,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== autotune nesting ===')
+    cat = category('autotune')
+    for fn in [
+        test_prune_configs_applies_the_filter,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== public entry ===')
+    cat = category('entry')
+    for fn in [
+        test_entry_int8_block_mask,
+        test_entry_int8_pv_block_mask,
+        test_entry_fp16_accum_block_mask,
+        test_entry_normalizes_bool_and_3d_block_masks,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== backward ===')
+    cat = category('backward')
+    for fn in [
+        test_backward_matches_the_expanded_token_mask,
+        test_backward_gqa_and_int8,
+        test_backward_zeroes_dropped_blocks,
+        test_backward_matches_fp32_autograd,
+        test_backward_ragged_sub_tile,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== consumers ===')
+    cat = category('consumers')
+    for fn in [
+        test_kernel_and_flex_agree_on_one_selection,
+        test_the_pin_left_the_autotuner_one_config,
+    ]:
+        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)

test/test-attention-sparse.py

@@ -0,0 +1,635 @@
+#!/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 router stage: the gates it applies (component role, mask, causal, cross attention, minimum
+  sequence), the per step budget schedule, and layout resolution with and without a publisher
+
+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_segments_from_live_splits_interior_padding():
+    from modules.attention.sparse import layout as layout_mod
+    live = torch.zeros(512, dtype=torch.bool, device=device)
+    live[:40] = True
+    live[-8:] = True
+    segments = layout_mod.segments_from_live(live, 'text')
+    assert segments == [('text', 40), ('pad', 464), ('text', 8)], segments
+    token_layout = layout_mod.layout_from_segments(segments + [('image', 1024)])
+    assert token_layout.length == 1536, token_layout.length
+    _, drops = layout_mod.block_pins(token_layout, 1536, 1536, 128, 64, device)
+    # only a key block that is padding all the way through is dropped, so the two straddling blocks survive
+    assert drops[..., 1:7].all(), 'whole padded key blocks are dropped'
+    assert not drops[..., 0].any() and not drops[..., 7].any(), 'a straddling block keeps its live tokens'
+    return True
+
+
+def test_layout_from_stream_ids_reads_the_joint_convention():
+    from modules.attention.sparse import layout as layout_mod
+    flat = (torch.zeros(512, 3, device=device), torch.zeros(4096, 3, device=device)) # flux1 passes 2d ids
+    batched = (torch.zeros(1, 512, 4, device=device), torch.zeros(1, 4096, 4, device=device)) # flux2 passes 3d
+    for text, image in (flat, batched):
+        token_layout = layout_mod.layout_from_stream_ids({'txt_ids': text, 'img_ids': image}, 'FluxTransformer2DModel')
+        assert token_layout is not None and token_layout.length == 4608, token_layout
+        assert [(s.kind, s.start, s.end) for s in token_layout.spans] == [('text', 0, 512), ('image', 512, 4608)], token_layout.spans
+    # an architecture whose packing order is not verified publishes nothing rather than pinning the wrong half dense
+    assert layout_mod.layout_from_stream_ids({'txt_ids': flat[0], 'img_ids': flat[1]}, 'HiDreamImageTransformer2DModel') is None
+    assert layout_mod.layout_from_stream_ids({'txt_ids': flat[0], 'img_ids': flat[1]}, None) is None
+    assert layout_mod.layout_from_stream_ids({}, 'FluxTransformer2DModel') is None
+    indices = {'video_indices': torch.arange(0, 64, device=device), 'txt_ids': flat[0], 'img_ids': flat[1]}
+    assert layout_mod.layout_from_kwargs(indices, 'FluxTransformer2DModel').source == 'indices', 'the index form wins when both are present'
+    return True
+
+
+def test_layout_hook_publishes_from_the_denoiser_kwargs():
+    from modules.attention import context as ctx
+
+    class FluxTransformer2DModel(torch.nn.Module): # the reader keys on the class name, so the fake carries a real one
+        def forward(self, hidden_states=None, txt_ids=None, img_ids=None): # pylint: disable=unused-argument
+            return hidden_states
+
+    class Pipe:
+        def __init__(self, transformer, second):
+            self.transformer = transformer
+            self.unconditional_transformer = second # ideogram runs a second denoiser, wan a14b a transformer_2
+
+    denoiser = FluxTransformer2DModel()
+    second = FluxTransformer2DModel()
+    pipe = Pipe(denoiser, second)
+    previous = getattr(shared.opts, 'sparse_attention_enabled', False)
+    try:
+        shared.opts.data['sparse_attention_enabled'] = False
+        ctx.install_layout_hook(pipe)
+        assert getattr(denoiser, 'sdnext_layout_hook', None) is None, 'nothing is hooked while the feature is off'
+        shared.opts.data['sparse_attention_enabled'] = True
+        ctx.install_layout_hook(pipe)
+        ctx.install_layout_hook(pipe)
+        assert getattr(denoiser, 'sdnext_layout_hook', None) is not None, 'the denoiser is hooked once'
+        assert getattr(second, 'sdnext_layout_hook', None) is not None, 'every denoiser slot is hooked, not just the first'
+        ctx.set_layout(None)
+        denoiser(hidden_states=torch.zeros(1, 4096, 4, device=device), txt_ids=torch.zeros(512, 3, device=device), img_ids=torch.zeros(4096, 3, device=device))
+        published = ctx.current.layout
+        assert published is not None and published.length == 4608 and published.source == 'stream-ids', published
+    finally:
+        shared.opts.data['sparse_attention_enabled'] = previous
+        ctx.set_layout(None)
+    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
+
+
+# ============================================================
+# Router stage
+# ============================================================
+
+def stage_options(**kwargs):
+    from modules.attention.sparse import stage as stage_mod
+    base = dict(enabled=True, budget=0.30, min_tokens=1024)
+    base.update(kwargs)
+    return stage_mod.StageOptions(**base)
+
+
+def with_context(fn):
+    from modules.attention import context as ctx
+    ctx.begin(None, steps=10)
+    try:
+        return fn()
+    finally:
+        ctx.end()
+
+
+def test_stage_is_none_when_disabled_or_at_full_budget():
+    from modules.attention.sparse import stage as stage_mod
+    assert stage_mod.make_stage(stage_options(enabled=False)) is None
+    assert stage_mod.make_stage(stage_options(budget=1.0)) is None
+    assert stage_mod.make_stage(stage_options()) is not None
+    return True
+
+
+def test_stage_gates():
+    from modules.attention.sparse import stage as stage_mod
+    stage = stage_mod.make_stage(stage_options())
+    q, k, v = qkv(heads=2, seq=2048)
+    short_q, short_k, short_v = qkv(heads=2, seq=512)
+    cross_k, cross_v = randn(1, 2, 77, 64), randn(1, 2, 77, 64)
+
+    def checks():
+        from modules.attention import context as ctx
+        assert stage(q, k, v, None, False) is not None, 'an eligible call must be selected'
+        mask = torch.zeros(1, 1, 2048, 2048, dtype=torch.bool, device=device)
+        assert stage(q, k, v, mask, False) is None, 'a masked call needs a backend that composes the two'
+        assert stage(q, k, v, mask, False, frozenset({'masked_block'})) is not None, 'a composing backend takes the masked call'
+        assert stage(q, k, v, None, True) is None, 'a causal call is not eligible'
+        assert stage(q, cross_k, cross_v, None, False) is None, 'cross attention is not eligible'
+        assert stage(short_q, short_k, short_v, None, False) is None, 'below the gate attention stays dense'
+        assert stage.last_skip == 'below the minimum sequence', stage.last_skip
+        ctx.set_role('vae')
+        assert stage(q, k, v, None, False) is None, 'only the denoiser is sparsified'
+        ctx.set_role('transformer')
+        return True
+    return with_context(checks)
+
+
+def test_minimum_sequence_of_zero_sparsifies_everything():
+    from modules.attention.sparse import stage as stage_mod
+    floored = stage_mod.make_stage(stage_options())
+    unfloored = stage_mod.make_stage(stage_options(min_tokens=0))
+    short_q, short_k, short_v = qkv(heads=2, seq=512)
+
+    def checks():
+        assert floored(short_q, short_k, short_v, None, False) is None, 'the default floor keeps a short sequence dense'
+        assert unfloored(short_q, short_k, short_v, None, False) is not None, 'a floor of zero sparsifies every eligible call'
+        return True
+    return with_context(checks)
+
+
+def test_exclusion_list_parsing():
+    from modules.attention.sparse import stage as stage_mod
+    assert stage_mod.parse_exclusions(' Anima , CosmosTransformer3DModel ,, ') == ('anima', 'cosmostransformer3dmodel')
+    assert stage_mod.parse_exclusions('') == ()
+    assert stage_mod.parse_exclusions(None) == ()
+    return True
+
+
+def test_an_exclusion_entry_matches_any_of_the_three_names():
+    """One entry, matched against whichever of the architecture, pipeline class or denoiser class the user knew."""
+    from modules.attention.sparse import stage as stage_mod
+    key = ('AnimaTextToImagePipeline', 'CosmosTransformer3DModel')
+    for entry in ('cosmostransformer3dmodel', 'animatexttoimagepipeline', 'anima'):
+        assert stage_mod.match_exclusion(key, 'anima', (entry,)) == entry, entry
+    assert stage_mod.match_exclusion(key, 'anima', ('krea2',)) == '', 'an unlisted model matches nothing'
+    assert stage_mod.match_exclusion(None, 'anima', ('anima',)) == 'anima', 'the architecture stands in when no model key is published'
+    assert stage_mod.match_exclusion(('Pipe', None), None, ('pipe',)) == 'pipe', 'an absent denoiser class is skipped, not matched'
+    return True
+
+
+def test_an_excluded_model_stays_dense():
+    from modules.attention import context as ctx
+    from modules.attention.sparse import stage as stage_mod
+    listed = stage_mod.make_stage(stage_options(exclude=('cosmostransformer3dmodel',)))
+    unlisted = stage_mod.make_stage(stage_options(exclude=('somethingelse',)))
+    q, k, v = qkv(heads=2, seq=2048)
+
+    def checks():
+        ctx.current.model_key = ('AnimaTextToImagePipeline', 'CosmosTransformer3DModel')
+        assert listed(q, k, v, None, False) is None, 'a listed denoiser class stays dense'
+        assert listed.last_skip == 'excluded', listed.last_skip
+        assert unlisted(q, k, v, None, False) is not None, 'a list that matches nothing changes nothing'
+        ctx.current.model_key = ('Krea2Pipeline', 'Krea2Transformer2DModel')
+        assert listed(q, k, v, None, False) is not None, 'the exclusion applies to the listed model, not to every model'
+        return True
+    return with_context(checks)
+
+
+def test_published_segments_reach_the_stage():
+    from modules.attention.sparse import layout as layout_mod
+    from modules.attention.sparse import stage as stage_mod
+    from modules.attention import context as ctx
+    stage = stage_mod.make_stage(stage_options(budget=0.30))
+    q, k, v = qkv(heads=2, seq=2048)
+
+    def checks():
+        layout_mod.publish_segments((('text', 256), ('image', 1536), ('pad', 256)), source='test')
+        published = ctx.current.layout
+        assert published.length == 2048 and published.source == 'test', published
+        assert published.kinds() == ('text', 'image', 'pad'), published.kinds()
+        selection = stage(q, k, v, None, False)
+        keep, block_kv = selection.keep, selection.block_kv
+        assert keep[..., :256 // block_kv].all(), 'conditioning key tiles stay dense'
+        assert not keep[..., 1792 // block_kv:].any(), 'padding key tiles are dropped'
+        return True
+    return with_context(checks)
+
+
+def test_stage_follows_the_step_schedule():
+    from modules.attention.sparse import stage as stage_mod
+    from modules.attention import context as ctx
+    stage = stage_mod.make_stage(stage_options(budget=0.30, schedule_steps=2, schedule_bump=0.40))
+    q, k, v = qkv(heads=2, seq=2048)
+
+    def checks():
+        densities = []
+        for step in range(10):
+            ctx.set_step(step)
+            selection = stage(q, k, v, None, False)
+            densities.append(selection.budget)
+        assert densities[0] == densities[1] > densities[5], densities
+        assert densities[-1] == densities[-2] > densities[5], densities
+        assert len(set(densities)) == 2, set(densities)
+        return True
+    return with_context(checks)
+
+
+def test_stage_uses_a_published_layout_and_falls_back_without_one():
+    from modules.attention.sparse import stage as stage_mod
+    from modules.attention import context as ctx
+    stage = stage_mod.make_stage(stage_options(budget=0.20))
+    q, k, v = qkv(heads=2, seq=2048)
+
+    def checks():
+        loose = stage(q, k, v, None, False)
+        ctx.set_layout(sparse.layout_from_segments([('text', 256), ('video', 1792)]))
+        pinned = stage(q, k, v, None, False)
+        assert pinned.density() > loose.density(), f'pinning conditioning must keep more tiles: {pinned.density()} vs {loose.density()}'
+        assert bool(pinned.keep[..., 0:4].all()), 'the pinned text columns must survive'
+        return True
+    return with_context(checks)
+
+
+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,
+        test_segments_from_live_splits_interior_padding,
+        test_layout_from_stream_ids_reads_the_joint_convention,
+        test_layout_hook_publishes_from_the_denoiser_kwargs,
+    ]:
+        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('=== stage ===')
+    cat = category('stage')
+    for fn in [
+        test_stage_is_none_when_disabled_or_at_full_budget,
+        test_stage_gates,
+        test_exclusion_list_parsing,
+        test_an_exclusion_entry_matches_any_of_the_three_names,
+        test_an_excluded_model_stays_dense,
+        test_minimum_sequence_of_zero_sparsifies_everything,
+        test_stage_follows_the_step_schedule,
+        test_published_segments_reach_the_stage,
+        test_stage_uses_a_published_layout_and_falls_back_without_one,
+    ]:
+        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)

test/test-xyz-attention.py

@@ -0,0 +1,252 @@
+#!/usr/bin/env python
+"""
+Offline unit tests for the attention and sparse axes of the xyz grid.
+
+Covers:
+
+- every [Attention] and [Sparse] axis resolves its choices and targets a registered option
+- applying an axis and then leaving the grid restores shared.opts.data exactly, keys the axis
+  introduced included, so a grid never leaks its last cell into the session
+- the axes backed by boolean options take the string the dropdown hands them
+- the sdp override axis turns a label, a plus joined pair or None into the option's list
+- an override axis apply reaches the router: the rebuilt chain contains the backend it named
+- the restore set covers every setting the sparse stage reads
+
+No running server required. Nothing is moved to the accelerator, and no axis value that would
+install a package is used.
+
+Usage:
+    python test/test-xyz-attention.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 # importing shared installs the configured hijacks in-process
+
+from modules.errors import log                                    # pylint: disable=wrong-import-position
+from modules import attention, shared                             # pylint: disable=wrong-import-position
+from modules.attention.sparse import stage as sparse_stage        # pylint: disable=wrong-import-position
+from scripts.xyz import xyz_grid_shared as xyz                    # pylint: disable=wrong-import-position
+from scripts.xyz.xyz_grid_classes import axis_options             # pylint: disable=wrong-import-position
+
+
+results: dict[str, dict] = {}
+
+
+def category(name: str):
+    if name not in results:
+        results[name] = {'passed': 0, 'failed': 0, 'tests': []}
+    return name
+
+
+def record(cat: str, passed: bool, name: str, detail: str = ''):
+    status = 'PASS' if passed else 'FAIL'
+    results[cat]['passed' if passed else 'failed'] += 1
+    results[cat]['tests'].append((status, name))
+    msg = f'  {status}: {name}'
+    if detail:
+        msg += f' ({detail})'
+    if passed:
+        log.info(msg)
+    else:
+        log.error(msg)
+
+
+def run_test(cat: str, fn):
+    try:
+        ok = fn()
+        record(cat, ok is not False, fn.__name__)
+    except AssertionError as e:
+        record(cat, False, fn.__name__, str(e))
+    except Exception as e:  # pylint: disable=broad-except
+        record(cat, False, fn.__name__, f'exception: {e}')
+        import traceback
+        traceback.print_exc()
+
+
+def attention_axes():
+    return [axis for axis in axis_options if axis.label.startswith('[Attention]') or axis.label.startswith('[Sparse]')]
+
+
+def sample_value(axis):
+    """A value the axis accepts that is not the current one, avoiding backends whose prepare installs a package."""
+    if axis.choices is None:
+        return int(shared.opts.get(option_of(axis) or 'sparse_attention_budget') or 0) + 5
+    choices = [choice for choice in axis.choices() if choice not in ['Sage attention', 'Flash attention', 'Triton Flash attention']]
+    current = str(shared.opts.get(option_of(axis)) if option_of(axis) else '')
+    return next((choice for choice in choices if str(choice) != current), choices[0])
+
+
+def option_of(axis):
+    """The option an axis writes, read back from the closure the axis factory built."""
+    closure = getattr(axis.apply, '__closure__', None) or ()
+    for cell in closure:
+        if isinstance(cell.cell_contents, str) and cell.cell_contents in shared.opts.data_labels:
+            return cell.cell_contents
+    return {'[Attention] SDP override': 'sdp_overrides', '[Attention] Dispatcher': 'hf_attention'}.get(axis.label, None)
+
+
+def bool_axes():
+    return [axis for axis in attention_axes() if isinstance(shared.opts.get(option_of(axis) or ''), bool)]
+
+
+# ============================================================
+# Tests
+# ============================================================
+
+def test_axes_are_registered():
+    axes = attention_axes()
+    assert len(axes) >= 10, f'only {len(axes)} attention axes'
+    for axis in axes:
+        option = option_of(axis)
+        assert option is not None, f'{axis.label} names no option'
+        assert option in shared.opts.data_labels, f'{axis.label} targets unknown option {option}'
+    return True
+
+
+def test_axis_choices_resolve():
+    for axis in attention_axes():
+        if axis.choices is None:
+            assert axis.type is int, f'{axis.label} has no choices and is not numeric'
+            continue
+        choices = axis.choices()
+        assert isinstance(choices, list) and len(choices) > 0, f'{axis.label} resolved {choices}'
+    return True
+
+
+def test_axes_write_and_restore_exactly():
+    for axis in attention_axes():
+        before = dict(shared.opts.data)
+        saved = xyz.save_attention()
+        axis.apply(None, sample_value(axis), [])
+        assert dict(shared.opts.data) != before, f'{axis.label} wrote nothing'
+        xyz.restore_attention(saved)
+        after = dict(shared.opts.data)
+        leaked = {key for key in set(before) | set(after) if before.get(key, '<absent>') != after.get(key, '<absent>')}
+        assert not leaked, f'{axis.label} leaked {sorted(leaked)}'
+    return True
+
+
+def test_bool_axes_coerce_the_string_the_dropdown_sends():
+    axes = bool_axes()
+    assert len(axes) >= 3, f'only {len(axes)} boolean axes'
+    for axis in axes:
+        option = option_of(axis)
+        saved = xyz.save_attention()
+        axis.apply(None, 'True', [])
+        assert shared.opts.data[option] is True, f'{axis.label} took "True" as {shared.opts.data[option]!r}'
+        axis.apply(None, 'False', [])
+        assert shared.opts.data[option] is False, f'{axis.label} took "False" as {shared.opts.data[option]!r}'
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_override_axis_parses_labels():
+    saved = xyz.save_attention()
+    try:
+        xyz.apply_attention_overrides(None, 'None', [])
+        assert shared.opts.data['sdp_overrides'] == [], shared.opts.data['sdp_overrides']
+        xyz.apply_attention_overrides(None, 'Flex attention', [])
+        assert shared.opts.data['sdp_overrides'] == ['Flex attention'], shared.opts.data['sdp_overrides']
+        xyz.apply_attention_overrides(None, 'Flex attention+SDNQ attention', [])
+        assert shared.opts.data['sdp_overrides'] == ['Flex attention', 'SDNQ attention'], shared.opts.data['sdp_overrides']
+    finally:
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_override_axis_rebuilds_the_chain():
+    saved = xyz.save_attention()
+    try:
+        xyz.apply_attention_overrides(None, 'Flex attention', [])
+        assert 'flex' in attention.get_plan().chain(), attention.get_plan().chain()
+        xyz.apply_attention_overrides(None, 'None', [])
+        assert 'flex' not in attention.get_plan().chain(), attention.get_plan().chain()
+    finally:
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_dispatcher_axis_clears_on_none():
+    saved = xyz.save_attention()
+    try:
+        xyz.apply_attention_dispatcher(None, 'native', [])
+        assert shared.opts.data['hf_attention'] == 'native', shared.opts.data['hf_attention']
+        xyz.apply_attention_dispatcher(None, 'None', [])
+        assert shared.opts.data['hf_attention'] == '', repr(shared.opts.data['hf_attention'])
+    finally:
+        xyz.restore_attention(saved)
+    return True
+
+
+def test_restore_set_covers_every_attention_setting():
+    covered = set(xyz.attention_options())
+    missing = set(sparse_stage.OPTION_NAMES) - covered
+    assert not missing, f'sparse settings outside the restore set: {sorted(missing)}'
+    for backend in attention.registry.backends.values():
+        outside = set(backend.options) - covered
+        assert not outside, f'{backend.label} settings outside the restore set: {sorted(outside)}'
+    return True
+
+
+def run_all():
+    log.warning('=== xyz attention axes ===')
+    cat = category('axes')
+    for fn in [
+        test_axes_are_registered,
+        test_axis_choices_resolve,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== apply and restore ===')
+    cat = category('apply')
+    for fn in [
+        test_axes_write_and_restore_exactly,
+        test_bool_axes_coerce_the_string_the_dropdown_sends,
+        test_override_axis_parses_labels,
+        test_override_axis_rebuilds_the_chain,
+        test_dispatcher_axis_clears_on_none,
+        test_restore_set_covers_every_attention_setting,
+    ]:
+        run_test(cat, fn)
+
+    log.warning('=== Results ===')
+    total_passed = 0
+    total_failed = 0
+    for cat_name, info in results.items():
+        status = 'PASS' if info['failed'] == 0 else 'FAIL'
+        log.info(f"  {cat_name}: {info['passed']} passed, {info['failed']} failed [{status}]")
+        total_passed += info['passed']
+        total_failed += info['failed']
+    log.warning(f'Total: {total_passed} passed, {total_failed} failed')
+    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)

ui/locale/locale_en.json

@@ -1465,15 +1465,22 @@
     {"id":"","label":"Sections","localized":"","hint":"","ui":"video"},
     {"id":"","label":"Samplers","localized":"","hint":"Samplers/schedulers advanced settings","ui":"tab_txt2img"},
     {"id":"","label":"SDP kernels","localized":"","hint":"Which of torch's built-in attention kernels torch is allowed to choose from. These are permissions rather than a selection: torch picks one per call from whatever is left enabled, preferring <b>Flash</b>, dropping to <b>Memory</b> for calls flash cannot serve such as those carrying an arbitrary attention mask, and to <b>Math</b> when neither fits. Clearing a box removes a candidate; it never pins the remaining one to every call.<br><br><b>Flash</b> is torch's own build of the FlashAttention kernel. It is not the same thing as the <b>Flash attention</b> entry in <b><i>SDP overrides</i></b>, which calls the separately installed flash-attn package and bypasses torch entirely.<br><b>Memory</b> is the memory-efficient kernel, which accepts arbitrary masks that flash does not.<br><b>Math</b> is the unfused reference path, the widest in what it accepts and the least optimized. Leaving it enabled keeps a fallback for calls the other two decline.<br><br>Applies while <b><i>Attention method</i></b> is <b>Scaled-Dot-Product</b>, and continues to govern the calls that an enabled override declines.<br><br>All three by default. ZLUDA starts with <b>Math</b> alone.","ui":"settings_cuda"},
-    {"id":"","label":"SDP overrides","localized":"","hint":"Replaces torch attention with another implementation. Each entry declares the shapes, dtypes and mask conditions it can serve; a call that fails them moves to the next entry and finally back to torch, so several can be enabled together and the chain resolves per call.<br><br><b>Flash attention</b> installs and calls the flash-attn package directly, for calls with no attention mask, half precision inputs and a head dimension of 128 or less.<br><b>Sage attention</b> computes attention with quantized matmuls, for head dimensions of 64, 96 or 128 and no attention mask.<br><b>SDNQ attention</b> is SD.Next's own quantized Triton kernel, configured in the section below. It takes attention masks, which the other quantized backends do not.<br><b>Flex attention</b> uses torch's compiled flex_attention.<br><b>Dynamic attention</b> slices attention to fit available memory and serves every call the others decline, standing in for the torch fallback.<br><b>Triton Flash attention</b> is a Triton implementation for ROCm and ZLUDA, listed only on those backends.<br><br>The quantized and compiled backends trade some numerical accuracy for throughput. How much of each arrives depends on the model, the sequence length and the GPU, so comparing them on the actual workload settles it faster than picking by reputation.<br><br>None by default on CUDA. ZLUDA, CPU and MPS start with <b>Dynamic attention</b>, as do ROCm GPUs older than RDNA3.","ui":"settings_cuda"},
-    {"id":"","label":"SDNQ Attention","localized":"","hint":"Settings for the <b>SDNQ attention</b> entry in <b><i>SDP overrides</i></b>. They do nothing until that override is enabled.<br><br>The kernel quantizes the two matmuls inside attention, computing them on lower precision operands and rescaling the result. It is written in Triton, so it needs a working Triton for the active device.<br>Short sequences and single-head calls are left to the rest of the chain, so text encoders and the VAE keep ordinary attention.","ui":"settings_cuda"},
+    {"id":"","label":"SDP overrides","localized":"","hint":"Replaces torch attention with another implementation. Each entry declares the shapes, dtypes and mask conditions it can serve; a call that fails them moves to the next entry and finally back to torch, so several can be enabled together and the chain resolves per call.<br><br><b>Flash attention</b> installs and calls the flash-attn package directly, for calls with no attention mask, half precision inputs and a head dimension of 128 or less.<br><b>Sage attention</b> computes attention with quantized matmuls, for head dimensions of 64, 96 or 128 and no attention mask.<br><b>SDNQ attention</b> is SD.Next's own quantized Triton kernel, configured in the section below. It takes attention masks, which the other quantized backends do not, and it is one of the two backends <b><i>Sparse Attention</i></b> can drive.<br><b>Flex attention</b> uses torch's compiled flex_attention, the other backend <b><i>Sparse Attention</i></b> can drive.<br><b>Dynamic attention</b> slices attention to fit available memory and serves every call the others decline, standing in for the torch fallback.<br><b>Triton Flash attention</b> is a Triton implementation for ROCm and ZLUDA, listed only on those backends.<br><br>The quantized and compiled backends trade some numerical accuracy for throughput. How much of each arrives depends on the model, the sequence length and the GPU, so comparing them on the actual workload settles it faster than picking by reputation.<br><br>None by default on CUDA. ZLUDA, CPU and MPS start with <b>Dynamic attention</b>, as do ROCm GPUs older than RDNA3.","ui":"settings_cuda"},
+    {"id":"","label":"SDNQ Attention","localized":"","hint":"Settings for the <b>SDNQ attention</b> entry in <b><i>SDP overrides</i></b>. They do nothing until that override is enabled.<br><br>The kernel quantizes the two matmuls inside attention, computing them on lower precision operands and rescaling the result. It is written in Triton, so it needs a working Triton for the active device.<br>It takes an attention mask and a block mask together, which is what lets <b><i>Sparse Attention</i></b> use it.<br>Short sequences and single-head calls are left to the rest of the chain, so text encoders and the VAE keep ordinary attention.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention use Smooth K","localized":"","hint":"Subtracts the mean of the keys before quantizing them. Keys carry a large offset that is shared across the sequence, which spends most of the quantized range representing a value identical for every key and leaves little of it for the differences that decide the attention.<br>Softmax ignores a constant shift applied to every score in a row, so removing that offset changes the quantization error and not the attention.<br><br>Costs one mean and one subtraction per call.<br>Enabled by default.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention use Hadamard","localized":"","hint":"Rotates queries and keys by a Hadamard transform before quantizing them. The rotation spreads a few oversized channels across all of them, which is the error shape quantization handles worst. The transform is orthogonal, so the scores it produces are the ones the unrotated tensors would produce, minus the quantization error it removes.<br>With <b><i>SDNQ Attention PV MatMul type</i></b> also set, the values are rotated as well and the output is rotated back.<br><br>Costs a rotation pass on every attention call, so it is worth enabling where a model shows quantization artifacts without it.<br>Idle while <b><i>SDNQ Attention MatMul type</i></b> is <b>disabled</b>, since nothing is quantized then.<br><br>Disabled by default.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention use FP16 Accumulation","localized":"","hint":"Accumulates the floating point matmuls in fp16 rather than fp32. Some tensor cores run fp16 accumulation at a higher rate than fp32, and on those the kernel is cheaper for it.<br>Operands are pre-scaled to keep products inside the fp16 range, which covers ordinary activations with less headroom than fp32 leaves.<br><br>Reaches the parts of the kernel that run in floating point. An int8 matmul accumulates in int32 and is unaffected, so at the default <b><i>SDNQ Attention MatMul type</i></b> this applies to the probability-value matmul alone.<br><br>Disabled by default.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention MatMul type","localized":"","hint":"Precision the query-key matmul is computed in, the first of the two matmuls in attention.<br><br><b>enabled</b> selects int8, and <b>int8</b> and <b>uint8</b> reach the same kernel.<br><b>float16</b> and <b>float8_e4m3fn</b> take the floating point path. fp8 needs a GPU with fp8 tensor cores and fails on hardware without them rather than falling back.<br><b>disabled</b> leaves queries and keys in the model's own precision, which also idles <b><i>SDNQ Attention use Smooth K</i></b> and <b><i>SDNQ Attention use Hadamard</i></b>.<br><br>Default enabled.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention PV MatMul type","localized":"","hint":"Precision the probability-value matmul is computed in, the second of the two matmuls in attention. Choices match <b><i>SDNQ Attention MatMul type</i></b>.<br><br>Quantizing this one as well takes out the floating point work the first setting leaves behind, and it is the more delicate of the two: its inputs are already normalized probabilities, and the small ones among them carry the fine detail.<br><b>disabled</b> keeps this matmul in the model's own precision.<br><br>Default disabled.","ui":"settings_cuda"},
     {"id":"","label":"SDNQ Attention Hadamard Group Size","localized":"","hint":"Width of the Hadamard rotation in channels. Wider groups mix more channels together and spread outliers further.<br><br>Clamped to the head dimension of the running model, rounded down to a power of two that divides it. On a model with 64 or 128 channels per head the upper part of this range resolves to that head dimension rather than to the number shown. Rotation is skipped below 4.<br>Applies while <b><i>SDNQ Attention use Hadamard</i></b> is enabled.<br><br>Default 256.","ui":"settings_cuda"},
-    {"id":"","label":"SDNQ Attention Quantize FP32","localized":"","hint":"Upcasts queries, keys and values to fp32 for the quantization step, meaning the mean subtraction, scale and rounding that produce the low precision operands. The matmuls themselves are unaffected, and the kernel applies the scales in fp32 either way.<br>Turned off, that arithmetic runs in the model's own precision. bf16 carries eight mantissa bits, so a scale derived in it is coarser than one derived in fp32, and <b><i>SDNQ Attention use Smooth K</i></b> loses the most from it, since a mean across the whole sequence is exactly the kind of sum that wants the extra bits.<br><br>Whether the upcast costs anything depends on how the GPU runs fp32 vector work against fp16 and bf16. NVIDIA and AMD run them at the same rate here, so there is nothing to save; Intel runs fp32 slower and takes a noticeable hit.<br><br>Enabled by default.","ui":"settings_cuda"}
+    {"id":"","label":"SDNQ Attention Quantize FP32","localized":"","hint":"Upcasts queries, keys and values to fp32 for the quantization step, meaning the mean subtraction, scale and rounding that produce the low precision operands. The matmuls themselves are unaffected, and the kernel applies the scales in fp32 either way.<br>Turned off, that arithmetic runs in the model's own precision. bf16 carries eight mantissa bits, so a scale derived in it is coarser than one derived in fp32, and <b><i>SDNQ Attention use Smooth K</i></b> loses the most from it, since a mean across the whole sequence is exactly the kind of sum that wants the extra bits.<br><br>Whether the upcast costs anything depends on how the GPU runs fp32 vector work against fp16 and bf16. NVIDIA and AMD run them at the same rate here, so there is nothing to save; Intel runs fp32 slower and takes a noticeable hit.<br><br>Enabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention","localized":"","hint":"Computes attention over a subset of the key tiles instead of all of them. Attention cost grows with the square of the sequence length, so on long sequences it dominates generation time, and skipping the tiles that contribute least buys much of it back.<br>The saving grows with sequence length: negligible on a short sequence, useful at high resolution, largest on video.<br><br>Requires an attention backend that accepts a block mask, <b>SDNQ attention</b> or <b>Flex attention</b> in <b><i>SDP overrides</i></b>; with neither enabled a warning is logged and attention stays dense. It also stays dense below <b><i>Sparse Attention minimum sequence</i></b>.<br>Not every architecture tolerates a reduced key set. The ones known to break are listed in <b><i>Sparse Attention excluded models</i></b> and stay dense; a model that breaks up rather than merely softening belongs on that list.<br><br>Disabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention KV budget","localized":"","hint":"Percentage of the eligible key tiles each query tile keeps. Lower budgets are faster and coarser, and the cost shows first in fine detail and in consistency across the image.<br>Text, conditioning and audio tokens are always kept, as are the tiles on the diagonal, so the budget applies only to the bulk image or video tokens.<br><br><b>100</b> keeps every tile, which is ordinary dense attention and a convenient comparison run.<br>Default 30.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention minimum sequence","localized":"","hint":"Shortest sequence that is sparsified. Below it attention stays dense, because choosing the tiles costs more than skipping them saves.<br>Sequence length is not resolution: a diffusion transformer sees roughly (width/16) x (height/16) tokens for an image, so 1024x1024 is about 4k tokens and 2048x2048 about 16k, and video multiplies that by the frame count.<br><br><b>0</b> sparsifies every sequence that reaches the stage.<br>Default 8192, around 1450x1450 for an image.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention dense steps","localized":"","hint":"Number of steps at the start and end of sampling that receive a larger budget, where composition and fine detail are set. Counted at each end, and capped at half the run.<br><br>Only takes effect when <b><i>Sparse Attention dense step bonus</i></b> is above 0.<br><b>0</b> applies one budget to every step.<br>Default 0.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention dense step bonus","localized":"","hint":"Percentage points added to the budget on the dense steps, capped at 100.<br><br>Only takes effect when <b><i>Sparse Attention dense steps</i></b> is above 0.<br>Default 30, so a budget of 30 rises to 60 on those steps.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention share selection across heads","localized":"","hint":"Computes one selection for all attention heads rather than one per head, by averaging the heads before scoring. Cheaper to select and coarser in what it keeps, since heads that attend to different regions are served by a single compromise.<br>Worth trying when selection is itself a visible share of the cost, which happens on models with many heads.<br>Also the first thing to try when a model breaks up into bands under sparse attention: some architectures need every head to see one consistent context, and per-head selection is what breaks them.<br><br>Disabled by default.","ui":"settings_cuda"},
+    {"id":"","label":"Sparse Attention excluded models","localized":"","hint":"Models that stay dense no matter how the rest of this section is set. Comma separated, matched case insensitively against the architecture, the pipeline class and the denoiser class, so whichever of those names is to hand works as an entry.<br>The class names appear in the model load log; the architecture is the short name used elsewhere in the settings, such as <b>f1</b> or <b>anima</b>.<br><br>Listed by default is <b>CosmosTransformer3DModel</b>, the transformer Anima runs, which collapses into banded noise when each head selects its own tiles and stays degraded even with <b><i>Sparse Attention share selection across heads</i></b> enabled. The class is listed rather than the architecture because the other models built on it have not been checked.<br>A model whose output breaks up rather than merely softening belongs here.<br><br>Default <b>CosmosTransformer3DModel</b>.","ui":"settings_cuda"}
   ],
   "t": [
     {"id":"txt2img_nav","label":"T2I","localized":"","hint":"Create image from text<br>Legacy interface that mimics original text-to-image interface and behavior"},