Repository navigation
Sparse decode attention: per-token reference lifecycle for async self-speculation - #157
kunkunblueberry wants to merge 32 commits into
Conversation
…-speculation Implements the asynchronous self-speculation cell of the sparse decode attention column (RFC InferMatrixLab#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.
|
https://github.com/Superone77/FlashLoop is an important reference, could you please add more e2e test result? |
|
Hi @kunkunblueberry, I’ve reviewed #157 and would like to reuse its reference plan/store for the basic and synchronous paths. Before wiring attention into it, could we align on three points?
If this sounds reasonable, I’ll build against the shared contract and cover retained-count boundaries, bucket overflow/full retention, and capacity accounting in the basic-path tests. |
Thanks,i will check it and add it carefully. |
…rate capacity Aligns with the review on InferMatrixLab#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.
|
Thanks @Levius-Fubuki — this is exactly the alignment I was hoping for. All three points made sense, and I pushed
Please build against the shared contract: the store publishes the per-(slot, layer, head) |
For precision test, please use flashloop_torch implementation as gold result, thanks! |
…dation 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.
|
Done in What I did: the lifecycle layer alone cannot show the optimization works, so I added How it validates:
Honest reading: the algorithmic saving is real (10% of keys touched, consume cost scales with Regression: CPU |
|
Could you provide TTFT result? To evaluate the prefill. Thanks! |
…tention Prefill rows never take the decode correction (RFC InferMatrixLab#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.
|
TTFT measurement added ( prompt 64, max_tokens 32
prompt 100, max_tokens 16 (prefill-dominated)
Reading: TTFT is unchanged at 1.00–1.01x, i.e. within noise, which is what §6.3 predicts — prefill rows never take the decode correction, and the first token still comes from the prefill/coda path, so enabling the feature does not tax the latency users feel first. With a longer prompt the absolute TTFT doubles while the delta stays ~0.4 ms, so this is not a coincidence of a short prompt. What it does cost today: ITL goes 15 ms → 120–128 ms (about 8x) because the reference path replaces the Triton kernel with plain PyTorch row loops (gather + argsort + per-step launch/sync). That is the reference implementation's cost, not the algorithm's — and it sets the bar for the fused kernel: it has to beat the 15 ms dense Triton ITL to be worth enabling. Note Until that kernel lands I would keep |
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
Requesting changes for the reproducible memory-capacity failure in the acceptance path. The current implementation can construct the engine and pass a short warmup, then OOM while processing an ordinary admitted concurrent batch. Please make acceptance/invalidation memory-bounded, account for the supported reference capacity before admitting work, and add coverage for the overlapping-round case below before merging.
The blocking inline finding is on AttentionReferenceStore.accept(). The performance observations and full-retention token divergence below are additional evidence: the token divergence has not been localized to a numerical or lifecycle cause, and this review does not require a plain-PyTorch reference to deliver an optimized-kernel speedup.
The sequential tiny-model TTFT observations in the latest comment also should not be generalized to concurrent serving. In the real-model B=64 case below, p95 engine first-token delivery changes from about 2.39 s (off) to 37.82/38.66 s (r=.1), even though prefill does not use sparse correction. These measurements include scheduling and output delivery and are not isolated prefill-kernel timings.
Supplementary real-model concurrency results for this PR.
I executed db5ce44642402102cc2a879196aac34d8a4026b5. While the campaign ran, HEAD advanced to f53c35d3bfa5b2412d4993018b1c07fd4bafd139; that commit only adds benchmarks/ttft_bench.py. The vllm_rlt/ and tests/ Git trees are identical between those two revisions. The new benchmark itself was not executed in this campaign.
Outcome: 10 measured cases completed, one measured case failed with CUDA OOM, and its repeat was skipped according to the frozen plan. The existing sparse test module passed all 28 tests, including its GPU tiny-model tests. The real-model results below expose additional capacity and output-equivalence limits.
Reproduction configuration
- One NVIDIA GH200, 97871 MiB visible HBM, ARM CPU; driver 595.58.03, Python 3.12.3, PyTorch 2.14.0+cu130, Triton 3.8.0.
- Native Ouro-1.4B, BF16 weights/activations/KV, unmodified
max_position_embeddings=65536. - Model weight SHA256:
58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af; config SHA256:ce9cc13da41591b8b4deca053d7dfee06424c0228628ee862ea86d725bc163f3. - Native StageEngine, not an HTTP server. For each batch size, enqueue all requests before calling
step(). - Requests
i=0..B-1use 256 deterministic prompt IDs:[1+(j*137+i*53)%49000 for j in range(256)]. - Greedy sampling,
max_tokens=16,max_loops=4,exit_threshold=1.0,ignore_eos=True. - Per case: fresh process/model/engine, same-concurrency 4-output-token warmup, then the 16-token measured batch. Startup and warmup are excluded from measurements. CPU/OMP threads=4; CPU affinity and GPU clocks were not pinned.
StageEngine(
model,
attention_backend="triton",
cache_config=CacheConfig(
num_blocks=2368 if B == 32 else 4672,
block_size=16,
incremental_allocation=True,
),
scheduler_config=SchedulerConfig(
max_num_seqs=B,
max_num_batched_tokens=4096,
prefill_chunk_size=256,
admission_scan_limit=128,
),
execution_config=ExecutionConfig(
async_scheduling=True,
cuda_graphs=False,
),
speculative_config=SpeculativeConfig(3), # draft 2 / target 4, block mode
sparse_decode_attention_config=(
None if mode == "off"
else SparseDecodeAttentionConfig(enabled=True, key_retention=r)
),
)key_bucket=None; prefix caching and preemption are disabled. The physical KV capacity covers all requests and four recurrent planes. Per batch size, the order was OFF1 / r=1 FULL1 / r=.1 SPARSE1 / SPARSE2 / FULL2 / OFF2. Each case had a 300-second wall limit including startup/warmup. A failed FULL1 was not repeated; independent configurations continued.
Measured observations
These are end-to-end native-engine output tokens/s, including enqueue, prefill, decode, output delivery and drain. They are not isolated attention-kernel timings. Both observations are retained because there is timing variability.
| Requests | Mode | Output tokens/s, run 1 / 2 | p95 engine TTFT seconds, run 1 / 2 | Peak allocated / reserved GiB |
|---|---|---|---|---|
| 32 | Off | 134.17 / 160.51 | 1.078 / 1.042 | 10.62 / 11.36 |
| 32 | r=1 | 7.43 / 6.97 | 6.854 / 6.985 | 58.09 / 80.45 |
| 32 | r=.1 | 14.45 / 10.44 | 5.339 / 8.070 | 14.89 / 17.46 |
| 64 | Off | 146.37 / 157.12 | 2.390 / 2.390 | 17.37 / 18.11 |
| 64 | r=1 | OOM / repeat not run | — | 82.94 allocated at OOM; another 18.00 requested |
| 64 | r=.1 | 9.40 / 9.94 | 37.819 / 38.664 | 26.45 / 35.02 |
Every successful case completed all B requests with exactly 16 output tokens. Both observed maximum resident requests and maximum requests in one speculative batch reached B=32/64. This was actual concurrency, not merely a larger pending queue. Sparse reference slots reached 256/256 at B=32 and 445/512 at B=64. Capture and consume path counters were nonzero. KV blocks and reference bindings drained to zero in every successful case, and same-mode repeats produced identical complete token/depth outputs.
The r=.1 observations do not show an acceleration on this short-context workload: throughput was approximately 6.3–6.4% of the paired dense controls at B=64. This is consistent with the PR's current scope as a plain PyTorch reference, not an optimized GPU serving kernel. Output changes also affect speculative acceptance, so the end-to-end ratio cannot isolate attention arithmetic cost.
Concrete B=64 full-retention failure
The 4-token warmup completed 64/64, but the 16-token measurement failed in AttentionReferenceStore.accept():
self.selected[slot_ids] = self.selected[slot_ids] * keep4.to(self.selected.dtype)CUDA OOM requested 18.00 GiB, with 11.06 GiB free. PyTorch reported 82.94 GiB allocated and only 213.42 MiB reserved but unallocated. The last completed step observation had 64 resident requests and 0 completed measurement requests. Successful warmup is therefore not a successful measured case.
The reference pool alone is 24.09/48.19 GiB for B=32/64 at r=1, versus 2.49/4.99 GiB at r=.1. Pool sizing follows the model's 65,536-token limit even for a 256-token prompt, and the indexed assignment above introduces substantial transient allocation. No model-context reduction or allocator workaround was applied. Avoiding full-bucket indexed temporaries and making the service context capacity explicit are useful follow-up areas, but were not implemented or validated here.
Output equivalence
Full retention was not token-identical to dense on this larger BF16 workload:
| Comparison to same-head dense | Matching token positions | Identical complete sequences |
|---|---|---|
| B=32, r=1, both repeats | 503/512 | 31/32 |
| B=32, r=.1, both repeats | 259/512 | 10/32 |
| B=64, r=.1, both repeats | 450/1024 | 16/64 |
For r=1, request i=23 first differs at zero-based output position 7: dense token 30 versus full-retention token 2578. Both dense repeats are identical; both full-retention repeats are identical; exit depths match. This is reproducible autoregressive output divergence, not yet a localization of its numerical cause. No same-input logits-margin or FlashLoop-gold comparison was performed, so this observation alone does not identify an arithmetic or lifecycle bug.
These are synthetic-prompt generation comparisons, not benchmark task accuracy. Long contexts, long generation, HTTP serving, CUDA graphs and production performance remain outside this finite campaign.
All raw results, including the OOM, warmups, per-step traces, complete token/depth arrays and source/model/package provenance, were preserved in a 51-file archive and independently SHA-verified both remotely and after download. The original GPU lock was released, all owned children were joined, and the final GPU compute-process list was empty. Archive SHA256: f01759c21b454e936fc26bcc83ba39b172ccb2d3b202bf19e64f8d22df9ef8db.
Standalone native-engine reproduction harness used for the measurements
Run this against the pinned runtime revision with the local native Ouro-1.4B checkpoint above. For example, the failure case is python concurrency_probe.py --model /path/to/Ouro-1.4B --out b64-full1.json --batch 64 --retention 1.0 --context 256 --decode 16 --warmup 4. Use retention 0 for the off control or 0.1 for the sparse case. Run only within an allocated/exclusively owned GPU slot. The campaign wrapped each invocation in a 300-second timeout and isolated cases in separate processes.
"""Frozen native-engine concurrency probe; no HTTP or dataset accuracy claims."""
import argparse
import collections
import json
import math
import time
import traceback
from pathlib import Path
import torch
from vllm_rlt import CacheConfig, ExecutionConfig, SamplingParams, SchedulerConfig
from vllm_rlt import SparseDecodeAttentionConfig, SpeculativeConfig
from vllm_rlt.engine.stage_engine import StageEngine
from vllm_rlt.models import OuroForCausalLM
from vllm_rlt.request import RequestPhase
import vllm_rlt.core.kv_cache_manager as kv_module
p = argparse.ArgumentParser()
p.add_argument('--model', required=True)
p.add_argument('--out', required=True)
p.add_argument('--batch', type=int, required=True)
p.add_argument('--retention', type=float, required=True, help='0 means feature off')
p.add_argument('--context', type=int, default=256)
p.add_argument('--decode', type=int, default=16)
p.add_argument('--warmup', type=int, default=4)
a = p.parse_args()
torch.set_num_threads(4)
torch.manual_seed(157)
out = Path(a.out)
record = dict(status='STARTING', args=vars(a), head='db5ce44642402102cc2a879196aac34d8a4026b5',
torch=torch.__version__, gpu=torch.cuda.get_device_name(),
scope='native StageEngine synthetic prompts, not HTTP; timing includes prefill; fixed 4 loops')
def save():
out.write_text(json.dumps(record, indent=2) + '\n')
save()
path_counts = collections.Counter()
original = kv_module.sparse_paged_attention
def counted(*args, **kwargs):
plan = args[5]
path_counts['calls'] += 1
for mode, rows in plan.groups().items():
path_counts[str(mode)] += len(rows)
return original(*args, **kwargs)
kv_module.sparse_paged_attention = counted
def generate(engine, decode, label):
path_counts.clear()
start = time.perf_counter()
done, timing, trace = {}, {}, []
for i in range(a.batch):
rid = str(i)
prompt = [1 + (j * 137 + i * 53) % 49000 for j in range(a.context)]
engine.add_request(rid, prompt, SamplingParams(max_tokens=decode, max_loops=4,
exit_threshold=1., temperature=0, ignore_eos=True))
timing[rid] = dict(first=None, finish=None, chunks=[])
max_resident = max_spec_batch = max_ref = max_blocks = 0
with out.with_suffix('.' + label + '.jsonl').open('w') as stream:
for step in range(10000):
if not engine.has_unfinished_requests():
break
resident_before = len(engine.cache_manager._allocations)
outputs = engine.step()
now = time.perf_counter() - start
resident = len(engine.cache_manager._allocations)
max_resident = max(max_resident, resident_before, resident)
schedule = engine.last_schedule
phase = str(schedule.phase) if schedule else None
scheduled = len(schedule.requests) if schedule else 0
if schedule and schedule.phase == RequestPhase.SPECULATIVE:
max_spec_batch = max(max_spec_batch, scheduled)
ref = engine.reference_store.outstanding if engine.reference_store else 0
max_ref = max(max_ref, ref)
blocks = engine.cache_manager.num_used_blocks
max_blocks = max(max_blocks, blocks)
for item in outputs:
t = timing[item.request_id]
if item.token_ids:
if t['first'] is None:
t['first'] = now
t['chunks'].append([now, len(item.token_ids)])
if item.finished:
if item.request_id in done:
raise AssertionError('duplicate completion')
done[item.request_id] = dict(tokens=list(item.token_ids), depths=list(item.exit_depths))
t['finish'] = now
row = dict(step=step, seconds=now, phase=phase, scheduled=scheduled,
resident=resident, ref_slots=ref, used_blocks=blocks, completed=len(done))
stream.write(json.dumps(row) + '\n')
stream.flush()
else:
raise RuntimeError('bounded step limit')
torch.cuda.synchronize()
seconds = time.perf_counter() - start
result = dict(seconds=seconds, tokens_s=sum(len(v['tokens']) for v in done.values())/seconds,
completed=len(done), outputs=done, timing=timing, steps=step,
max_resident=max_resident, max_speculative_batch=max_spec_batch,
max_reference_slots=max_ref, max_used_blocks=max_blocks,
drained_blocks=engine.cache_manager.num_used_blocks,
drained_references=engine.reference_store.outstanding if engine.reference_store else 0,
sparse_path_counts=dict(path_counts), peak_allocated=torch.cuda.max_memory_allocated(),
peak_reserved=torch.cuda.max_memory_reserved())
assert len(done) == a.batch and all(len(v['tokens']) == decode for v in done.values()), 'incomplete batch'
assert result['drained_blocks'] == result['drained_references'] == 0, 'leaked ownership'
if a.retention:
assert path_counts['1'] > 0 and path_counts['2'] > 0, 'sparse path did not execute'
return result
engine = None
try:
model = OuroForCausalLM.from_pretrained(a.model, device='cuda', dtype=torch.bfloat16)
# Sufficient physical pages for all requests and all four recurrent planes.
blocks = a.batch * math.ceil((a.context + a.decode + 4) / 16) * 4 + 64
record['num_blocks'] = blocks
record['model_max_position_embeddings'] = model.config.max_position_embeddings
engine = StageEngine(model, attention_backend='triton',
cache_config=CacheConfig(blocks, 16, incremental_allocation=True),
scheduler_config=SchedulerConfig(max_num_seqs=a.batch, max_num_batched_tokens=4096,
prefill_chunk_size=256, admission_scan_limit=128),
execution_config=ExecutionConfig(async_scheduling=True, cuda_graphs=False),
speculative_config=SpeculativeConfig(3),
sparse_decode_attention_config=SparseDecodeAttentionConfig(enabled=True, key_retention=a.retention)
if a.retention else None)
record['reference_bytes'] = engine.reference_store.bytes if engine.reference_store else 0
record['kv_bytes'] = engine.cache_manager.key_cache.nbytes + engine.cache_manager.value_cache.nbytes
record['status'] = 'WARMUP'
save()
record['warmup_result'] = generate(engine, a.warmup, 'warmup')
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
record['status'] = 'MEASURING'
save()
record['result'] = generate(engine, a.decode, 'measure')
record['status'] = 'COMPLETED'
except BaseException as exc:
record['status'] = 'FAILED'
record['error'] = repr(exc)
record['traceback'] = traceback.format_exc()
save()
raise
finally:
if engine is not None:
engine.close()
record['closed'] = True
save()The raw archive is retained locally; it is not publicly hosted by this review. The configuration, observed results, failure allocation sizes, source identity, and executable harness are included above so the failure can be independently reproduced without access to those private artifacts.
| 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) |
There was a problem hiding this comment.
[P2] Avoid full-bucket temporaries when applying acceptance
The RHS materializes self.selected[slot_ids] and then another tensor for the masked result across the complete key_bucket, including unused positions. With native Ouro-1.4B BF16, its unmodified 65,536-token model capacity, max_num_seqs=64, K=3 async block speculation, r=1, 256-token prompts and 16 generated tokens, the engine constructs successfully and a four-token warmup completes, but this line raises CUDA OOM during the measured batch. The failed allocation is 18.00 GiB with only 11.06 GiB free (82.94 GiB already allocated by PyTorch); the last completed step has 64 resident requests and no completed measurement requests. The reference pool alone is 48.19 GiB, and the otherwise-matched dense cases complete 64/64. Please avoid temporaries proportional to the entire selected bucket here and make the configured memory budget cover the peak acceptance path (or reject an unsupported capacity before admitting requests). Add coverage that grows beyond the short warmup into overlapping rounds; complete reproduction and raw failure figures are in the review body.
…ence 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.
|
Thanks — the OOM reproduction is solid and the finding is correct. Fixed in Blocking finding: retirement is now memory-bounded. You are right that rejected = slot_ids[rows > limit]
if rejected.numel():
self.valid[rejected] = False
self.counts[rejected] = 0
self.mass[rejected] = 0While fixing it I found the same pattern in A regression test now asserts the property directly rather than the symptom: for a 512 MiB bucket, Capacity is explicit and accounted before admission. New Overlapping rounds coverage added as requested: On the TTFT claim — you are right and I withdraw the generalization. What I measured was sequential, single-request, prefill-dominated latency; it says nothing about concurrent serving, and your p95 numbers (2.39 s → 37.8/38.7 s at B=64) are the relevant ones because they include scheduling and output delivery while the reference decode path is in the pipeline. The only defensible claim from my measurement is that prefill rows never take the correction (§6.3), not that enabling the feature preserves first-token latency under concurrency. On the full-retention divergence — not localized, and I do not want to claim it is benign. Two data points I can add: on the tiny model the sparse engine matches dense exactly in both fp32 and bf16 (B=4 and B=8, 48-token prompts, 100% of positions), and the torch reference vs the Triton kernel agree to 2e-7 in fp32 and ~2e-9 in bf16 on identical paged inputs. That makes a structural arithmetic error unlikely but does not prove it, and it does not explain a 30 → 2578 flip at scale. The plausible cause is that the sparse path replaces Triton with fp32-accumulating PyTorch only for capture/consume rows, so hidden state is produced by mixed kernels across the four loops and small deltas compound; at 1.4B with sharper distributions a greedy argmax can flip on a near-tie. The localization you suggested is the right next step — capture the logit margin at the first divergent step for request 23 and compare against a same-input dense run — and I will do that before claiming full-retention equivalence at scale. Defaults stay as they are: |
|
ok, I will add some high-concurrency e2e fix soon, just wait for my news. |
Thanks,waitting for u. |
|
@kunkunblueberry Please cherry-pick the high-concurrency follow-up from 0z5a/vllm-rlt#31. It lives in my fork and compares directly with your git fetch https://github.com/0z5a/vllm-rlt.git codex/pr157-batched-attention
git cherry-pick e8f8261a447f778a280e20e25dec7d6ec3b4701f 90120f2116f1b09994b90b3a8578068d850e9e76 398333dd08fcb41e60d7a616846fb9d5738fa7d4Please take all three in order. They batch equal-length gather/selection work, read consume counts once per layer, and prepare views once while retaining the scalar per-row cuBLAS shapes/strides. The first prototype's folded matmul batch changed BF16 rounding and generated tokens; the latter commits correct that regression. The rejected prototype's faster numbers are not used below. Your memory-bounded retirement fix is already in the baseline. Completed GPU A/B results
B=32's three scalar observations were collected in the preceding validated phase. A separate fresh scalar drift check reached 12.15 tok/s; the candidate median is still 3.62× that observation. B=64 scalar/candidate repetitions were interleaved. Clocks and CPU affinity were not pinned, so the observed ranges are included rather than choosing the fastest run. Correctness, memory and test status
Reproduction scope and remaining boundaryOne GH200 (97,871 MiB HBM), real Ouro-1.4B BF16, PyTorch 2.14.0+cu130 / Triton 3.8, OMP=4, seed=157. Native StageEngine, Triton attention, async scheduling, default multistream, CUDA graphs off; 3 draft tokens / 2 draft loops / 4 target loops, block mode. Each fresh process warms up 4 output tokens at the target concurrency, then measures 256 input / 16 output tokens, greedy, ignore EOS, fixed four loops. Measurement includes enqueue, prefill, decode, output delivery, drain and CUDA synchronization. Prompts for request i: These are short-burst native-engine measurements, not steady-state HTTP serving or task-quality results. The original full-retention-versus-dense discrepancy remains: current A matches dense on 503/512 tokens at B=32 and 1001/1024 at B=64. This patch preserves A exactly; it does not claim to resolve that separate numerical/quality issue. Closeout completed: all four separate CPU+CUDA profiles (A/B × ops/full) passed, with matching layer-0 capture/consume windows (32/128 rows). Full captures contain readable shapes, memory and stack evidence. Capture CUDA launches fell 1,058 → 378; consume launches 3,588 → 441, while the corresponding The complete output archive passed an independent local check of 126 files, and all 12 nested profile bundles passed their internal SHA256 inventories. No owned GPU children remain and the original GPU lock was released. Archive SHA256: |
|
@kunkunblueberry — progress update for my attention/backend, ordinary decode and synchronous self-speculation work: Draft collaboration PR: kunkunblueberry/vllm-rlt#1 targets It reuses the shared reference plan/store, integrates ordinary synchronous decode and synchronous block speculation at draft depths 1/2, and covers per-input reference ownership, rejection/replacement, cancellation, capacity checks and stream-safe retirement. Validation completed
Numerical contract and remaining quality workThe original discrepancy was reproduced and traced to attention backend rounding: identical Q/K/V at the watched loop-2 physical layer produce a small BF16 difference between Triton and the Torch reference, which later recurrence amplifies. Frozen-input replay reproduces both outputs. The fix preserves the selected backend for dense/capture outputs and stores the actual rounded recurrence input as Partial-retention task-quality validation is still pending. Retention .1 completes the synthetic execution tests, but this does not establish task accuracy. Remaining work includes final-path Detailed diagnosis and validation report · Reproduction scripts, frozen tensors, raw results and SHA256 inventory. The complete archive's 322 files were independently verified locally, and GitHub asset digests match. Could you please review the collaboration PR—especially the |
…tion Brings in the basic/synchronous cell from @Levius-Fubuki (PR #1 in this fork), based on 048b9d7: - ordinary decode and synchronous block speculation at draft depths 1/2, per-input reference ownership, rejection/replacement, cancellation, capacity checks and stream-safe retirement; - numerical contract fix: dense/capture outputs keep the selected backend and the actually rounded recurrence input is stored as O_ref; retention 1.0 uses the same dense backend as the control, which removes the full-retention token divergence traced to Triton-vs-Torch BF16 rounding.
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
Thanks for merging kunkunblueberry#2. 43d9cc9 has the same tree as d6c52eb, and CI is green, so I'm switching my review to approve.
Scope, as agreed:
- Sparse decode stays opt-in (
SparseDecodeAttentionConfig), and the default dense path is unchanged. - In eager end-to-end runs it does not beat dense yet. The gains we measured come with CUDA graphs, P/D and token pipelines, and they depend on the GPU. That part comes in the stacked follow-up I'll open on top of this branch.
|
Before this goes to approval: thank you for staying with this one. The review has been unusually useful — the OOM reproduction, the dense baseline switch, the CUDA-graph measurements and now the retirement race each changed what this PR actually is, and none of them were things I would have found from my side. The race in particular only shows up on a host that runs far enough ahead, which is exactly the kind of bug a reviewer with different hardware catches and a single-machine author does not. I appreciate the patience with the long iteration, and the willingness to keep the scope honest (opt-in, GPU-dependent, no serving claims) rather than let it ship on the best-looking number. |
|
We already validate model context length and total KV capacity at admission. Please remove |
|
Please rename |
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
_incarnations retains an entry for every (request_id, row) ever reserved, and neither release() nor discard() removes those entries. With unique request IDs, this dictionary grows indefinitely during long-running serving, even after all requests finish.
| incarnations = [] | ||
| for row, slot in enumerate(slots): | ||
| count = self._incarnations.get((request_id, row), 0) + 1 | ||
| self._incarnations[(request_id, row)] = count |
There was a problem hiding this comment.
Please track the incarnation counter per physical slot instead, incrementing it on each allocation without resetting it on release. This preserves stale-reference detection while keeping the number of counters bounded by the slot count.
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
The incarnation scheme also allows a stale plan to pass validation after a slot is reused by another request; details at _require_live() below.
| if slot == -1 and mode == DENSE: | ||
| continue | ||
| binding = self._bindings.get(slot) | ||
| if binding is None or binding[2] != incarnation: |
There was a problem hiding this comment.
_require_live() checks only the incarnation, but reserve() generates that counter per (request_id, row), so different requests can have the same value.
Reproduction at the store interface:
- A reserves slot 0 with incarnation 1 and builds a plan.
- A releases the slot.
- B reserves slot 0, also with incarnation 1.
- A's old plan passes this check. Calling
capture()with it can then modify B's reference state instead of rejecting the stale plan.
Please make all three entry points validate the same allocation identity:
| Entry point | Current check |
|---|---|
build_plan() |
request + generation + incarnation |
_require_live() |
incarnation only |
release() |
request + generation only |
A monotonically increasing allocation version per physical slot, carried by both leases and plans and checked consistently in all three places, would simplify the state and avoid these collisions. Do not reset that version on release.
Regression coverage should include at least:
- Different requests reuse the same slot: the old plan must be rejected before modifying the new owner's state.
- The same request and generation reuse the same slot: a delayed
release()using the old lease must be rejected without freeing or clearing the new allocation. The currentrelease()check cannot distinguish these two allocations.
This is a reproduced store-interface defect. The same sequence has not been reproduced through normal engine scheduling, so this does not establish that serving outputs are already incorrect.
…t bound Four review items, one of them a real allocation-identity bug. **Stale plans could reach a new owner's slot.** Incarnations were counted per (request_id, row), so the next request to borrow a physical slot started again at 1. A plan built by the previous owner therefore carried the same number as the new owner's, passed _require_live, and could write into a reference slot it no longer owned. Counters now belong to the slot and only increase, which is what makes a recycled slot's old plan unambiguously stale. The new test_stale_plan_from_a_previous_owner_cannot_touch_a_reused_slot reproduces the reported sequence (A reserves, releases; B reserves the same slot; A captures) and asserts both that the old plan is rejected and that the new owner's valid/counts/mass are unchanged. Verified to fail against the previous per-request counter with exactly the reported symptom: both owners see incarnation 1 on slot 0. **max_context_tokens is gone.** It imposed a request-length limit that the reference does not need: capacity is a buffer size, not an algorithm bound, and the index buffer already starts small and grows on demand. Removed the field, its validation, the admission check that duplicated the model-context check, the CLI option, the docs paragraph and the tests that existed only to exercise it. Capacity now follows the model's positional limit and the engine budgets the growth peak. **Naming.** The config was stored as self.sparse_decode (workers) and self.sparse_decode_attention (engine), which reads like an object rather than a configuration. Both are self.sparse_decode_config now, and the worker keyword follows. Regression: CPU 980 passed / 651 skipped; GPU (four sparse modules, async speculative, lifecycle) 217 passed / 4 skipped; ruff check and format clean under the CI-pinned 0.13.2.
|
All four items are in 1. Stale plans could reach a new owner's slot (P0)Confirmed exactly as described. Counters now belong to the physical slot and only increase: count = self._slot_incarnations[slot] + 1
self._slot_incarnations[slot] = count
The regression test is I verified the test fails against the old code, by temporarily restoring the per-request counter on the branch: it fails at the first assertion with Scope, stated honestly: I did not produce an engine-scheduling path that turns this into wrong output, and I am not claiming one. What changed is that the guard is now structural rather than dependent on request identity happening to differ. One test had to change with the semantics: 2.
|
|
Please remove the benchmark scripts and validation reports added by this PR, and summarize the test setup and results in the PR description instead. Keep the regression tests, and link any detailed evidence as attachments. |
|
I checked the FA3/FA4 follow-up PR. It keeps the regression tests and summarizes the H800 test setup and results in the PR description. It adds no benchmark scripts or standalone validation reports. The two earlier sync-decode reports from my contribution can be removed from #157 while keeping their regression tests. The numerical regression evidence is already available as a release attachment and can be linked from the PR description. |
Review asked for the evidence to live in the PR description and as attachments rather than as checked-in benchmark scripts and standalone reports. The regression tests stay; only the harnesses and the two write-ups go. - benchmarks/: sparse_decode_attention_bench, sparse_decode_batched_ab, sparse_decode_context_sweep, sparse_decode_engine_e2e, sparse_decode_engine_e2e_text, ttft_bench. - docs/validation/: sparse-decode-numerics-20261029.md and sparse-decode-sync-20261029.md. The numerical diagnosis is linked from docs/runtime.md as its release attachment instead. The test setup and the measured results now live in the PR description.
|
Done in Removed (six harnesses, two write-ups):
Kept: all four regression modules —
Description now carries the evidence, in two sections:
One thing I kept visible on purpose: the engine table is eager, so it shows the batched path at 0.78x / 0.40x of dense, and I say so in the same paragraph as the graph numbers. Removing the harnesses should not make the PR read as if it had a serving win. @bjf-frz the three earlier items are in |
|
Please simplify the naming and ownership model:
Keep regression coverage for cancellation, failed submissions, and overlapping asynchronous batches. |
|
Thanks again for all the work on this integration! One small suggestion: would you be open to mentioning FlashLoop as the source of the cross-loop sparse attention mechanism in the README feature entry or the relevant runtime documentation or comment around core code, with links to the paper and reference implementation? That would help users find the underlying method, alongside the kernel optimizations and serving integration developed here. Really appreciate the effort to bring this into vLLM-RLT! |
Review asked for the ownership model to use what the engine already guarantees rather than a parallel bookkeeping scheme. **Identity is the lease.** `reserve` binds a slot to the lease object it returns; `release` requires that same object; a plan is live only while every slot still points at the lease it was built from. A recycled slot is bound to a different lease, so the previous owner's plan cannot pass, with no per-slot counter and no `(request_id, row)` dictionary that grows for the life of the process. `ReferenceLease.generation`, `ReferenceLease.incarnations`, `SparseDecodePlan.incarnations`, the store's slot counters and its device `incarnation` tensor are gone; the plan carries `owners` instead. **Duplicate request ids are rejected at the boundary.** `_add_request` used to retire a still-retiring request and reuse its id; it now raises, including for requests that are still retiring. Together with the existing rule that slots are recycled only after their last reader completes, that is what makes lease identity sufficient. The EOS test now asserts the rejection and then reuses the id after retirement completes. **Renames.** `SpeculativeRoundTicket` is `AsyncSpeculativeOutput`. Also credits FlashLoop as the source of the cross-loop sparse attention mechanism in the README, docs/runtime.md and the sparse attention module docstring, with the paper and reference implementation linked. Regression: CPU 980 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 216 passed / 4 skipped; ruff 0.13.2 clean.
|
Both suggestions are in Ownership is the lease now, not a counter
Removed: This also removes the leak you flagged in review: there is no Duplicate request ids are rejected at the boundary. Rename. Coverage you asked to keep is intact and passing: That last one is the cross-request reuse regression from the previous round, rewritten for the new model: it still performs A reserve → release → B reserve the same slot → A capture, and still asserts the old plan is rejected and that B's Verification: CPU @Superone77 thank you for pointing at the origin — I had described the mechanism without crediting it, which is not a small omission for a paper-derived method. FlashLoop is now named in three places, each linking the paper and the reference implementation: the README feature entry, the sparse-decode section of |
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
Review comments on reference-slot exception safety, state ownership, naming, and type contracts.
| 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)) |
There was a problem hiding this comment.
Slots are removed from self.free before _index() and torch.arange() succeed. If either allocation raises, no lease is returned or registered, so cleanup cannot recover these slots. I reproduced this with an injected allocation failure on CPU: subsequent requests fail with “no retired attention reference slots”. Please make reservation exception-safe and add a regression test that submits another request after the failure.
| # GPU acceptance counts without a host readback. | ||
| self.reference_store = reference_store | ||
| self.sparse_decode_config = sparse_decode_config | ||
| self._generations: dict[str, int] = {} |
There was a problem hiding this comment.
_generations still retains every historical request ID, although the generation value no longer has a production reader. Please remove this dictionary, _AsyncRequestState.generation, and the corresponding assignments in _state().
| 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) |
There was a problem hiding this comment.
What production decisions consume valid, uses, position, and accepted? These fields are updated or cleared, but availability is determined by counts and lease ownership. Please remove unused state and its GPU writes. Also check whether post-verification invalidation in accept() has any remaining consumer before retirement.
| return len(self.slots) | ||
|
|
||
|
|
||
| class AttentionReferenceStore: |
There was a problem hiding this comment.
Please rename this class to SparseDecodeReferenceStore and update its references consistently, so its role in sparse decode attention is explicit.
|
|
||
| modes: tuple[int, ...] | ||
| slots: tuple[int, ...] | ||
| owners: tuple[object | None, ...] |
There was a problem hiding this comment.
Please give owners and store explicit types instead of object; imports under TYPE_CHECKING can avoid runtime dependency cycles. Also remove SparseDecodePlan.key_bucket if it remains unused by production callers.
| reservation upper bound and must never be used as the actual write position. | ||
| """ | ||
|
|
||
| # Sparse decode attention lifecycle; test doubles that skip __init__ |
There was a problem hiding this comment.
Please initialize these dependencies explicitly in the test doubles instead of adding class-level defaults for tests that skip __init__. sparse_decode is also stale after the rename to sparse_decode_config.
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
Two additional findings from the sparse/reference review.
| continue | ||
| # Contiguous slots turn these gathers into views, which is the | ||
| # difference between one index kernel per tensor and none. | ||
| if chunk_slots[0] + len(chunk_slots) == chunk_slots[-1] + 1: |
There was a problem hiding this comment.
[P1] Checking only the endpoints does not establish contiguous slot order. Normal reserve/release operations can produce [5, 3, 7], which passes this condition but selects [5, 6, 7]. CONSUME then reads another request's reference; the identical check at line 619 makes CAPTURE overwrite that request's live slot and leave slot 3 unwritten. CPU execution of the actual batched implementation reproduces both failures (consume output max absolute error 0.11385 versus the scalar oracle). A tiny-Ouro StageEngine run also naturally produced [1, 4, 3] with 8 concurrent requests and varying output lengths/loop counts. Please verify every slot == first + offset before taking a slice, otherwise use indexed access, and add regression coverage for fragmented slots with an unrelated live lease. CUDA execution has not been verified.
| 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 |
There was a problem hiding this comment.
These module-level caches (_WORKSPACE and _CONSUME_WORKSPACE) retain scratch tensors after their owning engine is closed and released; there is no cleanup path for their stream entries. A new stream can retain another workspace. The capture score buffer alone requires 128 MiB for 32 rows × 16 heads × 65,536 positions. Please give this scratch an explicit execution owner and release it after draining that owner's work, using the existing workspace argument or a stream-specific cleanup path. This is a lifetime finding; actual GPU memory retention has not been measured.
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
Clarify the helper's lifecycle responsibility.
| ): | ||
| self.reference_store.release(state[2]) | ||
|
|
||
| def _reference_metadata(self, batch, kv): |
There was a problem hiding this comment.
Please rename _reference_metadata() to _prepare_sparse_reference() and update its caller. It also reserves/releases reference leases and binds device positions, so the current name understates its side effects. Keeping this helper and _release_attention_reference() in the runner is appropriate.
Review comments on the reference store, the runner and the plan types. **reserve is exception-safe.** It popped slots from the pool before `_index()` and `torch.arange()` could fail, so an allocation error leaked the slots and the next request failed with "no retired attention reference slots". It now peeks, allocates, and only then removes them from the pool. test_failed_reservation_keeps_the_pool_intact injects the failure and checks that a later request still reserves. **Write-only state removed.** `valid`, `uses`, `position` and `accepted` were written by the store and by three kernels and never read back: availability is the retained count and ownership is the lease. The tensors, their writes, the `bind_positions` API and the corresponding kernel arguments (`VALID` in `_capture_kernel` and `_select_pass_kernel`, `USES` in `_consume_merge_kernel`) are gone, which also drops 12 bytes per slot from the accounted payload. `capture` and `consume` remain as liveness assertions. **Naming and types.** `AttentionReferenceStore` is `SparseDecodeReferenceStore`. `SparseDecodePlan.owners` and `.store` carry real types through a `TYPE_CHECKING` import, and the unused `key_bucket` field is gone. **Runner.** `_generations` and `SpeculativeRequestState.generation` are removed; the value had no production reader left. The class-level `reference_store` / `sparse_decode` defaults that existed for test doubles are gone, and the one double that skips `__init__` now sets both fields explicitly. Regression: CPU 981 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 217 passed / 4 skipped; ruff 0.13.2 clean.
|
All six are in 1.
|
Two review findings that are not style. **Fragmented slots were read as a span (P1).** The batched path accepted a chunk when `first + len == last + 1`, which `[5, 3, 7]` satisfies, and then used `slice(5, 8)`. CONSUME read slot 6 — another request's reference — and CAPTURE overwrote it while leaving slot 3 unwritten. `slot_span` now checks every slot against `first + offset` and returns None otherwise, so a fragmented chunk takes the indexed path. Both call sites use it. `test_fragmented_consume_reads_each_rows_own_reference` and `test_fragmented_capture_writes_only_the_slots_it_reserved` build the exact fragmentation from the report (release 7, then 3, then 5; reserve three rows gives `[5, 3, 7]`) with slots 4 and 6 held by unrelated live leases. Verified to fail against the endpoint check: slot 3 keeps its sentinel, which is the reported symptom. **Kernel scratch outlived its engine (P1).** `_WORKSPACE` and `_CONSUME_WORKSPACE` were module-level dicts keyed by (device, stream) that nothing released, so a closed engine's tensors stayed reachable — the capture score buffer alone is rows x heads x capacity floats. They move to `kernels/sparse_workspace.py` (Triton-free, so CPU-only installs can release them) and `SparseDecodeReferenceStore.close()` releases its device after the caller has drained; `StageEngine.close()` calls it after its own drain. **Cleanup.** `_reference_metadata` is `_prepare_sparse_reference` on both the runner and the speculative runner, since it reserves and releases leases rather than building metadata. `SparseDecodePlan.reference_depth` is gone: it returned a module constant and had no production reader. Regression: CPU 986 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 222 passed / 4 skipped; ruff 0.13.2 clean.
|
Both P1s and the cleanup items are in P1: fragmented slots were read as a span — confirmed, and worse than a read
def slot_span(slots: Sequence[int]) -> slice | None:
first = slots[0]
for offset, slot in enumerate(slots):
if slot != first + offset:
return None
return slice(first, first + len(slots))Both call sites use it (CONSUME read, CAPTURE write). The regression reproduces your exact fragmentation.
I verified both fail against the endpoint check by temporarily restoring it: Scope: the fragmentation regression runs on CPU. I did not build a GPU-specific fragmentation case, so I am not claiming GPU coverage beyond the existing suite. P1:
|
|
ROCm compatibility check: I merged #171 (ROCm platform support) into this branch as This is compatibility evidence, not a performance or accuracy claim: the branch's own numbers still show sparse slower than dense on this device. Branch: https://github.com/0z5a/vllm-rlt/tree/rocm-sparse |
Priority scheduling cuts high-priority completion 7.80s to 2.22s; incremental KV gives +3.8%; sparse decode InferMatrixLab#157 passes 137 tests on the merged ROCm branch.
Status: draft — the per-key work is solved, the per-layer fixed cost is not
This branch is no longer blocked on anyone else's work. It carries the lifecycle contract, the numerical contract, and now both fused kernels. What it does not do is beat dense attention end to end, and the reason has changed twice during review, so the current one is worth stating precisely.
Per-layer, one attention call (32 rows, 24L/16H/128D, BF16, r=0.1, median of 10, dense = Triton):
That is 3.0–3.4x faster than where this started, and the algorithm now behaves as advertised: capture reads the paged cache once instead of materialising an expanded FP32 row, and consume touches only the
kretained keys instead of gathering K and V densely.Engine level, tiny Ouro (2 layers), async self-speculation, prompt 32, max_tokens 16:
Full retention reproduces dense tokens exactly at every concurrency, which is the correctness result. Throughput is 5–19x worse, and the gap widens with concurrency, which is the diagnosis: the per-layer fixed cost dominates and it scales with rows, not with keys. With two layers and 48-token contexts there is no per-key work to amortise it against.
So the remaining work is a Python-overhead project, not an attention-kernel project: per-chunk index construction, per-row
note_count,plan.groups()/_require_livewalks,context_lengths.tolist()per layer, andq[row_ids]/block_tables[row_ids]gathers. A Triton launch in this environment is 0.015 ms and store views changed nothing measurable, so the cost is the Python around the kernels.A meaningful serving number needs a trained Ouro-1.4B checkpoint (24 layers, long context), which this host does not have; the tiny model cannot answer that question.
Required human contributor answers
Answers below reuse the decision recorded by the human contributor @kunkunblueberry in issue #156, claim comment; they are referenced, not invented.
1. What outcome are you trying to achieve?
As recorded in the claim comment: contribute the asynchronous self-speculation integration for sparse decode attention, i.e. the lifecycle contract for per-token reference state across outstanding rounds.
2. What is the highest-leverage variable for that outcome?
The acceptance criteria of the RFC #156 §8 matrix cell for "Asynchronous self-speculation × Sparse decode attention": isolate references for outstanding rounds; GPU acceptance drives retirement, and allocation, conversion and invalidation dependencies across outstanding rounds. The state ownership rules of §5 (one slot per live token, incarnation after same-position replacement, reuse ordered after the last reader/writer) are the contract every later kernel change must respect.
3. How will this work move that variable rather than create busywork?
It implements that contract, plus the numerical contract that makes the result checkable: a runner-owned reference store, a per-row plan carried on the existing execution metadata (§4 "small named views rather than a new engine"), runner integration that reserves and retires slots per round, two fused kernels that make the arithmetic affordable, and focused tests for draft depths 1/2, rejection at every candidate position, same-position replacement, delayed completion and slot reuse.
4. Who is the customer or user, and what do they actually need?
The shared sparse decode attention feature and its reviewers: they need the lifetime/ownership layer to exist and be tested, and they need to know whether the feature can pay for itself. The second question now has a measured answer with a named remaining cost.
PR necessity and bugfix evidence
This is not a bugfix; it is a new RFC cell implementation. Scope follows the recorded claim in issue #156 and RFC §11 (cells contribute commits to the shared feature PR).
Purpose
Implements the asynchronous self-speculation × sparse decode attention cell of RFC #156:
attention/sparse.py— the §3.2 arithmetic contract (reference statistics + sparse correction, plain PyTorch, checked against the RFC's worked example), the per-rowSparseDecodePlannamed view, and the dispatch between the scalar oracle, the bitwise-exact batched path and the fused path.worker/attention_reference.py—AttentionReferenceStore, the runner-owned per-token slot pool (§5):empty → source-ready → reusable → retired, one slot per live token, incarnation per(request, row)so a corrected token reusing a position rejects stale plans, device-side acceptance-driven invalidation with no host readback, metadata-only retirement, and host-side retained counts so a reader never syncs.kernels/sparse_capture.py— one program per(row, kv head, position block)computesq @ K^Tstraight out of the paged cache: no expanded K, no per-row cuBLAS, one launch per layer.kernels/sparse_consume.py—selected_scoresandselected_valueswalk the retained indices and never materialise K or V.kernels/sparse_decode.py— the opt-infusedimplementation: fixed per-layer kernels for capture and consume with no host readbacks, selected withSparseDecodeAttentionConfig.implementation(batchedby default). It changes consume-row numerics, so the bitwise contract stays withbatched. The reference index buffer starts small (initial_key_capacity) and grows on the owner stream; the engine budgets the growth peak, not the initial allocation.worker/attention_reference.py— retirement builds its index on the store's owner stream. Building it on the caller's stream let the host free and reuse that memory before the owner stream'sindex_fill_ran, which surfaced as an index-out-of-bounds device assert once the host ran ahead under CUDA graphs.worker/async_speculative.py— per round, the runner reserves one lease per verification row, publishesCAPTUREplans at depth index 1 andCONSUMEplans 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.py/entrypoints/stage_engine.py— thesparse_decode_attentiongroup (enabled, default false;key_retention;key_bucket;max_context_tokens) with explicit capability checks and reference-state memory accounted before admission.Numerical contract: dense and capture rows keep the selected backend's output, the actually rounded recurrence input is stored as
O_ref, and full-retention consume rows use the same dense backend as the control. The fused capture pass computes the same FP32 dot product as the oracle with a different reduction order: max score delta 1.1e-5 on a magnitude of 65, retained mass delta 9e-8, attention output delta 2e-8, identical selected sets on 511/512 row-head pairs.batched_paged_attentionwithoutfused_scoresstill reproduces the scalar oracle bit for bit.Explicitly out of scope: synchronous speculation, token pipelines, CUDA graph capture with sparse enabled, other matrix rows.
Test setup
Hardware and software. A100 (SM80), torch 2.13.0+cu130, Triton attention, ruff 0.13.2 (the version CI pins). GPU suites run with
--run-gpu -k "not flash_attn_4"; the deselected cases areflash_attn_4, which does not support SM80.Model.
ByteDance/Ouro-1.4Bat revision574fa66c, sha256 verified against the LFS digest (58872a72…): 24 layers, 16 attention heads, 16 KV heads, 65,536 context, 4 UT steps, 1.435B parameters, BF16. Properties that do not depend on trained weights are checked on the test-only tiny Ouro (2 layers, 128 context).Regression tests kept in this PR.
tests/test_sparse_decode_attention.py— plan and arithmetic contracts (RFC §3.2 worked example = 23.5, full retention ≡ dense, ceil retention rule, capacity), store lifecycle (one slot per live token, owner-checked release, acceptance-driven invalidation at accepted counts 0/1/2, root survival, stale-plan rejection after reincarnation and after cross-request slot reuse, disjoint outstanding rounds, metadata-only retirement allocation bound, byte accounting, plan carriage), end-to-end decode through the real model + paged KV + attention, and engine-level full-retention parity.tests/test_sparse_decode_batched.py(from @0z5a) — the batched path must be scalar-exact.tests/test_sparse_decode_sync.py(from @Levius-Fubuki) — ordinary decode and synchronous speculation at draft depths 1/2.tests/test_sparse_decode_fused.py(from @0z5a) — the opt-in fused kernels.vllm-rlt Version / Commit:
989921a(upstream/main, v0.2.0, merged here) + this branch.Test results
Regression on the current head. CPU
OMP_NUM_THREADS=1 pytest -q: 980 passed, 651 skipped, 0 failed. GPU (async speculative, async lifecycle, four sparse modules including the fused path): 217 passed, 4 skipped.ruff check .andruff format --check .clean under 0.13.2.Numerical contract. The fused capture computes the same FP32 dot product as the scalar oracle with a different reduction order: max score delta 1.1e-5 on a magnitude of 65, retained mass delta 9e-8, attention output delta 2e-8, identical selected sets on 511/512 row-head pairs at 8k context. Full retention reproduces the dense engine's tokens bit for bit at concurrency 1/8/32 on the tiny model, and 5904/5904 tokens across all exit depths on the real 1.4B model (@Levius-Fubuki's run for the merged
O_refcontract). On SM90 the split capture rounds FP32 to within 1.8e-7 of the dense kernel while BF16 stays bit-exact, which is why the FP32 test allows one ulp.Per-layer cost (32 rows, 24 layers, 16 heads, head_dim 128, BF16, r=0.1, median of 10, dense = Triton):
Engine end to end (Ouro-1.4B, 200-token real-text prompt from the checkpoint's tokenizer, 16 tokens out, median of 2, eager, A100):
Eager decoding on this host is host-bound, so these numbers describe the batched path, not the algorithm's ceiling. With CUDA Graphs the fused path reaches parity on GH200 (0.92x at 1k × 1, 1.00x at 1k × 8, 0.98x at 8k × 1, 1.02x at 8k × 4) and gains on lower-bandwidth GPUs (+9.3% and +9.6% on RTX PRO 6000 at 32k and 16k, +7.3% on L40S at 32k, +13.5% for a token pipeline at 22.5k). That work is a stacked follow-up, not this PR.
Retained mass — why token agreement is not the quality signal. On Ouro-1.4B, r=0.02 retains 10 of 200 keys and discards 17% of the attention mass on average (44% on the worst row), yet the 16 greedy tokens stay bit-identical to dense at every concurrency. A quality gate that compares short generations will pass trivially; retained mass, logit margins or task accuracy over a real cohort are the measurements that discriminate.
Evidence attachments. The numerical diagnosis and regression report is kept as a release attachment. The benchmark harnesses that produced the tables above were removed from this PR at review request; their setup is described here and their results are the numbers above.
Not claimed: a serving speedup on this PR's own path, TTFT under concurrency, or partial-retention task quality. The feature stays
enabled=falseby default, andimplementation="batched"keeps the bitwise contract.