perf(attention): drop the per call readbacks from block selection

Selecting blocks synchronized with the accelerator three times per
attention call: once to size the top-k, once to decide whether the
budget covered every candidate, and once to compute the density for
reporting. At the sdxl shape those stalls cost 1.6 ms against a 1.2 ms
attention, so the selector lost to dense at every budget.

The parts that depend only on geometry and layout are now built once and
cached per layout, ranking replaces top-k so no host side k is needed,
and density became a method that reads back only when something asks.
The block mask also stops sorting a mask of zeros to fill partial slots
it leaves empty. Selection overhead at sdxl drops from 1.84 ms to
0.23 ms, and a 30 percent budget moves from 0.36x of dense to 1.19x.
This commit is contained in:
CalamitousFelicitousness
2026-08-22 23:54:30 +01:00
parent 039227119c
commit fc0fd41ecd
4 changed files with 42 additions and 22 deletions
+2 -2
View File
@@ -258,7 +258,7 @@ 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}'
assert abs(control.density() - density) < 0.05, f'requested {density}, got {control.density()}'
return True
@@ -410,7 +410,7 @@ def test_stage_uses_a_published_layout_and_falls_back_without_one():
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 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)