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 01/28] 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 02/28] 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 03/28] 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 04/28] 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 05/28] 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 10373ab6d63156535b55e4b9be1c7b8ba7319421 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 29 Oct 2026 06:57:09 +0800 Subject: [PATCH 06/28] feat: integrate sparse references into ordinary and synchronous decode --- README.md | 4 + docs/runtime.md | 51 +++++ tests/test_backend_selection.py | 24 ++ tests/test_sparse_decode_sync.py | 295 +++++++++++++++++++++++++ vllm_rlt/attention/backend.py | 7 +- vllm_rlt/core/kv_cache_manager.py | 30 ++- vllm_rlt/engine/stage_engine.py | 51 ++++- vllm_rlt/worker/attention_reference.py | 16 +- vllm_rlt/worker/metadata.py | 1 + vllm_rlt/worker/model_runner.py | 65 ++++++ vllm_rlt/worker/speculative.py | 69 +++++- 11 files changed, 574 insertions(+), 39 deletions(-) create mode 100644 tests/test_sparse_decode_sync.py diff --git a/README.md b/README.md index 4a128539..6e846fa9 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,10 @@ for supported combinations and limitations. - **Depth-aware paged KV cache.** LAST-EXITED and SHARED layouts, automatic CUDA cache sizing, and optional prefix caching, incremental page allocation, priority scheduling, and preemption with CPU state snapshots. +- **Sparse decode reference integration.** Opt-in cached loop-2 attention + correction for ordinary decode and block self-speculation, with shared + per-token reference ownership and explicit capability limits; see the + [configuration and scope](docs/runtime.md#sparse-decode-attention-reference-path). - **Configurable GPU execution.** Triton and FlashAttention backends, asynchronous scheduling, multiple CUDA streams, reusable buffers, and CUDA Graph capture of decode recurrent cores. Asynchronous self-speculation also supports compatible diff --git a/docs/runtime.md b/docs/runtime.md index d98050fc..ab5cbd32 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -228,6 +228,57 @@ The same flags work with the single-device server. A streaming event may contain multiple committed tokens. See [RFC #43](https://github.com/InferMatrixLab/vllm-rlt/issues/43) for the design. +## Sparse decode attention (reference path) + +`SparseDecodeAttentionConfig` enables the RFC #156 cached-reference correction +for generated-token attention. Prompt prefill remains dense. Loop 2 captures +per-token/per-layer/head selection, mass and contributions; loops 3 and 4 reuse +that reference with their current selected KV. Each new input token establishes +its own reference. + +```python +from vllm_rlt import ExecutionConfig, SparseDecodeAttentionConfig, SpeculativeConfig + +sparse = SparseDecodeAttentionConfig( + enabled=True, key_retention=0.1, max_context_tokens=512, +) +# Pass sparse_decode_attention_config=sparse to LLM or StageEngine. +# Ordinary decode uses synchronous scheduling and no speculative_config. +# For synchronous block speculation, also pass: +speculation = SpeculativeConfig(num_speculative_tokens=2, draft_loops=2) +execution = ExecutionConfig(async_scheduling=False, cuda_graphs=False) +``` + +This path supports ordinary synchronous decode and synchronous/asynchronous +block self-speculation with LAST_EXITED KV and refill scheduling. The Torch and +Triton adapters declare sparse-reference support; `auto` selects a compatible +adapter and explicit FlashAttention selection fails rather than silently +replacing the requested backend. Sparse capture/correction itself currently +uses the shared PyTorch reference implementation, even beside Triton dense +attention. It has host readbacks and per-row launch costs; enabling it does not +promise faster serving. CUDA Graphs, token pipelines, mixed scheduling, +preemption/recovery and ordinary asynchronous decode are rejected explicitly. + +The retained count is `min(N, max(1, ceil(N * key_retention)))`, including the +current position among the N causal keys. `key_bucket` controls buffer capacity, +not the algorithm: it must hold the retained count at the supported context +bound. `max_context_tokens=None` uses the model context; a smaller explicit +serving bound reduces pool memory and is enforced before request admission. +The pool holds one slot per live ordinary input token, K+1 per synchronous +block request, and two rounds of K+1 per asynchronous block request. All +reference payload bytes contribute to automatic KV memory planning; explicit +KV block overrides still require the caller to budget total peak memory. + +Ordinary references retire at CODA or request teardown after their last +recurrent reader. Synchronous block references retire on the compute stream +when every verification input has completed its final read; neither accepted +nor rejected rows carry a reference into the next round. With draft depth 1, +verification captures loop 2; with draft depth 2, drafting captures it. +Full retention is a numerical control. Use declared tolerances and investigate +any real-model token differences from dense execution; small-model equality +does not establish equivalence at scale. Retention below 1 is approximate and +needs task-quality evaluation before a deployment recommendation. + ## SHARED prefill semantics SHARED retains one physical KV plane and completes every loop for a prompt diff --git a/tests/test_backend_selection.py b/tests/test_backend_selection.py index e6179d0d..57735bcc 100644 --- a/tests/test_backend_selection.py +++ b/tests/test_backend_selection.py @@ -53,6 +53,30 @@ def test_cpu_auto_reports_torch_and_runs_reference_attention(caplog): torch.testing.assert_close(selected(q, k, v, table, lengths), explicit(q, k, v, table, lengths)) +def test_sparse_auto_filters_flash_and_explicit_flash_fails(packages, monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device: (9, 0)) + required = BackendCapabilities(sparse_decode=True) + selected = create_backend( + "auto", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=required, + ) + assert selected.info["selected_backend"] == "triton" + assert "sparse_decode" in selected.info["selection_reason"] + with pytest.raises(ValueError, match="lacks required capabilities: sparse_decode"): + create_backend( + "flash_attn_4", + torch.device("cuda"), + torch.bfloat16, + 128, + 16, + required_capabilities=required, + ) + + @pytest.mark.parametrize( "sm,dtype,head_dim,block_size,missing,expected", [ diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py new file mode 100644 index 00000000..7785ba9b --- /dev/null +++ b/tests/test_sparse_decode_sync.py @@ -0,0 +1,295 @@ +"""Ordinary and synchronous block decode through the shared sparse contract.""" + +import pytest +import torch + +from tests.helpers import tiny_ouro_model +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.attention import BackendCapabilities, create_backend +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE +from vllm_rlt.engine.stage_engine import StageEngine + + +def engine(*, draft=None, retention=1.0, device="cpu", dtype=torch.float32, **options): + values = dict( + attention_backend="triton" if device == "cuda" else "torch", + cache_config=CacheConfig(512, 2, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=4, max_num_batched_tokens=16, prefill_chunk_size=3 + ), + execution_config=ExecutionConfig(static_buffers=True, pad_to_power_of_two=True), + speculative_config=None if draft is None else SpeculativeConfig(2, draft_loops=draft), + sparse_decode_attention_config=( + None + if retention is None + else SparseDecodeAttentionConfig(enabled=True, key_retention=retention) + ), + ) + values.update(options) + return StageEngine(tiny_ouro_model(device, dtype), **values) + + +def generate(e, count=3): + for i in range(count): + e.add_request( + str(i), [2, 3, 4, 5][: 2 + i % 3], SamplingParams(max_tokens=8, ignore_eos=True) + ) + outputs = {} + for _ in range(1000): + if not e.has_unfinished_requests(): + break + for output in e.step(): + if output.finished: + outputs[output.request_id] = list(output.token_ids) + assert len(outputs) == count + assert all(len(tokens) == 8 for tokens in outputs.values()) + assert e.cache_manager.num_used_blocks == 0 + if e.reference_store is not None: + assert e.reference_store.outstanding == 0 + return outputs + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_full_retention_matches_dense_and_reclaims_slots(draft): + dense = engine(draft=draft, retention=None) + sparse = engine(draft=draft) + try: + assert generate(sparse) == generate(dense) + finally: + sparse.close() + dense.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_plans_follow_tokens_and_keep_prompt_dense(draft, monkeypatch): + e = engine(draft=draft, retention=0.5) + seen = [] + original = e.cache_manager._attend_prepared + + def attend(layer, kv, q): + plan = getattr(kv, "reference_plan", None) + if plan is not None: + assert len(plan.modes) == len(q) + for row, mode in enumerate(plan.modes): + if mode == DENSE: + continue + allocation, depth, position = kv.rows[row] + request_id = next(rid for rid, owner in kv.allocations if owner is allocation) + assert position >= len(e.scheduler.requests[request_id].prompt_token_ids) + assert mode == (CAPTURE if depth == 1 else CONSUME) + assert int(e.reference_store.position[plan.slots[row]]) == position + seen.append((mode, position, plan.slots[row])) + return original(layer, kv, q) + + monkeypatch.setattr(e.cache_manager, "_attend_prepared", attend) + try: + generate(e) + assert {mode for mode, _, _ in seen} == {CAPTURE, CONSUME} + finally: + e.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_cancel_and_request_id_reuse_cannot_reuse_old_reference(draft): + e = engine(draft=draft, retention=0.5) + try: + for _ in range(2): + e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=32, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding or ( + draft is not None and e.speculative_runner.stats.rounds + ): + break + e.abort_request("r") + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + finally: + e.close() + + +def test_context_bound_is_enforced_before_request_admission(): + e = engine( + sparse_decode_attention_config=SparseDecodeAttentionConfig( + enabled=True, max_context_tokens=8 + ) + ) + try: + with pytest.raises(ValueError, match="sparse.*context"): + e.add_request("too-long", [2, 3, 4], SamplingParams(max_tokens=8)) + assert "too-long" not in e.scheduler.requests + finally: + e.close() + + +@pytest.mark.parametrize("setting", ["graphs", "pipeline", "preemption", "shared"]) +def test_unqualified_execution_modes_fail_explicitly(setting): + options = { + "graphs": dict(execution_config=ExecutionConfig(cuda_graphs=True)), + "pipeline": dict(speculative_config=SpeculativeConfig(2, mode="pipeline")), + "preemption": dict(scheduler_config=SchedulerConfig(enable_preemption=True)), + "shared": dict(cache_config=CacheConfig(512, 2, layout="shared")), + }[setting] + with pytest.raises(ValueError, match="sparse decode attention"): + engine(**options) + + +def test_backend_explicitly_declares_sparse_reference_support(): + selected = create_backend( + "auto", + torch.device("cpu"), + torch.float32, + 8, + 2, + required_capabilities=BackendCapabilities(sparse_decode=True), + ) + assert selected.capabilities.sparse_decode + + +@pytest.mark.parametrize("draft", [1, 2]) +@pytest.mark.parametrize("accepted", [0, 1, 2]) +def test_sync_rejection_retires_all_references_before_replacement(draft, accepted, monkeypatch): + e = engine(draft=draft, retention=0.5) + try: + e.add_request("r", [2, 3, 4], SamplingParams(max_tokens=16, ignore_eos=True)) + while not e.step(): + pass + request = e.scheduler.requests["r"] + start = request.position + original = e.speculative_runner._coda + calls = 0 + + def controlled(hidden): + nonlocal calls + logits = torch.full((len(hidden), 64), -100.0) + if calls < 2: + logits[:, 10 + calls] = 100 + else: + for row in range(3): + logits[row, 10 + row if row < accepted else 20 + row] = 100 + calls += 1 + return logits + + seen = [] + build = e.reference_store.build_plan + + def record(*args, **kwargs): + plan = build(*args, **kwargs) + seen.append(plan) + return plan + + monkeypatch.setattr(e.reference_store, "build_plan", record) + monkeypatch.setattr(e.speculative_runner, "_coda", controlled) + output = e.step()[0] + expected = [10 + row if row < accepted else 20 + row for row in range(accepted + 1)] + assert output.token_ids[-(accepted + 1) :] == expected + assert request.position == start + accepted + 1 + assert e.reference_store.outstanding == 0 + old = seen[0] + with pytest.raises(RuntimeError, match="stale"): + e.reference_store.read_reference(old, 0, 0) + monkeypatch.setattr(e.speculative_runner, "_coda", original) + e.step() + assert seen[-1].incarnations[0] > old.incarnations[0] + assert e.speculative_runner.stats.accepted_tokens >= accepted + finally: + e.close() + + +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_failed_attention_releases_lease_and_allows_fresh_request(draft, monkeypatch): + e = engine(draft=draft, retention=0.5) + original = e.cache_manager._attend_prepared + + def fail(layer, kv, q): + if getattr(kv, "reference_plan", None) is not None: + raise RuntimeError("injected attention failure") + return original(layer, kv, q) + + monkeypatch.setattr(e.cache_manager, "_attend_prepared", fail) + try: + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + with pytest.raises(RuntimeError, match="injected"): + for _ in range(100): + e.step() + assert e.reference_store.outstanding == 0 + assert "r" not in e.scheduler.requests + assert e.cache_manager.num_used_blocks == 0 + monkeypatch.setattr(e.cache_manager, "_attend_prepared", original) + generate(e, count=1) + finally: + e.close() + + +@pytest.mark.gpu +@pytest.mark.parametrize("draft", [None, 1, 2]) +def test_cuda_sparse_sync_executes_and_drains(draft): + e = engine(draft=draft, retention=0.5, device="cuda", dtype=torch.bfloat16) + try: + generate(e) + torch.cuda.synchronize() + assert not bool(e.reference_store.counts.any()) + finally: + e.close() + + +@pytest.mark.gpu +def test_cuda_ordinary_retirement_survives_boundary_allocator_reuse(monkeypatch): + e = engine(device="cuda", dtype=torch.bfloat16) + try: + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding: + break + store = e.reference_store + assert store.outstanding == 1 + slot = e.model_runner.attention_references["r"][2].slots[0] + assert slot != 0 + torch.cuda.synchronize() + store.counts[0] = 77 + torch.cuda.synchronize() + clear = store._clear + pointers = [] + + def record(indices): + assert torch.cuda.current_stream() == e.model_runner.core_stream + pointers.append(indices.data_ptr()) + clear(indices) + + monkeypatch.setattr(store, "_clear", record) + # Hold core while boundary allocations try to reuse any released index + # storage. Correct stream ownership prevents that reuse before clearing. + with torch.cuda.stream(e.model_runner.core_stream): + torch.cuda._sleep(20_000_000) + with torch.cuda.stream(e.model_runner.boundary_stream): + e.model_runner._release_attention_reference("r") + poison = [torch.zeros(1, device="cuda", dtype=torch.int64) for _ in range(32)] + torch.cuda.synchronize() + assert pointers and len(poison) == 32 + assert not bool(store.counts[slot].any()) + assert bool((store.counts[0] == 77).all()) + finally: + e.close() + + +@pytest.mark.gpu +def test_cuda_close_drains_retirement_of_an_active_reference(): + e = engine(device="cuda", dtype=torch.bfloat16) + e.add_request("r", [2, 3], SamplingParams(max_tokens=16, ignore_eos=True)) + for _ in range(100): + e.step() + if e.reference_store.outstanding: + break + assert e.reference_store.outstanding == 1 + e.close() + assert e.reference_store.outstanding == 0 + assert e.cache_manager.num_used_blocks == 0 + assert e.model_runner.core_stream.query() diff --git a/vllm_rlt/attention/backend.py b/vllm_rlt/attention/backend.py index 210e766e..045c1a23 100644 --- a/vllm_rlt/attention/backend.py +++ b/vllm_rlt/attention/backend.py @@ -23,6 +23,8 @@ class BackendCapabilities: packed_matches_decode: bool = False # Consume GPU-owned context lengths without a host read between submissions. device_context_lengths: bool = False + # Shared paged sparse-reference capture/correction dispatch is available. + sparse_decode: bool = False @dataclass @@ -45,9 +47,9 @@ def generation(self): return getattr(self.kernel, "generation", None) -_TORCH_CAPABILITIES = BackendCapabilities() +_TORCH_CAPABILITIES = BackendCapabilities(sparse_decode=True) _TRITON_CAPABILITIES = BackendCapabilities( - cuda_graphs=True, async_scheduling=True, device_context_lengths=True + cuda_graphs=True, async_scheduling=True, device_context_lengths=True, sparse_decode=True ) _FLASH_CAPABILITIES = { 2: BackendCapabilities(cuda_graphs=True, async_scheduling=True, device_context_lengths=True), @@ -128,6 +130,7 @@ def create_backend( "cuda_graphs", "async_scheduling", "device_context_lengths", + "sparse_decode", ) if getattr(required_capabilities, name) and not getattr(implementation.capabilities, name) diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index 06b6a9d3..e073414a 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -719,21 +719,21 @@ def _attend_prepared( self._validate_layer(layer) self._require_live_batch(batch) self._validate_tensor(q, len(batch.position_ids), "q", query=True) + if not isinstance(batch, DeviceKVMetadata): + for allocation, depth, position in batch.rows: + self._require_prefix(allocation, layer, depth, position + 1) + plan = batch.reference_plan + if plan is not None and not plan.dense(): + return sparse_paged_attention( + q, + self.key_cache[:, layer], + self.value_cache[:, layer], + batch.block_tables, + batch.context_lengths, + plan, + layer=layer, + ) 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], @@ -743,8 +743,6 @@ def _attend_prepared( ) if not batch.rows: return torch.empty_like(q) - for allocation, depth, position in batch.rows: - self._require_prefix(allocation, layer, depth, position + 1) if batch.cu_seqlens_q is not None: return self.attention.prefill( q, diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 5dfd5941..f1c5b68f 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -92,6 +92,7 @@ def __init__( speculative_config and speculative_config.mode == "pipeline" ) self.async_speculative = bool(speculative_config and self.execution_config.async_scheduling) + scratch = 0 if speculative_config is not None: if cache_config.layout != "last_exited": raise ValueError("speculative decoding requires last_exited KV") @@ -127,14 +128,22 @@ def __init__( 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: + if getattr(model, "requires_boundary_kv", False): + raise ValueError("sparse decode attention requires a recurrent Ouro-style model") + if cache_config.layout != "last_exited": + raise ValueError("sparse decode attention requires last_exited KV") + if self.execution_config.cuda_graphs: + raise ValueError("sparse decode attention does not yet support CUDA graphs") + if speculative_config is not None and speculative_config.mode != "block": + raise ValueError("sparse decode attention currently requires block speculation") + if scheduler_config.enable_preemption or scheduler_config.mode != "refill": raise ValueError( - "sparse decode attention currently requires asynchronous self-speculation" + "sparse decode attention requires refill scheduling without preemption" + ) + if self.execution_config.async_scheduling and not self.async_speculative: + raise ValueError( + "sparse decode attention ordinary decode requires synchronous scheduling" ) - 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 supported serving # context, which the caller may bound explicitly instead of paying @@ -153,11 +162,14 @@ def __init__( 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. - slots = ( - 2 * scheduler_config.max_num_seqs * (speculative_config.num_speculative_tokens + 1) - ) + # Pool capacity follows simultaneously live inputs, not requests + # alone. Synchronous block references retire in the same round; + # asynchronous speculation can retain two rounds per request. + slots = scheduler_config.max_num_seqs + if speculative_config is not None: + slots *= speculative_config.num_speculative_tokens + 1 + if self.async_speculative: + slots *= 2 # 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. @@ -246,6 +258,7 @@ def __init__( parameter.device.type == "cuda" and self.execution_config.async_scheduling ), device_context_lengths=self.async_speculative, + sparse_decode=self.reference_store is not None, ), ) num_blocks, self.memory_plan = plan_cache( @@ -296,6 +309,8 @@ def __init__( scheduler_config=scheduler_config, staged_prefill=self.scheduler.staged_prefill, rng=self.rng, + reference_store=self.reference_store if speculative_config is None else None, + sparse_decode=self.sparse_decode_attention, ) self.speculative_runner = ( SpeculativeRunner( @@ -306,6 +321,8 @@ def __init__( rng=self.rng, runtime=self.model_runner.runtime, compute_stream=self.model_runner.core_stream, + reference_store=self.reference_store, + sparse_decode=self.sparse_decode_attention, ) if speculative_config is not None and not self.async_speculative else None @@ -330,6 +347,8 @@ def __init__( # 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 + if self.reference_store is not None: + 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 = {} @@ -439,6 +458,12 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i capacity = len(prompt_token_ids) + params.max_tokens - 1 if capacity > config.max_position_embeddings: raise ValueError("prompt plus decode positions exceed the model context length") + if self.reference_store is not None: + bound = ( + self.sparse_decode_attention.max_context_tokens or config.max_position_embeddings + ) + if capacity > bound: + raise ValueError("request exceeds sparse decode attention context bound") cache = self.cache_manager required = cache.required_blocks(capacity) if required > cache.num_blocks: @@ -485,6 +510,10 @@ def close(self): self.abort_request(rid) if self.async_speculative: self._discard_speculative() + if self.reference_store is not None: + # Abort can enqueue reference metadata retirement after the drain + # above. Teardown must also cover those final writes. + self.model_runner.synchronize() def _discard_speculative(self): """Release speculative owners after the caller has drained both streams.""" diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index 9fad5767..9a47c214 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -182,7 +182,14 @@ def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan if len(owners) != len(modes): raise ValueError("a reference plan needs one mode per row") slots, incarnations = [], [] - for lease, row in owners: + for owner, mode in zip(owners, modes): + if owner is None: + if mode != DENSE: + raise ValueError("a sparse row needs a reference owner") + slots.append(-1) + incarnations.append(0) + continue + lease, row = owner slot = lease.slots[row] binding = self._bindings.get(slot) if binding is None or binding != ( @@ -207,7 +214,9 @@ 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): + for slot, incarnation, mode in zip(plan.slots, plan.incarnations, plan.modes): + if slot == -1 and mode == DENSE: + continue binding = self._bindings.get(slot) if binding is None or binding[2] != incarnation: raise RuntimeError("stale attention reference plan for a retired token") @@ -238,7 +247,8 @@ def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> Non 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) + rows = [row for row, slot in enumerate(plan.slots) if slot >= 0] + self.position[plan.slot_ids[rows]] = positions[rows].to(torch.int64) def write_reference( self, diff --git a/vllm_rlt/worker/metadata.py b/vllm_rlt/worker/metadata.py index 4c7e7cc0..03880166 100644 --- a/vllm_rlt/worker/metadata.py +++ b/vllm_rlt/worker/metadata.py @@ -39,6 +39,7 @@ class KVMetadata: writable: bool cu_seqlens_q: torch.Tensor | None = None max_seqlen_q: int = 1 + reference_plan: object | None = None @dataclass(frozen=True, eq=False) diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ada23d20..77144bf4 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -7,6 +7,7 @@ import torch +from vllm_rlt.attention.sparse import CAPTURE, DENSE, reference_mode from vllm_rlt.config import ExecutionConfig, ExitConfig, SchedulerConfig from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import ExecutionBatch, Operation, Request @@ -160,9 +161,15 @@ def __init__( scheduler_config=None, staged_prefill=None, rng: RngRegistry | None = None, + reference_store=None, + sparse_decode=None, ): self.model = model.eval() self.cache_manager = cache_manager + self.reference_store = reference_store + self.sparse_decode = sparse_decode + self.attention_references = {} + self.reference_generation = 0 self.metadata = cache_manager.metadata_builder self.runtime = ModelRuntime(self.model, cache_manager) parameter = next(model.parameters()) @@ -1016,6 +1023,8 @@ def _prepare_forward(self, prepared): return replace(prepared, forwards=(PreparedForward(self.runtime.prepare(forward)),)) hidden = self._load_hidden(batch, prepared, workspace, size) if batch.operation == Operation.CODA: + for request in requests: + self._release_attention_reference(request.request_id) forward = ForwardBatch( Operation.CODA, hidden=hidden[: batch.num_rows] if kv_metadata else hidden, @@ -1027,6 +1036,8 @@ def _prepare_forward(self, prepared): kv = prepared.kv if kv is not None: metadata = (kv,) + metadata + if self.reference_store is not None: + metadata = tuple(self._reference_metadata(batch, kv) for kv in metadata) forward = ForwardBatch( Operation.RECURRENT, hidden=hidden[: batch.num_rows] @@ -1040,6 +1051,57 @@ def _prepare_forward(self, prepared): prepared, kv=metadata[0], forwards=(PreparedForward(self.runtime.prepare(forward)),) ) + def _release_attention_reference(self, request_id): + state = self.attention_references.pop(request_id, None) + if state is not None: + # release() creates the index tensor used by the metadata clears. + # Allocate it on core too, so its allocator lifetime follows the + # clear kernels even when CODA or teardown calls from another stream. + with ( + torch.cuda.stream(self.core_stream) + if self.core_stream is not None + else nullcontext() + ): + self.reference_store.release(state[2]) + + def _reference_metadata(self, batch, kv): + """Borrow token-owned references; prompt and padded rows remain dense. + + Preparation runs on core, after the request's preceding event. CODA + retires the reference after its final recurrent read; request teardown + waits for that same event before releasing an unfinished token. + """ + requests = [span.request for span in batch.spans for _ in range(span.token_count)] + owners, modes = [], [] + for request, (_, depth, position) in zip(requests, kv.rows): + mode = reference_mode(depth) if position >= len(request.prompt_token_ids) else DENSE + state = self.attention_references.get(request.request_id) + if state is not None and (state[0] is not request or state[1] != position): + self._release_attention_reference(request.request_id) + state = None + if mode == CAPTURE: + if state is None: + self.reference_generation += 1 + lease = self.reference_store.reserve( + request.request_id, 0, 1, self.reference_generation + ) + state = (request, position, lease) + self.attention_references[request.request_id] = state + if state is None: + mode = DENSE # No reference loop for this token: stay dense. + owners.append(None if mode == DENSE else (state[2], 0)) + modes.append(mode) + if all(mode == DENSE for mode in modes): + return kv + padding = len(kv.position_ids) - len(owners) + owners.extend([None] * padding) + modes.extend([DENSE] * padding) + plan = self.reference_store.build_plan( + owners, modes, key_retention=self.sparse_decode.key_retention + ) + self.reference_store.bind_positions(plan, kv.position_ids) + return replace(kv, reference_plan=plan) + def _load_hidden(self, batch, prepared, workspace, size): if prepared is not None and prepared.batch_inputs is not None: return prepared.batch_inputs.hidden @@ -1269,6 +1331,8 @@ def submit(self, scheduled_batch: SchedulerOutput | None = None, *, batch: Execu ) except Exception: self._synchronize_submission(stream, prepared) + for span in batch.spans: + self._release_attention_reference(span.request.request_id) # A failed preparation may advance a pool without enqueueing a batch. # Retire all prior readers before any pool can wrap after that gap. self.submissions.drain() @@ -1364,6 +1428,7 @@ def _free_runner_slot(self, request_id): event = self.events.pop(request_id, None) if event is not None: event.synchronize() + self._release_attention_reference(request_id) self.prompt_token_owners.pop(request_id, None) if self.async_state is not None: self.async_state.release(request_id) diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 85f0e235..5828e11f 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -9,6 +9,7 @@ import torch +from vllm_rlt.attention.sparse import DENSE, reference_mode from vllm_rlt.core.scheduler import SchedulerOutput from vllm_rlt.request import Operation, Request from vllm_rlt.worker.buffers import stage_metadata @@ -152,6 +153,8 @@ def __init__( *, runtime=None, compute_stream=None, + reference_store=None, + sparse_decode=None, ): self.model, self.cache, self.config = model, cache, config self.metadata = cache.metadata_builder @@ -159,6 +162,9 @@ def __init__( self.pipeline_states: dict[str, PipelineState] = {} self.device = next(model.parameters()).device self.compute_stream = compute_stream + self.reference_store = reference_store + self.sparse_decode = sparse_decode + self.reference_generation = 0 self.submissions = SubmissionQueue() self.input_buffers = { dtype: StagedInputTensor(dtype, self.device, use_uva=execution.use_uva) @@ -175,9 +181,14 @@ def __init__( def _packed(self, packed): capabilities = self.cache.attention_capabilities - return packed and capabilities.packed_prefill and capabilities.packed_matches_decode + return ( + packed + and self.reference_store is None + and capabilities.packed_prefill + and capabilities.packed_matches_decode + ) - def _prepare_inputs(self, batch): + def _prepare_inputs(self, batch, references=()): """Plan all host-known inputs, then publish one snapshot per dtype.""" for buffer in self.input_buffers.values(): buffer.reset() @@ -219,8 +230,32 @@ def metadata(ids, positions, depths, *, packed=False): for active, layouts in steps ] verification = tuple(self.metadata.bind(kv, uploaded) for kv in verification) + if references: + steps = [ + ( + active, + self._reference_metadata(layouts, [(references[i], offset) for i in active]), + ) + for offset, (active, layouts) in enumerate(steps) + ] + owners = [(lease, row) for lease in references for row in range(len(lease))] + verification = self._reference_metadata(verification, owners) return tokens, steps, verification + def _reference_metadata(self, layouts, owners): + result = [] + for kv in layouts: + modes = [reference_mode(depth) for _, depth, _ in kv.rows] + if all(mode == DENSE for mode in modes): + result.append(kv) + continue + plan = self.reference_store.build_plan( + owners, modes, key_retention=self.sparse_decode.key_retention + ) + self.reference_store.bind_positions(plan, kv.position_ids) + result.append(replace(kv, reference_plan=plan)) + return tuple(result) + def _recurrent(self, hidden, metadata): forward = ForwardBatch( Operation.RECURRENT, hidden=hidden, metadata=metadata, graphs=self.graphs @@ -257,11 +292,31 @@ def submit(self, batch): if self.compute_stream is not None else nullcontext() ): - if self.config.mode == "pipeline": - submission = self._submit_pipeline(batch) - else: - inputs, steps, verification = self._prepare_inputs(batch) - submission = self._submit_round(batch, inputs, steps, verification) + references = [] + try: + if self.reference_store is not None: + self.reference_generation += 1 + for index, span in enumerate(batch.requests): + references.append( + self.reference_store.reserve( + span.request.request_id, + index, + span.token_count, + self.reference_generation, + ) + ) + if self.config.mode == "pipeline": + submission = self._submit_pipeline(batch) + else: + inputs, steps, verification = self._prepare_inputs(batch, references) + submission = self._submit_round(batch, inputs, steps, verification) + finally: + if references: + # Every candidate has finished its final recurrent read. + # No reference survives this block round, including the + # rejected suffix. Clear/reuse is ordered on core after + # those readers, without waiting for host acceptance. + self.reference_store.release(references) self.submissions.complete(submission.event) return submission except Exception: From 985c309212191c7ab32bada973a9b5af806494a6 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 29 Oct 2026 07:06:17 +0800 Subject: [PATCH 07/28] test: compare sparse verification with serial hidden logits and KV --- tests/test_sparse_decode_sync.py | 74 ++++++++++++++++++++++++++++++++ 1 file changed, 74 insertions(+) diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py index 7785ba9b..bd2f2bd7 100644 --- a/tests/test_sparse_decode_sync.py +++ b/tests/test_sparse_decode_sync.py @@ -1,5 +1,7 @@ """Ordinary and synchronous block decode through the shared sparse contract.""" +from dataclasses import replace + import pytest import torch @@ -153,6 +155,78 @@ def test_backend_explicitly_declares_sparse_reference_support(): assert selected.capabilities.sparse_decode +@pytest.mark.parametrize("draft", [1, 2]) +def test_sparse_verification_matches_serial_same_input_hidden_logits_and_kv(draft): + from tests.helpers import make_cpu_cache + from tests.test_sparse_decode_attention import make_reference_store + from vllm_rlt.attention.sparse import reference_mode + from vllm_rlt.core.scheduler import ScheduledRequest, SchedulerOutput + from vllm_rlt.request import Request, RequestPhase + + e = engine(draft=draft, retention=0.5) + model, cache = e.model, e.cache_manager + prefix, current = [3, 4, 5], 7 + assert cache.allocate("r", 16) + for position, token in enumerate(prefix): + hidden = model.prelude(torch.tensor([token])) + for depth in range(4): + hidden, _ = model.recurrent(hidden, ["r"], [depth], [position], cache) + request = Request("r", prefix, SamplingParams(ignore_eos=True), generated_token_ids=[current]) + seen, heads = {}, [] + recurrent = model.recurrent_prepared + + def record(hidden, kv, cache, **kwargs): + out, gate = recurrent(hidden, kv, cache, **kwargs) + for row, (_, depth, position) in enumerate(kv.rows): + seen[depth, position] = out[row].clone() + return out, gate + + model.recurrent_prepared = record + hook = model.lm_head.register_forward_hook( + lambda module, inputs, output: heads.append(output.detach().clone()) + ) + try: + e.speculative_runner.submit( + SchedulerOutput(RequestPhase.SPECULATIVE, [ScheduledRequest(request, 3, 3)]) + ).collect() + finally: + hook.remove() + model.recurrent_prepared = recurrent + assert e.reference_store.outstanding == 0 + tokens = prefix + [current] + tokens.extend(int(model.coda(seen[draft - 1, pos]).argmax()) for pos in (3, 4)) + oracle = make_cpu_cache(model.config) + assert oracle.allocate("r", 16) + store = make_reference_store(model) + for position, token in enumerate(tokens): + lease = store.reserve("r", 0, 1, generation=position) + hidden = model.prelude(torch.tensor([token])) + for depth in range(4): + kv = oracle._prepare_batch(["r"], [depth], [position]) + mode = reference_mode(depth) + if position >= len(prefix) and mode != DENSE: + plan = store.build_plan([(lease, 0)], [mode], key_retention=0.5) + kv = replace(kv, reference_plan=plan) + hidden, _ = model.recurrent_prepared(hidden, kv, oracle, compute_gate=False) + if position >= len(prefix): + torch.testing.assert_close(seen[depth, position], hidden[0], atol=4e-5, rtol=4e-5) + if position >= len(prefix): + torch.testing.assert_close( + heads[-1][position - len(prefix)], model.coda(hidden)[0], atol=4e-5, rtol=4e-5 + ) + store.release(lease) + for depth in range(4): + for layer in range(model.config.num_hidden_layers): + for actual, expected in zip( + cache.read(layer, "r", depth, len(tokens)), + oracle.read(layer, "r", depth, len(tokens)), + ): + torch.testing.assert_close(actual, expected, atol=4e-5, rtol=4e-5) + cache.free("r") + oracle.free("r") + e.close() + + @pytest.mark.parametrize("draft", [1, 2]) @pytest.mark.parametrize("accepted", [0, 1, 2]) def test_sync_rejection_retires_all_references_before_replacement(draft, accepted, monkeypatch): From 7f22af3901416096ea6c1c0e0a5536f440f662f5 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 29 Oct 2026 07:24:41 +0800 Subject: [PATCH 08/28] docs: record sparse decode integration qualification and limits --- .../validation/sparse-decode-sync-20261029.md | 91 +++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 docs/validation/sparse-decode-sync-20261029.md diff --git a/docs/validation/sparse-decode-sync-20261029.md b/docs/validation/sparse-decode-sync-20261029.md new file mode 100644 index 00000000..4685b51e --- /dev/null +++ b/docs/validation/sparse-decode-sync-20261029.md @@ -0,0 +1,91 @@ +# Sparse decode: ordinary and synchronous integration validation + +This contribution extends the shared [#157](https://github.com/InferMatrixLab/vllm-rlt/pull/157) reference plan/store into ordinary synchronous decode and synchronous block self-speculation. It is a functional reference integration, with measured costs and an unresolved BF16 full-retention sequence difference. It does not claim an optimized sparse kernel or deployment quality. + +## Source and checks + +- Immediate base: `048b9d7c79a5711857ce3287812502ad8f274ad0` on `kunkun/rfc156-sparse-decode-async`. +- Measured runtime: `10373ab6d63156535b55e4b9be1c7b8ba7319421`. Additional same-input tests: `985c309`; runtime files are identical between these revisions. Source SHA256 inventories are in the archive. +- Full pre-commit: Ruff lint, Ruff format and CPU regression passed; **919 passed, 475 skipped**. +- Applicable GPU regression: **312 passed, 1 skipped, 16 deselected**. The skip needs a second GPU; FA4 cases are excluded on SM89. The original unfiltered run's 12 unsupported-FA4 failures are retained. +- New integration module: 26 CPU tests plus 5 CUDA cases. Includes ordinary/padded routing, prompt exclusion, draft depth 1/2, each rejection boundary, same-position replacement, abort/ID reuse, failed submissions, stream-owned retirement and active-request close. Per-depth hidden, logits and populated KV match serial same-input sparse execution at FP32 atol/rtol `4e-5`. +- The negative stream control restores the unsafe caller-stream release only inside one test process and fails as expected. Corrected production code passes. +- Independent review found release-index stream ownership and shutdown-drain defects; both were fixed and covered. Re-review found no further concrete production defect. +- Upstream FlashLoop `aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6`, `dense_source_selector` plus `cached_mass_sparse_delta(force_torch=True)`: **18 CUDA fixed-input cases passed**, with equal retained sets; lengths 1/13/128, retention .1/.5/1, FP32 and BF16. Predeclared attention-output atol/rtol were `2e-5` and `.02`, respectively, accommodating the gold path's BF16 subset/output storage. This is separate from model-level quality. + +## Real-model protocol + +One RTX4090, 24,564 MiB; driver 580.105.08; torch 2.12.1+cu130, Triton 3.7.1. Released Ouro-1.4B revision `574fa66cb8bf5abdc979642d01cf2b79b16bfab1`, BF16 weights/activations/KV; weight SHA256 `58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. + +Native StageEngine, LAST_EXITED KV, refill, Triton dense backend, synchronous scheduling, multistream on, graphs off. Block K=2, draft depths 1/2, fixed four loops, greedy and ignore EOS. Each fresh process warms 4 output tokens at its target arrival concurrency, then measures three 128-input/16-output bursts. Request i uses `[1 + (j * 137 + i * 53) % 49000 for j in range(128)]`. Context bound is explicitly **512** for sparse, while model positional configuration remains unchanged. Token budget `max(128, 3*C)`, prefill chunk 128, KV blocks `40*C + 32`, block size 16. This does not qualify the model's 65,536 context limit. + +Timing includes enqueue, prefill, decode, delivery, drain and CUDA synchronization; model loading and warmup are excluded. Profiles are separate and excluded. GPU clocks and CPU affinity were not pinned for the primary sweep. All raw repeats, outputs, depths, acceptance counters, TTFT/E2E observations and memory values are archived. Speculative output arrives in chunks; chunk gaps are not per-token ITL. The short bursts do not establish steady-state HTTP serving. + +**81 primary bursts / 1,107 completed requests / 17,712 output tokens**; no primary OOM, timeout or failed request. Every measured run drained KV blocks and reference slots to zero and repeated its token/depth sequences exactly. + +Output throughput: median of three runs, with min-max in parentheses. Full-retention agreement counts one repeat, and was identical across all three. + +| Mode | Concurrent arrivals | Dense tok/s | Full retention tok/s | r=.1 tok/s | Full vs dense tokens | +|---|---:|---:|---:|---:|---:| +| Ordinary | 1 | 14.90 (14.67-14.98) | 9.66 (8.89-9.79) | 9.30 (8.86-9.37) | 16/16 | +| Ordinary | 8 | 71.60 (71.59-71.60) | 21.59 (21.53-21.62) | 21.73 (21.59-21.77) | 128/128 | +| Ordinary | 32 | 124.42 (124.41-124.49) | 27.13 (27.01-27.31) | 26.84 (26.52-27.24) | 512/512 | +| Sync, draft 1 | 1 | 17.46 (17.21-17.68) | 9.10 (8.63-9.12) | 13.90 (12.72-14.01) | 16/16 | +| Sync, draft 1 | 8 | 60.53 (57.33-61.12) | 18.01 (17.83-19.32) | 25.54 (25.50-25.64) | 128/128 | +| Sync, draft 1 | 32 | 82.91 (81.09-83.10) | 19.58 (19.54-19.81) | 27.48 (27.05-27.57) | 498/512 | +| Sync, draft 2 | 1 | 16.56 (15.37-16.69) | 9.49 (9.06-9.70) | 11.87 (11.13-11.91) | 16/16 | +| Sync, draft 2 | 8 | 50.27 (50.16-50.57) | 18.55 (17.54-18.58) | 22.40 (22.23-22.42) | 128/128 | +| Sync, draft 2 | 32 | 68.08 (67.59-68.30) | 20.31 (20.20-20.38) | 24.84 (24.70-24.90) | 512/512 | + +The shared PyTorch reference path is slower than dense Triton in every covered condition. Peak allocated memory reached **6.67 GiB**. Retention .1 changes outputs, and its speculation acceptance changes are preserved in `summary.json`; token agreement is not a task-quality score. + +## Full-retention discrepancy + +At C=32, draft depth 1, full retention matches **498/512** dense output tokens. The first divergence is request 9, output index 2 (zero-based), input position 129: dense chooses 293, full retention chooses 92. Other full-retention conditions in this sweep match all tokens; this does not prove broader equivalence. + +Instrumented replays reproduce the mismatch with **identical 130-token input prefixes**. Loop-1 hidden and loop-2 input hidden are exact. The first observed hidden drift is already at loop-2 capture output (max .01953125), before any reference consume; later-loop input differences follow. Recorded layer-0 Q/K/V and attention are exact at the capture loop; their same-input full-retention attention replays match Torch. Final logits differ by at most .125. Dense ranks `(293, 8.875)` above `(92, 8.8125)`; the reference path reverses those values. This localizes the observed drift to the capture computation before reuse, consistent with backend numerical differences, but does not localize the first physical-layer operator or establish task quality. No lossless/full-retention-at-scale claim is made. The existing async discrepancy in #157 is also not resolved by this contribution. + +## Disabled-feature drift control + +Initial unpinned C8 base/candidate controls showed timing drift, so a separate CPU-pinned **ABBA** check was run for sync draft 2, C8, feature off. Each arm uses a fresh process, the same workload and three measurements; CPUs are `[0, 1, 2, 3]`. Tokens and exit depths match the base exactly in all arms. + +- Base: median 52.06, range 51.90-53.91 tok/s. +- Candidate: median 52.19, range 50.63-52.54 tok/s. +- Candidate/base median ratio: 1.0027. These pinned observations are a separate control, not replacements for the primary unpinned numbers or a general performance guarantee. + +## Separate profiling + +Twelve readable CPU+CUDA captures cover ordinary C1/C8 and sync draft-2 C8, dense versus r=.1, each in independent ops-only/full processes. Full enables shapes, memory and stacks together. Every full capture has shape events, memory events and a nonempty exported stack inventory. The profile window spans a further output from each active request after the first outputs; raw shapes and row work remain available for checking scheduling/acceptance differences. Diagnostic count ratios are not throughput ratios. + +| C8 arrival workload | Active requests dense / r=.1 | GPU kernels dense / r=.1 | Scalar item reads dense / r=.1 | Argsorts dense / r=.1 | +|---|---:|---:|---:|---:| +| Ordinary | 8 / 8 | 5,600 / 26,220 | 0 / 384 | 0 / 192 | +| Sync, draft 2 | 7 / 5 | 11,319 / 43,242 | 0 / 720 | 0 / 360 | + +Ordinary C8 windows contain eight active requests in both arms. The synchronous windows contain seven versus five active requests because progress differs; they support operator diagnosis, not a matched-work count ratio. + +The source chain is `reference_plan -> sparse_paged_attention -> reference_statistics/read_reference`: per-row gather, selection and scalar count reads add launches and waits. Pool operations also add retirement work. The measured slowdown is therefore not explained by retained-key count alone. This contribution retains the common reference oracle; it does not duplicate or replace the pending attention optimization discussed on #157. + +![Aligned single-request diagnostic decode windows](https://github.com/Levius-Fubuki/vllm-rlt/releases/download/pr157-sync-validation-20261029/single-request-timeline.png) + +## Reproduction and evidence + +The [validation release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-sync-validation-20261029) contains scripts, original traces, stack exports, raw outputs, source inventories, the negative control, failed diagnostic-harness attempts and the first unsupported-FA4 run. The SHA256 inventory covers the package contents; `SHA256SUMS` covers the downloadable archive and timeline. + +From the corresponding source checkout with the recorded dependencies and a reserved GPU: + +```bash +pre-commit run --all-files +CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 python -m pytest \ + tests/test_sparse_decode_sync.py tests/test_sparse_decode_attention.py \ + tests/test_speculative.py tests/test_async_speculative.py \ + tests/test_async_speculative_lifecycle.py tests/test_backend_selection.py \ + --run-gpu -k 'not flash_attn_4' -q +PYTHONPATH=. CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 python scripts/real_model_bench.py \ + --model /path/to/pinned/Ouro-1.4B --mode sync2 --concurrency 32 \ + --retention .1 --repeats 3 --output /tmp/sync2-c32-r01.json +``` + +Use `--mode ordinary|sync1|sync2`, `--retention 0|1|.1`, `--concurrency 1|8|32`; retention 0 disables the feature. The script contains the exact engine startup and per-request parameters. Use `--profile ops|full --repeats 1` only for separate diagnostics. The orchestrators retain their original task directory for exact provenance; the individual benchmark CLI accepts arbitrary model/output locations. + +Graphs, token pipelines, mixed scheduling, preemption and ordinary async decode are explicitly rejected. PD, prefix-cache composition, adaptive-exit quality, long generations, held-out task quality and other hardware are not qualified here. Sparse remains disabled by default. This is a Draft collaboration contribution to the existing shared feature PR, with numerical/performance limits retained for review. From 1bbe6fdd24be3e9bbea6a9dd1440d4b64f8dd89b Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 29 Oct 2026 07:47:35 +0800 Subject: [PATCH 09/28] fix: preserve dense backend numerics in sparse capture and full retention --- tests/test_sparse_decode_attention.py | 24 ++++++++++ tests/test_sparse_decode_sync.py | 67 +++++++++++++++++++++++++++ vllm_rlt/attention/sparse.py | 25 +++++++++- vllm_rlt/core/kv_cache_manager.py | 1 + 4 files changed, 116 insertions(+), 1 deletion(-) diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index e7f85e59..7d359517 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -30,6 +30,7 @@ reference_statistics, retained_count, sparse_correction, + sparse_paged_attention, ) from vllm_rlt.worker.attention_reference import AttentionReferenceStore from vllm_rlt.worker.metadata import MetadataBuilder @@ -159,6 +160,29 @@ def test_full_retention_matches_dense_attention(): assert torch.allclose(output, dense, atol=1e-5) +@pytest.mark.parametrize("retention, expected", [(0.5, 23.5), (1.0, 40.0)]) +def test_paged_oracle_forces_correction_without_backend_shortcut(retention, expected): + store = make_store(slots=1, layers=1, heads=1, head_dim=1, key_bucket=3) + lease = store.reserve("r", 0, 1, generation=1) + q = torch.ones(1, 1, 1) + keys = torch.zeros(1, 4, 1, 1) + values = torch.zeros_like(keys) + keys[0, :3, 0, 0] = torch.log(torch.tensor([0.1, 0.3, 0.6])) + values[0, :3, 0, 0] = torch.tensor([10.0, 20.0, 30.0]) + tables, lengths = torch.tensor([[0]], dtype=torch.int32), torch.tensor([3]) + capture = store.build_plan([(lease, 0)], [CAPTURE], key_retention=retention) + # No dense callback: exercise the actual paged Torch oracle even at r=1. + output = sparse_paged_attention(q, keys, values, tables, lengths, capture, layer=0) + assert output.item() == pytest.approx(25.0) + keys[0, :3, 0, 0] = torch.log(torch.tensor([1.0, 3.0, 1.0])) + values[0, :3, 0, 0] = torch.tensor([100.0, 20.0, 40.0]) + consume = store.build_plan([(lease, 0)], [CONSUME], key_retention=retention) + output = sparse_paged_attention(q, keys, values, tables, lengths, consume, layer=0) + assert output.item() == pytest.approx(expected) + assert store.uses.item() == 1 + store.release(lease) + + # --------------------------------------------------------------------------- # Store lifecycle (RFC #156 section 5) # --------------------------------------------------------------------------- diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py index bd2f2bd7..cccbbc7f 100644 --- a/tests/test_sparse_decode_sync.py +++ b/tests/test_sparse_decode_sync.py @@ -155,6 +155,73 @@ def test_backend_explicitly_declares_sparse_reference_support(): assert selected.capabilities.sparse_decode +@pytest.mark.gpu +@pytest.mark.parametrize("kv_heads", [2, 16]) +@pytest.mark.parametrize("retention", [0.1, 1.0]) +def test_cuda_sparse_preserves_dense_backend_outputs(kv_heads, retention): + from vllm_rlt.core.kv_cache_manager import KVCacheManager + from vllm_rlt.worker.attention_reference import AttentionReferenceStore + + # BF16 rounding of Torch softmax and Triton split-KV can differ even for + # identical operands. Dense/capture rows and the full-retention control + # must preserve the configured backend, including in a mixed plan. + torch.manual_seed(157) + rows, heads, dim, length = 32, 16, 128, 130 + cache = KVCacheManager( + 1, + kv_heads, + dim, + rows * 9, + 16, + 1, + device="cuda", + dtype=torch.bfloat16, + backend="triton", + ) + store = AttentionReferenceStore( + slots=rows, layers=1, heads=heads, head_dim=dim, key_bucket=length, device="cuda" + ) + ids = [str(i) for i in range(rows)] + owners = [] + for rid in ids: + assert cache.allocate(rid, length) + keys = torch.randn(length, kv_heads, dim, device="cuda", dtype=torch.bfloat16) + values = torch.randn_like(keys) + cache.write(0, [rid] * length, [0] * length, list(range(length)), keys, values) + lease = store.reserve(rid, 0, 1, generation=1) + owners.append((lease, 0)) + batch = cache._prepare_batch(ids, [0] * rows, [length - 1] * rows) + q = torch.randn(rows, heads, dim, device="cuda", dtype=torch.bfloat16) + expected = cache._attend_prepared(0, batch, q) + modes = [DENSE] + [CAPTURE] * (rows - 1) + plan = store.build_plan(owners, modes, key_retention=retention) + actual = cache._attend_prepared(0, replace(batch, reference_plan=plan), q) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + for row in range(1, rows): + torch.testing.assert_close( + store.output[plan.slots[row], 0], expected[row].float(), atol=0, rtol=0 + ) + if retention == 1.0: + # Changed later-loop Q/K/V prevents an identity-only correction test. + q = torch.randn_like(q) + cache._write_prepared( + 0, batch, torch.randn_like(keys[:rows]), torch.randn_like(keys[:rows]) + ) + expected = cache._attend_prepared(0, batch, q) + consume = store.build_plan( + owners, [DENSE] + [CONSUME] * (rows - 1), key_retention=retention + ) + actual = cache._attend_prepared(0, replace(batch, reference_plan=consume), q) + torch.testing.assert_close(actual, expected, atol=0, rtol=0) + assert int(store.uses.sum()) == rows - 1 + for lease, _ in owners: + store.release(lease) + for rid in ids: + cache.free(rid) + torch.cuda.synchronize() + assert store.outstanding == cache.num_used_blocks == 0 + + @pytest.mark.parametrize("draft", [1, 2]) def test_sparse_verification_matches_serial_same_input_hidden_logits_and_kv(draft): from tests.helpers import make_cpu_cache diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index 80bc9e96..741d7fb0 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -28,6 +28,7 @@ """ import math +from collections.abc import Callable from dataclasses import dataclass import torch @@ -227,6 +228,7 @@ def sparse_paged_attention( plan: SparseDecodePlan, *, layer: int, + dense_attention: Callable | None = None, ) -> torch.Tensor: """Reference sparse decode attention over a paged KV cache (RFC #156 3.2). @@ -238,7 +240,10 @@ def sparse_paged_attention( ``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. + backend receives them. When supplied, ``dense_attention`` preserves that + backend's outputs for dense/capture rows and the full-retention control. + Omitting it forces the independent Torch arithmetic oracle, including the + full-retention correction; that path need not be bitwise equal to Triton. """ store = plan.store if store is None: @@ -248,6 +253,13 @@ def sparse_paged_attention( groups = heads // key_cache.shape[2] scale = 1.0 / math.sqrt(q.shape[-1]) output = torch.empty_like(q) + dense_output = None + if dense_attention is not None and ( + plan.key_retention == 1.0 or any(mode != CONSUME for mode in plan.modes) + ): + # Keep the original packed geometry: slicing rows can also change a + # backend's numerics. Only the rows needing dense output use this result. + dense_output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) captured = consumed = False for row, length in enumerate(context_lengths.tolist()): if length == 0: @@ -259,6 +271,10 @@ def sparse_paged_attention( mode = plan.modes[row] if mode == CONSUME: selected, mass, reference, contribution = store.read_reference(plan, layer, row) + if dense_output is not None and plan.key_retention == 1.0: + output[row] = dense_output[row] + consumed = True + continue keys, values = gather_retained( key_cache, value_cache, block_tables[row], selected, block_size, groups ) @@ -277,6 +293,10 @@ def sparse_paged_attention( key_retention=plan.key_retention, key_bucket=store.key_bucket, ) + if dense_output is not None: + # O_ref is the output actually fed into the recurrence. Keep + # selector probabilities/contributions in the Torch FP32 oracle. + reference = dense_output[row].float() output[row] = reference.to(q.dtype) store.write_reference( plan, @@ -291,6 +311,9 @@ def sparse_paged_attention( captured = True continue # Dense row: the ordinary attention path. + if dense_output is not None: + output[row] = dense_output[row] + continue 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 diff --git a/vllm_rlt/core/kv_cache_manager.py b/vllm_rlt/core/kv_cache_manager.py index e073414a..a6595f28 100644 --- a/vllm_rlt/core/kv_cache_manager.py +++ b/vllm_rlt/core/kv_cache_manager.py @@ -732,6 +732,7 @@ def _attend_prepared( batch.context_lengths, plan, layer=layer, + dense_attention=self.attention, ) if isinstance(batch, DeviceKVMetadata): return self.attention( From fbd8dea4fd6b66374411c6555a29150e4f57b935 Mon Sep 17 00:00:00 2001 From: levius <2114377220@qq.com> Date: Thu, 29 Oct 2026 07:58:58 +0800 Subject: [PATCH 10/28] docs: record sparse backend numerical diagnosis and fixed BF16 regressions --- docs/runtime.md | 18 ++- .../sparse-decode-numerics-20261029.md | 145 ++++++++++++++++++ .../validation/sparse-decode-sync-20261029.md | 5 + 3 files changed, 162 insertions(+), 6 deletions(-) create mode 100644 docs/validation/sparse-decode-numerics-20261029.md diff --git a/docs/runtime.md b/docs/runtime.md index ab5cbd32..d3958ebd 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -253,9 +253,12 @@ This path supports ordinary synchronous decode and synchronous/asynchronous block self-speculation with LAST_EXITED KV and refill scheduling. The Torch and Triton adapters declare sparse-reference support; `auto` selects a compatible adapter and explicit FlashAttention selection fails rather than silently -replacing the requested backend. Sparse capture/correction itself currently -uses the shared PyTorch reference implementation, even beside Triton dense -attention. It has host readbacks and per-row launch costs; enabling it does not +replacing the requested backend. Dense and capture rows preserve the selected +backend's output and original packed batch geometry. Selection, retained mass +and contributions, and partial-retention correction use the shared PyTorch +reference implementation. The captured output is the value actually entering +the recurrence, stored in FP32 after the backend's output-dtype conversion. +It has host readbacks and per-row launch costs; enabling it does not promise faster serving. CUDA Graphs, token pipelines, mixed scheduling, preemption/recovery and ordinary asynchronous decode are rejected explicitly. @@ -274,9 +277,12 @@ recurrent reader. Synchronous block references retire on the compute stream when every verification input has completed its final read; neither accepted nor rejected rows carry a reference into the next round. With draft depth 1, verification captures loop 2; with draft depth 2, drafting captures it. -Full retention is a numerical control. Use declared tolerances and investigate -any real-model token differences from dense execution; small-model equality -does not establish equivalence at scale. Retention below 1 is approximate and +Full retention uses the same dense backend at later loops as a control, while +still validating and consuming the captured reference lifecycle. The standalone +Torch arithmetic oracle is tested separately; different attention algorithms +need not round to identical BF16 values. See the +[numerical diagnosis and regression report](validation/sparse-decode-numerics-20261029.md). +Retention below 1 is approximate and needs task-quality evaluation before a deployment recommendation. ## SHARED prefill semantics diff --git a/docs/validation/sparse-decode-numerics-20261029.md b/docs/validation/sparse-decode-numerics-20261029.md new file mode 100644 index 00000000..106918c7 --- /dev/null +++ b/docs/validation/sparse-decode-numerics-20261029.md @@ -0,0 +1,145 @@ +# Sparse decode: backend numerical diagnosis + +This follow-up fixes the full-retention discrepancy in the +[ordinary/synchronous integration report](sparse-decode-sync-20261029.md), within +the same Draft collaboration PR targeting `kunkun/rfc156-sparse-decode-async`. +The immediate feature base remains `048b9d7c79a5711857ce3287812502ad8f274ad0`; +the pre-fix contribution head is `7f22af3901416096ea6c1c0e0a5536f440f662f5`. +The fixed runtime/test commit is `1bbe6fdd24be3e9bbea6a9dd1440d4b64f8dd89b`. + +## Reproduced cause + +The original BF16 Ouro-1.4B workload reproduces 498/512 token agreement at +concurrent arrivals 32 and synchronous draft depth 1. Request 9 first differs +at output index 2: dense chooses 293, the original sparse reference chooses 92. +The 14 differences include subsequent autoregressive continuation, rather than +14 independent initial errors. + +The new diagnostic retains the original 128-input/16-output protocol, including +four-token warmup. It freezes the complete accepted verification batch's input +hidden states, logical row/depth/position metadata, position IDs and context +lengths, all physical-layer module outputs, full packed queries and attention +outputs, and the watched row's causal K/V. It filters out rejected executions of +the same position. The accepted loop-2 batch here has **22 rows**, not the +theoretical 96-row maximum; both arms have exactly equal packed input states, +request histories, shapes, positions and lengths. A shorter-output exploratory +run changed the batch geometry and is retained separately, not substituted for +this reproduction. + +For the watched row, physical layers 0 and 1 agree. At physical layer **2** +(zero-based), Q, all 130 legal keys and values are identical, but the attention +output differs at one BF16 element by **0.000030517578125**. The next `o_proj` +differs at two elements, max **0.000244140625**. Across the complete batch, +attention differences already exist at physical layer 0. Later layers and +recurrences amplify these differences. This places the first observed watched +operator discrepancy at attention, before reference consumption. + +A separate replay uses the frozen operands with no model or scheduler. Triton +split-KV reproduces the captured dense output exactly; the Torch +einsum/softmax reference reproduces the captured sparse output exactly. Changing +only the backend reproduces the layer-2 difference. Disabling BF16 reduced +precision reduction does not change that reference result; its reductions +already use FP32 operands. TF32 is disabled throughout. Equal packed shapes +rule out a batch-shape change as the explanation for this captured attention +discrepancy. No global precision flags or model GEMM implementations are changed. + +## Fix and numerical contract + +`KVCacheManager` passes its selected dense attention adapter into the shared +`sparse_paged_attention` dispatcher. The adapter receives the original packed +queries, tables and lengths, without slicing into different batch shapes. + +- DENSE and CAPTURE rows use that backend's actual output. Capture stores that + output as FP32 **after** the backend's output-dtype rounding, so `O_ref` is + the value entering recurrence. Retained selection, mass and contribution + remain the FP32 Torch reference arithmetic. +- At retention 1, CONSUME rows also use the configured dense backend, while + still checking reference ownership/payload availability and recording + consumption. This is explicitly a dense control; it does not qualify the + Torch correction as bitwise equal to Triton. +- At retention below 1, later loops retain `O_ref - C_ref + m * C_new`. This + changes the original reference's source-output rounding contract, not the + selection or sparse delta formula. Approximate quality remains unqualified. +- Omitting the adapter when calling the standalone paged oracle forces the + original Torch capture and correction, including at retention 1. Known-answer + paged cases test changed later keys/values at retention .5 and 1; the full-retention + control cannot hide a regression in that independent arithmetic path. + +A full-retention-only bypass would leave partial-retention capture rows using +a different dense algorithm; preserving the selected backend at every dense +source row fixes that inconsistency too. No new kernel or global precision +setting is introduced. The fallback oracle remains independently available. + +## Regression evidence + +Four CUDA BF16 cases cover MHA/GQA, retention .1/1, a mixed dense/capture plan, +stored `O_ref`, changed later-loop operands and reference consumption. All four +fail on the original implementation with backend-output differences, then pass +with the fix using exact equality. The standalone paged known-answer cases +return 23.5 and 40.0 after changing later keys/values, exercising partial and +full Torch correction without the backend control. + +Full CPU regression: **921 passed, 479 skipped**. Full pre-commit passes Ruff +lint, Ruff format and the CPU suite. Independent read-only review found no +actionable defect in the four-file runtime/test delta; its targeted CPU check +passed 48 cases, with 22 GPU cases deselected. + +Applicable CUDA regression: **318 passed, 1 skipped, 16 deselected**. The skip +requires two GPUs; FA4 cases remain excluded on this SM89 device. A fixed +instrumented replay now matches the original dense replay exactly at every +recorded depth: complete packed inputs/outputs, all physical module outputs, +attention queries/outputs and the watched row's causal K/V. This comparison is +separate from the uninstrumented generation matrix. + +The uninstrumented BF16 regression matrix covers ordinary decode and synchronous +draft depths 1/2, simultaneous arrivals 1/8/32, and retention off/1/.1. Each of +the 27 fresh-process configurations warms four output tokens, then measures +three bursts using the original 128-input/16-output protocol. All **81 bursts, +1,107 requests and 17,712 output tokens** complete, with KV/reference state +drained. Tokens and exit depths repeat exactly within each configuration. + +| Mode | C=1 full vs dense | C=8 full vs dense | C=32 full vs dense | +|---|---:|---:|---:| +| Ordinary | 16/16 | 128/128 | 512/512 | +| Sync, draft 1 | 16/16 | 128/128 | 512/512 | +| Sync, draft 2 | 16/16 | 128/128 | 512/512 | + +Every cell matches in all three repetitions, including exit depths: **5,904/5,904 +full-retention tokens** in total. In particular, the original C32/draft-1 case +improves from 498/512 to 512/512 in each repetition. This verifies the explicit +dense-backend control; the forced Torch correction remains a separate oracle. +Retention .1 completes and drains, but its token agreement is not a quality score. + +## Reproduction and archive + +The [numerical validation release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-numerics-validation-20261029) +contains the frozen tensor snapshots, source and SHA256 inventories, Git bundle, +same-input replay, original failures, red/green regression logs, full raw matrix, +environment metadata and reproduction commands. Model weights are not included. +The original unsupported-FA4 failures and performance profiles remain in the +[earlier release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-sync-validation-20261029). + +Use `trace_first_operator.py` against `7f22af3` for the dense and pre-fix capture, +then against `1bbe6fd` for the fixed capture; `replay_attention.py` tests the +frozen operands independently. `compare_fixed_trace.py` checks every recorded +state, and `run_regression_matrix.py` / `summarize_matrix.py` reproduce and verify +the complete generation matrix. These local experiment scripts preserve their +recorded root directory; the archive README explains relocation and exact CLI +commands. Diagnostic timings are excluded from uninstrumented results. + +## Scope + +The model weights match the prior SHA256 +`58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. +The follow-up server reports one RTX4090 (SM89) with 49,140 MiB, driver +595.71.05, torch 2.12.1+cu130 and Triton 3.7.1. This is a different instance +from the original 24,564-MiB server with driver 580.105.08; +timings across these instances are not a matched performance comparison. + +These are numerical and execution regressions on synthetic fixed-token inputs, +not a held-out task-quality evaluation. Partial retention, adaptive exits, +long contexts, other GPUs and untested feature combinations remain outside the +qualification. Existing asynchronous lifecycle regressions are included, but +no new real-model asynchronous quality claim is made. The original performance +profiles remain historical evidence; +this correctness fix makes no serving-speed improvement claim. diff --git a/docs/validation/sparse-decode-sync-20261029.md b/docs/validation/sparse-decode-sync-20261029.md index 4685b51e..91d3a7a1 100644 --- a/docs/validation/sparse-decode-sync-20261029.md +++ b/docs/validation/sparse-decode-sync-20261029.md @@ -1,5 +1,10 @@ # Sparse decode: ordinary and synchronous integration validation +This is the historical report for runtime `10373ab`. The subsequent +[numerical diagnosis and regression report](sparse-decode-numerics-20261029.md) +documents the backend-preserving fix and its separate validation. The timings +and numerical limitation below describe the original implementation. + This contribution extends the shared [#157](https://github.com/InferMatrixLab/vllm-rlt/pull/157) reference plan/store into ordinary synchronous decode and synchronous block self-speculation. It is a functional reference integration, with measured costs and an unresolved BF16 full-retention sequence difference. It does not claim an optimized sparse kernel or deployment quality. ## Source and checks From 0aaab79b4a3b86731c89eb67f42036b794d9b05c Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 28 Oct 2026 15:00:37 +0100 Subject: [PATCH 11/28] 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 741d7fb0..bb90de69 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -219,7 +219,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, @@ -324,3 +324,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 903e1013e0f753715c9733061c487d40d5bb4ff9 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 28 Oct 2026 15:15:17 +0100 Subject: [PATCH 12/28] 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 bb90de69..09638d9c 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -336,6 +336,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 ): @@ -389,9 +399,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] @@ -405,9 +415,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) @@ -420,7 +430,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 a7859a5c695dbd009027a77f46b934633ecf4e0a Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 28 Oct 2026 15:30:01 +0100 Subject: [PATCH 13/28] 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 09638d9c..5ba89789 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -343,7 +343,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( From e4f89883551b40b9ff6ed926084f7ad53ba66f48 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Thu, 29 Oct 2026 14:30:13 +0800 Subject: [PATCH 14/28] sparse decode attention: restore the dense-output contract in the batched path The cherry-picked batching commits replaced sparse_paged_attention with a dispatcher and dropped the dense_attention handoff, so _attend_prepared would have raised TypeError and the numerical contract would have been lost. Both implementations now honour it: - the dispatcher forwards dense_attention to either implementation; - the batched path uses the backend output for dense rows, stores the actually rounded recurrence input as O_ref for capture rows, and returns the backend output directly for full-retention consume rows, while keeping selector probabilities and the retained contribution in the Torch FP32 oracle. benchmarks/sparse_decode_batched_ab.py isolates the batching gain from that contract by timing both implementations on the same plan and store. --- benchmarks/sparse_decode_batched_ab.py | 169 +++++++++++++++++++++++++ vllm_rlt/attention/sparse.py | 54 +++++++- 2 files changed, 220 insertions(+), 3 deletions(-) create mode 100644 benchmarks/sparse_decode_batched_ab.py diff --git a/benchmarks/sparse_decode_batched_ab.py b/benchmarks/sparse_decode_batched_ab.py new file mode 100644 index 00000000..99541948 --- /dev/null +++ b/benchmarks/sparse_decode_batched_ab.py @@ -0,0 +1,169 @@ +"""A/B the scalar and batched sparse decode paths inside the merged tree. + +Both implementations live in ``vllm_rlt.attention.sparse`` and share the plan and +store contract, so this isolates the batching gain from the numerical contract: + +* ``scalar`` ``reference_paged_attention`` (row loop, the oracle), +* ``batched`` ``batched_paged_attention`` (equal-length row groups, CUDA). + +Capture rows are also measured with and without the dense-backend handoff, which +the full-retention control requires. + +Usage: + python benchmarks/sparse_decode_batched_ab.py --device cuda --rows 32 --context 512 +""" + +import argparse +import json +import math +import statistics +import time + +import torch + +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + batched_paged_attention, + reference_paged_attention, +) +from vllm_rlt.kernels.paged_attention import torch_paged_attention +from vllm_rlt.worker.attention_reference import AttentionReferenceStore + + +def build(rows, context, heads, head_dim, kv_heads, layers, block_size, retention, device): + blocks = math.ceil(context / block_size) + generator = torch.Generator().manual_seed(0) + shape = (blocks, block_size, kv_heads, head_dim) + key = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + value = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + table = torch.arange(blocks, device=device, dtype=torch.int32).repeat(rows, 1) + lengths = torch.full((rows,), context, device=device, dtype=torch.int32) + q = torch.randn((rows, heads, head_dim), generator=generator).to(device, torch.bfloat16) + capacity = math.ceil(context * retention) + store = AttentionReferenceStore( + slots=rows, + layers=layers, + heads=heads, + head_dim=head_dim, + key_bucket=capacity, + device=device, + ) + return key, value, table, lengths, q, store + + +def timed(fn, repeats, device): + for _ in range(3): + fn() + torch.cuda.synchronize(device) + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize(device) + samples.append(time.perf_counter() - start) + return statistics.median(samples) * 1e3 + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--rows", type=int, default=32) + 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("--block-size", type=int, default=16) + parser.add_argument("--retention", type=float, default=0.1) + parser.add_argument("--repeats", type=int, default=10) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + device = torch.device(args.device) + + report = { + "device": str(device), + "rows": args.rows, + "context": args.context, + "retention": args.retention, + "layers": args.layers, + } + for mode, name in ((CAPTURE, "capture"), (CONSUME, "consume")): + key, value, table, lengths, q, store = build( + args.rows, + args.context, + args.heads, + args.head_dim, + args.kv_heads, + args.layers, + args.block_size, + args.retention, + device, + ) + lease = store.reserve("r", 0, args.rows, generation=1) + owners = [(lease, row) for row in range(args.rows)] + # Seed a capture pass so consume rows have state. + seed_plan = store.build_plan(owners, [CAPTURE] * args.rows, key_retention=args.retention) + reference_paged_attention( + q, + key, + value, + table, + lengths, + seed_plan, + layer=0, + dense_attention=torch_paged_attention, + ) + plan = store.build_plan(owners, [mode] * args.rows, key_retention=args.retention) + dense = torch_paged_attention + scalar = timed( + lambda: reference_paged_attention( + q, key, value, table, lengths, plan, layer=0, dense_attention=dense + ), + args.repeats, + device, + ) + batched = timed( + lambda: batched_paged_attention( + q, key, value, table, lengths, plan, layer=0, dense_attention=dense + ), + args.repeats, + device, + ) + batched_oracle = timed( + lambda: batched_paged_attention(q, key, value, table, lengths, plan, layer=0), + args.repeats, + device, + ) + report[name] = { + "scalar_ms": round(scalar, 3), + "batched_ms": round(batched, 3), + "batched_no_dense_handoff_ms": round(batched_oracle, 3), + "speedup": round(scalar / batched, 2), + } + store.release(lease) + + dense_ms = timed( + lambda: torch_paged_attention(q, key, value, table, lengths), args.repeats, device + ) + report["dense_attention_ms"] = round(dense_ms, 3) + if args.json: + print(json.dumps(report, indent=2)) + return + print( + f"sparse decode path A/B — {device}, rows={args.rows}, context={args.context}, " + f"retention={args.retention}, layers={args.layers}" + ) + for name in ("capture", "consume"): + values = report[name] + print( + f" {name:8s} scalar {values['scalar_ms']:7.3f} ms " + f"batched {values['batched_ms']:7.3f} ms " + f"speedup {values['speedup']:5.2f}x " + f"(batched w/o dense handoff {values['batched_no_dense_handoff_ms']:7.3f} ms)" + ) + print(f" dense attention reference: {dense_ms:.3f} ms") + + +if __name__ == "__main__": + main() diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index 5ba89789..3b8eb5d7 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -327,12 +327,27 @@ def reference_paged_attention( def sparse_paged_attention( - q, key_cache, value_cache, block_tables, context_lengths, plan, *, layer + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention: Callable | None = None, ): """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 + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, ) @@ -357,7 +372,15 @@ def _rowwise_product(equation, left, right): def batched_paged_attention( - q, key_cache, value_cache, block_tables, context_lengths, plan, *, layer + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention: Callable | None = None, ): """Batch equal-length rows without padding or changing retained-key selection. @@ -365,12 +388,25 @@ def batched_paged_attention( 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. + + ``dense_attention`` keeps the numerical contract of the scalar path: dense + and capture rows publish the selected backend's output, the actually + rounded recurrence input becomes ``O_ref``, and a full-retention consume row + returns that same backend output as the control. Selector probabilities and + the retained contribution stay in the Torch FP32 oracle. """ 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() + dense_output = None + if dense_attention is not None and ( + plan.key_retention == 1.0 or any(mode != CONSUME for mode in plan.modes) + ): + # Keep the original packed geometry: slicing rows can also change a + # backend's numerics. + dense_output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) consume_rows = [r for r, mode in enumerate(plan.modes) if mode == CONSUME and lengths[r]] counts = {} if consume_rows: @@ -403,6 +439,11 @@ def batched_paged_attention( tables = block_tables[row_ids] batch_index = torch.arange(len(rows), device=q.device)[:, None, None] if mode == CONSUME: + if dense_output is not None and plan.key_retention == 1.0: + # Full retention is the control: same backend as dense. + output[row_ids] = dense_output[row_ids] + consumed = True + continue 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] @@ -423,11 +464,18 @@ def batched_paged_attention( positions = torch.arange(extent, device=q.device) blocks = tables[:, positions // block_size] offsets = positions % block_size + if mode == DENSE and dense_output is not None: + output[row_ids] = dense_output[row_ids] + continue keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() scores = _rowwise_product("hd,lhd->hl", queries, keys) * scale probabilities = torch.softmax(scores, dim=-1) reference = _rowwise_product("hl,lhd->hd", probabilities, values) + if dense_output is not None: + # O_ref is the output actually fed into the recurrence; the + # selector state below stays in the Torch FP32 oracle. + reference = dense_output[row_ids].float() output[row_ids] = reference.to(q.dtype) if mode == CAPTURE: count = retained_count(extent, plan.key_retention) From f292b76f51466fc85d05b6671f17eaec7dc39846 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Thu, 29 Oct 2026 21:58:25 +0800 Subject: [PATCH 15/28] benchmarks: dense-baseline context sweep and memory breakdown Measures the four-loop decode cost against the dense Triton kernel (the acceptance baseline, not the scalar sparse reference) across 256-32768 tokens, plus the three memory components: KV storage, the persistent reference pool and the transient capture/consume workspace. The result is the diagnosis the review asked for: capture = dense + selector, and the selector costs about 25x per key relative to the fused kernel, so no context in the sweep is a net win. --- benchmarks/sparse_decode_context_sweep.py | 242 ++++++++++++++++++++++ 1 file changed, 242 insertions(+) create mode 100644 benchmarks/sparse_decode_context_sweep.py diff --git a/benchmarks/sparse_decode_context_sweep.py b/benchmarks/sparse_decode_context_sweep.py new file mode 100644 index 00000000..5ed1f667 --- /dev/null +++ b/benchmarks/sparse_decode_context_sweep.py @@ -0,0 +1,242 @@ +"""Dense-baseline context sweep and memory breakdown for sparse decode attention. + +The acceptance baseline is the dense Triton kernel, not the scalar sparse +reference: a four-loop decode costs ``4 x dense`` with the feature off and +``dense + capture + 2 x consume`` with it on. This script measures both on +identical paged KV at increasing context, reports the ratio, and breaks memory +down into KV storage, the persistent reference pool and transient workspace. + +Usage: + python benchmarks/sparse_decode_context_sweep.py --device cuda + python benchmarks/sparse_decode_context_sweep.py --contexts 512,2048,8192 --rows 32 +""" + +import argparse +import json +import math +import statistics +import time + +import torch + +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + batched_paged_attention, + reference_paged_attention, +) +from vllm_rlt.kernels.paged_attention import triton_paged_attention +from vllm_rlt.worker.attention_reference import AttentionReferenceStore + + +def build(context, rows, heads, head_dim, kv_heads, layers, block_size, retention, device): + blocks = math.ceil(context / block_size) + generator = torch.Generator().manual_seed(0) + shape = (blocks, block_size, kv_heads, head_dim) + key = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + value = torch.randn(shape, generator=generator).to(device, torch.bfloat16) + table = torch.arange(blocks, device=device, dtype=torch.int32).repeat(rows, 1) + lengths = torch.full((rows,), context, device=device, dtype=torch.int32) + q = torch.randn((rows, heads, head_dim), generator=generator).to(device, torch.bfloat16) + capacity = math.ceil(context * retention) + store = AttentionReferenceStore( + slots=rows, + layers=layers, + heads=heads, + head_dim=head_dim, + key_bucket=capacity, + device=device, + ) + return key, value, table, lengths, q, store + + +def timed(fn, repeats, device): + for _ in range(2): + fn() + torch.cuda.synchronize(device) + samples = [] + for _ in range(repeats): + start = time.perf_counter() + fn() + torch.cuda.synchronize(device) + samples.append(time.perf_counter() - start) + return statistics.median(samples) * 1e3 + + +def transient_bytes(fn, device): + torch.cuda.synchronize(device) + torch.cuda.reset_peak_memory_stats(device) + before = torch.cuda.memory_allocated(device) + fn() + torch.cuda.synchronize(device) + return max(0, torch.cuda.max_memory_allocated(device) - before) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--contexts", default="256,512,1024,2048,4096,8192,16384,32768") + parser.add_argument("--rows", type=int, default=32) + 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("--block-size", type=int, default=16) + parser.add_argument("--retention", type=float, default=0.1) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + device = torch.device(args.device) + contexts = [int(value) for value in args.contexts.split(",")] + + rows = [] + for context in contexts: + try: + key, value, table, lengths, q, store = build( + context, + args.rows, + args.heads, + args.head_dim, + args.kv_heads, + args.layers, + args.block_size, + args.retention, + device, + ) + lease = store.reserve("r", 0, args.rows, generation=1) + owners = [(lease, row) for row in range(args.rows)] + seed = store.build_plan(owners, [CAPTURE] * args.rows, key_retention=args.retention) + batched_paged_attention( + q, key, value, table, lengths, seed, layer=0, dense_attention=triton_paged_attention + ) + capture_plan = store.build_plan( + owners, [CAPTURE] * args.rows, key_retention=args.retention + ) + consume_plan = store.build_plan( + owners, [CONSUME] * args.rows, key_retention=args.retention + ) + dense_ms = timed( + lambda: triton_paged_attention(q, key, value, table, lengths), + args.repeats, + device, + ) + capture_ms = timed( + lambda: batched_paged_attention( + q, + key, + value, + table, + lengths, + capture_plan, + layer=0, + dense_attention=triton_paged_attention, + ), + args.repeats, + device, + ) + consume_ms = timed( + lambda: batched_paged_attention( + q, key, value, table, lengths, consume_plan, layer=0 + ), + args.repeats, + device, + ) + scalar_consume_ms = timed( + lambda: reference_paged_attention( + q, key, value, table, lengths, consume_plan, layer=0 + ), + args.repeats, + device, + ) + dense_total = 4 * dense_ms + sparse_total = dense_ms + capture_ms + 2 * consume_ms + rows.append( + dict( + context=context, + retained=math.ceil(context * args.retention), + dense_ms=round(dense_ms, 3), + capture_ms=round(capture_ms, 3), + consume_ms=round(consume_ms, 3), + scalar_consume_ms=round(scalar_consume_ms, 3), + dense_4loop_ms=round(dense_total, 3), + sparse_4loop_ms=round(sparse_total, 3), + sparse_over_dense=round(sparse_total / dense_total, 3), + capture_transient_mib=round( + transient_bytes( + lambda: batched_paged_attention( + q, + key, + value, + table, + lengths, + capture_plan, + layer=0, + dense_attention=triton_paged_attention, + ), + device, + ) + / 1024**2, + 1, + ), + consume_transient_mib=round( + transient_bytes( + lambda: batched_paged_attention( + q, key, value, table, lengths, consume_plan, layer=0 + ), + device, + ) + / 1024**2, + 1, + ), + reference_pool_mib=round(rows_pool_bytes(args, context) / 1024**2, 1), + kv_mib=round(key.numel() * key.element_size() * 2 / 1024**2, 1), + ) + ) + store.release(lease) + torch.cuda.empty_cache() + except torch.OutOfMemoryError: + rows.append(dict(context=context, status="oom")) + torch.cuda.empty_cache() + + if args.json: + print(json.dumps(rows, indent=2)) + return + print( + f"dense-baseline sweep — {device}, rows={args.rows}, layers={args.layers}, " + f"heads={args.heads}, head_dim={args.head_dim}, retention={args.retention}" + ) + print( + f"{'ctx':>7}{'k':>7}{'dense':>9}{'capture':>9}{'consume':>9}" + f"{'4-loop dense':>13}{'4-loop sparse':>14}{'ratio':>8}" + ) + for row in rows: + if row.get("status") == "oom": + print(f"{row['context']:>7} OOM") + continue + print( + f"{row['context']:>7}{row['retained']:>7}{row['dense_ms']:>9.2f}" + f"{row['capture_ms']:>9.2f}{row['consume_ms']:>9.2f}" + f"{row['dense_4loop_ms']:>13.2f}{row['sparse_4loop_ms']:>14.2f}" + f"{row['sparse_over_dense']:>8.2f}x" + ) + print() + print(f"{'ctx':>7}{'capture ws':>12}{'consume ws':>12}{'ref pool':>11}{'KV':>9}") + for row in rows: + if row.get("status") == "oom": + continue + print( + f"{row['context']:>7}{row['capture_transient_mib']:>10.1f}Mi" + f"{row['consume_transient_mib']:>10.1f}Mi" + f"{row['reference_pool_mib']:>9.1f}Mi{row['kv_mib']:>7.1f}Mi" + ) + + +def rows_pool_bytes(args, context): + """Persistent reference pool for this sweep's row count.""" + capacity = math.ceil(context * args.retention) + per_slot = capacity * 4 + 4 + 4 * 2 * args.head_dim + 8 + return args.rows * args.layers * args.heads * per_slot + + +if __name__ == "__main__": + main() From 25e15d7eb0948bef123cb0c9f8f02be76e4cb543 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Sat, 31 Oct 2026 23:08:14 +0800 Subject: [PATCH 16/28] sparse decode attention: fused paged score pass for capture The capture step was CPU-bound and paid for an expanded FP32 copy of the whole causal row. Measured at 32 rows / 24L / 16H / 128D / BF16, r=0.1, capture went from 59.6 ms to 20.9 ms at 8k context and from 221.0 ms to 43.9 ms at 32k. - kernels/sparse_capture.py: one program per (row, kv head, position block) computes q @ K^T for that KV head's query group straight out of the paged cache, writing FP32 scores. No expanded K, no per-row cuBLAS, one launch per layer. Accepts a 5-D cache or the per-layer view the cache manager hands the attention layer. Against the oracle at 8k: max score delta 1.1e-5 on a magnitude of 65, retained mass delta 9e-8, attention output delta 2e-8, and identical selected sets on 511/512 row-head pairs. - attention/sparse.py: the capture branch takes that pass when Triton is available and a dense backend is present; the scalar oracle remains the definition of the arithmetic, and batched_paged_attention without fused_scores still reproduces it bit for bit (0z5a's test is unchanged). - Chunk budgets follow the data each path actually materialises. The old 16M element budget was below a single 8k expanded row, so a group ran one row at a time: 32 small launches plus a host sync each. The fused path budgets the score matrix instead and keeps the whole group in one chunk. - The capture contribution gathers retained values straight from the paged cache rather than from the expanded FP32 copy. BF16 -> FP32 is exact, so the bits are identical while the read no longer scatters through a row-length tensor. - benchmarks/sparse_decode_context_sweep.py measures the production entry point instead of one implementation, so the sweep follows the dispatched path. Consume is now the bottleneck and is untouched here: about 10 ms of per-layer fixed cost (including a 2.45 ms host readback of the retained count) plus a gather that costs 16 ms at 32k, which is why the four-loop total is still 2.5-4x dense even though capture alone is now cheaper than dense at 512. --- benchmarks/sparse_decode_context_sweep.py | 12 +- vllm_rlt/attention/sparse.py | 68 ++++++++-- vllm_rlt/kernels/sparse_capture.py | 144 ++++++++++++++++++++++ 3 files changed, 208 insertions(+), 16 deletions(-) create mode 100644 vllm_rlt/kernels/sparse_capture.py diff --git a/benchmarks/sparse_decode_context_sweep.py b/benchmarks/sparse_decode_context_sweep.py index 5ed1f667..6eae0687 100644 --- a/benchmarks/sparse_decode_context_sweep.py +++ b/benchmarks/sparse_decode_context_sweep.py @@ -22,8 +22,8 @@ from vllm_rlt.attention.sparse import ( CAPTURE, CONSUME, - batched_paged_attention, reference_paged_attention, + sparse_paged_attention, ) from vllm_rlt.kernels.paged_attention import triton_paged_attention from vllm_rlt.worker.attention_reference import AttentionReferenceStore @@ -106,7 +106,7 @@ def main(): lease = store.reserve("r", 0, args.rows, generation=1) owners = [(lease, row) for row in range(args.rows)] seed = store.build_plan(owners, [CAPTURE] * args.rows, key_retention=args.retention) - batched_paged_attention( + sparse_paged_attention( q, key, value, table, lengths, seed, layer=0, dense_attention=triton_paged_attention ) capture_plan = store.build_plan( @@ -121,7 +121,7 @@ def main(): device, ) capture_ms = timed( - lambda: batched_paged_attention( + lambda: sparse_paged_attention( q, key, value, @@ -135,7 +135,7 @@ def main(): device, ) consume_ms = timed( - lambda: batched_paged_attention( + lambda: sparse_paged_attention( q, key, value, table, lengths, consume_plan, layer=0 ), args.repeats, @@ -163,7 +163,7 @@ def main(): sparse_over_dense=round(sparse_total / dense_total, 3), capture_transient_mib=round( transient_bytes( - lambda: batched_paged_attention( + lambda: sparse_paged_attention( q, key, value, @@ -180,7 +180,7 @@ def main(): ), consume_transient_mib=round( transient_bytes( - lambda: batched_paged_attention( + lambda: sparse_paged_attention( q, key, value, table, lengths, consume_plan, layer=0 ), device, diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index 3b8eb5d7..cedf9343 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -33,6 +33,8 @@ import torch +from vllm_rlt.kernels.sparse_capture import TRITON_AVAILABLE, fused_capture_scores + # Row modes for one recurrent submission. DENSE = 0 CAPTURE = 1 @@ -337,9 +339,25 @@ def sparse_paged_attention( layer, dense_attention: Callable | None = None, ): - """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( + """Use bounded row batches on CUDA, retaining the scalar numerical oracle. + + On CUDA with a dense backend available, the capture rows take the fused + paged score pass; the scalar oracle stays the definition of the arithmetic, + and ``batched_paged_attention`` without ``fused_scores`` still reproduces it + bit for bit. + """ + if not q.is_cuda: + return reference_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + ) + return batched_paged_attention( q, key_cache, value_cache, @@ -348,6 +366,7 @@ def sparse_paged_attention( plan, layer=layer, dense_attention=dense_attention, + fused_scores=TRITON_AVAILABLE and dense_attention is not None, ) @@ -381,6 +400,7 @@ def batched_paged_attention( *, layer, dense_attention: Callable | None = None, + fused_scores: bool = False, ): """Batch equal-length rows without padding or changing retained-key selection. @@ -429,8 +449,14 @@ def batched_paged_attention( 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))) + # The fused path materialises only the score matrix, so its budget is + # the scores themselves; the expanded path budgets an FP32 K row. Both + # must allow the whole group at once, or a 32k row degenerates into one + # launch per row. + if fused_scores: + chunk = max(1, min(32, 512 * 1024 * 1024 // (heads * extent * 4))) + else: + chunk = max(1, min(32, 64 * 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) @@ -467,15 +493,30 @@ def batched_paged_attention( if mode == DENSE and dense_output is not None: output[row_ids] = dense_output[row_ids] continue - keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() - values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() - scores = _rowwise_product("hd,lhd->hl", queries, keys) * scale + if fused_scores and dense_output is not None: + # One pass over the paged cache: no expanded FP32 K row, and + # no per-row cuBLAS. The dot product is still FP32; only the + # reduction order differs from the oracle. + scores = fused_capture_scores( + q[row_ids], + key_cache, + tables, + context_lengths[row_ids], + layer=layer, + block_size=block_size, + ) + scores = scores * scale + else: + keys = key_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + scores = _rowwise_product("hd,lhd->hl", queries, keys) * scale probabilities = torch.softmax(scores, dim=-1) - reference = _rowwise_product("hl,lhd->hd", probabilities, values) if dense_output is not None: # O_ref is the output actually fed into the recurrence; the # selector state below stays in the Torch FP32 oracle. reference = dense_output[row_ids].float() + else: + values = value_cache[blocks, offsets].repeat_interleave(groups, dim=2).float() + 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) @@ -487,7 +528,14 @@ def batched_paged_attention( 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] + # Gather the retained values straight from the paged cache: the + # expanded fp32 copy would be the same bits (bf16 -> fp32 is + # exact) but reading it scatters through a row-length tensor. + selected_values = value_cache[ + tables[batch_index, selected // block_size], + selected % block_size, + kv_index, + ].float() 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 diff --git a/vllm_rlt/kernels/sparse_capture.py b/vllm_rlt/kernels/sparse_capture.py new file mode 100644 index 00000000..88135673 --- /dev/null +++ b/vllm_rlt/kernels/sparse_capture.py @@ -0,0 +1,144 @@ +"""Fused paged score pass for the sparse-decode capture step. + +The selector needs a probability for every key of the reference loop, but the +Torch oracle materialises the whole causal row in FP32 with the GQA heads +expanded, which costs tens of MiB per layer and leaves the pass CPU-bound. This +kernel reads the paged cache once and writes one FP32 score per +``(row, head, position)``, so selection, softmax and the retained contribution +stay on the device without an expanded copy. + +It computes the same dot product as the oracle in FP32 (the cache is BF16, and +BF16 -> FP32 is exact); only the reduction order differs. Measured against the +oracle at 8k context: max score delta 1.1e-5 on a magnitude of 65, retained mass +delta 9e-8, full attention output delta 2e-8, and identical selected sets on +511/512 row-head pairs. +""" + +from __future__ import annotations + +import torch + +try: # pragma: no cover - exercised only without Triton + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None + tl = None + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _capture_scores_kernel( + q_ptr, + k_ptr, + t_ptr, + l_ptr, + out_ptr, + context, + q_stride_row, + k_stride_block, + DIM: tl.constexpr, + GROUPS: tl.constexpr, + KV: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_N: tl.constexpr, + ): + row = tl.program_id(0) + kv_head = tl.program_id(1) + chunk = tl.program_id(2) + + position = chunk * BLOCK_N + tl.arange(0, BLOCK_N) + row_length = tl.load(l_ptr + row) + mask = position < row_length + page = tl.load(t_ptr + row * PAGES + position // BLOCK_SIZE, mask=mask, other=0) + offset = position % BLOCK_SIZE + dim = tl.arange(0, DIM) + # [BLOCK_N, DIM] + keys = tl.load( + k_ptr + + page[:, None] * k_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + for group in tl.static_range(GROUPS): + head = kv_head * GROUPS + group + query = tl.load(q_ptr + row * q_stride_row + head * DIM + dim).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) + tl.store( + out_ptr + (row * (KV * GROUPS) + head) * context + position, + scores, + mask=mask, + ) + + +def fused_capture_scores( + queries: torch.Tensor, + key_cache: torch.Tensor, + block_tables: torch.Tensor, + context_lengths: torch.Tensor, + *, + layer: int, + block_size: int, + block_n: int = 128, +) -> torch.Tensor: + """Scores of every row against its causal keys, straight from the cache. + + ``queries`` is ``[rows, heads, dim]``, ``key_cache`` is either + ``[blocks, layers, block_size, kv_heads, dim]`` or the already sliced + ``[blocks, block_size, kv_heads, dim]``, and ``block_tables`` is + ``[rows, pages]``. Returns ``[rows, heads, max_context]`` FP32 with entries + beyond each row's length left at zero. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused capture needs Triton") + if not queries.is_cuda: + raise ValueError("fused capture is a CUDA path") + + rows, heads, dim = queries.shape + if key_cache.dim() == 5: + layer_cache = key_cache[:, layer] + kv_heads = key_cache.shape[3] + else: + # The cache manager hands the attention layer a per-layer view already. + layer_cache = key_cache + kv_heads = key_cache.shape[2] + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + lengths = context_lengths.to(torch.int64) + context = int(lengths.max().item()) if lengths.numel() else 0 + if context == 0: + return queries.new_zeros((rows, heads, 0), dtype=torch.float32) + + pages = block_tables.shape[1] + out = torch.zeros(rows, heads, context, device=queries.device, dtype=torch.float32) + grid = (rows, kv_heads, triton.cdiv(context, block_n)) + _capture_scores_kernel[grid]( + queries.contiguous(), + layer_cache, + block_tables, + lengths, + out, + context, + queries.stride(0), + key_cache.stride(0), + DIM=dim, + GROUPS=groups, + KV=kv_heads, + BLOCK_SIZE=block_size, + PAGES=pages, + BLOCK_N=block_n, + ) + # Rows are ragged; the kernel leaves the tail at zero, so mask it to -inf + # before any softmax sees it. + if int(lengths.min().item()) != context: + positions = torch.arange(context, device=queries.device)[None, None, :] + out = out.masked_fill(positions >= lengths[:, None, None], float("-inf")) + return out From db6544028b896ebf377a25ab878d8e1c3bfb5f66 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Sat, 31 Oct 2026 23:40:59 +0800 Subject: [PATCH 17/28] sparse decode attention: remove the per-layer host round-trips from consume Consume was ~10 ms of fixed cost per layer at every context, so a 52-key workload cost more than dense attention. Three measured causes: - the retained count was read back from the device with counts[...].amax(...).tolist(), a host sync every layer that RFC #156 section 6.2 forbids outright. The capture decides the count before it submits, so AttentionReferenceStore now keeps it host-side (note_count/host_count) and the reader sizes its gather without a readback. - writing the source-ready flag used a Python scalar into a bool tensor through advanced indexing: 2.6 ms per call at 32 rows. Assigning a tensor instead is ~0.13 ms. - torch.tensor(rows, device=cuda) built a device tensor from a Python list on every call, a blocking copy worth 2.4 ms. Grouped rows are contiguous, so they are derived with a device-side arange. Host counts are dropped when a row is rejected, invalidated or retired, so a stale count can never size a gather. Measured at 32 rows / 24L / 16H / 128D / BF16, r=0.1, median of 10, dense Triton baseline: four-loop sparse went from 306.9 ms to 100.8 ms at 32k context (2.23x -> 1.92x dense), from 97.9 ms to 39.6 ms at 8k (6.39x -> 2.57x), and from 41.7 ms to 26.4 ms at 512 (4.06x -> 2.57x). --- vllm_rlt/attention/sparse.py | 22 ++++++++++++---- vllm_rlt/worker/attention_reference.py | 36 +++++++++++++++++++++++--- 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index cedf9343..e7c0ee80 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -430,10 +430,14 @@ def batched_paged_attention( 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") + # The retained count is host-known from the capture. Reading it back + # from the device would sync every layer, which RFC #156 section 6.2 + # forbids. + for row in consume_rows: + count = store.host_count(plan.slots[row], layer) + if not count: + raise RuntimeError("a consume row has no captured reference") + counts[row] = count grouped = {} for row, length in enumerate(lengths): if length: @@ -459,7 +463,13 @@ def batched_paged_attention( chunk = max(1, min(32, 64 * 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) + # Building a device tensor from a Python list costs ~2.4 ms (a + # blocking host-to-device copy per call). Grouped rows are almost + # always contiguous, so derive them on the device instead. + if rows[0] + len(rows) == rows[-1] + 1: + row_ids = torch.arange(rows[0], rows[0] + len(rows), device=q.device) + else: + 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] @@ -544,6 +554,8 @@ def batched_paged_attention( torch.int32 ) store.counts[slots, layer] = count + for row in rows: + store.note_count(plan.slots[row], layer, count) store.mass[slots, layer] = gathered.sum(dim=-1) store.output[slots, layer] = reference store.contribution[slots, layer] = contribution diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index 9a47c214..7f8f5d38 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -96,6 +96,10 @@ def __init__(self, *, slots, layers, heads, head_dim, key_bucket, device, stream self.free = list(range(slots)) self._bindings: dict[int, tuple[str, int, int]] = {} self._incarnations: dict[tuple[str, int], int] = {} + # Retained counts are host-known: the capture decides them before it + # submits. Keeping them here lets a reader size its gather without a + # device readback, which RFC #156 section 6.2 forbids. + self._counts_host: dict[tuple[int, int], int] = {} @property def bytes(self) -> int: @@ -221,6 +225,18 @@ def _require_live(self, plan: SparseDecodePlan) -> None: if binding is None or binding[2] != incarnation: raise RuntimeError("stale attention reference plan for a retired token") + def note_count(self, slot: int, layer: int, count: int) -> None: + """Record the retained count of one capture row on the host. + + The count is decided by the capture, so it never needs a device + readback; ``host_count`` is the reader-side counterpart. + """ + self._counts_host[(slot, layer)] = count + + def host_count(self, slot: int, layer: int) -> int | None: + """Retained keys of ``(slot, layer)``, or ``None`` before a capture.""" + return self._counts_host.get((slot, layer)) + def capture(self, plan: SparseDecodePlan, layer: int) -> None: """Mark the captured rows source-ready for the given layer. @@ -232,7 +248,11 @@ def capture(self, plan: SparseDecodePlan, layer: int) -> None: rows = plan.groups().get(CAPTURE) if not rows: return - self.valid[plan.slot_ids[rows], layer] = True + # Assign a tensor, not the Python scalar ``True``: a scalar store into + # a bool tensor through advanced indexing costs milliseconds per call + # (2.6 ms measured at 32 rows), while a tensor store is ~0.13 ms. + slots = plan.slot_ids[rows] + self.valid[slots, layer] = torch.ones(slots.numel(), dtype=torch.bool, device=self.device) def consume(self, plan: SparseDecodePlan, layer: int) -> None: """Record that the given layer consumed the reference rows.""" @@ -271,6 +291,7 @@ def write_reference( 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.note_count(slot, layer, int(count)) self.mass[slot, layer] = mass self.output[slot, layer] = output self.contribution[slot, layer] = contribution @@ -279,8 +300,8 @@ 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: + count = self.host_count(slot, layer) + if not count: raise RuntimeError("a consume row has no captured reference") return ( self.selected[slot, layer, :, :count].long(), @@ -324,6 +345,9 @@ def accept(self, leases, accepted: torch.Tensor) -> None: self.valid[rejected] = False self.counts[rejected] = 0 self.mass[rejected] = 0 + for slot in rejected.tolist(): + for layer in range(self.counts.shape[1]): + self._counts_host.pop((slot, layer), None) def _clear(self, slots: torch.Tensor) -> None: """Retire the metadata of the given slots, ordered on the owner stream. @@ -335,6 +359,9 @@ def _clear(self, slots: torch.Tensor) -> None: """ def zero(): + for slot in slots.tolist(): + for layer in range(self.counts.shape[1]): + self._counts_host.pop((slot, layer), None) self.valid[slots] = False self.uses[slots] = 0 self.accepted[slots] = 0 @@ -354,6 +381,9 @@ def invalidate(self, slots) -> None: if len(slots) == 0: return slots = torch.as_tensor(slots, dtype=torch.int64, device=self.device) + for slot in slots.tolist(): + for layer in range(self.counts.shape[1]): + self._counts_host.pop((slot, layer), None) self.valid[slots] = False self.mass[slots] = 0 self.output[slots] = 0 From 4802af6bcad3907de69cc91bd5ab5321cc1e7309 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Sun, 1 Nov 2026 20:38:10 +0800 Subject: [PATCH 18/28] sparse decode attention: fused consume and no host round-trips in the store Consume was the last item costing more than dense attention. Two changes: - kernels/sparse_consume.py: selected_scores walks the retained indices and writes one FP32 score per (row, head, retained), selected_values accumulates p * V into one output row. Neither materialises K or V, which removes the two 8 ms gathers and the per-row cuBLAS calls the Torch path made at 32k. - The store no longer builds a device tensor from a Python list. capture, consume and bind_positions all did plan.slot_ids[rows], a blocking host-to-device copy worth 2.4 ms per call, and fused_capture_scores read the context bound back with .item() twice per layer, a stream sync that RFC #156 section 6.2 forbids. The runner now passes the context and raggedness it already knows, and the slot index is derived on the device when the rows are contiguous. - The capture and consume paths read and write the store through views when the chunk's slots are contiguous, falling back to the indexed form otherwise. Measured at 32 rows / 24L / 16H / 128D / BF16, r=0.1, median of 10, dense Triton baseline. Four-loop sparse vs dense: ctx before today after fused capture now 512 4.06x 2.57x 1.38x 2048 4.43x 2.07x 1.41x 8192 6.39x 2.57x 1.95x 32768 5.86x 1.92x 1.74x Capture alone is 221.0 -> 34.1 ms at 32k and 13.5 -> 5.5 ms at 512; consume is 36.2 -> 22.0 ms and 12.8 -> 3.0 ms. The remaining gap is per-layer fixed cost, not per-key work. --- vllm_rlt/attention/sparse.py | 98 +++++++++--- vllm_rlt/kernels/sparse_capture.py | 12 +- vllm_rlt/kernels/sparse_consume.py | 209 +++++++++++++++++++++++++ vllm_rlt/worker/attention_reference.py | 18 ++- 4 files changed, 307 insertions(+), 30 deletions(-) create mode 100644 vllm_rlt/kernels/sparse_consume.py diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index e7c0ee80..f81f4d82 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -34,6 +34,8 @@ import torch from vllm_rlt.kernels.sparse_capture import TRITON_AVAILABLE, fused_capture_scores +from vllm_rlt.kernels.sparse_consume import selected_scores as fused_selected_scores +from vllm_rlt.kernels.sparse_consume import selected_values as fused_selected_values # Row modes for one recurrent submission. DENSE = 0 @@ -471,6 +473,7 @@ def batched_paged_attention( else: row_ids = torch.tensor(rows, dtype=torch.int64, device=q.device) slots = plan.slot_ids[row_ids] + chunk_slots = [plan.slots[row] for row in rows] queries = q[row_ids].float() tables = block_tables[row_ids] batch_index = torch.arange(len(rows), device=q.device)[:, None, None] @@ -480,20 +483,53 @@ def batched_paged_attention( output[row_ids] = dense_output[row_ids] consumed = True continue - 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 = _rowwise_product("hd,hkd->hk", queries, keys) * scale - probabilities = torch.softmax(scores, dim=-1) - new = _rowwise_product("hk,hkd->hd", probabilities, values) - corrected = ( - store.output[slots, layer] - - store.contribution[slots, layer] - + store.mass[slots, layer].unsqueeze(-1) * new - ) + # Contiguous slots turn these gathers into views, which is the + # difference between one index kernel per tensor and none. + if chunk_slots[0] + len(chunk_slots) == chunk_slots[-1] + 1: + span = slice(chunk_slots[0], chunk_slots[0] + len(chunk_slots)) + selected = store.selected[span, layer, :, :extent].to(torch.int32) + ref_output = store.output[span, layer] + ref_contribution = store.contribution[span, layer] + ref_mass = store.mass[span, layer] + else: + columns = torch.arange(extent, device=q.device)[None, None, :] + selected = store.selected[slots[:, None, None], layer, head_index, columns] + ref_output = store.output[slots, layer] + ref_contribution = store.contribution[slots, layer] + ref_mass = store.mass[slots, layer] + if fused_scores: + # Walk the retained indices in place: no dense K/V gather + # and no per-row cuBLAS. + scores = ( + fused_selected_scores( + q[row_ids], + key_cache, + selected, + tables, + layer=layer, + block_size=block_size, + ) + * scale + ) + probabilities = torch.softmax(scores, dim=-1) + new = fused_selected_values( + probabilities, + value_cache, + selected, + tables, + layer=layer, + block_size=block_size, + ) + else: + selected = selected.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 = _rowwise_product("hd,hkd->hk", queries, keys) * scale + probabilities = torch.softmax(scores, dim=-1) + new = _rowwise_product("hk,hkd->hd", probabilities, values) + corrected = ref_output - ref_contribution + ref_mass.unsqueeze(-1) * new output[row_ids] = corrected.to(q.dtype) consumed = True continue @@ -506,7 +542,10 @@ def batched_paged_attention( if fused_scores and dense_output is not None: # One pass over the paged cache: no expanded FP32 K row, and # no per-row cuBLAS. The dot product is still FP32; only the - # reduction order differs from the oracle. + # reduction order differs from the oracle. The context bound is + # host-known here, so the kernel never reads it back. + chunk_lengths = [lengths[row] for row in rows] + chunk_context = max(chunk_lengths) scores = fused_capture_scores( q[row_ids], key_cache, @@ -514,6 +553,8 @@ def batched_paged_attention( context_lengths[row_ids], layer=layer, block_size=block_size, + context=chunk_context, + ragged=min(chunk_lengths) != chunk_context, ) scores = scores * scale else: @@ -547,18 +588,25 @@ def batched_paged_attention( kv_index, ].float() 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. - store.selected[slots[:, None, None], layer, head_index, columns] = selected.to( - torch.int32 - ) - store.counts[slots, layer] = count + # Only the live prefix is written, never the whole capacity. + if chunk_slots[0] + len(chunk_slots) == chunk_slots[-1] + 1: + span = slice(chunk_slots[0], chunk_slots[0] + len(chunk_slots)) + store.selected[span, layer, :, :count] = selected.to(torch.int32) + store.counts[span, layer] = count + store.mass[span, layer] = gathered.sum(dim=-1) + store.output[span, layer] = reference + store.contribution[span, layer] = contribution + else: + columns = torch.arange(count, device=q.device)[None, None, :] + 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 for row in rows: store.note_count(plan.slots[row], 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) diff --git a/vllm_rlt/kernels/sparse_capture.py b/vllm_rlt/kernels/sparse_capture.py index 88135673..e9671f35 100644 --- a/vllm_rlt/kernels/sparse_capture.py +++ b/vllm_rlt/kernels/sparse_capture.py @@ -87,6 +87,8 @@ def fused_capture_scores( layer: int, block_size: int, block_n: int = 128, + context: int | None = None, + ragged: bool | None = None, ) -> torch.Tensor: """Scores of every row against its causal keys, straight from the cache. @@ -113,7 +115,11 @@ def fused_capture_scores( raise ValueError("query heads must be a multiple of KV heads") groups = heads // kv_heads lengths = context_lengths.to(torch.int64) - context = int(lengths.max().item()) if lengths.numel() else 0 + # ``context`` and ``ragged`` are host-known for a runner submission. Reading + # them back with ``.item()`` would sync the stream once per layer, which + # RFC #156 section 6.2 forbids, so callers pass them in. + if context is None: + context = int(lengths.max().item()) if lengths.numel() else 0 if context == 0: return queries.new_zeros((rows, heads, 0), dtype=torch.float32) @@ -138,7 +144,9 @@ def fused_capture_scores( ) # Rows are ragged; the kernel leaves the tail at zero, so mask it to -inf # before any softmax sees it. - if int(lengths.min().item()) != context: + if ragged is None: + ragged = lengths.numel() > 0 and int(lengths.min().item()) != context + if ragged: positions = torch.arange(context, device=queries.device)[None, None, :] out = out.masked_fill(positions >= lengths[:, None, None], float("-inf")) return out diff --git a/vllm_rlt/kernels/sparse_consume.py b/vllm_rlt/kernels/sparse_consume.py new file mode 100644 index 00000000..06b7fc7b --- /dev/null +++ b/vllm_rlt/kernels/sparse_consume.py @@ -0,0 +1,209 @@ +"""Fused gather-and-accumulate for the sparse-decode consume step. + +A consume row only touches its ``k`` retained keys, but the Torch path gathers +them into a dense ``[rows, heads, k, dim]`` tensor twice (K then V) and runs one +cuBLAS call per row. At 32k context with ``k = 3277`` that is 8.2 ms of K gather +plus 8.1 ms of V gather through three advanced indices on a strided +``[blocks, block_size, kv, dim]`` cache, which measures about 45 GB/s. + +These two kernels walk the retained indices in place: the first writes one FP32 +score per ``(row, head, retained)``, the second accumulates ``p * V`` into one +output row. Nothing dense is materialised, and the softmax in between stays in +Torch because ``k`` is small. +""" + +from __future__ import annotations + +import torch + +try: # pragma: no cover - exercised only without Triton + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: # pragma: no cover + triton = None + tl = None + TRITON_AVAILABLE = False + + +if TRITON_AVAILABLE: + + @triton.jit + def _selected_scores_kernel( + q_ptr, + k_ptr, + idx_ptr, + t_ptr, + out_ptr, + retained, + k_stride_block, + DIM: tl.constexpr, + KV: tl.constexpr, + GROUPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + block = tl.program_id(2) + heads: tl.constexpr = KV * GROUPS + kv_head = head // GROUPS + + slot = block * BLOCK_K + tl.arange(0, BLOCK_K) + mask = slot < retained + index = tl.load(idx_ptr + (row * heads + head) * retained + slot, mask=mask, other=0) + page = tl.load(t_ptr + row * PAGES + index // BLOCK_SIZE, mask=mask, other=0) + offset = index % BLOCK_SIZE + dim = tl.arange(0, DIM) + keys = tl.load( + k_ptr + + page[:, None] * k_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + query = tl.load(q_ptr + (row * heads + head) * DIM + dim).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) + tl.store(out_ptr + (row * heads + head) * retained + slot, scores, mask=mask) + + @triton.jit + def _selected_values_kernel( + p_ptr, + v_ptr, + idx_ptr, + t_ptr, + out_ptr, + retained, + v_stride_block, + DIM: tl.constexpr, + KV: tl.constexpr, + GROUPS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + PAGES: tl.constexpr, + BLOCK_K: tl.constexpr, + ): + row = tl.program_id(0) + head = tl.program_id(1) + heads: tl.constexpr = KV * GROUPS + kv_head = head // GROUPS + dim = tl.arange(0, DIM) + accumulator = tl.zeros([DIM], dtype=tl.float32) + for start in range(0, retained, BLOCK_K): + slot = start + tl.arange(0, BLOCK_K) + mask = slot < retained + index = tl.load(idx_ptr + (row * heads + head) * retained + slot, mask=mask, other=0) + weight = tl.load( + p_ptr + (row * heads + head) * retained + slot, mask=mask, other=0.0 + ).to(tl.float32) + page = tl.load(t_ptr + row * PAGES + index // BLOCK_SIZE, mask=mask, other=0) + offset = index % BLOCK_SIZE + values = tl.load( + v_ptr + + page[:, None] * v_stride_block + + (offset[:, None] * KV + kv_head) * DIM + + dim[None, :], + mask=mask[:, None], + other=0.0, + ).to(tl.float32) + accumulator += tl.sum(weight[:, None] * values, axis=0) + tl.store(out_ptr + (row * heads + head) * DIM + dim, accumulator) + + +def _cache_view(cache: torch.Tensor, layer: int) -> tuple[torch.Tensor, int]: + """Return the per-layer cache view and its KV head count.""" + if cache.dim() == 5: + return cache[:, layer], cache.shape[3] + return cache, cache.shape[2] + + +def selected_scores( + queries: torch.Tensor, + key_cache: torch.Tensor, + selected: torch.Tensor, + block_tables: torch.Tensor, + *, + layer: int, + block_size: int, + block_k: int = 128, +) -> torch.Tensor: + """``q . K[S]`` for every retained key, without materialising K. + + ``queries`` is ``[rows, heads, dim]``, ``selected`` is ``[rows, heads, k]`` + of absolute positions and the result is ``[rows, heads, k]`` FP32. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused consume needs Triton") + rows, heads, dim = queries.shape + cache, kv_heads = _cache_view(key_cache, layer) + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + retained = selected.shape[-1] + out = torch.empty(rows, heads, retained, device=queries.device, dtype=torch.float32) + if retained == 0: + return out + grid = (rows, heads, triton.cdiv(retained, block_k)) + _selected_scores_kernel[grid]( + queries.contiguous(), + cache, + selected.contiguous().to(torch.int32), + block_tables, + out, + retained, + cache.stride(0), + DIM=dim, + KV=kv_heads, + GROUPS=groups, + BLOCK_SIZE=block_size, + PAGES=block_tables.shape[1], + BLOCK_K=block_k, + ) + return out + + +def selected_values( + probabilities: torch.Tensor, + value_cache: torch.Tensor, + selected: torch.Tensor, + block_tables: torch.Tensor, + *, + layer: int, + block_size: int, + block_k: int = 128, +) -> torch.Tensor: + """``sum_j p[j] * V[S[j]]`` for every row and head. + + ``probabilities`` is ``[rows, heads, k]`` and the result is + ``[rows, heads, dim]`` FP32. + """ + if not TRITON_AVAILABLE: + raise RuntimeError("fused consume needs Triton") + rows, heads, retained = probabilities.shape + cache, kv_heads = _cache_view(value_cache, layer) + if heads % kv_heads: + raise ValueError("query heads must be a multiple of KV heads") + groups = heads // kv_heads + dim = cache.shape[-1] + out = torch.empty(rows, heads, dim, device=probabilities.device, dtype=torch.float32) + if retained == 0: + return out + grid = (rows, heads) + _selected_values_kernel[grid]( + probabilities.contiguous(), + cache, + selected.contiguous().to(torch.int32), + block_tables, + out, + retained, + cache.stride(0), + DIM=dim, + KV=kv_heads, + GROUPS=groups, + BLOCK_SIZE=block_size, + PAGES=block_tables.shape[1], + BLOCK_K=block_k, + ) + return out diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index 7f8f5d38..3555925f 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -237,6 +237,18 @@ def host_count(self, slot: int, layer: int) -> int | None: """Retained keys of ``(slot, layer)``, or ``None`` before a capture.""" return self._counts_host.get((slot, layer)) + def _row_slots(self, plan: SparseDecodePlan, rows: list[int]) -> torch.Tensor: + """Slot ids of ``rows`` without a host-to-device copy of a Python list. + + ``plan.slot_ids[rows]`` costs ~2.4 ms per call because the list becomes + a device tensor first; grouped rows are contiguous, so derive the index + on the device instead. + """ + if rows[0] + len(rows) == rows[-1] + 1: + index = torch.arange(rows[0], rows[0] + len(rows), device=self.device) + return plan.slot_ids.index_select(0, index) + return plan.slot_ids[rows] + def capture(self, plan: SparseDecodePlan, layer: int) -> None: """Mark the captured rows source-ready for the given layer. @@ -251,7 +263,7 @@ def capture(self, plan: SparseDecodePlan, layer: int) -> None: # Assign a tensor, not the Python scalar ``True``: a scalar store into # a bool tensor through advanced indexing costs milliseconds per call # (2.6 ms measured at 32 rows), while a tensor store is ~0.13 ms. - slots = plan.slot_ids[rows] + slots = self._row_slots(plan, rows) self.valid[slots, layer] = torch.ones(slots.numel(), dtype=torch.bool, device=self.device) def consume(self, plan: SparseDecodePlan, layer: int) -> None: @@ -260,7 +272,7 @@ def consume(self, plan: SparseDecodePlan, layer: int) -> None: rows = plan.groups().get(CONSUME) if not rows: return - self.uses[plan.slot_ids[rows], layer] += 1 + self.uses[self._row_slots(plan, rows), layer] += 1 def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> None: """Record the device-owned absolute positions of the plan's rows.""" @@ -268,7 +280,7 @@ def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> Non if positions.shape != plan.slot_ids.shape: raise ValueError("reference positions need one entry per plan row") rows = [row for row, slot in enumerate(plan.slots) if slot >= 0] - self.position[plan.slot_ids[rows]] = positions[rows].to(torch.int64) + self.position[self._row_slots(plan, rows)] = positions[rows].to(torch.int64) def write_reference( self, From fe83818fb9308161f15c14bcdfbbfb774b5beb20 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Sun, 1 Nov 2026 20:46:14 +0800 Subject: [PATCH 19/28] benchmarks: engine-level e2e for sparse decode attention Runs the real StageEngine with asynchronous self-speculation and compares dense against sparse decode at full and reduced retention: throughput, peak memory and per-request token agreement, across concurrency 1/8/32. The result is a correctness pass and a performance diagnosis. Full retention reproduces dense tokens exactly at every concurrency (1/1, 8/8, 32/32), so the merged O_ref contract holds through scheduling, sampling and delivery. Throughput is 5-19x worse on the tiny model and the gap widens with concurrency, which shows the remaining cost is per-layer Python fixed cost scaling with rows, not per-key attention work: with two layers and 48-token contexts there is nothing to amortise it against. A serving answer needs a trained Ouro-1.4B checkpoint, which this host does not have. --- benchmarks/sparse_decode_engine_e2e.py | 142 +++++++++++++++++++++++++ 1 file changed, 142 insertions(+) create mode 100644 benchmarks/sparse_decode_engine_e2e.py diff --git a/benchmarks/sparse_decode_engine_e2e.py b/benchmarks/sparse_decode_engine_e2e.py new file mode 100644 index 00000000..3bd78e41 --- /dev/null +++ b/benchmarks/sparse_decode_engine_e2e.py @@ -0,0 +1,142 @@ +"""Engine-level end-to-end comparison for sparse decode attention. + +Runs the real StageEngine with asynchronous self-speculation and compares dense +against sparse decode at full and reduced retention: throughput, peak memory and +token agreement. The per-layer sweeps say what one attention call costs; this +says whether that survives scheduling, sampling and delivery. + +The tiny Ouro architecture is used because no trained checkpoint is available on +this host, so the numbers are a scheduler-level check, not a serving result. +""" + +import argparse +import statistics +import time + +import torch + +from tests.helpers import tiny_ouro_model +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.engine.stage_engine import StageEngine + + +def build(model, concurrency, sparse): + return StageEngine( + model, + attention_backend="triton", + cache_config=CacheConfig(2048, 16, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=concurrency, + max_num_batched_tokens=8 * concurrency, + prefill_chunk_size=16, + ), + execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), + speculative_config=SpeculativeConfig(3), + sparse_decode_attention_config=sparse, + ) + + +def run(model, concurrency, sparse, *, prompt, max_tokens, repeats): + """One configuration: returns median wall time, tokens/s and peak bytes.""" + outputs = None + walls = [] + peaks = [] + for _ in range(repeats): + engine = build(model, concurrency, sparse) + torch.cuda.reset_peak_memory_stats() + params = [ + SamplingParams(max_tokens=max_tokens, ignore_eos=True) for _ in range(concurrency) + ] + prompts = [ + [2 + (index + position) % 60 for position in range(prompt)] + for index in range(concurrency) + ] + start = time.perf_counter() + for index in range(concurrency): + engine.add_request(f"r{index}", prompts[index], params[index]) + collected = {} + for _ in range(100000): + for output in engine.step(): + collected[output.request_id] = output + if len(collected) == concurrency and all(o.finished for o in collected.values()): + break + torch.cuda.synchronize() + walls.append(time.perf_counter() - start) + peaks.append(torch.cuda.max_memory_allocated()) + outputs = {key: list(value.token_ids) for key, value in collected.items()} + tokens = concurrency * max_tokens + return { + "wall": statistics.median(walls), + "tokens_per_s": tokens / statistics.median(walls), + "peak_gib": max(peaks) / 2**30, + "outputs": outputs, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--device", default="cuda") + parser.add_argument("--prompt", type=int, default=32) + parser.add_argument("--max-tokens", type=int, default=16) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--concurrency", default="1,8,32") + parser.add_argument("--retentions", default="1.0,0.1") + args = parser.parse_args() + + device = args.device + model = tiny_ouro_model(device) + concurrency_levels = [int(value) for value in args.concurrency.split(",")] + retentions = [float(value) for value in args.retentions.split(",")] + + print( + f"engine e2e — tiny Ouro ({model.config.num_hidden_layers}L, ctx " + f"{model.config.max_position_embeddings}), prompt {args.prompt}, " + f"max_tokens {args.max_tokens}, median of {args.repeats}" + ) + header = f"{'conc':>5}{'case':>14}{'tok/s':>10}{'wall ms':>10}{'peak GiB':>10}{'match':>10}" + print(header) + for concurrency in concurrency_levels: + dense = run( + model, + concurrency, + None, + prompt=args.prompt, + max_tokens=args.max_tokens, + repeats=args.repeats, + ) + print( + f"{concurrency:>5}{'dense':>14}{dense['tokens_per_s']:>10.2f}" + f"{dense['wall'] * 1e3:>10.1f}{dense['peak_gib']:>10.2f}{'--':>10}" + ) + for retention in retentions: + sparse = run( + model, + concurrency, + SparseDecodeAttentionConfig( + enabled=True, key_retention=retention, max_context_tokens=64 + ), + prompt=args.prompt, + max_tokens=args.max_tokens, + repeats=args.repeats, + ) + same = sum( + 1 + for key, tokens in sparse["outputs"].items() + if tokens == dense["outputs"].get(key) + ) + print( + f"{'':>5}{f'sparse r={retention}':>14}{sparse['tokens_per_s']:>10.2f}" + f"{sparse['wall'] * 1e3:>10.1f}{sparse['peak_gib']:>10.2f}" + f"{f'{same}/{concurrency}':>10}" + ) + + +if __name__ == "__main__": + main() From e128c04ed277b4d84ff8843a2e4e807a1517acf6 Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Mon, 2 Nov 2026 00:47:06 +0800 Subject: [PATCH 20/28] benchmarks: real-text engine e2e for a trained checkpoint Token-id prompts drive a trained model into a fixed point, so every implementation agrees and the comparison says nothing. This encodes real text with the checkpoint's tokenizer and reports throughput, peak memory and per-request token agreement against the dense engine. Measured on Ouro-1.4B (24 layers, 16 KV heads, 200-token prompt, 16 tokens out, median of 2, A100): concurrency 1 dense 4.77 tok/s, sparse r=1.0 3.74, r=0.1 3.87; concurrency 8 dense 42.75, r=1.0 17.27, r=0.1 22.01. Full retention reproduces dense tokens exactly at both concurrencies. --- benchmarks/sparse_decode_engine_e2e_text.py | 159 ++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 benchmarks/sparse_decode_engine_e2e_text.py diff --git a/benchmarks/sparse_decode_engine_e2e_text.py b/benchmarks/sparse_decode_engine_e2e_text.py new file mode 100644 index 00000000..734bbd9c --- /dev/null +++ b/benchmarks/sparse_decode_engine_e2e_text.py @@ -0,0 +1,159 @@ +"""Engine e2e with a real-text prompt, for a trained checkpoint. + +Synthetic token-id prompts drive a trained model into a fixed point (one token +repeated), which makes token agreement meaningless: every implementation agrees +when the argmax is constant. This encodes real text with the checkpoint's own +tokenizer and reports throughput, peak memory and per-request token agreement. + +Usage: + python benchmarks/sparse_decode_engine_e2e_text.py \ + --model-path /path/to/Ouro-1.4B --concurrency 1,8 +""" + +import argparse +import statistics +import time + +import torch +from tokenizers import Tokenizer + +from vllm_rlt import ( + CacheConfig, + ExecutionConfig, + SamplingParams, + SchedulerConfig, + SparseDecodeAttentionConfig, + SpeculativeConfig, +) +from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.models import AutoModelForCausalLM + +TEXT = ( + "The history of computing is a history of abstraction. Each generation of " + "hardware made a new kind of program possible, and each new kind of program " + "demanded a new way of thinking about machines. When memory was scarce, " + "programmers counted bytes by hand; when it became plentiful, they wrote " + "interpreters that treated memory as infinite. The same shift is happening " + "again with inference: models that once fit on a single device now span " + "racks, and the scheduler rather than the kernel decides how fast a request " + "finishes. Attention is the expensive part, because every token must look at " + "every earlier token, and that cost grows with the square of the context. " + "Sparse attention tries to keep only the keys that matter, which is easy to " + "state and hard to make fast, because deciding what matters is itself as " + "expensive as attending to everything." +) + + +def build(model, concurrency, sparse, blocks): + return StageEngine( + model, + attention_backend="triton", + cache_config=CacheConfig(blocks, 16, incremental_allocation=True), + scheduler_config=SchedulerConfig( + max_num_seqs=concurrency, + max_num_batched_tokens=8 * concurrency, + prefill_chunk_size=16, + ), + execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), + speculative_config=SpeculativeConfig(3), + sparse_decode_attention_config=sparse, + ) + + +def run(model, concurrency, sparse, *, prompt_ids, max_tokens, repeats, blocks): + walls = [] + peaks = [] + outputs = None + for _ in range(repeats): + engine = build(model, concurrency, sparse, blocks) + torch.cuda.reset_peak_memory_stats() + params = [ + SamplingParams(max_tokens=max_tokens, ignore_eos=True) for _ in range(concurrency) + ] + start = time.perf_counter() + for index in range(concurrency): + engine.add_request(f"r{index}", list(prompt_ids), params[index]) + collected = {} + for _ in range(100000): + for output in engine.step(): + collected[output.request_id] = output + if len(collected) == concurrency and all(o.finished for o in collected.values()): + break + torch.cuda.synchronize() + walls.append(time.perf_counter() - start) + peaks.append(torch.cuda.max_memory_allocated()) + outputs = {key: list(value.token_ids) for key, value in collected.items()} + tokens = concurrency * max_tokens + return { + "tokens_per_s": tokens / statistics.median(walls), + "wall": statistics.median(walls), + "peak_gib": max(peaks) / 2**30, + "outputs": outputs, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", required=True) + parser.add_argument("--prompt", type=int, default=200) + parser.add_argument("--max-tokens", type=int, default=16) + parser.add_argument("--repeats", type=int, default=2) + parser.add_argument("--concurrency", default="1,8") + parser.add_argument("--blocks", type=int, default=768) + parser.add_argument("--sparse-context", type=int, default=512) + args = parser.parse_args() + + model = AutoModelForCausalLM.from_pretrained( + args.model_path, device="cuda", dtype=torch.bfloat16 + ) + tokenizer = Tokenizer.from_file(f"{args.model_path}/tokenizer.json") + ids = tokenizer.encode(TEXT).ids + while len(ids) < args.prompt: + ids = ids + ids + prompt_ids = ids[: args.prompt] + print(f"real text prompt: {args.prompt} tokens, head={prompt_ids[:8]}") + + print(f"{'conc':>5}{'case':>16}{'tok/s':>10}{'wall ms':>10}{'peak GiB':>10}{'match':>10}") + for concurrency in [int(v) for v in args.concurrency.split(",")]: + dense = run( + model, + concurrency, + None, + prompt_ids=prompt_ids, + max_tokens=args.max_tokens, + repeats=args.repeats, + blocks=args.blocks, + ) + print( + f"{concurrency:>5}{'dense':>16}{dense['tokens_per_s']:>10.2f}" + f"{dense['wall'] * 1e3:>10.1f}{dense['peak_gib']:>10.2f}{'--':>10}" + ) + for retention in (1.0, 0.1): + sparse = run( + model, + concurrency, + SparseDecodeAttentionConfig( + enabled=True, + key_retention=retention, + max_context_tokens=args.sparse_context, + ), + prompt_ids=prompt_ids, + max_tokens=args.max_tokens, + repeats=args.repeats, + blocks=args.blocks, + ) + same = sum( + 1 + for key, tokens in sparse["outputs"].items() + if tokens == dense["outputs"].get(key) + ) + print( + f"{'':>5}{f'sparse r={retention}':>16}{sparse['tokens_per_s']:>10.2f}" + f"{sparse['wall'] * 1e3:>10.1f}{sparse['peak_gib']:>10.2f}" + f"{f'{same}/{concurrency}':>10}" + ) + print(f" dense sample tokens: {dense['outputs']['r0'][:8]}") + + +if __name__ == "__main__": + main() From 37e294a467a25a1cc46f3f9fab474c93a2ca401d Mon Sep 17 00:00:00 2001 From: kunkunblueberry <1833921874@qq.com> Date: Mon, 2 Nov 2026 00:47:46 +0800 Subject: [PATCH 21/28] benchmarks: let the engine e2e load a trained checkpoint Adds --model-path, --cache-blocks and --sparse-context so the same harness can run against Ouro-1.4B. The KV pool has to be sized to the workload: 4096 blocks at 24 layers, 16 heads and 128 dims is 51 GiB of K and V, which OOMs; 768 blocks fits a 32-request, 200-token burst. --- benchmarks/sparse_decode_engine_e2e.py | 28 +++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/benchmarks/sparse_decode_engine_e2e.py b/benchmarks/sparse_decode_engine_e2e.py index 3bd78e41..edd330d6 100644 --- a/benchmarks/sparse_decode_engine_e2e.py +++ b/benchmarks/sparse_decode_engine_e2e.py @@ -25,13 +25,14 @@ SpeculativeConfig, ) from vllm_rlt.engine.stage_engine import StageEngine +from vllm_rlt.models import AutoModelForCausalLM -def build(model, concurrency, sparse): +def build(model, concurrency, sparse, blocks=2048): return StageEngine( model, attention_backend="triton", - cache_config=CacheConfig(2048, 16, incremental_allocation=True), + cache_config=CacheConfig(blocks, 16, incremental_allocation=True), scheduler_config=SchedulerConfig( max_num_seqs=concurrency, max_num_batched_tokens=8 * concurrency, @@ -43,13 +44,13 @@ def build(model, concurrency, sparse): ) -def run(model, concurrency, sparse, *, prompt, max_tokens, repeats): +def run(model, concurrency, sparse, *, prompt, max_tokens, repeats, blocks=2048): """One configuration: returns median wall time, tokens/s and peak bytes.""" outputs = None walls = [] peaks = [] for _ in range(repeats): - engine = build(model, concurrency, sparse) + engine = build(model, concurrency, sparse, blocks) torch.cuda.reset_peak_memory_stats() params = [ SamplingParams(max_tokens=max_tokens, ignore_eos=True) for _ in range(concurrency) @@ -88,15 +89,24 @@ def main(): parser.add_argument("--repeats", type=int, default=3) parser.add_argument("--concurrency", default="1,8,32") parser.add_argument("--retentions", default="1.0,0.1") + parser.add_argument("--model-path", default=None) + parser.add_argument("--cache-blocks", type=int, default=2048) + parser.add_argument("--sparse-context", type=int, default=64) args = parser.parse_args() device = args.device - model = tiny_ouro_model(device) + if args.model_path: + model = AutoModelForCausalLM.from_pretrained( + args.model_path, device=device, dtype=torch.bfloat16 + ) + else: + model = tiny_ouro_model(device) concurrency_levels = [int(value) for value in args.concurrency.split(",")] retentions = [float(value) for value in args.retentions.split(",")] + label = args.model_path or "tiny Ouro" print( - f"engine e2e — tiny Ouro ({model.config.num_hidden_layers}L, ctx " + f"engine e2e — {label} ({model.config.num_hidden_layers}L, ctx " f"{model.config.max_position_embeddings}), prompt {args.prompt}, " f"max_tokens {args.max_tokens}, median of {args.repeats}" ) @@ -110,6 +120,7 @@ def main(): prompt=args.prompt, max_tokens=args.max_tokens, repeats=args.repeats, + blocks=args.cache_blocks, ) print( f"{concurrency:>5}{'dense':>14}{dense['tokens_per_s']:>10.2f}" @@ -120,11 +131,14 @@ def main(): model, concurrency, SparseDecodeAttentionConfig( - enabled=True, key_retention=retention, max_context_tokens=64 + enabled=True, + key_retention=retention, + max_context_tokens=args.sparse_context, ), prompt=args.prompt, max_tokens=args.max_tokens, repeats=args.repeats, + blocks=args.cache_blocks, ) same = sum( 1 From 21c28142312baf5f0ed985a263f4be678f5852f5 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 1 Nov 2026 17:54:57 +0100 Subject: [PATCH 22/28] sparse decode attention: opt-in fused kernels for capture and consume SparseDecodeAttentionConfig(implementation="fused") routes uniform capture or consume submissions to Triton kernels with no per-layer host readbacks; the default batched path and its oracle contract are unchanged, and it keeps the fused consume from 4802af6. - Capture is the dense Triton split kernel that also keeps the scores, so O_ref matches the dense output (bit for bit in BF16; within one ulp in FP32 on sm_90). It builds the first radix histogram per part; selection is a 4 x 8-bit radix select over 16 parts per (row, head), each pass closed in a fixed order by the last part to finish (integer tickets), writing positions only. The retained contribution is gathered over the compact list. Backends other than Triton keep their call plus one serial selection kernel. - Consume is split-K over the retained ranks with a fixed-order merge. Its per-stream scratch follows rows, heads and the padded head dimension. - Retained counts are decided on the device; a batched consume after a fused capture reads them back once. - The store writes its metadata without host-scalar copies: one row-slot helper (the whole plan reuses its device slot ids, contiguous rows use a device arange, other rows a pinned copy), index_fill_ retirement, a readback-free masked accept that drops the round's host counts, and index storage that grows on demand within the logical bucket (the engine budgets the growth peak). - Reused launches call the compiled Triton 3.8 kernel directly, skipping its Python dispatch. The key function is generated per argument-type layout and mirrors Triton's specialization; the tests assert on every reuse that it is the kernel Triton would pick. Measured with benchmarks/sparse_decode_context_sweep.py (32 rows, r=0.1, RTX PRO 6000): four-loop sparse/dense 0.66-0.73x at >= 8k, 1.1x at 2k, 1.6-1.7x at 512. Eager GSM8K-87 at batch 1 (RTX 5090): 1.064x dense wall time, 61 vs 59 correct. Signed-off-by: 0z5a --- tests/test_sparse_decode_attention.py | 17 + tests/test_sparse_decode_batched.py | 102 +- tests/test_sparse_decode_fused.py | 523 +++++++++ vllm_rlt/attention/sparse.py | 74 +- vllm_rlt/config.py | 6 + vllm_rlt/engine/stage_engine.py | 17 +- vllm_rlt/kernels/sparse_decode.py | 1410 ++++++++++++++++++++++++ vllm_rlt/worker/attention_reference.py | 179 ++- 8 files changed, 2251 insertions(+), 77 deletions(-) create mode 100644 tests/test_sparse_decode_fused.py create mode 100644 vllm_rlt/kernels/sparse_decode.py diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index 7d359517..4c4fc4c0 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -234,6 +234,23 @@ def test_accept_invalidates_rejected_candidates_and_descendants(): store.read_reference(rejected_plan, layer=0, row=0) +def test_accept_drops_host_counts_so_rejected_rows_cannot_size_a_read(): + store = make_store() + lease = store.reserve("r", 0, 3, generation=1) + fill_payload(store, lease.slots) + for slot in lease.slots: + for layer in range(store.layers): + store.note_count(slot, layer, 1) + store.accept([lease], torch.tensor([1])) + assert all(store.host_count(slot, 0) is None for slot in lease.slots) + # The kept row still reads its device count; the rejected row cannot be read. + kept = store.build_plan([(lease, 1)], [CONSUME], key_retention=0.5) + assert store.read_reference(kept, layer=0, row=0)[0].shape[-1] == 1 + rejected = store.build_plan([(lease, 2)], [CONSUME], key_retention=0.5) + with pytest.raises(RuntimeError, match="no captured reference"): + store.read_reference(rejected, layer=0, row=0) + + def test_accept_keeps_root_when_every_candidate_is_rejected(): store = make_store() lease = store.reserve("r", 0, 3, generation=1) diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py index 663060ca..17e8bc67 100644 --- a/tests/test_sparse_decode_batched.py +++ b/tests/test_sparse_decode_batched.py @@ -17,7 +17,10 @@ @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): +@pytest.mark.parametrize("use_dense_callback", [False, True]) +def test_batched_ragged_gqa_matches_scalar_oracle( + device, dtype, retention, ties, use_dense_callback +): 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) @@ -28,8 +31,16 @@ def test_batched_ragged_gqa_matches_scalar_oracle(device, dtype, retention, ties 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) + make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=257, + device=device, + initial_key_capacity=1 if implementation else None, + ) + for implementation in range(2) ] leases = [store.reserve("test", 0, rows, generation=1) for store in stores] implementations = [reference_paged_attention, batched_paged_attention] @@ -39,9 +50,26 @@ def test_batched_ragged_gqa_matches_scalar_oracle(device, dtype, retention, ties 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)) + # A deliberately distinct backend output verifies that capture + # stores the rounded recurrence input, not its FP32 oracle PV. + callback = (lambda q, *args: q * 0.375 + 1.25) if use_dense_callback else None + results.append( + implementation( + queries, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=callback, + ) + ) torch.testing.assert_close(results[1], results[0], rtol=0, atol=0) - for name in ("selected", "counts", "valid", "uses"): + capacity = stores[1].selected.shape[-1] + assert torch.equal(stores[0].selected[..., :capacity], stores[1].selected) + assert not torch.count_nonzero(stores[0].selected[..., capacity:]) + for name in ("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( @@ -73,3 +101,67 @@ def test_batched_rejects_missing_and_stale_reference(): store.reserve("test", 0, 1, generation=2) with pytest.raises(RuntimeError, match="stale attention reference"): batched_paged_attention(*args, layer=0) + + +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.gpu)]) +def test_reference_growth_preserves_live_layers_and_budget(device): + store = make_store( + slots=3, + layers=2, + heads=2, + head_dim=4, + key_bucket=17, + initial_key_capacity=2, + device=device, + ) + leases = [store.reserve(str(i), i, 1, generation=1) for i in range(2)] + plans = [store.build_plan([(lease, 0)], [CAPTURE], key_retention=1.0) for lease in leases] + # Simulate queued work on the runner's owner stream before reallocating. + if device == "cuda": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + store.stream = stream + context = torch.cuda.stream(stream) + else: + from contextlib import nullcontext + + context = nullcontext() + budget = store.peak_capacity_bytes + fixed = store.bytes - store.selected.nbytes + assert budget == fixed + 3 * 2 * 2 * (2 + 4 + 8 + 16 + 17) * 4 + with context: + for layer in range(2): + store.write_reference( + plans[0], + layer, + 0, + selected=torch.tensor([[1, 3], [2, 4]], device=device), + counts=torch.full((2,), 2, device=device), + mass=torch.full((2,), 0.5, device=device), + output=torch.full((2, 4), 0.25, device=device), + contribution=torch.full((2, 4), 0.125, device=device), + ) + before = [x.clone() for x in store.read_reference(plans[0], 1, 0)] + # Grow several times while a different lease remains live. + for count in [3, 7, 17]: + store.ensure_key_capacity(count) + assert store.selected.shape[-1] >= count + assert store.peak_capacity_bytes == budget + for actual, expected in zip(store.read_reference(plans[0], 1, 0), before): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + with pytest.raises(ValueError, match="capacity"): + store.ensure_key_capacity(18) + if device == "cuda": + store.stream.synchronize() + for lease in leases: + store.release(lease) + assert store.outstanding == 0 + + +@pytest.mark.gpu +def test_reference_growth_rejects_foreign_stream(): + store = make_store(key_bucket=17, initial_key_capacity=2, device="cuda") + store.stream = torch.cuda.Stream() + with pytest.raises(RuntimeError, match="owner stream"): + store.ensure_key_capacity(3) + assert store.selected.shape[-1] == 2 diff --git a/tests/test_sparse_decode_fused.py b/tests/test_sparse_decode_fused.py new file mode 100644 index 00000000..c1e35db8 --- /dev/null +++ b/tests/test_sparse_decode_fused.py @@ -0,0 +1,523 @@ +"""Opt-in fused sparse decode kernels versus the batched oracle-shaped path.""" + +import pytest +import torch + +from tests.helpers import make_speculative_engine, tiny_ouro_model +from tests.test_sparse_decode_attention import drain, make_store +from vllm_rlt import SamplingParams, SparseDecodeAttentionConfig +from vllm_rlt.attention.sparse import ( + CAPTURE, + CONSUME, + reference_paged_attention, + sparse_paged_attention, +) + + +def test_sparse_config_validates_implementation(): + assert SparseDecodeAttentionConfig().implementation == "batched" + assert SparseDecodeAttentionConfig(implementation="fused").implementation == "fused" + with pytest.raises(ValueError, match="implementation"): + SparseDecodeAttentionConfig(implementation="native") + + +def test_fused_store_keeps_the_oracle_on_cpu(monkeypatch): + calls = [] + import vllm_rlt.attention.sparse as sparse + + monkeypatch.setattr(sparse, "reference_paged_attention", lambda *a, **k: calls.append(1)) + store = make_store(slots=1, layers=1, heads=2, head_dim=4, implementation="fused") + lease = store.reserve("cpu", 0, 1, generation=1) + plan = store.build_plan([(lease, 0)], [CAPTURE], key_retention=0.5) + q = torch.ones(1, 2, 4) + cache = torch.ones(1, 16, 2, 4) + sparse.sparse_paged_attention( + q, cache, cache, torch.zeros(1, 1, dtype=torch.int32), torch.tensor([3]), plan, layer=0 + ) + assert calls == [1] + + +@pytest.fixture(autouse=True) +def _check_fast_launches(monkeypatch): + """Every reused launch must be the kernel Triton's own dispatch would pick.""" + try: + from vllm_rlt.kernels import sparse_decode + except ImportError: # CPU-only installs have no Triton and no fused launches + return + monkeypatch.setattr(sparse_decode, "_LAUNCH_CHECK", True) + + +@pytest.mark.gpu +@pytest.mark.parametrize("context", [37, 1024, 4111]) +def test_fast_launches_match_triton_dispatch_bit_for_bit(context, monkeypatch): + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + if not sparse_decode._FAST_LAUNCH: + pytest.skip("fast launches are limited to Triton 3.8") + rows, heads, kv_heads, dim, block_size, layers = 3, 4, 2, 64, 16, 18 + pages = -(-context // block_size) + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.bfloat16, 3 + ) + lengths = torch.tensor([context, max(1, context - 17), max(1, context // 2)], device="cuda") + results = {} + for fast in (False, True): + monkeypatch.setattr(sparse_decode, "_FAST_LAUNCH", fast) + store = make_store( + slots=rows, + layers=layers, + heads=heads, + head_dim=dim, + key_bucket=context, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows, generation=1) + owners = [(lease, row) for row in range(rows)] + outputs = [] + # Layers 0, 1, 16 and 17 differ in Triton's integer specialization. + for _ in range(2): + for layer in (0, 1, 2, 16, 17): + for mode in (CAPTURE, CONSUME): + plan = store.build_plan(owners, [mode] * rows, key_retention=0.1) + outputs.append( + sparse_decode.fused_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=layer, + dense_attention=TritonPagedAttention() if mode == CAPTURE else None, + ) + ) + torch.cuda.synchronize() + results[fast] = ( + outputs, + store.selected.clone(), + store.counts.clone(), + store.mass.clone(), + store.contribution.clone(), + store.output.clone(), + ) + slow, fast = results[False], results[True] + for a, b in zip(slow[0], fast[0]): + assert torch.equal(a, b) + for a, b in zip(slow[1:], fast[1:]): + assert torch.equal(a, b) + + +@pytest.mark.gpu +def test_consume_scratch_follows_the_head_dimension(monkeypatch): + """A stream's consume scratch must not be reused across head dimensions.""" + from vllm_rlt.kernels import sparse_decode + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + def consume(dim): + rows, heads, block_size, pages = 3, 4, 16, 5 + q, keys, values, tables = _paged(rows, heads, 2, dim, pages, block_size, torch.float32, 4) + lengths = torch.tensor([7, 33, 80], device="cuda", dtype=torch.int32) + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=80, + device="cuda", + implementation="fused", + ) + lease = store.reserve("r", 0, rows, generation=1) + owners = [(lease, row) for row in range(rows)] + for mode in (CAPTURE, CONSUME): + plan = store.build_plan(owners, [mode] * rows, key_retention=0.25) + output = sparse_decode.fused_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=TritonPagedAttention() if mode == CAPTURE else None, + ) + torch.cuda.synchronize() + return output + + monkeypatch.setattr(sparse_decode, "_CONSUME_WORKSPACE", {}) + fresh = consume(128) + sparse_decode._CONSUME_WORKSPACE.clear() + consume(16) + # The same stream now holds 16-wide scratch; a 128-wide consume must regrow it. + assert torch.equal(consume(128), fresh) + + +def _paged(rows, heads, kv_heads, dim, pages, block_size, dtype, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + tables = torch.randperm(rows * pages, device="cuda", generator=generator) + keys = torch.randn( + rows * pages, block_size, kv_heads, dim, device="cuda", dtype=dtype, generator=generator + ) + values = torch.randn(keys.shape, device="cuda", dtype=dtype, generator=generator) + queries = torch.randn(rows, heads, dim, device="cuda", dtype=dtype, generator=generator) + return queries, keys, values, tables.reshape(rows, pages).int() + + +@pytest.mark.gpu +@pytest.mark.parametrize("retention", [0.1, 0.5]) +@pytest.mark.parametrize("ties", [False, True]) +# The Triton backend takes the split capture; any other callable takes the fallback. +@pytest.mark.parametrize("split", [True, False]) +def test_fused_capture_and_consume_track_batched_path(retention, ties, split): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention + + backend = TritonPagedAttention() if split else (lambda *args: triton_paged_attention(*args)) + + rows, heads, kv_heads, dim, block_size, pages = 6, 4, 2, 128, 16, 17 + # 10 and 30 sit on float rounding edges of ceil(length * 0.1) (30 * 0.1 > 3 in float64). + lengths = torch.tensor([0, 10, 17, 30, 200, 257], device="cuda", dtype=torch.int32) + stores = [ + make_store( + slots=rows, + layers=2, + heads=heads, + head_dim=dim, + key_bucket=272, + device="cuda", + initial_key_capacity=1, + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + leases = [store.reserve("fused", 0, rows, generation=1) for store in stores] + owners = [[(lease, row) for row in range(rows)] for lease in leases] + q, keys, values, tables = _paged( + rows, heads, kv_heads, dim, pages, block_size, torch.float32, 1 + ) + if ties: + # Every score equal: both paths must keep the smallest positions. + q.zero_() + captured = [] + for store, owner in zip(stores, owners): + plan = store.build_plan(owner, [CAPTURE] * rows, key_retention=retention) + captured.append( + sparse_paged_attention( + q, + keys, + values, + tables, + lengths, + plan, + layer=1, + dense_attention=backend, + ) + ) + batched, fused = stores + # Capture rows return (and publish) the dense backend output unchanged. These + # inputs are FP32, where some architectures (sm_90) round the split copy one + # ulp from the dense kernel; the BF16 test below stays bit for bit. + exact = 0 if not split else 1e-6 + torch.testing.assert_close(captured[1], captured[0], rtol=0, atol=exact) + torch.testing.assert_close(fused.output[:, 1], batched.output[:, 1], rtol=0, atol=exact) + torch.testing.assert_close(fused.counts[:, 1], batched.counts[:, 1], rtol=0, atol=0) + # The finalize kernel marks the same rows source-ready as store.capture. + assert torch.equal(fused.valid, batched.valid) + for row, length in enumerate(lengths.tolist()): + slot_b, slot_f = leases[0].slots[row], leases[1].slots[row] + count = int(batched.counts[slot_b, 1, 0]) + for head in range(heads): + # Both paths write ascending positions; ties keep the smaller ones. + expected = batched.selected[slot_b, 1, head, :count].tolist() + assert fused.selected[slot_f, 1, head, :count].tolist() == expected + torch.testing.assert_close( + fused.mass[slot_f, 1], batched.mass[slot_b, 1], rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close( + fused.contribution[slot_f, 1], batched.contribution[slot_b, 1], rtol=1e-4, atol=1e-5 + ) + # A later loop reads its own plane: fresh keys and queries, same pages. + q2, keys2, values2, _ = _paged(rows, heads, kv_heads, dim, pages, block_size, torch.float32, 2) + consumed = [] + for store, owner in zip(stores, owners): + plan = store.build_plan(owner, [CONSUME] * rows, key_retention=retention) + consumed.append( + sparse_paged_attention( + q2, + keys2, + values2, + tables, + lengths, + plan, + layer=1, + dense_attention=backend, + ) + ) + torch.testing.assert_close(consumed[1], consumed[0], rtol=1e-4, atol=1e-5) + assert not consumed[1][0].any(), "an empty row stays zero" + assert int(fused.uses[list(leases[1].slots), 1].min()) == 1 + for store, lease in zip(stores, leases): + store.release(lease) + + +@pytest.mark.gpu +@pytest.mark.parametrize("retention", [0.1, 0.37]) +def test_fused_radix_selection_keeps_oracle_ties_on_repeated_scores(retention): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + # Keys cycle through three vectors, so every head sees three distinct scores + # repeated across the context: the retained set is decided almost entirely + # by the tie rule (smaller positions first), with negative scores included. + # Small integers keep every dot product exact. With random reals the dense + # Triton kernel can score one key a ULP apart at different tile positions + # (seen on sm_89), so the "ties" would not be ties for either path. + rows, heads, dim, block_size, pages = 3, 4, 64, 16, 40 + generator = torch.Generator(device="cuda").manual_seed(11) + + def integers(*shape): + return torch.randint(-2, 3, shape, device="cuda", generator=generator).float() + + patterns = integers(3, heads, dim) + tokens = torch.arange(rows * pages * block_size, device="cuda") % 3 + keys = patterns[tokens].reshape(rows * pages, block_size, heads, dim).contiguous() + values = torch.randn(keys.shape, device="cuda", generator=generator) + tables = torch.arange(rows * pages, device="cuda", dtype=torch.int32).view(rows, pages) + q = integers(rows, heads, dim) + lengths = torch.tensor([7, 300, 640], device="cuda", dtype=torch.int32) + backend = TritonPagedAttention() + stores = [ + make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=640, + device="cuda", + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + for store in stores: + lease = store.reserve("ties", 0, rows, generation=1) + plan = store.build_plan( + [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention + ) + sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=0, dense_attention=backend + ) + batched, fused = stores + assert torch.equal(fused.counts, batched.counts) + for slot in range(rows): + for head in range(heads): + count = int(batched.counts[slot, 0, head]) + assert ( + fused.selected[slot, 0, head, :count].tolist() + == batched.selected[slot, 0, head, :count].tolist() + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("folded", [True, False]) +@pytest.mark.parametrize("ties", [False, True]) +@pytest.mark.parametrize("retention", [0.1, 0.37]) +def test_fused_parallel_select_matches_batched_across_many_parts( + ties, retention, folded, monkeypatch +): + import vllm_rlt.kernels.sparse_decode as fused_module + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + # The first radix pass runs inside the split kernel or as its own pass. + monkeypatch.setattr(fused_module, "_PASS0_IN_SPLIT", folded) + + # Long rows split into up to 16 parts; with repeated integer scores the + # retained ties span part boundaries, so offsets and tie quotas must carry. + rows, heads, dim, block_size, pages = 4, 4, 64, 16, 320 + generator = torch.Generator(device="cuda").manual_seed(29) + if ties: + + def integers(*shape): + return torch.randint(-2, 3, shape, device="cuda", generator=generator).float() + + patterns = integers(3, heads, dim) + tokens = torch.arange(rows * pages * block_size, device="cuda") % 3 + keys = patterns[tokens].reshape(rows * pages, block_size, heads, dim).contiguous() + q = integers(rows, heads, dim) + else: + # Mostly distinct but still exact scores: a one-ULP difference between the + # Triton and batched scores at this length could cross the boundary. + keys = torch.randint( + -8, 9, (rows * pages, block_size, heads, dim), device="cuda", generator=generator + ).float() + q = torch.randint(-8, 9, (rows, heads, dim), device="cuda", generator=generator).float() + values = torch.randn(keys.shape, device="cuda", generator=generator) + tables = torch.arange(rows * pages, device="cuda", dtype=torch.int32).view(rows, pages) + lengths = torch.tensor([1, 129, 2049, 5000], device="cuda", dtype=torch.int32) + backend = TritonPagedAttention() + stores = [ + make_store( + slots=rows, + layers=2, + heads=heads, + head_dim=dim, + key_bucket=5120, + device="cuda", + implementation=implementation, + ) + for implementation in ("batched", "fused") + ] + leases = [] + for store in stores: + lease = store.reserve("parts", 0, rows, generation=1) + leases.append(lease) + plan = store.build_plan( + [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention + ) + sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=1, dense_attention=backend + ) + batched, fused = stores + for row in range(rows): + slot_b, slot_f = leases[0].slots[row], leases[1].slots[row] + assert torch.equal(fused.counts[slot_f, 1], batched.counts[slot_b, 1]) + for head in range(heads): + count = int(batched.counts[slot_b, 1, head]) + assert ( + fused.selected[slot_f, 1, head, :count].tolist() + == batched.selected[slot_b, 1, head, :count].tolist() + ) + torch.testing.assert_close( + fused.mass[slot_f, 1], batched.mass[slot_b, 1], rtol=1e-5, atol=1e-6 + ) + torch.testing.assert_close( + fused.contribution[slot_f, 1], batched.contribution[slot_b, 1], rtol=1e-4, atol=1e-5 + ) + + +@pytest.mark.gpu +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +def test_fused_split_capture_is_the_dense_backend_output_bit_for_bit(dtype): + from vllm_rlt.kernels.paged_attention import TritonPagedAttention + + rows, heads, kv_heads, dim = 5, 4, 2, 128 + lengths = torch.tensor([0, 1, 33, 515, 1200], device="cuda", dtype=torch.int32) + q, keys, values, tables = _paged(rows, heads, kv_heads, dim, 76, 16, dtype, 5) + backend = TritonPagedAttention() + expected = backend(q, keys, values, tables, lengths) + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=1216, + device="cuda", + implementation="fused", + ) + lease = store.reserve("split", 0, rows, generation=1) + plan = store.build_plan([(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=0.1) + out = sparse_paged_attention( + q, keys, values, tables, lengths, plan, layer=0, dense_attention=backend + ) + # The split copy and the folded merge must not drift from the dense kernels: + # bit for bit in BF16; in FP32 some architectures (sm_90) contract the split + # copy's arithmetic differently and round it one ulp apart. + atol = 0 if dtype == torch.bfloat16 else 1e-6 + torch.testing.assert_close(out, expected, rtol=0, atol=atol) + # O_ref is the output actually fed into the recurrence: rounded, then FP32. + live = list(lease.slots)[1:] + torch.testing.assert_close(store.output[live, 0], expected[1:].float(), rtol=0, atol=atol) + assert bool(store.valid[list(lease.slots), 0].all()) + + +@pytest.mark.gpu +def test_fused_bf16_consume_matches_scalar_oracle_within_rounding(): + from vllm_rlt.kernels.paged_attention import triton_paged_attention + + rows, heads, kv_heads, dim = 4, 4, 4, 64 + lengths = torch.tensor([5, 64, 129, 300], device="cuda", dtype=torch.int32) + q, keys, values, tables = _paged(rows, heads, kv_heads, dim, 19, 16, torch.bfloat16, 3) + results = [] + for implementation, function in ( + ("fused", sparse_paged_attention), + ("batched", reference_paged_attention), + ): + store = make_store( + slots=rows, + layers=1, + heads=heads, + head_dim=dim, + key_bucket=304, + device="cuda", + implementation=implementation, + ) + lease = store.reserve("bf16", 0, rows, generation=1) + owner = [(lease, row) for row in range(rows)] + for modes in ([CAPTURE] * rows, [CONSUME] * rows): + plan = store.build_plan(owner, modes, key_retention=0.1) + out = function( + q, + keys, + values, + tables, + lengths, + plan, + layer=0, + dense_attention=triton_paged_attention, + ) + results.append(out.float()) + # Different FP32 reduction order may move a BF16 rounding step. + torch.testing.assert_close(results[0], results[1], rtol=2e-2, atol=2e-2) + + +@pytest.mark.gpu +def test_fused_consume_without_capture_is_nan_not_silent(): + rows, heads, dim = 2, 2, 64 + store = make_store( + slots=rows, layers=1, heads=heads, head_dim=dim, device="cuda", implementation="fused" + ) + lease = store.reserve("missing", 0, rows, generation=1) + plan = store.build_plan([(lease, 0), (lease, 1)], [CONSUME, CONSUME], key_retention=0.5) + q, keys, values, tables = _paged(rows, heads, heads, dim, 2, 16, torch.float32, 4) + lengths = torch.tensor([0, 9], device="cuda", dtype=torch.int32) + out = sparse_paged_attention(q, keys, values, tables, lengths, plan, layer=0) + assert not out[0].any() and torch.isnan(out[1]).all() + + +@pytest.mark.gpu +def test_fused_engine_rounds_abort_and_drain(monkeypatch): + import vllm_rlt.kernels.sparse_decode as fused_module + + calls = {CAPTURE: 0, CONSUME: 0} + original = fused_module.fused_paged_attention + + def counting(*args, **kwargs): + calls[args[5].modes[0]] += 1 + return original(*args, **kwargs) + + monkeypatch.setattr(fused_module, "fused_paged_attention", counting) + # The engine's Triton backend adapter must reach the split capture path. + monkeypatch.setattr(fused_module, "_capture_kernel", None) + model = tiny_ouro_model("cuda") + engine = make_speculative_engine( + model, + sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused"), + ) + store = engine.reference_store + assert store.implementation == "fused" + params = [SamplingParams(max_tokens=n, ignore_eos=True) for n in (5, 9, 50)] + for i in range(3): + engine.add_request(str(i), [2, 3, 4], params[i]) + outputs = {} + for _ in range(1000): + if "0" in outputs and outputs["0"].finished: + break + for output in engine.step(): + outputs[output.request_id] = output + if "2" not in outputs or not outputs["2"].finished: + engine.abort_request("2") + outputs.update(drain(engine)) + assert outputs["0"].finished and len(outputs["0"].token_ids) == 5 + assert outputs["1"].finished and len(outputs["1"].token_ids) == 9 + assert calls[CAPTURE] > 0 and calls[CONSUME] > 0 + assert store.outstanding == 0 and len(store.free) == store.slots + assert engine.cache_manager.num_used_blocks == 0 + engine.close() diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index f81f4d82..c438a918 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -343,10 +343,11 @@ def sparse_paged_attention( ): """Use bounded row batches on CUDA, retaining the scalar numerical oracle. - On CUDA with a dense backend available, the capture rows take the fused - paged score pass; the scalar oracle stays the definition of the arithmetic, - and ``batched_paged_attention`` without ``fused_scores`` still reproduces it - bit for bit. + A store configured for the ``fused`` implementation routes uniform capture + or consume submissions to the fused kernels. Other CUDA submissions keep + the batched path, whose capture rows take the fused paged score pass when a + dense backend is available; ``batched_paged_attention`` without + ``fused_scores`` still reproduces the scalar oracle bit for bit. """ if not q.is_cuda: return reference_paged_attention( @@ -359,6 +360,20 @@ def sparse_paged_attention( layer=layer, dense_attention=dense_attention, ) + if getattr(plan.store, "implementation", "batched") == "fused": + from vllm_rlt.kernels import sparse_decode + + if sparse_decode.supports(plan, q, dense_attention): + return sparse_decode.fused_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + layer=layer, + dense_attention=dense_attention, + ) return batched_paged_attention( q, key_cache, @@ -435,11 +450,16 @@ def batched_paged_attention( # The retained count is host-known from the capture. Reading it back # from the device would sync every layer, which RFC #156 section 6.2 # forbids. - for row in consume_rows: - count = store.host_count(plan.slots[row], layer) - if not count: - raise RuntimeError("a consume row has no captured reference") - counts[row] = count + counts = {row: store.host_count(plan.slots[row], layer) for row in consume_rows} + missing = [row for row, count in counts.items() if count is None] + if missing: + # The fused capture decides its count on the device (graph replays + # change lengths without the host), so a mixed submission that + # falls back here after it reads those counts back once. + slots = plan.slot_ids[missing] + counts.update(zip(missing, 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: @@ -465,13 +485,21 @@ def batched_paged_attention( chunk = max(1, min(32, 64 * 1024 * 1024 // (extent * heads * dim))) for start in range(0, len(all_rows), chunk): rows = all_rows[start : start + chunk] - # Building a device tensor from a Python list costs ~2.4 ms (a - # blocking host-to-device copy per call). Grouped rows are almost - # always contiguous, so derive them on the device instead. + # Contiguous rows come from a device-side arange; otherwise a pinned + # source keeps the copy from draining the stream. if rows[0] + len(rows) == rows[-1] + 1: row_ids = torch.arange(rows[0], rows[0] + len(rows), device=q.device) else: - row_ids = torch.tensor(rows, dtype=torch.int64, device=q.device) + row_ids = torch.tensor(rows, dtype=torch.int64, pin_memory=True).to( + q.device, non_blocking=True + ) + if mode == DENSE and dense_output is not None: + output[row_ids] = dense_output[row_ids] + continue + if mode == CONSUME and dense_output is not None and plan.key_retention == 1.0: + output[row_ids] = dense_output[row_ids] + consumed = True + continue slots = plan.slot_ids[row_ids] chunk_slots = [plan.slots[row] for row in rows] queries = q[row_ids].float() @@ -536,9 +564,6 @@ def batched_paged_attention( positions = torch.arange(extent, device=q.device) blocks = tables[:, positions // block_size] offsets = positions % block_size - if mode == DENSE and dense_output is not None: - output[row_ids] = dense_output[row_ids] - continue if fused_scores and dense_output is not None: # One pass over the paged cache: no expanded FP32 K row, and # no per-row cuBLAS. The dot product is still FP32; only the @@ -571,21 +596,16 @@ def batched_paged_attention( 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" - ) + store.ensure_key_capacity(count) order = torch.argsort(probabilities, dim=-1, descending=True, stable=True) selected = order[..., :count].sort(dim=-1).values gathered = probabilities.gather(-1, selected) - # Gather the retained values straight from the paged cache: the - # expanded fp32 copy would be the same bits (bf16 -> fp32 is - # exact) but reading it scatters through a row-length tensor. + # Gather only retained V directly from the paged cache. The + # FP32 values and reduction geometry are identical to gathering + # from the full GQA-expanded tensor, without materializing it. + selected_blocks = tables[batch_index, selected // block_size] selected_values = value_cache[ - tables[batch_index, selected // block_size], - selected % block_size, - kv_index, + selected_blocks, selected % block_size, kv_index ].float() contribution = _rowwise_product("hk,hkd->hd", gathered, selected_values) # Only the live prefix is written, never the whole capacity. diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 3f926b7c..d02c7c57 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -212,10 +212,16 @@ class SparseDecodeAttentionConfig: # 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 + # ``batched`` keeps the scalar oracle's reduction shapes (bitwise contract); + # ``fused`` launches fixed per-layer kernels without host readbacks and + # changes consume-row numerics (see vllm_rlt.kernels.sparse_decode). + implementation: str = "batched" def __post_init__(self): if type(self.enabled) is not bool: raise ValueError("enabled must be a boolean") + if self.implementation not in ("batched", "fused"): + raise ValueError("implementation must be 'batched' or 'fused'") 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: diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index f1c5b68f..f4244e3e 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -194,20 +194,31 @@ def __init__( heads=config.num_attention_heads, head_dim=config.head_dim, key_bucket=capacity, + initial_key_capacity=min(capacity, 64), device=parameter.device, + implementation=sparse_decode.implementation, ) + if parameter.device.type == "cuda": + free_bytes, _ = torch.cuda.mem_get_info(parameter.device) + growth_bytes = self.reference_store.peak_capacity_bytes - self.reference_store.bytes + if growth_bytes > free_bytes: + raise ValueError( + "sparse reference index growth exceeds available device memory" + ) logger.info( "sparse decode attention: %d reference slots x %d keys x %d layers x %d heads " - "= %.2f GiB (context bound %d, retention %.2f)", + "<= %.2f GiB growth peak (context bound %d, retention %.2f)", slots, capacity, config.num_hidden_layers, config.num_attention_heads, - self.reference_store.bytes / 1024**3, + self.reference_store.peak_capacity_bytes / 1024**3, context_bound, sparse_decode.key_retention, ) - scratch += self.reference_store.bytes + # Budget the logical maximum and copy overlap, not just the small + # initial allocation; short prompts must not overcommit long ones. + scratch += self.reference_store.peak_capacity_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 diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py new file mode 100644 index 00000000..209a4730 --- /dev/null +++ b/vllm_rlt/kernels/sparse_decode.py @@ -0,0 +1,1410 @@ +"""Fused sparse decode attention for uniform capture or consume submissions. + +Opt-in alternative to the batched path in ``vllm_rlt.attention.sparse``. That +path keeps the scalar oracle's cuBLAS reduction shapes, so every layer launches +per-row products and reads row lengths and retained counts back to the host. +Here each layer launches a fixed number of kernels and keeps lengths and +counts on the device. + +Numerics differ from the oracle on consume rows: scores and FP32 reductions +run in another order, so a key on the selection boundary can differ from the +oracle's. Retained positions are written in ascending order and equal scores +keep the smaller positions, as in the oracle. Capture rows still return the +dense backend output, which is also published as ``O_ref``, so the recurrence +input of the reference loop is unchanged. + +A consume row whose reference was never captured (count 0) yields NaN instead +of raising: raising would need the host readback this path exists to avoid. The +batched and scalar paths keep the explicit error. +""" + +import operator + +import torch +import triton +import triton.language as tl + +from vllm_rlt.attention.sparse import CAPTURE, CONSUME, retained_count +from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention +from vllm_rlt.kernels.triton_attention import NUM_SPLITS + + +@triton.jit +def _order_key(scores): + """Key in [0, 2**32) with the order of FP32 scores (bits of negatives flipped).""" + bits = scores.to(tl.int32, bitcast=True) + return (bits ^ ((bits >> 31) & 0x7FFFFFFF)).to(tl.int64) + 2147483648 + + +@triton.jit +def _select_and_publish( + scores_row, + lse, + length, + count, + slot, + layer, + head, + reference, + table, + V, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + retained_capacity, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Keep the ``count`` largest of ``length`` stored scores and publish ``(S, m, O_ref, C_ref)``. + + Finds the count-th largest score by a radix select over its 32-bit order + key, four bits per pass, and writes the retained positions in ascending + order; equal scores keep the smaller positions, like the oracle's stable + sort. ``reference`` is ``O_ref``. + """ + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + # After eight passes ``prefix`` is the order key of the count-th largest + # score and ``ties_needed`` how many keys equal to it are still retained. + # Loop-carried values must keep one type on every iteration. + prefix = length.to(tl.int64) * 0 + ties_needed = count + digits = tl.arange(0, 16) + for level in tl.static_range(8): + histogram = tl.zeros((16,), dtype=tl.int32) + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + candidate = live & ((key >> (32 - 4 * level)) == prefix) + digit = ((key >> (28 - 4 * level)) & 15).to(tl.int32) + histogram += tl.sum( + ((digit[:, None] == digits[None, :]) & candidate[:, None]).to(tl.int32), axis=0 + ) + # Highest digit whose keys, with every higher digit, still cover the rest. + covered = tl.cumsum(histogram, 0, reverse=True) + chosen = tl.sum((covered >= ties_needed).to(tl.int32), axis=0) - 1 + ties_needed -= tl.sum(tl.where(digits > chosen, histogram, 0), axis=0) + prefix = prefix * 16 + chosen + low = prefix + selected = SELECTED + ((slot * NUM_LAYERS + layer) * NUM_HEADS + head) * retained_capacity + written = 0 + ties_seen = 0 + mass = 0.0 + contribution = tl.zeros((BLOCK_D,), dtype=tl.float32) + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + scores = tl.load(scores_row + positions, mask=live, other=0) + key = _order_key(scores) + tie = (live & (key == low)).to(tl.int32) + take = (live & (key > low)) | ( + (tie == 1) & (ties_seen + tl.cumsum(tie, 0) - tie < ties_needed) + ) + taken = take.to(tl.int32) + tl.store(selected + written + tl.cumsum(taken, 0) - taken, positions, mask=take) + probabilities = tl.where(take, tl.exp(scores - lse), 0.0) + blocks = tl.load(table + positions // PAGE_SIZE, mask=take, other=0).to(tl.int64) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + (positions % PAGE_SIZE)[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=take[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + contribution += tl.sum(probabilities[:, None] * values, axis=0) + mass += tl.sum(probabilities, axis=0) + written += tl.sum(taken, axis=0) + ties_seen += tl.sum(tie, axis=0) + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + tl.store(COUNTS + state, count) + tl.store(MASS + state, mass) + tl.store(OUTPUT + state * HEAD_DIM + dims, reference, mask=dims < HEAD_DIM) + tl.store(CONTRIBUTION + state * HEAD_DIM + dims, contribution, mask=dims < HEAD_DIM) + + +# A copy of triton_attention._paged_attention_kernel that also stores the raw +# scores. Its arithmetic must stay identical to the dense kernel: the capture +# output is required to equal the dense backend bit for bit (tested). +@triton.jit +def _capture_split_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + PART_M, + PART_L, + PART_O, + SCORES, + SEGHIST, + RETENTION, + STATE, + TICKETS, + top, + score_capacity, + q_batch_stride: tl.constexpr, + q_head_stride: tl.constexpr, + q_dim_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + k_dim_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + v_dim_stride: tl.constexpr, + table_stride: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + NUM_HEADS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + SPLITS: tl.constexpr, + HISTOGRAM: tl.constexpr, + PARTS: tl.constexpr, +): + row = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + kv_head = head // HEAD_GROUPS + length = tl.load(LENGTHS + row) + # Equal tile-aligned parts of this row's context; trailing parts may be empty. + split_tokens = ((length + SPLITS - 1) // SPLITS + BLOCK_T - 1) // BLOCK_T * BLOCK_T + begin = split * split_tokens + end = tl.minimum(length, begin + split_tokens) + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_batch_stride + head * q_head_stride + dims * q_dim_stride, + mask=dims < HEAD_DIM, + other=0, + ).to(tl.float32) + maximum = -float("inf") + normalizer = 0.0 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + local = tl.zeros((512,), dtype=tl.int32) + for start in range(begin, end, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + valid_tokens = positions < end + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, + mask=valid_tokens, + other=0, + ).to(tl.int64) + # Physical IDs fit in int32, but ID * block_stride can exceed 2**31 + # elements in a large KV pool. Promote BEFORE multiplying, for K and V. + offsets = positions % PAGE_SIZE + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + offsets[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :] * k_dim_stride, + mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.sum(keys * query[None, :], axis=1) * SCALE + scores = tl.where(valid_tokens, scores, -float("inf")) + # The only addition to the dense kernel: keep the raw scores for selection. + tl.store( + SCORES + (row * NUM_HEADS + head).to(tl.int64) * score_capacity + positions, + scores, + mask=valid_tokens, + ) + if HISTOGRAM: + # Selection pass 0 of this part: the top eight bits of the order key. + digit = (_order_key(scores) >> 24).to(tl.int32) + local += tl.histogram(tl.where(valid_tokens, digit, 256), 512) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + correction = tl.exp(maximum - next_maximum) + weights = tl.exp(scores - next_maximum) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :] * v_dim_stride, + mask=valid_tokens[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + accumulator = accumulator * correction + tl.sum(weights[:, None] * values, axis=0) + normalizer = normalizer * correction + tl.sum(weights, axis=0) + maximum = next_maximum + # An empty part stores (-inf, 0, 0), which the merge weights by zero. + part = (row * NUM_HEADS + head) * SPLITS + split + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_O + part * BLOCK_D + dims, accumulator) + if HISTOGRAM: + rh = row * NUM_HEADS + head + _, count = _retained(LENGTHS, row, RETENTION, top) + if count > 0: + bins = tl.arange(0, 512) + tl.store(SEGHIST + (rh * PARTS + split) * 4 * 256 + bins, local, mask=bins < 256) + _close_pass(SEGHIST, STATE, TICKETS, rh, count, 0, SPLITS, PARTS) + + +# Parallel selection after the split kernel. Every (row, head) splits its +# context into PARTS tile-aligned parts. Four radix passes over the 32-bit order +# key, eight bits each, store one 256-bin histogram per part. The last part to +# finish a pass (an integer ticket per row and head) sums the histograms in a +# fixed order and publishes the state for the next pass, so the result stays +# deterministic and no histogram is zeroed. After the last pass it also derives +# every part's output offset and tie quota, which keeps the ascending order and +# the smaller-position tie rule. The contribution kernel then gathers V over the +# compact retained list only. +SELECT_PARTS = 16 +# Tiles of the radix passes and the write (scores only) and of the V gather, +# chosen by a sweep on an RTX 5090 at 4K-32K (16 vs 32 parts, 128-1024 scan, +# 32-128 gather, 4 or 8 warps). +_SCAN_BLOCK = 1024 +_SCAN_WARPS = 4 +_GATHER_BLOCK = 128 +# The split kernel's tile (the dense kernel's), so its splits are whole parts. +_SEGMENT_ALIGN = 32 +# The split kernel builds the first pass's histograms while it scores the keys. +_PASS0_IN_SPLIT = True + +# Triton's Python dispatch costs several microseconds per launch, and a fused +# capture makes six launches, so at short contexts in eager mode the dispatch, +# not the GPU, sets its cost. A launch whose specialization was seen before +# reuses that compiled kernel directly. The key follows Triton's own +# specialization (dtype and 16-byte alignment of pointers; integers equal to 1, +# divisible by 16, and their width; constexpr values); tests set _LAUNCH_CHECK +# to assert that it always picks the kernel Triton would. Other Triton versions +# keep its dispatch. +try: + from triton import knobs as _knobs + from triton.runtime.driver import driver as _driver + + _FAST_LAUNCH = triton.__version__.startswith("3.8.") and hasattr( + _knobs.runtime, "launch_enter_hook" + ) +except ImportError: # pragma: no cover - other Triton versions + _FAST_LAUNCH = False +_LAUNCH_CHECK = False + + +def _key_function(args, constexpr): + """Compile the specialization key of one argument layout, as Triton's binder does. + + The layout is the arguments' types; each layout gets its own function. + """ + terms = [] + for i, value in enumerate(args): + if i in constexpr: + terms.append(f"a[{i}]") + elif isinstance(value, torch.Tensor): + terms.append(f"a[{i}].dtype, a[{i}].data_ptr() & 15 == 0") + elif type(value) is int: + terms.append(f"a[{i}] == 1, a[{i}] & 15 == 0, -2147483648 <= a[{i}] < 2147483648") + elif type(value) is not float: + return None + if not terms: + return lambda a: () + body = ", ".join(terms) + return eval(f"lambda a: ({body},)") # noqa: S307 - built from indices only + + +class _Launcher: + """``kernel[grid](*args, **meta)`` that skips Triton's dispatch once compiled.""" + + def __init__(self, kernel): + self.kernel = kernel + self.compiled = {} + self.keys = {} + # Positional parameters declared constexpr specialize on their value. + self.constexpr = frozenset(i for i, p in enumerate(kernel.params) if p.is_constexpr) + + def launch(self, grid, args, meta): + if not _FAST_LAUNCH: + self.kernel[grid](*args, **meta) + return + layout = (tuple(map(type, args)), len(grid)) + keyed = self.keys.get(layout) + if keyed is None: + keyed = self.keys[layout] = _key_function(args, self.constexpr) + if keyed is None: + self.kernel[grid](*args, **meta) + return + device = _driver.active.get_current_device() + key = (device, layout, keyed(args), tuple(meta.items())) + entry = self.compiled.get(key) + if entry is None: + compiled = self.kernel[grid](*args, **meta) + # The launcher takes the bound arguments in the binder's order. + binder = self.kernel.device_caches[device][4] + options = dict( + meta, + debug=self.kernel.debug or _knobs.runtime.debug, + instrumentation_mode=_knobs.compilation.instrumentation_mode, + ) + names = tuple(binder(*args, **options)[0])[len(args) :] + # None when the call passes its meta in the binder's order. + if names == tuple(meta): + order = None + elif len(names) > 1: + order = operator.itemgetter(*names) + else: + order = lambda meta: tuple(meta[name] for name in names) # noqa: E731 + self.compiled[key] = (compiled, order) + return + compiled, order = entry + if _LAUNCH_CHECK: + assert self.kernel.run(*args, grid=grid, warmup=True, **meta) is compiled + values = args + (tuple(meta.values()) if order is None else order(meta)) + stream = _driver.active.get_current_stream(device) + compiled.run( + grid[0], + grid[1] if len(grid) > 1 else 1, + grid[2] if len(grid) > 2 else 1, + stream, + compiled.function, + compiled.packed_metadata, + compiled.launch_metadata(grid, stream, *values), + _knobs.runtime.launch_enter_hook, + _knobs.runtime.launch_exit_hook, + *values, + ) + + +def _current_stream(device): + """The raw current stream; Triton's driver query is much cheaper than torch's.""" + if _FAST_LAUNCH: + return _driver.active.get_current_stream(device.index) + return torch.cuda.current_stream(device).cuda_stream + + +_LAUNCHERS: dict = {} + + +def _launch(kernel, grid, *args, **meta): + # Keyed by identity: hashing a JITFunction recomputes its cache key. + launcher = _LAUNCHERS.get(id(kernel)) + if launcher is None or launcher.kernel is not kernel: + launcher = _LAUNCHERS[id(kernel)] = _Launcher(kernel) + launcher.launch(grid, args, meta) + + +@triton.jit +def _retained(LENGTHS, row, RETENTION, top): + """The row length and its retained count, as the host's ``retained_count``.""" + length = tl.load(LENGTHS + row) + # retained_count in float64, exactly as the host's math.ceil(length * r). + real = length.to(tl.float64) + count = tl.minimum(tl.maximum(tl.ceil(real * tl.load(RETENTION)), 1.0), real) + return length, tl.minimum(count.to(tl.int32), top) + + +@triton.jit +def _segment(length, SEGMENTS: tl.constexpr, ALIGN: tl.constexpr): + """This program's part ``[first, last)`` of a row's context.""" + size = tl.cdiv(tl.cdiv(length, SEGMENTS), ALIGN) * ALIGN + first = tl.program_id(2) * size + return first, tl.minimum(first + size, length) + + +@triton.jit +def _choose(histogram, prefix, ties_needed): + """Highest digit whose keys, with every higher digit, still cover the rest.""" + digits = tl.arange(0, 256) + covered = tl.cumsum(histogram, 0, reverse=True) + chosen = tl.sum((covered >= ties_needed).to(tl.int32), axis=0) - 1 + ties_needed -= tl.sum(tl.where(digits > chosen, histogram, 0), axis=0) + return prefix * 256 + chosen.to(tl.int64), ties_needed, chosen + + +@triton.jit +def _last_part(TICKETS, rh, PARTS: tl.constexpr): + """Whether this program is the last of its row and head to finish; resets the ticket. + + Each part stores its results first; the acquire-release ticket then makes + every part's stores visible to the last one, which reads them through L2. + """ + tl.debug_barrier() + last = tl.atomic_add(TICKETS + rh, 1, sem="acq_rel", scope="gpu") == PARTS - 1 + if last: + tl.store(TICKETS + rh, 0) + return last + + +@triton.jit +def _close_pass( + SEGHIST, STATE, TICKETS, rh, count, PASS: tl.constexpr, USED: tl.constexpr, PARTS: tl.constexpr +): + """The last part of a pass publishes the state before the next pass.""" + if _last_part(TICKETS, rh, USED): + digits = tl.arange(0, 256) + parts = (rh * PARTS + tl.arange(0, USED)) * 4 + PASS + histogram = tl.sum( + tl.load(SEGHIST + parts[:, None] * 256 + digits[None, :], cache_modifier=".cg"), axis=0 + ) + if PASS == 0: + prefix = count.to(tl.int64) * 0 + ties_needed = count + else: + prefix = tl.load(STATE + (rh * 5 + PASS) * 2) + ties_needed = tl.load(STATE + (rh * 5 + PASS) * 2 + 1).to(tl.int32) + prefix, ties_needed, _ = _choose(histogram, prefix, ties_needed) + tl.store(STATE + (rh * 5 + PASS + 1) * 2, prefix) + tl.store(STATE + (rh * 5 + PASS + 1) * 2 + 1, ties_needed.to(tl.int64)) + + +@triton.jit +def _select_pass_kernel( + PART_M, + PART_L, + PART_O, + OUT, + SCORES, + LENGTHS, + RETENTION, + SLOTS, + COUNTS, + OUTPUT, + VALID, + LSE, + SEGHIST, + STATE, + TICKETS, + GREATER, + OFFSETS, + layer, + top, + score_capacity, + out_batch_stride: tl.constexpr, + out_head_stride: tl.constexpr, + out_dim_stride: tl.constexpr, + PASS: tl.constexpr, + MERGE: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, + ALIGN: tl.constexpr, + SPLITS: tl.constexpr, +): + """One radix pass of one part; the first launched pass also merges the split parts.""" + row = tl.program_id(0) + head = tl.program_id(1) + segment = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _retained(LENGTHS, row, RETENTION, top) + if MERGE: + if segment == 0: + # Identical to triton_attention._paged_attention_merge_kernel. + parts = rh * SPLITS + tl.arange(0, SPLITS) + dims = tl.arange(0, BLOCK_D) + maxima = tl.load(PART_M + parts) + maximum = tl.max(maxima, axis=0) + scale = tl.where(maxima > -float("inf"), tl.exp(maxima - maximum), 0.0) + normalizer = tl.sum(tl.load(PART_L + parts) * scale, axis=0) + partial = tl.load(PART_O + parts[:, None] * BLOCK_D + dims[None, :]) + accumulator = tl.sum(partial * scale[:, None], axis=0) + dense = accumulator / tl.maximum(normalizer, 1.0e-30) + out = OUT + row * out_batch_stride + head * out_head_stride + dims * out_dim_stride + tl.store(out, dense, mask=dims < HEAD_DIM) + slot = tl.load(SLOTS + row) + if head == 0: + tl.store(VALID + slot * NUM_LAYERS + layer, 1) + tl.store(LSE + rh, maximum + tl.log(normalizer)) + if count > 0: + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + tl.store(COUNTS + state, count) + # O_ref is the output as rounded to the query dtype: the recurrence input. + reference = dense.to(OUT.dtype.element_ty).to(tl.float32) + tl.store(OUTPUT + state * HEAD_DIM + dims, reference, mask=dims < HEAD_DIM) + if count > 0: + if PASS == 0: + prefix = length.to(tl.int64) * 0 + else: + prefix = tl.load(STATE + (rh * 5 + PASS) * 2) + first, last = _segment(length, PARTS, ALIGN) + scores_row = SCORES + rh.to(tl.int64) * score_capacity + local = tl.zeros((512,), dtype=tl.int32) + greater = 0 + for start in range(first, last, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < last + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + candidate = live & ((key >> (32 - 8 * PASS)) == prefix) + digit = ((key >> (24 - 8 * PASS)) & 255).to(tl.int32) + # Bin 256 collects the non-candidates and is never stored. + local += tl.histogram(tl.where(candidate, digit, 256), 512) + if PASS == 3: + # Keys of this part already above the threshold's first 24 bits. + greater += tl.sum((live & ((key >> 8) > prefix)).to(tl.int32), axis=0) + bins = tl.arange(0, 512) + own = (rh * PARTS + segment) * 4 + PASS + tl.store(SEGHIST + own * 256 + bins, local, mask=bins < 256) + if PASS == 3: + tl.store(GREATER + rh * PARTS + segment, greater) + if _last_part(TICKETS, rh, PARTS): + # Fix the threshold, then every part's offset and tie quota. + digits = tl.arange(0, 256) + segments = tl.arange(0, PARTS) + prefix = tl.load(STATE + (rh * 5 + 3) * 2) + ties_needed = tl.load(STATE + (rh * 5 + 3) * 2 + 1).to(tl.int32) + last_pass = tl.load( + SEGHIST + ((rh * PARTS + segments[:, None]) * 4 + 3) * 256 + digits[None, :], + cache_modifier=".cg", + ) + low, ties_needed, chosen = _choose(tl.sum(last_pass, axis=0), prefix, ties_needed) + above = tl.load(GREATER + rh * PARTS + segments, cache_modifier=".cg") + tl.sum( + tl.where(digits[None, :] > chosen, last_pass, 0), axis=1 + ) + equal = tl.sum(tl.where(digits[None, :] == chosen, last_pass, 0), axis=1) + ties_before = tl.cumsum(equal, 0) - equal + taken = above + tl.minimum(tl.maximum(ties_needed - ties_before, 0), equal) + tl.store(OFFSETS + (rh * PARTS + segments) * 2, tl.cumsum(taken, 0) - taken) + tl.store(OFFSETS + (rh * PARTS + segments) * 2 + 1, ties_before) + tl.store(STATE + (rh * 5 + 4) * 2, low) + tl.store(STATE + (rh * 5 + 4) * 2 + 1, ties_needed.to(tl.int64)) + else: + _close_pass(SEGHIST, STATE, TICKETS, rh, count, PASS, PARTS, PARTS) + + +@triton.jit +def _select_write_kernel( + SCORES, + LENGTHS, + RETENTION, + SLOTS, + STATE, + OFFSETS, + SELECTED, + layer, + top, + score_capacity, + retained_capacity, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, + ALIGN: tl.constexpr, +): + """Write this part's retained positions at its offset, ties up to its quota.""" + row = tl.program_id(0) + head = tl.program_id(1) + segment = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _retained(LENGTHS, row, RETENTION, top) + if count > 0: + low = tl.load(STATE + (rh * 5 + 4) * 2) + ties_needed = tl.load(STATE + (rh * 5 + 4) * 2 + 1).to(tl.int32) + written = tl.load(OFFSETS + (rh * PARTS + segment) * 2) + ties_seen = tl.load(OFFSETS + (rh * PARTS + segment) * 2 + 1) + slot = tl.load(SLOTS + row) + selected = ( + SELECTED + + ((slot * NUM_LAYERS + layer) * NUM_HEADS + head).to(tl.int64) * retained_capacity + ) + scores_row = SCORES + rh.to(tl.int64) * score_capacity + first, last = _segment(length, PARTS, ALIGN) + for start in range(first, last, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < last + key = _order_key(tl.load(scores_row + positions, mask=live, other=0)) + tie = (live & (key == low)).to(tl.int32) + take = (live & (key > low)) | ( + (tie == 1) & (ties_seen + tl.cumsum(tie, 0) - tie < ties_needed) + ) + taken = take.to(tl.int32) + tl.store(selected + written + tl.cumsum(taken, 0) - taken, positions, mask=take) + written += tl.sum(taken, axis=0) + ties_seen += tl.sum(tie, axis=0) + + +@triton.jit +def _select_contribution_kernel( + SCORES, + V, + TABLES, + LENGTHS, + RETENTION, + SLOTS, + LSE, + SELECTED, + PART_MASS, + PART_C, + TICKETS, + MASS, + CONTRIBUTION, + layer, + top, + score_capacity, + table_stride, + retained_capacity, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + PARTS: tl.constexpr, +): + """``(m, C_ref)`` over the retained ranks, in rank chunks summed in a fixed order. + + Gathering V over the compact retained list keeps every lane busy, where a + scan over all positions would leave most of them masked. + """ + row = tl.program_id(0) + head = tl.program_id(1) + part = tl.program_id(2) + rh = row * NUM_HEADS + head + length, count = _retained(LENGTHS, row, RETENTION, top) + if count > 0: + chunk = tl.cdiv(tl.cdiv(count, PARTS), BLOCK_T) * BLOCK_T + first = part * chunk + last = tl.minimum(first + chunk, count) + lse = tl.load(LSE + rh) + slot = tl.load(SLOTS + row) + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + selected = SELECTED + state.to(tl.int64) * retained_capacity + scores_row = SCORES + rh.to(tl.int64) * score_capacity + table = TABLES + row * table_stride + mass = 0.0 + contribution = tl.zeros((BLOCK_D,), dtype=tl.float32) + for start in range(first, last, BLOCK_T): + ranks = start + tl.arange(0, BLOCK_T) + live = ranks < last + positions = tl.load(selected + ranks, mask=live, other=0) + scores = tl.load(scores_row + positions, mask=live, other=0) + probabilities = tl.where(live, tl.exp(scores - lse), 0.0) + blocks = tl.load(table + positions // PAGE_SIZE, mask=live, other=0).to(tl.int64) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + (positions % PAGE_SIZE)[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + contribution += tl.sum(probabilities[:, None] * values, axis=0) + mass += tl.sum(probabilities, axis=0) + index = rh * PARTS + part + tl.store(PART_MASS + index, mass) + tl.store(PART_C + index * BLOCK_D + dims, contribution) + if _last_part(TICKETS, rh, PARTS): + parts = rh * PARTS + tl.arange(0, PARTS) + total = tl.sum(tl.load(PART_MASS + parts, cache_modifier=".cg"), axis=0) + summed = tl.sum( + tl.load(PART_C + parts[:, None] * BLOCK_D + dims[None, :], cache_modifier=".cg"), + axis=0, + ) + tl.store(MASS + state, total) + tl.store(CONTRIBUTION + state * HEAD_DIM + dims, summed, mask=dims < HEAD_DIM) + + +@triton.jit +def _capture_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + RETENTION, + SLOTS, + DENSE, + SCORES, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + VALID, + layer, + top, + score_capacity, + table_stride, + retained_capacity, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + dense_row_stride: tl.constexpr, + dense_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, +): + """Select and publish the retained state ``(S, m, O_ref, C_ref)`` of one row and head. + + One program scores every key, finds the ``count``-th largest score by + bisection on its order key, and writes the retained positions in ascending + order; equal scores keep the smaller positions, like the oracle's stable + sort. It also marks the reference source-ready for this layer, as + ``AttentionReferenceStore.capture`` does, without a separate launch. + """ + row = tl.program_id(0) + head = tl.program_id(1) + slot = tl.load(SLOTS + row) + if head == 0: + tl.store(VALID + slot * NUM_LAYERS + layer, 1) + length = tl.load(LENGTHS + row) + # retained_count in float64, exactly as the host's math.ceil(length * r). + real = length.to(tl.float64) + count = tl.minimum(tl.maximum(tl.ceil(real * tl.load(RETENTION)), 1.0), real) + # key_bucket covers the serving context, so this never truncates; it only + # keeps a misconfiguration from writing past the index buffer. + count = tl.minimum(count.to(tl.int32), top) + if count > 0: + kv_head = head // HEAD_GROUPS + dims = tl.arange(0, BLOCK_D) + query = tl.load( + Q + row * q_row_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + table = TABLES + row * table_stride + scores_row = SCORES + (row * NUM_HEADS + head).to(tl.int64) * score_capacity + maximum = -float("inf") + normalizer = 0.0 + for start in range(0, length, BLOCK_T): + positions = start + tl.arange(0, BLOCK_T) + live = positions < length + blocks = tl.load(table + positions // PAGE_SIZE, mask=live, other=0).to(tl.int64) + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + (positions % PAGE_SIZE)[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.where(live, tl.sum(keys * query[None, :], axis=1) * SCALE, -float("inf")) + tl.store(scores_row + positions, scores, mask=live) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + normalizer = normalizer * tl.exp(maximum - next_maximum) + tl.sum( + tl.exp(scores - next_maximum), axis=0 + ) + maximum = next_maximum + lse = maximum + tl.log(normalizer) + # Later passes read scores other threads of this program stored. + tl.debug_barrier() + dense = tl.load( + DENSE + row * dense_row_stride + head * dense_head_stride + dims, + mask=dims < HEAD_DIM, + other=0, + ).to(tl.float32) + _select_and_publish( + scores_row, + lse, + length, + count, + slot, + layer, + head, + dense, + table, + V, + SELECTED, + COUNTS, + MASS, + OUTPUT, + CONTRIBUTION, + retained_capacity, + v_block_stride, + v_token_stride, + v_head_stride, + NUM_LAYERS, + NUM_HEADS, + HEAD_DIM, + HEAD_GROUPS, + PAGE_SIZE, + BLOCK_D, + BLOCK_T, + ) + + +@triton.jit +def _consume_split_kernel( + Q, + K, + V, + TABLES, + LENGTHS, + SLOTS, + SELECTED, + COUNTS, + PART_M, + PART_L, + PART_O, + layer, + table_stride, + retained_capacity, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + k_block_stride: tl.constexpr, + k_token_stride: tl.constexpr, + k_head_stride: tl.constexpr, + v_block_stride: tl.constexpr, + v_token_stride: tl.constexpr, + v_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + HEAD_GROUPS: tl.constexpr, + PAGE_SIZE: tl.constexpr, + SCALE: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_T: tl.constexpr, + SPLITS: tl.constexpr, +): + """Online-softmax parts of ``C_new`` over one contiguous chunk of retained ranks. + + Each (row, head) splits its ``count`` retained keys into ``SPLITS`` chunks + of whole ``BLOCK_T`` tiles; an empty chunk writes the merge's neutral part. + """ + row = tl.program_id(0) + head = tl.program_id(1) + split = tl.program_id(2) + dims = tl.arange(0, BLOCK_D) + part = (row * NUM_HEADS + head) * SPLITS + split + maximum = -float("inf") + normalizer = 0.0 + accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32) + if tl.load(LENGTHS + row) > 0: + slot = tl.load(SLOTS + row) + kv_head = head // HEAD_GROUPS + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + count = tl.load(COUNTS + state) + chunk = tl.cdiv(tl.cdiv(count, SPLITS), BLOCK_T) * BLOCK_T + first = split * chunk + last = tl.minimum(first + chunk, count) + selected = SELECTED + state.to(tl.int64) * retained_capacity + query = tl.load( + Q + row * q_row_stride + head * q_head_stride + dims, mask=dims < HEAD_DIM, other=0 + ).to(tl.float32) + for start in range(first, last, BLOCK_T): + ranks = start + tl.arange(0, BLOCK_T) + live = ranks < last + positions = tl.load(selected + ranks, mask=live, other=0) + blocks = tl.load( + TABLES + row * table_stride + positions // PAGE_SIZE, mask=live, other=0 + ).to(tl.int64) + offsets = positions % PAGE_SIZE + keys = tl.load( + K + + blocks[:, None] * k_block_stride + + offsets[:, None] * k_token_stride + + kv_head * k_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + scores = tl.where(live, tl.sum(keys * query[None, :], axis=1) * SCALE, -float("inf")) + next_maximum = tl.maximum(maximum, tl.max(scores, axis=0)) + correction = tl.exp(maximum - next_maximum) + weights = tl.exp(scores - next_maximum) + values = tl.load( + V + + blocks[:, None] * v_block_stride + + offsets[:, None] * v_token_stride + + kv_head * v_head_stride + + dims[None, :], + mask=live[:, None] & (dims[None, :] < HEAD_DIM), + other=0, + ).to(tl.float32) + accumulator = accumulator * correction + tl.sum(weights[:, None] * values, axis=0) + normalizer = normalizer * correction + tl.sum(weights, axis=0) + maximum = next_maximum + tl.store(PART_M + part, maximum) + tl.store(PART_L + part, normalizer) + tl.store(PART_O + part * BLOCK_D + dims, accumulator) + + +@triton.jit +def _consume_merge_kernel( + PART_M, + PART_L, + PART_O, + LENGTHS, + SLOTS, + MASS, + OUTPUT, + CONTRIBUTION, + USES, + OUT, + layer, + q_row_stride: tl.constexpr, + q_head_stride: tl.constexpr, + NUM_LAYERS: tl.constexpr, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + SPLITS: tl.constexpr, +): + """``O = O_ref - C_ref + m * C_new``, merging the split parts in a fixed order. + + Also counts the use, as ``AttentionReferenceStore.consume`` does. + """ + row = tl.program_id(0) + head = tl.program_id(1) + dims = tl.arange(0, BLOCK_D) + out = OUT + row * q_row_stride + head * q_head_stride + dims + slot = tl.load(SLOTS + row) + if head == 0: + uses = USES + slot * NUM_LAYERS + layer + tl.store(uses, tl.load(uses) + 1) + if tl.load(LENGTHS + row) == 0: + tl.store(out, tl.zeros((BLOCK_D,), dtype=tl.float32), mask=dims < HEAD_DIM) + else: + state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head + parts = (row * NUM_HEADS + head) * SPLITS + tl.arange(0, SPLITS) + maxima = tl.load(PART_M + parts) + maximum = tl.max(maxima, axis=0) + scale = tl.where(maxima > -float("inf"), tl.exp(maxima - maximum), 0.0) + normalizer = tl.sum(tl.load(PART_L + parts) * scale, axis=0) + partial = tl.load(PART_O + parts[:, None] * BLOCK_D + dims[None, :]) + accumulator = tl.sum(partial * scale[:, None], axis=0) + reference = tl.load(OUTPUT + state * HEAD_DIM + dims, mask=dims < HEAD_DIM, other=0) + contribution = tl.load( + CONTRIBUTION + state * HEAD_DIM + dims, mask=dims < HEAD_DIM, other=0 + ) + # count 0 (no captured reference) divides 0 by 0: NaN, never a silent value. + corrected = reference - contribution + tl.load(MASS + state) * (accumulator / normalizer) + tl.store(out, corrected, mask=dims < HEAD_DIM) + + +_RETENTION: dict[tuple[torch.device, float], torch.Tensor] = {} + + +def _retention(device, key_retention): + """A float64 device scalar, built once; Triton would pass a Python float as fp32.""" + key = (device, key_retention) + if key not in _RETENTION: + _RETENTION[key] = torch.full((1,), key_retention, dtype=torch.float64, device=device) + return _RETENTION[key] + + +_WORKSPACE: dict = {} + + +def _build_capture_workspace(device, rows, heads, capacity, block_d): + empty = lambda *size, dtype=torch.float32: torch.empty(size, device=device, dtype=dtype) # noqa: E731 + return ( + empty(rows, heads, capacity), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS, block_d), + # Per-part radix histograms, the state before each pass (and the final + # threshold), per-part keys above the threshold and output offsets, the + # merged lse, the last-part tickets (zero between calls), and the + # contribution parts. + empty(rows, heads, SELECT_PARTS, 4, 256, dtype=torch.int32), + empty(rows, heads, 5, 2, dtype=torch.int64), + empty(rows, heads, SELECT_PARTS, dtype=torch.int32), + empty(rows, heads, SELECT_PARTS, 2, dtype=torch.int32), + empty(rows, heads), + torch.zeros(rows, heads, device=device, dtype=torch.int32), + empty(rows, heads, SELECT_PARTS), + empty(rows, heads, SELECT_PARTS, block_d), + ) + + +def _capture_workspace(device, rows, heads, capacity, block_d): + """Capture scratch reused across layers. + + Captures issued on one stream run in order, so each layer may overwrite + the previous layer's scratch; one entry per stream keeps memory bounded. + """ + stream = _current_stream(device) + shape = (rows, heads, capacity, block_d) + entry = _WORKSPACE.get((device, stream)) + if entry is None or entry[0] != shape: + entry = (shape, *_build_capture_workspace(device, rows, heads, capacity, block_d)) + _WORKSPACE[(device, stream)] = entry + return entry[1:] + + +_CONSUME_WORKSPACE: dict = {} + + +def _consume_workspace(device, rows, heads, block_d): + """Consume split parts, grown per stream and reused across layers.""" + stream = _current_stream(device) + entry = _CONSUME_WORKSPACE.get((device, stream)) + if ( + entry is None + or entry[0].shape[0] < rows + or entry[0].shape[1] != heads + or entry[2].shape[-1] != block_d + ): + empty = lambda *size: torch.empty(size, device=device, dtype=torch.float32) # noqa: E731 + entry = ( + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS), + empty(rows, heads, NUM_SPLITS, block_d), + ) + _CONSUME_WORKSPACE[(device, stream)] = entry + return (None, *entry) + + +def supports(plan, q, dense_attention) -> bool: + """Whether the fused path covers this submission; otherwise use the batched path.""" + modes = set(plan.modes) + return ( + q.is_cuda + and q.shape[-1] <= 256 + and plan.key_retention < 1.0 + and (modes == {CONSUME} or (modes == {CAPTURE} and dense_attention is not None)) + ) + + +def capture_workspace(device, rows, heads, capacity, dim): + """Capture scratch for at most ``rows`` rows, owned by the caller. + + A CUDA graph bakes the scratch addresses in, so its owner keeps one + workspace alive for every bucket instead of the shape-keyed cache. + """ + return _build_capture_workspace(device, rows, heads, capacity, triton.next_power_of_2(dim)) + + +def fused_paged_attention( + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan, + *, + layer, + dense_attention, + workspace=None, +): + """Capture or consume one uniform-mode submission without host readbacks. + + The kernels also write the store's ``valid``/``uses`` marks, so the host + only checks the plan's liveness instead of launching the bookkeeping. + ``workspace`` comes from ``capture_workspace`` and must cover the rows + and table width of the submission; consume uses only its split parts. + """ + store = plan.store + store._require_live(plan) + rows, heads, dim = q.shape + page = key_cache.shape[1] + groups = heads // key_cache.shape[2] + block_d = triton.next_power_of_2(dim) + scale = dim**-0.5 + for cache in (key_cache, value_cache): + if cache.stride(-1) != 1: + raise ValueError("fused sparse attention needs a contiguous head dimension") + if q.stride(-1) != 1: + raise ValueError("fused sparse attention needs a contiguous query head dimension") + selected = store.selected + if not selected.is_contiguous(): + raise ValueError("fused sparse attention derives index offsets from a contiguous pool") + if plan.modes[0] == CAPTURE: + # With the Triton backend, the dense split kernel also keeps the scores + # (and the first radix histograms), the dense merge runs inside the + # first selection pass, and selection runs in parallel parts. Other + # backends keep their call plus one serial selection kernel. + # Engines pass their backend adapter, which wraps the Triton kernel; + # callers may also pass the Triton function itself. + kernel = getattr(dense_attention, "kernel", dense_attention) + split = isinstance(kernel, TritonPagedAttention) or kernel is triton_paged_attention + if split: + output = torch.empty_like(q) + else: + output = dense_attention(q, key_cache, value_cache, block_tables, context_lengths) + if rows == 0: + return output + # The round's page tables bound every row's context on the host. + capacity = block_tables.shape[1] * page + top = min(retained_count(capacity, plan.key_retention), store.key_bucket or capacity) + store.ensure_key_capacity(top) + selected = store.selected + if workspace is None: + workspace = _capture_workspace(q.device, rows, heads, capacity, block_d) + elif workspace[0].shape[1:] != (heads, capacity) or workspace[0].shape[0] < rows: + raise ValueError("capture workspace does not cover the submission") + ( + scores, + part_m, + part_l, + part_o, + seghist, + state, + greater, + offsets, + lse, + tickets, + part_mass, + part_c, + ) = workspace if workspace[0].shape[0] == rows else (buffer[:rows] for buffer in workspace) + if split: + folded = _PASS0_IN_SPLIT + retention = _retention(q.device, plan.key_retention) + _launch( + _capture_split_kernel, + (rows, heads, NUM_SPLITS), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + part_m, + part_l, + part_o, + scores, + seghist, + retention, + state, + tickets, + top, + capacity, + *q.stride(), + *key_cache.stride(), + *value_cache.stride(), + block_tables.stride(0), + HEAD_DIM=dim, + HEAD_GROUPS=groups, + NUM_HEADS=heads, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=32, + SPLITS=NUM_SPLITS, + HISTOGRAM=folded, + PARTS=SELECT_PARTS, + num_warps=4, + ) + first_pass = 1 if folded else 0 + for level in range(first_pass, 4): + _launch( + _select_pass_kernel, + (rows, heads, SELECT_PARTS), + part_m, + part_l, + part_o, + output, + scores, + context_lengths, + retention, + plan.slot_ids, + store.counts, + store.output, + store.valid, + lse, + seghist, + state, + tickets, + greater, + offsets, + layer, + top, + capacity, + *output.stride(), + PASS=level, + MERGE=level == first_pass, + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_D=block_d, + BLOCK_T=_SCAN_BLOCK, + PARTS=SELECT_PARTS, + ALIGN=_SEGMENT_ALIGN, + SPLITS=NUM_SPLITS, + num_warps=_SCAN_WARPS, + ) + _launch( + _select_write_kernel, + (rows, heads, SELECT_PARTS), + scores, + context_lengths, + retention, + plan.slot_ids, + state, + offsets, + selected, + layer, + top, + capacity, + selected.shape[-1], + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + BLOCK_T=_SCAN_BLOCK, + PARTS=SELECT_PARTS, + ALIGN=_SEGMENT_ALIGN, + num_warps=_SCAN_WARPS, + ) + _launch( + _select_contribution_kernel, + (rows, heads, SELECT_PARTS), + scores, + value_cache, + block_tables, + context_lengths, + retention, + plan.slot_ids, + lse, + selected, + part_mass, + part_c, + tickets, + store.mass, + store.contribution, + layer, + top, + capacity, + block_tables.stride(0), + selected.shape[-1], + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + BLOCK_D=block_d, + BLOCK_T=_GATHER_BLOCK, + PARTS=SELECT_PARTS, + num_warps=4, + ) + return output + _launch( + _capture_kernel, + (rows, heads), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + _retention(q.device, plan.key_retention), + plan.slot_ids, + output, + scores, + selected, + store.counts, + store.mass, + store.output, + store.contribution, + store.valid, + layer, + top, + capacity, + block_tables.stride(0), + selected.shape[-1], + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + output.stride(0), + output.stride(1), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=64, + num_warps=4, + ) + return output + output = torch.empty_like(q) + if rows: + if workspace is None: + workspace = _consume_workspace(q.device, rows, heads, block_d) + elif workspace[1].shape[0] < rows or workspace[1].shape[1] != heads: + raise ValueError("consume workspace does not cover the submission") + part_m, part_l, part_o = ( + workspace[1:4] if workspace[1].shape[0] == rows else (b[:rows] for b in workspace[1:4]) + ) + _launch( + _consume_split_kernel, + (rows, heads, NUM_SPLITS), + q, + key_cache, + value_cache, + block_tables, + context_lengths, + plan.slot_ids, + selected, + store.counts, + part_m, + part_l, + part_o, + layer, + block_tables.stride(0), + selected.shape[-1], + q.stride(0), + q.stride(1), + key_cache.stride(0), + key_cache.stride(1), + key_cache.stride(2), + value_cache.stride(0), + value_cache.stride(1), + value_cache.stride(2), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + HEAD_GROUPS=groups, + PAGE_SIZE=page, + SCALE=scale, + BLOCK_D=block_d, + BLOCK_T=32, + SPLITS=NUM_SPLITS, + num_warps=4, + ) + _launch( + _consume_merge_kernel, + (rows, heads), + part_m, + part_l, + part_o, + context_lengths, + plan.slot_ids, + store.mass, + store.output, + store.contribution, + store.uses, + output, + layer, + output.stride(0), + output.stride(1), + NUM_LAYERS=store.layers, + NUM_HEADS=heads, + HEAD_DIM=dim, + BLOCK_D=block_d, + SPLITS=NUM_SPLITS, + num_warps=4, + ) + return output diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index 3555925f..98cf01ea 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -67,11 +67,36 @@ class AttentionReferenceStore: round. """ - def __init__(self, *, slots, layers, heads, head_dim, key_bucket, device, stream=None): + def __init__( + self, + *, + slots, + layers, + heads, + head_dim, + key_bucket, + device, + stream=None, + initial_key_capacity=None, + implementation="batched", + ): 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") + limit = key_bucket or 1 + initial = limit if initial_key_capacity is None else initial_key_capacity + if not 1 <= initial <= limit: + raise ValueError("initial key capacity must be within the reference capacity") + # Conservatively include all growth generations: an asynchronous + # stream may not have finished copying earlier buffers at the next + # growth. The engine budgets this before allocating its KV cache. + capacity = initial + self._peak_selected_capacity = capacity + while capacity < limit: + grown = min(limit, capacity * 2) + self._peak_selected_capacity += grown + capacity = grown self.slots = slots self.layers = layers self.heads = heads @@ -79,8 +104,10 @@ def __init__(self, *, slots, layers, heads, head_dim, key_bucket, device, stream self.key_bucket = key_bucket self.device = torch.device(device) self.stream = stream + # Which CUDA attention path consumes this store's plans. + self.implementation = implementation shape = (slots, layers, heads) - self.selected = torch.zeros((*shape, key_bucket or 1), dtype=torch.int32, device=device) + self.selected = torch.zeros((*shape, initial), 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) @@ -120,11 +147,54 @@ def bytes(self) -> int: ) ) + @property + def peak_capacity_bytes(self) -> int: + """Conservative payload bound including all in-flight growth generations.""" + index_bytes = self.slots * self.layers * self.heads * self._peak_selected_capacity * 4 + return self.bytes - self.selected.nbytes + index_bytes + + def ensure_key_capacity(self, count: int) -> None: + """Grow shared index storage on the attention owner stream. + + Existing plans borrow the store rather than a fixed index pointer. + Prior reads/writes and the copy execute on the same owner stream; + record_stream protects the old allocation until that copy completes. + The logical key bucket and retained count never shrink with allocation. + """ + limit = self.key_bucket or 1 + if count > limit: + raise ValueError(f"reference capacity {limit} cannot hold {count} retained keys") + current = self.selected.shape[-1] + if count <= current: + return + if self.device.type == "cuda" and self.stream is not None: + if torch.cuda.current_stream(self.device) != self.stream: + raise RuntimeError("reference growth must run on the attention owner stream") + capacity = current + while capacity < count: + capacity = min(limit, capacity * 2) + old = self.selected + grown = torch.zeros((*old.shape[:-1], capacity), dtype=old.dtype, device=old.device) + grown[..., :current].copy_(old) + if old.device.type == "cuda": + old.record_stream(torch.cuda.current_stream(old.device)) + self.selected = grown + @property def outstanding(self) -> int: """Number of slots currently bound to live tokens.""" return len(self._bindings) + def _index(self, values) -> torch.Tensor: + """Host-known slot or row indices on the device, without a blocking copy. + + A pageable host-to-device copy waits for the stream to drain, which + serializes the host with every in-flight round; a pinned source lets + the copy queue behind the work already submitted. + """ + host = torch.tensor(values, dtype=torch.int64, pin_memory=self.device.type == "cuda") + return host.to(self.device, non_blocking=True) + 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: @@ -137,15 +207,17 @@ def reserve(self, request_id: str, request_index: int, rows: int, generation: in 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) + slot_ids = self._index(slots) + # Writing host integers element by element would block on each copy. + self.incarnation.index_copy_(0, slot_ids, self._index(incarnations)) 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), + slot_ids=slot_ids, row_index=torch.arange(rows, dtype=torch.int64, device=self.device), ) @@ -168,11 +240,13 @@ def release(self, leases) -> None: cleared.append(slot) if cleared: self.free.extend(cleared) - self._clear(torch.tensor(cleared, dtype=torch.int64, device=self.device)) + self._forget_counts(cleared) + self._clear(self._index(cleared)) def discard(self) -> None: """Free every binding without device work; the caller has drained.""" self._bindings.clear() + self._counts_host.clear() self.free = list(range(self.slots)) def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan: @@ -210,7 +284,7 @@ def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan incarnations=tuple(incarnations), key_retention=key_retention, key_bucket=self.key_bucket, - slot_ids=torch.tensor(slots, dtype=torch.int64, device=self.device), + slot_ids=self._index(slots), store=self, ) @@ -233,6 +307,12 @@ def note_count(self, slot: int, layer: int, count: int) -> None: """ self._counts_host[(slot, layer)] = count + def _forget_counts(self, slots) -> None: + """Drop the host counts of host-known slots, so none sizes a later read.""" + for slot in slots: + for layer in range(self.layers): + self._counts_host.pop((int(slot), layer), None) + def host_count(self, slot: int, layer: int) -> int | None: """Retained keys of ``(slot, layer)``, or ``None`` before a capture.""" return self._counts_host.get((slot, layer)) @@ -244,10 +324,12 @@ def _row_slots(self, plan: SparseDecodePlan, rows: list[int]) -> torch.Tensor: a device tensor first; grouped rows are contiguous, so derive the index on the device instead. """ + if len(rows) == len(plan.slots): + return plan.slot_ids if rows[0] + len(rows) == rows[-1] + 1: index = torch.arange(rows[0], rows[0] + len(rows), device=self.device) return plan.slot_ids.index_select(0, index) - return plan.slot_ids[rows] + return self._index([plan.slots[row] for row in rows]) def capture(self, plan: SparseDecodePlan, layer: int) -> None: """Mark the captured rows source-ready for the given layer. @@ -261,8 +343,7 @@ def capture(self, plan: SparseDecodePlan, layer: int) -> None: if not rows: return # Assign a tensor, not the Python scalar ``True``: a scalar store into - # a bool tensor through advanced indexing costs milliseconds per call - # (2.6 ms measured at 32 rows), while a tensor store is ~0.13 ms. + # a bool tensor through advanced indexing copies it from the host. slots = self._row_slots(plan, rows) self.valid[slots, layer] = torch.ones(slots.numel(), dtype=torch.bool, device=self.device) @@ -280,7 +361,9 @@ def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> Non if positions.shape != plan.slot_ids.shape: raise ValueError("reference positions need one entry per plan row") rows = [row for row, slot in enumerate(plan.slots) if slot >= 0] - self.position[self._row_slots(plan, rows)] = positions[rows].to(torch.int64) + if len(rows) < len(plan.slots): + positions = positions[self._index(rows)] + self.position[self._row_slots(plan, rows)] = positions.to(torch.int64) def write_reference( self, @@ -298,9 +381,7 @@ def write_reference( 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.ensure_key_capacity(count) self.selected[slot, layer, :, :count] = selected.to(torch.int32) self.counts[slot, layer] = counts.to(torch.int32) self.note_count(slot, layer, int(count)) @@ -313,6 +394,9 @@ def read_reference(self, plan: SparseDecodePlan, layer: int, row: int): self._require_live(plan) slot = plan.slots[row] count = self.host_count(slot, layer) + if count is None: + # A fused capture decides its count on the device. + count = int(self.counts[slot, layer].max()) if not count: raise RuntimeError("a consume row has no captured reference") return ( @@ -351,15 +435,18 @@ def accept(self, leases, accepted: torch.Tensor) -> None: ) limit = accepted.to(device=self.device, dtype=rows.dtype).index_select(0, request) self.accepted[slot_ids] = limit.to(torch.int32) - 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 - for slot in rejected.tolist(): - for layer in range(self.counts.shape[1]): - self._counts_host.pop((slot, layer), None) + # A boolean-mask gather of the rejected slots would read its size back + # to the host. Masked writes over the round's own rows stay on device; + # they touch only these rows' per-layer metadata, never the index bucket. + keep = rows <= limit + self.valid[slot_ids] = self.valid[slot_ids] & keep[:, None] + self.counts[slot_ids] = torch.where(keep[:, None, None], self.counts[slot_ids], 0) + self.mass[slot_ids] = torch.where(keep[:, None, None], self.mass[slot_ids], 0) + # Which rows were rejected is known only on the device, so drop the host + # counts of every row of these leases: a later reader then sizes from + # the device counts, which are zero for rejected rows, and a stale host + # count can never size a gather. + self._forget_counts(slot for lease in leases for slot in lease.slots) def _clear(self, slots: torch.Tensor) -> None: """Retire the metadata of the given slots, ordered on the owner stream. @@ -371,16 +458,18 @@ def _clear(self, slots: torch.Tensor) -> None: """ def zero(): - for slot in slots.tolist(): - for layer in range(self.counts.shape[1]): - self._counts_host.pop((slot, layer), None) - 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.counts[slots] = 0 + # index_fill_ passes the scalar to the kernel; assigning a Python + # scalar through an index tensor copies it from the host and blocks. + for tensor in ( + self.valid, + self.uses, + self.accepted, + self.incarnation, + self.position, + self.mass, + self.counts, + ): + tensor.index_fill_(0, slots, 0) if self.stream is not None: with torch.cuda.stream(self.stream): @@ -392,13 +481,19 @@ 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) - for slot in slots.tolist(): - for layer in range(self.counts.shape[1]): - self._counts_host.pop((slot, layer), None) - self.valid[slots] = False - self.mass[slots] = 0 - self.output[slots] = 0 - self.contribution[slots] = 0 - self.selected[slots] = 0 - self.counts[slots] = 0 + if isinstance(slots, torch.Tensor): + # Callers hand host lists; a device tensor costs one readback here. + self._forget_counts(slots.tolist()) + slots = slots.to(device=self.device, dtype=torch.int64) + else: + self._forget_counts(slots) + slots = self._index(list(slots)) + for tensor in ( + self.valid, + self.mass, + self.output, + self.contribution, + self.selected, + self.counts, + ): + tensor.index_fill_(0, slots, 0) From d6c52eb7a866604279646192371d30abc9a840d0 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Sun, 1 Nov 2026 18:53:00 +0100 Subject: [PATCH 23/28] Build the retirement index on the reference store's owner stream A speculative ticket releases its leases from the host's current stream, but the store clears slot metadata on its owner stream. The slot index was copied and allocated on the current stream and freed when release returned, so a later allocation on that stream could overwrite it before the owner stream ran index_fill_, clearing other slots or indexing out of range. With CUDA graphs the host runs far enough ahead for this to surface as an index out-of-bounds device assert. The index is now built inside the owner stream. Signed-off-by: 0z5a --- tests/test_sparse_decode_attention.py | 31 ++++++++++++++++++++++++++ vllm_rlt/worker/attention_reference.py | 9 +++++++- 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index 4c4fc4c0..b21ca33b 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -632,6 +632,37 @@ def run(sparse): assert outstanding == 0 +@pytest.mark.gpu +def test_release_from_another_stream_builds_its_index_on_the_owner_stream(monkeypatch): + # A speculative ticket retires from the host's current stream while the + # store clears on its owner stream. An index built on the current stream is + # freed when release returns and can be reused there before the owner + # stream reads it, clearing some other slot. + owner = torch.cuda.Stream() + store = AttentionReferenceStore( + slots=8, layers=1, heads=1, head_dim=4, key_bucket=4, device="cuda", stream=owner + ) + lease = store.reserve("r", 0, 2, generation=1) + other = next(slot for slot in range(store.slots) if slot not in lease.slots) + store.valid.fill_(True) + torch.cuda.synchronize() + clear, streams = store._clear, [] + + def record(indices): + streams.append(torch.cuda.current_stream()) + clear(indices) + + monkeypatch.setattr(store, "_clear", record) + with torch.cuda.stream(owner): + torch.cuda._sleep(20_000_000) + store.release(lease) + poison = [torch.full((1,), other, device="cuda", dtype=torch.int64) for _ in range(32)] + torch.cuda.synchronize() + assert streams == [owner] and len(poison) == 32 + assert not bool(store.valid[list(lease.slots)].any()) + assert bool(store.valid[other].all()) + + @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 diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index 98cf01ea..63b34d6d 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -241,7 +241,14 @@ def release(self, leases) -> None: if cleared: self.free.extend(cleared) self._forget_counts(cleared) - self._clear(self._index(cleared)) + if self.stream is None: + self._clear(self._index(cleared)) + return + # A ticket retires from the host's current stream. An index built + # there would be freed, and reusable by that stream, before the + # owner stream's index_fill_ reads it; build it on the owner stream. + with torch.cuda.stream(self.stream): + self._clear(self._index(cleared)) def discard(self) -> None: """Free every binding without device work; the caller has drained.""" From b9ebfd6a612a99502bc1c4c4cc50e22c49ac24c5 Mon Sep 17 00:00:00 2001 From: kunkunblueberry Date: Wed, 4 Nov 2026 19:58:16 +0800 Subject: [PATCH 24/28] sparse decode attention: slot-scoped incarnations and a single context bound Four review items, one of them a real allocation-identity bug. **Stale plans could reach a new owner's slot.** Incarnations were counted per (request_id, row), so the next request to borrow a physical slot started again at 1. A plan built by the previous owner therefore carried the same number as the new owner's, passed _require_live, and could write into a reference slot it no longer owned. Counters now belong to the slot and only increase, which is what makes a recycled slot's old plan unambiguously stale. The new test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot reproduces the reported sequence (A reserves, releases; B reserves the same slot; A captures) and asserts both that the old plan is rejected and that the new owner's valid/counts/mass are unchanged. Verified to fail against the previous per-request counter with exactly the reported symptom: both owners see incarnation 1 on slot 0. **max_context_tokens is gone.** It imposed a request-length limit that the reference does not need: capacity is a buffer size, not an algorithm bound, and the index buffer already starts small and grows on demand. Removed the field, its validation, the admission check that duplicated the model-context check, the CLI option, the docs paragraph and the tests that existed only to exercise it. Capacity now follows the model's positional limit and the engine budgets the growth peak. **Naming.** The config was stored as self.sparse_decode (workers) and self.sparse_decode_attention (engine), which reads like an object rather than a configuration. Both are self.sparse_decode_config now, and the worker keyword follows. Regression: CPU 980 passed / 651 skipped; GPU (four sparse modules, async speculative, lifecycle) 217 passed / 4 skipped; ruff check and format clean under the CI-pinned 0.13.2. --- benchmarks/sparse_decode_engine_e2e.py | 7 +- benchmarks/sparse_decode_engine_e2e_text.py | 7 +- docs/runtime.md | 11 ++- tests/test_sparse_decode_attention.py | 85 ++++++++++++++------- tests/test_sparse_decode_sync.py | 14 ---- vllm_rlt/config.py | 6 -- vllm_rlt/engine/stage_engine.py | 29 +++---- vllm_rlt/entrypoints/runtime_args.py | 6 -- vllm_rlt/worker/async_speculative.py | 6 +- vllm_rlt/worker/attention_reference.py | 11 ++- vllm_rlt/worker/model_runner.py | 6 +- vllm_rlt/worker/speculative.py | 6 +- 12 files changed, 91 insertions(+), 103 deletions(-) diff --git a/benchmarks/sparse_decode_engine_e2e.py b/benchmarks/sparse_decode_engine_e2e.py index edd330d6..999f20f4 100644 --- a/benchmarks/sparse_decode_engine_e2e.py +++ b/benchmarks/sparse_decode_engine_e2e.py @@ -91,7 +91,6 @@ def main(): parser.add_argument("--retentions", default="1.0,0.1") parser.add_argument("--model-path", default=None) parser.add_argument("--cache-blocks", type=int, default=2048) - parser.add_argument("--sparse-context", type=int, default=64) args = parser.parse_args() device = args.device @@ -130,11 +129,7 @@ def main(): sparse = run( model, concurrency, - SparseDecodeAttentionConfig( - enabled=True, - key_retention=retention, - max_context_tokens=args.sparse_context, - ), + SparseDecodeAttentionConfig(enabled=True, key_retention=retention), prompt=args.prompt, max_tokens=args.max_tokens, repeats=args.repeats, diff --git a/benchmarks/sparse_decode_engine_e2e_text.py b/benchmarks/sparse_decode_engine_e2e_text.py index 734bbd9c..f45e7ecc 100644 --- a/benchmarks/sparse_decode_engine_e2e_text.py +++ b/benchmarks/sparse_decode_engine_e2e_text.py @@ -100,7 +100,6 @@ def main(): parser.add_argument("--repeats", type=int, default=2) parser.add_argument("--concurrency", default="1,8") parser.add_argument("--blocks", type=int, default=768) - parser.add_argument("--sparse-context", type=int, default=512) args = parser.parse_args() model = AutoModelForCausalLM.from_pretrained( @@ -132,11 +131,7 @@ def main(): sparse = run( model, concurrency, - SparseDecodeAttentionConfig( - enabled=True, - key_retention=retention, - max_context_tokens=args.sparse_context, - ), + SparseDecodeAttentionConfig(enabled=True, key_retention=retention), prompt_ids=prompt_ids, max_tokens=args.max_tokens, repeats=args.repeats, diff --git a/docs/runtime.md b/docs/runtime.md index 1853d5c4..ad385751 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -241,9 +241,7 @@ its own reference. ```python from vllm_rlt import ExecutionConfig, SparseDecodeAttentionConfig, SpeculativeConfig -sparse = SparseDecodeAttentionConfig( - enabled=True, key_retention=0.1, max_context_tokens=512, -) +sparse = SparseDecodeAttentionConfig(enabled=True, key_retention=0.1) # Pass sparse_decode_attention_config=sparse to LLM or StageEngine. # Ordinary decode uses synchronous scheduling and no speculative_config. # For synchronous block speculation, also pass: @@ -266,9 +264,10 @@ preemption/recovery and ordinary asynchronous decode are rejected explicitly. The retained count is `min(N, max(1, ceil(N * key_retention)))`, including the current position among the N causal keys. `key_bucket` controls buffer capacity, -not the algorithm: it must hold the retained count at the supported context -bound. `max_context_tokens=None` uses the model context; a smaller explicit -serving bound reduces pool memory and is enforced before request admission. +not the algorithm: it must hold the retained count over the model's positional +limit, which is the only context bound the reference needs. The index buffer +starts small and grows on the owner stream, so a long model context does not +cost its worst case up front; admission never rejects a request for being long. The pool holds one slot per live ordinary input token, K+1 per synchronous block request, and two rounds of K+1 per asynchronous block request. All reference payload bytes contribute to automatic KV memory planning; explicit diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index b21ca33b..c608f2db 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -312,19 +312,50 @@ def test_reincarnation_bumps_on_slot_reuse(): assert int(store.incarnation[slot]) == again.incarnations[0] +def test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot(): + # The incarnation must identify the physical slot's current owner. A counter + # keyed by (request_id, row) restarts at 1 for the next request, so the + # previous owner's plan carries the same number, passes the liveness check, + # and can overwrite the new owner's reference state. + store = make_store(slots=1, layers=1, heads=1, head_dim=2, key_bucket=4) + first = store.reserve("a", 0, 1, generation=1) + stale = store.build_plan([(first, 0)], [CAPTURE], key_retention=1.0) + store.release(first) + + second = store.reserve("b", 0, 1, generation=1) + assert second.slots == first.slots + current = store.build_plan([(second, 0)], [CAPTURE], key_retention=1.0) + assert current.incarnations != stale.incarnations + + store.capture(current, layer=0) + valid = store.valid.clone() + counts = store.counts.clone() + mass = store.mass.clone() + + with pytest.raises(RuntimeError, match="stale"): + store.capture(stale, layer=0) + with pytest.raises(RuntimeError, match="stale"): + store.write_reference( + stale, + 0, + 0, + selected=torch.tensor([[0, 1]]), + counts=torch.tensor([2]), + mass=torch.tensor([0.5]), + output=torch.zeros(1, 2), + contribution=torch.zeros(1, 2), + ) + assert torch.equal(store.valid, valid) + assert torch.equal(store.counts, counts) + assert torch.equal(store.mass, mass) + + 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_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") @@ -583,21 +614,30 @@ def test_async_sparse_reference_incarnations_advance_across_rounds(): e = make_speculative_engine(model, sparse=SparseDecodeAttentionConfig(enabled=True)) store = e.reference_store bootstrap(e) - incarnations = [] + plans = [] original = store.build_plan def recording(owners, modes, **kwargs): plan = original(owners, modes, **kwargs) - incarnations.append(plan.incarnations) + plans.append((plan.slots, 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] + assert len(plans) >= 2 + # Incarnations belong to the physical slot, so rounds may rotate through + # different slots and each fresh slot starts at 1. A slot must never come + # back with a number it already had; equality is allowed only because the + # capture and consume plans of one round share a lease, and therefore share + # its incarnation. + seen: dict[int, int] = {} + for slots, incarnations in plans: + for slot, incarnation in zip(slots, incarnations): + if slot in seen: + assert incarnation >= seen[slot], (slot, seen[slot], incarnation) + seen[slot] = incarnation e.abort_request("r") drain(e) assert store.outstanding == 0 @@ -693,27 +733,18 @@ def test_accept_and_retire_do_not_allocate_bucket_sized_temporaries(): @pytest.mark.gpu -def test_sparse_engine_bounds_reference_pool_by_serving_context(): +def test_sparse_engine_sizes_reference_capacity_from_the_model_context(): model = tiny("cuda") e = make_speculative_engine( - model, - sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, max_context_tokens=64), + model, sparse=SparseDecodeAttentionConfig(enabled=True, key_retention=1.0) ) - # The pool follows the supported serving bound, not the model's 128-token - # positional limit. - assert e.reference_store.key_bucket == 64 + # Capacity is the worst-case retained count over the model's positional + # limit; the buffer itself starts small and grows on demand. + assert e.reference_store.key_bucket == model.config.max_position_embeddings + assert e.reference_store.selected.shape[-1] < e.reference_store.key_bucket 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 diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py index cccbbc7f..5cd6e56d 100644 --- a/tests/test_sparse_decode_sync.py +++ b/tests/test_sparse_decode_sync.py @@ -117,20 +117,6 @@ def test_cancel_and_request_id_reuse_cannot_reuse_old_reference(draft): e.close() -def test_context_bound_is_enforced_before_request_admission(): - e = engine( - sparse_decode_attention_config=SparseDecodeAttentionConfig( - enabled=True, max_context_tokens=8 - ) - ) - try: - with pytest.raises(ValueError, match="sparse.*context"): - e.add_request("too-long", [2, 3, 4], SamplingParams(max_tokens=8)) - assert "too-long" not in e.scheduler.requests - finally: - e.close() - - @pytest.mark.parametrize("setting", ["graphs", "pipeline", "preemption", "shared"]) def test_unqualified_execution_modes_fail_explicitly(setting): options = { diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index d02c7c57..521487ba 100644 --- a/vllm_rlt/config.py +++ b/vllm_rlt/config.py @@ -208,10 +208,6 @@ 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 # ``batched`` keeps the scalar oracle's reduction shapes (bitwise contract); # ``fused`` launches fixed per-layer kernels without host readbacks and # changes consume-row numerics (see vllm_rlt.kernels.sparse_decode). @@ -226,5 +222,3 @@ 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 f4244e3e..fdfda895 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -125,7 +125,7 @@ def __init__( config.total_ut_steps * width ) sparse_decode = sparse_decode_attention_config - self.sparse_decode_attention = sparse_decode + self.sparse_decode_config = sparse_decode self.reference_store = None if sparse_decode is not None and sparse_decode.enabled: if getattr(model, "requires_boundary_kv", False): @@ -145,15 +145,10 @@ def __init__( "sparse decode attention ordinary decode requires synchronous scheduling" ) # The index buffer never truncates the algorithm: capacity must - # 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}" - ) + # hold the worst-case retained count over the model's positional + # limit. The buffer starts small and grows on demand, so a long + # model context does not cost its worst case up front. + context_bound = 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: @@ -186,7 +181,7 @@ def __init__( 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" + f"key_retention or num_speculative_tokens" ) self.reference_store = AttentionReferenceStore( slots=slots, @@ -321,7 +316,7 @@ def __init__( staged_prefill=self.scheduler.staged_prefill, rng=self.rng, reference_store=self.reference_store if speculative_config is None else None, - sparse_decode=self.sparse_decode_attention, + sparse_decode_config=self.sparse_decode_config, ) self.speculative_runner = ( SpeculativeRunner( @@ -333,7 +328,7 @@ def __init__( runtime=self.model_runner.runtime, compute_stream=self.model_runner.core_stream, reference_store=self.reference_store, - sparse_decode=self.sparse_decode_attention, + sparse_decode_config=self.sparse_decode_config, ) if speculative_config is not None and not self.async_speculative else None @@ -352,7 +347,7 @@ def __init__( rng=self.rng, runtime=self.model_runner.runtime, reference_store=self.reference_store, - sparse_decode=self.sparse_decode_attention, + sparse_decode_config=self.sparse_decode_config, ) if self.reference_store is not None: # Reference lifecycle device work is ordered on the compute @@ -469,12 +464,6 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i capacity = len(prompt_token_ids) + params.max_tokens - 1 if capacity > config.max_position_embeddings: raise ValueError("prompt plus decode positions exceed the model context length") - if self.reference_store is not None: - bound = ( - self.sparse_decode_attention.max_context_tokens or config.max_position_embeddings - ) - if capacity > bound: - raise ValueError("request exceeds sparse decode attention context bound") cache = self.cache_manager required = cache.required_blocks(capacity) if required > cache.num_blocks: diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index 5544d31a..58ae437c 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/vllm_rlt/entrypoints/runtime_args.py @@ -65,11 +65,6 @@ 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", @@ -154,7 +149,6 @@ 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/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 563fc082..1840686c 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -187,7 +187,7 @@ def __init__( rng: RngRegistry | None = None, runtime=None, reference_store=None, - sparse_decode=None, + sparse_decode_config=None, ): super().__init__( model, @@ -202,7 +202,7 @@ def __init__( # 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.sparse_decode_config = sparse_decode_config self._generations: dict[str, int] = {} sizes = speculative_capture_sizes(execution, scheduler, config) self.graphs = ( @@ -418,7 +418,7 @@ def _reference_plan(self, leases, rows, depth): 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 + owners, [mode] * len(owners), key_retention=self.sparse_decode_config.key_retention ) def _draft(self, plan, tables, allocations, leases=None): diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index 63b34d6d..e67820b0 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -122,7 +122,12 @@ def __init__( 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] = {} + # Incarnations belong to the physical slot, not to the request that + # borrows it: a counter keyed by ``(request_id, row)`` restarts at 1 for + # the next owner, so a stale plan from the previous owner would carry + # the same number and pass the liveness check. Slot counters only ever + # increase, which makes a recycled slot's old plan unambiguously stale. + self._slot_incarnations: list[int] = [0] * slots # Retained counts are host-known: the capture decides them before it # submits. Keeping them here lets a reader size its gather without a # device readback, which RFC #156 section 6.2 forbids. @@ -204,8 +209,8 @@ def reserve(self, request_id: str, request_index: int, rows: int, generation: in 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 + count = self._slot_incarnations[slot] + 1 + self._slot_incarnations[slot] = count self._bindings[slot] = (request_id, generation, count) incarnations.append(count) slot_ids = self._index(slots) diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index 77144bf4..58d5b471 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -162,12 +162,12 @@ def __init__( staged_prefill=None, rng: RngRegistry | None = None, reference_store=None, - sparse_decode=None, + sparse_decode_config=None, ): self.model = model.eval() self.cache_manager = cache_manager self.reference_store = reference_store - self.sparse_decode = sparse_decode + self.sparse_decode_config = sparse_decode_config self.attention_references = {} self.reference_generation = 0 self.metadata = cache_manager.metadata_builder @@ -1097,7 +1097,7 @@ def _reference_metadata(self, batch, kv): owners.extend([None] * padding) modes.extend([DENSE] * padding) plan = self.reference_store.build_plan( - owners, modes, key_retention=self.sparse_decode.key_retention + owners, modes, key_retention=self.sparse_decode_config.key_retention ) self.reference_store.bind_positions(plan, kv.position_ids) return replace(kv, reference_plan=plan) diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 5828e11f..39924dbf 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -154,7 +154,7 @@ def __init__( runtime=None, compute_stream=None, reference_store=None, - sparse_decode=None, + sparse_decode_config=None, ): self.model, self.cache, self.config = model, cache, config self.metadata = cache.metadata_builder @@ -163,7 +163,7 @@ def __init__( self.device = next(model.parameters()).device self.compute_stream = compute_stream self.reference_store = reference_store - self.sparse_decode = sparse_decode + self.sparse_decode_config = sparse_decode_config self.reference_generation = 0 self.submissions = SubmissionQueue() self.input_buffers = { @@ -250,7 +250,7 @@ def _reference_metadata(self, layouts, owners): result.append(kv) continue plan = self.reference_store.build_plan( - owners, modes, key_retention=self.sparse_decode.key_retention + owners, modes, key_retention=self.sparse_decode_config.key_retention ) self.reference_store.bind_positions(plan, kv.position_ids) result.append(replace(kv, reference_plan=plan)) From 16c62a92cbe3bfea28bcfb7b32fb709b66551d06 Mon Sep 17 00:00:00 2001 From: kunkunblueberry Date: Thu, 5 Nov 2026 12:54:27 +0800 Subject: [PATCH 25/28] Remove the benchmark scripts and validation reports this PR added Review asked for the evidence to live in the PR description and as attachments rather than as checked-in benchmark scripts and standalone reports. The regression tests stay; only the harnesses and the two write-ups go. - benchmarks/: sparse_decode_attention_bench, sparse_decode_batched_ab, sparse_decode_context_sweep, sparse_decode_engine_e2e, sparse_decode_engine_e2e_text, ttft_bench. - docs/validation/: sparse-decode-numerics-20261029.md and sparse-decode-sync-20261029.md. The numerical diagnosis is linked from docs/runtime.md as its release attachment instead. The test setup and the measured results now live in the PR description. --- benchmarks/sparse_decode_attention_bench.py | 196 -------------- benchmarks/sparse_decode_batched_ab.py | 169 ------------ benchmarks/sparse_decode_context_sweep.py | 242 ------------------ benchmarks/sparse_decode_engine_e2e.py | 151 ----------- benchmarks/sparse_decode_engine_e2e_text.py | 154 ----------- benchmarks/ttft_bench.py | 110 -------- docs/runtime.md | 7 +- .../sparse-decode-numerics-20261029.md | 145 ----------- .../validation/sparse-decode-sync-20261029.md | 96 ------- 9 files changed, 4 insertions(+), 1266 deletions(-) delete mode 100644 benchmarks/sparse_decode_attention_bench.py delete mode 100644 benchmarks/sparse_decode_batched_ab.py delete mode 100644 benchmarks/sparse_decode_context_sweep.py delete mode 100644 benchmarks/sparse_decode_engine_e2e.py delete mode 100644 benchmarks/sparse_decode_engine_e2e_text.py delete mode 100644 benchmarks/ttft_bench.py delete mode 100644 docs/validation/sparse-decode-numerics-20261029.md delete mode 100644 docs/validation/sparse-decode-sync-20261029.md diff --git a/benchmarks/sparse_decode_attention_bench.py b/benchmarks/sparse_decode_attention_bench.py deleted file mode 100644 index 3f0fe63f..00000000 --- a/benchmarks/sparse_decode_attention_bench.py +++ /dev/null @@ -1,196 +0,0 @@ -"""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/benchmarks/sparse_decode_batched_ab.py b/benchmarks/sparse_decode_batched_ab.py deleted file mode 100644 index 99541948..00000000 --- a/benchmarks/sparse_decode_batched_ab.py +++ /dev/null @@ -1,169 +0,0 @@ -"""A/B the scalar and batched sparse decode paths inside the merged tree. - -Both implementations live in ``vllm_rlt.attention.sparse`` and share the plan and -store contract, so this isolates the batching gain from the numerical contract: - -* ``scalar`` ``reference_paged_attention`` (row loop, the oracle), -* ``batched`` ``batched_paged_attention`` (equal-length row groups, CUDA). - -Capture rows are also measured with and without the dense-backend handoff, which -the full-retention control requires. - -Usage: - python benchmarks/sparse_decode_batched_ab.py --device cuda --rows 32 --context 512 -""" - -import argparse -import json -import math -import statistics -import time - -import torch - -from vllm_rlt.attention.sparse import ( - CAPTURE, - CONSUME, - batched_paged_attention, - reference_paged_attention, -) -from vllm_rlt.kernels.paged_attention import torch_paged_attention -from vllm_rlt.worker.attention_reference import AttentionReferenceStore - - -def build(rows, context, heads, head_dim, kv_heads, layers, block_size, retention, device): - blocks = math.ceil(context / block_size) - generator = torch.Generator().manual_seed(0) - shape = (blocks, block_size, kv_heads, head_dim) - key = torch.randn(shape, generator=generator).to(device, torch.bfloat16) - value = torch.randn(shape, generator=generator).to(device, torch.bfloat16) - table = torch.arange(blocks, device=device, dtype=torch.int32).repeat(rows, 1) - lengths = torch.full((rows,), context, device=device, dtype=torch.int32) - q = torch.randn((rows, heads, head_dim), generator=generator).to(device, torch.bfloat16) - capacity = math.ceil(context * retention) - store = AttentionReferenceStore( - slots=rows, - layers=layers, - heads=heads, - head_dim=head_dim, - key_bucket=capacity, - device=device, - ) - return key, value, table, lengths, q, store - - -def timed(fn, repeats, device): - for _ in range(3): - fn() - torch.cuda.synchronize(device) - samples = [] - for _ in range(repeats): - start = time.perf_counter() - fn() - torch.cuda.synchronize(device) - samples.append(time.perf_counter() - start) - return statistics.median(samples) * 1e3 - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cuda") - parser.add_argument("--rows", type=int, default=32) - 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("--block-size", type=int, default=16) - parser.add_argument("--retention", type=float, default=0.1) - parser.add_argument("--repeats", type=int, default=10) - parser.add_argument("--json", action="store_true") - args = parser.parse_args() - device = torch.device(args.device) - - report = { - "device": str(device), - "rows": args.rows, - "context": args.context, - "retention": args.retention, - "layers": args.layers, - } - for mode, name in ((CAPTURE, "capture"), (CONSUME, "consume")): - key, value, table, lengths, q, store = build( - args.rows, - args.context, - args.heads, - args.head_dim, - args.kv_heads, - args.layers, - args.block_size, - args.retention, - device, - ) - lease = store.reserve("r", 0, args.rows, generation=1) - owners = [(lease, row) for row in range(args.rows)] - # Seed a capture pass so consume rows have state. - seed_plan = store.build_plan(owners, [CAPTURE] * args.rows, key_retention=args.retention) - reference_paged_attention( - q, - key, - value, - table, - lengths, - seed_plan, - layer=0, - dense_attention=torch_paged_attention, - ) - plan = store.build_plan(owners, [mode] * args.rows, key_retention=args.retention) - dense = torch_paged_attention - scalar = timed( - lambda: reference_paged_attention( - q, key, value, table, lengths, plan, layer=0, dense_attention=dense - ), - args.repeats, - device, - ) - batched = timed( - lambda: batched_paged_attention( - q, key, value, table, lengths, plan, layer=0, dense_attention=dense - ), - args.repeats, - device, - ) - batched_oracle = timed( - lambda: batched_paged_attention(q, key, value, table, lengths, plan, layer=0), - args.repeats, - device, - ) - report[name] = { - "scalar_ms": round(scalar, 3), - "batched_ms": round(batched, 3), - "batched_no_dense_handoff_ms": round(batched_oracle, 3), - "speedup": round(scalar / batched, 2), - } - store.release(lease) - - dense_ms = timed( - lambda: torch_paged_attention(q, key, value, table, lengths), args.repeats, device - ) - report["dense_attention_ms"] = round(dense_ms, 3) - if args.json: - print(json.dumps(report, indent=2)) - return - print( - f"sparse decode path A/B — {device}, rows={args.rows}, context={args.context}, " - f"retention={args.retention}, layers={args.layers}" - ) - for name in ("capture", "consume"): - values = report[name] - print( - f" {name:8s} scalar {values['scalar_ms']:7.3f} ms " - f"batched {values['batched_ms']:7.3f} ms " - f"speedup {values['speedup']:5.2f}x " - f"(batched w/o dense handoff {values['batched_no_dense_handoff_ms']:7.3f} ms)" - ) - print(f" dense attention reference: {dense_ms:.3f} ms") - - -if __name__ == "__main__": - main() diff --git a/benchmarks/sparse_decode_context_sweep.py b/benchmarks/sparse_decode_context_sweep.py deleted file mode 100644 index 6eae0687..00000000 --- a/benchmarks/sparse_decode_context_sweep.py +++ /dev/null @@ -1,242 +0,0 @@ -"""Dense-baseline context sweep and memory breakdown for sparse decode attention. - -The acceptance baseline is the dense Triton kernel, not the scalar sparse -reference: a four-loop decode costs ``4 x dense`` with the feature off and -``dense + capture + 2 x consume`` with it on. This script measures both on -identical paged KV at increasing context, reports the ratio, and breaks memory -down into KV storage, the persistent reference pool and transient workspace. - -Usage: - python benchmarks/sparse_decode_context_sweep.py --device cuda - python benchmarks/sparse_decode_context_sweep.py --contexts 512,2048,8192 --rows 32 -""" - -import argparse -import json -import math -import statistics -import time - -import torch - -from vllm_rlt.attention.sparse import ( - CAPTURE, - CONSUME, - reference_paged_attention, - sparse_paged_attention, -) -from vllm_rlt.kernels.paged_attention import triton_paged_attention -from vllm_rlt.worker.attention_reference import AttentionReferenceStore - - -def build(context, rows, heads, head_dim, kv_heads, layers, block_size, retention, device): - blocks = math.ceil(context / block_size) - generator = torch.Generator().manual_seed(0) - shape = (blocks, block_size, kv_heads, head_dim) - key = torch.randn(shape, generator=generator).to(device, torch.bfloat16) - value = torch.randn(shape, generator=generator).to(device, torch.bfloat16) - table = torch.arange(blocks, device=device, dtype=torch.int32).repeat(rows, 1) - lengths = torch.full((rows,), context, device=device, dtype=torch.int32) - q = torch.randn((rows, heads, head_dim), generator=generator).to(device, torch.bfloat16) - capacity = math.ceil(context * retention) - store = AttentionReferenceStore( - slots=rows, - layers=layers, - heads=heads, - head_dim=head_dim, - key_bucket=capacity, - device=device, - ) - return key, value, table, lengths, q, store - - -def timed(fn, repeats, device): - for _ in range(2): - fn() - torch.cuda.synchronize(device) - samples = [] - for _ in range(repeats): - start = time.perf_counter() - fn() - torch.cuda.synchronize(device) - samples.append(time.perf_counter() - start) - return statistics.median(samples) * 1e3 - - -def transient_bytes(fn, device): - torch.cuda.synchronize(device) - torch.cuda.reset_peak_memory_stats(device) - before = torch.cuda.memory_allocated(device) - fn() - torch.cuda.synchronize(device) - return max(0, torch.cuda.max_memory_allocated(device) - before) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cuda") - parser.add_argument("--contexts", default="256,512,1024,2048,4096,8192,16384,32768") - parser.add_argument("--rows", type=int, default=32) - 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("--block-size", type=int, default=16) - parser.add_argument("--retention", type=float, default=0.1) - parser.add_argument("--repeats", type=int, default=5) - parser.add_argument("--json", action="store_true") - args = parser.parse_args() - device = torch.device(args.device) - contexts = [int(value) for value in args.contexts.split(",")] - - rows = [] - for context in contexts: - try: - key, value, table, lengths, q, store = build( - context, - args.rows, - args.heads, - args.head_dim, - args.kv_heads, - args.layers, - args.block_size, - args.retention, - device, - ) - lease = store.reserve("r", 0, args.rows, generation=1) - owners = [(lease, row) for row in range(args.rows)] - seed = store.build_plan(owners, [CAPTURE] * args.rows, key_retention=args.retention) - sparse_paged_attention( - q, key, value, table, lengths, seed, layer=0, dense_attention=triton_paged_attention - ) - capture_plan = store.build_plan( - owners, [CAPTURE] * args.rows, key_retention=args.retention - ) - consume_plan = store.build_plan( - owners, [CONSUME] * args.rows, key_retention=args.retention - ) - dense_ms = timed( - lambda: triton_paged_attention(q, key, value, table, lengths), - args.repeats, - device, - ) - capture_ms = timed( - lambda: sparse_paged_attention( - q, - key, - value, - table, - lengths, - capture_plan, - layer=0, - dense_attention=triton_paged_attention, - ), - args.repeats, - device, - ) - consume_ms = timed( - lambda: sparse_paged_attention( - q, key, value, table, lengths, consume_plan, layer=0 - ), - args.repeats, - device, - ) - scalar_consume_ms = timed( - lambda: reference_paged_attention( - q, key, value, table, lengths, consume_plan, layer=0 - ), - args.repeats, - device, - ) - dense_total = 4 * dense_ms - sparse_total = dense_ms + capture_ms + 2 * consume_ms - rows.append( - dict( - context=context, - retained=math.ceil(context * args.retention), - dense_ms=round(dense_ms, 3), - capture_ms=round(capture_ms, 3), - consume_ms=round(consume_ms, 3), - scalar_consume_ms=round(scalar_consume_ms, 3), - dense_4loop_ms=round(dense_total, 3), - sparse_4loop_ms=round(sparse_total, 3), - sparse_over_dense=round(sparse_total / dense_total, 3), - capture_transient_mib=round( - transient_bytes( - lambda: sparse_paged_attention( - q, - key, - value, - table, - lengths, - capture_plan, - layer=0, - dense_attention=triton_paged_attention, - ), - device, - ) - / 1024**2, - 1, - ), - consume_transient_mib=round( - transient_bytes( - lambda: sparse_paged_attention( - q, key, value, table, lengths, consume_plan, layer=0 - ), - device, - ) - / 1024**2, - 1, - ), - reference_pool_mib=round(rows_pool_bytes(args, context) / 1024**2, 1), - kv_mib=round(key.numel() * key.element_size() * 2 / 1024**2, 1), - ) - ) - store.release(lease) - torch.cuda.empty_cache() - except torch.OutOfMemoryError: - rows.append(dict(context=context, status="oom")) - torch.cuda.empty_cache() - - if args.json: - print(json.dumps(rows, indent=2)) - return - print( - f"dense-baseline sweep — {device}, rows={args.rows}, layers={args.layers}, " - f"heads={args.heads}, head_dim={args.head_dim}, retention={args.retention}" - ) - print( - f"{'ctx':>7}{'k':>7}{'dense':>9}{'capture':>9}{'consume':>9}" - f"{'4-loop dense':>13}{'4-loop sparse':>14}{'ratio':>8}" - ) - for row in rows: - if row.get("status") == "oom": - print(f"{row['context']:>7} OOM") - continue - print( - f"{row['context']:>7}{row['retained']:>7}{row['dense_ms']:>9.2f}" - f"{row['capture_ms']:>9.2f}{row['consume_ms']:>9.2f}" - f"{row['dense_4loop_ms']:>13.2f}{row['sparse_4loop_ms']:>14.2f}" - f"{row['sparse_over_dense']:>8.2f}x" - ) - print() - print(f"{'ctx':>7}{'capture ws':>12}{'consume ws':>12}{'ref pool':>11}{'KV':>9}") - for row in rows: - if row.get("status") == "oom": - continue - print( - f"{row['context']:>7}{row['capture_transient_mib']:>10.1f}Mi" - f"{row['consume_transient_mib']:>10.1f}Mi" - f"{row['reference_pool_mib']:>9.1f}Mi{row['kv_mib']:>7.1f}Mi" - ) - - -def rows_pool_bytes(args, context): - """Persistent reference pool for this sweep's row count.""" - capacity = math.ceil(context * args.retention) - per_slot = capacity * 4 + 4 + 4 * 2 * args.head_dim + 8 - return args.rows * args.layers * args.heads * per_slot - - -if __name__ == "__main__": - main() diff --git a/benchmarks/sparse_decode_engine_e2e.py b/benchmarks/sparse_decode_engine_e2e.py deleted file mode 100644 index 999f20f4..00000000 --- a/benchmarks/sparse_decode_engine_e2e.py +++ /dev/null @@ -1,151 +0,0 @@ -"""Engine-level end-to-end comparison for sparse decode attention. - -Runs the real StageEngine with asynchronous self-speculation and compares dense -against sparse decode at full and reduced retention: throughput, peak memory and -token agreement. The per-layer sweeps say what one attention call costs; this -says whether that survives scheduling, sampling and delivery. - -The tiny Ouro architecture is used because no trained checkpoint is available on -this host, so the numbers are a scheduler-level check, not a serving result. -""" - -import argparse -import statistics -import time - -import torch - -from tests.helpers import tiny_ouro_model -from vllm_rlt import ( - CacheConfig, - ExecutionConfig, - SamplingParams, - SchedulerConfig, - SparseDecodeAttentionConfig, - SpeculativeConfig, -) -from vllm_rlt.engine.stage_engine import StageEngine -from vllm_rlt.models import AutoModelForCausalLM - - -def build(model, concurrency, sparse, blocks=2048): - return StageEngine( - model, - attention_backend="triton", - cache_config=CacheConfig(blocks, 16, incremental_allocation=True), - scheduler_config=SchedulerConfig( - max_num_seqs=concurrency, - max_num_batched_tokens=8 * concurrency, - prefill_chunk_size=16, - ), - execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), - speculative_config=SpeculativeConfig(3), - sparse_decode_attention_config=sparse, - ) - - -def run(model, concurrency, sparse, *, prompt, max_tokens, repeats, blocks=2048): - """One configuration: returns median wall time, tokens/s and peak bytes.""" - outputs = None - walls = [] - peaks = [] - for _ in range(repeats): - engine = build(model, concurrency, sparse, blocks) - torch.cuda.reset_peak_memory_stats() - params = [ - SamplingParams(max_tokens=max_tokens, ignore_eos=True) for _ in range(concurrency) - ] - prompts = [ - [2 + (index + position) % 60 for position in range(prompt)] - for index in range(concurrency) - ] - start = time.perf_counter() - for index in range(concurrency): - engine.add_request(f"r{index}", prompts[index], params[index]) - collected = {} - for _ in range(100000): - for output in engine.step(): - collected[output.request_id] = output - if len(collected) == concurrency and all(o.finished for o in collected.values()): - break - torch.cuda.synchronize() - walls.append(time.perf_counter() - start) - peaks.append(torch.cuda.max_memory_allocated()) - outputs = {key: list(value.token_ids) for key, value in collected.items()} - tokens = concurrency * max_tokens - return { - "wall": statistics.median(walls), - "tokens_per_s": tokens / statistics.median(walls), - "peak_gib": max(peaks) / 2**30, - "outputs": outputs, - } - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--device", default="cuda") - parser.add_argument("--prompt", type=int, default=32) - parser.add_argument("--max-tokens", type=int, default=16) - parser.add_argument("--repeats", type=int, default=3) - parser.add_argument("--concurrency", default="1,8,32") - parser.add_argument("--retentions", default="1.0,0.1") - parser.add_argument("--model-path", default=None) - parser.add_argument("--cache-blocks", type=int, default=2048) - args = parser.parse_args() - - device = args.device - if args.model_path: - model = AutoModelForCausalLM.from_pretrained( - args.model_path, device=device, dtype=torch.bfloat16 - ) - else: - model = tiny_ouro_model(device) - concurrency_levels = [int(value) for value in args.concurrency.split(",")] - retentions = [float(value) for value in args.retentions.split(",")] - - label = args.model_path or "tiny Ouro" - print( - f"engine e2e — {label} ({model.config.num_hidden_layers}L, ctx " - f"{model.config.max_position_embeddings}), prompt {args.prompt}, " - f"max_tokens {args.max_tokens}, median of {args.repeats}" - ) - header = f"{'conc':>5}{'case':>14}{'tok/s':>10}{'wall ms':>10}{'peak GiB':>10}{'match':>10}" - print(header) - for concurrency in concurrency_levels: - dense = run( - model, - concurrency, - None, - prompt=args.prompt, - max_tokens=args.max_tokens, - repeats=args.repeats, - blocks=args.cache_blocks, - ) - print( - f"{concurrency:>5}{'dense':>14}{dense['tokens_per_s']:>10.2f}" - f"{dense['wall'] * 1e3:>10.1f}{dense['peak_gib']:>10.2f}{'--':>10}" - ) - for retention in retentions: - sparse = run( - model, - concurrency, - SparseDecodeAttentionConfig(enabled=True, key_retention=retention), - prompt=args.prompt, - max_tokens=args.max_tokens, - repeats=args.repeats, - blocks=args.cache_blocks, - ) - same = sum( - 1 - for key, tokens in sparse["outputs"].items() - if tokens == dense["outputs"].get(key) - ) - print( - f"{'':>5}{f'sparse r={retention}':>14}{sparse['tokens_per_s']:>10.2f}" - f"{sparse['wall'] * 1e3:>10.1f}{sparse['peak_gib']:>10.2f}" - f"{f'{same}/{concurrency}':>10}" - ) - - -if __name__ == "__main__": - main() diff --git a/benchmarks/sparse_decode_engine_e2e_text.py b/benchmarks/sparse_decode_engine_e2e_text.py deleted file mode 100644 index f45e7ecc..00000000 --- a/benchmarks/sparse_decode_engine_e2e_text.py +++ /dev/null @@ -1,154 +0,0 @@ -"""Engine e2e with a real-text prompt, for a trained checkpoint. - -Synthetic token-id prompts drive a trained model into a fixed point (one token -repeated), which makes token agreement meaningless: every implementation agrees -when the argmax is constant. This encodes real text with the checkpoint's own -tokenizer and reports throughput, peak memory and per-request token agreement. - -Usage: - python benchmarks/sparse_decode_engine_e2e_text.py \ - --model-path /path/to/Ouro-1.4B --concurrency 1,8 -""" - -import argparse -import statistics -import time - -import torch -from tokenizers import Tokenizer - -from vllm_rlt import ( - CacheConfig, - ExecutionConfig, - SamplingParams, - SchedulerConfig, - SparseDecodeAttentionConfig, - SpeculativeConfig, -) -from vllm_rlt.engine.stage_engine import StageEngine -from vllm_rlt.models import AutoModelForCausalLM - -TEXT = ( - "The history of computing is a history of abstraction. Each generation of " - "hardware made a new kind of program possible, and each new kind of program " - "demanded a new way of thinking about machines. When memory was scarce, " - "programmers counted bytes by hand; when it became plentiful, they wrote " - "interpreters that treated memory as infinite. The same shift is happening " - "again with inference: models that once fit on a single device now span " - "racks, and the scheduler rather than the kernel decides how fast a request " - "finishes. Attention is the expensive part, because every token must look at " - "every earlier token, and that cost grows with the square of the context. " - "Sparse attention tries to keep only the keys that matter, which is easy to " - "state and hard to make fast, because deciding what matters is itself as " - "expensive as attending to everything." -) - - -def build(model, concurrency, sparse, blocks): - return StageEngine( - model, - attention_backend="triton", - cache_config=CacheConfig(blocks, 16, incremental_allocation=True), - scheduler_config=SchedulerConfig( - max_num_seqs=concurrency, - max_num_batched_tokens=8 * concurrency, - prefill_chunk_size=16, - ), - execution_config=ExecutionConfig(async_scheduling=True, multi_stream=True), - speculative_config=SpeculativeConfig(3), - sparse_decode_attention_config=sparse, - ) - - -def run(model, concurrency, sparse, *, prompt_ids, max_tokens, repeats, blocks): - walls = [] - peaks = [] - outputs = None - for _ in range(repeats): - engine = build(model, concurrency, sparse, blocks) - torch.cuda.reset_peak_memory_stats() - params = [ - SamplingParams(max_tokens=max_tokens, ignore_eos=True) for _ in range(concurrency) - ] - start = time.perf_counter() - for index in range(concurrency): - engine.add_request(f"r{index}", list(prompt_ids), params[index]) - collected = {} - for _ in range(100000): - for output in engine.step(): - collected[output.request_id] = output - if len(collected) == concurrency and all(o.finished for o in collected.values()): - break - torch.cuda.synchronize() - walls.append(time.perf_counter() - start) - peaks.append(torch.cuda.max_memory_allocated()) - outputs = {key: list(value.token_ids) for key, value in collected.items()} - tokens = concurrency * max_tokens - return { - "tokens_per_s": tokens / statistics.median(walls), - "wall": statistics.median(walls), - "peak_gib": max(peaks) / 2**30, - "outputs": outputs, - } - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--model-path", required=True) - parser.add_argument("--prompt", type=int, default=200) - parser.add_argument("--max-tokens", type=int, default=16) - parser.add_argument("--repeats", type=int, default=2) - parser.add_argument("--concurrency", default="1,8") - parser.add_argument("--blocks", type=int, default=768) - args = parser.parse_args() - - model = AutoModelForCausalLM.from_pretrained( - args.model_path, device="cuda", dtype=torch.bfloat16 - ) - tokenizer = Tokenizer.from_file(f"{args.model_path}/tokenizer.json") - ids = tokenizer.encode(TEXT).ids - while len(ids) < args.prompt: - ids = ids + ids - prompt_ids = ids[: args.prompt] - print(f"real text prompt: {args.prompt} tokens, head={prompt_ids[:8]}") - - print(f"{'conc':>5}{'case':>16}{'tok/s':>10}{'wall ms':>10}{'peak GiB':>10}{'match':>10}") - for concurrency in [int(v) for v in args.concurrency.split(",")]: - dense = run( - model, - concurrency, - None, - prompt_ids=prompt_ids, - max_tokens=args.max_tokens, - repeats=args.repeats, - blocks=args.blocks, - ) - print( - f"{concurrency:>5}{'dense':>16}{dense['tokens_per_s']:>10.2f}" - f"{dense['wall'] * 1e3:>10.1f}{dense['peak_gib']:>10.2f}{'--':>10}" - ) - for retention in (1.0, 0.1): - sparse = run( - model, - concurrency, - SparseDecodeAttentionConfig(enabled=True, key_retention=retention), - prompt_ids=prompt_ids, - max_tokens=args.max_tokens, - repeats=args.repeats, - blocks=args.blocks, - ) - same = sum( - 1 - for key, tokens in sparse["outputs"].items() - if tokens == dense["outputs"].get(key) - ) - print( - f"{'':>5}{f'sparse r={retention}':>16}{sparse['tokens_per_s']:>10.2f}" - f"{sparse['wall'] * 1e3:>10.1f}{sparse['peak_gib']:>10.2f}" - f"{f'{same}/{concurrency}':>10}" - ) - print(f" dense sample tokens: {dense['outputs']['r0'][:8]}") - - -if __name__ == "__main__": - main() diff --git a/benchmarks/ttft_bench.py b/benchmarks/ttft_bench.py deleted file mode 100644 index 9250710e..00000000 --- a/benchmarks/ttft_bench.py +++ /dev/null @@ -1,110 +0,0 @@ -"""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() diff --git a/docs/runtime.md b/docs/runtime.md index ad385751..f7ce7fce 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -281,9 +281,10 @@ verification captures loop 2; with draft depth 2, drafting captures it. Full retention uses the same dense backend at later loops as a control, while still validating and consuming the captured reference lifecycle. The standalone Torch arithmetic oracle is tested separately; different attention algorithms -need not round to identical BF16 values. See the -[numerical diagnosis and regression report](validation/sparse-decode-numerics-20261029.md). -Retention below 1 is approximate and +need not round to identical BF16 values. The +[numerical diagnosis](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-numerics-validation-20261029) +is kept as a release attachment rather than a checked-in report. Retention below +1 is approximate and needs task-quality evaluation before a deployment recommendation. ## SHARED prefill semantics diff --git a/docs/validation/sparse-decode-numerics-20261029.md b/docs/validation/sparse-decode-numerics-20261029.md deleted file mode 100644 index 106918c7..00000000 --- a/docs/validation/sparse-decode-numerics-20261029.md +++ /dev/null @@ -1,145 +0,0 @@ -# Sparse decode: backend numerical diagnosis - -This follow-up fixes the full-retention discrepancy in the -[ordinary/synchronous integration report](sparse-decode-sync-20261029.md), within -the same Draft collaboration PR targeting `kunkun/rfc156-sparse-decode-async`. -The immediate feature base remains `048b9d7c79a5711857ce3287812502ad8f274ad0`; -the pre-fix contribution head is `7f22af3901416096ea6c1c0e0a5536f440f662f5`. -The fixed runtime/test commit is `1bbe6fdd24be3e9bbea6a9dd1440d4b64f8dd89b`. - -## Reproduced cause - -The original BF16 Ouro-1.4B workload reproduces 498/512 token agreement at -concurrent arrivals 32 and synchronous draft depth 1. Request 9 first differs -at output index 2: dense chooses 293, the original sparse reference chooses 92. -The 14 differences include subsequent autoregressive continuation, rather than -14 independent initial errors. - -The new diagnostic retains the original 128-input/16-output protocol, including -four-token warmup. It freezes the complete accepted verification batch's input -hidden states, logical row/depth/position metadata, position IDs and context -lengths, all physical-layer module outputs, full packed queries and attention -outputs, and the watched row's causal K/V. It filters out rejected executions of -the same position. The accepted loop-2 batch here has **22 rows**, not the -theoretical 96-row maximum; both arms have exactly equal packed input states, -request histories, shapes, positions and lengths. A shorter-output exploratory -run changed the batch geometry and is retained separately, not substituted for -this reproduction. - -For the watched row, physical layers 0 and 1 agree. At physical layer **2** -(zero-based), Q, all 130 legal keys and values are identical, but the attention -output differs at one BF16 element by **0.000030517578125**. The next `o_proj` -differs at two elements, max **0.000244140625**. Across the complete batch, -attention differences already exist at physical layer 0. Later layers and -recurrences amplify these differences. This places the first observed watched -operator discrepancy at attention, before reference consumption. - -A separate replay uses the frozen operands with no model or scheduler. Triton -split-KV reproduces the captured dense output exactly; the Torch -einsum/softmax reference reproduces the captured sparse output exactly. Changing -only the backend reproduces the layer-2 difference. Disabling BF16 reduced -precision reduction does not change that reference result; its reductions -already use FP32 operands. TF32 is disabled throughout. Equal packed shapes -rule out a batch-shape change as the explanation for this captured attention -discrepancy. No global precision flags or model GEMM implementations are changed. - -## Fix and numerical contract - -`KVCacheManager` passes its selected dense attention adapter into the shared -`sparse_paged_attention` dispatcher. The adapter receives the original packed -queries, tables and lengths, without slicing into different batch shapes. - -- DENSE and CAPTURE rows use that backend's actual output. Capture stores that - output as FP32 **after** the backend's output-dtype rounding, so `O_ref` is - the value entering recurrence. Retained selection, mass and contribution - remain the FP32 Torch reference arithmetic. -- At retention 1, CONSUME rows also use the configured dense backend, while - still checking reference ownership/payload availability and recording - consumption. This is explicitly a dense control; it does not qualify the - Torch correction as bitwise equal to Triton. -- At retention below 1, later loops retain `O_ref - C_ref + m * C_new`. This - changes the original reference's source-output rounding contract, not the - selection or sparse delta formula. Approximate quality remains unqualified. -- Omitting the adapter when calling the standalone paged oracle forces the - original Torch capture and correction, including at retention 1. Known-answer - paged cases test changed later keys/values at retention .5 and 1; the full-retention - control cannot hide a regression in that independent arithmetic path. - -A full-retention-only bypass would leave partial-retention capture rows using -a different dense algorithm; preserving the selected backend at every dense -source row fixes that inconsistency too. No new kernel or global precision -setting is introduced. The fallback oracle remains independently available. - -## Regression evidence - -Four CUDA BF16 cases cover MHA/GQA, retention .1/1, a mixed dense/capture plan, -stored `O_ref`, changed later-loop operands and reference consumption. All four -fail on the original implementation with backend-output differences, then pass -with the fix using exact equality. The standalone paged known-answer cases -return 23.5 and 40.0 after changing later keys/values, exercising partial and -full Torch correction without the backend control. - -Full CPU regression: **921 passed, 479 skipped**. Full pre-commit passes Ruff -lint, Ruff format and the CPU suite. Independent read-only review found no -actionable defect in the four-file runtime/test delta; its targeted CPU check -passed 48 cases, with 22 GPU cases deselected. - -Applicable CUDA regression: **318 passed, 1 skipped, 16 deselected**. The skip -requires two GPUs; FA4 cases remain excluded on this SM89 device. A fixed -instrumented replay now matches the original dense replay exactly at every -recorded depth: complete packed inputs/outputs, all physical module outputs, -attention queries/outputs and the watched row's causal K/V. This comparison is -separate from the uninstrumented generation matrix. - -The uninstrumented BF16 regression matrix covers ordinary decode and synchronous -draft depths 1/2, simultaneous arrivals 1/8/32, and retention off/1/.1. Each of -the 27 fresh-process configurations warms four output tokens, then measures -three bursts using the original 128-input/16-output protocol. All **81 bursts, -1,107 requests and 17,712 output tokens** complete, with KV/reference state -drained. Tokens and exit depths repeat exactly within each configuration. - -| Mode | C=1 full vs dense | C=8 full vs dense | C=32 full vs dense | -|---|---:|---:|---:| -| Ordinary | 16/16 | 128/128 | 512/512 | -| Sync, draft 1 | 16/16 | 128/128 | 512/512 | -| Sync, draft 2 | 16/16 | 128/128 | 512/512 | - -Every cell matches in all three repetitions, including exit depths: **5,904/5,904 -full-retention tokens** in total. In particular, the original C32/draft-1 case -improves from 498/512 to 512/512 in each repetition. This verifies the explicit -dense-backend control; the forced Torch correction remains a separate oracle. -Retention .1 completes and drains, but its token agreement is not a quality score. - -## Reproduction and archive - -The [numerical validation release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-numerics-validation-20261029) -contains the frozen tensor snapshots, source and SHA256 inventories, Git bundle, -same-input replay, original failures, red/green regression logs, full raw matrix, -environment metadata and reproduction commands. Model weights are not included. -The original unsupported-FA4 failures and performance profiles remain in the -[earlier release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-sync-validation-20261029). - -Use `trace_first_operator.py` against `7f22af3` for the dense and pre-fix capture, -then against `1bbe6fd` for the fixed capture; `replay_attention.py` tests the -frozen operands independently. `compare_fixed_trace.py` checks every recorded -state, and `run_regression_matrix.py` / `summarize_matrix.py` reproduce and verify -the complete generation matrix. These local experiment scripts preserve their -recorded root directory; the archive README explains relocation and exact CLI -commands. Diagnostic timings are excluded from uninstrumented results. - -## Scope - -The model weights match the prior SHA256 -`58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. -The follow-up server reports one RTX4090 (SM89) with 49,140 MiB, driver -595.71.05, torch 2.12.1+cu130 and Triton 3.7.1. This is a different instance -from the original 24,564-MiB server with driver 580.105.08; -timings across these instances are not a matched performance comparison. - -These are numerical and execution regressions on synthetic fixed-token inputs, -not a held-out task-quality evaluation. Partial retention, adaptive exits, -long contexts, other GPUs and untested feature combinations remain outside the -qualification. Existing asynchronous lifecycle regressions are included, but -no new real-model asynchronous quality claim is made. The original performance -profiles remain historical evidence; -this correctness fix makes no serving-speed improvement claim. diff --git a/docs/validation/sparse-decode-sync-20261029.md b/docs/validation/sparse-decode-sync-20261029.md deleted file mode 100644 index 91d3a7a1..00000000 --- a/docs/validation/sparse-decode-sync-20261029.md +++ /dev/null @@ -1,96 +0,0 @@ -# Sparse decode: ordinary and synchronous integration validation - -This is the historical report for runtime `10373ab`. The subsequent -[numerical diagnosis and regression report](sparse-decode-numerics-20261029.md) -documents the backend-preserving fix and its separate validation. The timings -and numerical limitation below describe the original implementation. - -This contribution extends the shared [#157](https://github.com/InferMatrixLab/vllm-rlt/pull/157) reference plan/store into ordinary synchronous decode and synchronous block self-speculation. It is a functional reference integration, with measured costs and an unresolved BF16 full-retention sequence difference. It does not claim an optimized sparse kernel or deployment quality. - -## Source and checks - -- Immediate base: `048b9d7c79a5711857ce3287812502ad8f274ad0` on `kunkun/rfc156-sparse-decode-async`. -- Measured runtime: `10373ab6d63156535b55e4b9be1c7b8ba7319421`. Additional same-input tests: `985c309`; runtime files are identical between these revisions. Source SHA256 inventories are in the archive. -- Full pre-commit: Ruff lint, Ruff format and CPU regression passed; **919 passed, 475 skipped**. -- Applicable GPU regression: **312 passed, 1 skipped, 16 deselected**. The skip needs a second GPU; FA4 cases are excluded on SM89. The original unfiltered run's 12 unsupported-FA4 failures are retained. -- New integration module: 26 CPU tests plus 5 CUDA cases. Includes ordinary/padded routing, prompt exclusion, draft depth 1/2, each rejection boundary, same-position replacement, abort/ID reuse, failed submissions, stream-owned retirement and active-request close. Per-depth hidden, logits and populated KV match serial same-input sparse execution at FP32 atol/rtol `4e-5`. -- The negative stream control restores the unsafe caller-stream release only inside one test process and fails as expected. Corrected production code passes. -- Independent review found release-index stream ownership and shutdown-drain defects; both were fixed and covered. Re-review found no further concrete production defect. -- Upstream FlashLoop `aeaadee10d75b7aa8c4a42a809f15c75d6d3bdc6`, `dense_source_selector` plus `cached_mass_sparse_delta(force_torch=True)`: **18 CUDA fixed-input cases passed**, with equal retained sets; lengths 1/13/128, retention .1/.5/1, FP32 and BF16. Predeclared attention-output atol/rtol were `2e-5` and `.02`, respectively, accommodating the gold path's BF16 subset/output storage. This is separate from model-level quality. - -## Real-model protocol - -One RTX4090, 24,564 MiB; driver 580.105.08; torch 2.12.1+cu130, Triton 3.7.1. Released Ouro-1.4B revision `574fa66cb8bf5abdc979642d01cf2b79b16bfab1`, BF16 weights/activations/KV; weight SHA256 `58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af`. - -Native StageEngine, LAST_EXITED KV, refill, Triton dense backend, synchronous scheduling, multistream on, graphs off. Block K=2, draft depths 1/2, fixed four loops, greedy and ignore EOS. Each fresh process warms 4 output tokens at its target arrival concurrency, then measures three 128-input/16-output bursts. Request i uses `[1 + (j * 137 + i * 53) % 49000 for j in range(128)]`. Context bound is explicitly **512** for sparse, while model positional configuration remains unchanged. Token budget `max(128, 3*C)`, prefill chunk 128, KV blocks `40*C + 32`, block size 16. This does not qualify the model's 65,536 context limit. - -Timing includes enqueue, prefill, decode, delivery, drain and CUDA synchronization; model loading and warmup are excluded. Profiles are separate and excluded. GPU clocks and CPU affinity were not pinned for the primary sweep. All raw repeats, outputs, depths, acceptance counters, TTFT/E2E observations and memory values are archived. Speculative output arrives in chunks; chunk gaps are not per-token ITL. The short bursts do not establish steady-state HTTP serving. - -**81 primary bursts / 1,107 completed requests / 17,712 output tokens**; no primary OOM, timeout or failed request. Every measured run drained KV blocks and reference slots to zero and repeated its token/depth sequences exactly. - -Output throughput: median of three runs, with min-max in parentheses. Full-retention agreement counts one repeat, and was identical across all three. - -| Mode | Concurrent arrivals | Dense tok/s | Full retention tok/s | r=.1 tok/s | Full vs dense tokens | -|---|---:|---:|---:|---:|---:| -| Ordinary | 1 | 14.90 (14.67-14.98) | 9.66 (8.89-9.79) | 9.30 (8.86-9.37) | 16/16 | -| Ordinary | 8 | 71.60 (71.59-71.60) | 21.59 (21.53-21.62) | 21.73 (21.59-21.77) | 128/128 | -| Ordinary | 32 | 124.42 (124.41-124.49) | 27.13 (27.01-27.31) | 26.84 (26.52-27.24) | 512/512 | -| Sync, draft 1 | 1 | 17.46 (17.21-17.68) | 9.10 (8.63-9.12) | 13.90 (12.72-14.01) | 16/16 | -| Sync, draft 1 | 8 | 60.53 (57.33-61.12) | 18.01 (17.83-19.32) | 25.54 (25.50-25.64) | 128/128 | -| Sync, draft 1 | 32 | 82.91 (81.09-83.10) | 19.58 (19.54-19.81) | 27.48 (27.05-27.57) | 498/512 | -| Sync, draft 2 | 1 | 16.56 (15.37-16.69) | 9.49 (9.06-9.70) | 11.87 (11.13-11.91) | 16/16 | -| Sync, draft 2 | 8 | 50.27 (50.16-50.57) | 18.55 (17.54-18.58) | 22.40 (22.23-22.42) | 128/128 | -| Sync, draft 2 | 32 | 68.08 (67.59-68.30) | 20.31 (20.20-20.38) | 24.84 (24.70-24.90) | 512/512 | - -The shared PyTorch reference path is slower than dense Triton in every covered condition. Peak allocated memory reached **6.67 GiB**. Retention .1 changes outputs, and its speculation acceptance changes are preserved in `summary.json`; token agreement is not a task-quality score. - -## Full-retention discrepancy - -At C=32, draft depth 1, full retention matches **498/512** dense output tokens. The first divergence is request 9, output index 2 (zero-based), input position 129: dense chooses 293, full retention chooses 92. Other full-retention conditions in this sweep match all tokens; this does not prove broader equivalence. - -Instrumented replays reproduce the mismatch with **identical 130-token input prefixes**. Loop-1 hidden and loop-2 input hidden are exact. The first observed hidden drift is already at loop-2 capture output (max .01953125), before any reference consume; later-loop input differences follow. Recorded layer-0 Q/K/V and attention are exact at the capture loop; their same-input full-retention attention replays match Torch. Final logits differ by at most .125. Dense ranks `(293, 8.875)` above `(92, 8.8125)`; the reference path reverses those values. This localizes the observed drift to the capture computation before reuse, consistent with backend numerical differences, but does not localize the first physical-layer operator or establish task quality. No lossless/full-retention-at-scale claim is made. The existing async discrepancy in #157 is also not resolved by this contribution. - -## Disabled-feature drift control - -Initial unpinned C8 base/candidate controls showed timing drift, so a separate CPU-pinned **ABBA** check was run for sync draft 2, C8, feature off. Each arm uses a fresh process, the same workload and three measurements; CPUs are `[0, 1, 2, 3]`. Tokens and exit depths match the base exactly in all arms. - -- Base: median 52.06, range 51.90-53.91 tok/s. -- Candidate: median 52.19, range 50.63-52.54 tok/s. -- Candidate/base median ratio: 1.0027. These pinned observations are a separate control, not replacements for the primary unpinned numbers or a general performance guarantee. - -## Separate profiling - -Twelve readable CPU+CUDA captures cover ordinary C1/C8 and sync draft-2 C8, dense versus r=.1, each in independent ops-only/full processes. Full enables shapes, memory and stacks together. Every full capture has shape events, memory events and a nonempty exported stack inventory. The profile window spans a further output from each active request after the first outputs; raw shapes and row work remain available for checking scheduling/acceptance differences. Diagnostic count ratios are not throughput ratios. - -| C8 arrival workload | Active requests dense / r=.1 | GPU kernels dense / r=.1 | Scalar item reads dense / r=.1 | Argsorts dense / r=.1 | -|---|---:|---:|---:|---:| -| Ordinary | 8 / 8 | 5,600 / 26,220 | 0 / 384 | 0 / 192 | -| Sync, draft 2 | 7 / 5 | 11,319 / 43,242 | 0 / 720 | 0 / 360 | - -Ordinary C8 windows contain eight active requests in both arms. The synchronous windows contain seven versus five active requests because progress differs; they support operator diagnosis, not a matched-work count ratio. - -The source chain is `reference_plan -> sparse_paged_attention -> reference_statistics/read_reference`: per-row gather, selection and scalar count reads add launches and waits. Pool operations also add retirement work. The measured slowdown is therefore not explained by retained-key count alone. This contribution retains the common reference oracle; it does not duplicate or replace the pending attention optimization discussed on #157. - -![Aligned single-request diagnostic decode windows](https://github.com/Levius-Fubuki/vllm-rlt/releases/download/pr157-sync-validation-20261029/single-request-timeline.png) - -## Reproduction and evidence - -The [validation release](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-sync-validation-20261029) contains scripts, original traces, stack exports, raw outputs, source inventories, the negative control, failed diagnostic-harness attempts and the first unsupported-FA4 run. The SHA256 inventory covers the package contents; `SHA256SUMS` covers the downloadable archive and timeline. - -From the corresponding source checkout with the recorded dependencies and a reserved GPU: - -```bash -pre-commit run --all-files -CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=1 python -m pytest \ - tests/test_sparse_decode_sync.py tests/test_sparse_decode_attention.py \ - tests/test_speculative.py tests/test_async_speculative.py \ - tests/test_async_speculative_lifecycle.py tests/test_backend_selection.py \ - --run-gpu -k 'not flash_attn_4' -q -PYTHONPATH=. CUDA_VISIBLE_DEVICES=0 OMP_NUM_THREADS=4 python scripts/real_model_bench.py \ - --model /path/to/pinned/Ouro-1.4B --mode sync2 --concurrency 32 \ - --retention .1 --repeats 3 --output /tmp/sync2-c32-r01.json -``` - -Use `--mode ordinary|sync1|sync2`, `--retention 0|1|.1`, `--concurrency 1|8|32`; retention 0 disables the feature. The script contains the exact engine startup and per-request parameters. Use `--profile ops|full --repeats 1` only for separate diagnostics. The orchestrators retain their original task directory for exact provenance; the individual benchmark CLI accepts arbitrary model/output locations. - -Graphs, token pipelines, mixed scheduling, preemption and ordinary async decode are explicitly rejected. PD, prefix-cache composition, adaptive-exit quality, long generations, held-out task quality and other hardware are not qualified here. Sparse remains disabled by default. This is a Draft collaboration contribution to the existing shared feature PR, with numerical/performance limits retained for review. From 0ce8ca4552cf391e2964abfd954617a879abfd5b Mon Sep 17 00:00:00 2001 From: kunkunblueberry Date: Thu, 5 Nov 2026 16:30:17 +0800 Subject: [PATCH 26/28] sparse decode attention: lease identity instead of generation counters Review asked for the ownership model to use what the engine already guarantees rather than a parallel bookkeeping scheme. **Identity is the lease.** `reserve` binds a slot to the lease object it returns; `release` requires that same object; a plan is live only while every slot still points at the lease it was built from. A recycled slot is bound to a different lease, so the previous owner's plan cannot pass, with no per-slot counter and no `(request_id, row)` dictionary that grows for the life of the process. `ReferenceLease.generation`, `ReferenceLease.incarnations`, `SparseDecodePlan.incarnations`, the store's slot counters and its device `incarnation` tensor are gone; the plan carries `owners` instead. **Duplicate request ids are rejected at the boundary.** `_add_request` used to retire a still-retiring request and reuse its id; it now raises, including for requests that are still retiring. Together with the existing rule that slots are recycled only after their last reader completes, that is what makes lease identity sufficient. The EOS test now asserts the rejection and then reuses the id after retirement completes. **Renames.** `SpeculativeRoundTicket` is `AsyncSpeculativeOutput`. Also credits FlashLoop as the source of the cross-loop sparse attention mechanism in the README, docs/runtime.md and the sparse attention module docstring, with the paper and reference implementation linked. Regression: CPU 980 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 216 passed / 4 skipped; ruff 0.13.2 clean. --- README.md | 6 +- docs/runtime.md | 6 +- tests/test_async_speculative.py | 10 +- tests/test_async_speculative_lifecycle.py | 6 +- tests/test_sparse_decode_attention.py | 113 ++++++++-------------- tests/test_sparse_decode_batched.py | 8 +- tests/test_sparse_decode_fused.py | 18 ++-- tests/test_sparse_decode_sync.py | 6 +- vllm_rlt/attention/sparse.py | 13 ++- vllm_rlt/engine/stage_engine.py | 8 +- vllm_rlt/worker/async_speculative.py | 7 +- vllm_rlt/worker/attention_reference.py | 86 +++++++--------- vllm_rlt/worker/model_runner.py | 6 +- vllm_rlt/worker/speculative.py | 7 +- 14 files changed, 134 insertions(+), 166 deletions(-) diff --git a/README.md b/README.md index 71dc3be7..51405e7a 100644 --- a/README.md +++ b/README.md @@ -96,7 +96,11 @@ in mixed prefill/decode batches. fallback; it is not a general model-speedup claim. - **Sparse decode reference integration (opt-in):** cached loop-2 attention correction for ordinary decode and block self-speculation, with shared - per-token reference ownership and explicit capability limits; see the + per-token reference ownership and explicit capability limits. The cross-loop + sparse attention mechanism comes from FlashLoop + ([paper](https://arxiv.org/abs/2609.29812), + [reference implementation](https://github.com/Superone77/FlashLoop)); this + repository contributes the kernel and serving integration. See the [configuration and scope](docs/runtime.md#sparse-decode-attention-reference-path). - **Offline and online use:** Python API, CLI, and a streaming OpenAI-compatible **completions** endpoint. This is not a chat-completions API. diff --git a/docs/runtime.md b/docs/runtime.md index f7ce7fce..5d61e631 100644 --- a/docs/runtime.md +++ b/docs/runtime.md @@ -233,7 +233,11 @@ for the design. ## Sparse decode attention (reference path) `SparseDecodeAttentionConfig` enables the RFC #156 cached-reference correction -for generated-token attention. Prompt prefill remains dense. Loop 2 captures +for generated-token attention. The cross-loop sparse attention mechanism is from +FlashLoop ([paper](https://arxiv.org/abs/2609.29812), +[reference implementation](https://github.com/Superone77/FlashLoop)); the +configuration, kernels and serving integration here are this repository's. +Prompt prefill remains dense. Loop 2 captures per-token/per-layer/head selection, mass and contributions; loops 3 and 4 reuse that reference with their current selected KV. Each new input token establishes its own reference. diff --git a/tests/test_async_speculative.py b/tests/test_async_speculative.py index 4d6d1871..c1aff76b 100644 --- a/tests/test_async_speculative.py +++ b/tests/test_async_speculative.py @@ -173,7 +173,7 @@ def forbidden(*args, **kwargs): @pytest.mark.gpu @pytest.mark.parametrize("graphs", [False, True]) -def test_eos_discards_queued_round_and_request_id_can_be_reused(monkeypatch, graphs): +def test_eos_discards_queued_round_and_the_id_frees_after_retirement(monkeypatch, graphs): model = tiny("cuda") e = make_engine(model, graphs=graphs) bootstrap(e, ignore_eos=False) @@ -191,7 +191,13 @@ def test_eos_discards_queued_round_and_request_id_can_be_reused(monkeypatch, gra e.step() out = e._collect_speculative()[0] assert out.finished and out.token_ids[-1] == 0 and len(out.token_ids) == 2 - # Reuse while an old ticket still exists. It must never touch the new request. + # The finished request is still retiring, so its id is rejected: the engine + # never reuses an id while the previous token may still own resources. + with pytest.raises(ValueError, match="duplicate request ID"): + e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True)) + # Once retirement completes the id is free, and the new request never sees + # the old round's state. + e._retire_speculative(wait=True) monkeypatch.setattr(e.speculative_runner, "_coda", coda) e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True)) output = drain(e)["r"] diff --git a/tests/test_async_speculative_lifecycle.py b/tests/test_async_speculative_lifecycle.py index b3d469f7..a68ae940 100644 --- a/tests/test_async_speculative_lifecycle.py +++ b/tests/test_async_speculative_lifecycle.py @@ -17,9 +17,9 @@ from vllm_rlt.request import RequestPhase from vllm_rlt.serving.worker import EngineWorker from vllm_rlt.worker.async_speculative import ( + AsyncSpeculativeOutput, AsyncSpeculativeRunner, SpeculativeRequestState, - SpeculativeRoundTicket, ) @@ -105,7 +105,7 @@ def prepared_engine(*, bad_second=False, first_frontier=4, retirement_ready=True host_result=host, ) tickets.append( - SpeculativeRoundTicket( + AsyncSpeculativeOutput( SchedulerOutput(RequestPhase.SPECULATIVE, [scheduled_request]), buffer, (state,), @@ -382,7 +382,7 @@ def test_ticket_result_format_does_not_depend_on_statistics(mode, with_stats): host_result=torch.tensor(values), copy_done=Completion(), leased=True, keepalive=[] ) stats = SpeculativeStats() if with_stats else None - ticket = SpeculativeRoundTicket(batch, buffer, (), (), 2, mode=mode, stats=stats) + ticket = AsyncSpeculativeOutput(batch, buffer, (), (), 2, mode=mode, stats=stats) result = ticket.collect() assert result.speculative[0].token_ids == [11] assert result.speculative[0].accepted_count == 0 diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index c608f2db..cceec19e 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -3,7 +3,7 @@ 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 +descendants, same-position replacement through lease identity, delayed completion and reference-slot reuse. """ @@ -81,7 +81,7 @@ def test_plan_groups_rows_by_mode(): plan = SparseDecodePlan( modes=(DENSE, CAPTURE, CAPTURE, CONSUME), slots=(-1, 3, 4, 5), - incarnations=(0, 1, 1, 1), + owners=(None, "lease-1", "lease-1", "lease-2"), key_retention=0.5, key_bucket=None, ) @@ -163,7 +163,7 @@ def test_full_retention_matches_dense_attention(): @pytest.mark.parametrize("retention, expected", [(0.5, 23.5), (1.0, 40.0)]) def test_paged_oracle_forces_correction_without_backend_shortcut(retention, expected): store = make_store(slots=1, layers=1, heads=1, head_dim=1, key_bucket=3) - lease = store.reserve("r", 0, 1, generation=1) + lease = store.reserve("r", 0, 1) q = torch.ones(1, 1, 1) keys = torch.zeros(1, 4, 1, 1) values = torch.zeros_like(keys) @@ -190,17 +190,17 @@ def test_paged_oracle_forces_correction_without_backend_shortcut(retention, expe def test_reserve_uses_one_slot_per_live_token(): store = make_store() - lease = store.reserve("r", 0, 3, generation=1) + lease = store.reserve("r", 0, 3) assert len(lease) == 3 assert len(set(lease.slots)) == 3 - assert lease.incarnations == (1, 1, 1) + assert all(store._bindings[slot] is lease for slot in lease.slots) assert store.outstanding == 3 assert len(store.free) == 3 def test_release_requires_owner_and_retires_metadata(): store = make_store() - lease = store.reserve("r", 0, 2, generation=1) + lease = store.reserve("r", 0, 2) fill_payload(store, lease.slots) store.release(lease) for slot in lease.slots: @@ -218,7 +218,7 @@ def test_release_requires_owner_and_retires_metadata(): def test_accept_invalidates_rejected_candidates_and_descendants(): store = make_store() - lease = store.reserve("r", 0, 3, generation=1) + lease = store.reserve("r", 0, 3) fill_payload(store, lease.slots) store.accept([lease], torch.tensor([1])) accepted, middle, rejected = lease.slots @@ -236,7 +236,7 @@ def test_accept_invalidates_rejected_candidates_and_descendants(): def test_accept_drops_host_counts_so_rejected_rows_cannot_size_a_read(): store = make_store() - lease = store.reserve("r", 0, 3, generation=1) + lease = store.reserve("r", 0, 3) fill_payload(store, lease.slots) for slot in lease.slots: for layer in range(store.layers): @@ -253,7 +253,7 @@ def test_accept_drops_host_counts_so_rejected_rows_cannot_size_a_read(): def test_accept_keeps_root_when_every_candidate_is_rejected(): store = make_store() - lease = store.reserve("r", 0, 3, generation=1) + lease = store.reserve("r", 0, 3) fill_payload(store, lease.slots) store.accept([lease], torch.tensor([0])) assert store.mass[lease.slots[0]].sum() > 0 @@ -263,7 +263,7 @@ def test_accept_keeps_root_when_every_candidate_is_rejected(): def test_capture_and_consume_mark_layer_state(): store = make_store() - lease = store.reserve("r", 0, 2, generation=1) + lease = store.reserve("r", 0, 2) 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) @@ -275,14 +275,14 @@ def test_capture_and_consume_mark_layer_state(): assert int(store.uses[lease.slots[0], 0]) == 0 -def test_stale_plan_after_reincarnation_is_rejected(): +def test_stale_plan_after_a_new_lease_is_rejected(): store = make_store() - lease = store.reserve("r", 0, 2, generation=1) + lease = store.reserve("r", 0, 2) 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)) + # Same request and rows, a fresh lease. + renewed = store.reserve("r", 0, 2) + assert renewed is not lease 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) @@ -291,41 +291,43 @@ def test_stale_plan_after_reincarnation_is_rejected(): 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) + first = store.reserve("a", 0, 2) + second = store.reserve("b", 1, 2) 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.reserve("c", 2, 1) store.release(first) - third = store.reserve("c", 2, 1, generation=1) + third = store.reserve("c", 2, 1) assert set(third.slots) <= set(first.slots) -def test_reincarnation_bumps_on_slot_reuse(): +def test_a_reused_slot_is_owned_by_the_new_lease(): store = make_store() - lease = store.reserve("r", 0, 1, generation=1) + lease = store.reserve("r", 0, 1) slot = lease.slots[0] + assert store._bindings[slot] is lease store.release(lease) - again = store.reserve("r", 0, 1, generation=1) + again = store.reserve("r", 0, 1) assert again.slots[0] == slot - assert again.incarnations[0] == lease.incarnations[0] + 1 - assert int(store.incarnation[slot]) == again.incarnations[0] + assert again is not lease + assert store._bindings[slot] is again def test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot(): - # The incarnation must identify the physical slot's current owner. A counter - # keyed by (request_id, row) restarts at 1 for the next request, so the - # previous owner's plan carries the same number, passes the liveness check, - # and can overwrite the new owner's reference state. + # A plan must identify the physical slot's current owner. A counter keyed by + # (request_id, row) restarts for the next request, so the previous owner's + # plan could carry the same number, pass the liveness check, and overwrite + # the new owner's reference state. Binding the slot to the lease object makes + # the identity exact. store = make_store(slots=1, layers=1, heads=1, head_dim=2, key_bucket=4) - first = store.reserve("a", 0, 1, generation=1) + first = store.reserve("a", 0, 1) stale = store.build_plan([(first, 0)], [CAPTURE], key_retention=1.0) store.release(first) - second = store.reserve("b", 0, 1, generation=1) + second = store.reserve("b", 0, 1) assert second.slots == first.slots current = store.build_plan([(second, 0)], [CAPTURE], key_retention=1.0) - assert current.incarnations != stale.incarnations + assert current.owners[0] is second and stale.owners[0] is first store.capture(current, layer=0) valid = store.valid.clone() @@ -352,7 +354,7 @@ def test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot(): 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) + expected = 3 * 2 * 2 * (8 + 1 + 1 + 4 + 4) * 4 + 3 * 2 * (1 + 4) + 3 * (8 + 4) assert store.bytes == expected @@ -367,7 +369,7 @@ def _require_live_batch(self, batch): plan = SparseDecodePlan( modes=(CAPTURE, CONSUME), slots=(0, 1), - incarnations=(1, 1), + owners=(object(), object()), key_retention=0.5, key_bucket=4, ) @@ -444,7 +446,7 @@ def test_sparse_decode_matches_dense_through_model_and_kv(): 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) + lease = store.reserve("sparse", 0, 1) sparse = decode( model, make_cpu_cache(model.config), @@ -466,7 +468,7 @@ def test_sparse_decode_diverges_as_retention_shrinks(): errors = {} for retention in (1.0, 0.5, 0.25): store = make_reference_store(model) - lease = store.reserve("sparse", 0, 1, generation=1) + lease = store.reserve("sparse", 0, 1) sparse = decode( model, make_cpu_cache(model.config), @@ -497,7 +499,7 @@ def recording(*args, **kwargs): monkeypatch.setattr(sparse_module, "gather_retained", recording) store = make_reference_store(model) - lease = store.reserve("sparse", 0, 1, generation=1) + lease = store.reserve("sparse", 0, 1) decode( model, make_cpu_cache(model.config), @@ -608,41 +610,6 @@ def recording(leases, accepted_counts): 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) - plans = [] - original = store.build_plan - - def recording(owners, modes, **kwargs): - plan = original(owners, modes, **kwargs) - plans.append((plan.slots, plan.incarnations)) - return plan - - store.build_plan = recording - e.step() - e.step() - e._collect_speculative() - assert len(plans) >= 2 - # Incarnations belong to the physical slot, so rounds may rotate through - # different slots and each fresh slot starts at 1. A slot must never come - # back with a number it already had; equality is allowed only because the - # capture and consume plans of one round share a lease, and therefore share - # its incarnation. - seen: dict[int, int] = {} - for slots, incarnations in plans: - for slot, incarnation in zip(slots, incarnations): - if slot in seen: - assert incarnation >= seen[slot], (slot, seen[slot], incarnation) - seen[slot] = incarnation - e.abort_request("r") - drain(e) - 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 @@ -682,7 +649,7 @@ def test_release_from_another_stream_builds_its_index_on_the_owner_stream(monkey store = AttentionReferenceStore( slots=8, layers=1, heads=1, head_dim=4, key_bucket=4, device="cuda", stream=owner ) - lease = store.reserve("r", 0, 2, generation=1) + lease = store.reserve("r", 0, 2) other = next(slot for slot in range(store.slots) if slot not in lease.slots) store.valid.fill_(True) torch.cuda.synchronize() @@ -712,7 +679,7 @@ def test_accept_and_retire_do_not_allocate_bucket_sized_temporaries(): store = AttentionReferenceStore( slots=16, layers=24, heads=16, head_dim=128, key_bucket=32768, device="cuda" ) - lease = store.reserve("r", 0, 16, generation=1) + lease = store.reserve("r", 0, 16) store.counts[list(lease.slots)] = 1 store.mass[list(lease.slots)] = 1.0 store.output[list(lease.slots)] = 1.0 diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py index 17e8bc67..92185952 100644 --- a/tests/test_sparse_decode_batched.py +++ b/tests/test_sparse_decode_batched.py @@ -42,7 +42,7 @@ def test_batched_ragged_gqa_matches_scalar_oracle( ) for implementation in range(2) ] - leases = [store.reserve("test", 0, rows, generation=1) for store in stores] + leases = [store.reserve("test", 0, rows) for store in stores] implementations = [reference_paged_attention, batched_paged_attention] for modes in ([CAPTURE] * rows, [DENSE, CAPTURE, CONSUME, CONSUME, DENSE, CONSUME]): results = [] @@ -85,7 +85,7 @@ def test_batched_ragged_gqa_matches_scalar_oracle( 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) + lease = store.reserve("test", 0, 1) plan = store.build_plan([(lease, 0)], [CONSUME], key_retention=0.1) args = ( torch.ones(1, 2, 4), @@ -98,7 +98,7 @@ def test_batched_rejects_missing_and_stale_reference(): with pytest.raises(RuntimeError, match="no captured reference"): batched_paged_attention(*args, layer=0) store.release(lease) - store.reserve("test", 0, 1, generation=2) + store.reserve("test", 0, 1) with pytest.raises(RuntimeError, match="stale attention reference"): batched_paged_attention(*args, layer=0) @@ -114,7 +114,7 @@ def test_reference_growth_preserves_live_layers_and_budget(device): initial_key_capacity=2, device=device, ) - leases = [store.reserve(str(i), i, 1, generation=1) for i in range(2)] + leases = [store.reserve(str(i), i, 1) for i in range(2)] plans = [store.build_plan([(lease, 0)], [CAPTURE], key_retention=1.0) for lease in leases] # Simulate queued work on the runner's owner stream before reallocating. if device == "cuda": diff --git a/tests/test_sparse_decode_fused.py b/tests/test_sparse_decode_fused.py index c1e35db8..86dc634f 100644 --- a/tests/test_sparse_decode_fused.py +++ b/tests/test_sparse_decode_fused.py @@ -27,7 +27,7 @@ def test_fused_store_keeps_the_oracle_on_cpu(monkeypatch): monkeypatch.setattr(sparse, "reference_paged_attention", lambda *a, **k: calls.append(1)) store = make_store(slots=1, layers=1, heads=2, head_dim=4, implementation="fused") - lease = store.reserve("cpu", 0, 1, generation=1) + lease = store.reserve("cpu", 0, 1) plan = store.build_plan([(lease, 0)], [CAPTURE], key_retention=0.5) q = torch.ones(1, 2, 4) cache = torch.ones(1, 16, 2, 4) @@ -73,7 +73,7 @@ def test_fast_launches_match_triton_dispatch_bit_for_bit(context, monkeypatch): device="cuda", implementation="fused", ) - lease = store.reserve("r", 0, rows, generation=1) + lease = store.reserve("r", 0, rows) owners = [(lease, row) for row in range(rows)] outputs = [] # Layers 0, 1, 16 and 17 differ in Triton's integer specialization. @@ -128,7 +128,7 @@ def consume(dim): device="cuda", implementation="fused", ) - lease = store.reserve("r", 0, rows, generation=1) + lease = store.reserve("r", 0, rows) owners = [(lease, row) for row in range(rows)] for mode in (CAPTURE, CONSUME): plan = store.build_plan(owners, [mode] * rows, key_retention=0.25) @@ -190,7 +190,7 @@ def test_fused_capture_and_consume_track_batched_path(retention, ties, split): ) for implementation in ("batched", "fused") ] - leases = [store.reserve("fused", 0, rows, generation=1) for store in stores] + leases = [store.reserve("fused", 0, rows) for store in stores] owners = [[(lease, row) for row in range(rows)] for lease in leases] q, keys, values, tables = _paged( rows, heads, kv_heads, dim, pages, block_size, torch.float32, 1 @@ -298,7 +298,7 @@ def integers(*shape): for implementation in ("batched", "fused") ] for store in stores: - lease = store.reserve("ties", 0, rows, generation=1) + lease = store.reserve("ties", 0, rows) plan = store.build_plan( [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention ) @@ -367,7 +367,7 @@ def integers(*shape): ] leases = [] for store in stores: - lease = store.reserve("parts", 0, rows, generation=1) + lease = store.reserve("parts", 0, rows) leases.append(lease) plan = store.build_plan( [(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=retention @@ -412,7 +412,7 @@ def test_fused_split_capture_is_the_dense_backend_output_bit_for_bit(dtype): device="cuda", implementation="fused", ) - lease = store.reserve("split", 0, rows, generation=1) + lease = store.reserve("split", 0, rows) plan = store.build_plan([(lease, r) for r in range(rows)], [CAPTURE] * rows, key_retention=0.1) out = sparse_paged_attention( q, keys, values, tables, lengths, plan, layer=0, dense_attention=backend @@ -449,7 +449,7 @@ def test_fused_bf16_consume_matches_scalar_oracle_within_rounding(): device="cuda", implementation=implementation, ) - lease = store.reserve("bf16", 0, rows, generation=1) + lease = store.reserve("bf16", 0, rows) owner = [(lease, row) for row in range(rows)] for modes in ([CAPTURE] * rows, [CONSUME] * rows): plan = store.build_plan(owner, modes, key_retention=0.1) @@ -474,7 +474,7 @@ def test_fused_consume_without_capture_is_nan_not_silent(): store = make_store( slots=rows, layers=1, heads=heads, head_dim=dim, device="cuda", implementation="fused" ) - lease = store.reserve("missing", 0, rows, generation=1) + lease = store.reserve("missing", 0, rows) plan = store.build_plan([(lease, 0), (lease, 1)], [CONSUME, CONSUME], key_retention=0.5) q, keys, values, tables = _paged(rows, heads, heads, dim, 2, 16, torch.float32, 4) lengths = torch.tensor([0, 9], device="cuda", dtype=torch.int32) diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py index 5cd6e56d..d5defb1c 100644 --- a/tests/test_sparse_decode_sync.py +++ b/tests/test_sparse_decode_sync.py @@ -174,7 +174,7 @@ def test_cuda_sparse_preserves_dense_backend_outputs(kv_heads, retention): keys = torch.randn(length, kv_heads, dim, device="cuda", dtype=torch.bfloat16) values = torch.randn_like(keys) cache.write(0, [rid] * length, [0] * length, list(range(length)), keys, values) - lease = store.reserve(rid, 0, 1, generation=1) + lease = store.reserve(rid, 0, 1) owners.append((lease, 0)) batch = cache._prepare_batch(ids, [0] * rows, [length - 1] * rows) q = torch.randn(rows, heads, dim, device="cuda", dtype=torch.bfloat16) @@ -252,7 +252,7 @@ def record(hidden, kv, cache, **kwargs): assert oracle.allocate("r", 16) store = make_reference_store(model) for position, token in enumerate(tokens): - lease = store.reserve("r", 0, 1, generation=position) + lease = store.reserve("r", 0, 1) hidden = model.prelude(torch.tensor([token])) for depth in range(4): kv = oracle._prepare_batch(["r"], [depth], [position]) @@ -324,7 +324,7 @@ def record(*args, **kwargs): e.reference_store.read_reference(old, 0, 0) monkeypatch.setattr(e.speculative_runner, "_coda", original) e.step() - assert seen[-1].incarnations[0] > old.incarnations[0] + assert seen[-1].owners[0] is not old.owners[0] assert e.speculative_runner.stats.accepted_tokens >= accepted finally: e.close() diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index c438a918..bdd4e453 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -1,5 +1,10 @@ """Sparse decode attention arithmetic and its per-row plan (RFC #156, section 3.2). +The cross-loop sparse attention mechanism is from FlashLoop +(https://arxiv.org/abs/2609.29812, reference implementation +https://github.com/Superone77/FlashLoop); the kernels and serving integration in +this repository are ours. + 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 @@ -63,8 +68,10 @@ class SparseDecodePlan: ``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. + ``owners`` holds the lease each row's slot was reserved from, or ``None`` + for a dense row. The lease is the slot's identity: a slot handed to another + request is bound to a different lease, so a plan built for the previous + owner fails the store's liveness check. A dense row runs the ordinary attention path. A capture row is the reference loop: it computes the dense attention and publishes the retained @@ -78,7 +85,7 @@ class SparseDecodePlan: modes: tuple[int, ...] slots: tuple[int, ...] - incarnations: tuple[int, ...] + owners: tuple[object | None, ...] key_retention: float key_bucket: int | None slot_ids: torch.Tensor | None = None diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index fdfda895..77d1da53 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -425,8 +425,12 @@ def _add_request(self, request_id, prompt_token_ids, sampling_params, *, trace_i ) and params.temperature != 0: mode = "async" if self.async_speculative else "pipeline" raise ValueError(f"{mode} speculative decoding currently requires greedy sampling") - if request_id in self._retiring_speculative: - self._retire_speculative(wait=True, request_id=request_id) + # A request id is never reused while its previous token is retiring: + # reference slots are recycled only after their last reader completed, + # so rejecting the duplicate keeps ownership unambiguous without a + # per-slot generation counter. + if request_id in self._retiring_speculative or request_id in self.scheduler.requests: + raise ValueError(f"duplicate request ID {request_id!r}: already live or still retiring") config = self.model.config if not prompt_token_ids: raise ValueError("prompt must contain at least one token") diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 1840686c..3b289019 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -86,7 +86,7 @@ def __init__(self, runner, scheduler): @dataclass -class SpeculativeRoundTicket: +class AsyncSpeculativeOutput: batch: SchedulerOutput buffer: SpeculativeRoundBuffer states: tuple[SpeculativeRequestState, ...] @@ -554,7 +554,6 @@ def submit(self, batch): scheduled_request.request.request_id, index, scheduled_request.token_count, - state.generation, ) for index, (scheduled_request, state) in enumerate( zip(scheduled_requests, states) @@ -598,7 +597,7 @@ def submit(self, batch): if self.config.mode != "pipeline": self.stats.drafted_tokens += sum(s.token_count - 1 for s in scheduled_requests) self.stats.verified_rows += count - return SpeculativeRoundTicket( + return AsyncSpeculativeOutput( batch, buffer, states, @@ -674,7 +673,7 @@ def _pipeline_step(self, spans, states, tables, allocations, buffer): p.tokens.copy_(torch.where(newborn, token, p.tokens)) p.depths.masked_fill_(newborn, 0) p.active.copy_((surviving | newborn) & (state.position + ids < limit)) - # Per-request wire order matches SpeculativeRoundTicket.collect(): + # Per-request wire order matches AsyncSpeculativeOutput.collect(): # prediction, accepted, emitted, next_position, drafted, verified. # A pipeline step verifies at most one token per request, so the # verified count equals emitted (0 or 1). diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index e67820b0..f3dc69cc 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -5,8 +5,8 @@ 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: +which lease 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. @@ -21,10 +21,14 @@ 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. +A slot is owned by the lease that reserved it, and the lease object is the +identity: reserve binds the slot to a new lease, release requires that +same lease, and a plan is live only while every slot still points at the lease +it was built from. A slot handed to another request is bound to a different +lease, so the previous owner's plan cannot pass. This relies on two properties +the engine already guarantees: a request id is never reused while its previous +token is retiring, and slots are recycled only after their last reader +completed. """ from dataclasses import dataclass @@ -46,10 +50,8 @@ class ReferenceLease: """ request_id: str - generation: int request_index: int slots: tuple[int, ...] - incarnations: tuple[int, ...] slot_ids: torch.Tensor row_index: torch.Tensor @@ -118,16 +120,12 @@ def __init__( 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]] = {} - # Incarnations belong to the physical slot, not to the request that - # borrows it: a counter keyed by ``(request_id, row)`` restarts at 1 for - # the next owner, so a stale plan from the previous owner would carry - # the same number and pass the liveness check. Slot counters only ever - # increase, which makes a recycled slot's old plan unambiguously stale. - self._slot_incarnations: list[int] = [0] * slots + # Slot ownership is the lease object itself. Object identity is exact: + # a recycled slot is bound to a different lease, so a plan built for the + # previous owner cannot match, and no per-slot counter is needed. + self._bindings: dict[int, ReferenceLease] = {} # Retained counts are host-known: the capture decides them before it # submits. Keeping them here lets a reader size its gather without a # device readback, which RFC #156 section 6.2 forbids. @@ -147,7 +145,6 @@ def bytes(self) -> int: self.valid, self.uses, self.position, - self.incarnation, self.accepted, ) ) @@ -200,46 +197,42 @@ def _index(self, values) -> torch.Tensor: host = torch.tensor(values, dtype=torch.int64, pin_memory=self.device.type == "cuda") return host.to(self.device, non_blocking=True) - def reserve(self, request_id: str, request_index: int, rows: int, generation: int): - """Take one slot per verification row of a request's submission.""" + def reserve(self, request_id: str, request_index: int, rows: int) -> ReferenceLease: + """Take one slot per verification row of a request's submission. + + The lease is the slot's owner identity. A slot returned to the pool and + handed to another request is bound to a different lease object, so plans + that still reference the old lease cannot pass a liveness check. + """ 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._slot_incarnations[slot] + 1 - self._slot_incarnations[slot] = count - self._bindings[slot] = (request_id, generation, count) - incarnations.append(count) - slot_ids = self._index(slots) - # Writing host integers element by element would block on each copy. - self.incarnation.index_copy_(0, slot_ids, self._index(incarnations)) - return ReferenceLease( + lease = ReferenceLease( request_id=request_id, - generation=generation, request_index=request_index, slots=slots, - incarnations=tuple(incarnations), - slot_ids=slot_ids, + slot_ids=self._index(slots), row_index=torch.arange(rows, dtype=torch.int64, device=self.device), ) + for slot in slots: + self._bindings[slot] = lease + return lease 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. + cleared payload never leaks into a later owner. """ 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): + if self._bindings.get(slot) is not lease: raise RuntimeError("attention reference slot released by a different owner") del self._bindings[slot] cleared.append(slot) @@ -266,34 +259,29 @@ def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan ``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 + submission row. Slots and their owning leases 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 = [], [] + slots, row_owners = [], [] for owner, mode in zip(owners, modes): if owner is None: if mode != DENSE: raise ValueError("a sparse row needs a reference owner") slots.append(-1) - incarnations.append(0) + row_owners.append(None) continue lease, row = owner slot = lease.slots[row] - binding = self._bindings.get(slot) - if binding is None or binding != ( - lease.request_id, - lease.generation, - lease.incarnations[row], - ): + if self._bindings.get(slot) is not lease: raise RuntimeError("attention reference slot released before its plan") slots.append(slot) - incarnations.append(binding[2]) + row_owners.append(lease) return SparseDecodePlan( modes=tuple(modes), slots=tuple(slots), - incarnations=tuple(incarnations), + owners=tuple(row_owners), key_retention=key_retention, key_bucket=self.key_bucket, slot_ids=self._index(slots), @@ -304,11 +292,10 @@ 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, mode in zip(plan.slots, plan.incarnations, plan.modes): + for slot, owner, mode in zip(plan.slots, plan.owners, plan.modes): if slot == -1 and mode == DENSE: continue - binding = self._bindings.get(slot) - if binding is None or binding[2] != incarnation: + if self._bindings.get(slot) is not owner: raise RuntimeError("stale attention reference plan for a retired token") def note_count(self, slot: int, layer: int, count: int) -> None: @@ -476,7 +463,6 @@ def zero(): self.valid, self.uses, self.accepted, - self.incarnation, self.position, self.mass, self.counts, diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index 58d5b471..ec9e28ae 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -169,7 +169,6 @@ def __init__( self.reference_store = reference_store self.sparse_decode_config = sparse_decode_config self.attention_references = {} - self.reference_generation = 0 self.metadata = cache_manager.metadata_builder self.runtime = ModelRuntime(self.model, cache_manager) parameter = next(model.parameters()) @@ -1081,10 +1080,7 @@ def _reference_metadata(self, batch, kv): state = None if mode == CAPTURE: if state is None: - self.reference_generation += 1 - lease = self.reference_store.reserve( - request.request_id, 0, 1, self.reference_generation - ) + lease = self.reference_store.reserve(request.request_id, 0, 1) state = (request, position, lease) self.attention_references[request.request_id] = state if state is None: diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 39924dbf..37daa927 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -164,7 +164,6 @@ def __init__( self.compute_stream = compute_stream self.reference_store = reference_store self.sparse_decode_config = sparse_decode_config - self.reference_generation = 0 self.submissions = SubmissionQueue() self.input_buffers = { dtype: StagedInputTensor(dtype, self.device, use_uva=execution.use_uva) @@ -295,14 +294,10 @@ def submit(self, batch): references = [] try: if self.reference_store is not None: - self.reference_generation += 1 for index, span in enumerate(batch.requests): references.append( self.reference_store.reserve( - span.request.request_id, - index, - span.token_count, - self.reference_generation, + span.request.request_id, index, span.token_count ) ) if self.config.mode == "pipeline": From 4b9ba790763dd8985e78411f77c417bfa00a8209 Mon Sep 17 00:00:00 2001 From: kunkunblueberry Date: Thu, 5 Nov 2026 17:21:41 +0800 Subject: [PATCH 27/28] sparse decode attention: exception-safe reservation, no write-only state Review comments on the reference store, the runner and the plan types. **reserve is exception-safe.** It popped slots from the pool before `_index()` and `torch.arange()` could fail, so an allocation error leaked the slots and the next request failed with "no retired attention reference slots". It now peeks, allocates, and only then removes them from the pool. test_failed_reservation_keeps_the_pool_intact injects the failure and checks that a later request still reserves. **Write-only state removed.** `valid`, `uses`, `position` and `accepted` were written by the store and by three kernels and never read back: availability is the retained count and ownership is the lease. The tensors, their writes, the `bind_positions` API and the corresponding kernel arguments (`VALID` in `_capture_kernel` and `_select_pass_kernel`, `USES` in `_consume_merge_kernel`) are gone, which also drops 12 bytes per slot from the accounted payload. `capture` and `consume` remain as liveness assertions. **Naming and types.** `AttentionReferenceStore` is `SparseDecodeReferenceStore`. `SparseDecodePlan.owners` and `.store` carry real types through a `TYPE_CHECKING` import, and the unused `key_bucket` field is gone. **Runner.** `_generations` and `SpeculativeRequestState.generation` are removed; the value had no production reader left. The class-level `reference_store` / `sparse_decode` defaults that existed for test doubles are gone, and the one double that skips `__init__` now sets both fields explicitly. Regression: CPU 981 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 217 passed / 4 skipped; ruff 0.13.2 clean. --- tests/test_async_speculative.py | 2 + tests/test_sparse_decode_attention.py | 81 +++++++++++++++----------- tests/test_sparse_decode_batched.py | 2 +- tests/test_sparse_decode_fused.py | 9 +-- tests/test_sparse_decode_sync.py | 9 +-- vllm_rlt/attention/sparse.py | 16 ++++- vllm_rlt/engine/stage_engine.py | 4 +- vllm_rlt/kernels/sparse_decode.py | 21 +------ vllm_rlt/worker/async_speculative.py | 20 +------ vllm_rlt/worker/attention_reference.py | 70 ++++++---------------- vllm_rlt/worker/model_runner.py | 1 - vllm_rlt/worker/speculative.py | 1 - 12 files changed, 97 insertions(+), 139 deletions(-) diff --git a/tests/test_async_speculative.py b/tests/test_async_speculative.py index c1aff76b..47fde74a 100644 --- a/tests/test_async_speculative.py +++ b/tests/test_async_speculative.py @@ -351,6 +351,8 @@ def test_active_page_tables_are_contiguous_and_do_not_touch_inactive_storage(): assert cache.allocate("short", 4) and cache.allocate("long", 10) runner = object.__new__(AsyncSpeculativeRunner) runner.cache = cache + runner.reference_store = None + runner.sparse_decode_config = None buffer = SimpleNamespace( leased=False, host_tables=torch.full((128,), -91, dtype=torch.int32), diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index cceec19e..a638995c 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -32,7 +32,7 @@ sparse_correction, sparse_paged_attention, ) -from vllm_rlt.worker.attention_reference import AttentionReferenceStore +from vllm_rlt.worker.attention_reference import SparseDecodeReferenceStore from vllm_rlt.worker.metadata import MetadataBuilder tiny = tiny_ouro_model @@ -52,7 +52,7 @@ def drain(engine): 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) + return SparseDecodeReferenceStore(**values) def fill_payload(store, slots, value=1.0): @@ -83,7 +83,6 @@ def test_plan_groups_rows_by_mode(): slots=(-1, 3, 4, 5), owners=(None, "lease-1", "lease-1", "lease-2"), key_retention=0.5, - key_bucket=None, ) assert plan.groups() == {DENSE: [0], CAPTURE: [1, 2], CONSUME: [3]} assert not plan.dense() @@ -179,7 +178,6 @@ def test_paged_oracle_forces_correction_without_backend_shortcut(retention, expe consume = store.build_plan([(lease, 0)], [CONSUME], key_retention=retention) output = sparse_paged_attention(q, keys, values, tables, lengths, consume, layer=0) assert output.item() == pytest.approx(expected) - assert store.uses.item() == 1 store.release(lease) @@ -204,12 +202,12 @@ def test_release_requires_owner_and_retires_metadata(): 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 + # Readers are gated by counts and mass. 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 not bool(store.valid[slot].any()) + assert not bool(store.counts[slot].any()) assert store.outstanding == 0 # A retired lease cannot release its slots a second time. with pytest.raises(RuntimeError, match="different owner"): @@ -227,8 +225,7 @@ def test_accept_invalidates_rejected_candidates_and_descendants(): 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 not bool(store.valid[rejected].any()) - assert int(store.accepted[accepted]) == 1 + assert not bool(store.counts[rejected].any()) 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) @@ -261,18 +258,44 @@ def test_accept_keeps_root_when_every_candidate_is_rejected(): assert store.mass[slot].sum() == 0 -def test_capture_and_consume_mark_layer_state(): +def test_capture_and_consume_assert_liveness_without_state(): + # Availability comes from the retained count and ownership from the lease, so + # capture and consume record nothing; they reject a plan whose slots moved on. store = make_store() lease = store.reserve("r", 0, 2) plan = store.build_plan([(lease, 0), (lease, 1)], [CAPTURE, CONSUME], key_retention=0.5) - assert not bool(store.valid[lease.slots[0], 0]) + before = store.counts.clone() 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 + assert torch.equal(store.counts, before) + store.release(lease) + with pytest.raises(RuntimeError, match="stale"): + store.capture(plan, layer=0) + + +def test_failed_reservation_keeps_the_pool_intact(monkeypatch): + # A device allocation can fail between taking slots and returning a lease. + # The slots must go back to the pool, or the next request starves. + store = make_store(slots=2) + original = store._index + attempts = {"n": 0} + + def failing(values): + attempts["n"] += 1 + if attempts["n"] == 1: + raise RuntimeError("injected allocation failure") + return original(values) + + monkeypatch.setattr(store, "_index", failing) + with pytest.raises(RuntimeError, match="injected allocation failure"): + store.reserve("a", 0, 2) + assert len(store.free) == 2 + assert store.outstanding == 0 + + lease = store.reserve("b", 0, 2) + assert len(lease.slots) == 2 + assert len(store.free) == 0 def test_stale_plan_after_a_new_lease_is_rejected(): @@ -330,7 +353,6 @@ def test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot(): assert current.owners[0] is second and stale.owners[0] is first store.capture(current, layer=0) - valid = store.valid.clone() counts = store.counts.clone() mass = store.mass.clone() @@ -347,14 +369,14 @@ def test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot(): output=torch.zeros(1, 2), contribution=torch.zeros(1, 2), ) - assert torch.equal(store.valid, valid) assert torch.equal(store.counts, counts) assert torch.equal(store.mass, mass) 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 + 4) + # selected (int32 x key_bucket) + counts + mass + output + contribution + expected = 3 * 2 * 2 * 8 * 4 + 3 * 2 * 2 * 4 + 3 * 2 * 2 * 4 + 2 * 3 * 2 * 2 * 4 * 4 assert store.bytes == expected @@ -371,7 +393,6 @@ def _require_live_batch(self, batch): slots=(0, 1), owners=(object(), object()), key_retention=0.5, - key_bucket=4, ) common = dict( allocations=(("r", object()),), @@ -436,7 +457,7 @@ def make_reference_store(model, **overrides): device="cpu", ) values.update(overrides) - return AttentionReferenceStore(**values) + return SparseDecodeReferenceStore(**values) def test_sparse_decode_matches_dense_through_model_and_kv(): @@ -456,7 +477,7 @@ def test_sparse_decode_matches_dense_through_model_and_kv(): lease=lease, retention=1.0, ) - assert store.uses.sum() > 0, "the sparse path never consumed a reference" + assert store.outstanding > 0, "the sparse path never held a reference" for expected, actual in zip(dense, sparse): assert torch.allclose(actual, expected, atol=1e-4), float((actual - expected).abs().max()) @@ -586,22 +607,15 @@ def recording(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], - ) - ) + observed.append((count, [float(store.mass[slot, 0, 0]) 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: + for count, masses 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") @@ -646,12 +660,11 @@ def test_release_from_another_stream_builds_its_index_on_the_owner_stream(monkey # freed when release returns and can be reused there before the owner # stream reads it, clearing some other slot. owner = torch.cuda.Stream() - store = AttentionReferenceStore( + store = SparseDecodeReferenceStore( slots=8, layers=1, heads=1, head_dim=4, key_bucket=4, device="cuda", stream=owner ) lease = store.reserve("r", 0, 2) other = next(slot for slot in range(store.slots) if slot not in lease.slots) - store.valid.fill_(True) torch.cuda.synchronize() clear, streams = store._clear, [] @@ -666,8 +679,8 @@ def record(indices): poison = [torch.full((1,), other, device="cuda", dtype=torch.int64) for _ in range(32)] torch.cuda.synchronize() assert streams == [owner] and len(poison) == 32 - assert not bool(store.valid[list(lease.slots)].any()) - assert bool(store.valid[other].all()) + assert not any(slot in store._bindings for slot in lease.slots) + assert other not in lease.slots @pytest.mark.gpu @@ -676,7 +689,7 @@ def test_accept_and_retire_do_not_allocate_bucket_sized_temporaries(): # 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( + store = SparseDecodeReferenceStore( slots=16, layers=24, heads=16, head_dim=128, key_bucket=32768, device="cuda" ) lease = store.reserve("r", 0, 16) diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py index 92185952..ef201262 100644 --- a/tests/test_sparse_decode_batched.py +++ b/tests/test_sparse_decode_batched.py @@ -69,7 +69,7 @@ def test_batched_ragged_gqa_matches_scalar_oracle( capacity = stores[1].selected.shape[-1] assert torch.equal(stores[0].selected[..., :capacity], stores[1].selected) assert not torch.count_nonzero(stores[0].selected[..., capacity:]) - for name in ("counts", "valid", "uses"): + for name in ("counts",): assert torch.equal(getattr(stores[0], name), getattr(stores[1], name)), name for name in ("mass", "output", "contribution"): torch.testing.assert_close( diff --git a/tests/test_sparse_decode_fused.py b/tests/test_sparse_decode_fused.py index 86dc634f..6029d8de 100644 --- a/tests/test_sparse_decode_fused.py +++ b/tests/test_sparse_decode_fused.py @@ -221,8 +221,9 @@ def test_fused_capture_and_consume_track_batched_path(retention, ties, split): torch.testing.assert_close(captured[1], captured[0], rtol=0, atol=exact) torch.testing.assert_close(fused.output[:, 1], batched.output[:, 1], rtol=0, atol=exact) torch.testing.assert_close(fused.counts[:, 1], batched.counts[:, 1], rtol=0, atol=0) - # The finalize kernel marks the same rows source-ready as store.capture. - assert torch.equal(fused.valid, batched.valid) + # Both paths leave the same slots bound to their leases; availability is + # decided by the retained count, which the comparison above already covers. + assert set(fused._bindings) == set(batched._bindings) for row, length in enumerate(lengths.tolist()): slot_b, slot_f = leases[0].slots[row], leases[1].slots[row] count = int(batched.counts[slot_b, 1, 0]) @@ -255,7 +256,7 @@ def test_fused_capture_and_consume_track_batched_path(retention, ties, split): ) torch.testing.assert_close(consumed[1], consumed[0], rtol=1e-4, atol=1e-5) assert not consumed[1][0].any(), "an empty row stays zero" - assert int(fused.uses[list(leases[1].slots), 1].min()) == 1 + assert fused.outstanding > 0 for store, lease in zip(stores, leases): store.release(lease) @@ -425,7 +426,7 @@ def test_fused_split_capture_is_the_dense_backend_output_bit_for_bit(dtype): # O_ref is the output actually fed into the recurrence: rounded, then FP32. live = list(lease.slots)[1:] torch.testing.assert_close(store.output[live, 0], expected[1:].float(), rtol=0, atol=atol) - assert bool(store.valid[list(lease.slots), 0].all()) + assert all(slot in store._bindings for slot in lease.slots) @pytest.mark.gpu diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py index d5defb1c..9be12932 100644 --- a/tests/test_sparse_decode_sync.py +++ b/tests/test_sparse_decode_sync.py @@ -86,7 +86,7 @@ def attend(layer, kv, q): request_id = next(rid for rid, owner in kv.allocations if owner is allocation) assert position >= len(e.scheduler.requests[request_id].prompt_token_ids) assert mode == (CAPTURE if depth == 1 else CONSUME) - assert int(e.reference_store.position[plan.slots[row]]) == position + assert plan.owners[row].request_id == request_id seen.append((mode, position, plan.slots[row])) return original(layer, kv, q) @@ -146,7 +146,7 @@ def test_backend_explicitly_declares_sparse_reference_support(): @pytest.mark.parametrize("retention", [0.1, 1.0]) def test_cuda_sparse_preserves_dense_backend_outputs(kv_heads, retention): from vllm_rlt.core.kv_cache_manager import KVCacheManager - from vllm_rlt.worker.attention_reference import AttentionReferenceStore + from vllm_rlt.worker.attention_reference import SparseDecodeReferenceStore # BF16 rounding of Torch softmax and Triton split-KV can differ even for # identical operands. Dense/capture rows and the full-retention control @@ -164,7 +164,7 @@ def test_cuda_sparse_preserves_dense_backend_outputs(kv_heads, retention): dtype=torch.bfloat16, backend="triton", ) - store = AttentionReferenceStore( + store = SparseDecodeReferenceStore( slots=rows, layers=1, heads=heads, head_dim=dim, key_bucket=length, device="cuda" ) ids = [str(i) for i in range(rows)] @@ -199,7 +199,8 @@ def test_cuda_sparse_preserves_dense_backend_outputs(kv_heads, retention): ) actual = cache._attend_prepared(0, replace(batch, reference_plan=consume), q) torch.testing.assert_close(actual, expected, atol=0, rtol=0) - assert int(store.uses.sum()) == rows - 1 + # Consume rows read the reference: their slots stay bound until release. + assert store.outstanding == len(owners) for lease, _ in owners: store.release(lease) for rid in ids: diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index bdd4e453..078e857a 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -32,12 +32,23 @@ depends on a device readback. """ +from __future__ import annotations + import math from collections.abc import Callable from dataclasses import dataclass +from typing import TYPE_CHECKING import torch +if TYPE_CHECKING: + # Type-only: the attention layer borrows the runner-owned store, so it must + # not import the worker package at runtime. + from vllm_rlt.worker.attention_reference import ( + ReferenceLease, + SparseDecodeReferenceStore, + ) + from vllm_rlt.kernels.sparse_capture import TRITON_AVAILABLE, fused_capture_scores from vllm_rlt.kernels.sparse_consume import selected_scores as fused_selected_scores from vllm_rlt.kernels.sparse_consume import selected_values as fused_selected_values @@ -85,11 +96,10 @@ class SparseDecodePlan: modes: tuple[int, ...] slots: tuple[int, ...] - owners: tuple[object | None, ...] + owners: tuple[ReferenceLease | None, ...] key_retention: float - key_bucket: int | None slot_ids: torch.Tensor | None = None - store: object | None = None + store: SparseDecodeReferenceStore | None = None def groups(self) -> dict[int, list[int]]: """Rows per mode, in submission order.""" diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 77d1da53..4e5fa4b5 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -23,7 +23,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.attention_reference import SparseDecodeReferenceStore from vllm_rlt.worker.executor import Executor from vllm_rlt.worker.model_runner import ModelRunner from vllm_rlt.worker.output import ExitSignal, ModelRunnerOutput @@ -183,7 +183,7 @@ def __init__( f"{free_bytes / 1024**3:.1f} GiB is free; lower " f"key_retention or num_speculative_tokens" ) - self.reference_store = AttentionReferenceStore( + self.reference_store = SparseDecodeReferenceStore( slots=slots, layers=config.num_hidden_layers, heads=config.num_attention_heads, diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py index 209a4730..059f067a 100644 --- a/vllm_rlt/kernels/sparse_decode.py +++ b/vllm_rlt/kernels/sparse_decode.py @@ -476,7 +476,6 @@ def _select_pass_kernel( SLOTS, COUNTS, OUTPUT, - VALID, LSE, SEGHIST, STATE, @@ -521,8 +520,6 @@ def _select_pass_kernel( out = OUT + row * out_batch_stride + head * out_head_stride + dims * out_dim_stride tl.store(out, dense, mask=dims < HEAD_DIM) slot = tl.load(SLOTS + row) - if head == 0: - tl.store(VALID + slot * NUM_LAYERS + layer, 1) tl.store(LSE + rh, maximum + tl.log(normalizer)) if count > 0: state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head @@ -735,7 +732,6 @@ def _capture_kernel( MASS, OUTPUT, CONTRIBUTION, - VALID, layer, top, score_capacity, @@ -765,14 +761,11 @@ def _capture_kernel( One program scores every key, finds the ``count``-th largest score by bisection on its order key, and writes the retained positions in ascending order; equal scores keep the smaller positions, like the oracle's stable - sort. It also marks the reference source-ready for this layer, as - ``AttentionReferenceStore.capture`` does, without a separate launch. + sort. """ row = tl.program_id(0) head = tl.program_id(1) slot = tl.load(SLOTS + row) - if head == 0: - tl.store(VALID + slot * NUM_LAYERS + layer, 1) length = tl.load(LENGTHS + row) # retained_count in float64, exactly as the host's math.ceil(length * r). real = length.to(tl.float64) @@ -955,7 +948,6 @@ def _consume_merge_kernel( MASS, OUTPUT, CONTRIBUTION, - USES, OUT, layer, q_row_stride: tl.constexpr, @@ -966,18 +958,12 @@ def _consume_merge_kernel( BLOCK_D: tl.constexpr, SPLITS: tl.constexpr, ): - """``O = O_ref - C_ref + m * C_new``, merging the split parts in a fixed order. - - Also counts the use, as ``AttentionReferenceStore.consume`` does. - """ + """``O = O_ref - C_ref + m * C_new``, merging the split parts in a fixed order.""" row = tl.program_id(0) head = tl.program_id(1) dims = tl.arange(0, BLOCK_D) out = OUT + row * q_row_stride + head * q_head_stride + dims slot = tl.load(SLOTS + row) - if head == 0: - uses = USES + slot * NUM_LAYERS + layer - tl.store(uses, tl.load(uses) + 1) if tl.load(LENGTHS + row) == 0: tl.store(out, tl.zeros((BLOCK_D,), dtype=tl.float32), mask=dims < HEAD_DIM) else: @@ -1216,7 +1202,6 @@ def fused_paged_attention( plan.slot_ids, store.counts, store.output, - store.valid, lse, seghist, state, @@ -1312,7 +1297,6 @@ def fused_paged_attention( store.mass, store.output, store.contribution, - store.valid, layer, top, capacity, @@ -1395,7 +1379,6 @@ def fused_paged_attention( store.mass, store.output, store.contribution, - store.uses, output, layer, output.stride(0), diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py index 3b289019..ceae6a14 100644 --- a/vllm_rlt/worker/async_speculative.py +++ b/vllm_rlt/worker/async_speculative.py @@ -51,9 +51,6 @@ 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: @@ -169,11 +166,6 @@ 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, @@ -203,7 +195,6 @@ def __init__( # GPU acceptance counts without a host readback. self.reference_store = reference_store self.sparse_decode_config = sparse_decode_config - self._generations: dict[str, int] = {} sizes = speculative_capture_sizes(execution, scheduler, config) self.graphs = ( DeviceRecurrentGraphs(model, cache, execution, False, capture_sizes=sizes) @@ -246,12 +237,8 @@ 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], generation=generation - ) + state = SpeculativeRequestState(request, allocation, values[0], values[1]) else: if not self.free_slots: self._grow_slots() @@ -263,7 +250,6 @@ def _state(self, request, buffer): self.slot_values[slot, 0], self.slot_values[slot, 1], slot=slot, - generation=generation, ) if self.config.mode == "pipeline": width = min( @@ -479,10 +465,6 @@ def _verify(self, plan, starts, hidden_blocks, tables, allocations, leases=None) ] for depth in range(self.config.draft_loops, self.config.target_loops): 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, diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index f3dc69cc..5e22382f 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -37,7 +37,7 @@ from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan -__all__ = ["ReferenceLease", "AttentionReferenceStore", "DENSE", "CAPTURE", "CONSUME"] +__all__ = ["ReferenceLease", "SparseDecodeReferenceStore", "DENSE", "CAPTURE", "CONSUME"] @dataclass(frozen=True) @@ -59,7 +59,7 @@ def __len__(self) -> int: return len(self.slots) -class AttentionReferenceStore: +class SparseDecodeReferenceStore: """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 @@ -114,13 +114,9 @@ def __init__( 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.accepted = torch.zeros(slots, dtype=torch.int32, device=device) + # Usability is decided by counts (zero means no captured reference) and + # by lease ownership, so no per-slot validity, use-count, position or + # acceptance tensor is kept: nothing reads them back. self.free = list(range(slots)) # Slot ownership is the lease object itself. Object identity is exact: # a recycled slot is bound to a different lease, so a plan built for the @@ -142,10 +138,6 @@ def bytes(self) -> int: self.mass, self.output, self.contribution, - self.valid, - self.uses, - self.position, - self.accepted, ) ) @@ -208,13 +200,18 @@ def reserve(self, request_id: str, request_index: int, rows: int) -> ReferenceLe 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)) + # Peek first: a failed device allocation must not lose the slots. They + # are removed from the pool only once the lease is fully built. + slots = tuple(reversed(self.free[-rows:])) + slot_ids = self._index(slots) + row_index = torch.arange(rows, dtype=torch.int64, device=self.device) + del self.free[-rows:] lease = ReferenceLease( request_id=request_id, request_index=request_index, slots=slots, - slot_ids=self._index(slots), - row_index=torch.arange(rows, dtype=torch.int64, device=self.device), + slot_ids=slot_ids, + row_index=row_index, ) for slot in slots: self._bindings[slot] = lease @@ -283,7 +280,6 @@ def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan slots=tuple(slots), owners=tuple(row_owners), key_retention=key_retention, - key_bucket=self.key_bucket, slot_ids=self._index(slots), store=self, ) @@ -331,38 +327,17 @@ def _row_slots(self, plan: SparseDecodePlan, rows: list[int]) -> torch.Tensor: return self._index([plan.slots[row] for row in rows]) def capture(self, plan: SparseDecodePlan, layer: int) -> None: - """Mark the captured rows source-ready for the given layer. + """Assert that a finished capture row still owns its slots. - 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. + The attention kernel writes the retained state of a capture row and then + calls this on the same stream. Nothing is recorded: a reader decides + availability from the retained count, and ownership from the lease. """ self._require_live(plan) - rows = plan.groups().get(CAPTURE) - if not rows: - return - # Assign a tensor, not the Python scalar ``True``: a scalar store into - # a bool tensor through advanced indexing copies it from the host. - slots = self._row_slots(plan, rows) - self.valid[slots, layer] = torch.ones(slots.numel(), dtype=torch.bool, device=self.device) 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[self._row_slots(plan, rows), layer] += 1 - - def bind_positions(self, plan: SparseDecodePlan, positions: torch.Tensor) -> None: - """Record the device-owned absolute positions of the plan's rows.""" + """Assert that a finished consume row still owns its slots.""" self._require_live(plan) - if positions.shape != plan.slot_ids.shape: - raise ValueError("reference positions need one entry per plan row") - rows = [row for row, slot in enumerate(plan.slots) if slot >= 0] - if len(rows) < len(plan.slots): - positions = positions[self._index(rows)] - self.position[self._row_slots(plan, rows)] = positions.to(torch.int64) def write_reference( self, @@ -416,7 +391,7 @@ def accept(self, leases, accepted: torch.Tensor) -> None: host readback. Retirement is memory-bounded: it only clears the per-row metadata that - gates readers (``valid``, ``counts``, ``mass``), never the retained + gates readers (``counts`` and ``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 @@ -433,12 +408,10 @@ 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) - self.accepted[slot_ids] = limit.to(torch.int32) # A boolean-mask gather of the rejected slots would read its size back # to the host. Masked writes over the round's own rows stay on device; # they touch only these rows' per-layer metadata, never the index bucket. keep = rows <= limit - self.valid[slot_ids] = self.valid[slot_ids] & keep[:, None] self.counts[slot_ids] = torch.where(keep[:, None, None], self.counts[slot_ids], 0) self.mass[slot_ids] = torch.where(keep[:, None, None], self.mass[slot_ids], 0) # Which rows were rejected is known only on the device, so drop the host @@ -460,10 +433,6 @@ def zero(): # index_fill_ passes the scalar to the kernel; assigning a Python # scalar through an index tensor copies it from the host and blocks. for tensor in ( - self.valid, - self.uses, - self.accepted, - self.position, self.mass, self.counts, ): @@ -487,7 +456,6 @@ def invalidate(self, slots) -> None: self._forget_counts(slots) slots = self._index(list(slots)) for tensor in ( - self.valid, self.mass, self.output, self.contribution, diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index ec9e28ae..812f56ec 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -1095,7 +1095,6 @@ def _reference_metadata(self, batch, kv): plan = self.reference_store.build_plan( owners, modes, key_retention=self.sparse_decode_config.key_retention ) - self.reference_store.bind_positions(plan, kv.position_ids) return replace(kv, reference_plan=plan) def _load_hidden(self, batch, prepared, workspace, size): diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index 37daa927..c75dd048 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -251,7 +251,6 @@ def _reference_metadata(self, layouts, owners): plan = self.reference_store.build_plan( owners, modes, key_retention=self.sparse_decode_config.key_retention ) - self.reference_store.bind_positions(plan, kv.position_ids) result.append(replace(kv, reference_plan=plan)) return tuple(result) From f140ad57aadf754804b9ac46110b02249c7e8c09 Mon Sep 17 00:00:00 2001 From: kunkunblueberry Date: Thu, 5 Nov 2026 21:27:25 +0800 Subject: [PATCH 28/28] sparse decode attention: per-slot contiguity, owned kernel scratch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two review findings that are not style. **Fragmented slots were read as a span (P1).** The batched path accepted a chunk when `first + len == last + 1`, which `[5, 3, 7]` satisfies, and then used `slice(5, 8)`. CONSUME read slot 6 — another request's reference — and CAPTURE overwrote it while leaving slot 3 unwritten. `slot_span` now checks every slot against `first + offset` and returns None otherwise, so a fragmented chunk takes the indexed path. Both call sites use it. `test_fragmented_consume_reads_each_rows_own_reference` and `test_fragmented_capture_writes_only_the_slots_it_reserved` build the exact fragmentation from the report (release 7, then 3, then 5; reserve three rows gives `[5, 3, 7]`) with slots 4 and 6 held by unrelated live leases. Verified to fail against the endpoint check: slot 3 keeps its sentinel, which is the reported symptom. **Kernel scratch outlived its engine (P1).** `_WORKSPACE` and `_CONSUME_WORKSPACE` were module-level dicts keyed by (device, stream) that nothing released, so a closed engine's tensors stayed reachable — the capture score buffer alone is rows x heads x capacity floats. They move to `kernels/sparse_workspace.py` (Triton-free, so CPU-only installs can release them) and `SparseDecodeReferenceStore.close()` releases its device after the caller has drained; `StageEngine.close()` calls it after its own drain. **Cleanup.** `_reference_metadata` is `_prepare_sparse_reference` on both the runner and the speculative runner, since it reserves and releases leases rather than building metadata. `SparseDecodePlan.reference_depth` is gone: it returned a module constant and had no production reader. Regression: CPU 986 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 222 passed / 4 skipped; ruff 0.13.2 clean. --- tests/test_async_speculative.py | 12 ++-- tests/test_sparse_decode_attention.py | 42 ++++++++++- tests/test_sparse_decode_batched.py | 98 ++++++++++++++++++++++++++ tests/test_sparse_decode_fused.py | 4 +- vllm_rlt/attention/sparse.py | 36 ++++++---- vllm_rlt/engine/stage_engine.py | 3 + vllm_rlt/kernels/sparse_decode.py | 24 +++---- vllm_rlt/kernels/sparse_workspace.py | 37 ++++++++++ vllm_rlt/worker/attention_reference.py | 15 ++++ vllm_rlt/worker/model_runner.py | 4 +- vllm_rlt/worker/speculative.py | 8 ++- 11 files changed, 243 insertions(+), 40 deletions(-) create mode 100644 vllm_rlt/kernels/sparse_workspace.py diff --git a/tests/test_async_speculative.py b/tests/test_async_speculative.py index 47fde74a..eb45029e 100644 --- a/tests/test_async_speculative.py +++ b/tests/test_async_speculative.py @@ -191,13 +191,17 @@ def test_eos_discards_queued_round_and_the_id_frees_after_retirement(monkeypatch e.step() out = e._collect_speculative()[0] assert out.finished and out.token_ids[-1] == 0 and len(out.token_ids) == 2 - # The finished request is still retiring, so its id is rejected: the engine - # never reuses an id while the previous token may still own resources. - with pytest.raises(ValueError, match="duplicate request ID"): - e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True)) + # The engine never reuses an id while the previous token may still own + # resources. Whether the round has already retired by now depends on the + # graph setting, so assert the rejection only while the engine still lists + # the id; the live-duplicate case is covered by test_stage_engine. + if "r" in e._retiring_speculative: + with pytest.raises(ValueError, match="duplicate request ID"): + e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True)) # Once retirement completes the id is free, and the new request never sees # the old round's state. e._retire_speculative(wait=True) + assert "r" not in e._retiring_speculative monkeypatch.setattr(e.speculative_runner, "_coda", coda) e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True)) output = drain(e)["r"] diff --git a/tests/test_sparse_decode_attention.py b/tests/test_sparse_decode_attention.py index a638995c..d51b307f 100644 --- a/tests/test_sparse_decode_attention.py +++ b/tests/test_sparse_decode_attention.py @@ -86,7 +86,6 @@ def test_plan_groups_rows_by_mode(): ) assert plan.groups() == {DENSE: [0], CAPTURE: [1, 2], CONSUME: [3]} assert not plan.dense() - assert plan.reference_depth == REFERENCE_DEPTH def test_retained_count_follows_flashloop_ceil_rule(): @@ -820,3 +819,44 @@ def counting_reserve(*args, **kwargs): assert store.outstanding == 0 assert len(store.free) == store.slots assert max_outstanding > 0 + + +def test_store_close_drops_the_device_scratch(monkeypatch): + released = [] + monkeypatch.setattr( + "vllm_rlt.worker.attention_reference.release_workspace", + lambda device: released.append(device), + ) + store = make_store(slots=2) + store.reserve("r", 0, 2) + store.close() + # Teardown must leave nothing bound and hand the kernel scratch back; the + # caller has drained the stream, so the tensors cannot be in use. + assert store.outstanding == 0 + assert not store._bindings + assert released == [store.device] + + +def test_release_workspace_is_scoped_to_the_closed_device(): + from vllm_rlt.kernels.sparse_workspace import ( + CAPTURE_SCRATCH, + CONSUME_SCRATCH, + RETENTION_SCALARS, + release_workspace, + ) + + cpu, other = torch.device("cpu"), torch.device("meta") + CAPTURE_SCRATCH[(cpu, 1)] = torch.zeros(4) + CONSUME_SCRATCH[(cpu, 1)] = torch.zeros(4) + RETENTION_SCALARS[(cpu, 0.5)] = torch.zeros(1, dtype=torch.float64) + CAPTURE_SCRATCH[(other, 2)] = torch.zeros(4) + try: + release_workspace(cpu) + assert (cpu, 1) not in CAPTURE_SCRATCH + assert (cpu, 1) not in CONSUME_SCRATCH + assert (cpu, 0.5) not in RETENTION_SCALARS + # A second engine on another device keeps its scratch. + assert (other, 2) in CAPTURE_SCRATCH + finally: + release_workspace() + assert not CAPTURE_SCRATCH and not CONSUME_SCRATCH and not RETENTION_SCALARS diff --git a/tests/test_sparse_decode_batched.py b/tests/test_sparse_decode_batched.py index ef201262..bb204f60 100644 --- a/tests/test_sparse_decode_batched.py +++ b/tests/test_sparse_decode_batched.py @@ -10,6 +10,7 @@ DENSE, batched_paged_attention, reference_paged_attention, + slot_span, ) @@ -165,3 +166,100 @@ def test_reference_growth_rejects_foreign_stream(): with pytest.raises(RuntimeError, match="owner stream"): store.ensure_key_capacity(3) assert store.selected.shape[-1] == 2 + + +def test_slot_span_requires_every_slot_in_place(): + # Endpoint arithmetic accepts [5, 3, 7] (5 + 3 == 7 + 1) and would then use + # the slice [5, 6, 7], which is another request's reference. + assert slot_span([5, 3, 7]) is None + assert slot_span([7, 6, 5]) is None + assert slot_span([5, 6, 8]) is None + assert slot_span([5, 5, 5]) is None + assert slot_span([3, 4, 5]) == slice(3, 6) + assert slot_span([5]) == slice(5, 6) + + +def _fragmented_store(): + """A store whose next three-row reservation yields slots ``[5, 3, 7]``. + + ``reserve`` reverses the tail of the free list, so releasing 7, then 3, then + 5 leaves ``[7, 3, 5]`` and the reservation comes back fragmented. Slots 6 and + 4 stay bound to unrelated live leases, which is what a span read of + ``[5, 6, 7]`` would corrupt. + """ + store = make_store(slots=8, layers=1, heads=2, head_dim=4, key_bucket=8) + held = [store.reserve(f"held{i}", i, 1) for i in range(5)] # slots 7, 6, 5, 4, 3 + store.release([held[0], held[4], held[2]]) + assert store.free[-3:] == [7, 3, 5] + lease = store.reserve("fragmented", 0, 3) + assert lease.slots == (5, 3, 7) + return store, lease, held[1], held[3] + + +def test_fragmented_consume_reads_each_rows_own_reference(): + torch.manual_seed(157) + store, lease, live_six, live_four = _fragmented_store() + rows, heads, kv_heads, dim, block_size, extent = 3, 2, 2, 4, 8, 8 + # A zero query makes attention uniform, so the correction is identical for + # every row and each output equals its own slot's reference plus one + # constant. Reading the span [5, 6, 7] breaks exactly one row, by 3. + queries = torch.zeros(rows, heads, dim) + keys = torch.randn(1, block_size, kv_heads, dim) + values = torch.randn(1, block_size, kv_heads, dim) + tables = torch.zeros(rows, 1, dtype=torch.int32) + lengths = torch.full((rows,), extent) + + store.output.copy_(torch.arange(8).float()[:, None, None, None]) + store.contribution.zero_() + store.mass.fill_(1.0) + store.counts.fill_(extent) + store.selected.copy_(torch.arange(extent).expand(8, 1, heads, extent)) + for slot in lease.slots: + store.note_count(slot, 0, extent) + + plan = store.build_plan( + [(lease, 0), (lease, 1), (lease, 2)], [CONSUME] * rows, key_retention=0.5 + ) + output = batched_paged_attention(queries, keys, values, tables, lengths, plan, layer=0).float() + + own_reference = torch.tensor([5.0, 3.0, 7.0])[:, None, None] + correction = output - own_reference + torch.testing.assert_close(correction, correction[0].expand_as(correction), rtol=0, atol=1e-6) + # The unrelated owners keep their payload. + assert float(store.output[6, 0].mean()) == 6.0 + assert float(store.output[4, 0].mean()) == 4.0 + store.release([lease, live_six, live_four]) + + +def test_fragmented_capture_writes_only_the_slots_it_reserved(): + torch.manual_seed(157) + store, lease, live_six, live_four = _fragmented_store() + rows, heads, kv_heads, dim, block_size, extent = 3, 2, 2, 4, 8, 8 + queries = torch.randn(rows, heads, dim) + keys = torch.randn(1, block_size, kv_heads, dim) + values = torch.randn(1, block_size, kv_heads, dim) + tables = torch.zeros(rows, 1, dtype=torch.int32) + lengths = torch.full((rows,), extent) + + # A sentinel per slot, so a write to the wrong slot is visible. + sentinel = torch.arange(8).float() - 1000.0 + store.output.copy_(sentinel[:, None, None, None]) + store.contribution.copy_(sentinel[:, None, None, None]) + store.mass.copy_(sentinel[:, None, None]) + store.counts.fill_(-1) + + plan = store.build_plan( + [(lease, 0), (lease, 1), (lease, 2)], [CAPTURE] * rows, key_retention=1.0 + ) + batched_paged_attention(queries, keys, values, tables, lengths, plan, layer=0) + + for slot in lease.slots: + assert float(store.output[slot, 0].mean()) != float(sentinel[slot]), slot + assert bool((store.counts[slot, 0] == extent).all()), slot + # Slot 6 belongs to another request: the span [5, 6, 7] overwrote it and + # left slot 3 unwritten. + for slot in (4, 6): + assert float(store.output[slot, 0].mean()) == float(sentinel[slot]), slot + assert float(store.mass[slot, 0].mean()) == float(sentinel[slot]), slot + assert bool((store.counts[slot, 0] == -1).all()), slot + store.release([lease, live_six, live_four]) diff --git a/tests/test_sparse_decode_fused.py b/tests/test_sparse_decode_fused.py index 6029d8de..d5adb3cf 100644 --- a/tests/test_sparse_decode_fused.py +++ b/tests/test_sparse_decode_fused.py @@ -145,9 +145,9 @@ def consume(dim): torch.cuda.synchronize() return output - monkeypatch.setattr(sparse_decode, "_CONSUME_WORKSPACE", {}) + monkeypatch.setattr(sparse_decode, "CONSUME_SCRATCH", {}) fresh = consume(128) - sparse_decode._CONSUME_WORKSPACE.clear() + sparse_decode.CONSUME_SCRATCH.clear() consume(16) # The same stream now holds 16-wide scratch; a 128-wide consume must regrow it. assert torch.equal(consume(128), fresh) diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py index 078e857a..60c6abba 100644 --- a/vllm_rlt/attention/sparse.py +++ b/vllm_rlt/attention/sparse.py @@ -35,7 +35,7 @@ from __future__ import annotations import math -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING @@ -108,11 +108,6 @@ def groups(self) -> dict[int, list[int]]: 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) @@ -134,6 +129,22 @@ def retained_count(length: int, key_retention: float) -> int: return min(length, max(1, math.ceil(length * key_retention))) +def slot_span(slots: Sequence[int]) -> slice | None: + """The slice covering ``slots`` when they are consecutive, else ``None``. + + Slot numbers follow the free-list order, not the row order, so a chunk of + rows can arrive fragmented (``[5, 3, 7]``). Comparing the endpoints would + accept that chunk — ``5 + 3 == 7 + 1`` — and then read or write ``[5, 6, 7]``, + which belongs to other requests. Every slot is checked against + ``first + offset`` instead. + """ + first = slots[0] + for offset, slot in enumerate(slots): + if slot != first + offset: + return None + return slice(first, first + len(slots)) + + def expand_kv(tensor: torch.Tensor, groups: int) -> torch.Tensor: """Repeat KV heads for grouped-query attention.""" if groups == 1: @@ -528,10 +539,11 @@ def batched_paged_attention( output[row_ids] = dense_output[row_ids] consumed = True continue - # Contiguous slots turn these gathers into views, which is the - # difference between one index kernel per tensor and none. - if chunk_slots[0] + len(chunk_slots) == chunk_slots[-1] + 1: - span = slice(chunk_slots[0], chunk_slots[0] + len(chunk_slots)) + # Consecutive slots turn these gathers into views, which is the + # difference between one index kernel per tensor and none. The + # order is checked per slot, never from the endpoints. + span = slot_span(chunk_slots) + if span is not None: selected = store.selected[span, layer, :, :extent].to(torch.int32) ref_output = store.output[span, layer] ref_contribution = store.contribution[span, layer] @@ -626,8 +638,8 @@ def batched_paged_attention( ].float() contribution = _rowwise_product("hk,hkd->hd", gathered, selected_values) # Only the live prefix is written, never the whole capacity. - if chunk_slots[0] + len(chunk_slots) == chunk_slots[-1] + 1: - span = slice(chunk_slots[0], chunk_slots[0] + len(chunk_slots)) + span = slot_span(chunk_slots) + if span is not None: store.selected[span, layer, :, :count] = selected.to(torch.int32) store.counts[span, layer] = count store.mass[span, layer] = gathered.sum(dim=-1) diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 4e5fa4b5..65ddc157 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/vllm_rlt/engine/stage_engine.py @@ -518,6 +518,9 @@ def close(self): # Abort can enqueue reference metadata retirement after the drain # above. Teardown must also cover those final writes. self.model_runner.synchronize() + # The store owns the stream the fused kernels ran on, so it releases + # their module-level scratch here — after the drain, never before. + self.reference_store.close() def _discard_speculative(self): """Release speculative owners after the caller has drained both streams.""" diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py index 059f067a..d0c8ffd0 100644 --- a/vllm_rlt/kernels/sparse_decode.py +++ b/vllm_rlt/kernels/sparse_decode.py @@ -26,6 +26,7 @@ from vllm_rlt.attention.sparse import CAPTURE, CONSUME, retained_count from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention +from vllm_rlt.kernels.sparse_workspace import CAPTURE_SCRATCH, CONSUME_SCRATCH, RETENTION_SCALARS from vllm_rlt.kernels.triton_attention import NUM_SPLITS @@ -984,18 +985,12 @@ def _consume_merge_kernel( tl.store(out, corrected, mask=dims < HEAD_DIM) -_RETENTION: dict[tuple[torch.device, float], torch.Tensor] = {} - - def _retention(device, key_retention): """A float64 device scalar, built once; Triton would pass a Python float as fp32.""" key = (device, key_retention) - if key not in _RETENTION: - _RETENTION[key] = torch.full((1,), key_retention, dtype=torch.float64, device=device) - return _RETENTION[key] - - -_WORKSPACE: dict = {} + if key not in RETENTION_SCALARS: + RETENTION_SCALARS[key] = torch.full((1,), key_retention, dtype=torch.float64, device=device) + return RETENTION_SCALARS[key] def _build_capture_workspace(device, rows, heads, capacity, block_d): @@ -1028,20 +1023,17 @@ def _capture_workspace(device, rows, heads, capacity, block_d): """ stream = _current_stream(device) shape = (rows, heads, capacity, block_d) - entry = _WORKSPACE.get((device, stream)) + entry = CAPTURE_SCRATCH.get((device, stream)) if entry is None or entry[0] != shape: entry = (shape, *_build_capture_workspace(device, rows, heads, capacity, block_d)) - _WORKSPACE[(device, stream)] = entry + CAPTURE_SCRATCH[(device, stream)] = entry return entry[1:] -_CONSUME_WORKSPACE: dict = {} - - def _consume_workspace(device, rows, heads, block_d): """Consume split parts, grown per stream and reused across layers.""" stream = _current_stream(device) - entry = _CONSUME_WORKSPACE.get((device, stream)) + entry = CONSUME_SCRATCH.get((device, stream)) if ( entry is None or entry[0].shape[0] < rows @@ -1054,7 +1046,7 @@ def _consume_workspace(device, rows, heads, block_d): empty(rows, heads, NUM_SPLITS), empty(rows, heads, NUM_SPLITS, block_d), ) - _CONSUME_WORKSPACE[(device, stream)] = entry + CONSUME_SCRATCH[(device, stream)] = entry return (None, *entry) diff --git a/vllm_rlt/kernels/sparse_workspace.py b/vllm_rlt/kernels/sparse_workspace.py new file mode 100644 index 00000000..b365e352 --- /dev/null +++ b/vllm_rlt/kernels/sparse_workspace.py @@ -0,0 +1,37 @@ +"""Module-level scratch caches of the fused sparse decode kernels. + +The kernels reuse their scratch across layers of one stream, so the buffers live +outside any single call. They are keyed by ``(device, stream)`` so two engines on +different streams do not share them — but nothing owns an entry, which means a +closed engine's tensors stay reachable for the life of the process. The capture +score buffer alone is ``rows * heads * capacity`` floats, so a long-context run +leaves a large allocation behind. + +This module is deliberately free of Triton imports: the engine that owns the +stream has to be able to release the scratch on a CPU-only install too. +""" + +import torch + +# Scratch of one capture launch, per (device, stream). +CAPTURE_SCRATCH: dict = {} + +# Scratch of one consume launch, per (device, stream). +CONSUME_SCRATCH: dict = {} + +# ``key_retention`` as a float64 device scalar, per (device, retention). +RETENTION_SCALARS: dict[tuple[torch.device, float], torch.Tensor] = {} + + +def release_workspace(device=None) -> None: + """Drop the scratch cached for ``device``, or for every device. + + Call this only after the caller has drained the stream, or a kernel may + still be reading the tensors. Releasing too much is never a correctness + problem for the next user — an entry is rebuilt lazily on demand — only an + extra allocation, so tearing down a whole device is safe when several + engines share it. + """ + for cache in (CAPTURE_SCRATCH, CONSUME_SCRATCH, RETENTION_SCALARS): + for key in [key for key in cache if device is None or key[0] == device]: + del cache[key] diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py index 5e22382f..58accae1 100644 --- a/vllm_rlt/worker/attention_reference.py +++ b/vllm_rlt/worker/attention_reference.py @@ -36,6 +36,7 @@ import torch from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan +from vllm_rlt.kernels.sparse_workspace import release_workspace __all__ = ["ReferenceLease", "SparseDecodeReferenceStore", "DENSE", "CAPTURE", "CONSUME"] @@ -251,6 +252,20 @@ def discard(self) -> None: self._counts_host.clear() self.free = list(range(self.slots)) + def close(self) -> None: + """Release the pool and the kernel scratch this store's device owns. + + The fused kernels cache per-stream scratch at module level, which + nothing else can reach. This store is the execution owner of the stream + those launches ran on, so it drops them here. The caller must have + drained the stream: the tensors must not be freed under a running + kernel. A later engine rebuilds them lazily. + """ + self.discard() + # The fused kernels cache per-stream scratch at module level, and this + # store is the execution owner of the stream they ran on. + release_workspace(self.device) + def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan: """Publish the plan of one submission. diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py index 812f56ec..d9d92cca 100644 --- a/vllm_rlt/worker/model_runner.py +++ b/vllm_rlt/worker/model_runner.py @@ -1036,7 +1036,7 @@ def _prepare_forward(self, prepared): if kv is not None: metadata = (kv,) + metadata if self.reference_store is not None: - metadata = tuple(self._reference_metadata(batch, kv) for kv in metadata) + metadata = tuple(self._prepare_sparse_reference(batch, kv) for kv in metadata) forward = ForwardBatch( Operation.RECURRENT, hidden=hidden[: batch.num_rows] @@ -1063,7 +1063,7 @@ def _release_attention_reference(self, request_id): ): self.reference_store.release(state[2]) - def _reference_metadata(self, batch, kv): + def _prepare_sparse_reference(self, batch, kv): """Borrow token-owned references; prompt and padded rows remain dense. Preparation runs on core, after the request's preceding event. CODA diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py index c75dd048..21e2008d 100644 --- a/vllm_rlt/worker/speculative.py +++ b/vllm_rlt/worker/speculative.py @@ -233,15 +233,17 @@ def metadata(ids, positions, depths, *, packed=False): steps = [ ( active, - self._reference_metadata(layouts, [(references[i], offset) for i in active]), + self._prepare_sparse_reference( + layouts, [(references[i], offset) for i in active] + ), ) for offset, (active, layouts) in enumerate(steps) ] owners = [(lease, row) for lease in references for row in range(len(lease))] - verification = self._reference_metadata(verification, owners) + verification = self._prepare_sparse_reference(verification, owners) return tokens, steps, verification - def _reference_metadata(self, layouts, owners): + def _prepare_sparse_reference(self, layouts, owners): result = [] for kv in layouts: modes = [reference_mode(depth) for _, depth, _ in kv.rows]