diff --git a/spyre_inference/envs.py b/spyre_inference/envs.py index 618d24cf2..38bd4efe2 100644 --- a/spyre_inference/envs.py +++ b/spyre_inference/envs.py @@ -33,7 +33,7 @@ SPYRE_ATTN_RECORD: bool = True SPYRE_ATTN_KV_BUCKETS: str | None = None SPYRE_ATTN_QUERY_BUCKETS: str | None = None - SPYRE_BUCKETED_DECODE: bool = False + SPYRE_BATCHED_DECODE: bool = False SPYRE_NUM_CPUS: int = 0 SPYRE_UPDATE_THREAD_CONFIG: bool = True @@ -63,10 +63,10 @@ # 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"), - # When "1", enables the bucketed multi-sequence decode kernel. Off by default + # 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_BUCKETED_DECODE": lambda: bool(int(os.getenv("SPYRE_BUCKETED_DECODE", "0"))), + "SPYRE_BATCHED_DECODE": lambda: bool(int(os.getenv("SPYRE_BATCHED_DECODE", "0"))), # CPU budget used to size thread pools. "0" (default) auto-detects the budget # (cgroup CPU quota, then physical core count). "SPYRE_NUM_CPUS": lambda: int(os.getenv("SPYRE_NUM_CPUS", "0")), diff --git a/spyre_inference/v1/attention/backends/spyre_attn.py b/spyre_inference/v1/attention/backends/spyre_attn.py index 3191733d3..30dbc0309 100644 --- a/spyre_inference/v1/attention/backends/spyre_attn.py +++ b/spyre_inference/v1/attention/backends/spyre_attn.py @@ -44,11 +44,12 @@ from spyre_inference.v1.attention.spyre_attn_bucketer import ( SpyreAttnBucket, SpyreAttnBucketer, + SpyreBatchedAttnBucket, ) logger = init_logger(__name__) -# When set, wraps forward(), _online_softmax_attention() and the bucketed +# When set, wraps forward(), _online_softmax_attention() and the batched # decode K/V/mask gather blocks in torch.profiler.record_function spans for # kineto trace capture. Off by default: the spans are not free, so a profiled # run is not wall-clock comparable to a default one. @@ -90,9 +91,9 @@ def _record_block(name: str): INT32_ELEMS_PER_STICK = 32 -# Batches below this fall back to the per-seq loop: the bucketed matmul's +# 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_SEQS_BUCKET = 4 +_MIN_BATCHED_SEQS = 4 def _powers_of_two_up_to(n: int, start: int = 1) -> tuple[int, ...]: @@ -366,7 +367,7 @@ def specialized_paged_attn_kernel( return specialized_paged_attn_kernel -def _create_compilable_bucketed_decode_attn( +def _create_compilable_batched_decode_attn( num_seqs: int, num_blocks: int, num_kv_heads: int, @@ -377,7 +378,7 @@ def _create_compilable_bucketed_decode_attn( needs_gather: bool = True, store_out: bool = False, ): - """Bucketed decode kernel factory; gathers K/V and the query in-graph. + """Batched decode kernel factory; gathers K/V and the query in-graph. Gathers one block at a time; block_ids rows must stay stick-aligned, since a flat per-block slice does not compile. @@ -385,13 +386,13 @@ def _create_compilable_bucketed_decode_attn( `logits_soft_cap` and `needs_gather` are closure constants resolved at trace time, so each distinct value produces a different compiled graph. `logits_soft_cap` is fixed per ``SpyreAttentionImpl`` instance and therefore - not part of ``_get_bucketed_decode_kernel``'s cache key; `needs_gather` varies + not part of ``_get_batched_decode_kernel``'s cache key; `needs_gather` varies per step and is. """ num_heads = num_kv_heads * num_queries_per_kv - def specialized_bucketed_decode_kernel( + def specialized_batched_decode_kernel( query, query_row_ids, k_pages, v_pages, block_ids, mask_by_block, scale, out ): # Q=1 puts the sequences in rows 0..num_seqs-1; lanes past the batch are @@ -455,7 +456,7 @@ def specialized_bucketed_decode_kernel( return out return attn - return specialized_bucketed_decode_kernel + return specialized_batched_decode_kernel @dataclass @@ -548,12 +549,12 @@ class SpyreAttentionMetadata(AttentionMetadata): # Device mirror of attention_mask_tiles, filled once per step by forward(). attention_mask_tiles_device: list[list[torch.Tensor]] | None = None - # Bucketed-decode precomputes. None-valued when the batch is ineligible + # Batched-decode precomputes. None-valued when the batch is ineligible # (callers fall back to the per-seq loop). query_row_ids is int64 because # Spyre's index_copy_ requires int64. mask_by_block is pre-permuted for cheap # axis-0 slicing in the dispatch. - bucket_num_seqs: int | None = None - bucket_num_blocks: int | None = None + padded_num_seqs: int | None = None + padded_batch_blocks: int | None = None query_row_ids_cpu: torch.Tensor | None = None # [B_seqs] int64 query_row_ids_dev: torch.Tensor | None = None block_ids_padded_cpu: torch.Tensor | None = None # [B_blocks, padded B_seqs] int32 @@ -623,21 +624,19 @@ def __init__( static_ctx[name] for name in layer_names if name in static_ctx ) - # Buckets for the bucketed decode fast path. One compiled kernel + # 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) + # Both ladders come from the bucketer, so the set build() dispatches onto and + # the set warmup records are the same by construction. + 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_max_query_len: int) -> torch.Tensor: """Return (or create) the shared all-zero mask tile for interior blocks. @@ -1009,16 +1008,16 @@ def build( # False, so the traced write keeps one shape per bucket, not one per token count. self._slot_mapping.publish(slot_mapping) - # Bucketed-decode precomputes: only when Q=1 and num_seqs is within the + # Batched-decode precomputes: only when Q=1 and num_seqs is within the # buckets. None-valued fields signal fallback. Sliding-window batches are # eligible: the kernel reads a precomputed mask, so a window only shrinks # the active block set. - bucket_num_seqs = None - bucket_num_blocks = None + padded_num_seqs = None + padded_batch_blocks = None query_row_ids_cpu = None block_ids_padded_cpu = None mask_by_block_cpu = None - if max_query_len == 1 and num_seqs >= _MIN_SEQS_BUCKET: + if max_query_len == 1 and num_seqs >= _MIN_BATCHED_SEQS: # Real counts, not padded: this path has its own buckets, so an # inflated count would only push it onto a larger bucket for no # reason. Safe because padding only appends blocks. Under a window @@ -1028,13 +1027,13 @@ def build( b_seqs = _find_bucket(num_seqs, self._num_seqs_buckets) b_blocks = _find_bucket(max(blocks_per_seq), self._num_blocks_buckets) if b_seqs is not None and b_blocks is not None: - bucket_num_seqs = b_seqs - bucket_num_blocks = b_blocks + padded_num_seqs = b_seqs + padded_batch_blocks = b_blocks # int64 (not int32): Spyre's index_copy_ requires int64 indices. query_row_ids_cpu = torch.zeros(b_seqs, dtype=torch.int64) query_row_ids_cpu[:num_seqs] = query_start_loc[:num_seqs].to(torch.int64) - # Guards the identity scatter used by _run_bucketed_decode_dispatch. + # Guards the identity scatter used by _run_batched_decode_dispatch. assert query_row_ids_cpu[:num_seqs].tolist() == list(range(num_seqs)) # Rows padded to the stick width: a narrower inner dim emits a @@ -1098,8 +1097,8 @@ def build( page_index_tables_cpu=page_index_tables_cpu, aligned_max_query_len=aligned_max_query_len, padded_num_blocks=padded_num_blocks, - bucket_num_seqs=bucket_num_seqs, - bucket_num_blocks=bucket_num_blocks, + padded_num_seqs=padded_num_seqs, + padded_batch_blocks=padded_batch_blocks, query_row_ids_cpu=query_row_ids_cpu, block_ids_padded_cpu=block_ids_padded_cpu, mask_by_block_cpu=mask_by_block_cpu, @@ -1241,7 +1240,7 @@ def __init__( self._kv_slots: SpyrePagedKVCache | None = None - # Keyed by (bucket_num_seqs, bucket_num_blocks, needs_gather, store_out). + # Keyed by (padded_num_seqs, padded_batch_blocks, needs_gather, store_out). self._decode_fns: dict[tuple[int, int, bool, bool], object] = {} # Constant for the run, so the kernel's arguments never carry the model @@ -1307,10 +1306,10 @@ def _get_attn_fn(self, num_blocks: int, padded_query_len: int): ) return self._attn_fns[key] - def _get_bucketed_decode_kernel( + def _get_batched_decode_kernel( self, - bucket_num_seqs: int, - bucket_num_blocks: int, + padded_num_seqs: int, + padded_batch_blocks: int, block_size: int, needs_gather: bool, store_out: bool, @@ -1319,12 +1318,12 @@ def _get_bucketed_decode_kernel( # KV cache spec, logits_soft_cap by __init__), so they are passed to the # factory but not keyed on. needs_gather and store_out are closure # constants, so they are. - key = (bucket_num_seqs, bucket_num_blocks, needs_gather, store_out) + key = (padded_num_seqs, padded_batch_blocks, needs_gather, store_out) if key not in self._decode_fns: self._decode_fns[key] = _maybe_compile( - _create_compilable_bucketed_decode_attn( - bucket_num_seqs, - bucket_num_blocks, + _create_compilable_batched_decode_attn( + padded_num_seqs, + padded_batch_blocks, self.num_kv_heads, self.num_queries_per_kv, block_size, @@ -1337,16 +1336,16 @@ def _get_bucketed_decode_kernel( ) return self._decode_fns[key] - def _bucketed_decode_preconditions_met(self, attn_metadata: "SpyreAttentionMetadata") -> bool: - # Off by default: the bucketed matmul pads every sequence row up to the + def _batched_decode_preconditions_met(self, attn_metadata: "SpyreAttentionMetadata") -> bool: + # Off by default: the batched matmul pads every sequence row up to the # bucket width, and that overhead is uncharacterised at the smallest - # bucket (num_seqs == _MIN_SEQS_BUCKET), where there is no headroom. - # Set SPYRE_BUCKETED_DECODE=1 to restore the path. - if not envs.SPYRE_BUCKETED_DECODE: + # bucket (num_seqs == _MIN_BATCHED_SEQS), where there is no headroom. + # Set SPYRE_BATCHED_DECODE=1 to restore the path. + if not envs.SPYRE_BATCHED_DECODE: return False # Layer 0's builder gates on max_query_len and the bucket lattice; - # we add ALiBi, which the bucketed kernel doesn't implement. - if attn_metadata.bucket_num_seqs is None: + # we add ALiBi, which the batched kernel doesn't implement. + if attn_metadata.padded_num_seqs is None: return False return self.alibi_slopes is None @@ -1394,11 +1393,11 @@ def forward( # The KV write is not here: attn_layer.py traces it for the layers it splits, # and upstream's own unified_kv_cache_update op covers the rest. - # Mirror bucketed-decode precomputes to device once per step, only for - # layers whose impl can actually use the bucketed kernel (skips ALiBi + # Mirror batched-decode precomputes to device once per step, only for + # layers whose impl can actually use the batched kernel (skips ALiBi # and soft-cap layers). if ( - self._bucketed_decode_preconditions_met(attn_metadata) + self._batched_decode_preconditions_met(attn_metadata) and attn_metadata.query_row_ids_dev is None ): assert attn_metadata.query_row_ids_cpu is not None @@ -1447,25 +1446,33 @@ def record_graphs( k_pages, v_pages = kv_cache num_pages, block_size = k_pages.shape[0], k_pages.shape[1] variants = bucketer.variants() + # The batched kernel keys on (num_seqs, num_blocks) rather than the per-sequence + # tuple, so it needs its own enumeration. + batched = bucketer.batched_variants() if envs.SPYRE_BATCHED_DECODE else [] + total_variants = len(variants) + len(batched) t_start = time.time() # Belt-and-suspenders: platform._raise_dynamo_recompile_limits already # raises this globally, but bump it here too in case that hasn't run. prev_limit = torch._dynamo.config.accumulated_recompile_limit torch._dynamo.config.accumulated_recompile_limit = max( # ty: ignore[invalid-assignment] - prev_limit, 4 * len(variants) + 64 + prev_limit, 4 * total_variants + 64 ) - logger.info("Recording %d attention variants for layer...", len(variants)) + logger.info("Recording %d attention variants for layer...", total_variants) try: recorded = self._record_all(variants, k_pages, v_pages, num_pages, block_size, device) + if batched: + recorded += self._record_batched_all( + batched, k_pages, v_pages, num_pages, block_size, device + ) finally: torch._dynamo.config.accumulated_recompile_limit = prev_limit # ty: ignore[invalid-assignment] logger.info( "Recorded %d/%d attention variants in %.2fs.", recorded, - len(variants), + total_variants, time.time() - t_start, ) return recorded @@ -1509,6 +1516,78 @@ def _record_all( ) return recorded + def _record_batched_all( + self, + variants: "list[SpyreBatchedAttnBucket]", + k_pages: torch.Tensor, + v_pages: torch.Tensor, + num_pages: int, + block_size: int, + device: torch.device, + ) -> int: + """Compile every batched decode variant, mirroring _record_all's contract.""" + recorded = 0 + for bucket in variants: + if bucket.num_blocks > num_pages: + # The buckets are sized from max_model_len; a small KV allocation + # cannot host that many distinct pages to gather. + continue + if bucket.key in self._decode_fns: + continue + try: + self._record_batched_one(bucket, k_pages, v_pages, block_size, device) + except Exception: + self._decode_fns.pop(bucket.key, None) + logger.warning( + "Batched decode variant %s failed to record; it will compile on " + "first use instead.", + bucket.key, + exc_info=True, + ) + continue + recorded += 1 + return recorded + + def _record_batched_one( + self, + bucket: "SpyreBatchedAttnBucket", + k_pages: torch.Tensor, + v_pages: torch.Tensor, + block_size: int, + device: torch.device, + ) -> None: + """Trace one batched variant on dummy args matching the kernel's contract.""" + b_seqs, b_blocks = bucket.num_seqs, bucket.num_blocks + kernel = self._get_batched_decode_kernel( + b_seqs, b_blocks, block_size, bucket.needs_gather, bucket.store_out + ) + + # The very buffers attn_layer hands over, not fresh tensors of the same shape: + # Dynamo guards on the argument's shape, and since #789 that shape is the + # staging width, not b_seqs. Recording on anything else compiles a graph the + # serving path cannot reuse. + q_staging, out_staging = self._staging_buffers(device) + query = q_staging + row_ids = None + if bucket.needs_gather: + row_ids = convert(torch.zeros(b_seqs, dtype=torch.int64), device=device) + + block_ids = convert( + torch.zeros(b_blocks, _stick_aligned_len(b_seqs), dtype=torch.int32), + device=device, + ) + # All-zero additive mask: zero is the one choice that cannot leave a row fully + # masked, which would make tile_sum 0 and the result NaN. + mask_by_block = convert( + torch.zeros( + b_blocks, b_seqs * self.num_kv_heads, 1, block_size, dtype=self.model_dtype + ), + device=device, + ) + out = out_staging if bucket.store_out else None + + kernel(query, row_ids, k_pages, v_pages, block_ids, mask_by_block, self.scale, out) + def _record_one( self, bucket: "SpyreAttnBucket", @@ -1606,7 +1685,7 @@ def do_kv_cache_update( # kernel, so ordering the read after it covers the V write too. return k_slots - def _run_bucketed_decode_dispatch( + def _run_batched_decode_dispatch( self, query_dev: torch.Tensor, k_pages: torch.Tensor, @@ -1622,8 +1701,8 @@ def _run_bucketed_decode_dispatch( # Mod(d0, num_blocks) for a runtime .select); (4) result scatter is a # single contiguous copy_ at offset 0, valid because Q=1 forces # query_row_ids_cpu[:num_seqs] == range(num_seqs) (asserted in builder). - b_seqs = attn_metadata.bucket_num_seqs - b_blocks = attn_metadata.bucket_num_blocks + b_seqs = attn_metadata.padded_num_seqs + b_blocks = attn_metadata.padded_batch_blocks num_seqs = attn_metadata.num_seqs num_heads = self.num_heads head_size = self.head_size @@ -1639,9 +1718,11 @@ def _run_bucketed_decode_dispatch( # (torch-spyre#3770), so each block needed a .clone(). block_ids = attn_metadata.block_ids_padded_dev - # Short of b_seqs rows only when the runner's compile bucket is tighter than - # the power-of-two seq bucket; the kernel slices a prefix otherwise. - needs_gather = query_dev.shape[0] < b_seqs + # Padding rows are present exactly when the real batch is short of the + # bucket; the kernel slices a prefix otherwise. Derived from the seq counts, + # not the query width: since #789 the width is a constant staging buffer + # (max_num_batched_tokens + 1) and carries no information about the batch. + needs_gather = num_seqs < b_seqs # The kernel writes b_seqs rows, so output must have them. Re-checked per call: # vLLM hands out a fresh buffer per layer. store_out = ( @@ -1652,7 +1733,7 @@ def _run_bucketed_decode_dispatch( and output.storage_offset() == 0 and output.is_contiguous() ) - kernel = self._get_bucketed_decode_kernel( + kernel = self._get_batched_decode_kernel( b_seqs, b_blocks, block_size, needs_gather, store_out ) result = kernel( @@ -1719,8 +1800,8 @@ def _online_softmax_attention( ) assert page_index_tables is not None, "page_index_tables must be mirrored by forward()" - if self._bucketed_decode_preconditions_met(attn_metadata): - self._run_bucketed_decode_dispatch(query_dev, k_pages, v_pages, attn_metadata, output) + if self._batched_decode_preconditions_met(attn_metadata): + self._run_batched_decode_dispatch(query_dev, k_pages, v_pages, attn_metadata, output) return output # Mirrors the batch layout row for row, so the absolute query_start_loc diff --git a/spyre_inference/v1/attention/spyre_attn_bucketer.py b/spyre_inference/v1/attention/spyre_attn_bucketer.py index 58848751b..062d093e4 100644 --- a/spyre_inference/v1/attention/spyre_attn_bucketer.py +++ b/spyre_inference/v1/attention/spyre_attn_bucketer.py @@ -103,6 +103,24 @@ def _resolve_buckets( return buckets +@dataclass(frozen=True) +class SpyreBatchedAttnBucket: + """One recordable batched decode kernel variant. + + Fields mirror ``SpyreAttentionImpl._get_batched_decode_kernel``'s cache key exactly, + so a recorded bucket and a runtime dispatch are the same tuple. + """ + + num_seqs: int + num_blocks: int + needs_gather: bool + store_out: bool + + @property + def key(self) -> tuple[int, int, bool, bool]: + return (self.num_seqs, self.num_blocks, self.needs_gather, self.store_out) + + class SpyreAttnBucketer: """Enumerates the attention variants to record, and rounds lengths onto them. @@ -161,6 +179,11 @@ def __init__(self, vllm_config: VllmConfig) -> None: # 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. + # The batched decode kernel adds a sequence axis, so it needs a second ladder. + # Owned here so one place defines every bucket either kernel can dispatch onto. + max_num_seqs = vllm_config.scheduler_config.max_num_seqs + self._num_seqs_buckets: list[int] = list(_powers_of_two_up_to(max_num_seqs)) + self._num_blocks_buckets: list[int] = sorted( {(kv + block_size - 1) // block_size for kv in self._kv_buckets} ) @@ -189,6 +212,44 @@ 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 batched_variants(self) -> list[SpyreBatchedAttnBucket]: + """Every batched decode variant worth recording, largest first. + + The kernel specialises on (num_seqs, num_blocks) and on two flags. + ``needs_gather`` is True when the query buffer holds fewer rows than the seqs + bucket. At the smallest bucket that cannot happen -- ``build()`` only enters at + ``_MIN_BATCHED_SEQS`` and that rounds to itself -- so it is enumerated only above. + ``store_out`` additionally needs a fused-store-eligible output buffer, and the + dispatch only takes it when ``needs_gather`` is False. + """ + out: list[SpyreBatchedAttnBucket] = [] + # build() only takes this path at or above _MIN_BATCHED_SEQS, so smaller seqs + # buckets are unreachable and recording them would compile dead variants. + from spyre_inference.v1.attention.backends.spyre_attn import _MIN_BATCHED_SEQS + + reachable = [n for n in self._num_seqs_buckets if n >= _MIN_BATCHED_SEQS] + for num_seqs in sorted(reachable, reverse=True): + for num_blocks in sorted(self._num_blocks_buckets, reverse=True): + # At the smallest bucket num_seqs == b_seqs always, so no gather. + gathers = (False,) if num_seqs == _MIN_BATCHED_SEQS else (False, True) + for needs_gather in gathers: + # store_out is only reachable without a gather (dispatch requires it). + store_outs = (False,) if needs_gather else (True, False) + for store_out in store_outs: + out.append( + SpyreBatchedAttnBucket( + num_seqs=num_seqs, + num_blocks=num_blocks, + needs_gather=needs_gather, + store_out=store_out, + ) + ) + return out + 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.py b/tests/attention/test_spyre_attn.py index 387dfd20c..47772fd61 100644 --- a/tests/attention/test_spyre_attn.py +++ b/tests/attention/test_spyre_attn.py @@ -30,7 +30,7 @@ SpyreAttentionMetadataBuilder, SpyrePagedKVCache, _build_query_row_tables, - _create_compilable_bucketed_decode_attn, + _create_compilable_batched_decode_attn, _mirror_mask_tiles, _stick_aligned_len, ) @@ -40,16 +40,16 @@ @pytest.fixture() -def enable_bucketed_decode(monkeypatch): - """Enable the bucketed decode kernel for tests that exercise it. +def enable_batched_decode(monkeypatch): + """Enable the batched decode kernel for tests that exercise it. - The path ships gated off (``SPYRE_BUCKETED_DECODE``, default "0") pending + The path ships gated off (``SPYRE_BATCHED_DECODE``, default "0") pending performance characterisation at the smallest bucket. Without this fixture the - bucketed tests would silently fall back to the per-seq loop and pass while + batched tests would silently fall back to the per-seq loop and pass while testing nothing. The autouse cache-clearing fixture in ``tests/conftest.py`` makes the monkeypatched value visible to ``envs``. """ - monkeypatch.setenv("SPYRE_BUCKETED_DECODE", "1") + monkeypatch.setenv("SPYRE_BATCHED_DECODE", "1") @pytest.fixture() @@ -1715,14 +1715,14 @@ def test_install_patches_layers_not_the_attention_class(): ), ], ) -def test_spyre_attn_bucketed_decode_correctness( +def test_spyre_attn_batched_decode_correctness( default_vllm_config, - enable_bucketed_decode, + enable_batched_decode, seq_lens: list[tuple[int, int]], configure_compilation: str, configure_device: str, ) -> None: - """Bucketed decode fast path: bit-exact vs the per-seq reference.""" + """Batched decode fast path: bit-exact vs the per-seq reference.""" _run_spyre_attn_test( seq_lens=seq_lens, block_size=128, @@ -1770,15 +1770,15 @@ def test_spyre_attn_bucketed_decode_correctness( ], ) @pytest.mark.parametrize("soft_cap", [pytest.param(50.0, id="soft_cap(50)")]) -def test_spyre_attn_bucketed_decode_soft_cap( +def test_spyre_attn_batched_decode_soft_cap( default_vllm_config, - enable_bucketed_decode, + enable_batched_decode, seq_lens: list[tuple[int, int]], soft_cap: float, configure_compilation: str, configure_device: str, ) -> None: - """Bucketed decode with logits soft-cap, vs the per-seq reference.""" + """Batched decode with logits soft-cap, vs the per-seq reference.""" _run_spyre_attn_test( seq_lens=seq_lens, block_size=128, @@ -1789,7 +1789,7 @@ def test_spyre_attn_bucketed_decode_soft_cap( ) -def test_bucketed_decode_soft_cap_changes_the_kernel() -> None: +def test_batched_decode_soft_cap_changes_the_kernel() -> None: """The capped kernel must actually clamp, not silently ignore the cap.""" torch.set_default_device("cpu") set_random_seed(0) @@ -1798,7 +1798,7 @@ def test_bucketed_decode_soft_cap_changes_the_kernel() -> None: lead = num_seqs * num_kv_heads def build(cap: float): - return _create_compilable_bucketed_decode_attn( + return _create_compilable_batched_decode_attn( num_seqs=num_seqs, num_blocks=num_blocks, num_kv_heads=num_kv_heads, @@ -1859,9 +1859,9 @@ def build(cap: float): ), ], ) -def test_spyre_attn_bucketed_decode_fallback( +def test_spyre_attn_batched_decode_fallback( default_vllm_config, - enable_bucketed_decode, + enable_batched_decode, seq_lens: list[tuple[int, int]], configure_compilation: str, configure_device: str, @@ -1902,15 +1902,15 @@ def test_spyre_attn_bucketed_decode_fallback( pytest.param([(1, 256)] * 8, 4096, id="window_covers_all(N=8)"), ], ) -def test_spyre_attn_bucketed_decode_sliding_window( +def test_spyre_attn_batched_decode_sliding_window( default_vllm_config, - enable_bucketed_decode, + enable_batched_decode, seq_lens: list[tuple[int, int]], sliding_window: int, configure_compilation: str, configure_device: str, ) -> None: - """Bucketed decode with a sliding window: matches the per-seq reference.""" + """Batched decode with a sliding window: matches the per-seq reference.""" _run_spyre_attn_test( seq_lens=seq_lens, block_size=128, diff --git a/tests/attention/test_spyre_attn_recorder.py b/tests/attention/test_spyre_attn_recorder.py index 1997bbb39..7c3594786 100644 --- a/tests/attention/test_spyre_attn_recorder.py +++ b/tests/attention/test_spyre_attn_recorder.py @@ -66,11 +66,14 @@ 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 + # MagicMock returns a Mock for anything unset, which the bucket arithmetic + # then compares against an int; every field it reads has to be real. + 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 9848d8f48..02636a0fd 100644 --- a/tests/runtime/test_spyre_attn_bucketer.py +++ b/tests/runtime/test_spyre_attn_bucketer.py @@ -29,13 +29,20 @@ ) BLOCK_SIZE = 64 +MAX_NUM_SEQS = 8 -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=MAX_NUM_SEQS, +): 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 @@ -185,6 +192,51 @@ def test_every_rounded_size_lands_on_a_recorded_variant(self, bucketer, kv_len, sizes = {(v.num_blocks, v.padded_query_len) for v in bucketer.variants()} assert (num_blocks, padded_query_len) in sizes + @pytest.mark.parametrize("num_seqs", [4, 5, 7, 8]) + @pytest.mark.parametrize("kv_len", [64, 300, 1025, 2048]) + def test_every_batched_size_lands_on_a_recorded_variant(self, bucketer, num_seqs, kv_len): + """Same guarantee for the batched kernel: no runtime batch may miss the cache. + + Drives the two lookups the batched dispatch uses -- _round_up onto + num_seqs_buckets and onto num_blocks_buckets. + """ + b_seqs = bucketer._round_up(num_seqs, bucketer.num_seqs_buckets) + b_blocks = bucketer._round_up( + (kv_len + BLOCK_SIZE - 1) // BLOCK_SIZE, bucketer.num_blocks_buckets + ) + assert b_seqs is not None and b_blocks is not None + sizes = {(v.num_seqs, v.num_blocks) for v in bucketer.batched_variants()} + assert (b_seqs, b_blocks) in sizes + + def test_batched_variants_cover_both_flag_states(self, bucketer): + """needs_gather varies per step, so both values must be recorded at every size. + + store_out additionally needs a fused-store-eligible output buffer, and the + dispatch only takes it when needs_gather is False. + """ + by_size: dict[tuple[int, int], set[tuple[bool, bool]]] = {} + for v in bucketer.batched_variants(): + by_size.setdefault((v.num_seqs, v.num_blocks), set()).add((v.needs_gather, v.store_out)) + from spyre_inference.v1.attention.backends.spyre_attn import _MIN_BATCHED_SEQS + + assert by_size + for (num_seqs, _), flags in by_size.items(): + assert (False, False) in flags + assert (False, True) in flags + # A gather needs a query buffer narrower than the bucket, which cannot happen + # at the smallest one: build() only enters there and it rounds to itself. + assert ((True, False) in flags) == (num_seqs > _MIN_BATCHED_SEQS) + # store_out with a gather is unreachable: the dispatch requires not needs_gather. + assert (True, True) not in flags + + def test_batched_variants_skip_unreachable_seqs_buckets(self, bucketer): + """build() only takes the batched path at or above _MIN_BATCHED_SEQS.""" + from spyre_inference.v1.attention.backends.spyre_attn import _MIN_BATCHED_SEQS + + assert {v.num_seqs for v in bucketer.batched_variants()} == { + n for n in bucketer.num_seqs_buckets if n >= _MIN_BATCHED_SEQS + } + def test_count_stays_tractable_at_long_context(self): """Dense buckets here would be tens of thousands of Inductor compiles.""" b = SpyreAttnBucketer(make_config(32768, 2048))