[dexhunter]

TL;DR

• This PR speeds up _scatter_diagonals in the minference CPU reference (vertical_and_slash_attention).
• It marks the chosen diagonals in a small per-head table, then reads the table once with diag_id.
• A mixed benchmark of reference calls went from 104.40 ms to 25.59 ms (75.5% less time, about 4.1x faster).
• Tested with pytest tests/sparse/...: same result as main (72 passed; the 5 failures are the same CUDA/Triton-only tests).

What & Why

Before, _scatter_diagonals looped over every chosen diagonal. Each pass compared the full (B, H, Lq, Lk) id grid with one id and merged the result into the mask. The budget is at least 50 diagonals, so the code read the whole grid 50 to 500 times.

Now the function builds one bool table per head over the diagonal ids. It sets True at the chosen ids with scatter_. Then it reads the table once with diag_id. When diag_id is shared by all heads (the normal call site), it indexes only the last table dimension.

The output stays identical. An entry is True exactly when its diagonal id is one of the chosen ids, as before. The table covers the full id range, so negative or large ids also work.

Changes

• angelslim/compressor/sparsity/algorithms/minference/reference.py: replace the per-diagonal compare loop in _scatter_diagonals with a table lookup.
• Return an all-False mask early when there are no chosen diagonals or the grid is empty.
• No other function or public signature changes.

Results

Median time per vertical_and_slash_attention call, CPU, one thread. The outputs are bit-identical (torch.equal) in every row.

Situation	Before	After	Change
Repo test config (L=512, 1 head, 120 diagonals)	17.24 ms	5.43 ms	3.2x faster
L=1024, 4 heads, 300 diagonals	660.46 ms	101.30 ms	6.5x faster
L=2048, 2 heads, 500 diagonals	2445.17 ms	265.16 ms	9.2x faster
L=512, 1 head, 50 diagonals (minimum budget)	8.03 ms	3.20 ms	2.5x faster
a_shape variant, L=1024, 4 heads (does not use this code)	32.28 ms	32.28 ms	no change

The gain grows with sequence length and with the number of diagonals, because the old loop did one full pass per diagonal.

Why this is safe and easy to maintain

The change is small: one function, about 15 lines, with plain scatter_ and indexing. The mask has the same values, dtype and shape as before. test_vertical_and_slash_diagonal_alignment and the other tests/sparse tests cover this path, and they pass.

Validation

• Native tests on this branch and on main (same result on both):
  python -m pytest -q tests/sparse/test_algorithms_integration.py tests/sparse/test_stem_integration.py tests/sparse/test_guards.py tests/sparse/test_save_load.py
  -> 72 passed, 5 failed, 57 skipped. The 5 failures are the xattention/flexprefill/flashprefill head_dim and block-size checks and test_sparsity_registered_in_compressor_factory. They also fail on main on a CPU-only host without CUDA/Triton kernels, and they do not use this code.
• Benchmark: a mixed set of reference calls (the repo test config, two random single-head shapes and a two-head minference_reference_prefill dispatch). The before and after runs were interleaved on an isolated CPU core with OMP_NUM_THREADS=1, and every output was checked to be exactly equal. 104.40 ms -> 25.59 ms per run, 3 repeats (ratios 0.280, 0.243, 0.243).
• Quick reproduction:

  import time, torch
  from angelslim.compressor.sparsity.algorithms.minference.reference import vertical_and_slash_attention
  torch.set_num_threads(1)
  q, k, v = (torch.randn(1, 4, 1024, 16) for _ in range(3))
  t = time.perf_counter(); vertical_and_slash_attention(q, k, v, 64, 300, 64); print(time.perf_counter() - t)

AI assistance

OpenAI Codex (gpt-6.1-sol) generated candidate versions of this change, and Weco Observe tracked them. Claude Code picked the simplest fast candidate and ran the tests and benchmarks above. The diff is the candidate source, unchanged.

Weco share link: https://dashboard.weco.ai/share/81ec1Qup7kZzZJiWFCcWekBExmbj_ZDz