refactor(attn): let the bucketer own the batched decode buckets

spyre_inference/envs.py

@@ -33,6 +33,7 @@
     SPYRE_ATTN_RECORD: bool = True
     SPYRE_ATTN_KV_BUCKETS: str | None = None
     SPYRE_ATTN_QUERY_BUCKETS: str | None = None
+    SPYRE_ATTN_NUM_SEQS_BUCKETS: str | None = None
     SPYRE_BATCHED_DECODE: bool = False
     SPYRE_KERNEL_CACHE: bool = False
     SPYRE_NUM_CPUS: int = 0
@@ -64,6 +65,9 @@
     # Comma-separated query_len buckets to record, unset uses the default buckets
     # [1] + multiples of min(512, max_num_batched_tokens) up to max_num_batched_tokens.
     "SPYRE_ATTN_QUERY_BUCKETS": lambda: os.getenv("SPYRE_ATTN_QUERY_BUCKETS"),
+    # Comma-separated num_seqs buckets for the batched decode kernel, unset uses the
+    # default buckets of powers of two from 4 up to max_num_seqs.
+    "SPYRE_ATTN_NUM_SEQS_BUCKETS": lambda: os.getenv("SPYRE_ATTN_NUM_SEQS_BUCKETS"),
     # When "1", enables the batched multi-sequence decode kernel. Off by default
     # pending performance characterisation at small batch sizes (num_seqs <= 4).
     # Re-enable to measure the path or to restore it after calibration.

spyre_inference/v1/attention/backends/spyre_attn.py

@@ -43,6 +43,7 @@
 from spyre_inference.custom_ops.utils import convert
 from spyre_inference.v1.attention import attn_layer
 from spyre_inference.v1.attention.spyre_attn_bucketer import (
+    _MIN_BATCHED_SEQS,
     SpyreAttnBucket,
     SpyreAttnBucketer,
 )
@@ -91,30 +92,6 @@ def _record_block(name: str):
 INT32_ELEMS_PER_STICK = 32
 
 
-# Batches below this fall back to the per-seq loop: the batched matmul's
-# padded-row overhead exceeds the per-seq cost at small N.
-_MIN_BATCHED_SEQS = 4
-
-
-def _powers_of_two_up_to(n: int, start: int = 1) -> tuple[int, ...]:
-    """Powers of 2 in [start, n], plus n itself if it is not already a power of 2.
-
-    ``start`` is rounded up to a power of 2 first, keeping a pure doubling
-    sequence. A ``start`` above ``n`` yields just ``(n,)``.
-    """
-    if n < 1:
-        return ()
-    v = 1
-    while v < start:
-        v *= 2
-    result = []
-    while v < n:
-        result.append(v)
-        v *= 2
-    result.append(n)
-    return tuple(result)
-
-
 def _find_bucket(n: int, buckets: tuple[int, ...]) -> int | None:
     """Smallest bucket >= n, or None when n exceeds the top bucket."""
     idx = bisect.bisect_left(buckets, n)
@@ -657,15 +634,6 @@ def __init__(
             static_ctx[name] for name in layer_names if name in static_ctx
         )
 
-        # Buckets for the batched decode fast path. One compiled kernel
-        # per bucket. TODO: expose as engine args if configurability is needed.
-        max_num_seqs = vllm_config.scheduler_config.max_num_seqs
-        max_num_blocks_per_seq = (
-            model_config.max_model_len + self.block_size - 1
-        ) // self.block_size
-        self._num_seqs_buckets: tuple[int, ...] = _powers_of_two_up_to(max_num_seqs)
-        self._num_blocks_buckets: tuple[int, ...] = _powers_of_two_up_to(max_num_blocks_per_seq)
-
         # Owned here, not by the recorder, so a bucket build() can emit is
         # always a bucket that was compiled: the warmup recorder reads this
         # same instance back (spyre_model_runner._record_attention_graphs)
