Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 17:24:32 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(attention): sdnq backend consumes the block selection

Browse Source 
The sdnq Triton kernel gains a block mask input (on the submodule's feat/block-mask branch; the pointer here is unchanged), so the backend advertises the block_mask cap and passes the router's selection as block_mask plus its block sizes. Prepare probes the installed kernel's signature and narrows the cap when the input is missing, warning when sparse attention is on, so an older submodule stays dense instead of failing. PlanEntry carries the narrowed caps, and the router, the sparse stage and the report read those rather than the declaration.

- benchmark: int8-sparse100/50/30/15 and int8-radial30 rows feed the same producer as the flex rows into the quantized kernel, skipped on builds without block_mask; a --configs filter selects attention rows the way --block-configs selects block rows; sparse rows of either backend are held out of the settings advice
- test/test-attention-sdnq-sparse.py: the kernel contract against the token-expanded mask, the nesting filter, the public entry on the quantized paths, the backward against the same selection and against fp32 autograd, and the flex consumer on one selection, pinned to one tile per run with a --tiles sweep
- locale: the sparse attention hint names both capable backends
This commit is contained in:
CalamitousFelicitousness
2026-08-23 04:11:43 +01:00
parent c7b18bc926
commit b2de8f53c4
6 changed files with 811 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/sdnq-attention-benchmark.py
+56 -6
View File
@@ -249,6 +249,11 @@ bench_configs = [
("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)),
@@ -269,7 +274,8 @@ bench_configs = [
# tail lives outside what mean error can see
# 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", "flex-sparse100", "flex-sparse50", "flex-sparse30", "flex-sparse15", "flex-radial30")
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
@@ -357,6 +363,7 @@ 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("--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)")
@@ -504,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
@@ -530,6 +547,25 @@ def triton_mm_fp16_accum():
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():
@@ -983,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)}")
@@ -1000,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(),
@@ -1485,7 +1524,7 @@ 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=None, 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"]
@@ -1507,7 +1546,7 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=None, fp8_r
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
@@ -1517,6 +1556,8 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=None, fp8_r
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]):
@@ -1579,6 +1620,8 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=None, fp8_r
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)
@@ -2890,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)
@@ -2924,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
@@ -3743,6 +3786,13 @@ def main():
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":
@@ -3876,7 +3926,7 @@ 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.timeout_flag, 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"))
modules/attention/backends/sdnq.py
+18 -2
View File
@@ -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'}),
)
modules/attention/router.py
+5 -4
View File
@@ -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:
@@ -69,7 +70,7 @@ def make_router(plan: Plan, observer: Callable | None = None, stage: Callable |
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.backend.caps:
if stage is not None and 'block_mask' in entry.caps:
selection = stage(query, key, value, attn_mask, is_causal)
if selection is not None:
if observer is not None:
@@ -112,7 +113,7 @@ def build_sparse_stage(plan: Plan):
return None
if not options.enabled:
return None
capable = [entry.backend.name for entry in plan.entries if 'block_mask' in entry.backend.caps]
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')
@@ -158,7 +159,7 @@ def report() -> dict:
'enabled': options.enabled,
'budget': options.budget,
'gate': options.gate,
'capable': [entry.backend.name for entry in plan.entries if 'block_mask' in entry.backend.caps] if plan is not None else [],
'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,
test/test-attention-router.py
+35
View File
@@ -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
@@ -515,6 +549,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
+696
View File
@@ -0,0 +1,696 @@
#!/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, batch: int, heads: int):
"""The count and ascending index list per query block that the kernel walks, built the way the entry builds them: size-1 batch and head dims kept, padded to the descriptor chunks."""
return atten_module.get_block_mask_input(keep, batch, heads)
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, q.shape[0], q.shape[1])
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, 2, 4)
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, 2, 4)
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, 2, 4)
wide_count, wide_index = ordered(ones.expand(5, 4, -1, -1).contiguous(), 5, 4)
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_nesting_filter_keeps_only_nesting_tiles():
configs = configs_for((32, 64, 128, 256), (16, 32, 64, 128))
args = {'do_block_mask': 1, 'BLOCK_MASK_M': BLOCK_M, 'BLOCK_MASK_N': BLOCK_N}
kept = atten_module.nest_block_mask_configs(configs, args)
assert {(c.kwargs['BLOCK_SIZE_M'], c.kwargs['BLOCK_SIZE_N']) for c in kept} == {(m, n) for m in (32, 64, 128) for n in (16, 32, 64)}
assert atten_module.nest_block_mask_configs(configs, {'do_block_mask': 0}) is configs
assert atten_module.nest_block_mask_configs(configs, {}) is configs
try:
atten_module.nest_block_mask_configs(configs_for((256,), (128,)), args)
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')
def test_prune_configs_applies_the_filter():
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])
return True
# ============================================================
# 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_nesting_filter_keeps_only_nesting_tiles,
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)
ui/locale/locale_en.json
+1 -1
View File
@@ -1471,7 +1471,7 @@
{"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":"Sparse Attention","localized":"","hint":"Attention is computed on a subset of key tiles instead of all of them, which lowers cost on long sequences. Requires an attention backend that accepts a block mask, currently <b>Flex attention</b> in <b><i>SDP overrides</i></b>. Below the minimum sequence the setting stays inactive and attention is dense.","ui":"settings_cuda"},
{"id":"","label":"Sparse Attention","localized":"","hint":"Attention is computed on a subset of key tiles instead of all of them, which lowers cost on long sequences. 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>. Below the minimum sequence the setting stays inactive and attention is dense.","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. 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.","ui":"settings_cuda"},
{"id":"","label":"Sparse Attention minimum sequence","localized":"","hint":"Shortest sequence that is sparsified. Below it attention stays dense, because the selection costs more than it saves. <b>0</b> uses the built in threshold.","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 detail are set. <b>0</b> applies one budget to every step.","ui":"settings_cuda"},