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

spyre_inference/v1/attention/backends/spyre_attn.py

@@ -96,25 +96,6 @@ def _record_block(name: str):
 _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)
@@ -654,21 +635,15 @@ 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)
         # 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)
+        self._num_blocks_buckets: tuple[int, ...] = tuple(self._attn_bucketer.num_blocks_buckets)
+
     def _get_zero_tile(self, aligned_query_len: int) -> torch.Tensor:
         """Return (or create) the shared all-zero mask tile for interior blocks.
 

spyre_inference/v1/attention/spyre_attn_bucketer.py

@@ -69,6 +69,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 +128,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
@@ -161,6 +172,11 @@ def __init__(self, vllm_config: VllmConfig) -> None:
             {(kv + block_size - 1) // block_size for kv in self._kv_buckets}
         )
 
+        # The batched decode kernel adds a sequence axis, so it needs a second ladder.
+        self._num_seqs_buckets: list[int] = list(
+            _powers_of_two_up_to(vllm_config.scheduler_config.max_num_seqs)
+        )
+
         logger.info(
             "SpyreAttnBucketer: %d kv buckets [%d..%d], %d query buckets [%d..%d], "
             "max num_blocks=%d",
@@ -185,6 +201,10 @@ 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)
 

tests/attention/test_spyre_attn_recorder.py

@@ -81,11 +81,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,23 @@ 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):
+        """A kv override moves the ladder the attention impl dispatches onto, so the
+        impl cannot land on a low block count that warmup never recorded."""
+        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):