Sparse decode attention: per-token reference lifecycle for async self-speculation

README.md

@@ -80,8 +80,10 @@ currently **same-host CUDA** with supported Triton/FA4 paths; it is not a claim
 of universal multi-node or model support. Speculative steps are not included
 in mixed prefill/decode batches.
 
+
 <a id="features"></a>
 
+
 ### Designed for recurrent inference
 
 - **Depth-aware paged KV cache:** LAST_EXITED and SHARED layouts, with optional
@@ -92,6 +94,14 @@ in mixed prefill/decode batches.
   CUDA Graphs, and optional asynchronous/speculative paths. An optional native
   CUDA SiLU-and-Mul operator accelerates the Ouro activation path, with a Torch
   fallback; it is not a general model-speedup claim.
+- **Sparse decode reference integration (opt-in):** cached loop-2 attention
+  correction for ordinary decode and block self-speculation, with shared
+  per-token reference ownership and explicit capability limits. The cross-loop
+  sparse attention mechanism comes from FlashLoop
+  ([paper](https://arxiv.org/abs/2609.29812),
+  [reference implementation](https://github.com/Superone77/FlashLoop)); this
+  repository contributes the kernel and serving integration. See the
+  [configuration and scope](docs/runtime.md#sparse-decode-attention-reference-path).
 - **Offline and online use:** Python API, CLI, and a streaming
   OpenAI-compatible **completions** endpoint. This is not a chat-completions API.
 

docs/runtime.md

@@ -230,6 +230,67 @@ 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. The cross-loop sparse attention mechanism is from
+FlashLoop ([paper](https://arxiv.org/abs/2609.29812),
+[reference implementation](https://github.com/Superone77/FlashLoop)); the
+configuration, kernels and serving integration here are this repository's.
+Prompt prefill remains dense. Loop 2 captures
+per-token/per-layer/head selection, mass and contributions; loops 3 and 4 reuse
+that reference with their current selected KV. Each new input token establishes
+its own reference.
+
+```python
+from vllm_rlt import ExecutionConfig, SparseDecodeAttentionConfig, SpeculativeConfig
+
+sparse = SparseDecodeAttentionConfig(enabled=True, key_retention=0.1)
+# Pass sparse_decode_attention_config=sparse to LLM or StageEngine.
+# Ordinary decode uses synchronous scheduling and no speculative_config.
+# For synchronous block speculation, also pass:
+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. 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.
+
+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 over the model's positional
+limit, which is the only context bound the reference needs. The index buffer
+starts small and grows on the owner stream, so a long model context does not
+cost its worst case up front; admission never rejects a request for being long.
+The pool holds one slot per live ordinary input token, K+1 per synchronous
+block request, and two rounds of K+1 per asynchronous block request. All
+reference payload bytes contribute to automatic KV memory planning; explicit
+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 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. The
+[numerical diagnosis](https://github.com/Levius-Fubuki/vllm-rlt/releases/tag/pr157-numerics-validation-20261029)
+is kept as a release attachment rather than a checked-in report. Retention below
+1 is approximate and
+needs task-quality evaluation before a deployment recommendation.
+
 ## SHARED prefill semantics
 
 SHARED retains one physical KV plane and completes every loop for a prompt

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,
     )
 
 

tests/test_async_speculative.py

@@ -173,7 +173,7 @@ def forbidden(*args, **kwargs):
 
 @pytest.mark.gpu
 @pytest.mark.parametrize("graphs", [False, True])
-def test_eos_discards_queued_round_and_request_id_can_be_reused(monkeypatch, graphs):
+def test_eos_discards_queued_round_and_the_id_frees_after_retirement(monkeypatch, graphs):
     model = tiny("cuda")
     e = make_engine(model, graphs=graphs)
     bootstrap(e, ignore_eos=False)
@@ -191,7 +191,17 @@ def test_eos_discards_queued_round_and_request_id_can_be_reused(monkeypatch, gra
     e.step()
     out = e._collect_speculative()[0]
     assert out.finished and out.token_ids[-1] == 0 and len(out.token_ids) == 2
-    # Reuse while an old ticket still exists. It must never touch the new request.
+    # The engine never reuses an id while the previous token may still own
+    # resources. Whether the round has already retired by now depends on the
+    # graph setting, so assert the rejection only while the engine still lists
+    # the id; the live-duplicate case is covered by test_stage_engine.
+    if "r" in e._retiring_speculative:
+        with pytest.raises(ValueError, match="duplicate request ID"):
+            e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True))
+    # Once retirement completes the id is free, and the new request never sees
+    # the old round's state.
+    e._retire_speculative(wait=True)
+    assert "r" not in e._retiring_speculative
     monkeypatch.setattr(e.speculative_runner, "_coda", coda)
     e.add_request("r", [7, 8], SamplingParams(max_tokens=5, ignore_eos=True))
     output = drain(e)["r"]
@@ -345,6 +355,8 @@ def test_active_page_tables_are_contiguous_and_do_not_touch_inactive_storage():
     assert cache.allocate("short", 4) and cache.allocate("long", 10)
     runner = object.__new__(AsyncSpeculativeRunner)
     runner.cache = cache
+    runner.reference_store = None
+    runner.sparse_decode_config = None
     buffer = SimpleNamespace(
         leased=False,
         host_tables=torch.full((128,), -91, dtype=torch.int32),

tests/test_async_speculative_lifecycle.py

@@ -17,9 +17,9 @@
 from vllm_rlt.request import RequestPhase
 from vllm_rlt.serving.worker import EngineWorker
 from vllm_rlt.worker.async_speculative import (
+    AsyncSpeculativeOutput,
     AsyncSpeculativeRunner,
     SpeculativeRequestState,
-    SpeculativeRoundTicket,
 )
 
 
@@ -105,7 +105,7 @@ def prepared_engine(*, bad_second=False, first_frontier=4, retirement_ready=True
             host_result=host,
         )
         tickets.append(
-            SpeculativeRoundTicket(
+            AsyncSpeculativeOutput(
                 SchedulerOutput(RequestPhase.SPECULATIVE, [scheduled_request]),
                 buffer,
                 (state,),
@@ -382,7 +382,7 @@ def test_ticket_result_format_does_not_depend_on_statistics(mode, with_stats):
         host_result=torch.tensor(values), copy_done=Completion(), leased=True, keepalive=[]
     )
     stats = SpeculativeStats() if with_stats else None
-    ticket = SpeculativeRoundTicket(batch, buffer, (), (), 2, mode=mode, stats=stats)
+    ticket = AsyncSpeculativeOutput(batch, buffer, (), (), 2, mode=mode, stats=stats)
     result = ticket.collect()
     assert result.speculative[0].token_ids == [11]
     assert result.speculative[0].accepted_count == 0

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",
     [