refactor(attn): let the bucketer define the batched decode kernel's buckets

[jvlunteren]

Description

Makes SpyreAttnBucketer the single source of truth for attention bucket sizes, and has warmup pre-compile the batched decode kernel as well as the per-sequence one. Per review feedback on #772.

The metadata builder computed its own seqs and blocks buckets in parallel with the bucketer computing the same thing, so the set build() dispatched onto and the set warmup recorded agreed only by coincidence. It now reads both from the bucketer. The batched kernel keys on a different tuple from the per-sequence one, so it gets its own bucket type and enumeration, and record_graphs compiles both. Previously the batched kernel compiled lazily in the serving path, the stall warmup exists to prevent. Buckets dispatch cannot reach are pruned.

Rebased onto #789, which routes query and output through fixed-size staging buffers: needs_gather now comes from the sequence counts rather than the query width, which staging made constant, and the recorder records on those buffers instead of fresh tensors.

Stacked on #791, since it uses the new names throughout.

Related Issues

#772 #789 #791

Test Plan

• pytest tests/runtime/test_spyre_attn_bucketer.py : 69 passed
• pytest tests/attention/test_spyre_attn.py --attn-shards=4 --attn-shard-id=0 : 46 passed
• Graph reuse verified on device, with a mismatched control call to confirm the check can fail

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]

Look good and clean! just a two small comments

[sducouedic]

should we print a warning here, or even an assert? I think vllm upstream enforces --num-gpu-blocks-override to serve at least one max-model-len request

[sducouedic]

as in the other PR, is there anything preventing us to set _num_seqs_buckets with values already >= _MIN_BATCHED_SEQS

[sducouedic]

Superseded by #810.

Most of this has landed independently since the merge-base (0e83fc8): the rename is on main as #791, SPYRE_BATCHED_DECODE is on main, and #799 replaced the per-variant kernel factory with Dynamo-managed compilation — which removes _get_batched_decode_kernel and _decode_fns, the cache batched_variants() was built to mirror. With Dynamo keying on the flags itself, enumerating needs_gather/store_out now traces the same compiled function several times per (num_seqs, num_blocks) pair.

The idea that's still un-done is the one this PR opens with: the bucketer should own the ladders. #810 does that, with you as co-author.

It also fixes a related divergence I hit while extracting it — the builder's num_blocks ladder disagrees with the bucketer's under SPYRE_ATTN_KV_BUCKETS:

default KV buckets:                     both ladders → [1,2,4,8,16,32]   agree
SPYRE_ATTN_KV_BUCKETS=512,1024,2048:    bucketer     → [8,16,32]
                                        old builder  → [1,2,4,8,16,32]   diverge

The builder dispatches onto [1,2,4], which warmup never records — an Inductor compile mid-request, exactly what the recorder exists to prevent. Sharing one ladder makes that unrepresentable.

Leaving this open for you to close, in case you see something in batched_variants() worth keeping.