perf(attn): step KV buckets by 4/3 instead of powers of two

[jvlunteren]

Description

Steps the KV buckets by 4/3 instead of powers of two.

A sequence past a bucket boundary rounds up to nearly double its real block count. Padded blocks are real KV reads and decode attention is data-movement bound, so that round-up costs latency. A 4/3 step bounds the worst-case at a quarter of the bucket's block count rather than a half.

_token_buckets_up_to anchors at block_size so every step is a whole number of blocks. A non-power-of-two block_size needs no rounding, the warning about that case is gone.

The bucketer also gains the num_seqs ladder for the batched decode kernel and batched_variants() for pre-warming.

The cost is more recorded variants at warmup (13 vs 7 at max_model_len=8192, block_size=128). Since SPYRE_KERNEL_CACHE (#811) landed, repeated restarts do not re-pay the compile cost. SPYRE_ATTN_KV_BUCKETS set to powers of two restores the previous behaviour.

Related Issues

#771 #772 #791 #810 #811

Test Plan

• pytest tests/attention/test_spyre_attn.py tests/runtime/test_spyre_attn_bucketer.py tests/attention/test_spyre_attn_recorder.py -m 'not upstream' : all passed

Checklist

• I have read the contributing guidelines
• My code follows the project's code style (run bash format.sh)
• I have added tests for my changes (if applicable)
• I have updated the documentation (if applicable)
• My commits include a Signed-off-by: line (DCO compliance)

[github-actions]

👋 Hi! Thank you for contributing.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, run ./format.sh.
Now you are good to go 🚀.

We also recommend installing prek and configuring it to check your code before every local commit.

[sducouedic]

so for max-model-len 8192 we use the bucketting for kv-axis (block_size=128):

(128, 256, 384, 512, 768, 1024, 1408, 1920, 2560, 3456, 4608, 6144, 8192), len=13
instead of:
(128, 256, 512, 1024, 2048, 4096, 8192), len=7

given that this is bucketting is shared by both the pure-decode path and the mixed/prefill for-loop path, are we quite convinced that the performance gain is worth spending the additional time compiling? knowing that this kv-axis bucketting composes with the other dimensions bucketting (kind of combinatorial)

@tdoublep @jvlunteren

[sducouedic]

would it work to directly set self._num_seqs_buckets with values that are reachable (ie. >= _MIN_BATCHED_SEQS, so we don't need to filter them here

[jvlunteren]

Done. MIN_BATCHED_SEQS is now defined in spyre_attn_bucketer.py (where it semantically belongs, it is a bucketing threshold, not an attention-backend detail) and _num_seqs_buckets is filtered at construction time:

self._num_seqs_buckets: list[int] = [
    n for n in _powers_of_two_up_to(max_num_seqs) if n >= MIN_BATCHED_SEQS
]

batched_variants() now iterates self._num_seqs_buckets directly with no late import and no inline filter. spyre_attn.py keeps a local _MIN_BATCHED_SEQS = MIN_BATCHED_SEQS alias for its two internal uses.

[sducouedic]

For a future PR: might be worth checking #789, which eliminated needs_gather for the non-batch version

[jvlunteren]

Agreed, good follow-on. One detail to consider before opening it: the sequential path's staging trick makes needs_gather always True (buffer always wider than the query bucket). The batched path mirrors this in the opposite direction. needs_gather = query_dev.shape[0] < b_seqs would become always False if the buffer is always wider than b_seqs. Since store_out is only reachable when needs_gather=False, this collapses both axes and roughly halves the batched variant count.

The key constraint is that needs_gather=False assumes decode sequences occupy rows 0 to num_seqs-1 contiguously, which is not guaranteed in the mixed-batch case. The follow-on PR should gate needs_gather=False on is_pure_decode (i.e. decode_seq_indices is None) and leave the mixed-batch path on the existing gather variant.

[jvlunteren]

so for max-model-len 8192 we use the bucketting for kv-axis (block_size=128):

(128, 256, 384, 512, 768, 1024, 1408, 1920, 2560, 3456, 4608, 6144, 8192), len=13 instead of: (128, 256, 512, 1024, 2048, 4096, 8192), len=7

given that this is bucketting is shared by both the pure-decode path and the mixed/prefill for-loop path, are we quite convinced that the performance gain is worth spending the additional time compiling? knowing that this kv-axis bucketting composes with the other dimensions bucketting (kind of combinatorial)

@tdoublep @jvlunteren

Good point. I have updated the PR to cap the 4/3 ladder at 4096 tokens and use powers-of-two above that. The asymmetry justifies this: at short KV the round-up is a large fraction of real work (128 to 256 wastes 50% of a real index_select and matmul iteration), so dense buckets save the most there. At long KV each extra padded block is a small fraction of the total, so powers-of-two keep the variant count down without meaningful cost. For max_model_len=8192 with block_size=128 the new ladder is 13 dense buckets up to 4096 plus [8192], 14 total. SPYRE_ATTN_KV_BUCKETS remains the escape hatch for full density if needed.

On the combinatorial concern: both _num_seqs_buckets and _num_blocks_buckets on the builder are sourced from SpyreAttnBucketer, so there is no independent derivation and no drift between what warmup records and what dispatch requests.