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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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 <dezhen.lu@student.uni-tuebingen.de>
Date: Wed, 28 Oct 2026 15:00:37 +0100
Subject: [PATCH 11/32] 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 <dezhen.lu@student.uni-tuebingen.de>
Date: Wed, 28 Oct 2026 15:15:17 +0100
Subject: [PATCH 12/32] 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 <dezhen.lu@student.uni-tuebingen.de>
Date: Wed, 28 Oct 2026 15:30:01 +0100
Subject: [PATCH 13/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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/32] 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 <dezhen.lu@student.uni-tuebingen.de>
Date: Sun, 1 Nov 2026 17:54:57 +0100
Subject: [PATCH 22/32] 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 <dezhen.lu@student.uni-tuebingen.de>
---
 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 <dezhen.lu@student.uni-tuebingen.de>
Date: Sun, 1 Nov 2026 18:53:00 +0100
Subject: [PATCH 23/32] 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 <dezhen.lu@student.uni-tuebingen.de>
---
 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 1d98a84897b2e558d9747571ca850979f42f0c50 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Sat, 31 Oct 2026 21:09:35 +0100
Subject: [PATCH 24/32] Replay fused sparse decode in async speculative CUDA
 graphs and P/D decode

On top of the eager fused path:

- Async speculative CUDA graphs keep one entry per bucket and sparse mode.
  Their slot ids form a static plan refilled before each replay, with a
  never-bound padding slot for padded rows, and the graph owns its capture
  workspace with the index pool allocated whole. The engine accepts sparse
  with graphs only for async block speculation, the fused implementation and
  r < 1, and budgets the graph workspace and selection scratch.
- P/D passes the sparse decode configuration to the decode stage only.

Measured end to end with Ouro-1.4B, async self-speculation K=3, 512 greedy
tokens, interleaved dense/sparse engines: B=2 x 32k +9.3% (r=0.1) / +6.1%
(r=0.2), B=4 x 16k +9.6% / +7.2% (RTX PRO 6000); B=2 x 14.6k +3% to +8%
(RTX 5090); batch 1 parity.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_pd.py                       | 144 ++++++++++++++++++-
 tests/test_sparse_decode_fused.py      | 110 ++++++++++++++
 vllm_rlt/engine/stage_engine.py        |  41 +++++-
 vllm_rlt/engine/worker.py              |   9 +-
 vllm_rlt/pd/engine.py                  |   6 +
 vllm_rlt/worker/async_speculative.py   |  10 +-
 vllm_rlt/worker/attention_reference.py |  11 +-
 vllm_rlt/worker/cuda_graph.py          | 190 ++++++++++++++++++++-----
 8 files changed, 475 insertions(+), 46 deletions(-)

diff --git a/tests/test_pd.py b/tests/test_pd.py
index 90a4b38f..f98f2b48 100644
--- a/tests/test_pd.py
+++ b/tests/test_pd.py
@@ -13,6 +13,7 @@
     ExitConfig,
     SamplingParams,
     SchedulerConfig,
+    SparseDecodeAttentionConfig,
     SpeculativeConfig,
 )
 from vllm_rlt.core.kv_cache_manager import KVCacheManager
@@ -118,6 +119,47 @@ def test_pd_rejects_huginn_before_spawning(monkeypatch):
     clients.assert_not_called()
 
 
+def test_pd_forwards_sparse_decode_only_to_decode_workers(monkeypatch):
+    from unittest.mock import Mock
+
+    from vllm_rlt.pd.engine import PDEngine
+
+    config = tiny_ouro_config(head_dim=64)
+    sparse = SparseDecodeAttentionConfig(enabled=True, key_retention=0.5)
+
+    def ready_client(model, *, stage_config, **options):
+        return Mock(
+            ready=True,
+            connected=True,
+            stopped=True,
+            info=dict(
+                kind="ready",
+                stage_id=stage_config.stage_id,
+                role=stage_config.role,
+                fingerprint="same-model-and-cache",
+                info={"agent": stage_config.stage_id},
+                num_blocks=64,
+                block_size=16,
+                depths=model.total_ut_steps,
+            ),
+        )
+
+    clients = Mock(side_effect=ready_client)
+    monkeypatch.setattr("vllm_rlt.pd.engine.StageClient", clients)
+    with PDEngine(
+        config,
+        attention_backend="triton",
+        speculative_config=SpeculativeConfig(3),
+        sparse_decode_attention_config=sparse,
+    ):
+        roles = {
+            call.kwargs["stage_config"].role: call.kwargs["sparse_decode_attention_config"]
+            for call in clients.call_args_list
+        }
+    # Prefill rows are always dense: only the decode stage owns reference state.
+    assert roles == {"prefill": None, "decode": sparse}
+
+
 def test_pd_nanbeige_startup_uses_entry_config_for_request_validation(monkeypatch):
     from unittest.mock import Mock
 
@@ -317,6 +359,8 @@ def create_pd(
     speculative=False,
     asynchronous=True,
     prefix_caching=False,
+    sparse=None,
+    decode_batched_tokens=2,
 ):
     pytest.importorskip("nixl")
     from vllm_rlt.pd.engine import PDEngine
@@ -343,11 +387,14 @@ def create_pd(
             prefill_chunk_size=prefill_chunk_size,
             wavefront_chunked_prefill=wavefront_chunked_prefill,
         ),
-        decode_scheduler_config=SchedulerConfig(max_num_seqs=2, max_num_batched_tokens=2),
+        decode_scheduler_config=SchedulerConfig(
+            max_num_seqs=2, max_num_batched_tokens=decode_batched_tokens
+        ),
         exit_config=ExitConfig("ouro")
         if speculative
         else ExitConfig("trace", depths_by_request={"t": [4, 2, 3, 1] * 16}),
         speculative_config=SpeculativeConfig(3) if speculative else None,
