From e2f9f733764ebcbb902811f38bba66780a48f503 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Wed, 28 Oct 2026 01:08:59 +0800 Subject: [PATCH 1/8] sparse decode attention: per-token reference lifecycle for async self-speculation Implements the asynchronous self-speculation cell of the sparse decode attention column (RFC #156 section 8): the lifecycle contract for per-token reference state across outstanding rounds. - attention/sparse.py: the RFC 3.2 arithmetic contract and the per-row SparseDecodePlan named view (modes/slots/retention) that travels from the runner to the attention layer on the execution metadata. - worker/attention_reference.py: AttentionReferenceStore, the runner-owned per-token slot pool. One slot per live token, empty -> source-ready -> reusable -> retired lifecycle, incarnation per (request, row) so a corrected token that reuses a position rejects stale plans, and acceptance-driven invalidation computed on the device without a host readback. Slots return to the pool only after the final reader completes (ticket retire). - worker/async_speculative.py: the async speculative runner reserves one lease per verification row per round, publishes CAPTURE plans at depth 1 (draft phase for draft depth >=2, verify phase for draft depth 1) and CONSUME plans at deeper loops, retires rejected candidates and their dependent descendants from the GPU acceptance counts, and releases slots on ticket retirement. Outstanding rounds hold disjoint slots. - config/entrypoints/engine: the sparse_decode_attention group (enabled, key_retention, key_bucket), wired through LLM/CLI/StageEngine with explicit capability checks (async self-speculation, block mode) and reference-state memory accounted into the KV reserve. The attention kernel consumption of the plan (dense source statistics and the sparse correction) is the basic-implementation cell owned by @Levius-Fubuki; the store publishes the slot views it will borrow. --- tests/helpers.py | 3 +- tests/test_sparse_decode_attention.py | 395 +++++++++++++++++++++++++ vllm_rlt/__init__.py | 2 + vllm_rlt/attention/sparse.py | 195 ++++++++++++ vllm_rlt/config.py | 25 ++ vllm_rlt/engine/stage_engine.py | 34 +++ vllm_rlt/entrypoints/llm.py | 2 + vllm_rlt/entrypoints/runtime_args.py | 25 +- vllm_rlt/worker/async_speculative.py | 151 ++++++++-- vllm_rlt/worker/attention_reference.py | 307 +++++++++++++++++++ vllm_rlt/worker/metadata.py | 8 + 11 files changed, 1128 insertions(+), 19 deletions(-) create mode 100644 tests/test_sparse_decode_attention.py create mode 100644 vllm_rlt/attention/sparse.py create mode 100644 vllm_rlt/worker/attention_reference.py diff --git a/tests/helpers.py b/tests/helpers.py index cf4a4236..ba4f3064 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -146,7 +146,7 @@ def prefill_submitted(self, request, **submission): def make_speculative_engine( - model, *, asynchronous=True, k=3, budget=16, prefix=False, multi_stream=True + model, *, asynchronous=True, k=3, budget=16, prefix=False, multi_stream=True, sparse=None ): return StageEngine( model, @@ -157,6 +157,7 @@ def make_speculative_engine( ), execution_config=ExecutionConfig(async_scheduling=asynchronous, multi_stream=multi_stream), speculative_config=SpeculativeConfig(k), + sparse_decode_attention_config=sparse, ) diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py new file mode 100644 index 00000000..6c77289c --- /dev/null +++ b/tests/test_sparse_decode_attention.py @@ -0,0 +1,395 @@ +"""Sparse decode attention: reference lifecycle and plan contracts (RFC #156). + +Covers the asynchronous self-speculation cell of the sparse decode attention +column: one reference slot per live token, plan modes for the draft depths, +GPU acceptance driving retirement of rejected candidates and their dependent +descendants, same-position replacement through incarnations, delayed +completion and reference-slot reuse. +""" + +import pytest +import torch + +from tests.helpers import bootstrap_speculation, make_speculative_engine, tiny_ouro_model +from vllm_rlt import SparseDecodeAttentionConfig +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + DENSE, + REFERENCE_DEPTH, + SparseDecodePlan, + reference_mode, + reference_statistics, + retained_count, + sparse_correction, +) +from vllm_rlt.worker.attention_reference import AttentionReferenceStore +from vllm_rlt.worker.metadata import MetadataBuilder + +tiny = tiny_ouro_model +bootstrap = bootstrap_speculation + + +def drain(engine): + outputs = {} + for _ in range(1000): + if not engine.has_unfinished_requests(): + return outputs + for output in engine.step(): + outputs[output.request_id] = output + pytest.fail("engine failed to drain") + + +def make_store(**overrides): + values = dict(slots=6, layers=2, heads=2, head_dim=4, key_bucket=8, device="cpu") + values.update(overrides) + return AttentionReferenceStore(**values) + + +def fill_payload(store, slots, value=1.0): + for slot in slots: + store.mass[slot] = value + store.output[slot] = value + store.contribution[slot] = value + store.selected[slot] = 1 + store.counts[slot] = 1 + + +# --------------------------------------------------------------------------- +# Plan and arithmetic contracts (RFC #156 sections 3.2 and 6.3) +# --------------------------------------------------------------------------- + + +def test_reference_mode_follows_reference_depth(): + # Loop 2 (depth index 1) captures; later loops consume; loop 1 stays dense. + assert reference_mode(0) == DENSE + assert reference_mode(REFERENCE_DEPTH) == CAPTURE + assert reference_mode(2) == CONSUME + assert reference_mode(3) == CONSUME + + +def test_plan_groups_rows_by_mode(): + plan = SparseDecodePlan( + modes=(DENSE, CAPTURE, CAPTURE, CONSUME), + slots=(-1, 3, 4, 5), + incarnations=(0, 1, 1, 1), + key_retention=0.5, + key_bucket=None, + ) + assert plan.groups() == {DENSE: [0], CAPTURE: [1, 2], CONSUME: [3]} + assert not plan.dense() + assert plan.reference_depth == REFERENCE_DEPTH + + +def test_retained_count_rounds_and_caps(): + assert retained_count(3, 0.667, None) == 2 + assert retained_count(4, 1.0, None) == 4 + assert retained_count(4, 1.0, 2) == 2 + assert retained_count(0, 0.5, None) == 0 + with pytest.raises(ValueError, match="key_retention"): + retained_count(4, 1.5, None) + with pytest.raises(ValueError, match="key_bucket"): + retained_count(4, 0.5, 0) + + +def test_reference_statistics_match_rfc_worked_example(): + # RFC #156 section 3.2: probs [0.1, 0.3, 0.6], values [10, 20, 30]; + # full output 25; retaining the last two keys gives m=0.9, C_ref=24. + # A later loop with local softmax [0.75, 0.25] over values [20, 40] + # corrects to 25 - 24 + 0.9 * (0.75*20 + 0.25*40) = 23.5. + q = torch.ones(1, 1) + keys = torch.log(torch.tensor([0.1, 0.3, 0.6])).reshape(3, 1, 1) + values = torch.tensor([10.0, 20.0, 30.0]).reshape(3, 1, 1) + output, contribution, mass, positions = reference_statistics( + q, keys, values, key_retention=0.667 + ) + assert output.item() == pytest.approx(25.0) + assert mass.item() == pytest.approx(0.9) + assert contribution.item() == pytest.approx(24.0) + assert positions.tolist() == [[1, 2]] + + later_keys = torch.log(torch.tensor([3.0, 1.0])).reshape(1, 2, 1) + later_values = torch.tensor([20.0, 40.0]).reshape(1, 2, 1) + corrected = sparse_correction(q, later_keys, later_values, mass, output, contribution) + assert corrected.item() == pytest.approx(23.5) + + +def test_full_retention_matches_dense_attention(): + torch.manual_seed(7) + length, kv_heads, heads, dim = 9, 2, 4, 8 + q = torch.randn(heads, dim) + keys = torch.randn(length, kv_heads, dim) + values = torch.randn(length, kv_heads, dim) + output, contribution, mass, positions = reference_statistics(q, keys, values, key_retention=1.0) + # Full retention keeps every key, so the correction reproduces the dense + # reference output exactly. + dense = reference_statistics(q, keys, values, key_retention=1.0)[0] + groups = heads // kv_heads + corrected = sparse_correction( + q, + keys.permute(1, 0, 2).repeat_interleave(groups, dim=0), + values.permute(1, 0, 2).repeat_interleave(groups, dim=0), + mass, + output, + contribution, + ) + assert positions.shape[1] == length + assert torch.allclose(corrected, dense, atol=1e-5) + assert torch.allclose(output, dense, atol=1e-5) + + +# --------------------------------------------------------------------------- +# Store lifecycle (RFC #156 section 5) +# --------------------------------------------------------------------------- + + +def test_reserve_uses_one_slot_per_live_token(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + assert len(lease) == 3 + assert len(set(lease.slots)) == 3 + assert lease.incarnations == (1, 1, 1) + assert store.outstanding == 3 + assert len(store.free) == 3 + + +def test_release_requires_owner_and_clears_payload(): + store = make_store() + lease = store.reserve("r", 0, 2, generation=1) + fill_payload(store, lease.slots) + store.release(lease) + for slot in lease.slots: + assert store.mass[slot].sum() == 0 + assert store.output[slot].sum() == 0 + assert store.selected[slot].sum() == 0 + assert store.outstanding == 0 + # A retired lease cannot release its slots a second time. + with pytest.raises(RuntimeError, match="different owner"): + store.release(lease) + + +def test_accept_invalidates_rejected_candidates_and_descendants(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + fill_payload(store, lease.slots) + store.accept([lease], torch.tensor([1])) + accepted, middle, rejected = lease.slots + # Rows 0..1 survive the accepted prefix; the rejected candidate and every + # descendant lose the reference payload. + assert store.mass[accepted].sum() > 0 and store.mass[middle].sum() > 0 + assert store.mass[rejected].sum() == 0 + assert store.output[rejected].sum() == 0 + assert store.selected[rejected].sum() == 0 + assert int(store.accepted[accepted]) == 1 + + +def test_accept_keeps_root_when_every_candidate_is_rejected(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + fill_payload(store, lease.slots) + store.accept([lease], torch.tensor([0])) + assert store.mass[lease.slots[0]].sum() > 0 + for slot in lease.slots[1:]: + assert store.mass[slot].sum() == 0 + + +def test_capture_and_consume_mark_layer_state(): + store = make_store() + lease = store.reserve("r", 0, 2, generation=1) + plan = store.build_plan([(lease, 0), (lease, 1)], [CAPTURE, CONSUME], key_retention=0.5) + assert not bool(store.valid[lease.slots[0], 0]) + store.capture(plan, layer=0) + store.capture(plan, layer=1) + assert bool(store.valid[lease.slots[0], 0]) and bool(store.valid[lease.slots[0], 1]) + assert not bool(store.valid[lease.slots[1], 0]) + store.consume(plan, layer=0) + assert int(store.uses[lease.slots[1], 0]) == 1 + assert int(store.uses[lease.slots[0], 0]) == 0 + + +def test_stale_plan_after_reincarnation_is_rejected(): + store = make_store() + lease = store.reserve("r", 0, 2, generation=1) + stale = store.build_plan([(lease, 0), (lease, 1)], [CAPTURE, CONSUME], key_retention=0.5) + store.release(lease) + # Same request and rows, a fresh token incarnation. + renewed = store.reserve("r", 0, 2, generation=1) + assert all(new > old for new, old in zip(renewed.incarnations, lease.incarnations)) + with pytest.raises(RuntimeError, match="stale"): + store.capture(stale, layer=0) + current = store.build_plan([(renewed, 0), (renewed, 1)], [CAPTURE, CONSUME], key_retention=0.5) + store.capture(current, layer=0) + + +def test_outstanding_rounds_hold_disjoint_slots(): + store = make_store(slots=4) + first = store.reserve("a", 0, 2, generation=1) + second = store.reserve("b", 1, 2, generation=1) + assert not set(first.slots) & set(second.slots) + with pytest.raises(RuntimeError, match="no retired attention reference slots"): + store.reserve("c", 2, 1, generation=1) + store.release(first) + third = store.reserve("c", 2, 1, generation=1) + assert set(third.slots) <= set(first.slots) + + +def test_reincarnation_bumps_on_slot_reuse(): + store = make_store() + lease = store.reserve("r", 0, 1, generation=1) + slot = lease.slots[0] + store.release(lease) + again = store.reserve("r", 0, 1, generation=1) + assert again.slots[0] == slot + assert again.incarnations[0] == lease.incarnations[0] + 1 + assert int(store.incarnation[slot]) == again.incarnations[0] + + +def test_bytes_matches_layout(): + store = make_store(slots=3, layers=2, heads=2, head_dim=4, key_bucket=8) + expected = 3 * 2 * 2 * (8 + 1 + 1 + 4 + 4) * 4 + 3 * 2 * (1 + 4) + 3 * (8 + 8 + 4) + assert store.bytes == expected + + +def test_device_metadata_carries_reference_plan(): + class StubCache: + device = torch.device("cpu") + + def _require_live_batch(self, batch): + return None + + builder = MetadataBuilder(StubCache()) + plan = SparseDecodePlan( + modes=(CAPTURE, CONSUME), + slots=(0, 1), + incarnations=(1, 1), + key_retention=0.5, + key_bucket=4, + ) + common = dict( + allocations=(("r", object()),), + position_ids=torch.zeros(2, dtype=torch.int64), + write_blocks=torch.zeros(2, dtype=torch.int64), + write_offsets=torch.zeros(2, dtype=torch.int64), + block_tables=torch.zeros((2, 1), dtype=torch.int32), + context_lengths=torch.zeros(2, dtype=torch.int32), + ) + batch = builder.prepare_device(**common, reference_plan=plan) + assert batch.reference_plan is plan + padded = {**common, "position_ids": torch.zeros(3, dtype=torch.int64)} + with pytest.raises(ValueError, match="one row mode per query row"): + builder.prepare_device(**padded, reference_plan=plan) + + +# --------------------------------------------------------------------------- +# Async speculative integration (GPU) +# --------------------------------------------------------------------------- + + +@pytest.mark.gpu +def test_async_sparse_reference_plans_cover_capture_and_consume(): + model = tiny("cuda") + e = make_speculative_engine(model, sparse=SparseDecodeAttentionConfig(enabled=True)) + store = e.reference_store + assert store is not None + observed = [] + original = store.build_plan + + def recording(owners, modes, **kwargs): + plan = original(owners, modes, **kwargs) + observed.append(plan.modes) + return plan + + store.build_plan = recording + bootstrap(e) + drain(e) + assert any(CAPTURE in modes for modes in observed) + assert any(CONSUME in modes for modes in observed) + # Dense steps never publish a plan. + assert all(DENSE not in modes for modes in observed) + assert store.outstanding == 0 + assert len(store.free) == store.slots + + +@pytest.mark.gpu +@pytest.mark.parametrize("accepted", [0, 1, 2]) +def test_async_sparse_reference_rejection_invalidates_rejected_rows(accepted, monkeypatch): + model = tiny("cuda") + e = make_speculative_engine(model, sparse=SparseDecodeAttentionConfig(enabled=True)) + store = e.reference_store + bootstrap(e) + calls = 0 + + def controlled(hidden): + nonlocal calls + phase = calls % 4 + calls += 1 + logits = torch.full((len(hidden), model.config.vocab_size), -100.0, device="cuda") + if phase < 3: + logits[:, 10 + phase] = 100 + else: + for row in range(4): + logits[row, 10 + row if row < accepted else 20 + row] = 100 + return logits + + monkeypatch.setattr(e.speculative_runner, "_coda", controlled) + + observed = [] + original = store.accept + + def recording(leases, accepted_counts): + # Pretend the reference loop published state, then let the device-side + # retirement decide which rows keep it. + for lease in leases: + fill_payload(store, lease.slots) + original(leases, accepted_counts) + torch.cuda.synchronize() + counts = [int(a) for a in accepted_counts.tolist()] + for lease, count in zip(leases, counts): + observed.append( + ( + count, + [float(store.mass[slot, 0, 0]) for slot in lease.slots], + [int(store.accepted[slot]) for slot in lease.slots], + ) + ) + + store.accept = recording + e.step() + e.step() + e._collect_speculative() + assert observed, "acceptance-driven retirement never ran" + for count, masses, limits in observed: + assert count == accepted + assert all(limit == accepted for limit in limits) + for row, mass in enumerate(masses): + assert mass == (1.0 if row <= accepted else 0.0) + e.abort_request("r") + drain(e) + assert store.outstanding == 0 + assert len(store.free) == store.slots + + +@pytest.mark.gpu +def test_async_sparse_reference_incarnations_advance_across_rounds(): + model = tiny("cuda") + e = make_speculative_engine(model, sparse=SparseDecodeAttentionConfig(enabled=True)) + store = e.reference_store + bootstrap(e) + incarnations = [] + original = store.build_plan + + def recording(owners, modes, **kwargs): + plan = original(owners, modes, **kwargs) + incarnations.append(plan.incarnations) + return plan + + store.build_plan = recording + e.step() + e.step() + e._collect_speculative() + assert len(incarnations) >= 2 + # Every new round is a new token incarnation, even at the same position. + assert incarnations[-1][0] > incarnations[0][0] + e.abort_request("r") + drain(e) + assert store.outstanding == 0 diff --git a/vllm_rlt/__init__.py b/vllm_rlt/__init__.py index d0b346bc..c813eec0 100644 --- a/vllm_rlt/__init__.py +++ b/vllm_rlt/__init__.py @@ -5,6 +5,7 @@ ExecutionConfig, ExitConfig, SchedulerConfig, + SparseDecodeAttentionConfig, SpeculativeConfig, ) from vllm_rlt.entrypoints.llm import LLM @@ -19,6 +20,7 @@ "ExecutionConfig", "ExitConfig", "SchedulerConfig", + "SparseDecodeAttentionConfig", "SpeculativeConfig", "SamplingParams", "RequestOutput", diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py new file mode 100644 index 00000000..29db4fd7 --- /dev/null +++ b/vllm_rlt/attention/sparse.py @@ -0,0 +1,195 @@ +"""Sparse decode attention arithmetic and its per-row plan (RFC #156, section 3.2). + +The sparse decode algorithm keeps the attention output of a reference loop and +replaces only the contribution of the retained keys in a later loop. The +reference loop is loop 2 in the RFC's one-based numbering, which is depth index +1 at runtime. + +This module holds the arithmetic and the plan that describes one submission: + +``reference_statistics`` + the dense reference attention of one row, plus the retained state + ``(S, m, O_ref, C_ref)``, +``sparse_correction`` + the corrected output of a later loop, ``O = O_ref - C_ref + m * C_new``, +``SparseDecodePlan`` + the row modes and reference slots of one submission. + +The arithmetic is deliberately plain PyTorch. The retained state is small, so a +straightforward implementation is easier to check against the worked example in +the RFC than an optimized kernel, and it keeps the contract that the attention +kernels must satisfy visible in one place. + +The plan is the named view between the runner and the attention layer. The +runner owns the reference slots and their lifetime; the attention layer only +borrows the plan for the rows of one submission. Row modes and slots are +host-known: the runner decides both before it submits, so a submission never +depends on a device readback. +""" + +from dataclasses import dataclass + +import torch + +# Row modes for one recurrent submission. +DENSE = 0 +CAPTURE = 1 +CONSUME = 2 + +# The reference loop is loop 2 in the RFC's one-based numbering, which is +# zero-based depth index 1 at runtime. +REFERENCE_DEPTH = 1 + + +def reference_mode(depth: int) -> int: + """Row mode of a recurrent step at the given zero-based loop depth.""" + if depth == REFERENCE_DEPTH: + return CAPTURE + if depth > REFERENCE_DEPTH: + return CONSUME + return DENSE + + +@dataclass(frozen=True) +class SparseDecodePlan: + """The reference plan of one recurrent submission. + + ``modes`` and ``slots`` are host-known: the runner decides both before it + submits, so a submission never depends on a device readback. ``slots`` holds + one reference slot per row, or ``-1`` for a row that owns no reference. + ``incarnations`` identifies the token that occupies each slot; a plan built + for a superseded incarnation is stale and must be rejected. + + A dense row runs the ordinary attention path. A capture row is the + reference loop: it computes the dense attention and publishes the retained + state. A consume row is a later loop: it corrects the reference output with + the retained keys. + + ``store`` is the runner-owned reference pool, borrowed by the attention + layer through duck typing so that model code needs no import of the worker + package. + """ + + modes: tuple[int, ...] + slots: tuple[int, ...] + incarnations: tuple[int, ...] + key_retention: float + key_bucket: int | None + slot_ids: torch.Tensor | None = None + store: object | None = None + + def groups(self) -> dict[int, list[int]]: + """Rows per mode, in submission order.""" + grouped: dict[int, list[int]] = {} + for row, mode in enumerate(self.modes): + grouped.setdefault(mode, []).append(row) + return {mode: rows for mode, rows in grouped.items() if rows} + + @property + def reference_depth(self) -> int: + """The zero-based depth that captures the reference state.""" + return REFERENCE_DEPTH + + def dense(self) -> bool: + """Whether the submission has no sparse work at all.""" + return all(mode == DENSE for mode in self.modes) + + +def retained_count(length: int, key_retention: float, key_bucket: int | None) -> int: + """How many keys a query with ``length`` legal keys retains. + + The retention fraction is rounded to the nearest key. ``key_bucket`` caps + the retained count; ``None`` leaves the retention budget unbounded so that + ``key_retention=1.0`` stays a true full-retention control. + """ + if key_retention <= 0 or key_retention > 1: + raise ValueError("key_retention must be in (0, 1]") + if key_bucket is not None and key_bucket < 1: + raise ValueError("key_bucket must be positive when set") + if length < 1: + return 0 + budget = max(1, int(length * key_retention + 0.5)) + if key_bucket is not None: + budget = min(budget, key_bucket) + return min(budget, length) + + +def expand_kv(tensor: torch.Tensor, groups: int) -> torch.Tensor: + """Repeat KV heads for grouped-query attention.""" + if groups == 1: + return tensor + return tensor.repeat_interleave(groups, dim=1) + + +def select_keys(probabilities: torch.Tensor, count: int) -> torch.Tensor: + """Pick the retained key positions of every head. + + ``probabilities`` is ``[heads, length]``. Equal probabilities are broken by + the smaller absolute position, so the selection does not depend on the order + of the rows in a batch. + """ + if count < 1: + raise ValueError("a selection needs at least one key") + if count > probabilities.shape[1]: + raise ValueError("cannot retain more keys than the query has") + # A stable descending sort keeps equal probabilities in position order. + order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) + return order[:, :count].sort(dim=-1).values + + +def reference_statistics( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + *, + key_retention: float, + key_bucket: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Dense reference attention and the retained state of one row. + + ``q`` is ``[heads, dim]``, ``keys`` and ``values`` are + ``[length, kv_heads, dim]``. Returns ``(output, contribution, mass, + positions)`` where output and contribution are ``[heads, dim]``, mass is + ``[heads]`` and positions are the retained absolute key positions. + """ + heads, dim = q.shape + length = keys.shape[0] + groups = heads // keys.shape[1] + scale = 1.0 / dim**0.5 + expanded_keys = expand_kv(keys, groups).float() + expanded_values = expand_kv(values, groups).float() + scores = torch.einsum("hd,lhd->hl", q.float(), expanded_keys) * scale + probabilities = torch.softmax(scores, dim=-1) + output = torch.einsum("hl,lhd->hd", probabilities, expanded_values) + count = retained_count(length, key_retention, key_bucket) + positions = select_keys(probabilities, count) + gathered = probabilities.gather(1, positions) + # Index the per-head value rows: positions is [heads, count], so the head + # index broadcasts against it and selects each head's own key rows. + head_index = torch.arange(heads, device=q.device)[:, None] + selected_values = expanded_values[positions, head_index] + contribution = torch.einsum("hk,hkd->hd", gathered, selected_values) + mass = gathered.sum(dim=-1) + return output, contribution, mass, positions + + +def sparse_correction( + q: torch.Tensor, + keys: torch.Tensor, + values: torch.Tensor, + mass: torch.Tensor, + output: torch.Tensor, + contribution: torch.Tensor, +) -> torch.Tensor: + """Correct the reference output with the retained keys of a later loop. + + ``keys`` and ``values`` are ``[heads, retained, dim]`` and already carry the + keys of the later loop. ``mass``, ``output`` and ``contribution`` are the + retained reference state. + """ + dim = q.shape[-1] + scale = 1.0 / dim**0.5 + scores = torch.einsum("hd,hkd->hk", q.float(), keys.float()) * scale + probabilities = torch.softmax(scores, dim=-1) + new_contribution = torch.einsum("hk,hkd->hd", probabilities, values.float()) + return output - contribution + mass.float().unsqueeze(-1) * new_contribution diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 648781d7..f12d3368 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -190,3 +190,28 @@ def __post_init__(self): raise ValueError("pipeline draft_loops must divide target_loops") if self.draft_loops >= self.target_loops: raise ValueError("draft_loops must be smaller than target_loops") + + +@dataclass(frozen=True) +class SparseDecodeAttentionConfig: + """Sparse decode attention group (RFC #156, section 7.2). + + ``enabled`` recomputes attention contributions only for the retained keys + of the loop-2 reference; ``key_retention`` is the retained key fraction per + query and head relative to that query's causal key count, where ``1.0`` + provides the full-retention control. ``key_bucket`` caps the retained + count and the reference index buffer; ``None`` derives it from the model + context so that full retention stays exact. + """ + + enabled: bool = False + key_retention: float = 1.0 + key_bucket: int | None = None + + def __post_init__(self): + if type(self.enabled) is not bool: + raise ValueError("enabled must be a boolean") + if not math.isfinite(self.key_retention) or not 0 < self.key_retention <= 1: + raise ValueError("key_retention must be in (0, 1]") + if self.key_bucket is not None: + _positive("key_bucket", self.key_bucket) diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index a83d5101..d53abc55 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -20,6 +20,7 @@ from vllm_rlt.profiling import Profiler from vllm_rlt.request import FinishReason, Request, RequestOutput, RequestPhase from vllm_rlt.sampling_params import SamplingParams +from vllm_rlt.worker.attention_reference import AttentionReferenceStore from vllm_rlt.worker.executor import Executor from vllm_rlt.worker.model_runner import ModelRunner from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput @@ -51,6 +52,7 @@ def __init__( exit_config=None, execution_config=None, speculative_config=None, + sparse_decode_attention_config=None, stage_config=None, connector: StageConnector | None = None, ): @@ -118,6 +120,32 @@ def __init__( scratch = 2 * scratch + 8 * scheduler_config.max_num_seqs * ( config.total_ut_steps * width ) + sparse_decode = sparse_decode_attention_config + self.sparse_decode_attention = sparse_decode + self.reference_store = None + if sparse_decode is not None and sparse_decode.enabled: + # This PR covers the asynchronous self-speculation cell; other + # integration rows are explicit capability failures for now. + if not self.async_speculative: + raise ValueError( + "sparse decode attention currently requires asynchronous self-speculation" + ) + if speculative_config.mode != "block": + raise ValueError("sparse decode attention currently requires block speculation") + # Two in-flight rounds may each hold every verification row, so the + # pool covers both rounds' live tokens with room for slot reuse. + self.reference_store = AttentionReferenceStore( + slots=2 + * scheduler_config.max_num_seqs + * (speculative_config.num_speculative_tokens + 1), + layers=config.num_hidden_layers, + heads=config.num_attention_heads, + head_dim=config.head_dim, + key_bucket=sparse_decode.key_bucket or config.max_position_embeddings, + device=parameter.device, + ) + scratch += self.reference_store.bytes + if speculative_config is not None or self.reference_store is not None: cache_config = replace( cache_config, memory_reserve_bytes=cache_config.memory_reserve_bytes + scratch ) @@ -244,7 +272,13 @@ def __init__( copy_stream=self.model_runner.copy_stream, rng=self.rng, runtime=self.model_runner.runtime, + reference_store=self.reference_store, + sparse_decode=self.sparse_decode_attention, ) + if self.reference_store is not None: + # Reference lifecycle device work is ordered on the compute + # stream, after the readers that retire a slot. + self.reference_store.stream = self.model_runner.core_stream self.executor = Executor(self.model_runner, speculation=self.speculative_runner) self._pending_speculative = [] self._retiring_speculative = {} diff --git a/vllm_rlt/entrypoints/llm.py b/vllm_rlt/entrypoints/llm.py index 55e46104..8c4c83b5 100644 --- a/vllm_rlt/entrypoints/llm.py +++ b/vllm_rlt/entrypoints/llm.py @@ -28,6 +28,7 @@ def __init__( exit_config=None, execution_config=None, speculative_config=None, + sparse_decode_attention_config=None, ): tokenizer_source = resolve_model_source(model) if isinstance(model, str) else None model, revision, model_config = prepare_model(model, revision=revision) @@ -50,6 +51,7 @@ def __init__( exit_config=exit_config, execution_config=execution_config, speculative_config=speculative_config, + sparse_decode_attention_config=sparse_decode_attention_config, ) self._next_request_id = 0 diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index 883d884b..58ae437c 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -5,7 +5,13 @@ from pathlib import Path from types import SimpleNamespace -from vllm_rlt.config import CacheConfig, ExecutionConfig, ExitConfig, SpeculativeConfig +from vllm_rlt.config import ( + CacheConfig, + ExecutionConfig, + ExitConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) from vllm_rlt.profiling import ProfileConfig @@ -48,6 +54,17 @@ def add_runtime_args(parser): parser.add_argument("--exit-trace", help="JSON file with depths_by_request") parser.add_argument("--lookahead-seed", type=int, default=0) parser.add_argument("--async-scheduling", action="store_true") + parser.add_argument( + "--sparse-decode-attention", + action="store_true", + help="Recompute decode attention only for retained reference-loop keys", + ) + parser.add_argument( + "--sparse-decode-key-retention", + type=float, + default=1.0, + help="Retained key fraction per query/head; 1.0 is the full-retention control", + ) parser.add_argument( "--cuda-graphs", action="store_true", @@ -129,6 +146,12 @@ def runtime_configs(args): static_buffers=args.static_buffers, pad_to_power_of_two=args.pad_to_power_of_two, ), + sparse_decode_attention_config=SparseDecodeAttentionConfig( + enabled=args.sparse_decode_attention, + key_retention=args.sparse_decode_key_retention, + ) + if args.sparse_decode_attention + else None, ) diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 6b708df2..563fc082 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -8,6 +8,7 @@ import numpy as np import torch +from vllm_rlt.attention.sparse import DENSE, reference_mode from vllm_rlt.core.kv_cache_manager import _Allocation from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import Operation, Request @@ -50,6 +51,9 @@ class SpeculativeRequestState: event: torch.cuda.Event | None = None # Row in the runner's device slot table holding (token, position). slot: int | None = None + # Bumped per request instance so reference store bindings reject a reused + # request ID across generations. + generation: int = 0 class SpeculativeRoundBuffer: @@ -92,6 +96,10 @@ class SpeculativeRoundTicket: stats: SpeculativeStats | None = None cached: ModelRunnerOutput | None = None retired: bool = False + # Reference leases of this round. Released on retire, after the final + # reader completed, so no later submission can observe or overwrite them. + references: tuple = () + reference_store: object | None = None @property def has_result(self): @@ -148,6 +156,9 @@ def retire(self, *, discard=False): # collect() already waited for copy_done; normal delivery adds no wait. self.buffer.keepalive.clear() self.buffer.leased = False + if self.reference_store is not None and self.references: + self.reference_store.release(self.references) + self.references = () self.retired = True @@ -158,6 +169,11 @@ class AsyncSpeculativeRunner(SpeculativeRunner): reservation upper bound and must never be used as the actual write position. """ + # Sparse decode attention lifecycle; test doubles that skip __init__ + # inherit the disabled defaults. + reference_store = None + sparse_decode = None + def __init__( self, model, @@ -170,6 +186,8 @@ def __init__( copy_stream, rng: RngRegistry | None = None, runtime=None, + reference_store=None, + sparse_decode=None, ): super().__init__( model, @@ -180,6 +198,12 @@ def __init__( runtime=runtime, compute_stream=compute_stream, ) + # Sparse decode attention lifecycle: the runner owns the per-token + # reference pool (RFC #156 section 5) and drives retirement from the + # GPU acceptance counts without a host readback. + self.reference_store = reference_store + self.sparse_decode = sparse_decode + self._generations: dict[str, int] = {} sizes = speculative_capture_sizes(execution, scheduler, config) self.graphs = ( DeviceRecurrentGraphs(model, cache, execution, False, capture_sizes=sizes) @@ -222,15 +246,24 @@ def _state(self, request, buffer): host = torch.tensor([request.input_token_id, request.position], pin_memory=True) values = host.to(self.device, non_blocking=True) buffer.keepalive.extend((host, values)) + generation = self._generations.get(rid, 0) + 1 + self._generations[rid] = generation if self.config.mode == "pipeline": - state = SpeculativeRequestState(request, allocation, values[0], values[1]) + state = SpeculativeRequestState( + request, allocation, values[0], values[1], generation=generation + ) else: if not self.free_slots: self._grow_slots() slot = self.free_slots.pop() self.slot_values[slot].copy_(values) state = SpeculativeRequestState( - request, allocation, self.slot_values[slot, 0], self.slot_values[slot, 1], slot=slot + request, + allocation, + self.slot_values[slot, 0], + self.slot_values[slot, 1], + slot=slot, + generation=generation, ) if self.config.mode == "pipeline": width = min( @@ -264,7 +297,9 @@ def _grow_slots(self): for state in self.states.values(): state.token, state.position = new[state.slot, 0], new[state.slot, 1] - def _execute_recurrent_step(self, hidden, positions, tables, depth, allocations, active=None): + def _execute_recurrent_step( + self, hidden, positions, tables, depth, allocations, active=None, reference_plan=None + ): if isinstance(depth, torch.Tensor): table = tables.gather(1, depth[:, None, None].expand(-1, 1, tables.shape[2])).squeeze(1) else: @@ -280,6 +315,7 @@ def _execute_recurrent_step(self, hidden, positions, tables, depth, allocations, (positions + 1) if active is None else torch.where(active, positions + 1, 0) ).int(), write_mask=active, + reference_plan=reference_plan, ) forward = ForwardBatch( Operation.RECURRENT, hidden=hidden, metadata=(metadata,), graphs=self.graphs @@ -365,9 +401,27 @@ def add(name, values): values = host.to(self.device, non_blocking=True) buffer.keepalive.extend((host, values)) views = {name: values[start : start + size] for name, (start, size) in layout.items()} - return SimpleNamespace(n=n, width=width, rows=len(request_rows), offsets=offsets, **views) + return SimpleNamespace( + n=n, width=width, rows=len(request_rows), offsets=offsets, counts=counts, **views + ) + + def _reference_plan(self, leases, rows, depth): + """Build the sparse decode plan of one recurrent step. + + ``rows`` holds one ``(lease_index, row)`` pair per step row in step + order. Dense steps return ``None`` so their metadata stays untouched. + """ + if self.reference_store is None or leases is None: + return None + mode = reference_mode(depth) + if mode == DENSE: + return None + owners = [(leases[lease_index], row) for lease_index, row in rows] + return self.reference_store.build_plan( + owners, [mode] * len(owners), key_retention=self.sparse_decode.key_retention + ) - def _draft(self, plan, tables, allocations): + def _draft(self, plan, tables, allocations, leases=None): from vllm_rlt.kernels.routing import gather_kernel slot_values = self.slot_values.new_empty((plan.n, 2)) @@ -394,9 +448,18 @@ def _draft(self, plan, tables, allocations): positions = starts[active] + offset row_tables = tables[active] hidden = self.runtime.forward(ForwardBatch(Operation.PRELUDE, tokens=row_tokens)).hidden + # A drafting offset covers one verification row per active request: + # the token at column `offset`, which the verifier later reuses. + rows = range(plan.n) if active is None else [int(index) for index in active] for depth in range(self.config.draft_loops): + reference = self._reference_plan(leases, [(i, offset) for i in rows], depth) hidden = self._execute_recurrent_step( - hidden, positions, row_tables, depth, allocations + hidden, + positions, + row_tables, + depth, + allocations, + reference_plan=reference, ) hidden_blocks.append(hidden) if drafting_count: @@ -406,15 +469,31 @@ def _draft(self, plan, tables, allocations): tokens[getattr(plan, f"drafting{offset}"), offset + 1] = proposed return starts, tokens, hidden_blocks - def _verify(self, plan, starts, hidden_blocks, tables, allocations): + def _verify(self, plan, starts, hidden_blocks, tables, allocations, leases=None): positions = starts[plan.request_rows] + plan.row_offsets row_tables = tables[plan.request_rows] hidden = torch.cat(hidden_blocks)[plan.hidden_rows] + # Verification rows are request-major, matching the round plan. + verification_rows = [ + (i, j) for i, count in enumerate(plan.counts) for j in range(int(count)) + ] for depth in range(self.config.draft_loops, self.config.target_loops): - hidden = self._execute_recurrent_step(hidden, positions, row_tables, depth, allocations) + reference = self._reference_plan(leases, verification_rows, depth) + if reference is not None and depth == self.config.draft_loops: + # One device-owned absolute position per reference slot; the + # verification rows cover every slot of the round exactly once. + self.reference_store.bind_positions(reference, positions) + hidden = self._execute_recurrent_step( + hidden, + positions, + row_tables, + depth, + allocations, + reference_plan=reference, + ) return self._coda(hidden).argmax(-1) - def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer): + def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer, leases=None): from vllm_rlt.kernels.routing import scatter_kernel count = len(target) @@ -426,6 +505,10 @@ def _accept_and_update_state(self, plan, states, starts, tokens, target, buffer) drafts = torch.arange(plan.width - 1, device=self.device)[None, :] < plan.drafts[:, None] matches = (tokens[:, 1:] == targets[:, :-1]) & drafts accepted = matches.long().cumprod(1).sum(1) + if leases is not None and self.reference_store is not None: + # GPU acceptance drives retirement: rejected candidates and their + # dependent descendants lose their references on the device. + self.reference_store.accept(leases, accepted) # All target IDs through the first rejection equal the accepted # candidates followed by the correction (or final bonus). next_tokens = targets.gather(1, accepted[:, None]).squeeze(1) @@ -462,14 +545,40 @@ def submit(self, batch): rounds = tuple(s.submitted for s in states) if any(s.submitted > s.delivered for s in states): self.submitted_before_collect += 1 - tables.copy_(host_tables, non_blocking=True) - if self.config.mode == "pipeline": - count = self._pipeline_step(scheduled_requests, states, tables, allocations, buffer) - else: - plan = self._round_plan(scheduled_requests, states, buffer) - starts, tokens, hidden_blocks = self._draft(plan, tables, allocations) - target = self._verify(plan, starts, hidden_blocks, tables, allocations) - count = self._accept_and_update_state(plan, states, starts, tokens, target, buffer) + # One reference slot per verification row: the round's live tokens. + # Outstanding rounds keep disjoint slots until their tickets retire. + leases = None + if self.reference_store is not None: + leases = tuple( + self.reference_store.reserve( + scheduled_request.request.request_id, + index, + scheduled_request.token_count, + state.generation, + ) + for index, (scheduled_request, state) in enumerate( + zip(scheduled_requests, states) + ) + ) + try: + tables.copy_(host_tables, non_blocking=True) + if self.config.mode == "pipeline": + count = self._pipeline_step( + scheduled_requests, states, tables, allocations, buffer + ) + else: + plan = self._round_plan(scheduled_requests, states, buffer) + starts, tokens, hidden_blocks = self._draft(plan, tables, allocations, leases) + target = self._verify(plan, starts, hidden_blocks, tables, allocations, leases) + count = self._accept_and_update_state( + plan, states, starts, tokens, target, buffer, leases + ) + except BaseException: + # The engine drains both streams before discarding, so releasing + # here cannot race a later submission for these slots. + if leases is not None: + self.reference_store.release(leases) + raise buffer.compute_done = torch.cuda.Event() buffer.compute_done.record(self.compute_stream) for state in states: @@ -497,6 +606,8 @@ def submit(self, batch): count, mode=self.config.mode, stats=self.stats, + references=leases or (), + reference_store=self.reference_store, ) def _pipeline_step(self, spans, states, tables, allocations, buffer): @@ -685,6 +796,12 @@ def discard_buffers(self): for buffer in self.buffers: buffer.keepalive.clear() buffer.leased = False + # Test doubles borrow this method without the runner __init__. + reference_store = getattr(self, "reference_store", None) + if reference_store is not None: + # Both streams have drained and every ticket retired; any leftover + # binding can only come from a failed submission. + reference_store.discard() def synchronize(self): self.compute_stream.synchronize() diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py new file mode 100644 index 00000000..eb0bafbe --- /dev/null +++ b/vllm_rlt/worker/attention_reference.py @@ -0,0 +1,307 @@ +"""Per-token attention reference state for sparse decode attention. + +A sparse decode loop reuses the reference loop's attention output and only +recomputes the contribution of the retained keys. The retained state is +per token, per layer and per head, so it cannot be stored per request. + +This module owns that state. The store keeps one slot per live token, tracks +which request and generation owns it, and publishes the slot views that the +attention kernels borrow. Nothing here reads device state back to the host: +ownership, retirement and invalidation are expressed as device operations so +that an asynchronous round can retire its slots without a host readback. + +Slot lifecycle +-------------- +``empty`` -> ``source-ready`` -> ``reusable`` -> ``retired`` + +A slot becomes source-ready when its capture round finished writing the +reference, and reusable once a later loop consumed it. Retirement is ordered +after the last reader of the slot, so a later submission cannot observe a slot +that an earlier round still reads. + +Slot identity +------------- +A slot is identified by its request generation, its row inside the request and +an incarnation counter. The incarnation changes whenever the token of that row +changes, including a corrected token that reuses the same position after a +rejection. A plan built for a superseded incarnation is stale and rejected. +""" + +from dataclasses import dataclass + +import torch + +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan + +__all__ = ["ReferenceLease", "AttentionReferenceStore", "DENSE", "CAPTURE", "CONSUME"] + + +@dataclass(frozen=True) +class ReferenceLease: + """One request's reference allocation for one speculative submission. + + ``slots`` maps the request's verification row (token position) to its + reference slot. The device tensors mirror them so that acceptance-driven + retirement stays on the device. + """ + + request_id: str + generation: int + request_index: int + slots: tuple[int, ...] + incarnations: tuple[int, ...] + slot_ids: torch.Tensor + row_index: torch.Tensor + + def __len__(self) -> int: + return len(self.slots) + + +class AttentionReferenceStore: + """Device-resident reference state for the live tokens of one stage. + + The store is sized by the engine, which knows how many tokens can be live + at once. It never allocates per request: a speculative round holds one slot + per verification row, and several rounds of the same request may be + outstanding at the same time, so the capacity must cover every outstanding + round. + """ + + def __init__(self, *, slots, layers, heads, head_dim, key_bucket, device, stream=None): + if slots < 1 or layers < 1 or heads < 1 or head_dim < 1: + raise ValueError("reference store dimensions must be positive") + if key_bucket is not None and key_bucket < 1: + raise ValueError("key_bucket must be positive when set") + self.slots = slots + self.layers = layers + self.heads = heads + self.head_dim = head_dim + self.key_bucket = key_bucket + self.device = torch.device(device) + self.stream = stream + shape = (slots, layers, heads) + self.selected = torch.zeros((*shape, key_bucket or 1), dtype=torch.int32, device=device) + self.counts = torch.zeros(shape, dtype=torch.int32, device=device) + self.mass = torch.zeros(shape, dtype=torch.float32, device=device) + self.output = torch.zeros((*shape, head_dim), dtype=torch.float32, device=device) + self.contribution = torch.zeros((*shape, head_dim), dtype=torch.float32, device=device) + # A reference is usable per layer while its slot is live. + self.valid = torch.zeros((slots, layers), dtype=torch.bool, device=device) + # How many later loops consumed the reference of each slot and layer. + self.uses = torch.zeros((slots, layers), dtype=torch.int32, device=device) + # Identity of the token occupying each slot. + self.position = torch.zeros(slots, dtype=torch.int64, device=device) + self.incarnation = torch.zeros(slots, dtype=torch.int64, device=device) + self.accepted = torch.zeros(slots, dtype=torch.int32, device=device) + self.free = list(range(slots)) + self._bindings: dict[int, tuple[str, int, int]] = {} + self._incarnations: dict[tuple[str, int], int] = {} + + @property + def bytes(self) -> int: + """Payload bytes held by the store, excluding allocator overhead.""" + return sum( + tensor.numel() * tensor.element_size() + for tensor in ( + self.selected, + self.counts, + self.mass, + self.output, + self.contribution, + self.valid, + self.uses, + self.position, + self.incarnation, + self.accepted, + ) + ) + + @property + def outstanding(self) -> int: + """Number of slots currently bound to live tokens.""" + return len(self._bindings) + + def reserve(self, request_id: str, request_index: int, rows: int, generation: int): + """Take one slot per verification row of a request's submission.""" + if rows < 1: + raise ValueError("a reference reservation needs at least one row") + if len(self.free) < rows: + raise RuntimeError("no retired attention reference slots") + slots = tuple(self.free.pop() for _ in range(rows)) + incarnations = [] + for row, slot in enumerate(slots): + count = self._incarnations.get((request_id, row), 0) + 1 + self._incarnations[(request_id, row)] = count + self._bindings[slot] = (request_id, generation, count) + self.incarnation[slot] = count + incarnations.append(count) + return ReferenceLease( + request_id=request_id, + generation=generation, + request_index=request_index, + slots=slots, + incarnations=tuple(incarnations), + slot_ids=torch.tensor(slots, dtype=torch.int64, device=self.device), + row_index=torch.arange(rows, dtype=torch.int64, device=self.device), + ) + + def release(self, leases) -> None: + """Return slots to the pool once their last reader completed. + + The caller orders this after the final reader/writer, e.g. by releasing + from a ticket that already waited for the round's completion event. The + cleared payload never leaks into a later incarnation. + """ + if isinstance(leases, ReferenceLease): + leases = (leases,) + cleared = [] + for lease in leases: + for slot in lease.slots: + binding = self._bindings.get(slot) + if binding is None or binding[:2] != (lease.request_id, lease.generation): + raise RuntimeError("attention reference slot released by a different owner") + del self._bindings[slot] + cleared.append(slot) + if cleared: + self.free.extend(cleared) + self._clear(torch.tensor(cleared, dtype=torch.int64, device=self.device)) + + def discard(self) -> None: + """Free every binding without device work; the caller has drained.""" + self._bindings.clear() + self.free = list(range(self.slots)) + + def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan: + """Publish the plan of one submission. + + ``owners`` holds one ``(lease, row)`` pair per submission row, in the + row order of the recurrent step; ``modes`` holds the row mode of every + submission row. Slots and incarnations are re-read from the live + bindings, so a plan never aliases a superseded token. + """ + if len(owners) != len(modes): + raise ValueError("a reference plan needs one mode per row") + slots, incarnations = [], [] + for lease, row in owners: + slot = lease.slots[row] + binding = self._bindings.get(slot) + if binding is None or binding != ( + lease.request_id, + lease.generation, + lease.incarnations[row], + ): + raise RuntimeError("attention reference slot released before its plan") + slots.append(slot) + incarnations.append(binding[2]) + return SparseDecodePlan( + modes=tuple(modes), + slots=tuple(slots), + incarnations=tuple(incarnations), + key_retention=key_retention, + key_bucket=self.key_bucket, + slot_ids=torch.tensor(slots, dtype=torch.int64, device=self.device), + store=self, + ) + + def _require_live(self, plan: SparseDecodePlan) -> None: + """Reject a plan whose slots were recycled or superseded.""" + if plan.store is not self: + raise RuntimeError("attention reference plan belongs to a different store") + for slot, incarnation in zip(plan.slots, plan.incarnations): + binding = self._bindings.get(slot) + if binding is None or binding[2] != incarnation: + raise RuntimeError("stale attention reference plan for a retired token") + + def capture(self, plan: SparseDecodePlan, layer: int) -> None: + """Mark the captured rows source-ready for the given layer. + + The attention kernel writes the retained state of a capture row and + then calls this on the same stream, so a later reader observes a + complete payload. + """ + self._require_live(plan) + rows = plan.groups().get(CAPTURE) + if not rows: + return + self.valid[plan.slot_ids[rows], layer] = True + + def consume(self, plan: SparseDecodePlan, layer: int) -> None: + """Record that the given layer consumed the reference rows.""" + self._require_live(plan) + rows = plan.groups().get(CONSUME) + if not rows: + return + self.uses[plan.slot_ids[rows], layer] += 1 + + def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> None: + """Record the device-owned absolute positions of the plan's rows.""" + self._require_live(plan) + if positions.shape != plan.slot_ids.shape: + raise ValueError("reference positions need one entry per plan row") + self.position[plan.slot_ids] = positions.to(torch.int64) + + def accept(self, leases, accepted: torch.Tensor) -> None: + """Invalidate the references of rejected rows on the device. + + ``accepted`` is the number of accepted candidates of every request of + the submission. A verification row of a request survives exactly when + its row index is at most the accepted count: the root row is always + kept, and a rejected candidate invalidates its dependent descendants. + The comparison stays on the device so that acceptance never requires a + host readback. + """ + if isinstance(leases, ReferenceLease): + leases = (leases,) + leases = tuple(leases) + if not leases: + return + slot_ids = torch.cat([lease.slot_ids for lease in leases]) + rows = torch.cat([lease.row_index for lease in leases]) + request = torch.cat( + [torch.full_like(lease.row_index, lease.request_index) for lease in leases] + ) + limit = accepted.to(device=self.device, dtype=rows.dtype).index_select(0, request) + keep = rows <= limit + keep_rows = keep.reshape(-1) + keep3 = keep_rows.view(-1, 1, 1) + keep4 = keep_rows.view(-1, 1, 1, 1) + self.accepted[slot_ids] = limit.to(torch.int32) + self.valid[slot_ids] = self.valid[slot_ids] & keep_rows.view(-1, 1) + self.mass[slot_ids] = self.mass[slot_ids] * keep3 + self.output[slot_ids] = self.output[slot_ids] * keep4 + self.contribution[slot_ids] = self.contribution[slot_ids] * keep4 + self.selected[slot_ids] = self.selected[slot_ids] * keep4.to(self.selected.dtype) + self.counts[slot_ids] = self.counts[slot_ids] * keep3.to(self.counts.dtype) + + def _clear(self, slots: torch.Tensor) -> None: + """Zero the payload of the given slots, ordered on the owner stream.""" + + def zero(): + self.valid[slots] = False + self.uses[slots] = 0 + self.accepted[slots] = 0 + self.incarnation[slots] = 0 + self.position[slots] = 0 + self.mass[slots] = 0 + self.output[slots] = 0 + self.contribution[slots] = 0 + self.selected[slots] = 0 + self.counts[slots] = 0 + + if self.stream is not None: + with torch.cuda.stream(self.stream): + zero() + else: + zero() + + def invalidate(self, slots) -> None: + """Invalidate the references of the given slots without freeing them.""" + if len(slots) == 0: + return + slots = torch.as_tensor(slots, dtype=torch.int64, device=self.device) + self.valid[slots] = False + self.mass[slots] = 0 + self.output[slots] = 0 + self.contribution[slots] = 0 + self.selected[slots] = 0 + self.counts[slots] = 0 diff --git a/vllm_rlt/worker/metadata.py b/vllm_rlt/worker/metadata.py index 8f2c020b..4c7e7cc0 100644 --- a/vllm_rlt/worker/metadata.py +++ b/vllm_rlt/worker/metadata.py @@ -62,6 +62,10 @@ class DeviceKVMetadata: context_lengths: torch.Tensor # Inactive pipeline slots must not overwrite any live KV address. write_mask: torch.Tensor | None = None + # Sparse decode attention plan of this submission: one row mode and + # reference slot per row. Host-known; the attention layer borrows it and + # the runner-owned store keeps the referenced state alive. + reference_plan: object | None = None # The device-owned path currently uses independent single-query rows. # These layout fields also let recurrent graphs stage the same contract. cu_seqlens_q: torch.Tensor | None = None @@ -100,6 +104,7 @@ def prepare_device( block_tables, context_lengths, write_mask=None, + reference_plan=None, ) -> DeviceKVMetadata: """Borrow stream-ordered metadata without reading device positions on CPU. @@ -115,7 +120,10 @@ def prepare_device( block_tables=block_tables, context_lengths=context_lengths, write_mask=write_mask, + reference_plan=reference_plan, ) + if reference_plan is not None and len(reference_plan.modes) != len(position_ids): + raise ValueError("reference plan needs one row mode per query row") if write_mask is not None and ( write_mask.shape != position_ids.shape or write_mask.dtype != torch.bool From ccf0aef5201a0128dbc65c91b5eafbc70160789a Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Wed, 28 Oct 2026 14:26:56 +0800 Subject: [PATCH 2/8] sparse decode attention: align retained count with FlashLoop and separate capacity Aligns with the review on #157 and the FlashLoop reference (aeaadee select_decode_keys): - retained_count follows the reference rule k = min(N, max(1, ceil(N * key_retention))) instead of rounding to nearest; boundary tests include the N=13 / retention=0.1 case (2 keys, not 1). - key_bucket is buffer capacity only, never an algorithm knob: the retained count is a pure function of key_retention, and a count above the capacity fails loudly instead of silently truncating retention. - key_bucket=None sizes the index buffer as ceil(context x retention) rather than always max_position_embeddings, cutting per-slot index memory proportionally to retention (e.g. 24L/16H/65k at retention 0.1 costs ~1/10 of the indices); explicit buckets below the worst-case count are rejected at startup. - engine-level e2e: key buffer sizing, insufficient-bucket rejection, and a multi-request round with a mid-flight abort and full slot reclamation. --- tests/test_sparse_decode_attention.py | 99 +++++++++++++++++++++++---- vllm_rlt/attention/sparse.py | 24 ++++--- vllm_rlt/config.py | 7 +- vllm_rlt/engine/stage_engine.py | 13 +++- 4 files changed, 116 insertions(+), 27 deletions(-) diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index 6c77289c..245540b2 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -7,11 +7,13 @@ completion and reference-slot reuse. """ +import math + import pytest import torch from tests.helpers import bootstrap_speculation, make_speculative_engine, tiny_ouro_model -from vllm_rlt import SparseDecodeAttentionConfig +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig from vllm_rlt.attention.sparse import ( CAPTURE, CONSUME, @@ -81,15 +83,30 @@ def test_plan_groups_rows_by_mode(): assert plan.reference_depth == REFERENCE_DEPTH -def test_retained_count_rounds_and_caps(): - assert retained_count(3, 0.667, None) == 2 - assert retained_count(4, 1.0, None) == 4 - assert retained_count(4, 1.0, 2) == 2 - assert retained_count(0, 0.5, None) == 0 +def test_retained_count_follows_flashloop_ceil_rule(): + # FlashLoop aeaadee select_decode_keys: k = min(N, max(1, ceil(N * r))). + assert retained_count(13, 0.1) == 2 + assert retained_count(10, 0.1) == 1 + assert retained_count(10, 0.5) == 5 + assert retained_count(9, 0.5) == 5 + assert retained_count(1, 0.1) == 1 + assert retained_count(4, 1.0) == 4 + assert retained_count(0, 0.5) == 0 with pytest.raises(ValueError, match="key_retention"): - retained_count(4, 1.5, None) - with pytest.raises(ValueError, match="key_bucket"): - retained_count(4, 0.5, 0) + retained_count(4, 1.5) + + +def test_reference_statistics_fail_when_capacity_cannot_hold_retained_keys(): + # Capacity never truncates the algorithm: a count above it fails loudly. + q = torch.ones(1, 1) + keys = torch.log(torch.tensor([0.2, 0.3, 0.5])).reshape(3, 1, 1) + values = torch.ones(3, 1, 1) + with pytest.raises(ValueError, match="cannot hold"): + reference_statistics(q, keys, values, key_retention=0.5, key_bucket=1) + # The same query passes with sufficient capacity. + output, _, _, positions = reference_statistics(q, keys, values, key_retention=0.5, key_bucket=2) + assert positions.shape[1] == 2 + assert output.shape == (1, 1) def test_reference_statistics_match_rfc_worked_example(): @@ -100,9 +117,7 @@ def test_reference_statistics_match_rfc_worked_example(): q = torch.ones(1, 1) keys = torch.log(torch.tensor([0.1, 0.3, 0.6])).reshape(3, 1, 1) values = torch.tensor([10.0, 20.0, 30.0]).reshape(3, 1, 1) - output, contribution, mass, positions = reference_statistics( - q, keys, values, key_retention=0.667 - ) + output, contribution, mass, positions = reference_statistics(q, keys, values, key_retention=0.5) assert output.item() == pytest.approx(25.0) assert mass.item() == pytest.approx(0.9) assert contribution.item() == pytest.approx(24.0) @@ -393,3 +408,63 @@ def recording(owners, modes, **kwargs): e.abort_request("r") drain(e) assert store.outstanding == 0 + + +@pytest.mark.gpu +def test_sparse_engine_sizes_key_buffer_from_retention(): + model = tiny("cuda") + e = make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.25) + ) + # None key_bucket derives the capacity from context x retention, not the + # full model context (FlashLoop ceil rule at the worst-case key count). + assert e.reference_store.key_bucket == math.ceil(model.config.max_position_embeddings * 0.25) + + +@pytest.mark.gpu +def test_sparse_engine_rejects_insufficient_key_bucket(): + model = tiny("cuda") + with pytest.raises(ValueError, match="key_bucket"): + make_speculative_engine( + model, + sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, key_bucket=4), + ) + + +@pytest.mark.gpu +def test_sparse_engine_e2e_multiple_requests_abort_and_drain(): + # Engine-level lifecycle e2e: concurrent rounds across requests, an abort + # in flight, and full slot reclamation without stale references. + model = tiny("cuda") + e = make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.5) + ) + store = e.reference_store + max_outstanding = 0 + original_reserve = store.reserve + + def counting_reserve(*args, **kwargs): + nonlocal max_outstanding + lease = original_reserve(*args, **kwargs) + max_outstanding = max(max_outstanding, store.outstanding) + return lease + + store.reserve = counting_reserve + params = [SamplingParams(max_tokens=n, ignore_eos=True) for n in (5, 9, 50)] + for i in range(3): + e.add_request(str(i), [2, 3, 4], params[i]) + outputs = {} + for _ in range(1000): + if outputs.get("0") is not None and outputs["0"].finished: + break + for output in e.step(): + outputs[output.request_id] = output + # Abort one request mid-flight; the rest must finish cleanly. + if "2" not in outputs or not outputs["2"].finished: + e.abort_request("2") + outputs.update(drain(e)) + assert outputs["0"].finished and len(outputs["0"].token_ids) == 5 + assert outputs["1"].finished and len(outputs["1"].token_ids) == 9 + assert store.outstanding == 0 + assert len(store.free) == store.slots + assert max_outstanding > 0 diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index 29db4fd7..aa1654a6 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -27,6 +27,7 @@ depends on a device readback. """ +import math from dataclasses import dataclass import torch @@ -95,23 +96,20 @@ def dense(self) -> bool: return all(mode == DENSE for mode in self.modes) -def retained_count(length: int, key_retention: float, key_bucket: int | None) -> int: +def retained_count(length: int, key_retention: float) -> int: """How many keys a query with ``length`` legal keys retains. - The retention fraction is rounded to the nearest key. ``key_bucket`` caps - the retained count; ``None`` leaves the retention budget unbounded so that - ``key_retention=1.0`` stays a true full-retention control. + Follows the FlashLoop reference (``aeaadee``, ``select_decode_keys``): + ``k = min(length, max(1, ceil(length * key_retention)))``. The count is a + pure function of the retention fraction, so ``key_retention=1.0`` stays an + exact full-retention control; index-buffer capacity lives in the store and + never changes the algorithm. """ if key_retention <= 0 or key_retention > 1: raise ValueError("key_retention must be in (0, 1]") - if key_bucket is not None and key_bucket < 1: - raise ValueError("key_bucket must be positive when set") if length < 1: return 0 - budget = max(1, int(length * key_retention + 0.5)) - if key_bucket is not None: - budget = min(budget, key_bucket) - return min(budget, length) + return min(length, max(1, math.ceil(length * key_retention))) def expand_kv(tensor: torch.Tensor, groups: int) -> torch.Tensor: @@ -161,7 +159,11 @@ def reference_statistics( scores = torch.einsum("hd,lhd->hl", q.float(), expanded_keys) * scale probabilities = torch.softmax(scores, dim=-1) output = torch.einsum("hl,lhd->hd", probabilities, expanded_values) - count = retained_count(length, key_retention, key_bucket) + count = retained_count(length, key_retention) + # Buffer capacity never truncates the algorithm: a count above it fails + # loudly instead of silently degrading retention. + if key_bucket is not None and count > key_bucket: + raise ValueError(f"reference capacity {key_bucket} cannot hold {count} retained keys") positions = select_keys(probabilities, count) gathered = probabilities.gather(1, positions) # Index the per-head value rows: positions is [heads, count], so the head diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index f12d3368..ceb80daa 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -199,9 +199,10 @@ class SparseDecodeAttentionConfig: ``enabled`` recomputes attention contributions only for the retained keys of the loop-2 reference; ``key_retention`` is the retained key fraction per query and head relative to that query's causal key count, where ``1.0`` - provides the full-retention control. ``key_bucket`` caps the retained - count and the reference index buffer; ``None`` derives it from the model - context so that full retention stays exact. + provides the full-retention control. ``key_bucket`` is only the reference + index-buffer capacity, never an algorithm knob: the retained count follows + ``key_retention`` alone, and ``None`` sizes the buffer from the model + context and the retention fraction. """ enabled: bool = False diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index d53abc55..049431ac 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -1,4 +1,5 @@ import logging +import math from dataclasses import dataclass, replace from typing import TYPE_CHECKING, cast @@ -132,6 +133,16 @@ def __init__( ) if speculative_config.mode != "block": raise ValueError("sparse decode attention currently requires block speculation") + # The index buffer never truncates the algorithm: capacity must + # hold the worst-case retained count over the model context. + required_keys = math.ceil(config.max_position_embeddings * sparse_decode.key_retention) + capacity = sparse_decode.key_bucket or required_keys + if capacity < required_keys: + raise ValueError( + f"sparse decode attention key_bucket {capacity} cannot hold " + f"{required_keys} retained keys for key_retention " + f"{sparse_decode.key_retention}; raise key_bucket or lower key_retention" + ) # Two in-flight rounds may each hold every verification row, so the # pool covers both rounds' live tokens with room for slot reuse. self.reference_store = AttentionReferenceStore( @@ -141,7 +152,7 @@ def __init__( layers=config.num_hidden_layers, heads=config.num_attention_heads, head_dim=config.head_dim, - key_bucket=sparse_decode.key_bucket or config.max_position_embeddings, + key_bucket=capacity, device=parameter.device, ) scratch += self.reference_store.bytes From db5ce44642402102cc2a879196aac34d8a4026b5 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Wed, 28 Oct 2026 20:50:58 +0800 Subject: [PATCH 3/8] sparse decode attention: reference attention path and end-to-end validation Makes the sparse decode column executable end to end so the optimization can be measured instead of only asserted: - attention/sparse.py: sparse_paged_attention, the RFC 3.2 reference path over the paged KV cache. Dense rows reproduce ordinary attention, capture rows publish (S, m, O_ref, C_ref) into the plan's store, consume rows correct the stored reference with the current loop's retained keys. Plain PyTorch, no backend dependency, so it is the numerical oracle for the optimized kernels. - worker/attention_reference.py: write_reference/read_reference views for that path, with capacity enforced loudly. - core/kv_cache_manager.py: a plan-carrying device metadata routes to the reference path; metadata without a plan keeps the existing backend dispatch unchanged. - tests: end-to-end decode through the real model, paged cache and attention (full retention matches dense; divergence grows as retention shrinks; the consume path reads only the retained budget), plus an engine-level test that a sparse engine reproduces the dense engine's tokens exactly at retention 1.0. - benchmarks/sparse_decode_attention_bench.py: reference-path measurement of the dense/capture/consume costs, keys touched, reference-state bytes and the same-loop identity vs later-loop divergence. --- benchmarks/sparse_decode_attention_bench.py | 196 ++++++++++++++++++++ tests/test_sparse_decode_attention.py | 170 ++++++++++++++++- vllm_rlt/attention/sparse.py | 106 +++++++++++ vllm_rlt/core/kv_cache_manager.py | 15 ++ vllm_rlt/worker/attention_reference.py | 39 ++++ 5 files changed, 525 insertions(+), 1 deletion(-) create mode 100644 benchmarks/sparse_decode_attention_bench.py diff --git a/benchmarks/sparse_decode_attention_bench.py b/benchmarks/sparse_decode_attention_bench.py new file mode 100644 index 00000000..3f0fe63f --- /dev/null +++ b/benchmarks/sparse_decode_attention_bench.py @@ -0,0 +1,196 @@ +"""Sparse decode attention reference-path benchmark (RFC #156, section 3.2). + +Measures the reference implementation of the decode correction against the +dense reference path on identical paged-KV inputs: + +* ``dense`` one loop of full attention over ``N`` keys, +* ``capture`` the reference loop: dense attention plus retained statistics, +* ``consume`` a later loop: correction over ``k = ceil(N * retention)`` keys. + +A four-loop decode therefore costs ``4 x dense`` in the baseline and +``dense + capture + 2 x consume`` with sparse decode attention. This measures +the algorithm on the reference path, not an optimized kernel: the reported +speedup is a work proxy, not a serving claim (RFC #156 section 9). + +Usage: + python benchmarks/sparse_decode_attention_bench.py --device cuda + python benchmarks/sparse_decode_attention_bench.py --device cpu --context 512 +""" + +import argparse +import json +import math +import time + +import torch + +from vllm_rlt.attention.sparse import ( + reference_statistics, + retained_count, + sparse_correction, +) +from vllm_rlt.kernels.paged_attention import torch_paged_attention + + +def build_paged(context, heads, head_dim, kv_heads, block_size, device, dtype): + blocks = math.ceil(context / block_size) + shape = (blocks, block_size, kv_heads, head_dim) + generator = torch.Generator(device="cpu").manual_seed(0) + key_cache = torch.randn(shape, generator=generator, dtype=torch.float32).to(device, dtype) + value_cache = torch.randn(shape, generator=generator, dtype=torch.float32).to(device, dtype) + table = torch.arange(blocks, device=device, dtype=torch.int32).reshape(1, blocks) + lengths = torch.tensor([context], device=device, dtype=torch.int32) + q = torch.randn((1, heads, head_dim), generator=generator, dtype=torch.float32) + return key_cache, value_cache, table, lengths, q.to(device, dtype) + + +def timed(function, repeats, device): + """Median wall time of ``repeats`` calls, after warmup.""" + for _ in range(5): + function() + if device.type == "cuda": + torch.cuda.synchronize() + samples = [] + for _ in range(repeats): + start = time.perf_counter() + function() + if device.type == "cuda": + torch.cuda.synchronize() + samples.append(time.perf_counter() - start) + samples.sort() + return samples[len(samples) // 2] * 1e3 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cpu") + parser.add_argument("--context", type=int, default=512) + parser.add_argument("--heads", type=int, default=16) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--kv-heads", type=int, default=4) + parser.add_argument("--layers", type=int, default=24) + parser.add_argument("--loops", type=int, default=4) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--retention", type=float, default=0.1) + parser.add_argument("--repeats", type=int, default=20) + parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"]) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + + device = torch.device(args.device) + dtype = getattr(torch, args.dtype) + key_cache, value_cache, table, lengths, q = build_paged( + args.context, args.heads, args.head_dim, args.kv_heads, args.block_size, device, dtype + ) + block_size = args.block_size + positions = torch.arange(args.context, device=device) + blocks = table[0, positions // block_size] + offsets = positions % block_size + keys = key_cache[blocks, offsets] + values = value_cache[blocks, offsets] + retained = retained_count(args.context, args.retention) + + dense_ms = timed( + lambda: torch_paged_attention(q, key_cache, value_cache, table, lengths), + args.repeats, + device, + ) + capture_ms = timed( + lambda: reference_statistics( + q[0], keys, values, key_retention=args.retention, key_bucket=retained + ), + args.repeats, + device, + ) + output, contribution, mass, selected = reference_statistics( + q[0], keys, values, key_retention=args.retention, key_bucket=retained + ) + groups = args.heads // args.kv_heads + retained_keys = key_cache[ + table[0, selected.long() // block_size], + selected.long() % block_size, + (torch.arange(args.heads, device=device) // groups)[:, None].expand_as(selected), + ] + retained_values = value_cache[ + table[0, selected.long() // block_size], + selected.long() % block_size, + (torch.arange(args.heads, device=device) // groups)[:, None].expand_as(selected), + ] + consume_ms = timed( + lambda: sparse_correction(q[0], retained_keys, retained_values, mass, output, contribution), + args.repeats, + device, + ) + corrected = sparse_correction(q[0], retained_keys, retained_values, mass, output, contribution) + # The reference loop and the later loop hold different KV, so the retained + # set chosen at loop 2 is applied to loop 3 keys. Compare against dense + # attention over that later KV to measure the approximation, not the + # same-loop identity. + later_keys = keys * 1.05 + 0.02 + later_values = values * 0.97 - 0.01 + later_retained_keys = retained_keys * 1.05 + 0.02 + later_retained_values = retained_values * 0.97 - 0.01 + dense_later = reference_statistics(q[0], later_keys, later_values, key_retention=1.0)[0] + dense_reference = reference_statistics(q[0], keys, values, key_retention=1.0)[0] + corrected_later = sparse_correction( + q[0], later_retained_keys, later_retained_values, mass, output, contribution + ) + divergence = float((corrected_later - dense_later).abs().max()) + same_loop_error = float((corrected - dense_reference).abs().max()) + + dense_total = args.layers * args.loops * dense_ms + sparse_total = args.layers * (dense_ms + capture_ms + (args.loops - 2) * consume_ms) + element = 2 if args.dtype == "bfloat16" else 4 + store_bytes = args.layers * args.heads * (retained * 4 + 4 + element * args.head_dim * 2) + kv_bytes = args.layers * 2 * args.kv_heads * args.head_dim * element + + report = dict( + device=str(device), + dtype=args.dtype, + context=args.context, + retained=retained, + retention=args.retention, + layers=args.layers, + heads=args.heads, + head_dim=args.head_dim, + dense_ms=round(dense_ms, 4), + capture_ms=round(capture_ms, 4), + consume_ms=round(consume_ms, 4), + dense_decode_ms=round(dense_total, 3), + sparse_decode_ms=round(sparse_total, 3), + decode_speedup=round(dense_total / sparse_total, 3), + keys_dense=args.context, + keys_sparse=retained, + key_ratio=round(retained / args.context, 4), + reference_state_bytes_per_slot=store_bytes, + kv_bytes_per_token=kv_bytes, + state_over_kv_ratio=round(store_bytes / kv_bytes, 3), + same_loop_identity_error=same_loop_error, + later_loop_max_abs_divergence=divergence, + ) + if args.json: + print(json.dumps(report, indent=2)) + return + print(f"sparse decode attention reference path — {report['device']} {args.dtype}") + print( + f" context={args.context} retained={retained} retention={args.retention} " + f"layers={args.layers} heads={args.heads} head_dim={args.head_dim}" + ) + print(f" dense row {dense_ms:8.4f} ms") + print(f" capture row {capture_ms:8.4f} ms") + print(f" consume row {consume_ms:8.4f} ms (k={retained})") + print( + f" {args.loops}-loop decode: dense {dense_total:8.3f} ms " + f"sparse {sparse_total:8.3f} ms speedup {dense_total / sparse_total:5.2f}x" + ) + print(f" keys touched per row: {args.context} -> {retained} ({retained / args.context:.1%})") + print( + f" reference state {store_bytes / 1024:.1f} KiB/slot vs KV " + f"{kv_bytes / 1024:.1f} KiB/token ({store_bytes / kv_bytes:.2f}x)" + ) + print(f" same-loop identity error (must be ~0): {same_loop_error:.3e}") + print(f" later-loop divergence at this retention: {divergence:.3e}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index 245540b2..2a73ff42 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -12,8 +12,14 @@ import pytest import torch -from tests.helpers import bootstrap_speculation, make_speculative_engine, tiny_ouro_model +from tests.helpers import ( + bootstrap_speculation, + make_cpu_cache, + make_speculative_engine, + tiny_ouro_model, +) from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig +from vllm_rlt.attention import sparse as sparse_module from vllm_rlt.attention.sparse import ( CAPTURE, CONSUME, @@ -295,6 +301,139 @@ def _require_live_batch(self, batch): builder.prepare_device(**padded, reference_plan=plan) +# --------------------------------------------------------------------------- +# End-to-end decode through the model, the paged KV cache and attention +# --------------------------------------------------------------------------- + + +def step_metadata(cache, request_id, position, depth, plan=None): + """Device-style recurrent metadata for one token at one loop depth.""" + allocation = cache._get_allocation(request_id) + page = position // cache.block_size + table = torch.tensor([allocation.block_tables[depth][: page + 1]], dtype=torch.int32) + return cache.metadata_builder.prepare_device( + allocations=((request_id, allocation),), + position_ids=torch.tensor([position], dtype=torch.int64), + write_blocks=table[:, page].long(), + write_offsets=torch.tensor([position % cache.block_size], dtype=torch.int64), + block_tables=table, + context_lengths=torch.tensor([position + 1], dtype=torch.int32), + reference_plan=plan, + ) + + +def decode(model, cache, request_id, tokens, *, store=None, lease=None, retention=1.0): + """Serial greedy decode: prelude, every loop depth, coda, next token.""" + assert cache.allocate(request_id, len(tokens) + 1) + logits = [] + for position, token in enumerate(tokens): + hidden = model.prelude(torch.tensor([token])) + for depth in range(model.config.total_ut_steps): + plan = None + mode = reference_mode(depth) + if store is not None and mode != DENSE: + plan = store.build_plan([(lease, 0)], [mode], key_retention=retention) + metadata = step_metadata(cache, request_id, position, depth, plan) + hidden, _ = model.recurrent_prepared(hidden, metadata, cache, compute_gate=False) + logits.append(model.coda(hidden)) + return logits + + +def make_reference_store(model, **overrides): + values = dict( + slots=1, + layers=model.config.num_hidden_layers, + heads=model.config.num_attention_heads, + head_dim=model.config.head_dim, + key_bucket=model.config.max_position_embeddings, + device="cpu", + ) + values.update(overrides) + return AttentionReferenceStore(**values) + + +def test_sparse_decode_matches_dense_through_model_and_kv(): + # Full retention is the RFC's control: the corrected output must reproduce + # the dense path through the real model, paged KV cache and attention. + model = tiny_ouro_model() + tokens = [2, 3, 4, 5, 7] + dense = decode(model, make_cpu_cache(model.config), "dense", tokens) + store = make_reference_store(model) + lease = store.reserve("sparse", 0, 1, generation=1) + sparse = decode( + model, + make_cpu_cache(model.config), + "sparse", + tokens, + store=store, + lease=lease, + retention=1.0, + ) + assert store.uses.sum() > 0, "the sparse path never consumed a reference" + for expected, actual in zip(dense, sparse): + assert torch.allclose(actual, expected, atol=1e-4), float((actual - expected).abs().max()) + + +def test_sparse_decode_diverges_as_retention_shrinks(): + model = tiny_ouro_model() + tokens = [2, 3, 4, 5, 7, 9, 11] + dense = decode(model, make_cpu_cache(model.config), "dense", tokens) + errors = {} + for retention in (1.0, 0.5, 0.25): + store = make_reference_store(model) + lease = store.reserve("sparse", 0, 1, generation=1) + sparse = decode( + model, + make_cpu_cache(model.config), + "sparse", + tokens, + store=store, + lease=lease, + retention=retention, + ) + errors[retention] = max( + float((actual - expected).abs().max()) for expected, actual in zip(dense, sparse) + ) + assert errors[1.0] < 1e-4 + assert errors[0.25] > errors[1.0] + assert errors[0.5] <= errors[0.25] * 4 + + +def test_sparse_consume_reads_only_retained_keys(monkeypatch): + model = tiny_ouro_model() + tokens = [2, 3, 4, 5] + read_counts = [] + original = sparse_module.gather_retained + + def recording(*args, **kwargs): + keys, values = original(*args, **kwargs) + read_counts.append(keys.shape[1]) + return keys, values + + monkeypatch.setattr(sparse_module, "gather_retained", recording) + store = make_reference_store(model) + lease = store.reserve("sparse", 0, 1, generation=1) + decode( + model, + make_cpu_cache(model.config), + "sparse", + tokens, + store=store, + lease=lease, + retention=0.5, + ) + # Two consuming loops per token, each reading only the retained budget, and + # every layer gathers its own retained keys. + per_token = model.config.num_hidden_layers * 2 + expected = [ + retained_count(position + 1, 0.5) + for position in range(len(tokens)) + for _ in range(per_token) + ] + assert read_counts == expected + assert max(read_counts) < len(tokens) + + # --------------------------------------------------------------------------- # Async speculative integration (GPU) # --------------------------------------------------------------------------- @@ -410,6 +549,35 @@ def recording(owners, modes, **kwargs): assert store.outstanding == 0 +@pytest.mark.gpu +def test_sparse_engine_matches_dense_tokens_at_full_retention(): + # End-to-end through the serving engine: full retention must reproduce the + # dense engine's tokens exactly, including async speculative rounds. + def run(sparse): + model = tiny("cuda") + engine = make_speculative_engine(model, sparse=sparse) + prompts = [[2, 3, 4], [5, 6]] + for i, prompt in enumerate(prompts): + engine.add_request( + str(i), prompt, SamplingParams(max_tokens=12 - 4 * i, ignore_eos=True) + ) + outputs = {} + for _ in range(1000): + if not engine.has_unfinished_requests(): + break + for output in engine.step(): + outputs[output.request_id] = output + tokens = [outputs[str(i)].token_ids for i in range(len(prompts))] + outstanding = engine.reference_store.outstanding if engine.reference_store else 0 + engine.close() + return tokens, outstanding + + dense, _ = run(None) + sparse, outstanding = run(SparseDecodeAttentionConfig(enabled=True, key_retention=1.0)) + assert sparse == dense + assert outstanding == 0 + + @pytest.mark.gpu def test_sparse_engine_sizes_key_buffer_from_retention(): model = tiny("cuda") diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index aa1654a6..80bc9e96 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -195,3 +195,109 @@ def sparse_correction( probabilities = torch.softmax(scores, dim=-1) new_contribution = torch.einsum("hk,hkd->hd", probabilities, values.float()) return output - contribution + mass.float().unsqueeze(-1) * new_contribution + + +def gather_retained( + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_table: torch.Tensor, + selected: torch.Tensor, + block_size: int, + groups: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather each query head's retained keys and values from the paged cache. + + ``selected`` holds absolute key positions per query head, so every head + reads its own retained set. Returns ``[heads, retained, dim]`` tensors. + """ + index = selected.long() + blocks = block_table[index // block_size] + offsets = index % block_size + kv_index = (torch.arange(index.shape[0], device=index.device) // groups)[:, None] + kv_index = kv_index.expand_as(index) + return key_cache[blocks, offsets, kv_index], value_cache[blocks, offsets, kv_index] + + +def sparse_paged_attention( + q: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + plan: SparseDecodePlan, + *, + layer: int, +) -> torch.Tensor: + """Reference sparse decode attention over a paged KV cache (RFC #156 3.2). + + This is the numerical oracle for the sparse decode column, not an optimized + kernel: it walks rows on the host, reads the same paged tensors the dense + path reads, and publishes or consumes the retained state through the plan's + store. Dense rows reproduce the ordinary attention path exactly. + + ``key_cache``/``value_cache`` are one layer of + ``[physical block, token, KV head, dimension]``; ``block_tables`` and + ``context_lengths`` describe the submission rows exactly as the dense + backend receives them. + """ + store = plan.store + if store is None: + raise ValueError("a sparse decode plan needs its reference store") + block_size = key_cache.shape[1] + heads = q.shape[1] + groups = heads // key_cache.shape[2] + scale = 1.0 / math.sqrt(q.shape[-1]) + output = torch.empty_like(q) + captured = consumed = False + for row, length in enumerate(context_lengths.tolist()): + if length == 0: + output[row].zero_() + continue + positions = torch.arange(length, device=q.device) + blocks = block_tables[row, positions // block_size] + offsets = positions % block_size + mode = plan.modes[row] + if mode == CONSUME: + selected, mass, reference, contribution = store.read_reference(plan, layer, row) + keys, values = gather_retained( + key_cache, value_cache, block_tables[row], selected, block_size, groups + ) + output[row] = sparse_correction(q[row], keys, values, mass, reference, contribution).to( + q.dtype + ) + consumed = True + continue + keys = key_cache[blocks, offsets] + values = value_cache[blocks, offsets] + if mode == CAPTURE: + reference, contribution, mass, selected = reference_statistics( + q[row], + keys, + values, + key_retention=plan.key_retention, + key_bucket=store.key_bucket, + ) + output[row] = reference.to(q.dtype) + store.write_reference( + plan, + layer, + row, + selected=selected, + counts=torch.full((heads,), selected.shape[-1], dtype=torch.int32, device=q.device), + mass=mass, + output=reference, + contribution=contribution, + ) + captured = True + continue + # Dense row: the ordinary attention path. + expanded_keys = expand_kv(keys, groups).float() + expanded_values = expand_kv(values, groups).float() + scores = torch.einsum("hd,lhd->hl", q[row].float(), expanded_keys) * scale + probabilities = torch.softmax(scores, dim=-1) + output[row] = torch.einsum("hl,lhd->hd", probabilities, expanded_values).to(q.dtype) + if captured: + store.capture(plan, layer) + if consumed: + store.consume(plan, layer) + return output diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index dfcdf0df..06b6a9d3 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -23,6 +23,7 @@ create_backend, validate_backend_name, ) +from vllm_rlt.attention.sparse import sparse_paged_attention from vllm_rlt.request import ExecutionBatch from vllm_rlt.worker.buffers import stage_metadata from vllm_rlt.worker.metadata import DeviceKVMetadata, KVMetadata, MetadataBuilder @@ -719,6 +720,20 @@ def _attend_prepared( self._require_live_batch(batch) self._validate_tensor(q, len(batch.position_ids), "q", query=True) if isinstance(batch, DeviceKVMetadata): + plan = getattr(batch, "reference_plan", None) + if plan is not None and not plan.dense(): + # Reference sparse decode path: it reads the same paged tensors + # and publishes/consumes retained state through the plan's + # store. Optimized kernels consume the same plan later. + return sparse_paged_attention( + q, + self.key_cache[:, layer], + self.value_cache[:, layer], + batch.block_tables, + batch.context_lengths, + plan, + layer=layer, + ) return self.attention( q, self.key_cache[:, layer], diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index eb0bafbe..cf11e83a 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -240,6 +240,45 @@ def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> Non raise ValueError("reference positions need one entry per plan row") self.position[plan.slot_ids] = positions.to(torch.int64) + def write_reference( + self, + plan: SparseDecodePlan, + layer: int, + row: int, + *, + selected, + counts, + mass, + output, + contribution, + ) -> None: + """Publish the retained state of one capture row (reference path).""" + self._require_live(plan) + slot = plan.slots[row] + count = selected.shape[-1] + capacity = self.selected.shape[-1] + if count > capacity: + raise ValueError(f"reference capacity {capacity} cannot hold {count} retained keys") + self.selected[slot, layer, :, :count] = selected.to(torch.int32) + self.counts[slot, layer] = counts.to(torch.int32) + self.mass[slot, layer] = mass + self.output[slot, layer] = output + self.contribution[slot, layer] = contribution + + def read_reference(self, plan: SparseDecodePlan, layer: int, row: int): + """Borrow the retained state of one consume row (reference path).""" + self._require_live(plan) + slot = plan.slots[row] + count = int(self.counts[slot, layer].max()) + if count == 0: + raise RuntimeError("a consume row has no captured reference") + return ( + self.selected[slot, layer, :, :count].long(), + self.mass[slot, layer], + self.output[slot, layer], + self.contribution[slot, layer], + ) + def accept(self, leases, accepted: torch.Tensor) -> None: """Invalidate the references of rejected rows on the device. From f53c35d3bfa5b2412d4993018b1c07fd4bafd139 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Wed, 28 Oct 2026 21:21:34 +0800 Subject: [PATCH 4/8] benchmarks: engine-level TTFT and ITL comparison for sparse decode attention Prefill rows never take the decode correction (RFC #156 section 6.3), so the feature must not move TTFT; this measures that plus the reference path's inter-token cost against the dense Triton baseline. --- benchmarks/ttft_bench.py | 110 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 benchmarks/ttft_bench.py diff --git a/benchmarks/ttft_bench.py b/benchmarks/ttft_bench.py new file mode 100644 index 00000000..9250710e --- /dev/null +++ b/benchmarks/ttft_bench.py @@ -0,0 +1,110 @@ +"""Engine-level TTFT / ITL / throughput comparison: dense vs sparse decode. + +Prefill rows never take the sparse decode correction (RFC #156 section 6.3), so +enabling the feature must not change prefill work. Any TTFT difference therefore +comes from the first speculative round, which does run the decode path. + +Usage: + python benchmarks/ttft_bench.py --device cuda --prompt 64 --max-tokens 32 + python benchmarks/ttft_bench.py --device cuda --retention 0.5 +""" + +import argparse +import json +import statistics +import time + +import torch + +from tests.helpers import make_speculative_engine, tiny_ouro_model +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig + + +def measure(engine, prompt, max_tokens, repeats): + """TTFT and inter-token latency of sequential single-request decodes.""" + ttfts, itls, totals = [], [], [] + for repetition in range(repeats): + request_id = f"r{repetition}" + start = time.perf_counter() + engine.add_request( + request_id, + prompt, + SamplingParams(max_tokens=max_tokens, ignore_eos=True), + ) + first = None + finished = None + while True: + if not engine.has_unfinished_requests(): + break + for output in engine.step(): + if output.request_id != request_id: + continue + if first is None: + first = time.perf_counter() - start + if output.finished: + finished = time.perf_counter() - start + assert first is not None and finished is not None + ttfts.append(first * 1e3) + totals.append(finished * 1e3) + itls.append((finished - first) * 1e3 / max(1, max_tokens - 1)) + return dict( + ttft_median_ms=round(statistics.median(ttfts), 2), + ttft_min_ms=round(min(ttfts), 2), + itl_median_ms=round(statistics.median(itls), 3), + e2e_median_ms=round(statistics.median(totals), 2), + tokens_per_s=round(max_tokens / (statistics.median(totals) / 1e3), 1), + ) + + +def run(device, sparse, prompt, max_tokens, repeats, seed=123): + torch.manual_seed(seed) + model = tiny_ouro_model(device) + engine = make_speculative_engine(model, sparse=sparse) + # Warmup: compile-free but pays one-time kernel/allocator costs. + measure(engine, prompt, 4, 1) + result = measure(engine, prompt, max_tokens, repeats) + engine.close() + return result + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--prompt", type=int, default=64) + parser.add_argument("--max-tokens", type=int, default=32) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + + prompt = [2 + (index % 60) for index in range(args.prompt)] + cases = { + "dense": None, + "sparse r=1.0": SparseDecodeAttentionConfig(enabled=True, key_retention=1.0), + "sparse r=0.5": SparseDecodeAttentionConfig(enabled=True, key_retention=0.5), + } + report = {} + for name, sparse in cases.items(): + report[name] = run(args.device, sparse, prompt, args.max_tokens, args.repeats) + if args.json: + print(json.dumps(report, indent=2)) + return + header = f"{'case':<14}{'TTFT med':>10}{'TTFT min':>10}{'ITL':>9}{'e2e':>9}{'tok/s':>9}" + print( + f"TTFT / ITL — {args.device}, prompt {args.prompt}, max_tokens {args.max_tokens}, " + f"repeats {args.repeats}" + ) + print(header) + for name, values in report.items(): + print( + f"{name:<14}{values['ttft_median_ms']:>9.2f}ms{values['ttft_min_ms']:>9.2f}ms" + f"{values['itl_median_ms']:>8.3f}ms{values['e2e_median_ms']:>8.1f}ms" + f"{values['tokens_per_s']:>9.1f}" + ) + base = report["dense"]["ttft_median_ms"] + for name in list(report)[1:]: + delta = report[name]["ttft_median_ms"] / base + print(f" TTFT {name} vs dense: {delta:.2f}x") + + +if __name__ == "__main__": + main() From 048b9d7c79a5711857ce3287812502ad8f274ad0 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Wed, 28 Oct 2026 21:51:22 +0800 Subject: [PATCH 5/8] sparse decode attention: memory-bounded retirement and explicit reference capacity Addresses the review's reproducible OOM and capacity findings. Retirement is memory-bounded (blocking finding): - AttentionReferenceStore.accept() no longer reads-modifies-writes the retained index buffer. That indexed multiply materialised a bucket-sized temporary per batch (24 layers x 16 heads x 65,536 keys is tens of GiB), which is the OOM reported at B=64 full retention. Acceptance now clears only the per-row metadata that gates readers (valid, counts, mass) for rejected rows. - _clear() likewise retires metadata only. Wiping the bucket on every release also cost tens of GiB of writes per round, which is part of why full retention measured 6.9 tokens/s. - Readers gate on counts, and a capture overwrites the used prefix, so a stale payload is never observable. - Regression test asserts retirement allocates less than an eighth of the bucket for a 512 MiB bucket (the old code allocated a full copy). Capacity is explicit and accounted before admission: - SparseDecodeAttentionConfig.max_context_tokens declares the supported serving context bound; the index buffer is sized from it instead of always from max_position_embeddings. At 24L/16H, B=32/64 and 512-token serving, the pool drops from 24.09/48.19 GiB to 0.28/0.56 GiB at retention 1.0. - The engine now fails fast with an actionable message when the reference pool cannot fit in free device memory, and logs the resolved budget. - Coverage for overlapping rounds: two submitted rounds of the same request must hold disjoint slots, and retiring the older round must not disturb the newer round's references. --- tests/test_sparse_decode_attention.py | 115 +++++++++++++++++++++++-- vllm_rlt/config.py | 6 ++ vllm_rlt/engine/stage_engine.py | 50 +++++++++-- vllm_rlt/entrypoints/runtime_args.py | 6 ++ vllm_rlt/worker/attention_reference.py | 34 +++++--- 5 files changed, 185 insertions(+), 26 deletions(-) diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index 2a73ff42..e7f85e59 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -174,15 +174,18 @@ def test_reserve_uses_one_slot_per_live_token(): assert len(store.free) == 3 -def test_release_requires_owner_and_clears_payload(): +def test_release_requires_owner_and_retires_metadata(): store = make_store() lease = store.reserve("r", 0, 2, generation=1) fill_payload(store, lease.slots) store.release(lease) for slot in lease.slots: + # Readers are gated by counts/mass/valid. The retained index, output and + # contribution buffers stay allocated and are overwritten by the next + # capture's used prefix, so retiring never writes the whole bucket. + assert int(store.counts[slot].sum()) == 0 assert store.mass[slot].sum() == 0 - assert store.output[slot].sum() == 0 - assert store.selected[slot].sum() == 0 + assert not bool(store.valid[slot].any()) assert store.outstanding == 0 # A retired lease cannot release its slots a second time. with pytest.raises(RuntimeError, match="different owner"): @@ -196,12 +199,15 @@ def test_accept_invalidates_rejected_candidates_and_descendants(): store.accept([lease], torch.tensor([1])) accepted, middle, rejected = lease.slots # Rows 0..1 survive the accepted prefix; the rejected candidate and every - # descendant lose the reference payload. - assert store.mass[accepted].sum() > 0 and store.mass[middle].sum() > 0 + # descendant stop being readable. + assert int(store.counts[accepted].max()) > 0 and int(store.counts[middle].max()) > 0 + assert int(store.counts[rejected].max()) == 0 assert store.mass[rejected].sum() == 0 - assert store.output[rejected].sum() == 0 - assert store.selected[rejected].sum() == 0 + assert not bool(store.valid[rejected].any()) assert int(store.accepted[accepted]) == 1 + rejected_plan = store.build_plan([(lease, 2)], [CONSUME], key_retention=0.5) + with pytest.raises(RuntimeError, match="no captured reference"): + store.read_reference(rejected_plan, layer=0, row=0) def test_accept_keeps_root_when_every_candidate_is_rejected(): @@ -271,6 +277,13 @@ def test_bytes_matches_layout(): assert store.bytes == expected +def test_sparse_config_validates_serving_context_bound(): + assert SparseDecodeAttentionConfig(max_context_tokens=512).max_context_tokens == 512 + assert SparseDecodeAttentionConfig().max_context_tokens is None + with pytest.raises(ValueError, match="max_context_tokens"): + SparseDecodeAttentionConfig(max_context_tokens=0) + + def test_device_metadata_carries_reference_plan(): class StubCache: device = torch.device("cpu") @@ -578,6 +591,94 @@ def run(sparse): assert outstanding == 0 +@pytest.mark.gpu +def test_accept_and_retire_do_not_allocate_bucket_sized_temporaries(): + # Reviewer-reported OOM: acceptance used to read-modify-write the retained + # index buffer, materialising a bucket-sized temporary per batch. Retirement + # must stay bounded by the rejected rows, not by slot count x bucket. + torch.cuda.empty_cache() + store = AttentionReferenceStore( + slots=16, layers=24, heads=16, head_dim=128, key_bucket=32768, device="cuda" + ) + lease = store.reserve("r", 0, 16, generation=1) + store.counts[list(lease.slots)] = 1 + store.mass[list(lease.slots)] = 1.0 + store.output[list(lease.slots)] = 1.0 + store.contribution[list(lease.slots)] = 1.0 + torch.cuda.synchronize() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + store.accept([lease], torch.tensor([0], device="cuda")) + store.release(lease) + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() + bucket_bytes = 16 * 24 * 16 * 32768 * 4 + assert bucket_bytes > 512 * 1024**2, "test needs a bucket large enough to detect the bug" + assert peak - baseline < bucket_bytes // 8, ( + f"retirement allocated {(peak - baseline) / 1024**2:.1f} MiB " + f"for a {bucket_bytes / 1024**2:.1f} MiB bucket" + ) + + +@pytest.mark.gpu +def test_sparse_engine_bounds_reference_pool_by_serving_context(): + model = tiny("cuda") + e = make_speculative_engine( + model, + sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, max_context_tokens=64), + ) + # The pool follows the supported serving bound, not the model's 128-token + # positional limit. + assert e.reference_store.key_bucket == 64 + e.close() + + +@pytest.mark.gpu +def test_sparse_engine_rejects_context_beyond_model_limit(): + model = tiny("cuda") + with pytest.raises(ValueError, match="exceeds the model context"): + make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, max_context_tokens=4096) + ) + + +@pytest.mark.gpu +def test_sparse_async_overlapping_rounds_keep_isolated_references(): + # Reviewer-requested coverage: two submitted rounds of the same request are + # in flight at once. They must hold disjoint reference slots, and retiring + # the older round's rejected candidates must not disturb the newer round. + model = tiny("cuda") + e = make_speculative_engine( + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.5) + ) + store = e.reference_store + bootstrap(e) + accepted_rounds = [] + original = store.accept + + def recording(leases, accepted): + accepted_rounds.append(tuple(lease.slots for lease in leases)) + return original(leases, accepted) + + store.accept = recording + e.step() + e.step() + assert len(e._pending_speculative) == 2, "expected two overlapping rounds" + assert len(accepted_rounds) >= 2 + first, second = accepted_rounds[0][0], accepted_rounds[1][0] + assert first and second + assert not set(first) & set(second), "outstanding rounds share reference slots" + assert store.outstanding == len(first) + len(second) + e._collect_speculative() + # Only the older round retired; the newer round's references survive. + assert store.outstanding == len(second) + e.abort_request("r") + drain(e) + assert store.outstanding == 0 + assert len(store.free) == store.slots + e.close() + + @pytest.mark.gpu def test_sparse_engine_sizes_key_buffer_from_retention(): model = tiny("cuda") diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index ceb80daa..3f926b7c 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -208,6 +208,10 @@ class SparseDecodeAttentionConfig: enabled: bool = False key_retention: float = 1.0 key_bucket: int | None = None + # Supported serving context bound for the reference index buffer. ``None`` + # falls back to the model context; set it to the real serving bound so the + # pool is not sized for a context the deployment never uses. + max_context_tokens: int | None = None def __post_init__(self): if type(self.enabled) is not bool: @@ -216,3 +220,5 @@ def __post_init__(self): raise ValueError("key_retention must be in (0, 1]") if self.key_bucket is not None: _positive("key_bucket", self.key_bucket) + if self.max_context_tokens is not None: + _positive("max_context_tokens", self.max_context_tokens) diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 049431ac..5dfd5941 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -3,6 +3,8 @@ from dataclasses import dataclass, replace from typing import TYPE_CHECKING, cast +import torch + from vllm_rlt.attention import BackendCapabilities, create_backend from vllm_rlt.config import ( CacheConfig, @@ -134,8 +136,16 @@ def __init__( if speculative_config.mode != "block": raise ValueError("sparse decode attention currently requires block speculation") # The index buffer never truncates the algorithm: capacity must - # hold the worst-case retained count over the model context. - required_keys = math.ceil(config.max_position_embeddings * sparse_decode.key_retention) + # hold the worst-case retained count over the supported serving + # context, which the caller may bound explicitly instead of paying + # for the model's full positional limit. + context_bound = sparse_decode.max_context_tokens or config.max_position_embeddings + if context_bound > config.max_position_embeddings: + raise ValueError( + f"max_context_tokens {context_bound} exceeds the model context " + f"{config.max_position_embeddings}" + ) + required_keys = math.ceil(context_bound * sparse_decode.key_retention) capacity = sparse_decode.key_bucket or required_keys if capacity < required_keys: raise ValueError( @@ -145,16 +155,46 @@ def __init__( ) # Two in-flight rounds may each hold every verification row, so the # pool covers both rounds' live tokens with room for slot reuse. + slots = ( + 2 * scheduler_config.max_num_seqs * (speculative_config.num_speculative_tokens + 1) + ) + # Account for the reference pool before admitting any work: it is + # allocated up front, so a budget that cannot fit must fail here + # with an actionable message instead of OOM-ing mid-batch. + reference_bytes = ( + slots + * config.num_hidden_layers + * config.num_attention_heads + * (capacity * 4 + 4 + 4 * 2 * config.head_dim + 8) + ) + if parameter.device.type == "cuda": + free_bytes, _ = torch.cuda.mem_get_info(parameter.device) + if reference_bytes > free_bytes: + raise ValueError( + f"sparse decode attention reference pool needs " + f"{reference_bytes / 1024**3:.1f} GiB but only " + f"{free_bytes / 1024**3:.1f} GiB is free; lower " + f"max_context_tokens, key_retention or num_speculative_tokens" + ) self.reference_store = AttentionReferenceStore( - slots=2 - * scheduler_config.max_num_seqs - * (speculative_config.num_speculative_tokens + 1), + slots=slots, layers=config.num_hidden_layers, heads=config.num_attention_heads, head_dim=config.head_dim, key_bucket=capacity, device=parameter.device, ) + logger.info( + "sparse decode attention: %d reference slots x %d keys x %d layers x %d heads " + "= %.2f GiB (context bound %d, retention %.2f)", + slots, + capacity, + config.num_hidden_layers, + config.num_attention_heads, + self.reference_store.bytes / 1024**3, + context_bound, + sparse_decode.key_retention, + ) scratch += self.reference_store.bytes if speculative_config is not None or self.reference_store is not None: cache_config = replace( diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index 58ae437c..5544d31a 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -65,6 +65,11 @@ def add_runtime_args(parser): default=1.0, help="Retained key fraction per query/head; 1.0 is the full-retention control", ) + parser.add_argument( + "--sparse-decode-max-context", + type=int, + help="Supported serving context bound for the reference index buffer", + ) parser.add_argument( "--cuda-graphs", action="store_true", @@ -149,6 +154,7 @@ def runtime_configs(args): sparse_decode_attention_config=SparseDecodeAttentionConfig( enabled=args.sparse_decode_attention, key_retention=args.sparse_decode_key_retention, + max_context_tokens=args.sparse_decode_max_context, ) if args.sparse_decode_attention else None, diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index cf11e83a..9fad5767 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -288,6 +288,13 @@ def accept(self, leases, accepted: torch.Tensor) -> None: kept, and a rejected candidate invalidates its dependent descendants. The comparison stays on the device so that acceptance never requires a host readback. + + Retirement is memory-bounded: it only clears the per-row metadata that + gates readers (``valid``, ``counts``, ``mass``), never the retained + index buffer. An indexed read-modify-write of ``selected`` would + materialise the whole bucket — at 24 layers, 16 heads and 65,536 keys + that is tens of GiB for one batch — while a zero ``count`` already makes + the payload unreadable, and a later capture overwrites the used prefix. """ if isinstance(leases, ReferenceLease): leases = (leases,) @@ -300,20 +307,22 @@ def accept(self, leases, accepted: torch.Tensor) -> None: [torch.full_like(lease.row_index, lease.request_index) for lease in leases] ) limit = accepted.to(device=self.device, dtype=rows.dtype).index_select(0, request) - keep = rows <= limit - keep_rows = keep.reshape(-1) - keep3 = keep_rows.view(-1, 1, 1) - keep4 = keep_rows.view(-1, 1, 1, 1) self.accepted[slot_ids] = limit.to(torch.int32) - self.valid[slot_ids] = self.valid[slot_ids] & keep_rows.view(-1, 1) - self.mass[slot_ids] = self.mass[slot_ids] * keep3 - self.output[slot_ids] = self.output[slot_ids] * keep4 - self.contribution[slot_ids] = self.contribution[slot_ids] * keep4 - self.selected[slot_ids] = self.selected[slot_ids] * keep4.to(self.selected.dtype) - self.counts[slot_ids] = self.counts[slot_ids] * keep3.to(self.counts.dtype) + rejected = slot_ids[rows > limit] + if rejected.numel(): + # Scalar fills touch only the rejected rows and allocate nothing. + self.valid[rejected] = False + self.counts[rejected] = 0 + self.mass[rejected] = 0 def _clear(self, slots: torch.Tensor) -> None: - """Zero the payload of the given slots, ordered on the owner stream.""" + """Retire the metadata of the given slots, ordered on the owner stream. + + Only the tensors that gate readers are cleared. The retained index, + output and contribution buffers stay allocated and are overwritten by + the next capture's used prefix, so retiring a slot never writes the + whole bucket (which would be tens of GiB per batch at full retention). + """ def zero(): self.valid[slots] = False @@ -322,9 +331,6 @@ def zero(): self.incarnation[slots] = 0 self.position[slots] = 0 self.mass[slots] = 0 - self.output[slots] = 0 - self.contribution[slots] = 0 - self.selected[slots] = 0 self.counts[slots] = 0 if self.stream is not None: From e8f8261a447f778a280e20e25dec7d6ec3b4701f Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 28 Oct 2026 15:00:37 +0100 Subject: [PATCH 6/8] Batch equal-length sparse decode rows and retain scalar oracle --- tests/test_sparse_decode_batched.py | 75 ++++++++++++++++++ vllm_rlt/attention/sparse.py | 115 +++++++++++++++++++++++++++- 2 files changed, 189 insertions(+), 1 deletion(-) create mode 100644 tests/test_sparse_decode_batched.py diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py new file mode 100644 index 00000000..75b6942f --- /dev/null +++ b/tests/test_sparse_decode_batched.py @@ -0,0 +1,75 @@ +"""Batched CUDA arithmetic versus the original row-wise sparse oracle.""" + +import pytest +import torch + +from tests.test_sparse_decode_attention import make_store +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + DENSE, + batched_paged_attention, + reference_paged_attention, +) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("retention", [0.1, 1.0]) +@pytest.mark.parametrize("ties", [False, True]) +def test_batched_ragged_gqa_matches_scalar_oracle(device, dtype, retention, ties): + torch.manual_seed(157) + rows, heads, kv_heads, dim, block_size = 6, 4, 2, 128, 16 + lengths = torch.tensor([0, 17, 17, 31, 31, 257], device=device) + tables = torch.randperm(rows * 17, device=device).reshape(rows, 17).int() + keys = torch.randn(rows * 17, block_size, kv_heads, dim, device=device, dtype=dtype) + values = torch.randn_like(keys) + queries = torch.randn(rows, heads, dim, device=device, dtype=dtype) + if ties: + queries.zero_() + stores = [ + make_store(slots=rows, layers=1, heads=heads, head_dim=dim, key_bucket=257, device=device) + for _ in range(2) + ] + leases = [store.reserve("test", 0, rows, generation=1) for store in stores] + implementations = [reference_paged_attention, batched_paged_attention] + for modes in ([CAPTURE] * rows, [DENSE, CAPTURE, CONSUME, CONSUME, DENSE, CONSUME]): + results = [] + for store, lease, implementation in zip(stores, leases, implementations): + plan = store.build_plan( + [(lease, row) for row in range(rows)], modes, key_retention=retention + ) + results.append(implementation(queries, keys, values, tables, lengths, plan, layer=0)) + torch.testing.assert_close(results[1], results[0], rtol=1e-5, atol=1e-6) + for name in ("selected", "counts", "valid", "uses"): + assert torch.equal(getattr(stores[0], name), getattr(stores[1], name)), name + for name in ("mass", "output", "contribution"): + torch.testing.assert_close( + getattr(stores[0], name), getattr(stores[1], name), rtol=1e-5, atol=1e-6 + ) + queries = queries + 0.25 + keys = keys + 0.125 + values = values - 0.125 + for store, lease in zip(stores, leases): + store.release(lease) + assert store.outstanding == 0 + + +def test_batched_rejects_missing_and_stale_reference(): + store = make_store(slots=1, layers=1, heads=2, head_dim=4) + lease = store.reserve("test", 0, 1, generation=1) + plan = store.build_plan([(lease, 0)], [CONSUME], key_retention=0.1) + args = ( + torch.ones(1, 2, 4), + torch.ones(1, 16, 2, 4), + torch.ones(1, 16, 2, 4), + torch.zeros(1, 1, dtype=torch.int32), + torch.tensor([1]), + plan, + ) + with pytest.raises(RuntimeError, match="no captured reference"): + batched_paged_attention(*args, layer=0) + store.release(lease) + store.reserve("test", 0, 1, generation=2) + with pytest.raises(RuntimeError, match="stale attention reference"): + batched_paged_attention(*args, layer=0) diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index 80bc9e96..7d4b3234 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -218,7 +218,7 @@ def gather_retained( return key_cache[blocks, offsets, kv_index], value_cache[blocks, offsets, kv_index] -def sparse_paged_attention( +def reference_paged_attention( q: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor, @@ -301,3 +301,116 @@ def sparse_paged_attention( if consumed: store.consume(plan, layer) return output + + +def sparse_paged_attention( + q, key_cache, value_cache, block_tables, context_lengths, plan, *, layer +): + """Use bounded row batches on CUDA, retaining the scalar numerical oracle.""" + implementation = batched_paged_attention if q.is_cuda else reference_paged_attention + return implementation( + q, key_cache, value_cache, block_tables, context_lengths, plan, layer=layer + ) + + +def batched_paged_attention( + q, key_cache, value_cache, block_tables, context_lengths, plan, *, layer +): + """Batch equal-length rows without padding or changing retained-key selection. + + Each head still uses the same FP32 matrix products, stable sort and softmax + as the scalar oracle. Grouping exact lengths avoids padded reductions and + bounds gather scratch independently of the reference pool's key capacity. + There is one count readback per layer, rather than one per consume row. + """ + store = plan.store + if store is None: + raise ValueError("a sparse decode plan needs its reference store") + store._require_live(plan) + lengths = context_lengths.tolist() + consume_rows = [r for r, mode in enumerate(plan.modes) if mode == CONSUME and lengths[r]] + counts = {} + if consume_rows: + slots = plan.slot_ids[consume_rows] + counts = dict(zip(consume_rows, store.counts[slots, layer].amax(dim=-1).tolist())) + if any(count == 0 for count in counts.values()): + raise RuntimeError("a consume row has no captured reference") + grouped = {} + for row, length in enumerate(lengths): + if length: + mode = plan.modes[row] + extent = counts[row] if mode == CONSUME else length + grouped.setdefault((mode, extent), []).append(row) + output = torch.zeros_like(q) + heads, dim = q.shape[1:] + groups = heads // key_cache.shape[2] + block_size = key_cache.shape[1] + scale = 1.0 / math.sqrt(dim) + head_index = torch.arange(heads, device=q.device)[None, :, None] + kv_index = head_index // groups + captured = consumed = False + for (mode, extent), all_rows in grouped.items(): + # At most 16M elements per gathered K/V tensor (or one very long row). + chunk = max(1, min(32, 16 * 1024 * 1024 // (extent * heads * dim))) + for start in range(0, len(all_rows), chunk): + rows = all_rows[start : start + chunk] + row_ids = torch.tensor(rows, dtype=torch.int64, device=q.device) + slots = plan.slot_ids[row_ids] + queries = q[row_ids].float() + tables = block_tables[row_ids] + batch_index = torch.arange(len(rows), device=q.device)[:, None, None] + if mode == CONSUME: + columns = torch.arange(extent, device=q.device)[None, None, :] + selected = store.selected[slots[:, None, None], layer, head_index, columns].long() + blocks = tables[batch_index, selected // block_size] + offsets = selected % block_size + keys = key_cache[blocks, offsets, kv_index].float() + values = value_cache[blocks, offsets, kv_index].float() + scores = torch.einsum("bhd,bhkd->bhk", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + new = torch.einsum("bhk,bhkd->bhd", probabilities, values) + corrected = ( + store.output[slots, layer] + - store.contribution[slots, layer] + + store.mass[slots, layer].unsqueeze(-1) * new + ) + output[row_ids] = corrected.to(q.dtype) + consumed = True + continue + positions = torch.arange(extent, device=q.device) + blocks = tables[:, positions // block_size] + offsets = positions % block_size + keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + scores = torch.einsum("bhd,blhd->bhl", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + reference = torch.einsum("bhl,blhd->bhd", probabilities, values) + output[row_ids] = reference.to(q.dtype) + if mode == CAPTURE: + count = retained_count(extent, plan.key_retention) + capacity = store.selected.shape[-1] + if count > capacity: + raise ValueError( + f"reference capacity {capacity} cannot hold {count} retained keys" + ) + order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) + selected = order[..., :count].sort(dim=-1).values + gathered = probabilities.gather(-1, selected) + selected_values = values[batch_index, selected, head_index] + contribution = torch.einsum("bhk,bhkd->bhd", gathered, selected_values) + columns = torch.arange(count, device=q.device)[None, None, :] + # Explicit indices write only the live prefix, never gather the + # entire capacity-sized selected tensor to obtain a prefix view. + store.selected[slots[:, None, None], layer, head_index, columns] = selected.to( + torch.int32 + ) + store.counts[slots, layer] = count + store.mass[slots, layer] = gathered.sum(dim=-1) + store.output[slots, layer] = reference + store.contribution[slots, layer] = contribution + captured = True + if captured: + store.capture(plan, layer) + if consumed: + store.consume(plan, layer) + return output From 90120f2116f1b09994b90b3a8578068d850e9e76 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 28 Oct 2026 15:15:17 +0100 Subject: [PATCH 7/8] Preserve scalar cuBLAS reduction shapes when batching reference rows --- tests/test_sparse_decode_batched.py | 4 ++-- vllm_rlt/attention/sparse.py | 20 +++++++++++++++----- 2 files changed, 17 insertions(+), 7 deletions(-) diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py index 75b6942f..663060ca 100644 --- a/tests/test_sparse_decode_batched.py +++ b/tests/test_sparse_decode_batched.py @@ -40,12 +40,12 @@ def test_batched_ragged_gqa_matches_scalar_oracle(device, dtype, retention, ties [(lease, row) for row in range(rows)], modes, key_retention=retention ) results.append(implementation(queries, keys, values, tables, lengths, plan, layer=0)) - torch.testing.assert_close(results[1], results[0], rtol=1e-5, atol=1e-6) + torch.testing.assert_close(results[1], results[0], rtol=0, atol=0) for name in ("selected", "counts", "valid", "uses"): assert torch.equal(getattr(stores[0], name), getattr(stores[1], name)), name for name in ("mass", "output", "contribution"): torch.testing.assert_close( - getattr(stores[0], name), getattr(stores[1], name), rtol=1e-5, atol=1e-6 + getattr(stores[0], name), getattr(stores[1], name), rtol=0, atol=0 ) queries = queries + 0.25 keys = keys + 0.125 diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index 7d4b3234..acdd37bd 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -313,6 +313,16 @@ def sparse_paged_attention( ) +def _rowwise_product(equation, left, right): + """Keep each row's cuBLAS problem shape equal to the scalar oracle. + + Folding rows into the head dimension can choose a different FP32 reduction + kernel. Its small error can cross a BF16 rounding boundary and change a + greedy token, so batch the gathers/selection but retain row-wise products. + """ + return torch.stack([torch.einsum(equation, a, b) for a, b in zip(left, right)]) + + def batched_paged_attention( q, key_cache, value_cache, block_tables, context_lengths, plan, *, layer ): @@ -366,9 +376,9 @@ def batched_paged_attention( offsets = selected % block_size keys = key_cache[blocks, offsets, kv_index].float() values = value_cache[blocks, offsets, kv_index].float() - scores = torch.einsum("bhd,bhkd->bhk", queries, keys) * scale + scores = _rowwise_product("hd,hkd->hk", queries, keys) * scale probabilities = torch.softmax(scores, dim=-1) - new = torch.einsum("bhk,bhkd->bhd", probabilities, values) + new = _rowwise_product("hk,hkd->hd", probabilities, values) corrected = ( store.output[slots, layer] - store.contribution[slots, layer] @@ -382,9 +392,9 @@ def batched_paged_attention( offsets = positions % block_size keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() - scores = torch.einsum("bhd,blhd->bhl", queries, keys) * scale + scores = _rowwise_product("hd,lhd->hl", queries, keys) * scale probabilities = torch.softmax(scores, dim=-1) - reference = torch.einsum("bhl,blhd->bhd", probabilities, values) + reference = _rowwise_product("hl,lhd->hd", probabilities, values) output[row_ids] = reference.to(q.dtype) if mode == CAPTURE: count = retained_count(extent, plan.key_retention) @@ -397,7 +407,7 @@ def batched_paged_attention( selected = order[..., :count].sort(dim=-1).values gathered = probabilities.gather(-1, selected) selected_values = values[batch_index, selected, head_index] - contribution = torch.einsum("bhk,bhkd->bhd", gathered, selected_values) + contribution = _rowwise_product("hk,hkd->hd", gathered, selected_values) columns = torch.arange(count, device=q.device)[None, None, :] # Explicit indices write only the live prefix, never gather the # entire capacity-sized selected tensor to obtain a prefix view. From 398333dd08fcb41e60d7a616846fb9d5738fa7d4 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 28 Oct 2026 15:30:01 +0100 Subject: [PATCH 8/8] Prepare batched views once before identical per-row bmm calls --- vllm_rlt/attention/sparse.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index acdd37bd..4701efc0 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -320,7 +320,17 @@ def _rowwise_product(equation, left, right): kernel. Its small error can cross a BF16 rounding boundary and change a greedy token, so batch the gathers/selection but retain row-wise products. """ - return torch.stack([torch.einsum(equation, a, b) for a, b in zip(left, right)]) + # Prepare the exact views einsum uses once for the group, instead of + # repeating its permutation/reshape bookkeeping for every row. + if equation == "hd,lhd->hl": + right = right.permute(0, 2, 3, 1) + elif equation == "hl,lhd->hd": + right = right.transpose(1, 2) + elif equation == "hd,hkd->hk": + right = right.transpose(-1, -2) + elif equation != "hk,hkd->hd": + raise ValueError(f"unsupported reference product {equation}") + return torch.stack([torch.bmm(a, b).squeeze(1) for a, b in zip(left.unsqueeze(-2), right)]) def batched_paged_attention(