refactor(attn): let the bucketer own the batched decode buckets - #810
refactor(attn): let the bucketer own the batched decode buckets#810sducouedic wants to merge 3 commits into
Conversation
|
👋 Hi! Thank you for contributing. We also recommend installing prek and configuring it to check your code before every local commit. |
SpyreAttentionMetadataBuilder built its own num_seqs and num_blocks ladders while SpyreAttnBucketer derived a num_blocks ladder of its own from the KV buckets. The two agree only on the default geometric KV buckets; under SPYRE_ATTN_KV_BUCKETS the builder dispatches onto low block counts that warmup never records, putting an Inductor compile in the serving path. Both ladders now come from the bucketer, so the set build() dispatches onto and the set warmup records are the same by construction. Signed-off-by: Sophie du Couédic <sop@zurich.ibm.com> Co-authored-by: Jan van Lunteren <161835099+jvlunteren@users.noreply.github.com>
731f817 to
7e6e61a
Compare
tdoublep
left a comment
tdoublep
left a comment
There was a problem hiding this comment.
Review by Claude Code, posted by me (@tdoublep).
Ran it locally: bucketer 53 passed, recorder 15 passed, and -k batched_decode 18 passed (that last one isn't in your test plan but it's the path being changed). Default buckets come out identical to before, so the default path is unaffected.
One correction for the description: warmup doesn't record any batched-decode variants at all. _record_one only traces _attn_fn, and SpyreAttnBucket has no num_seqs field, so variants() can't describe a decode variant. Those kernels still compile on first use in the serving path, before and after this change. The real win is fewer of those compiles plus a single source of truth — worth saying that instead, otherwise the next reader will assume decode is covered by warmup.
Rest is inline.
| self._attn_bucketer = SpyreAttnBucketer(vllm_config) | ||
|
|
||
| self._num_seqs_buckets: tuple[int, ...] = tuple(self._attn_bucketer.num_seqs_buckets) | ||
| self._num_blocks_buckets: tuple[int, ...] = tuple(self._attn_bucketer.num_blocks_buckets) |
There was a problem hiding this comment.
This is the line with the real consequence. The decode path now inherits the KV-derived block ladder, so its floor moves with SPYRE_ATTN_KV_BUCKETS: default stays 1,2,4,8,16,32, but 512,1024,2048 gives 8,16,32, and a single 2048 gives 32. The kernel does for i in range(num_blocks), so a 1-block decode would then do 32 gathers + matmuls instead of 1.
The KV ladder is coarse because the recorded set is a product of both axes — that reason doesn't apply here, since nothing records decode variants. Either keep this ladder independent, or call out the trade-off.
| # rather than constructing a second one that could drift. | ||
| self._attn_bucketer = SpyreAttnBucketer(vllm_config) | ||
|
|
||
| self._num_seqs_buckets: tuple[int, ...] = tuple(self._attn_bucketer.num_seqs_buckets) |
There was a problem hiding this comment.
These two copies exist only to satisfy _find_bucket's tuple type, and _find_bucket is now the same lookup as SpyreAttnBucketer._round_up, which line 672 already calls on the bucketer's lists. Dropping both attributes and _find_bucket, and calling the bucketer directly at 1037-1038, removes the last place these can drift.
| ) | ||
|
|
||
| # The batched decode kernel adds a sequence axis, so it needs a second ladder. | ||
| self._num_seqs_buckets: list[int] = list( |
There was a problem hiding this comment.
Buckets 1 and 2 can never be reached — line 1030 only takes this path when num_seqs >= _MIN_BATCHED_SEQS (4). start=_MIN_BATCHED_SEQS would drop them.
| assert b.num_seqs_buckets[-1] == 6 | ||
| assert b.num_seqs_buckets == [1, 2, 4, 6] | ||
|
|
||
| def test_num_blocks_buckets_follow_the_kv_buckets(self, monkeypatch): |
There was a problem hiding this comment.
This checks the bucketer, but not the thing the PR fixes. A test that fails before and passes after: with this same override plus SPYRE_BATCHED_DECODE=1, build a 4-block decode batch and assert padded_batch_blocks in bucketer.num_blocks_buckets — today it's 4, which isn't in [8,16,32]. TestBuilderAttnBucketer looks like the natural home.
2 participants
Description
The batched decode path built its own
num_seqsandnum_blocksladders inSpyreAttentionMetadataBuilder, whileSpyreAttnBucketerderived anum_blocksladder of its own from the KV buckets. The two only agree on the default geometric KV buckets — underSPYRE_ATTN_KV_BUCKETSthe builder dispatches onto low block counts that warmup never records, so an Inductor compile lands in the serving path.Both ladders now come from the bucketer, making the set
build()dispatches onto and the set warmup records the same by construction.Related Issues
Test Plan
uv run pytest tests/runtime/test_spyre_attn_bucketer.py -m "not upstream"— 53 passed, 5 skippeduv run pytest tests/attention/test_spyre_attn_recorder.py -m "not upstream"— 15 passedbash format.sh— ruff check/format pass (tyis red onmaintoo, in this environment)Checklist
bash format.sh)Signed-off-by:line (DCO compliance)