diff --git a/spyre_inference/v1/attention/backends/spyre_attn.py b/spyre_inference/v1/attention/backends/spyre_attn.py index 5b60242f1..5604d6cb6 100644 --- a/spyre_inference/v1/attention/backends/spyre_attn.py +++ b/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) @@ -657,21 +638,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. diff --git a/spyre_inference/v1/attention/spyre_attn_bucketer.py b/spyre_inference/v1/attention/spyre_attn_bucketer.py index 1124f738f..ce48fcd46 100644 --- a/spyre_inference/v1/attention/spyre_attn_bucketer.py +++ b/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) diff --git a/tests/attention/test_spyre_attn_recorder.py b/tests/attention/test_spyre_attn_recorder.py index 270584bbf..b7a0567f9 100644 --- a/tests/attention/test_spyre_attn_recorder.py +++ b/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) diff --git a/tests/runtime/test_spyre_attn_bucketer.py b/tests/runtime/test_spyre_attn_bucketer.py index 53b99b727..8d5917156 100644 --- a/tests/runtime/test_spyre_attn_bucketer.py +++ b/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):