Skip to content

refactor(attn): rename the bucketed decode kernel to batched decode #791

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tdoublep merged 1 commit into from
Sep 7, 2026
+67 −67
Merged

refactor(attn): rename the bucketed decode kernel to batched decode #791

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions spyre_inference/envs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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")),
Expand Down
90 changes: 45 additions & 45 deletions spyre_inference/v1/attention/backends/spyre_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@

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.
Expand Down Expand Up @@ -90,9 +90,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, ...]:
Expand Down Expand Up @@ -366,7 +366,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,
Expand All @@ -377,21 +377,21 @@ 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.

`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
Expand Down Expand Up @@ -455,7 +455,7 @@ def specialized_bucketed_decode_kernel(
return out
return attn

return specialized_bucketed_decode_kernel
return specialized_batched_decode_kernel


@dataclass
Expand Down Expand Up @@ -548,12 +548,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
Expand Down Expand Up @@ -623,7 +623,7 @@ 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 = (
Expand Down Expand Up @@ -1009,16 +1009,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
Expand All @@ -1028,13 +1028,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
Expand Down Expand Up @@ -1098,8 +1098,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,
Expand Down Expand Up @@ -1241,7 +1241,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
Expand Down Expand Up @@ -1307,10 +1307,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,
Expand All @@ -1319,12 +1319,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,
Expand All @@ -1337,16 +1337,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

Expand Down Expand Up @@ -1394,11 +1394,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
Expand Down Expand Up @@ -1606,7 +1606,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,
Expand All @@ -1622,8 +1622,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
Expand Down Expand Up @@ -1652,7 +1652,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(
Expand Down Expand Up @@ -1719,8 +1719,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
Expand Down
38 changes: 19 additions & 19 deletions tests/attention/test_spyre_attn.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand All @@ -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()
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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)
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
Loading