Repository navigation
Sparse decode attention: per-token reference lifecycle for async self-speculation - #157
kunkunblueberry wants to merge 19 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.
…ched 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.
|
Both external contributions are now integrated into the feature branch as I had to repair one semantic conflictThe cherry-picks auto-merged without text conflicts, but they were not semantically compatible: 0z5a's commits replaced
Joint regression on the merged head
Batching gain, isolated on the merged tree
Per-layer four-loop totals (one dense + one capture + two consume): ctx512 r=0.1 229.2 → 68.6 ms (3.34x), ctx512 r=1.0 248.4 → 67.7 ms (3.67x), ctx2048 r=0.1 177.7 → 88.3 ms (2.01x). The batched consume path now costs about one dense attention at r=0.1, which is the algorithmic win finally showing up in wall clock. Two honest notes on this table. First, the end-to-end gain is lower than 0z5a's 4.23x because that campaign predates the Boundaries I am not claiming
Next steps if this direction is accepted: fix the long-context capture regression, then run the combined-revision concurrency A/B that @0z5a asked to re-measure, and start the partial-retention quality work. |
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
LGTM for the current opt-in lifecycle/correctness implementation, reviewed and retested at e4f8988.
I re-ran the current head on one GH200 with the real Ouro-1.4B BF16 weights, including B=32 and B=64. This resolves my previous full-retention concern on the tested Triton path: full retention now produces exactly the same generated tokens and exit depths as feature-disabled dense attention. At retention 0.1, the batched path exactly matches the same-head scalar reference. All three repetitions are output-identical within each mode.
Concurrent execution results
24/24 unprofiled cases, 1,152/1,152 requests and 18,432/18,432 generated tokens, with no timeout, duplicate completion or residual KV/reference ownership. All requests reached the intended resident batch and speculative batch size, and sparse capture/consume counters were nonzero for every enabled case.
| Batch | Mode | Median output tok/s | Three-run range | Max of three PyTorch peak allocated GiB |
|---|---|---|---|---|
| 32 | Feature disabled | 98.81 | 75.60–101.68 | 10.62 |
| 32 | Full retention (1.0) | 24.77 | 23.15–32.22 | 34.71 |
| 32 | Scalar reference (0.1) | 9.54 | 9.40–10.02 | 13.11 |
| 32 | Batched implementation (0.1) | 35.01 | 31.96–37.81 | 13.11 |
| 64 | Feature disabled | 123.96 | 111.83–149.11 | 17.37 |
| 64 | Full retention (1.0) | 34.83 | 24.80–36.47 | 65.56 |
| 64 | Scalar reference (0.1) | 7.82 | 7.78–8.47 | 22.36 |
| 64 | Batched implementation (0.1) | 27.23 | 24.55–27.46 | 22.36 |
The batched/scalar median ratios are 3.67× at B=32 and 3.48× at B=64. These are observations for this finite-burst protocol, not a sparse-over-dense speedup: feature-disabled dense remains substantially faster, and full retention is particularly expensive in memory. I do not recommend changing the default based on these data.
Reproduction configuration and limits
- Current-head native
StageEngine, Triton backend, async scheduling,SpeculativeConfig(3), CUDA Graph off; one GH200, Python 3.13, PyTorch 2.14 + CUDA 13.0, Triton 3.8. - Greedy fixed four-loop generation (
max_loops=4,exit_threshold=1,temperature=0,ignore_eos=True), 256 input / 16 output tokens. Request i uses token j =1 + (j*137 + i*53) % 49000; seed 157. These are synthetic performance inputs, not task-quality evaluation. max_num_seqs=B,max_num_batched_tokens=4096,prefill_chunk_size=256,admission_scan_limit=128, block size 16, incremental allocation. Physical blocks =B*ceil((256+16+4)/16)*4+64(2368 / 4672). The model's maximum position configuration stays at 65536.- Four modes: disabled, retention 1.0, same-head scalar reference at 0.1, batched at 0.1. Three repetitions, second repetition in reverse mode order, fresh process/model per case. The scalar comparison routes the sparse call to
reference_paged_attention, preserving this head's dense backend callback. - The timer covers simultaneous request admission through completion, including prefill. A 4-output-token warmup precedes the 16-token measurement; late speculative shapes may remain cold. The variability above is retained, so these results do not establish fully warmed steady-state throughput, HTTP saturation, SLO goodput, or dataset accuracy. Profiler runs are excluded from the table.
- Model safetensors SHA256:
58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af.
Additional validation
The six applicable sparse/speculation regression modules yielded 303 passed, 2 skipped; the same invocation also had 12 failures caused by the unavailable optional flash-attn-4 dependency. Those failures are preserved and FA4 is not qualified by this review. The modules were test_sparse_decode_batched, test_sparse_decode_attention, test_sparse_decode_sync, test_async_speculative, test_async_speculative_lifecycle, and test_speculative, with pytest -q --run-gpu -p no:cacheprovider.
Separate scalar/batched × ops-only/full sessions produced 12 valid CPU+CUDA traces covering matching layer-0 capture (32 rows), consume (128 rows), and reference acceptance windows. All six full traces contain actual shapes, allocation events, and call stacks. In the matching ops-only windows, scalar → batched kernel counts were capture 1092 → 392, consume 3588 → 441, and accept 46 → 46. This supports reduced launch/work overhead in the batched attention path; these counts are not an attribution of the entire end-to-end speedup.
The complete raw archive (159 files, including nested trace archives) was independently hash/inventory checked, and request counts, generated tokens, timing denominators, output equivalence and teardown were independently audited. Archive SHA256: 371f603fbee7f8245ea4569cbcfc7f41eb998a54c68e83e5b29e91ad6359031e.
This approval is scoped to the tested correctness/lifecycle behavior. Reduced-retention task quality, FA4, fully warmed stable load, and production saturation remain outside this qualification.
0z5a
left a comment
0z5a
left a comment
There was a problem hiding this comment.
Request changes on e4f8988. This supersedes my earlier correctness-scoped approval: the lifecycle/output checks pass, but the end-to-end performance and memory acceptance criteria are not met.
The independently audited current-head results and full reproduction configuration are in the preceding review. All 24 concurrent cases completed, full retention matched dense tokens/depths, and batched retention 0.1 matched the scalar reference. Those positive correctness results remain valid.
However, the relevant dense baseline shows a substantial regression:
| Concurrent requests | Feature-disabled dense, median output tok/s | Batched sparse at retention 0.1, median output tok/s | Dense peak allocated GiB | Batched sparse peak allocated GiB |
|---|---|---|---|---|
| 32 | 98.81 | 35.01 | 10.61 | 13.11 |
| 64 | 123.96 | 27.23 | 17.37 | 22.36 |
Throughput is the median of three unprofiled runs; memory is the maximum PyTorch allocated peak across those runs. The 3.67× / 3.48× gains are against the scalar sparse reference, with essentially unchanged peak memory against that reference. They do not establish an improvement over dense. At retention 1.0, peak allocated memory is higher still: 34.71 / 65.56 GiB for B=32 / B=64.
Before approval, please address these points:
- Use dense as the primary end-to-end acceptance baseline, while retaining scalar as the implementation-correctness/optimization baseline. Diagnose and reduce the added reference-cache, capture/consume and acceptance overhead. Include an explicit memory breakdown for persistent references, temporary/workspace buffers and KV storage; reduced attention retention alone is not evidence of lower total memory.
- Provide a fully warmed, matched-load follow-up with at least three repetitions, identical model, target loop/exit policy, hardware and output lengths. The present protocol warms up 4 output tokens and measures 16; it exposes a negative result but does not establish fully warmed steady-state performance. Include a longer-context/load sweep that identifies an actual useful operating region relative to dense, and preserve cases that still regress. Keep profiler runs separate from throughput measurements.
- Qualify the approximation at retention 0.1 against dense on a fixed real task cohort, with paired outputs, the quality metric and an explicit accuracy budget. Full-retention equality and equality to the scalar approximation do not establish reduced-retention task quality. Any speed/memory claim must state the associated quality tradeoff.
The feature remaining opt-in is appropriate, but does not by itself satisfy these performance/memory/quality gates. This request is based on measured costs and missing approximation-quality evidence, not a claim that the current-head token ownership fix failed. No HTTP saturation or production-capacity claim is supported by the finite-burst measurements above.
|
Another iteration is ongoing, will post it soon. |
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.
|
You are right to switch the baseline to dense, and the measurement answers the question directly: there is no operating region where this implementation wins.
DiagnosisThe budget for sparse decode to beat dense is The reason is a per-key cost asymmetry: capture's selector costs about 10 ms per 1k keys (29.9 → 53.2 → 98.4 → 211.8 → 323.3 across 1k/2k/8k/16k/32k), while the fused dense kernel costs about 0.4 ms per 1k keys (2.6 → 2.8 → 1.7 → 7.7 → 13.0). That is ~25x. It is structural, not tunable: the selector must scan all So batching was never the blocker. It already gave 6.8–11.1x on consume ( Memory breakdown (three components, 32 rows, 24 layers)
This reproduces your delta exactly: your +2.5 / +5.0 GiB at B=32/64 is the persistent pool sized to the full 65,536 context (24.09 / 48.19 GiB at r=1.0), not the attention path. With On the third gateReduced-retention task quality is genuinely missing, and I will not fabricate it. Proposed protocol: a fixed real cohort (the repo's GSM8K harness, plus a long-context retrieval set), paired dense/sparse outputs at r=0.1, contexts 512 / 2048 / 8192, a pre-declared accuracy budget, and the speed/memory tradeoff reported next to it. Given the sweep above, the honest expectation is that this cohort fails the speed side before accuracy even matters. What this means for the PRThe performance and memory gates cannot be met by this PR as scoped. I suggest reframing it as what it actually is — the lifecycle contract, the numerical oracle, the |
|
Agreed on the direction, and I want to make the ordering explicit rather than leave it implicit: this cell is downstream, so the basic path should land first. I have converted the PR to draft and updated the description to say so. Why the order matters hereThe sweep above shows the blocker is not the lifecycle layer, not batching, and not the numerical contract — all three are done and validated. The blocker is that selection has to move into the kernel. That is the basic-implementation cell, and it is upstream of everything this PR does: the plan, the store, and the retirement rules only pay off once a fused capture/consume path exists that beats dense attention. Landing this first would put a well-tested layer in front of a kernel that does not yet exist in a form worth using. So @Levius-Fubuki, the basic path has priority. The target from my measurements is concrete:
What I will do when the kernel lands
What is available to build against right nowThe store and plan surface is stable and tested; treat it as frozen unless it is genuinely inconvenient:
If any of that is the wrong shape for the fused kernel, say which and I will change it in this PR instead of asking you to work around it. And if you would rather this branch not carry your ordinary-decode/sync cell plus @0z5a's batching while it is a draft, I can split those into a separate feature branch — just tell me which shape you want to review against. |
|
Thanks for the sweep and for converting to draft. I agree that this PR cannot meet the performance gate as scoped. Below are three end-to-end results that complement the kernel-level sweep, plus one point where I disagree: the lifecycle layer is not yet free of host readbacks, and I would fix that before freezing the store/plan surface. All runs use the same engine settings for dense and sparse: async self-speculation with 3 draft tokens, fixed 4 loops, Triton, greedy, Ouro-1.4B BF16, 1. The store still synchronizes with the hostWith sparse enabled, a batch-1 decode step performed about 683 host synchronizations, versus 5.5 for dense. About 235 of them come from
All of these can be fixed without changing any value written:
With the copy and acceptance changes, GSM8K tokens and exit depths were identical to e4f8988 on all 87 questions. The scalar-write change writes the same values and passes the GPU lifecycle suites. After all of them, the store's reserve/invalidate/release path performs no synchronization. The wall-time effect at batch 1 is small (about 1%). The point is the contract the fused kernel will build on, so I would land this before treating the surface as frozen. 2. A short-context quality data point for the third gateSetup: the fixed 87-question GSM8K CoT subset the repo already uses (lm-eval strict-match), batch 1, on an RTX 5090.
Accuracy holds at r=0.1 on this short-context set. Long-context retrieval is still missing, as you said. 3. Selection inside the kernel reaches parity end-to-end, not a winI prototyped an opt-in fused path locally, with the default left as the current batched path and its bitwise contract:
Gates were fixed before the runs, and dense and fused were interleaved three times on an RTX 5090:
Two things limit the end-to-end gain. These may help size the follow-up kernel.
4. MemoryAgreed that The commits are local; I can push them as a branch to build on if that is useful: the store sync fixes, the index-pool growth, and the fused prototype with its tests. |
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.
|
Levius's basic path landed and is merged here ( What the profile actually saidI had reported "the selector costs 25x per key". That was directionally right but the mechanism was wrong, and the correction matters:
The fix
Numerics against the oracle at 8k: max score delta 1.1e-5 on a magnitude of 65, retained mass delta 9e-8, full attention output delta 2e-8, identical selected sets on 511/512 row-head pairs (the one difference is a near-tie). That is a new numerical reference, not a bitwise-identical one — same trade Levius made for Measured (32 rows, 24L/16H/128D, BF16, r=0.1, median of 10, dense = Triton)
Capture is 2.9x faster at 8k and 5.0x at 32k, and at 512 it is now cheaper than dense attention itself — the algorithmic claim ("selecting 10% of keys can cost less than attending to all of them") holds once selection is fused. Total four-loop is still 2.5–4.3x dense because consume has not been touched. Note the small-context rows: capture is ~18 ms flat from 1k to 4k while dense is 2.6–3.1 ms, and the 512 measurement (1.97 ms, below the 2.49 ms dense call it contains) is inside run-to-run noise. Those contexts are dominated by fixed per-layer overhead, not by the score pass. Consume is the remaining bottleneck, and it is all overheadAt 512 with k=52 keys — a workload whose GPU work is trivial — consume still costs 12.8 ms per layer. Bisected:
Fixing the first two removes ~8 ms of the ~10 ms fixed cost at every context; that alone puts the four-loop total near or below dense at 512–2k. The 32k gather needs a layout fix (indexing State of the PRStill draft, |
Brings in InferMatrixLab#159 (scheduler wavefront selection / prefill progress) and InferMatrixLab#160 (native CUDA SiLU-and-Mul) so the branch is measured against current main.
…onsume 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 InferMatrixLab#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).
|
Consume is fixed and the branch is rebased on current main. Three measured causes, all of them overheadAt 512 context with k=52 keys, consume cost 12.8 ms per layer for work whose GPU time is negligible. Bisected:
Measured (32 rows, 24L/16H/128D, BF16, r=0.1, median of 10, dense = Triton)
Consume alone: 12.8 → 5.5 ms at 512, 17.3 → 9.9 ms at 8k, 36.2 → 24.4 ms at 32k. Together with the fused capture pass, the four-loop total is 3.0x faster than this morning at 32k (306.9 → 100.8 ms) and 2.5x at 8k (97.9 → 39.6 ms). Still 1.9–2.6x slower than dense, so the feature remains opt-in and unqualified for serving. What is left, and it is now structural rather than silly
Doing both would put the four-loop total at roughly Regression on the rebased branch: CPU |
|
Thanks for the two rounds today, and for keeping the sweep as the single harness. On your closing question: I already have both of your next items on a local branch. One is a selection with no fixed per-layer cost. The other is a fused consume that walks the selected indices without materialising K/V. I ran them against your head so we don't build the same kernels twice. What the local branch doesThe branch is your head
Your sweep, both implementations on one GPUI ran
Eager GSM8K-87, batch 1, about 1k contextThree interleaved pairs on an RTX 5090, Ouro-1.4B, async self-speculation K=3, r=0.1:
So in eager mode the fused path pays off from roughly 8k context. At short context with batch 1 it is about 6% slower than dense. Two notes on
|
Status: draft — this is the downstream half, the basic path should land first
I am converting this PR to draft and handing priority to the basic implementation. The measurement below is the reason: the blocker for sparse decode attention is not the lifecycle layer, it is selection, and selection is the basic-implementation cell. Keeping this open as a ready PR would invert the dependency order and invite reviewers to spend time on the wrong layer.
Proposed sequence:
O_refrounding contract, and the batching work. Its tests are the harness that step 1 gets validated against, which is why they are worth landing in some form regardless of order.The branch currently also carries @Levius-Fubuki's ordinary-decode + synchronous-speculation work (his PR #1 in this fork, merged as
efe5bc7) and @0z5a's three batching commits. That is what makes it reviewable as a whole feature branch, but it also means the branch should be rebased or split once the basic path lands upstream; I am not asking anyone to merge this in its present shape.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 exactly that contract and nothing else: a runner-owned reference store, a per-row plan named view carried on the existing execution metadata (§4 "small named views rather than a new engine"), runner integration that reserves/retires slots per round, and focused tests for draft depths 1/2, rejection at every candidate position, same-position replacement, delayed completion and slot reuse. It deliberately does not duplicate the basic implementation / synchronous cells claimed by @Levius-Fubuki.
4. Who is the customer or user, and what do they actually need?
The shared sparse decode attention feature PR and its reviewers: they need the lifetime/ownership layer to exist, be tested, and expose a stable interface, so the attention kernel work can consume the plan without redesigning state ownership.
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). It is downstream of the basic path by construction, hence draft.
Purpose
Implements the asynchronous self-speculation × sparse decode attention cell of RFC #156:
vllm_rlt/attention/sparse.py— the §3.2 arithmetic contract (reference statistics + sparse correction, plain PyTorch, checked against the RFC's worked example) and the per-rowSparseDecodePlannamed view: row modes (DENSE/CAPTURE/CONSUME), reference slots, incarnations, retention. The plan rides onDeviceKVMetadataso the attention layer reads it where it already receives its batch. The plain-PyTorch path here is deliberately the oracle, not the serving path.vllm_rlt/worker/attention_reference.py—AttentionReferenceStore, the runner-owned per-token slot pool (§5 "Runner token-state pool"):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 (no host readback), and metadata-only retirement so a batch never materialises a bucket-sized temporary.vllm_rlt/worker/async_speculative.py— per round, the runner reserves one lease per verification row, publishesCAPTUREplans at depth index 1 (draft phase for draft depth ≥ 2, verify phase for draft depth 1, per §6.3) andCONSUMEplans at deeper loops, retires rejected candidates and their dependent descendants from the GPU acceptance counts, and releases slots on ticket retirement (after the final reader completes). 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 (from @Levius-Fubuki's cell, now honoured by both the scalar and the batched path): 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. This removed the full-retention token divergence that was traced to Triton-vs-Torch BF16 rounding (503/512 → 5904/5904 tokens and all exit depths on the real 1.4B model).Explicitly out of scope (capability errors for now): synchronous speculation, token pipelines, CUDA graph capture with sparse enabled, other matrix rows. The fused selector kernel is out of scope for this PR — it is the basic-path cell.
Why this is blocked (measured)
benchmarks/sparse_decode_context_sweep.py, four-loop decode per layer, 32 rows, 24L/16H/128D, BF16, r=0.1, median of 10, dense baseline = the Triton kernel with the feature off:The break-even condition is
S + 2C < 2D(selector, consume, dense). At 32k:S = 310.3,C = 60.1, soS + 2C = 430.5against2D = 26.0— 16.6x over budget, and even a free consume leavesSat 12x the budget. The cause is a per-key cost asymmetry: the selector costs ~10 ms per 1k keys in PyTorch, the fused dense kernel ~0.4 ms per 1k keys, because the Triton kernel returns no probabilities and the selector has to repeat the same O(N) scan 25x slower. Batching cannot fix that (it already gave 6.8–11.1x on consume) — only moving selection into the kernel can.That is the handoff to the basic path: ≥12x faster selection at 32k is the minimum to break even, and the sweep is the harness that will verify it.
Test Plan
tests/test_sparse_decode_attention.py: plan/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, disjoint outstanding rounds, metadata-only retirement allocation bound, byte accounting, plan carriage on device metadata), and end-to-end decode through the real model + paged KV + attention.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.vllm-rlt Version / Commit:
6b164d3(upstream/main) + this branch (f292b76).Test Result
A100 (SM80), torch 2.13.0+cu130:
pytest tests/ -q: 930 passed, 487 skipped, 0 failed.flash_attn_4(FlashAttention-4 is not supported on SM80) — hardware, not a regression.ruff check/ruff format --checkclean on all touched files.Not claimed: serving speedup, TTFT under concurrency, or partial-retention task quality (the quality protocol is proposed but unrun; the sweep suggests it would fail the speed side first).