@@ -1062,8 +1030,12 @@ def build(
             # real_num_blocks is empty and the tiles are the unpadded active
             # blocks, so num_active is already the real count.
             blocks_per_seq = real_num_blocks if active_block_indices is None else num_active
-            b_seqs = _find_bucket(num_seqs, self._num_seqs_buckets)
-            b_blocks = _find_bucket(max(blocks_per_seq), self._num_blocks_buckets)
+
+            b_seqs = SpyreAttnBucketer._round_up(num_seqs, self._attn_bucketer._num_seqs_buckets)
+            b_blocks = SpyreAttnBucketer._round_up(
+                max(blocks_per_seq), self._attn_bucketer._num_blocks_buckets
+            )
+
             if b_seqs is not None and b_blocks is not None:
                 padded_num_seqs = b_seqs
                 padded_batch_blocks = b_blocks

spyre_inference/v1/attention/spyre_attn_bucketer.py

@@ -46,6 +46,11 @@
 # multiple of this.
 _DEFAULT_QUERY_BUCKET_STEP = 512
 
+# Batches below this fall back to the per-seq loop: the batched matmul's
+# padded-row overhead exceeds the per-seq cost at small N. So the num_seqs ladder
+# starts here -- smaller batches never dispatch to a batched variant.
+_MIN_BATCHED_SEQS = 4
+
 
 @dataclass(frozen=True)
 class SpyreAttnBucket:
@@ -69,6 +74,21 @@ def _parse_buckets(raw: str | None) -> list[int] | None:
     return values
 
 
+def _powers_of_two_up_to(n: int, start: int = 1) -> tuple[int, ...]:
+    """Powers of 2 in [start, n] (start rounded up to a power of 2), plus n itself."""
+    if n < 1:
+        return ()
+    v = 1
+    while v < start:
+        v *= 2
+    result = []
+    while v < n:
+        result.append(v)
+        v *= 2
+    result.append(n)
+    return tuple(result)
+
+
 def _resolve_buckets(
     raw: str | None, limit: int, name: str, default: Callable[[], list[int]]
 ) -> list[int]:
@@ -113,10 +133,6 @@ def __init__(self, vllm_config: VllmConfig) -> None:
         max_model_len = vllm_config.model_config.max_model_len
         max_batched = vllm_config.scheduler_config.max_num_batched_tokens
 
-        # Imported at call time, not module scope: spyre_attn imports this
-        # module, so a top-level import back into it would be circular.
-        from spyre_inference.v1.attention.backends.spyre_attn import _powers_of_two_up_to
-
         if block_size & (block_size - 1):
             # Not fatal: _powers_of_two_up_to rounds the start up to a power of
             # two, just coarser at the bottom. Reachable because the platform
@@ -128,24 +144,19 @@ def __init__(self, vllm_config: VllmConfig) -> None:
                 block_size,
             )
 
-        # Default: powers of two from block_size up to max_model_len. The
-        # recorded set is a product of both axes, so a bucket per KV token at a
-        # 32k context would be tens of thousands of variants; doubling keeps it
-        # affordable, with each bucket's extra padding absorbed by the mask.
-        # Starting at block_size rather than 1 skips buckets that would dedupe
-        # away anyway, since num_blocks = ceil(kv / block_size).
-        self._kv_buckets: list[int] = _resolve_buckets(
-            envs.SPYRE_ATTN_KV_BUCKETS,
-            max_model_len,
-            "SPYRE_ATTN_KV_BUCKETS",
-            lambda: list(_powers_of_two_up_to(max_model_len, start=block_size)),
+        # Default: powers of two from _MIN_BATCHED_SEQS up to max_num_seqs, the
+        # batch sizes the batched decode kernel can be asked for.
+        max_num_seqs = vllm_config.scheduler_config.max_num_seqs
+        self._num_seqs_buckets: list[int] = _resolve_buckets(
+            envs.SPYRE_ATTN_NUM_SEQS_BUCKETS,
+            max_num_seqs,
+            "SPYRE_ATTN_NUM_SEQS_BUCKETS",
+            lambda: list(_powers_of_two_up_to(max_num_seqs, start=_MIN_BATCHED_SEQS)),
         )
 
         # Default: [1] (the decode-only batch, exempt from query padding by
-        # build()) then multiples of a step up to max_num_batched_tokens. Coarse
-        # bucketing: a prefill pays padding up to the next bucket, which the mask
-        # discards. The step is capped at 512 so a large max_num_batched_tokens
-        # doesn't make the one non-decode bucket enormous.
+        # build()) then multiples of a step up to max_num_batched_tokens, the
+        # query lengths a prefill pads up to.
         step = min(_DEFAULT_QUERY_BUCKET_STEP, max_batched)
         self._query_buckets: list[int] = _resolve_buckets(
             envs.SPYRE_ATTN_QUERY_BUCKETS,
@@ -154,6 +165,16 @@ def __init__(self, vllm_config: VllmConfig) -> None:
             lambda: sorted({1, *range(step, max_batched + 1, step), max_batched}),
         )
 
+        # Default: powers of two from block_size up to max_model_len. Geometric
+        # because the recorded set is a product of both axes; the extra padding
+        # each bucket costs is absorbed by the mask.
+        self._kv_buckets: list[int] = _resolve_buckets(
+            envs.SPYRE_ATTN_KV_BUCKETS,
+            max_model_len,
+            "SPYRE_ATTN_KV_BUCKETS",
+            lambda: list(_powers_of_two_up_to(max_model_len, start=block_size)),
+        )
+
         # num_blocks is what the kernel specializes on. Derived from the kv
         # buckets, one block count per kv bucket, rather than enumerating every
         # integer up to max_model_len / block_size.
@@ -185,12 +206,19 @@ def query_buckets(self) -> list[int]:
     def num_blocks_buckets(self) -> list[int]:
         return self._num_blocks_buckets
 
+    @property
+    def num_seqs_buckets(self) -> list[int]:
+        return self._num_seqs_buckets
+
     def find_kv_bucket(self, kv_len: int) -> int | None:
         return self._round_up(kv_len, self._kv_buckets)
 
     def find_query_bucket(self, query_len: int) -> int | None:
         return self._round_up(query_len, self._query_buckets)
 
+    def find_sequence_bucket(self, num_seqs: int) -> int | None:
+        return self._round_up(num_seqs, self._num_seqs_buckets)
+
     @staticmethod
     def _round_up(n: int, buckets: list[int]) -> int | None:
         idx = bisect.bisect_left(buckets, n)

tests/attention/test_spyre_attn_recorder.py

@@ -84,11 +84,12 @@ def kv_cache():
     )
 
 