+        sparse_decode_attention_config=sparse,
         execution_config=ExecutionConfig(
             async_scheduling=asynchronous,
             static_buffers=graph,
@@ -682,6 +729,101 @@ def test_pd_speculative_matches_local_and_retires_cancelled_rounds(
     assert all(metrics["used_blocks"] == 0 for metrics in engine.worker_metrics.values())
 
 
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "sparse",
+    [
+        SparseDecodeAttentionConfig(enabled=True, key_retention=1.0),
+        SparseDecodeAttentionConfig(enabled=True, key_retention=0.5),
+        SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused"),
+    ],
+)
+def test_pd_sparse_decode_on_the_decode_stage(sparse):
+    params = SamplingParams(max_tokens=9, ignore_eos=True)
+    prompts = [[1, 2, 3, 4, 5], [2, 3], [4], [5, 6, 7]]
+    with create_pd(speculative=True, asynchronous=True, sparse=sparse) as engine:
+        for index, prompt in enumerate(prompts):
+            engine.add_request(str(index), prompt, params)
+        outputs = drain(engine)
+        assert sorted(outputs) == ["0", "1", "2", "3"]
+        assert all(len(output.token_ids) == 9 for output in outputs.values())
+        if sparse.key_retention == 1.0:
+            # Full retention is the exact control: it must reproduce local dense decoding.
+            torch.manual_seed(123)
+            reference = LLMEngine(
+                OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda:0", torch.bfloat16),
+                cache_config=CacheConfig(128, 4),
+                attention_backend="triton",
+            )
+            try:
+                for index, prompt in enumerate(prompts):
+                    reference.add_request(str(index), prompt, params)
+                assert outputs == drain(reference)
+            finally:
+                reference.close()
+        # A request aborted mid-decode must still release its reference slots.
+        engine.add_request("abort", prompts[0], replace(params, max_tokens=60))
+        deadline = time.monotonic() + 60
+        while time.monotonic() < deadline:
+            if any(len(output.token_ids) > 1 for output in engine.step()):
+                break
+        else:
+            pytest.fail("sparse speculative decode did not start")
+        engine.abort_request("abort")
+        drain(engine)
+        assert all(peer.slots == 0 and peer.reserved_blocks == 0 for peer in engine.peers.values())
+    assert all(metrics["used_blocks"] == 0 for metrics in engine.worker_metrics.values())
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("multi", [False, True])
+def test_pd_sparse_decode_replays_graphs_on_the_decode_stage(multi):
+    fused = SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused")
+    params = SamplingParams(max_tokens=9, ignore_eos=True)
+    prompts = [[1, 2, 3, 4, 5], [2, 3], [4], [5, 6, 7]]
+    results = {}
+    for graph, sparse in [(False, None), (True, None), (False, fused), (True, fused)]:
+        # Room for full verification rounds, so consume steps replay graphs too.
+        with create_pd(
+            graph=graph,
+            multi=multi,
+            speculative=True,
+            asynchronous=True,
+            sparse=sparse,
+            decode_batched_tokens=8,
+        ) as engine:
+            for index, prompt in enumerate(prompts):
+                engine.add_request(str(index), prompt, params)
+            outputs = drain(engine)
+            assert sorted(outputs) == ["0", "1", "2", "3"]
+            results[graph, sparse is not None] = outputs
+            engine.add_request("abort", prompts[0], replace(params, max_tokens=60))
+            deadline = time.monotonic() + 60
+            while time.monotonic() < deadline:
+                if any(len(output.token_ids) > 1 for output in engine.step()):
+                    break
+            else:
+                pytest.fail("speculative decode did not start")
+            engine.abort_request("abort")
+            drain(engine)
+            assert all(
+                peer.slots == 0 and peer.reserved_blocks == 0 for peer in engine.peers.values()
+            )
+        metrics = engine.worker_metrics
+        assert all(m["used_blocks"] == 0 for m in metrics.values())
+        decode = [m for name, m in metrics.items() if "-decode-" in name]
+        assert len(decode) == (2 if multi else 1)
+        if graph:
+            assert all(m["graph_replays"] > 0 for m in decode)
+        # Every decode worker owns its store and replays its sparse graphs.
+        assert all(
+            (m["sparse_graph_replays"] > 0) == (graph and sparse is not None) for m in decode
+        )
+    assert results[True, False] == results[False, False]
+    # The replayed fused kernels are the eager ones, row for row.
+    assert results[True, True] == results[False, True]
+
+
 def test_pd_reservations_are_reported_before_worker_acknowledgments():
     from types import SimpleNamespace
 
diff --git a/tests/test_sparse_decode_fused.py b/tests/test_sparse_decode_fused.py
index c1e35db8..ecae5374 100644
--- a/tests/test_sparse_decode_fused.py
+++ b/tests/test_sparse_decode_fused.py
@@ -521,3 +521,113 @@ def counting(*args, **kwargs):
     assert store.outstanding == 0 and len(store.free) == store.slots
     assert engine.cache_manager.num_used_blocks == 0
     engine.close()
+
+
+def test_padding_slot_is_allocated_but_never_bound():
+    store = make_store(slots=3, padding_slot=True)
+    assert store.padding_slot == 3 and store.slots == 3
+    assert store.valid.shape[0] == store.uses.shape[0] == 4
+    lease = store.reserve("r", 0, 3, generation=1)
+    assert store.padding_slot not in lease.slots
+    with pytest.raises(RuntimeError, match="no retired"):
+        store.reserve("s", 1, 1, generation=1)
+    store.release(lease)
+    store.discard()
+    assert sorted(store.free) == [0, 1, 2]
+    assert make_store(slots=3).padding_slot is None
+
+
+def _graph_engine(*, graphs, max_rows=32):
+    from tests.helpers import tiny_ouro_config
+    from vllm_rlt import CacheConfig, ExecutionConfig, SchedulerConfig, SpeculativeConfig
+    from vllm_rlt.engine.stage_engine import StageEngine
+    from vllm_rlt.models import OuroForCausalLM
+
+    torch.manual_seed(123)
+    model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.bfloat16)
+    return StageEngine(
+        model,
+        attention_backend="triton",
+        cache_config=CacheConfig(128, 16, incremental_allocation=True),
+        scheduler_config=SchedulerConfig(
+            max_num_seqs=3, max_num_batched_tokens=16, prefill_chunk_size=8
+        ),
+        execution_config=ExecutionConfig(
+            async_scheduling=True,
+            static_buffers=True,
+            pad_to_power_of_two=True,
+            cuda_graphs=graphs,
+            cuda_graph_max_batch_size=max_rows,
+        ),
+        speculative_config=SpeculativeConfig(3),
+        sparse_decode_attention_config=SparseDecodeAttentionConfig(
+            enabled=True, key_retention=0.5, implementation="fused"
+        ),
+    )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("max_rows", [32, 2])
+def test_fused_graphs_replay_sparse_and_match_eager(max_rows, monkeypatch):
+    import vllm_rlt.kernels.sparse_decode as fused_module
+
+    # Graph and eager captures must both take the split path, not the serial fallback.
+    monkeypatch.setattr(fused_module, "_capture_kernel", None)
+    results = []
+    for graphs in (False, True):
+        engine = _graph_engine(graphs=graphs, max_rows=max_rows)
+        store = engine.reference_store
+        try:
+            used = 0
+            for i, (prompt_length, output_length) in enumerate([(15, 23), (17, 9), (31, 1)]):
+                engine.add_request(
+                    str(i),
+                    [2 + i, 3, 4] * (prompt_length // 3) + [5] * (prompt_length % 3),
+                    SamplingParams(max_tokens=output_length, 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
+                used = max(used, int(store.uses[: store.slots].sum()))
+            assert len(outputs) == 3
+            results.append({rid: o.token_ids for rid, o in outputs.items()})
+            # Consumers ran on live slots, not only on the padding slot.
+            assert used > 0
+            assert store.outstanding == 0 and len(store.free) == store.slots
+            assert engine.cache_manager.num_used_blocks == 0
+            graphs_owner = engine.speculative_runner.graphs
+            if graphs:
+                assert graphs_owner.sparse_replays > 0
+                if max_rows == 32:
+                    assert graphs_owner.fallbacks == 0
+                else:
+                    assert graphs_owner.fallbacks > 0
+        finally:
+            engine.close()
+    # The same per-row kernels run eager and replayed.
+    assert results[0] == results[1]
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "sparse",
+    [
+        SparseDecodeAttentionConfig(enabled=True, key_retention=0.5),
+        SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, implementation="fused"),
+    ],
+)
+def test_graphs_reject_sparse_paths_that_cannot_be_replayed(sparse):
+    from vllm_rlt import ExecutionConfig, SpeculativeConfig
+    from vllm_rlt.engine.stage_engine import StageEngine
+
+    with pytest.raises(ValueError, match="CUDA graphs requires"):
+        StageEngine(
+            tiny_ouro_model("cuda"),
+            attention_backend="triton",
+            execution_config=ExecutionConfig(async_scheduling=True, cuda_graphs=True),
+            speculative_config=SpeculativeConfig(3),
+            sparse_decode_attention_config=sparse,
+        )
diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py
index f4244e3e..5bb2b23f 100644
--- a/vllm_rlt/engine/stage_engine.py
+++ b/vllm_rlt/engine/stage_engine.py
@@ -132,8 +132,18 @@ def __init__(
                 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")
+            # Graph replays run only the fused kernels, which need no host
+            # work per row; other paths would silently fall back to dense.
+            graphs = self.execution_config.cuda_graphs
+            if graphs and (
+                not self.async_speculative
+                or sparse_decode.implementation != "fused"
+                or sparse_decode.key_retention >= 1.0
+            ):
+                raise ValueError(
+                    "sparse decode attention with CUDA graphs requires async block "
+                    "speculation, the fused implementation and key_retention < 1"
+                )
             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":
@@ -174,7 +184,7 @@ def __init__(
             # 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
+                (slots + graphs)
                 * config.num_hidden_layers
                 * config.num_attention_heads
                 * (capacity * 4 + 4 + 4 * 2 * config.head_dim + 8)
@@ -194,9 +204,11 @@ def __init__(
                 heads=config.num_attention_heads,
                 head_dim=config.head_dim,
                 key_bucket=capacity,
-                initial_key_capacity=min(capacity, 64),
+                # Graphs bake the index pool's address in: allocate it whole.
+                initial_key_capacity=capacity if graphs else min(capacity, 64),
                 device=parameter.device,
                 implementation=sparse_decode.implementation,
+                padding_slot=graphs,
             )
             if parameter.device.type == "cuda":
                 free_bytes, _ = torch.cuda.mem_get_info(parameter.device)
@@ -219,6 +231,27 @@ def __init__(
             # 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 graphs:
+                from vllm_rlt.kernels.sparse_decode import SELECT_PARTS
+                from vllm_rlt.kernels.triton_attention import NUM_SPLITS
+                from vllm_rlt.worker.cuda_graph import speculative_capture_sizes
+
+                # One capture workspace serves every bucket of the graphs.
+                rows = speculative_capture_sizes(
+                    self.execution_config, scheduler_config, speculative_config
+                )[-1]
+                width = math.ceil(config.max_position_embeddings / cache_config.block_size)
+                head_dim = 1 << (config.head_dim - 1).bit_length()
+                scratch += (
+                    4
+                    * rows
+                    * config.num_attention_heads
+                    * (width * cache_config.block_size + NUM_SPLITS * (2 + head_dim))
+                )
+                # Select scratch per row and head: radix histograms, part counts,
+                # offsets and contribution parts, plus state, lse and ticket.
+                parts = SELECT_PARTS * (4 * 256 + 1 + 2 + 1 + head_dim) * 4
+                scratch += rows * config.num_attention_heads * (parts + 5 * 2 * 8 + 8)
         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/engine/worker.py b/vllm_rlt/engine/worker.py
index 9a2b2805..e04429d8 100644
--- a/vllm_rlt/engine/worker.py
+++ b/vllm_rlt/engine/worker.py
@@ -194,9 +194,16 @@ def run(self):
                 and not self.channel.poll()
             ):
                 time.sleep(0.001)
+        graphs = getattr(getattr(self.engine, "speculative_runner", None), "graphs", None)
+        replays = dict(
+            graph_replays=getattr(graphs, "replays", 0),
+            sparse_graph_replays=getattr(graphs, "sparse_replays", 0),
+        )
         self.engine.close()
         metrics = self.connector.close() if self.connector is not None else {}
-        self.send("stopped", used_blocks=self.engine.cache_manager.num_used_blocks, **metrics)
+        self.send(
+            "stopped", used_blocks=self.engine.cache_manager.num_used_blocks, **replays, **metrics
+        )
 
 
 def worker_main(
diff --git a/vllm_rlt/pd/engine.py b/vllm_rlt/pd/engine.py
index 47f58ad1..3814f8f0 100644
--- a/vllm_rlt/pd/engine.py
+++ b/vllm_rlt/pd/engine.py
@@ -69,6 +69,7 @@ def __init__(
         model_config=None,
         seed=0,
         speculative_config=None,
+        sparse_decode_attention_config=None,
     ):
         self.speculative_config = speculative_config
         model, revision, self.model_config = prepare_model(
@@ -138,6 +139,11 @@ def __init__(
                             else self.exit_config
                         ),
                         speculative_config=speculative_config if role == "decode" else None,
+                        # Reference state is captured and consumed while decoding;
+                        # prefill rows are always dense, so only D owns a store.
+                        sparse_decode_attention_config=(
+                            sparse_decode_attention_config if role == "decode" else None
+                        ),
                         execution_config=execution,
                         attention_backend=attention_backend,
                     )
diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py
index 563fc082..f1bba19b 100644
--- a/vllm_rlt/worker/async_speculative.py
+++ b/vllm_rlt/worker/async_speculative.py
@@ -206,7 +206,15 @@ def __init__(
         self._generations: dict[str, int] = {}
         sizes = speculative_capture_sizes(execution, scheduler, config)
         self.graphs = (
-            DeviceRecurrentGraphs(model, cache, execution, False, capture_sizes=sizes)
+            DeviceRecurrentGraphs(
+                model,
+                cache,
+                execution,
+                False,
+                capture_sizes=sizes,
+                reference_store=reference_store,
+                key_retention=None if sparse_decode is None else sparse_decode.key_retention,
+            )
             if execution.cuda_graphs
             else None
         )
diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py
index 63b34d6d..6d036713 100644
--- a/vllm_rlt/worker/attention_reference.py
+++ b/vllm_rlt/worker/attention_reference.py
@@ -79,6 +79,7 @@ def __init__(
         stream=None,
         initial_key_capacity=None,
         implementation="batched",
+        padding_slot=False,
     ):
         if slots < 1 or layers < 1 or heads < 1 or head_dim < 1:
             raise ValueError("reference store dimensions must be positive")
@@ -106,6 +107,10 @@ def __init__(
         self.stream = stream
         # Which CUDA attention path consumes this store's plans.
         self.implementation = implementation
+        # CUDA graph replays pad a submission to its bucket. Padding rows read
+        # and write this extra slot, which is never bound to a token.
+        self.padding_slot = slots if padding_slot else None
+        slots += bool(padding_slot)
         shape = (slots, layers, heads)
         self.selected = torch.zeros((*shape, initial), dtype=torch.int32, device=device)
         self.counts = torch.zeros(shape, dtype=torch.int32, device=device)
@@ -120,7 +125,7 @@ def __init__(
         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.free = list(range(self.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
@@ -150,7 +155,9 @@ 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
+        index_bytes = (
+            self.selected.shape[0] * 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:
diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py
index 10149516..3b70d141 100644
--- a/vllm_rlt/worker/cuda_graph.py
+++ b/vllm_rlt/worker/cuda_graph.py
@@ -8,6 +8,7 @@
 
 import torch
 
+from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan
 from vllm_rlt.config import cuda_graph_rows
 from vllm_rlt.worker.metadata import DeviceKVMetadata
 
@@ -102,8 +103,10 @@ def _capture(hidden, stream, pool, run, *, startup=False):
 class _DeviceCache:
     """Capture only GPU cache operations, never Python allocation/written state."""
 
-    def __init__(self, cache):
+    def __init__(self, cache, workspace=None):
         self.cache = cache
+        # Sparse capture scratch shared by every bucket; replays run in order.
+        self.workspace = workspace
 
     def _write_prepared(self, layer, batch, k, v):
         cache = self.cache
@@ -121,6 +124,21 @@ def _write_prepared(self, layer, batch, k, v):
 
     def _attend_prepared(self, layer, batch, q):
         cache = self.cache
+        plan = getattr(batch, "reference_plan", None)
+        if plan is not None:
+            from vllm_rlt.kernels.sparse_decode import fused_paged_attention
+
+            return fused_paged_attention(
+                q,
+                cache.key_cache[:, layer],
+                cache.value_cache[:, layer],
+                batch.block_tables,
+                batch.context_lengths,
+                plan,
+                layer=layer,
+                dense_attention=cache.attention,
+                workspace=self.workspace,
+            )
         if batch.cu_seqlens_q is not None:
             return cache.attention.prefill(
                 q,
@@ -219,10 +237,17 @@ def __init__(self, model, cache, execution, compute_gate):
         self.last_event = None
         self.captures = self.replays = self.fallbacks = 0
         self.warned_cache_full = False
+        self.proxy = _DeviceCache(cache)
 
     def select(self, hidden, batch):
         count = _query_count(batch)
         packed = batch.cu_seqlens_q is not None
+        plan = getattr(batch, "reference_plan", None)
+        if plan is not None and not plan.dense():
+            # These entries carry no reference plan; replaying one would
+            # silently run dense attention.
+            self.fallbacks += 1
+            return PreparedGraph(batch)
         if len(batch.position_ids) != count:
             batch = replace(
                 batch,
@@ -287,13 +312,12 @@ def prepare_inputs(self, call, hidden):
                 stream.wait_event(self.last_event)
             _copy_graph_inputs(entry, hidden, batch)
             if entry.graph is None:
-                proxy = _DeviceCache(cache)
                 entry.graph, entry.output = _capture(
                     entry.hidden,
                     self.capture_stream,
                     self.pool,
                     lambda: self.model.recurrent_prepared(
-                        entry.hidden, entry.metadata, proxy, compute_gate=self.compute_gate
+                        entry.hidden, entry.metadata, self.proxy, compute_gate=self.compute_gate
                     ),
                 )
                 self.captures += 1
@@ -328,11 +352,83 @@ class DeviceRecurrentGraphs(RecurrentGraphs):
 
     Allocation identity is checked on every call, outside capture. Entries retain
     only tensor metadata, never request allocations or CPU written-position state.
+
+    With a fused sparse reference store, every bucket is captured once per row
+    mode. A sparse entry owns a static plan whose slot ids are refilled before
+    each replay; the live plan's liveness is checked on the host at that point,
+    since replay runs no Python.
     """
 
-    def __init__(self, model, cache, execution, compute_gate, capture_sizes=None):
+    def __init__(
+        self,
+        model,
+        cache,
+        execution,
+        compute_gate,
+        capture_sizes=None,
+        reference_store=None,
+        key_retention=None,
+    ):
         super().__init__(model, cache, execution, compute_gate)
         self.capture_sizes = capture_sizes or graph_capture_sizes(execution)
+        self.reference_store = reference_store
+        self.key_retention = key_retention
+        self.sparse_replays = 0
+        self.modes = (DENSE,)
+        if reference_store is not None:
+            from vllm_rlt.kernels.sparse_decode import capture_workspace
+
+            if reference_store.padding_slot is None:
+                raise ValueError("sparse CUDA graphs need a reference store padding slot")
+            if reference_store.selected.shape[-1] != (reference_store.key_bucket or 1):
+                # A replay must never see a grown index pool at a new address.
+                raise ValueError("sparse CUDA graphs need the reference index pool at full size")
+            self.modes = (DENSE, CAPTURE, CONSUME)
+            config = model.config
+            width = math.ceil(config.max_position_embeddings / cache.block_size)
+            self.proxy = _DeviceCache(
+                cache,
+                capture_workspace(
+                    cache.device,
+                    self.capture_sizes[-1],
+                    config.num_attention_heads,
+                    width * cache.block_size,
+                    config.head_dim,
+                ),
+            )
+
+    @staticmethod
+    def _key(size, mode):
+        # Dense entries keep their bucket-size keys; sparse ones add the mode.
+        return size if mode == DENSE else (size, mode)
+
+    def _entry(self, hidden, batch, size, mode):
+        entry = _graph_entry(
+            hidden,
+            batch,
+            self.cache,
+            self.model.config.max_position_embeddings,
+            count=size,
+            bucketed=True,
+        )
+        entry.mode = mode
+        if mode != DENSE:
+            store = self.reference_store
+            entry.slot_ids = torch.full(
+                (size,), store.padding_slot, device=self.cache.device, dtype=torch.int64
+            )
+            # Host liveness is checked against the live plan before replay,
+            # so the captured plan names no host slots.
+            entry.metadata.reference_plan = SparseDecodePlan(
+                modes=(mode,) * size,
+                slots=(),
+                incarnations=(),
+                key_retention=self.key_retention,
+                key_bucket=store.key_bucket,
+                slot_ids=entry.slot_ids,
+                store=store,
+            )
+        return entry
 
     @torch.inference_mode()
     def prepare(self):
@@ -347,50 +443,70 @@ def prepare(self):
             block_tables=torch.empty((0, 1), device=self.cache.device, dtype=torch.int32),
             cu_seqlens_q=None,
         )
-        proxy = _DeviceCache(self.cache)
-        for size in self.capture_sizes:
-            entry = _graph_entry(
-                hidden,
-                batch,
-                self.cache,
-                self.model.config.max_position_embeddings,
-                count=size,
-                bucketed=True,
-            )
-            _copy_graph_inputs(entry, hidden, batch, count=0)
-            entry.graph, entry.output = _capture(
-                entry.hidden,
-                self.capture_stream,
-                self.pool,
-                lambda: self.model.recurrent_prepared(
-                    entry.hidden, entry.metadata, proxy, compute_gate=self.compute_gate
-                ),
-                startup=True,
-            )
-            self.entries[size] = entry
-            self.captures += 1
+        for mode in self.modes:
+            for size in self.capture_sizes:
+                entry = self._entry(hidden, batch, size, mode)
+                _copy_graph_inputs(entry, hidden, batch, count=0)
+                entry.graph, entry.output = _capture(
+                    entry.hidden,
+                    self.capture_stream,
+                    self.pool,
+                    lambda entry=entry: self.model.recurrent_prepared(
+                        entry.hidden, entry.metadata, self.proxy, compute_gate=self.compute_gate
+                    ),
+                    startup=True,
+                )
+                self.entries[self._key(size, mode)] = entry
+                self.captures += 1
+
+    def _mode(self, batch):
+        plan = getattr(batch, "reference_plan", None)
+        if plan is None or plan.dense():
+            return DENSE
+        modes = set(plan.modes)
+        if self.reference_store is None or plan.store is not self.reference_store:
+            return None
+        if len(modes) != 1 or plan.key_retention != self.key_retention:
+            return None
+        return plan.modes[0]
 
     def select(self, hidden, batch):
         self.cache._require_live_batch(batch)
         count = _query_count(batch)
-        if not count or count > self.capture_sizes[-1]:
+        mode = self._mode(batch)
+        if not count or count > self.capture_sizes[-1] or mode is None:
             self.fallbacks += 1
             return PreparedGraph(batch)
         capacity = _bucket_size(self.capture_sizes, count)
-        entry = self.entries.get(capacity)
+        key = self._key(capacity, mode)
+        entry = self.entries.get(key)
         if entry is None:
-            entry = _graph_entry(
-                hidden,
-                batch,
-                self.cache,
-                self.model.config.max_position_embeddings,
-                count=capacity,
-                bucketed=True,
-            )
+            entry = self._entry(hidden, batch, capacity, mode)
             entry.graph = None
-            self.entries[capacity] = entry
+            self.entries[key] = entry
         return PreparedGraph(batch, entry)
 
+    @torch.inference_mode()
+    def prepare_inputs(self, call, hidden):
+        entry = call.entry
+        if entry is not None and entry.mode != DENSE:
+            plan = call.metadata.reference_plan
+            self.reference_store._require_live(plan)
+            count = _query_count(call.metadata)
+            stream = torch.cuda.current_stream(self.cache.device)
+            if self.last_event is not None:
+                stream.wait_event(self.last_event)
+            entry.slot_ids[count:].fill_(self.reference_store.padding_slot)
+            entry.slot_ids[:count].copy_(plan.slot_ids[:count])
+        super().prepare_inputs(call, hidden)
+
+    @torch.inference_mode()
+    def execute(self, call, hidden):
+        sparse = call.entry is not None and call.entry.mode != DENSE
+        output = super().execute(call, hidden)
+        self.sparse_replays += sparse
+        return output
+
 
 class CodaGraphs:
     """Select and stage coda inputs before replay; commit KV after submission."""

From bc0622f4817ff2ffb6171fe0a838a602c3b00c86 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Sat, 31 Oct 2026 23:07:01 +0100
Subject: [PATCH 25/32] Allow sparse decode attention with asynchronous
 ordinary decode

Ordinary decode builds its reference plans from host-known depths and
positions, and every reference read, write and retirement is ordered on the
core stream, so submitting a round before collecting the previous one needs no
change to the lifecycle; the pool stays one slot per scheduled request.

The runner now publishes each capture row's retained count when it builds the
plan, since the row's causal length is known there. A batched consume after a
fused capture then sizes its gather without reading the device count back.

Tests: asynchronous and synchronous ordinary decode produce the same tokens
(batched on CPU; batched and fused on CUDA, where every consume is also checked
to be sized from a published count equal to the device count); the pool never
exceeds one live reference per scheduled request, including abort and id reuse.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_sparse_decode_sync.py       | 94 ++++++++++++++++++++++++++
 vllm_rlt/engine/stage_engine.py        |  4 --
 vllm_rlt/worker/attention_reference.py |  9 +++
 vllm_rlt/worker/model_runner.py        | 13 +++-
 4 files changed, 114 insertions(+), 6 deletions(-)

diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py
index cccbbc7f..e22f8695 100644
--- a/tests/test_sparse_decode_sync.py
+++ b/tests/test_sparse_decode_sync.py
@@ -9,6 +9,7 @@
 from vllm_rlt import (
     CacheConfig,
     ExecutionConfig,
+    ExitConfig,
     SamplingParams,
     SchedulerConfig,
     SparseDecodeAttentionConfig,
@@ -69,6 +70,99 @@ def test_full_retention_matches_dense_and_reclaims_slots(draft):
         dense.close()
 
 
+def _ordinary(asynchronous, retention, device="cpu", dtype=torch.float32, **options):
+    return engine(
+        retention=retention,
+        device=device,
+        dtype=dtype,
+        exit_config=ExitConfig("ouro_delayed"),
+        execution_config=ExecutionConfig(
+            async_scheduling=asynchronous, static_buffers=True, pad_to_power_of_two=True
+        ),
+        **options,
+    )
+
+
+def _bounded_pool(e, monkeypatch):
+    """Fail if live references ever exceed one per scheduled request."""
+    store = e.reference_store
+    reserve = store.reserve
+
+    def checked(*args, **kwargs):
+        lease = reserve(*args, **kwargs)
+        assert store.outstanding <= e.scheduler.config.max_num_seqs
+        return lease
+
+    monkeypatch.setattr(store, "reserve", checked)
+
+
+@pytest.mark.parametrize("retention", [0.5, 1.0])
+def test_async_ordinary_decode_matches_synchronous(retention, monkeypatch):
+    sync = _ordinary(False, retention)
+    asynchronous = _ordinary(True, retention)
+    _bounded_pool(asynchronous, monkeypatch)
+    try:
+        # Staggered prompt lengths keep requests at different depths per batch.
+        assert generate(asynchronous, count=4) == generate(sync, count=4)
+    finally:
+        sync.close()
+        asynchronous.close()
+
+
+def test_async_ordinary_decode_reclaims_after_abort_and_reuse(monkeypatch):
+    e = _ordinary(True, 0.5)
+    _bounded_pool(e, monkeypatch)
+    try:
+        e.add_request("a", [2, 3, 4], SamplingParams(max_tokens=12, ignore_eos=True))
+        for _ in range(100):
+            e.step()
+            if e.reference_store.outstanding:
+                break
+        assert e.reference_store.outstanding == 1
+        e.abort_request("a")
+        for _ in range(100):
+            if not e.has_unfinished_requests():
+                break
+            e.step()
+        assert e.reference_store.outstanding == 0
+        # The same id restarts with a fresh incarnation.
+        assert generate(e, count=2)
+    finally:
+        e.close()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("implementation", ["batched", "fused"])
+def test_cuda_async_ordinary_decode_matches_synchronous(implementation, monkeypatch):
+    sparse = SparseDecodeAttentionConfig(
+        enabled=True, key_retention=0.5, implementation=implementation
+    )
+    options = dict(device="cuda", dtype=torch.bfloat16, sparse_decode_attention_config=sparse)
+    sync = _ordinary(False, 0.5, **options)
+    asynchronous = _ordinary(True, 0.5, **options)
+    _bounded_pool(asynchronous, monkeypatch)
+    store = asynchronous.reference_store
+    host_count = store.host_count
+    published = []
+
+    def checked(slot, layer):
+        # Every consume is sized from a count the host published at capture.
+        count = host_count(slot, layer)
+        assert count is not None
+        # ...and it is the count the capture actually wrote.
+        assert int(store.counts[slot, layer].max()) == count
+        published.append((slot, layer, count))
+        return count
+
+    monkeypatch.setattr(store, "host_count", checked)
+    try:
+        assert generate(asynchronous, count=4) == generate(sync, count=4)
+        assert published
+    finally:
+        sync.close()
+        asynchronous.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)
diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py
index 5bb2b23f..ff483819 100644
--- a/vllm_rlt/engine/stage_engine.py
+++ b/vllm_rlt/engine/stage_engine.py
@@ -150,10 +150,6 @@ def __init__(
                 raise ValueError(
                     "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"
-                )
             # 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
diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py
index 6d036713..833f20f2 100644
--- a/vllm_rlt/worker/attention_reference.py
+++ b/vllm_rlt/worker/attention_reference.py
@@ -321,6 +321,15 @@ def note_count(self, slot: int, layer: int, count: int) -> None:
         """
         self._counts_host[(slot, layer)] = count
 
+    def note_counts(self, slot: int, count: int) -> None:
+        """Record one capture row's retained count for every layer on the host.
+
+        The count follows the row's causal length alone, so a caller that
+        knows the length can publish it before any layer runs.
+        """
+        for layer in range(self.layers):
+            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:
diff --git a/vllm_rlt/worker/model_runner.py b/vllm_rlt/worker/model_runner.py
index 77144bf4..d82af179 100644
--- a/vllm_rlt/worker/model_runner.py
+++ b/vllm_rlt/worker/model_runner.py
@@ -7,7 +7,7 @@
 
 import torch
 
-from vllm_rlt.attention.sparse import CAPTURE, DENSE, reference_mode
+from vllm_rlt.attention.sparse import CAPTURE, DENSE, reference_mode, retained_count
 from vllm_rlt.config import ExecutionConfig, ExitConfig, SchedulerConfig
 from vllm_rlt.core.scheduler import SchedulerOutput
 from vllm_rlt.request import ExecutionBatch, Operation, Request
@@ -1072,7 +1072,7 @@ def _reference_metadata(self, batch, kv):
         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 = [], []
+        owners, modes, captured = [], [], []
         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)
@@ -1089,6 +1089,8 @@ def _reference_metadata(self, batch, kv):
                     self.attention_references[request.request_id] = state
             if state is None:
                 mode = DENSE  # No reference loop for this token: stay dense.
+            elif mode == CAPTURE:
+                captured.append((state[2].slots[0], position + 1))
             owners.append(None if mode == DENSE else (state[2], 0))
             modes.append(mode)
         if all(mode == DENSE for mode in modes):
@@ -1100,6 +1102,13 @@ def _reference_metadata(self, batch, kv):
             owners, modes, key_retention=self.sparse_decode.key_retention
         )
         self.reference_store.bind_positions(plan, kv.position_ids)
+        # A capture row's retained count follows its causal length, which the
+        # host knows here. Publishing it lets a later batched consume size its
+        # gather without reading a fused capture's device count back.
+        for slot, length in captured:
+            self.reference_store.note_counts(
+                slot, retained_count(length, self.sparse_decode.key_retention)
+            )
         return replace(kv, reference_plan=plan)
 
     def _load_hidden(self, batch, prepared, workspace, size):

From 9cb7407b868b813af226a8d490b799c3225ec380 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Sat, 31 Oct 2026 23:50:31 +0100
Subject: [PATCH 26/32] Run sparse decode attention in asynchronous token
 pipelines

Token pipelines advance, shift and restart their rows on the GPU, so the host
knows neither a row's loop nor which reference slot it owns. They now carry
both on the device:

- DeviceSparseDecodePlan holds one int8 mode and one slot id per row as device
  tensors, plus the owning leases, which the store checks live on the host.
  Only the fused kernels read it.
- The fused kernels take a per-row mode pointer for such plans: the split and
  selection kernels compute dense and capture rows (consume rows do no dense
  work, dense rows keep no scores), the consume kernels compute consume rows,
  and each writes only its own rows. Dense and capture rows stay the dense
  kernel's output; uniform host plans compile without the mode branches.
- Each pipeline request reserves one slot per pipeline row for its lifetime.
  The slots form a device permutation moved with the rows on every emission, so
  a row consumes the slot its loop-1 capture wrote; the request's release
  returns them, and buffer discard drops them from surviving states.
- Each recurrent offset derives its modes from the rows' loops (1 captures,
  later loops consume, inactive rows stay dense). With CUDA Graphs, pipeline
  runners capture device-plan entries whose slot ids and modes are refilled
  before replay.
- The engine accepts token pipelines with sparse decode only for asynchronous
  scheduling, the fused implementation and key_retention < 1.

GPU tests (to run on CUDA): one mixed-mode launch equals the uniform launches
row for row and slot for slot; pipeline sparse tokens equal block sparse tokens
(dense pipeline equals dense block as the precondition), with and without
graphs and for draft depths 1 and 2; slots follow requests through abort and id
reuse; pipelines reject the batched implementation and full retention.
docs/runtime.md lists the supported sparse modes.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 docs/runtime.md                        |  50 ++++--
 tests/test_sparse_decode_pipeline.py   | 221 +++++++++++++++++++++++
 vllm_rlt/attention/sparse.py           |  41 ++++-
 vllm_rlt/engine/stage_engine.py        |  13 +-
 vllm_rlt/kernels/sparse_decode.py      | 239 +++++++++++++++++--------
 vllm_rlt/worker/async_speculative.py   |  56 +++++-
 vllm_rlt/worker/attention_reference.py |  36 +++-
 vllm_rlt/worker/cuda_graph.py          |  45 ++++-
 8 files changed, 595 insertions(+), 106 deletions(-)
 create mode 100644 tests/test_sparse_decode_pipeline.py

diff --git a/docs/runtime.md b/docs/runtime.md
index 1853d5c4..3869d9f0 100644
--- a/docs/runtime.md
+++ b/docs/runtime.md
@@ -245,24 +245,43 @@ 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:
+# Ordinary decode uses no speculative_config, with synchronous or
+# asynchronous scheduling. For block speculation, also pass:
 speculation = SpeculativeConfig(num_speculative_tokens=2, draft_loops=2)
 execution = ExecutionConfig(async_scheduling=False, cuda_graphs=False)
+# Token pipelines, and CUDA Graphs with speculation, need the fused kernels:
+fused = SparseDecodeAttentionConfig(enabled=True, key_retention=0.1, implementation="fused")
+pipeline = SpeculativeConfig(num_speculative_tokens=3, draft_loops=2, mode="pipeline")
+execution = ExecutionConfig(async_scheduling=True, cuda_graphs=True)
 ```
 
-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. 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.
+This path runs with LAST_EXITED KV and refill scheduling in these modes:
+
+| Mode | Implementations |
+| --- | --- |
+| Ordinary decode, synchronous or asynchronous scheduling | `batched`, `fused` |
+| Block self-speculation, synchronous or asynchronous | `batched`, `fused` |
+| Asynchronous block self-speculation with CUDA Graphs | `fused`, `key_retention < 1` |
+| Asynchronous token pipelines, with or without CUDA Graphs | `fused`, `key_retention < 1` |
+
+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. Dense and capture rows preserve the
+selected backend's output and original packed batch geometry. In the default
+`batched` implementation, selection, retained mass and contributions, and
+partial-retention correction use the shared PyTorch reference implementation;
+it has host readbacks and per-row launch costs. The opt-in `fused`
+implementation (`vllm_rlt/kernels/sparse_decode.py`) keeps lengths and counts
+on the device and launches a fixed set of Triton kernels per layer, so it can
+be replayed inside CUDA Graphs; its consume rows differ from the reference in
+FP32 reduction order. The captured output is the value actually entering the
+recurrence, stored in FP32 after the backend's output-dtype conversion.
+Enabling sparse decode does not by itself promise faster serving. Token
+pipelines decide each row's depth on the device, so their plans carry
+per-row modes and slots on the device: every request owns one reference slot
+per pipeline row for its lifetime, and the slots move with the rows. Mixed
+scheduling, preemption/recovery, the synchronous token pipeline and SHARED
+KV 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,
@@ -270,7 +289,8 @@ 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
+block request, two rounds of K+1 per asynchronous block request, and one per
+pipeline row (at most K+1) per token-pipeline request. All
 reference payload bytes contribute to automatic KV memory planning; explicit
 KV block overrides still require the caller to budget total peak memory.
 
diff --git a/tests/test_sparse_decode_pipeline.py b/tests/test_sparse_decode_pipeline.py
new file mode 100644
index 00000000..6072e244
--- /dev/null
+++ b/tests/test_sparse_decode_pipeline.py
@@ -0,0 +1,221 @@
+"""Sparse decode attention with token pipelines: device-decided row modes and slots."""
+
+import pytest
+import torch
+
+from tests.helpers import tiny_ouro_config
+from tests.test_sparse_decode_attention import drain, make_store
+from tests.test_sparse_decode_fused import _paged
+from vllm_rlt import (
+    CacheConfig,
+    ExecutionConfig,
+    SamplingParams,
+    SchedulerConfig,
+    SparseDecodeAttentionConfig,
+    SpeculativeConfig,
+)
+from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE
+
+PROMPTS = [[2, 3, 4, 5, 6, 7, 8], [9, 10, 11], [12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22]]
+LENGTHS = (23, 9, 17)
+
+
+@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
+def test_mixed_modes_match_uniform_fused_kernels():
+    """One launch set with every mode equals the uniform launches row for row."""
+    from vllm_rlt.kernels import sparse_decode
+    from vllm_rlt.kernels.paged_attention import TritonPagedAttention
+
+    rows, heads, kv_heads, dim, block_size, pages = 6, 4, 2, 64, 16, 9
+    dense = TritonPagedAttention()
+    q, keys, values, tables = _paged(
+        rows, heads, kv_heads, dim, pages, block_size, torch.float32, 7
+    )
+    q2, _, _, _ = _paged(rows, heads, kv_heads, dim, pages, block_size, torch.float32, 8)
+    lengths = torch.tensor([5, 33, 70, 129, 140, 1], device="cuda", dtype=torch.int32)
+    modes = [DENSE, CAPTURE, CONSUME, CAPTURE, CONSUME, DENSE]
+    stores, leases = [], []
+    for _ in range(2):
+        store = make_store(
+            slots=rows,
+            layers=2,
+            heads=heads,
+            head_dim=dim,
+            key_bucket=pages * block_size,
+            device="cuda",
+            implementation="fused",
+        )
+        lease = store.reserve("r", 0, rows, generation=1)
+        # The consume rows need a reference from an earlier loop.
+        earlier = [row for row, mode in enumerate(modes) if mode == CONSUME]
+        plan = store.build_plan(
+            [(lease, row) for row in earlier], [CAPTURE] * len(earlier), key_retention=0.25
+        )
+        sparse_decode.fused_paged_attention(
+            q[earlier],
+            keys,
+            values,
+            tables[earlier],
+            lengths[earlier],
+            plan,
+            layer=1,
+            dense_attention=dense,
+        )
+        stores.append(store)
+        leases.append(lease)
+    mixed_store, uniform_store = stores
+    plan = mixed_store.build_device_plan(
+        [leases[0]],
+        leases[0].slot_ids,
+        torch.tensor(modes, device="cuda", dtype=torch.int8),
+        key_retention=0.25,
+    )
+    mixed = sparse_decode.fused_paged_attention(
+        q2, keys, values, tables, lengths, plan, layer=1, dense_attention=dense
+    )
+    expected = torch.empty_like(q2)
+    for mode in (DENSE, CAPTURE, CONSUME):
+        subset = [row for row, m in enumerate(modes) if m == mode]
+        if mode == DENSE:
+            expected[subset] = dense(q2[subset], keys, values, tables[subset], lengths[subset])
+            continue
+        host = uniform_store.build_plan(
+            [(leases[1], row) for row in subset], [mode] * len(subset), key_retention=0.25
+        )
+        expected[subset] = sparse_decode.fused_paged_attention(
+            q2[subset],
+            keys,
+            values,
+            tables[subset],
+            lengths[subset],
+            host,
+            layer=1,
+            dense_attention=dense if mode == CAPTURE else None,
+        )
+    torch.cuda.synchronize()
+    # Dense and capture rows are the dense kernel's output; consume rows the
+    # uniform consume's; every store field matches slot for slot.
+    assert torch.equal(mixed, expected)
+    for name in ("selected", "counts", "mass", "output", "contribution", "valid", "uses"):
+        assert torch.equal(getattr(mixed_store, name), getattr(uniform_store, name)), name
+
+
+def _engine(mode, sparse, *, draft_loops=2, graphs=False, retention=0.5):
+    from vllm_rlt.engine.stage_engine import StageEngine
+    from vllm_rlt.models import OuroForCausalLM
+
+    torch.manual_seed(123)
+    model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.float32)
+    return StageEngine(
+        model,
+        attention_backend="triton",
+        cache_config=CacheConfig(256, 16, incremental_allocation=True),
+        scheduler_config=SchedulerConfig(
+            max_num_seqs=3, max_num_batched_tokens=16, prefill_chunk_size=8
+        ),
+        execution_config=ExecutionConfig(
+            async_scheduling=True,
+            static_buffers=graphs,
+            pad_to_power_of_two=graphs,
+            cuda_graphs=graphs,
+        ),
+        speculative_config=SpeculativeConfig(3, draft_loops=draft_loops, mode=mode),
+        sparse_decode_attention_config=(
+            SparseDecodeAttentionConfig(
+                enabled=True, key_retention=retention, implementation="fused"
+            )
+            if sparse
+            else None
+        ),
+    )
+
+
+def _generate(e):
+    for i, (prompt, length) in enumerate(zip(PROMPTS, LENGTHS)):
+        e.add_request(str(i), prompt, SamplingParams(max_tokens=length, ignore_eos=True))
+    outputs = drain(e)
+    assert sorted(outputs) == ["0", "1", "2"]
+    assert e.cache_manager.num_used_blocks == 0
+    if e.reference_store is not None:
+        assert e.reference_store.outstanding == 0
+    return {rid: list(output.token_ids) for rid, output in outputs.items()}
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("draft_loops", [1, 2])
+@pytest.mark.parametrize("graphs", [False, True])
+def test_pipeline_sparse_matches_block_sparse(draft_loops, graphs):
+    """Both capture at loop 1 and consume at loops 2 and 3 of every token."""
+    results = {}
+    for mode in ("block", "pipeline"):
+        for sparse in (False, True):
+            e = _engine(mode, sparse, draft_loops=draft_loops, graphs=graphs)
+            try:
+                results[mode, sparse] = _generate(e)
+                if sparse and mode == "pipeline" and graphs:
+                    runner = e.speculative_runner
+                    assert runner.graphs.sparse_replays > 0
+            finally:
+                e.close()
+    # Precondition: the dense pipeline already reproduces block speculation.
+    assert results["pipeline", False] == results["block", False]
+    assert results["pipeline", True] == results["block", True]
+
+
+@pytest.mark.gpu
+def test_pipeline_sparse_slots_follow_requests_through_abort_and_reuse():
+    e = _engine("pipeline", True)
+    store = e.reference_store
+    width = min(4, 4 // 2)
+    try:
+        e.add_request("a", PROMPTS[0], SamplingParams(max_tokens=40, ignore_eos=True))
+        for _ in range(200):
+            e.step()
+            state = e.speculative_runner.states.get("a")
+            if state is not None and state.reference is not None:
+                break
+        else:
+            pytest.fail("pipeline state was not created")
+        # One lease of `width` slots for the request's lifetime.
+        assert store.outstanding == width
+        assert sorted(state.reference.slots) == sorted(state.pipeline.slots.tolist())
+        e.abort_request("a")
+        drain(e)
+        assert store.outstanding == 0
+        # A reused id gets a new generation and fresh slots.
+        assert _generate(e)
+    finally:
+        e.close()
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "sparse",
+    [
+        SparseDecodeAttentionConfig(enabled=True, key_retention=0.5),
+        SparseDecodeAttentionConfig(enabled=True, key_retention=1.0, implementation="fused"),
+    ],
+)
+def test_pipeline_rejects_sparse_paths_without_device_plans(sparse):
+    from vllm_rlt.engine.stage_engine import StageEngine
+    from vllm_rlt.models import OuroForCausalLM
+
+    model = OuroForCausalLM(tiny_ouro_config(head_dim=64)).to("cuda", torch.float32)
+    with pytest.raises(ValueError, match="token pipelines"):
+        StageEngine(
+            model,
+            attention_backend="triton",
+            execution_config=ExecutionConfig(async_scheduling=True),
+            speculative_config=SpeculativeConfig(3, mode="pipeline"),
+            sparse_decode_attention_config=sparse,
+        )
diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py
index c438a918..4fb5a981 100644
--- a/vllm_rlt/attention/sparse.py
+++ b/vllm_rlt/attention/sparse.py
@@ -101,6 +101,29 @@ def dense(self) -> bool:
         return all(mode == DENSE for mode in self.modes)
 
 
+@dataclass(frozen=True)
+class DeviceSparseDecodePlan:
+    """A reference plan whose row modes and slots are decided on the device.
+
+    Token pipelines advance, shift and restart their rows on the GPU, so the
+    host knows neither a row's depth nor which of its request's slots it owns.
+    ``modes`` and ``slot_ids`` are device tensors with one entry per row;
+    ``leases`` own every slot a row can name and are checked live on the host.
+    Only the fused kernels read such a plan.
+    """
+
+    modes: torch.Tensor
+    slot_ids: torch.Tensor
+    leases: tuple
+    key_retention: float
+    key_bucket: int | None
+    store: object | None = None
+
+    def dense(self) -> bool:
+        """Unknown on the host: the device may route any row to sparse work."""
+        return False
+
+
 def retained_count(length: int, key_retention: float) -> int:
     """How many keys a query with ``length`` legal keys retains.
 
@@ -344,11 +367,27 @@ def sparse_paged_attention(
     """Use bounded row batches on CUDA, retaining the scalar numerical oracle.
 
     A store configured for the ``fused`` implementation routes uniform capture
-    or consume submissions to the fused kernels. Other CUDA submissions keep
+    or consume submissions, and every device plan, 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 isinstance(plan, DeviceSparseDecodePlan):
+        from vllm_rlt.kernels import sparse_decode
+
+        if not sparse_decode.supports(plan, q, dense_attention):
+            raise ValueError("a device sparse plan needs the fused CUDA kernels")
+        return sparse_decode.fused_paged_attention(
+            q,
+            key_cache,
+            value_cache,
+            block_tables,
+            context_lengths,
+            plan,
+            layer=layer,
+            dense_attention=dense_attention,
+        )
     if not q.is_cuda:
         return reference_paged_attention(
             q,
diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py
index ff483819..ece18800 100644
--- a/vllm_rlt/engine/stage_engine.py
+++ b/vllm_rlt/engine/stage_engine.py
@@ -144,8 +144,17 @@ def __init__(
                     "sparse decode attention with CUDA graphs requires async block "
                     "speculation, the fused implementation and key_retention < 1"
                 )
-            if speculative_config is not None and speculative_config.mode != "block":
-                raise ValueError("sparse decode attention currently requires block speculation")
+            if self.pipeline_speculative and (
+                not self.async_speculative
+                or sparse_decode.implementation != "fused"
+                or sparse_decode.key_retention >= 1.0
+            ):
+                # Pipeline rows move and restart on the device; only the fused
+                # kernels read their device-decided modes and slots.
+                raise ValueError(
+                    "sparse decode attention with token pipelines requires asynchronous "
+                    "scheduling, the fused implementation and key_retention < 1"
+                )
             if scheduler_config.enable_preemption or scheduler_config.mode != "refill":
                 raise ValueError(
                     "sparse decode attention requires refill scheduling without preemption"
diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py
index 209a4730..7a2854f3 100644
--- a/vllm_rlt/kernels/sparse_decode.py
+++ b/vllm_rlt/kernels/sparse_decode.py
@@ -24,10 +24,14 @@
 import triton
 import triton.language as tl
 
-from vllm_rlt.attention.sparse import CAPTURE, CONSUME, retained_count
+from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DeviceSparseDecodePlan, retained_count
 from vllm_rlt.kernels.paged_attention import TritonPagedAttention, triton_paged_attention
 from vllm_rlt.kernels.triton_attention import NUM_SPLITS
 
+# Row modes as the kernels read them from a device plan.
+CAPTURE_MODE = tl.constexpr(CAPTURE)
+CONSUME_MODE = tl.constexpr(CONSUME)
+
 
 @triton.jit
 def _order_key(scores):
@@ -153,6 +157,7 @@ def _capture_split_kernel(
     RETENTION,
     STATE,
     TICKETS,
+    MODES,
     top,
     score_capacity,
     q_batch_stride: tl.constexpr,
@@ -177,12 +182,17 @@ def _capture_split_kernel(
     SPLITS: tl.constexpr,
     HISTOGRAM: tl.constexpr,
     PARTS: tl.constexpr,
+    MIXED: 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)
+    if MIXED:
+        # Per-row modes: consume rows do no dense work, dense rows keep no scores.
+        mode = tl.load(MODES + row)
+        length = tl.where(mode == CONSUME_MODE, 0, length)
     # 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
@@ -220,10 +230,13 @@ def _capture_split_kernel(
         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.
+        keep = valid_tokens
+        if MIXED:
+            keep = keep & (mode == CAPTURE_MODE)
         tl.store(
             SCORES + (row * NUM_HEADS + head).to(tl.int64) * score_capacity + positions,
             scores,
-            mask=valid_tokens,
+            mask=keep,
         )
         if HISTOGRAM:
             # Selection pass 0 of this part: the top eight bits of the order key.
@@ -251,7 +264,7 @@ def _capture_split_kernel(
     tl.store(PART_O + part * BLOCK_D + dims, accumulator)
     if HISTOGRAM:
         rh = row * NUM_HEADS + head
-        _, count = _retained(LENGTHS, row, RETENTION, top)
+        _, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED)
         if count > 0:
             bins = tl.arange(0, 512)
             tl.store(SEGHIST + (rh * PARTS + split) * 4 * 256 + bins, local, mask=bins < 256)
@@ -410,6 +423,15 @@ def _retained(LENGTHS, row, RETENTION, top):
     return length, tl.minimum(count.to(tl.int32), top)
 
 
+@triton.jit
+def _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED: tl.constexpr):
+    """``_retained`` for a capture row; a mixed submission's other rows retain nothing."""
+    length, count = _retained(LENGTHS, row, RETENTION, top)
+    if MIXED:
+        count = tl.where(tl.load(MODES + row) == CAPTURE_MODE, count, 0)
+    return length, count
+
+
 @triton.jit
 def _segment(length, SEGMENTS: tl.constexpr, ALIGN: tl.constexpr):
     """This program's part ``[first, last)`` of a row's context."""
@@ -483,6 +505,7 @@ def _select_pass_kernel(
     TICKETS,
     GREATER,
     OFFSETS,
+    MODES,
     layer,
     top,
     score_capacity,
@@ -499,15 +522,23 @@ def _select_pass_kernel(
     PARTS: tl.constexpr,
     ALIGN: tl.constexpr,
     SPLITS: tl.constexpr,
+    MIXED: 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)
+    length, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED)
     if MERGE:
-        if segment == 0:
+        merging = segment == 0
+        marking = True
+        if MIXED:
+            # Consume rows get their output from the consume kernels.
+            mode = tl.load(MODES + row)
+            merging = merging & (mode != CONSUME_MODE)
+            marking = mode == CAPTURE_MODE
+        if merging:
             # Identical to triton_attention._paged_attention_merge_kernel.
             parts = rh * SPLITS + tl.arange(0, SPLITS)
             dims = tl.arange(0, BLOCK_D)
@@ -522,7 +553,8 @@ def _select_pass_kernel(
             tl.store(out, dense, mask=dims < HEAD_DIM)
             slot = tl.load(SLOTS + row)
             if head == 0:
-                tl.store(VALID + slot * NUM_LAYERS + layer, 1)
+                if marking:
+                    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
@@ -589,6 +621,7 @@ def _select_write_kernel(
     STATE,
     OFFSETS,
     SELECTED,
+    MODES,
     layer,
     top,
     score_capacity,
@@ -598,13 +631,14 @@ def _select_write_kernel(
     BLOCK_T: tl.constexpr,
     PARTS: tl.constexpr,
     ALIGN: tl.constexpr,
+    MIXED: 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)
+    length, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED)
     if count > 0:
         low = tl.load(STATE + (rh * 5 + 4) * 2)
         ties_needed = tl.load(STATE + (rh * 5 + 4) * 2 + 1).to(tl.int32)
@@ -646,6 +680,7 @@ def _select_contribution_kernel(
     TICKETS,
     MASS,
     CONTRIBUTION,
+    MODES,
     layer,
     top,
     score_capacity,
@@ -662,6 +697,7 @@ def _select_contribution_kernel(
     BLOCK_D: tl.constexpr,
     BLOCK_T: tl.constexpr,
     PARTS: tl.constexpr,
+    MIXED: tl.constexpr,
 ):
     """``(m, C_ref)`` over the retained ranks, in rank chunks summed in a fixed order.
 
@@ -672,7 +708,7 @@ def _select_contribution_kernel(
     head = tl.program_id(1)
     part = tl.program_id(2)
     rh = row * NUM_HEADS + head
-    length, count = _retained(LENGTHS, row, RETENTION, top)
+    length, count = _capture_count(LENGTHS, row, RETENTION, top, MODES, MIXED)
     if count > 0:
         chunk = tl.cdiv(tl.cdiv(count, PARTS), BLOCK_T) * BLOCK_T
         first = part * chunk
@@ -861,6 +897,7 @@ def _consume_split_kernel(
     PART_M,
     PART_L,
     PART_O,
+    MODES,
     layer,
     table_stride,
     retained_capacity,
@@ -881,6 +918,7 @@ def _consume_split_kernel(
     BLOCK_D: tl.constexpr,
     BLOCK_T: tl.constexpr,
     SPLITS: tl.constexpr,
+    MIXED: tl.constexpr,
 ):
     """Online-softmax parts of ``C_new`` over one contiguous chunk of retained ranks.
 
@@ -895,7 +933,10 @@ def _consume_split_kernel(
     maximum = -float("inf")
     normalizer = 0.0
     accumulator = tl.zeros((BLOCK_D,), dtype=tl.float32)
-    if tl.load(LENGTHS + row) > 0:
+    reads = tl.load(LENGTHS + row) > 0
+    if MIXED:
+        reads = reads & (tl.load(MODES + row) == CONSUME_MODE)
+    if reads:
         slot = tl.load(SLOTS + row)
         kv_head = head // HEAD_GROUPS
         state = (slot * NUM_LAYERS + layer) * NUM_HEADS + head
@@ -957,6 +998,7 @@ def _consume_merge_kernel(
     CONTRIBUTION,
     USES,
     OUT,
+    MODES,
     layer,
     q_row_stride: tl.constexpr,
     q_head_stride: tl.constexpr,
@@ -965,6 +1007,7 @@ def _consume_merge_kernel(
     HEAD_DIM: tl.constexpr,
     BLOCK_D: tl.constexpr,
     SPLITS: tl.constexpr,
+    MIXED: tl.constexpr,
 ):
     """``O = O_ref - C_ref + m * C_new``, merging the split parts in a fixed order.
 
@@ -974,28 +1017,34 @@ def _consume_merge_kernel(
     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)
+    # A mixed submission writes only its consume rows here.
+    consuming = True
+    if MIXED:
+        consuming = tl.load(MODES + row) == CONSUME_MODE
+    if consuming:
+        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.
+            mass = tl.load(MASS + state)
+            corrected = reference - contribution + mass * (accumulator / normalizer)
+            tl.store(out, corrected, mask=dims < HEAD_DIM)
 
 
 _RETENTION: dict[tuple[torch.device, float], torch.Tensor] = {}
@@ -1074,6 +1123,14 @@ def _consume_workspace(device, rows, heads, block_d):
 
 def supports(plan, q, dense_attention) -> bool:
     """Whether the fused path covers this submission; otherwise use the batched path."""
+    if isinstance(plan, DeviceSparseDecodePlan):
+        # Dense and capture rows run the split kernel; it needs the dense backend.
+        return (
+            q.is_cuda
+            and q.shape[-1] <= 256
+            and plan.key_retention < 1.0
+            and dense_attention is not None
+        )
     modes = set(plan.modes)
     return (
         q.is_cuda
@@ -1104,7 +1161,12 @@ def fused_paged_attention(
     dense_attention,
     workspace=None,
 ):
-    """Capture or consume one uniform-mode submission without host readbacks.
+    """Capture or consume one submission without host readbacks.
+
+    A host plan has one mode for every row. A device plan
+    (``DeviceSparseDecodePlan``) carries each row's mode on the device: the
+    capture kernels then compute dense and capture rows, the consume kernels
+    consume rows, and each writes only its own rows of the output.
 
     The kernels also write the store's ``valid``/``uses`` marks, so the host
     only checks the plan's liveness instead of launching the bookkeeping.
@@ -1126,7 +1188,10 @@ def fused_paged_attention(
     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:
+    mixed = isinstance(plan, DeviceSparseDecodePlan)
+    # Uniform plans never read the mode pointer; any device tensor stands in.
+    modes = plan.modes if mixed else context_lengths
+    if mixed or 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
@@ -1135,6 +1200,8 @@ def fused_paged_attention(
         # 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 mixed and not split:
+            raise ValueError("a device sparse plan needs the Triton dense attention backend")
         if split:
             output = torch.empty_like(q)
         else:
@@ -1183,6 +1250,7 @@ def fused_paged_attention(
                 retention,
                 state,
                 tickets,
+                modes,
                 top,
                 capacity,
                 *q.stride(),
@@ -1199,6 +1267,7 @@ def fused_paged_attention(
                 SPLITS=NUM_SPLITS,
                 HISTOGRAM=folded,
                 PARTS=SELECT_PARTS,
+                MIXED=mixed,
                 num_warps=4,
             )
             first_pass = 1 if folded else 0
@@ -1223,6 +1292,7 @@ def fused_paged_attention(
                     tickets,
                     greater,
                     offsets,
+                    modes,
                     layer,
                     top,
                     capacity,
@@ -1237,6 +1307,7 @@ def fused_paged_attention(
                     PARTS=SELECT_PARTS,
                     ALIGN=_SEGMENT_ALIGN,
                     SPLITS=NUM_SPLITS,
+                    MIXED=mixed,
                     num_warps=_SCAN_WARPS,
                 )
             _launch(
@@ -1249,6 +1320,7 @@ def fused_paged_attention(
                 state,
                 offsets,
                 selected,
+                modes,
                 layer,
                 top,
                 capacity,
@@ -1258,6 +1330,7 @@ def fused_paged_attention(
                 BLOCK_T=_SCAN_BLOCK,
                 PARTS=SELECT_PARTS,
                 ALIGN=_SEGMENT_ALIGN,
+                MIXED=mixed,
                 num_warps=_SCAN_WARPS,
             )
             _launch(
@@ -1276,6 +1349,7 @@ def fused_paged_attention(
                 tickets,
                 store.mass,
                 store.contribution,
+                modes,
                 layer,
                 top,
                 capacity,
@@ -1292,54 +1366,59 @@ def fused_paged_attention(
                 BLOCK_D=block_d,
                 BLOCK_T=_GATHER_BLOCK,
                 PARTS=SELECT_PARTS,
+                MIXED=mixed,
+                num_warps=4,
+            )
+            if not mixed:
+                return output
+        else:
+            # Serial selection after another backend's dense call (uniform plans).
+            _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
-        _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 not mixed:
+        output = torch.empty_like(q)
     if rows:
         if workspace is None:
             workspace = _consume_workspace(q.device, rows, heads, block_d)
@@ -1362,6 +1441,7 @@ def fused_paged_attention(
             part_m,
             part_l,
             part_o,
+            modes,
             layer,
             block_tables.stride(0),
             selected.shape[-1],
@@ -1382,6 +1462,7 @@ def fused_paged_attention(
             BLOCK_D=block_d,
             BLOCK_T=32,
             SPLITS=NUM_SPLITS,
+            MIXED=mixed,
             num_warps=4,
         )
         _launch(
@@ -1397,6 +1478,7 @@ def fused_paged_attention(
             store.contribution,
             store.uses,
             output,
+            modes,
             layer,
             output.stride(0),
             output.stride(1),
@@ -1405,6 +1487,7 @@ def fused_paged_attention(
             HEAD_DIM=dim,
             BLOCK_D=block_d,
             SPLITS=NUM_SPLITS,
+            MIXED=mixed,
             num_warps=4,
         )
     return output
diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py
index f1bba19b..d6642884 100644
--- a/vllm_rlt/worker/async_speculative.py
+++ b/vllm_rlt/worker/async_speculative.py
@@ -8,7 +8,7 @@
 import numpy as np
 import torch
 
-from vllm_rlt.attention.sparse import DENSE, reference_mode
+from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, REFERENCE_DEPTH, 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
@@ -37,6 +37,9 @@ class DevicePipelineState:
     hidden: torch.Tensor
     candidates: torch.Tensor
     active: torch.Tensor
+    # Sparse decode: a permutation of the request's reference slots, moved with
+    # the tokens, so each row keeps the slot its loop-1 capture wrote.
+    slots: torch.Tensor | None = None
 
 
 @dataclass
@@ -54,6 +57,9 @@ class SpeculativeRequestState:
     # Bumped per request instance so reference store bindings reject a reused
     # request ID across generations.
     generation: int = 0
+    # Token pipelines own their reference slots for the request's lifetime: a
+    # row captures in one submission and consumes in the next.
+    reference: object | None = None
 
 
 class SpeculativeRoundBuffer:
@@ -214,6 +220,7 @@ def __init__(
                 capture_sizes=sizes,
                 reference_store=reference_store,
                 key_retention=None if sparse_decode is None else sparse_decode.key_retention,
+                device_plans=config.mode == "pipeline",
             )
             if execution.cuda_graphs
             else None
@@ -290,6 +297,9 @@ def _state(self, request, buffer):
                 torch.zeros(width, dtype=torch.long, device=self.device),
                 torch.arange(width, device=self.device) == 0,
             )
+            if self.reference_store is not None:
+                state.reference = self.reference_store.reserve(rid, 0, width, generation)
+                state.pipeline.slots = state.reference.slot_ids.clone()
         self.states[rid] = state
         return state
 
@@ -429,6 +439,22 @@ def _reference_plan(self, leases, rows, depth):
             owners, [mode] * len(owners), key_retention=self.sparse_decode.key_retention
         )
 
+    def _pipeline_reference_plan(self, states, slots, depth, active):
+        """Device modes of one pipeline step: each row's loop decides its mode."""
+        if slots is None:
+            return None
+        if any(state.reference is None for state in states):
+            raise RuntimeError("pipeline reference slots were discarded with their buffers")
+        capture = active & (depth == REFERENCE_DEPTH)
+        consume = active & (depth > REFERENCE_DEPTH)
+        modes = torch.where(capture, CAPTURE, torch.where(consume, CONSUME, DENSE)).to(torch.int8)
+        return self.reference_store.build_device_plan(
+            [state.reference for state in states],
+            slots,
+            modes,
+            key_retention=self.sparse_decode.key_retention,
+        )
+
     def _draft(self, plan, tables, allocations, leases=None):
         from vllm_rlt.kernels.routing import gather_kernel
 
@@ -556,7 +582,7 @@ def submit(self, batch):
             # 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:
+            if self.reference_store is not None and self.config.mode != "pipeline":
                 leases = tuple(
                     self.reference_store.reserve(
                         scheduled_request.request.request_id,
@@ -637,10 +663,19 @@ def _pipeline_step(self, spans, states, tables, allocations, buffer):
         prelude = self.runtime.forward(ForwardBatch(Operation.PRELUDE, tokens=tokens)).hidden
         hidden = torch.where((depths == 0)[:, None], prelude, retained)
         row_tables = torch.cat([tables[i : i + 1].expand(w, -1, -1) for i, w in enumerate(widths)])
+        slots = None
+        if self.reference_store is not None:
+            slots = torch.cat([state.pipeline.slots[:w] for state, w in zip(states, widths)])
         for offset in range(self.config.draft_loops):
             depth = torch.where(active, depths + offset, 0)
             hidden = self._execute_recurrent_step(
-                hidden, positions, row_tables, depth, allocations, active
+                hidden,
+                positions,
+                row_tables,
+                depth,
+                allocations,
+                active,
+                reference_plan=self._pipeline_reference_plan(states, slots, depth, active),
             )
         choices = self._coda(hidden).argmax(-1)
         start = 0
@@ -665,6 +700,10 @@ def _pipeline_step(self, spans, states, tables, allocations, buffer):
             source = (ids + emitted.long()).clamp_max(len(ids) - 1)
             remaining = valid.long().sum() - emitted.long()
             surviving = (ids < remaining) & ~reject
+            if p.slots is not None:
+                # A permutation, unlike `source`: a freed root slot moves to the
+                # tail, where only inactive rows or newborns can take it.
+                p.slots.copy_(p.slots.gather(0, (ids + emitted.long()) % len(ids)))
             p.tokens.copy_(p.tokens.gather(0, source))
             p.depths.copy_(new_depths.gather(0, source))
             p.hidden.copy_(next_hidden[source])
@@ -738,6 +777,9 @@ def restore_pipeline(self, request, snapshot):
             return
         if snapshot.request is not request or request.request_id in self.states:
             raise RuntimeError("invalid pipeline restoration generation")
+        if self.reference_store is not None:
+            # Sparse decode excludes preemption; a snapshot carries no references.
+            raise RuntimeError("sparse decode attention cannot restore a pipeline snapshot")
         if snapshot.pending or snapshot.position != request.position:
             raise RuntimeError("pipeline snapshot disagrees with committed request position")
         width = min(
@@ -793,6 +835,10 @@ def release(self, request):
             del self.states[request.request_id]
             if state.slot is not None:
                 self.free_slots.append(state.slot)
+            if state.reference is not None:
+                # The last reader completed with the event above.
+                with torch.cuda.stream(self.compute_stream):
+                    self.reference_store.release(state.reference)
 
     def discard_buffers(self):
         """Reset round storage after both streams drain and all tickets retire.
@@ -810,6 +856,10 @@ def discard_buffers(self):
             # Both streams have drained and every ticket retired; any leftover
             # binding can only come from a failed submission.
             reference_store.discard()
+            for state in getattr(self, "states", {}).values():
+                # Discard freed every binding; a surviving pipeline must not
+                # release or read these slots again.
+                state.reference = None
 
     def synchronize(self):
         self.compute_stream.synchronize()
diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py
index 833f20f2..63c7d20e 100644
--- a/vllm_rlt/worker/attention_reference.py
+++ b/vllm_rlt/worker/attention_reference.py
@@ -31,7 +31,13 @@
 
 import torch
 
-from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan
+from vllm_rlt.attention.sparse import (
+    CAPTURE,
+    CONSUME,
+    DENSE,
+    DeviceSparseDecodePlan,
+    SparseDecodePlan,
+)
 
 __all__ = ["ReferenceLease", "AttentionReferenceStore", "DENSE", "CAPTURE", "CONSUME"]
 
@@ -302,10 +308,36 @@ def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan
             store=self,
         )
 
-    def _require_live(self, plan: SparseDecodePlan) -> None:
+    def build_device_plan(self, leases, slot_ids, modes, *, key_retention: float):
+        """Publish a plan whose rows name their slots and modes on the device.
+
+        Every slot in ``slot_ids`` must belong to one of ``leases``; the host
+        checks those leases live before each use, as for a host plan.
+        """
+        if slot_ids.shape != modes.shape or slot_ids.dim() != 1:
+            raise ValueError("a device reference plan needs one slot and mode per row")
+        plan = DeviceSparseDecodePlan(
+            modes=modes,
+            slot_ids=slot_ids,
+            leases=tuple(leases),
+            key_retention=key_retention,
+            key_bucket=self.key_bucket,
+            store=self,
+        )
+        self._require_live(plan)
+        return plan
+
+    def _require_live(self, plan) -> 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")
+        if isinstance(plan, DeviceSparseDecodePlan):
+            for lease in plan.leases:
+                for slot, incarnation in zip(lease.slots, lease.incarnations):
+                    binding = self._bindings.get(slot)
+                    if binding != (lease.request_id, lease.generation, incarnation):
+                        raise RuntimeError("stale attention reference plan for a retired token")
+            return
         for slot, incarnation, mode in zip(plan.slots, plan.incarnations, plan.modes):
             if slot == -1 and mode == DENSE:
                 continue
diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py
index 3b70d141..4fd04733 100644
--- a/vllm_rlt/worker/cuda_graph.py
+++ b/vllm_rlt/worker/cuda_graph.py
@@ -8,7 +8,13 @@
 
 import torch
 
-from vllm_rlt.attention.sparse import CAPTURE, CONSUME, DENSE, SparseDecodePlan
+from vllm_rlt.attention.sparse import (
+    CAPTURE,
+    CONSUME,
+    DENSE,
+    DeviceSparseDecodePlan,
+    SparseDecodePlan,
+)
 from vllm_rlt.config import cuda_graph_rows
 from vllm_rlt.worker.metadata import DeviceKVMetadata
 
@@ -347,6 +353,10 @@ def execute(self, call, hidden):
         return output
 
 
+# Graph key of entries that replay a device plan (token pipelines); not a row mode.
+DEVICE_PLAN = "device"
+
+
 class DeviceRecurrentGraphs(RecurrentGraphs):
     """Core graphs whose actual positions and KV frontier remain on device.
 
@@ -356,7 +366,9 @@ class DeviceRecurrentGraphs(RecurrentGraphs):
     With a fused sparse reference store, every bucket is captured once per row
     mode. A sparse entry owns a static plan whose slot ids are refilled before
     each replay; the live plan's liveness is checked on the host at that point,
-    since replay runs no Python.
+    since replay runs no Python. With ``device_plans`` (token pipelines) the
+    sparse entries instead replay a device plan, whose per-row modes are
+    refilled alongside the slot ids.
     """
 
     def __init__(
@@ -368,6 +380,7 @@ def __init__(
         capture_sizes=None,
         reference_store=None,
         key_retention=None,
+        device_plans=False,
     ):
         super().__init__(model, cache, execution, compute_gate)
         self.capture_sizes = capture_sizes or graph_capture_sizes(execution)
@@ -383,7 +396,7 @@ def __init__(
             if reference_store.selected.shape[-1] != (reference_store.key_bucket or 1):
                 # A replay must never see a grown index pool at a new address.
                 raise ValueError("sparse CUDA graphs need the reference index pool at full size")
-            self.modes = (DENSE, CAPTURE, CONSUME)
+            self.modes = (DENSE, DEVICE_PLAN) if device_plans else (DENSE, CAPTURE, CONSUME)
             config = model.config
             width = math.ceil(config.max_position_embeddings / cache.block_size)
             self.proxy = _DeviceCache(
@@ -412,7 +425,22 @@ def _entry(self, hidden, batch, size, mode):
             bucketed=True,
         )
         entry.mode = mode
-        if mode != DENSE:
+        if mode == DEVICE_PLAN:
+            store = self.reference_store
+            entry.slot_ids = torch.full(
+                (size,), store.padding_slot, device=self.cache.device, dtype=torch.int64
+            )
+            # Padded rows stay dense on the never-bound padding slot.
+            entry.modes = torch.full((size,), DENSE, device=self.cache.device, dtype=torch.int8)
+            entry.metadata.reference_plan = DeviceSparseDecodePlan(
+                modes=entry.modes,
+                slot_ids=entry.slot_ids,
+                leases=(),
+                key_retention=self.key_retention,
+                key_bucket=store.key_bucket,
+                store=store,
+            )
+        elif mode != DENSE:
             store = self.reference_store
             entry.slot_ids = torch.full(
                 (size,), store.padding_slot, device=self.cache.device, dtype=torch.int64
@@ -463,9 +491,13 @@ def _mode(self, batch):
         plan = getattr(batch, "reference_plan", None)
         if plan is None or plan.dense():
             return DENSE
-        modes = set(plan.modes)
         if self.reference_store is None or plan.store is not self.reference_store:
             return None
+        if isinstance(plan, DeviceSparseDecodePlan):
+            if DEVICE_PLAN not in self.modes or plan.key_retention != self.key_retention:
+                return None
+            return DEVICE_PLAN
+        modes = set(plan.modes)
         if len(modes) != 1 or plan.key_retention != self.key_retention:
             return None
         return plan.modes[0]
@@ -498,6 +530,9 @@ def prepare_inputs(self, call, hidden):
                 stream.wait_event(self.last_event)
             entry.slot_ids[count:].fill_(self.reference_store.padding_slot)
             entry.slot_ids[:count].copy_(plan.slot_ids[:count])
+            if entry.mode == DEVICE_PLAN:
+                entry.modes[count:].fill_(DENSE)
+                entry.modes[:count].copy_(plan.modes[:count])
         super().prepare_inputs(call, hidden)
 
     @torch.inference_mode()

From 6f5ebdd9f25d81bca024d62de2f57ad7bcd1255e Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Sun, 1 Nov 2026 02:04:54 +0100
Subject: [PATCH 27/32] Cover token pipelines with sparse decode on the P/D
 decode stage

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_pd.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/tests/test_pd.py b/tests/test_pd.py
index f98f2b48..f1f553c4 100644
--- a/tests/test_pd.py
+++ b/tests/test_pd.py
@@ -361,6 +361,7 @@ def create_pd(
     prefix_caching=False,
     sparse=None,
     decode_batched_tokens=2,
+    speculative_mode="block",
 ):
     pytest.importorskip("nixl")
     from vllm_rlt.pd.engine import PDEngine
@@ -393,7 +394,7 @@ def create_pd(
         exit_config=ExitConfig("ouro")
         if speculative
         else ExitConfig("trace", depths_by_request={"t": [4, 2, 3, 1] * 16}),
-        speculative_config=SpeculativeConfig(3) if speculative else None,
+        speculative_config=SpeculativeConfig(3, mode=speculative_mode) if speculative else None,
         sparse_decode_attention_config=sparse,
         execution_config=ExecutionConfig(
             async_scheduling=asynchronous,
@@ -776,8 +777,9 @@ def test_pd_sparse_decode_on_the_decode_stage(sparse):
 
 
 @pytest.mark.gpu
+@pytest.mark.parametrize("mode", ["block", "pipeline"])
 @pytest.mark.parametrize("multi", [False, True])
-def test_pd_sparse_decode_replays_graphs_on_the_decode_stage(multi):
+def test_pd_sparse_decode_replays_graphs_on_the_decode_stage(multi, mode):
     fused = SparseDecodeAttentionConfig(enabled=True, key_retention=0.5, implementation="fused")
     params = SamplingParams(max_tokens=9, ignore_eos=True)
     prompts = [[1, 2, 3, 4, 5], [2, 3], [4], [5, 6, 7]]
@@ -791,6 +793,7 @@ def test_pd_sparse_decode_replays_graphs_on_the_decode_stage(multi):
             asynchronous=True,
             sparse=sparse,
             decode_batched_tokens=8,
+            speculative_mode=mode,
         ) as engine:
             for index, prompt in enumerate(prompts):
                 engine.add_request(str(index), prompt, params)

From 1aba42ac4c39f7b0d0368bd587927aaf16f43222 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Sun, 1 Nov 2026 02:33:10 +0100
Subject: [PATCH 28/32] Compare mixed-mode dense rows to fp32 rounding

On an L40S (sm_89) the mixed-mode kernel variant rounds dense and capture rows
up to one ulp away from the dense kernel (1.8e-7), while consume rows and every
store field stay bit-identical; the RTX 5090 matched exactly. The mixed-mode
test now requires exact consume rows and store fields, and dense and capture
rows within 1e-6; uniform launches keep their bit-for-bit test.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_sparse_decode_pipeline.py | 10 +++++++---
 vllm_rlt/kernels/sparse_decode.py    |  6 ++++--
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/tests/test_sparse_decode_pipeline.py b/tests/test_sparse_decode_pipeline.py
index 6072e244..1a5fc986 100644
--- a/tests/test_sparse_decode_pipeline.py
+++ b/tests/test_sparse_decode_pipeline.py
@@ -103,9 +103,13 @@ def test_mixed_modes_match_uniform_fused_kernels():
             dense_attention=dense if mode == CAPTURE else None,
         )
     torch.cuda.synchronize()
-    # Dense and capture rows are the dense kernel's output; consume rows the
-    # uniform consume's; every store field matches slot for slot.
-    assert torch.equal(mixed, expected)
+    # Consume rows are the uniform consume's, bit for bit. Dense and capture
+    # rows are the dense kernel's up to fp32 rounding: the mixed-mode variant is
+    # compiled separately, and some architectures (sm_89) round it differently
+    # by an ulp. Every store field matches slot for slot.
+    consume = [row for row, mode in enumerate(modes) if mode == CONSUME]
+    assert torch.equal(mixed[consume], expected[consume])
+    torch.testing.assert_close(mixed, expected, rtol=0, atol=1e-6)
     for name in ("selected", "counts", "mass", "output", "contribution", "valid", "uses"):
         assert torch.equal(getattr(mixed_store, name), getattr(uniform_store, name)), name
 
diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py
index 7a2854f3..a51f4e63 100644
--- a/vllm_rlt/kernels/sparse_decode.py
+++ b/vllm_rlt/kernels/sparse_decode.py
@@ -140,8 +140,10 @@ def _select_and_publish(
 
 
 # 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).
+# scores. Its arithmetic must stay identical to the dense kernel: a uniform
+# capture's output is required to equal the dense backend bit for bit (tested).
+# The mixed-mode variant (device plans) is compiled separately and matches it to
+# fp32 rounding; some architectures (sm_89) round it differently by an ulp.
 @triton.jit
 def _capture_split_kernel(
     Q,

From b0a73bcb56d1ae7696619832fc1683f3cbef32a0 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Sun, 1 Nov 2026 15:03:46 +0100
Subject: [PATCH 29/32] Launch no selection kernels in pipeline steps that
 cannot capture

Pipeline rows sit at multiples of draft_loops, so whether any row can be in its
capture loop at a given offset is known on the host: with the default two
draft loops, offset 0 runs loops 0 and 2 and never captures. Device plans now
carry that as may_capture; when it is false the fused path launches the dense
split, the first pass (which merges the dense rows) and the consume kernels,
four launches instead of eight, and graphs key such steps separately. Outputs
and store fields are identical (tested).

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_sparse_decode_pipeline.py   |  61 ++++++++++++
 vllm_rlt/attention/sparse.py           |   3 +
 vllm_rlt/kernels/sparse_decode.py      | 124 +++++++++++++------------
 vllm_rlt/worker/async_speculative.py   |   8 +-
 vllm_rlt/worker/attention_reference.py |   5 +-
 vllm_rlt/worker/cuda_graph.py          |  17 +++-
 6 files changed, 150 insertions(+), 68 deletions(-)

diff --git a/tests/test_sparse_decode_pipeline.py b/tests/test_sparse_decode_pipeline.py
index 1a5fc986..9d459e86 100644
--- a/tests/test_sparse_decode_pipeline.py
+++ b/tests/test_sparse_decode_pipeline.py
@@ -114,6 +114,67 @@ def test_mixed_modes_match_uniform_fused_kernels():
         assert torch.equal(getattr(mixed_store, name), getattr(uniform_store, name)), name
 
 
+@pytest.mark.gpu
+def test_steps_without_capture_skip_selection_kernels(monkeypatch):
+    """A device plan that cannot capture launches four kernels, with identical results."""
+    from vllm_rlt.kernels import sparse_decode
+    from vllm_rlt.kernels.paged_attention import TritonPagedAttention
+
+    rows, heads, kv_heads, dim, block_size, pages = 4, 4, 2, 64, 16, 9
+    dense = TritonPagedAttention()
+    q, keys, values, tables = _paged(
+        rows, heads, kv_heads, dim, pages, block_size, torch.float32, 3
+    )
+    lengths = torch.tensor([40, 129, 7, 90], device="cuda", dtype=torch.int32)
+    modes = torch.tensor([DENSE, CONSUME, DENSE, CONSUME], device="cuda", dtype=torch.int8)
+    results, launches = [], []
+    for may_capture in (True, False):
+        store = make_store(
+            slots=rows,
+            layers=1,
+            heads=heads,
+            head_dim=dim,
+            key_bucket=pages * block_size,
+            device="cuda",
+            implementation="fused",
+        )
+        lease = store.reserve("r", 0, rows, generation=1)
+        consume_rows = [1, 3]
+        capture = store.build_plan(
+            [(lease, row) for row in consume_rows], [CAPTURE] * 2, key_retention=0.25
+        )
+        sparse_decode.fused_paged_attention(
+            q[consume_rows],
+            keys,
+            values,
+            tables[consume_rows],
+            lengths[consume_rows],
+            capture,
+            layer=0,
+            dense_attention=dense,
+        )
+        plan = store.build_device_plan(
+            [lease], lease.slot_ids, modes, key_retention=0.25, may_capture=may_capture
+        )
+        count = []
+        launch = sparse_decode._launch
+        monkeypatch.setattr(
+            sparse_decode,
+            "_launch",
+            lambda kernel, *a, **k: (count.append(kernel), launch(kernel, *a, **k)),
+        )
+        output = sparse_decode.fused_paged_attention(
+            q, keys, values, tables, lengths, plan, layer=0, dense_attention=dense
+        )
+        monkeypatch.setattr(sparse_decode, "_launch", launch)
+        torch.cuda.synchronize()
+        launches.append(len(count))
+        results.append((output, store.uses.clone(), store.counts.clone(), store.valid.clone()))
+    assert launches == [8, 4]
+    for a, b in zip(*results):
+        assert torch.equal(a, b)
+
+
 def _engine(mode, sparse, *, draft_loops=2, graphs=False, retention=0.5):
     from vllm_rlt.engine.stage_engine import StageEngine
     from vllm_rlt.models import OuroForCausalLM
diff --git a/vllm_rlt/attention/sparse.py b/vllm_rlt/attention/sparse.py
index 4fb5a981..668b25a5 100644
--- a/vllm_rlt/attention/sparse.py
+++ b/vllm_rlt/attention/sparse.py
@@ -118,6 +118,9 @@ class DeviceSparseDecodePlan:
     key_retention: float
     key_bucket: int | None
     store: object | None = None
+    # False when no row can be in its capture loop (the host knows which loops a
+    # step runs); the fused path then launches no selection kernels.
+    may_capture: bool = True
 
     def dense(self) -> bool:
         """Unknown on the host: the device may route any row to sparse work."""
diff --git a/vllm_rlt/kernels/sparse_decode.py b/vllm_rlt/kernels/sparse_decode.py
index a51f4e63..7af6f64f 100644
--- a/vllm_rlt/kernels/sparse_decode.py
+++ b/vllm_rlt/kernels/sparse_decode.py
@@ -1273,7 +1273,10 @@ def fused_paged_attention(
                 num_warps=4,
             )
             first_pass = 1 if folded else 0
-            for level in range(first_pass, 4):
+            # A device plan whose step has no capture loop still merges its dense
+            # rows in the first pass; it launches no further selection kernels.
+            selecting = not mixed or plan.may_capture
+            for level in range(first_pass, 4 if selecting else first_pass + 1):
                 _launch(
                     _select_pass_kernel,
                     (rows, heads, SELECT_PARTS),
@@ -1312,65 +1315,66 @@ def fused_paged_attention(
                     MIXED=mixed,
                     num_warps=_SCAN_WARPS,
                 )
-            _launch(
-                _select_write_kernel,
-                (rows, heads, SELECT_PARTS),
-                scores,
-                context_lengths,
-                retention,
-                plan.slot_ids,
-                state,
-                offsets,
-                selected,
-                modes,
-                layer,
-                top,
-                capacity,
-                selected.shape[-1],
-                NUM_LAYERS=store.layers,
-                NUM_HEADS=heads,
-                BLOCK_T=_SCAN_BLOCK,
-                PARTS=SELECT_PARTS,
-                ALIGN=_SEGMENT_ALIGN,
-                MIXED=mixed,
-                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,
-                modes,
-                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,
-                MIXED=mixed,
-                num_warps=4,
-            )
+            if selecting:
+                _launch(
+                    _select_write_kernel,
+                    (rows, heads, SELECT_PARTS),
+                    scores,
+                    context_lengths,
+                    retention,
+                    plan.slot_ids,
+                    state,
+                    offsets,
+                    selected,
+                    modes,
+                    layer,
+                    top,
+                    capacity,
+                    selected.shape[-1],
+                    NUM_LAYERS=store.layers,
+                    NUM_HEADS=heads,
+                    BLOCK_T=_SCAN_BLOCK,
+                    PARTS=SELECT_PARTS,
+                    ALIGN=_SEGMENT_ALIGN,
+                    MIXED=mixed,
+                    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,
+                    modes,
+                    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,
+                    MIXED=mixed,
+                    num_warps=4,
+                )
             if not mixed:
                 return output
         else:
diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py
index d6642884..f89df1f8 100644
--- a/vllm_rlt/worker/async_speculative.py
+++ b/vllm_rlt/worker/async_speculative.py
@@ -439,7 +439,7 @@ def _reference_plan(self, leases, rows, depth):
             owners, [mode] * len(owners), key_retention=self.sparse_decode.key_retention
         )
 
-    def _pipeline_reference_plan(self, states, slots, depth, active):
+    def _pipeline_reference_plan(self, states, slots, depth, active, offset):
         """Device modes of one pipeline step: each row's loop decides its mode."""
         if slots is None:
             return None
@@ -448,11 +448,15 @@ def _pipeline_reference_plan(self, states, slots, depth, active):
         capture = active & (depth == REFERENCE_DEPTH)
         consume = active & (depth > REFERENCE_DEPTH)
         modes = torch.where(capture, CAPTURE, torch.where(consume, CONSUME, DENSE)).to(torch.int8)
+        # Rows sit at multiples of draft_loops, so whether any row can be in its
+        # capture loop at this offset is known on the host.
+        depths = range(0, self.config.target_loops, self.config.draft_loops)
         return self.reference_store.build_device_plan(
             [state.reference for state in states],
             slots,
             modes,
             key_retention=self.sparse_decode.key_retention,
+            may_capture=REFERENCE_DEPTH - offset in depths,
         )
 
     def _draft(self, plan, tables, allocations, leases=None):
@@ -675,7 +679,7 @@ def _pipeline_step(self, spans, states, tables, allocations, buffer):
                 depth,
                 allocations,
                 active,
-                reference_plan=self._pipeline_reference_plan(states, slots, depth, active),
+                reference_plan=self._pipeline_reference_plan(states, slots, depth, active, offset),
             )
         choices = self._coda(hidden).argmax(-1)
         start = 0
diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py
index 63c7d20e..d88f7291 100644
--- a/vllm_rlt/worker/attention_reference.py
+++ b/vllm_rlt/worker/attention_reference.py
@@ -308,7 +308,9 @@ def build_plan(self, owners, modes, *, key_retention: float) -> SparseDecodePlan
             store=self,
         )
 
-    def build_device_plan(self, leases, slot_ids, modes, *, key_retention: float):
+    def build_device_plan(
+        self, leases, slot_ids, modes, *, key_retention: float, may_capture: bool = True
+    ):
         """Publish a plan whose rows name their slots and modes on the device.
 
         Every slot in ``slot_ids`` must belong to one of ``leases``; the host
@@ -323,6 +325,7 @@ def build_device_plan(self, leases, slot_ids, modes, *, key_retention: float):
             key_retention=key_retention,
             key_bucket=self.key_bucket,
             store=self,
+            may_capture=may_capture,
         )
         self._require_live(plan)
         return plan
diff --git a/vllm_rlt/worker/cuda_graph.py b/vllm_rlt/worker/cuda_graph.py
index 4fd04733..87ae0ec6 100644
--- a/vllm_rlt/worker/cuda_graph.py
+++ b/vllm_rlt/worker/cuda_graph.py
@@ -353,8 +353,10 @@ def execute(self, call, hidden):
         return output
 
 
-# Graph key of entries that replay a device plan (token pipelines); not a row mode.
+# Graph keys of entries that replay a device plan (token pipelines); not row
+# modes. A step in which no row can capture records no selection kernels.
 DEVICE_PLAN = "device"
+DEVICE_PLAN_NO_CAPTURE = "device-no-capture"
 
 
 class DeviceRecurrentGraphs(RecurrentGraphs):
@@ -396,7 +398,11 @@ def __init__(
             if reference_store.selected.shape[-1] != (reference_store.key_bucket or 1):
                 # A replay must never see a grown index pool at a new address.
                 raise ValueError("sparse CUDA graphs need the reference index pool at full size")
-            self.modes = (DENSE, DEVICE_PLAN) if device_plans else (DENSE, CAPTURE, CONSUME)
+            self.modes = (
+                (DENSE, DEVICE_PLAN, DEVICE_PLAN_NO_CAPTURE)
+                if device_plans
+                else (DENSE, CAPTURE, CONSUME)
+            )
             config = model.config
             width = math.ceil(config.max_position_embeddings / cache.block_size)
             self.proxy = _DeviceCache(
@@ -425,7 +431,7 @@ def _entry(self, hidden, batch, size, mode):
             bucketed=True,
         )
         entry.mode = mode
-        if mode == DEVICE_PLAN:
+        if mode in (DEVICE_PLAN, DEVICE_PLAN_NO_CAPTURE):
             store = self.reference_store
             entry.slot_ids = torch.full(
                 (size,), store.padding_slot, device=self.cache.device, dtype=torch.int64
@@ -439,6 +445,7 @@ def _entry(self, hidden, batch, size, mode):
                 key_retention=self.key_retention,
                 key_bucket=store.key_bucket,
                 store=store,
+                may_capture=mode == DEVICE_PLAN,
             )
         elif mode != DENSE:
             store = self.reference_store
@@ -496,7 +503,7 @@ def _mode(self, batch):
         if isinstance(plan, DeviceSparseDecodePlan):
             if DEVICE_PLAN not in self.modes or plan.key_retention != self.key_retention:
                 return None
-            return DEVICE_PLAN
+            return DEVICE_PLAN if plan.may_capture else DEVICE_PLAN_NO_CAPTURE
         modes = set(plan.modes)
         if len(modes) != 1 or plan.key_retention != self.key_retention:
             return None
@@ -530,7 +537,7 @@ def prepare_inputs(self, call, hidden):
                 stream.wait_event(self.last_event)
             entry.slot_ids[count:].fill_(self.reference_store.padding_slot)
             entry.slot_ids[:count].copy_(plan.slot_ids[:count])
-            if entry.mode == DEVICE_PLAN:
+            if entry.mode in (DEVICE_PLAN, DEVICE_PLAN_NO_CAPTURE):
                 entry.modes[count:].fill_(DENSE)
                 entry.modes[:count].copy_(plan.modes[:count])
         super().prepare_inputs(call, hidden)

From 40056e8fa5eca550aa19219ef1d96fb8844ae1f0 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Sun, 1 Nov 2026 23:23:22 +0100
Subject: [PATCH 30/32] Preempt and resume ordinary sparse decode with its live
 reference

A token suspended after its reference loop must consume that same reference
when it resumes. Suspension used to drop it, so the remaining loops fell back
to dense attention; sparse decode therefore rejected preemption outright.

Suspension now copies the token's reference slot (retained indices, counts,
mass, output, contribution and the per-slot bookkeeping, plus the published
host counts) into storage the saved state owns, resident or on the CPU with
the offloaded KV, including a later resident-to-offload conversion. Resumption
binds a fresh slot and generation and copies it back on the owner stream.

Preemption is lifted for ordinary decode, synchronous and asynchronous.
Speculative rows hold per-round references with no snapshot, so speculative
decoding with sparse attention still rejects preemption explicitly.

Tests: a suspension right after the capture loop (resident, offload, and
resident then offload) resumes with a sparse consume and the same tokens as
an uninterrupted run; KV-pressure preemption matches the unpreempted run,
synchronous and asynchronous, with references carried across suspensions;
abort while suspended frees everything and the id is reusable. CUDA variant
for the batched and fused kernels.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_sparse_decode_preemption.py | 176 +++++++++++++++++++++++++
 tests/test_sparse_decode_sync.py       |   5 +-
 vllm_rlt/engine/preemption.py          |  10 ++
 vllm_rlt/engine/stage_engine.py        |   8 +-
 vllm_rlt/worker/attention_reference.py | 124 ++++++++++++++++-
 vllm_rlt/worker/model_runner.py        |  29 ++++
 6 files changed, 348 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_sparse_decode_preemption.py

diff --git a/tests/test_sparse_decode_preemption.py b/tests/test_sparse_decode_preemption.py
new file mode 100644
index 00000000..d7d09c88
--- /dev/null
+++ b/tests/test_sparse_decode_preemption.py
@@ -0,0 +1,176 @@
+"""Preemption of ordinary sparse decode: a suspended token keeps its reference."""
+
+import pytest
+import torch
+
+from tests.test_sparse_decode_sync import _bounded_pool, _ordinary, engine, generate
+from vllm_rlt import CacheConfig, SamplingParams, SchedulerConfig, SpeculativeConfig
+from vllm_rlt.attention.sparse import CAPTURE, CONSUME
+
+
+def _scheduler(**options):
+    values = dict(
+        max_num_seqs=1, max_num_batched_tokens=16, prefill_chunk_size=3, enable_preemption=True
+    )
+    values.update(options)
+    return SchedulerConfig(**values)
+
+
+def _record_modes(e, monkeypatch):
+    """Record (request, mode) of every sparse plan the runner builds."""
+    seen = []
+    build = e.reference_store.build_plan
+
+    def spy(owners, modes, **kwargs):
+        seen.append([(None if o is None else o[0].request_id, m) for o, m in zip(owners, modes)])
+        return build(owners, modes, **kwargs)
+
+    monkeypatch.setattr(e.reference_store, "build_plan", spy)
+    return seen
+
+
+def _drain(e, outputs):
+    for _ in range(1000):
+        if not e.has_unfinished_requests():
+            return outputs
+        for output in e.step():
+            if output.finished:
+                outputs[output.request_id] = list(output.token_ids)
+    raise AssertionError("engine did not drain")
+
+
+@pytest.mark.parametrize("mode", ["resident", "offload", "resident_then_offload"])
+def test_suspension_after_capture_resumes_with_the_same_reference(mode, monkeypatch):
+    victim = SamplingParams(max_tokens=8, ignore_eos=True, priority=1)
+    arrival = SamplingParams(max_tokens=4, ignore_eos=True)
+    results = []
+    for preempt in (False, True):
+        e = _ordinary(False, 0.5, scheduler_config=_scheduler())
+        seen = _record_modes(e, monkeypatch)
+        try:
+            e.add_request("a", [2, 3, 4, 5], victim)
+            outputs, mark = {}, None
+            for _ in range(1000):
+                for output in e.step():
+                    if output.finished:
+                        outputs[output.request_id] = list(output.token_ids)
+                # Stop right after a capture of a generated token: its later
+                # loops have not consumed the reference yet.
+                if seen and ("a", CAPTURE) in seen[-1] and e.reference_store.outstanding == 1:
+                    break
+            assert "a" not in outputs
+            e.add_request("b", [2, 3], arrival)
+            if preempt:
+                waiting = e.scheduler.requests["b"]
+                assert e.scheduler.preempt(waiting, priority_only=mode != "offload")
+                saved = e.preemption.snapshots["a"]["references"]
+                assert saved is not None and e.reference_store.outstanding == 0
+                if mode == "resident_then_offload":
+                    assert saved[1].selected.device == e.cache_manager.device
+                    assert e.scheduler.preempt(waiting, waiting_only=True)
+                    saved = e.preemption.snapshots["a"]["references"]
+                expected = e.cache_manager.device if mode == "resident" else torch.device("cpu")
+                assert saved[1].selected.device == expected
+                mark = len(seen)
+            results.append(_drain(e, outputs))
+            assert e.preemption.resumptions == int(preempt)
+            assert not e.preemption.snapshots
+            assert e.reference_store.outstanding == 0
+            assert e.cache_manager.num_used_blocks == 0
+            if preempt:
+                # The resumed token consumes its restored reference, not dense.
+                after = next(plan for plan in seen[mark:] if any(r == "a" for r, _ in plan))
+                assert ("a", CONSUME) in after
+        finally:
+            e.close()
+    assert results[0] == results[1]
+
+
+@pytest.mark.parametrize("asynchronous", [False, True])
+def test_kv_pressure_preemption_matches_unpreempted_sparse(asynchronous, monkeypatch):
+    def run(cache_blocks, preemption):
+        e = _ordinary(
+            asynchronous,
+            0.5,
+            cache_config=CacheConfig(cache_blocks, 2, incremental_allocation=True),
+            scheduler_config=_scheduler(max_num_seqs=3, enable_preemption=preemption),
+        )
+        _bounded_pool(e, monkeypatch)
+        snapshot = e.model_runner.snapshot_attention_reference
+        carried = []
+
+        def counted(request, device):
+            saved = snapshot(request, device)
+            carried.append(saved is not None)
+            return saved
+
+        monkeypatch.setattr(e.model_runner, "snapshot_attention_reference", counted)
+        try:
+            outputs = generate(e, count=3)
+            return outputs, e.preemption.preemptions, sum(carried)
+        finally:
+            e.close()
+
+    baseline, none, _ = run(512, False)
+    pressured, preemptions, carried = run(40, True)
+    assert none == 0 and preemptions > 0 and carried > 0
+    assert pressured == baseline
+
+
+def test_abort_while_suspended_frees_reference_and_id_is_reusable(monkeypatch):
+    e = _ordinary(False, 0.5, scheduler_config=_scheduler())
+    seen = _record_modes(e, monkeypatch)
+    try:
+        e.add_request("a", [2, 3, 4], SamplingParams(max_tokens=8, ignore_eos=True, priority=1))
+        for _ in range(1000):
+            e.step()
+            if seen and ("a", CAPTURE) in seen[-1]:
+                break
+        e.add_request("b", [2, 3], SamplingParams(max_tokens=2, ignore_eos=True))
+        assert e.scheduler.preempt(e.scheduler.requests["b"], priority_only=True)
+        assert e.preemption.snapshots["a"]["references"] is not None
+        e.abort_request("a")
+        assert "a" not in e.preemption.snapshots
+        _drain(e, {})
+        assert e.reference_store.outstanding == 0
+        assert generate(e, count=2)
+    finally:
+        e.close()
+
+
+@pytest.mark.parametrize("draft", [1, 2])
+def test_speculative_preemption_stays_rejected(draft):
+    with pytest.raises(ValueError, match="preemption only for ordinary decode"):
+        engine(
+            speculative_config=SpeculativeConfig(2, draft_loops=draft),
+            scheduler_config=SchedulerConfig(enable_preemption=True),
+        )
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("implementation", ["batched", "fused"])
+@pytest.mark.parametrize("asynchronous", [False, True])
+def test_cuda_kv_pressure_preemption_matches_unpreempted(implementation, asynchronous):
+    from vllm_rlt import SparseDecodeAttentionConfig
+
+    def run(cache_blocks, preemption):
+        e = _ordinary(
+            asynchronous,
+            None,
+            device="cuda",
+            cache_config=CacheConfig(cache_blocks, 2, incremental_allocation=True),
+            scheduler_config=_scheduler(max_num_seqs=3, enable_preemption=preemption),
+            sparse_decode_attention_config=SparseDecodeAttentionConfig(
+                enabled=True, key_retention=0.5, implementation=implementation
+            ),
+        )
+        try:
+            return generate(e, count=3), e.preemption.preemptions
+        finally:
+            e.close()
+
+    baseline, _ = run(512, False)
+    pressured, preemptions = run(40, True)
+    torch.cuda.synchronize()
+    assert preemptions > 0
+    assert pressured == baseline
diff --git a/tests/test_sparse_decode_sync.py b/tests/test_sparse_decode_sync.py
index e22f8695..a554dd9b 100644
--- a/tests/test_sparse_decode_sync.py
+++ b/tests/test_sparse_decode_sync.py
@@ -230,7 +230,10 @@ 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)),
+        "preemption": dict(
+            speculative_config=SpeculativeConfig(2),
+            scheduler_config=SchedulerConfig(enable_preemption=True),
+        ),
         "shared": dict(cache_config=CacheConfig(512, 2, layout="shared")),
     }[setting]
     with pytest.raises(ValueError, match="sparse decode attention"):
diff --git a/vllm_rlt/engine/preemption.py b/vllm_rlt/engine/preemption.py
index f8be4a75..3c9d977b 100644
--- a/vllm_rlt/engine/preemption.py
+++ b/vllm_rlt/engine/preemption.py
@@ -43,10 +43,18 @@ def suspend(self, victim, *, offload_kv=True):
             pipeline = runner.snapshot_pipeline(victim, device)
         elif pipeline is not None:
             pipeline = pipeline.to(device)
+        # Sparse decode: a token past its reference loop owns a reference that
+        # its remaining loops read. It moves with the rest of the saved state.
+        references = previous.get("references") if previous is not None else None
+        if previous is None and e.reference_store is not None:
+            references = e.model_runner.snapshot_attention_reference(victim, device)
+        elif references is not None:
+            references = (references[0], references[1].to(device))
         # Build a new record before releasing resources. In particular, a failed
         # offload must leave the prior resident state available for resumption.
         snapshot = dict(
             pipeline=pipeline,
+            references=references,
             phase=phase,
             kv=cache.snapshot(victim.request_id) if offload_kv else None,
             hidden=None if hidden is None else hidden.to(device).clone(),
@@ -85,6 +93,8 @@ def resume(self, request):
         )
         if self.engine.speculative_runner is not None:
             self.engine.speculative_runner.restore_pipeline(request, state.get("pipeline"))
+        if state.get("references") is not None:
+            self.engine.model_runner.restore_attention_reference(request, state["references"])
         # Resident copies completed before suspension released their runner slot.
         # Restored CPU buffers must be visible to every stream before resumption.
         if kv is not None and cache.device.type == "cuda":
diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py
index ece18800..6c8d62e3 100644
--- a/vllm_rlt/engine/stage_engine.py
+++ b/vllm_rlt/engine/stage_engine.py
@@ -155,9 +155,13 @@ def __init__(
                     "sparse decode attention with token pipelines requires asynchronous "
                     "scheduling, the fused implementation and key_retention < 1"
                 )
-            if scheduler_config.enable_preemption or scheduler_config.mode != "refill":
+            if scheduler_config.mode != "refill":
+                raise ValueError("sparse decode attention requires refill scheduling")
+            if scheduler_config.enable_preemption and speculative_config is not None:
+                # Ordinary decode snapshots a suspended token's reference;
+                # speculative rows hold per-round references with no snapshot.
                 raise ValueError(
-                    "sparse decode attention requires refill scheduling without preemption"
+                    "sparse decode attention supports preemption only for ordinary decode"
                 )
             # The index buffer never truncates the algorithm: capacity must
             # hold the worst-case retained count over the supported serving
diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py
index d88f7291..de1d8bb0 100644
--- a/vllm_rlt/worker/attention_reference.py
+++ b/vllm_rlt/worker/attention_reference.py
@@ -27,6 +27,7 @@
 rejection. A plan built for a superseded incarnation is stale and rejected.
 """
 
+from contextlib import nullcontext
 from dataclasses import dataclass
 
 import torch
@@ -39,7 +40,14 @@
     SparseDecodePlan,
 )
 
-__all__ = ["ReferenceLease", "AttentionReferenceStore", "DENSE", "CAPTURE", "CONSUME"]
+__all__ = [
+    "ReferenceLease",
+    "ReferenceSnapshot",
+    "AttentionReferenceStore",
+    "DENSE",
+    "CAPTURE",
+    "CONSUME",
+]
 
 
 @dataclass(frozen=True)
@@ -63,6 +71,48 @@ def __len__(self) -> int:
         return len(self.slots)
 
 
+@dataclass(frozen=True)
+class ReferenceSnapshot:
+    """Owned copy of a lease's slots, for suspension across preemption.
+
+    Every per-slot field a later reader or the retirement path looks at is
+    copied, so a restored slot is indistinguishable from the original one
+    except for its identity (slot, generation and incarnation are new).
+    ``host_counts`` keeps the published retained counts of each row.
+    """
+
+    selected: torch.Tensor
+    counts: torch.Tensor
+    mass: torch.Tensor
+    output: torch.Tensor
+    contribution: torch.Tensor
+    valid: torch.Tensor
+    uses: torch.Tensor
+    position: torch.Tensor
+    accepted: torch.Tensor
+    host_counts: tuple[dict[int, int], ...]
+
+    def __len__(self) -> int:
+        return self.counts.shape[0]
+
+    def to(self, device) -> "ReferenceSnapshot":
+        tensors = {
+            name: getattr(self, name).to(device)
+            for name in (
+                "selected",
+                "counts",
+                "mass",
+                "output",
+                "contribution",
+                "valid",
+                "uses",
+                "position",
+                "accepted",
+            )
+        }
+        return ReferenceSnapshot(**tensors, host_counts=self.host_counts)
+
+
 class AttentionReferenceStore:
     """Device-resident reference state for the live tokens of one stage.
 
@@ -269,6 +319,78 @@ def discard(self) -> None:
         self._counts_host.clear()
         self.free = list(range(self.slots))
 
+    def _owner(self):
+        if self.stream is None:
+            return nullcontext()
+        return torch.cuda.stream(self.stream)
+
+    def snapshot(self, lease: ReferenceLease, device) -> ReferenceSnapshot:
+        """Copy a live lease's slots into storage the snapshot owns.
+
+        Suspension calls this after the runner drained, before it releases the
+        lease. The copy runs on the owner stream and completes before return,
+        so the released slots can be reused at once.
+        """
+        for slot, incarnation in zip(lease.slots, lease.incarnations):
+            if self._bindings.get(slot) != (lease.request_id, lease.generation, incarnation):
+                raise RuntimeError("cannot snapshot a retired attention reference")
+        with self._owner():
+            ids = lease.slot_ids
+            snapshot = ReferenceSnapshot(
+                selected=self.selected.index_select(0, ids).to(device),
+                counts=self.counts.index_select(0, ids).to(device),
+                mass=self.mass.index_select(0, ids).to(device),
+                output=self.output.index_select(0, ids).to(device),
+                contribution=self.contribution.index_select(0, ids).to(device),
+                valid=self.valid.index_select(0, ids).to(device),
+                uses=self.uses.index_select(0, ids).to(device),
+                position=self.position.index_select(0, ids).to(device),
+                accepted=self.accepted.index_select(0, ids).to(device),
+                host_counts=tuple(
+                    {
+                        layer: self._counts_host[(slot, layer)]
+                        for layer in range(self.layers)
+                        if (slot, layer) in self._counts_host
+                    }
+                    for slot in lease.slots
+                ),
+            )
+        if self.stream is not None:
+            self.stream.synchronize()
+        return snapshot
+
+    def restore(
+        self, snapshot: ReferenceSnapshot, request_id: str, request_index: int, generation: int
+    ) -> ReferenceLease:
+        """Bind fresh slots for a resumed token and copy its snapshot back."""
+        lease = self.reserve(request_id, request_index, len(snapshot), generation)
+        with self._owner():
+            self.ensure_key_capacity(snapshot.selected.shape[-1])
+            ids = lease.slot_ids
+            width = snapshot.selected.shape[-1]
+            # Only the snapshot's prefix is meaningful; a later capture
+            # overwrites what it uses, exactly as for a never-suspended slot.
+            selected = self.selected[ids]
+            selected[..., :width] = snapshot.selected.to(self.device)
+            self.selected.index_copy_(0, ids, selected)
+            for name in (
+                "counts",
+                "mass",
+                "output",
+                "contribution",
+                "valid",
+                "uses",
+                "position",
+                "accepted",
+            ):
+                getattr(self, name).index_copy_(0, ids, getattr(snapshot, name).to(self.device))
+        if self.stream is not None:
+            self.stream.synchronize()
+        for slot, counts in zip(lease.slots, snapshot.host_counts):
+            for layer, count in counts.items():
+                self._counts_host[(slot, layer)] = count
+        return lease
+
     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 d82af179..41a084b1 100644
--- a/vllm_rlt/worker/model_runner.py
+++ b/vllm_rlt/worker/model_runner.py
@@ -1446,6 +1446,35 @@ def _free_runner_slot(self, request_id):
         if slot is not None:
             self.free_state_slots.append(slot)
 
+    def snapshot_attention_reference(self, request, device):
+        """Copy the live reference of a token suspended between its loops.
+
+        Call before ``suspend``. A token suspended after its reference loop
+        must consume that same reference when it resumes; restarting it dense
+        would be safe but would change its output.
+        """
+        state = self.attention_references.get(request.request_id)
+        if state is None:
+            return None
+        if state[0] is not request:
+            raise RuntimeError("attention reference belongs to another request generation")
+        event = self.events.get(request.request_id)
+        if event is not None:
+            event.synchronize()
+        return state[1], self.reference_store.snapshot(state[2], device)
+
+    def restore_attention_reference(self, request, saved):
+        if saved is None:
+            return
+        if request.request_id in self.attention_references:
+            raise RuntimeError("cannot restore over a live attention reference")
+        position, snapshot = saved
+        self.reference_generation += 1
+        lease = self.reference_store.restore(
+            snapshot, request.request_id, 0, self.reference_generation
+        )
+        self.attention_references[request.request_id] = (request, position, lease)
+
     def suspend(self, request_id):
         """Free the request's runner resources but keep its sampling RNG.
 

From 0f0569c974a836eb1b7557d92ae4189cf53e357a Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Mon, 2 Nov 2026 10:11:08 +0100
Subject: [PATCH 31/32] Preempt and resume sparse token pipelines with their
 reference slots

A suspended token pipeline keeps several live rows, each past or before its
capture loop, and each row reads the reference slot the device permutation
gives it. The pipeline snapshot now carries the request's whole reference
lease and, per row, which lease slot it reads; restoration binds a fresh lease
and generation, copies the slots back and rebuilds the permutation, so every
row resumes with the reference it captured. Live rows must form a prefix, as
the restored pipeline puts them back in rows 0..k-1.

Restoring a reference now reserves its slots on the owner stream as well:
the slot ids are copied on the current stream and must be on the stream that
indexes with them.

Preemption is lifted for token pipelines; block speculation with sparse
attention still rejects it.

Test (CUDA): a priority arrival suspends a sparse pipeline with rows past
their capture loop, with and without graphs, resident and offloaded; it
resumes to the same tokens as an uninterrupted run and releases every slot.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_sparse_decode_preemption.py | 76 ++++++++++++++++++++++++++
 vllm_rlt/engine/stage_engine.py        | 13 +++--
 vllm_rlt/worker/async_speculative.py   | 32 ++++++++++-
 vllm_rlt/worker/attention_reference.py |  4 +-
 vllm_rlt/worker/speculative.py         |  6 ++
 5 files changed, 123 insertions(+), 8 deletions(-)

diff --git a/tests/test_sparse_decode_preemption.py b/tests/test_sparse_decode_preemption.py
index d7d09c88..9210c7ca 100644
--- a/tests/test_sparse_decode_preemption.py
+++ b/tests/test_sparse_decode_preemption.py
@@ -174,3 +174,79 @@ def run(cache_blocks, preemption):
     torch.cuda.synchronize()
     assert preemptions > 0
     assert pressured == baseline
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("graphs", [False, True])
+@pytest.mark.parametrize("offload", [False, True])
+def test_cuda_pipeline_preemption_restores_each_rows_reference(graphs, offload):
+    """A token pipeline suspends with rows past their capture loop and resumes them sparse."""
+    from vllm_rlt import ExecutionConfig, SparseDecodeAttentionConfig
+
+    def make():
+        return engine(
+            retention=None,
+            device="cuda",
+            speculative_config=SpeculativeConfig(4, draft_loops=1, mode="pipeline"),
+            scheduler_config=_scheduler(policy="priority"),
+            execution_config=ExecutionConfig(
+                async_scheduling=True,
+                cuda_graphs=graphs,
+                static_buffers=True,
+                pad_to_power_of_two=True,
+            ),
+            sparse_decode_attention_config=SparseDecodeAttentionConfig(
+                enabled=True, key_retention=0.5, implementation="fused"
+            ),
+        )
+
+    prompt = [1, 2, 3]
+    low = SamplingParams(max_tokens=12, ignore_eos=True, priority=10)
+    e = make()
+    try:
+        e.add_request("low", prompt, low)
+        expected = _drain(e, {})["low"]
+    finally:
+        e.close()
+
+    e = make()
+    try:
+        e.add_request("low", prompt, low)
+        request = e.scheduler.requests["low"]
+        runner = e.speculative_runner
+        for _ in range(200):
+            e.step()
+            state = runner.states.get("low")
+            # Rows reach their consume loops a few submissions after the root.
+            if state is not None and state.submitted >= 4:
+                break
+        else:
+            pytest.fail("did not submit pipeline steps")
+        e.add_request("high", [5], SamplingParams(max_tokens=3, ignore_eos=True, priority=0))
+        for _ in range(200):
+            e.step()
+            if "low" in e.preemption.snapshots:
+                break
+        else:
+            pytest.fail("priority arrival did not suspend the pipeline")
+        saved = e.preemption.snapshots["low"]["pipeline"]
+        assert saved.references is not None
+        references, permutation = saved.references
+        assert sorted(permutation) == list(range(len(references)))
+        # Some live row is past its capture loop, so its reference matters.
+        assert any(t.depth > 1 for t in saved.tokens)
+        assert "low" not in runner.states
+        if offload:
+            e.preemption.suspend(request, offload_kv=True)
+            saved = e.preemption.snapshots["low"]["pipeline"]
+            assert saved.references[0].selected.device.type == "cpu"
+        outputs = _drain(e, {})
+        assert outputs["low"] == expected
+        assert e.preemption.resumptions >= 1
+        assert not e.preemption.snapshots
+        assert e.reference_store.outstanding == 0
+        assert e.cache_manager.num_used_blocks == 0
+        if graphs:
+            assert runner.graphs.replays > 0
+    finally:
+        e.close()
diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py
index 6c8d62e3..0140f73f 100644
--- a/vllm_rlt/engine/stage_engine.py
+++ b/vllm_rlt/engine/stage_engine.py
@@ -157,11 +157,16 @@ def __init__(
                 )
             if scheduler_config.mode != "refill":
                 raise ValueError("sparse decode attention requires refill scheduling")
-            if scheduler_config.enable_preemption and speculative_config is not None:
-                # Ordinary decode snapshots a suspended token's reference;
-                # speculative rows hold per-round references with no snapshot.
+            if (
+                scheduler_config.enable_preemption
+                and speculative_config is not None
+                and not self.pipeline_speculative
+            ):
+                # Ordinary decode and token pipelines snapshot their live
+                # references; block rounds hold per-round leases with none.
                 raise ValueError(
-                    "sparse decode attention supports preemption only for ordinary decode"
+                    "sparse decode attention supports preemption only for ordinary "
+                    "decode and token pipelines"
                 )
             # The index buffer never truncates the algorithm: capacity must
             # hold the worst-case retained count over the supported serving
diff --git a/vllm_rlt/worker/async_speculative.py b/vllm_rlt/worker/async_speculative.py
index f89df1f8..95857fa5 100644
--- a/vllm_rlt/worker/async_speculative.py
+++ b/vllm_rlt/worker/async_speculative.py
@@ -760,6 +760,17 @@ def snapshot_pipeline(self, request, device):
             (pipeline.tokens, pipeline.depths, pipeline.candidates, pipeline.active.long())
         ).tolist()
         tokens, depths, candidates, active = values
+        references = None
+        if self.reference_store is not None:
+            if state.reference is None:
+                raise RuntimeError("sparse pipeline lost its reference slots")
+            # Restoration puts the k live tokens back in rows 0..k-1; keep the
+            # slot of each row by recording it relative to the lease.
+            if not all(active[: sum(active)]):
+                raise RuntimeError("live pipeline rows must form a prefix")
+            lease_rows = {slot: row for row, slot in enumerate(state.reference.slots)}
+            permutation = tuple(lease_rows[slot] for slot in pipeline.slots.tolist())
+            references = (self.reference_store.snapshot(state.reference, device), permutation)
         return PipelineState(
             request,
             position,
@@ -774,6 +785,7 @@ def snapshot_pipeline(self, request, device):
                 for i, token in enumerate(tokens)
                 if active[i]
             ],
+            references=references,
         )
 
     def restore_pipeline(self, request, snapshot):
@@ -781,9 +793,8 @@ def restore_pipeline(self, request, snapshot):
             return
         if snapshot.request is not request or request.request_id in self.states:
             raise RuntimeError("invalid pipeline restoration generation")
-        if self.reference_store is not None:
-            # Sparse decode excludes preemption; a snapshot carries no references.
-            raise RuntimeError("sparse decode attention cannot restore a pipeline snapshot")
+        if (self.reference_store is not None) != (snapshot.references is not None):
+            raise RuntimeError("pipeline snapshot and sparse decode configuration disagree")
         if snapshot.pending or snapshot.position != request.position:
             raise RuntimeError("pipeline snapshot disagrees with committed request position")
         width = min(
@@ -817,12 +828,27 @@ def restore_pipeline(self, request, snapshot):
                 values[2],
                 torch.arange(width, device=self.device) < len(snapshot.tokens),
             )
+            generation = self._generations.get(request.request_id, 0) + 1
+            self._generations[request.request_id] = generation
+            reference = None
+            if snapshot.references is not None:
+                saved, permutation = snapshot.references
+                if len(saved) != width or sorted(permutation) != list(range(width)):
+                    raise RuntimeError("pipeline reference snapshot does not match the window")
+                # Fresh slots and generation; each row reads the slot it read
+                # before suspension.
+                reference = self.reference_store.restore(saved, request.request_id, 0, generation)
+                pipeline.slots = reference.slot_ids[
+                    torch.tensor(permutation, device=self.device)
+                ].clone()
             self.states[request.request_id] = SpeculativeRequestState(
                 request,
                 self.cache._get_allocation(request.request_id),
                 torch.tensor(request.input_token_id, device=self.device),
                 torch.tensor(snapshot.position, device=self.device),
                 pipeline,
+                generation=generation,
+                reference=reference,
             )
 
     def request_ready(self, request):
diff --git a/vllm_rlt/worker/attention_reference.py b/vllm_rlt/worker/attention_reference.py
index de1d8bb0..cadc311c 100644
--- a/vllm_rlt/worker/attention_reference.py
+++ b/vllm_rlt/worker/attention_reference.py
@@ -363,8 +363,10 @@ def restore(
         self, snapshot: ReferenceSnapshot, request_id: str, request_index: int, generation: int
     ) -> ReferenceLease:
         """Bind fresh slots for a resumed token and copy its snapshot back."""
-        lease = self.reserve(request_id, request_index, len(snapshot), generation)
         with self._owner():
+            # Reserve here too: its slot ids are copied on the current stream,
+            # which must be the stream that indexes with them below.
+            lease = self.reserve(request_id, request_index, len(snapshot), generation)
             self.ensure_key_capacity(snapshot.selected.shape[-1])
             ids = lease.slot_ids
             width = snapshot.selected.shape[-1]
diff --git a/vllm_rlt/worker/speculative.py b/vllm_rlt/worker/speculative.py
index 5828e11f..06bfe557 100644
--- a/vllm_rlt/worker/speculative.py
+++ b/vllm_rlt/worker/speculative.py
@@ -120,6 +120,9 @@ class PipelineState:
     position: int
     tokens: list[PipelineToken]
     pending: bool = False
+    # Sparse decode: the request's reference slots and, per pipeline row, the
+    # index of the slot it reads inside that snapshot.
+    references: tuple | None = None
 
     def to(self, device):
         """Own hidden copies for suspension; never borrow graph/runner buffers."""
@@ -131,6 +134,9 @@ def to(self, device):
                 )
                 for token in self.tokens
             ],
+            references=None
+            if self.references is None
+            else (self.references[0].to(device), self.references[1]),
         )
 
 

From 44737aa2db3135978b52ab3d009b63eefd4704a6 Mon Sep 17 00:00:00 2001
From: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Date: Mon, 2 Nov 2026 11:53:15 +0100
Subject: [PATCH 32/32] Walk pipeline suspensions through the cycle in the
 sparse preemption test

The schedule is deterministic and a random tiny model rejects most drafts, so
a single suspension landed at the same point of the pipeline cycle in every
parametrization, with no row past its capture loop; the test stopped at its
precondition before checking any output. It now suspends repeatedly, one more
submission later each time, until a suspension holds such a row, checks the
saved permutation at every suspension, and then compares the output with an
uninterrupted run.

Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
---
 tests/test_sparse_decode_preemption.py | 69 ++++++++++++++++----------
 1 file changed, 42 insertions(+), 27 deletions(-)

diff --git a/tests/test_sparse_decode_preemption.py b/tests/test_sparse_decode_preemption.py
index 9210c7ca..a8ef0d1f 100644
--- a/tests/test_sparse_decode_preemption.py
+++ b/tests/test_sparse_decode_preemption.py
@@ -201,7 +201,7 @@ def make():
         )
 
     prompt = [1, 2, 3]
-    low = SamplingParams(max_tokens=12, ignore_eos=True, priority=10)
+    low = SamplingParams(max_tokens=64, ignore_eos=True, priority=10)
     e = make()
     try:
         e.add_request("low", prompt, low)
@@ -214,33 +214,48 @@ def make():
         e.add_request("low", prompt, low)
         request = e.scheduler.requests["low"]
         runner = e.speculative_runner
-        for _ in range(200):
-            e.step()
-            state = runner.states.get("low")
-            # Rows reach their consume loops a few submissions after the root.
-            if state is not None and state.submitted >= 4:
-                break
-        else:
-            pytest.fail("did not submit pipeline steps")
-        e.add_request("high", [5], SamplingParams(max_tokens=3, ignore_eos=True, priority=0))
-        for _ in range(200):
-            e.step()
-            if "low" in e.preemption.snapshots:
-                break
-        else:
-            pytest.fail("priority arrival did not suspend the pipeline")
-        saved = e.preemption.snapshots["low"]["pipeline"]
-        assert saved.references is not None
-        references, permutation = saved.references
-        assert sorted(permutation) == list(range(len(references)))
-        # Some live row is past its capture loop, so its reference matters.
-        assert any(t.depth > 1 for t in saved.tokens)
-        assert "low" not in runner.states
-        if offload:
-            e.preemption.suspend(request, offload_kv=True)
+        outputs, deep = {}, 0
+
+        def step():
+            for output in e.step():
+                if output.finished:
+                    outputs[output.request_id] = list(output.token_ids)
+
+        # A random tiny model rejects most drafts, so a pipeline often restarts
+        # from its root. Suspend repeatedly, one more submission later each time
+        # so the suspension point walks through the pipeline's cycle, until one
+        # suspension holds a row past its capture loop.
+        for attempt in range(8):
+            for _ in range(200):
+                step()
+                state = runner.states.get("low")
+                if state is not None and state.submitted >= 2 + attempt:
+                    break
+            else:
+                pytest.fail("the low request did not run pipeline steps")
+            e.add_request(
+                f"high{attempt}", [5], SamplingParams(max_tokens=3, ignore_eos=True, priority=0)
+            )
+            for _ in range(200):
+                step()
+                if "low" in e.preemption.snapshots:
+                    break
+            else:
+                pytest.fail("priority arrival did not suspend the pipeline")
             saved = e.preemption.snapshots["low"]["pipeline"]
-            assert saved.references[0].selected.device.type == "cpu"
-        outputs = _drain(e, {})
+            references, permutation = saved.references
+            assert sorted(permutation) == list(range(len(references)))
+            assert "low" not in runner.states
+            if offload:
+                e.preemption.suspend(request, offload_kv=True)
+                saved = e.preemption.snapshots["low"]["pipeline"]
+                assert saved.references[0].selected.device.type == "cpu"
+            deep += any(t.depth > 1 for t in saved.tokens)
+            if deep:
+                break
+        assert deep, "no suspension held a row past its capture loop"
+        assert "low" not in outputs
+        outputs.update(_drain(e, {}))
         assert outputs["low"] == expected
         assert e.preemption.resumptions >= 1
         assert not e.preemption.snapshots