-def make_bucketer(max_model_len=256, max_num_batched_tokens=64):
+def make_bucketer(max_model_len=256, max_num_batched_tokens=64, max_num_seqs=8):
     config = MagicMock()
     config.cache_config.block_size = BLOCK_SIZE
     config.model_config.max_model_len = max_model_len
     config.scheduler_config.max_num_batched_tokens = max_num_batched_tokens
+    config.scheduler_config.max_num_seqs = max_num_seqs
     return SpyreAttnBucketer(config)
 
 

tests/runtime/test_spyre_attn_bucketer.py

@@ -21,20 +21,23 @@
 import pytest
 
 from spyre_inference import envs
-from spyre_inference.v1.attention.backends.spyre_attn import _powers_of_two_up_to
 from spyre_inference.v1.attention.spyre_attn_bucketer import (
     SpyreAttnBucketer,
     _parse_buckets,
+    _powers_of_two_up_to,
 )
 
 BLOCK_SIZE = 64
 
 
-def make_config(max_model_len=2048, max_num_batched_tokens=512, block_size=BLOCK_SIZE):
+def make_config(
+    max_model_len=2048, max_num_batched_tokens=512, block_size=BLOCK_SIZE, max_num_seqs=8
+):
     config = MagicMock()
     config.cache_config.block_size = block_size
     config.model_config.max_model_len = max_model_len
     config.scheduler_config.max_num_batched_tokens = max_num_batched_tokens
+    config.scheduler_config.max_num_seqs = max_num_seqs
     return config
 
 
@@ -189,6 +192,21 @@ def test_count_stays_tractable_at_long_context(self):
         b = SpyreAttnBucketer(make_config(32768, 2048))
         assert len(b.variants()) < 500
 
+    def test_num_seqs_buckets_are_powers_of_two_to_max_num_seqs(self):
+        b = SpyreAttnBucketer(make_config(max_num_seqs=8))
+        assert b.num_seqs_buckets == [1, 2, 4, 8]
+
+    def test_num_seqs_buckets_top_out_at_max_num_seqs(self):
+        b = SpyreAttnBucketer(make_config(max_num_seqs=6))
+        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):
+        monkeypatch.setenv("SPYRE_ATTN_KV_BUCKETS", "512,1024,2048")
+        envs.clear_env_cache()
+        b = SpyreAttnBucketer(make_config())
+        assert b.num_blocks_buckets == [8, 16, 32]
+
 
 class TestEnvOverride:
     def test_kv_buckets_override(self, monkeypatch):
@@ -307,3 +325,29 @@ def test_skips_builders_without_a_bucketer(self):
         bucketer = SpyreAttnBucketer(make_config())
         runner = self._runner([None, bucketer])
         assert runner._resolve_builder_attn_bucketer() is bucketer
+
+    def test_batched_decode_dispatches_onto_a_recorded_block_count(
+        self, monkeypatch, default_vllm_config
+    ):
+        """The regression this guards: ``build()`` and warmup must agree."""
+        from tests.attention.test_spyre_attn import _padded_mask_metadata
+
+        monkeypatch.setenv("SPYRE_ATTN_KV_BUCKETS", "512,1024,2048")
+        monkeypatch.setenv("SPYRE_BATCHED_DECODE", "1")
+        envs.clear_env_cache()
+
+        from vllm.config import get_current_vllm_config
+
+        # block_size pinned to match what _padded_mask_metadata builds with, so
+        # the override resolves onto the same block counts build() produces.
+        vllm_config = get_current_vllm_config()
+        vllm_config.cache_config.block_size = BLOCK_SIZE
+        bucketer = SpyreAttnBucketer(vllm_config)
+
+        # 4 blocks of real KV, and enough sequences to clear _MIN_BATCHED_SEQS.
+        metadata = _padded_mask_metadata(
+            [(1, 4 * BLOCK_SIZE)] * 4, max_num_blocks=bucketer.num_blocks_buckets[-1]
+        )
+
+        assert metadata.padded_batch_blocks in bucketer.num_blocks_buckets
+        assert metadata.padded_num_seqs in bucketer.num_seqs_buckets