Repository navigation
Sparse decode attention: per-token reference lifecycle for async self-speculation - #157
kunkunblueberry wants to merge 27 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.
|
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
|
… store Consume was the last item costing more than dense attention. Two changes: - kernels/sparse_consume.py: selected_scores walks the retained indices and writes one FP32 score per (row, head, retained), selected_values accumulates p * V into one output row. Neither materialises K or V, which removes the two 8 ms gathers and the per-row cuBLAS calls the Torch path made at 32k. - The store no longer builds a device tensor from a Python list. capture, consume and bind_positions all did plan.slot_ids[rows], a blocking host-to-device copy worth 2.4 ms per call, and fused_capture_scores read the context bound back with .item() twice per layer, a stream sync that RFC InferMatrixLab#156 section 6.2 forbids. The runner now passes the context and raggedness it already knows, and the slot index is derived on the device when the rows are contiguous. - The capture and consume paths read and write the store through views when the chunk's slots are contiguous, falling back to the indexed form otherwise. Measured at 32 rows / 24L / 16H / 128D / BF16, r=0.1, median of 10, dense Triton baseline. Four-loop sparse vs dense: ctx before today after fused capture now 512 4.06x 2.57x 1.38x 2048 4.43x 2.07x 1.41x 8192 6.39x 2.57x 1.95x 32768 5.86x 1.92x 1.74x Capture alone is 221.0 -> 34.1 ms at 32k and 13.5 -> 5.5 ms at 512; consume is 36.2 -> 22.0 ms and 12.8 -> 3.0 ms. The remaining gap is per-layer fixed cost, not per-key work.
|
Fused consume is in, and the whole column is now within 1.4–2.0x of dense instead of 4–6x. Two changes
No host round-trips left in the store — three of them, all measurable:
Measured (32 rows, 24L/16H/128D, BF16, r=0.1, median of 10, dense = Triton)
Per step: capture 221.0 → 34.1 ms at 32k and 13.5 → 5.5 ms at 512; consume 36.2 → 22.0 ms and 12.8 → 3.0 ms. Cumulative over the day the four-loop total is 3.4x faster at 32k (306.9 → 91.2 ms), 3.3x at 8k, and 3.0x at 512. What is leftThe gap is now per-layer fixed cost, not per-key work. At 512 with Closing the last 1.4x needs either a single kernel that does capture end to end (scores, selection, contribution, store write) or a CUDA-graph capture of the per-layer Python sequence — both are larger changes than anything in this branch so far, and the second one interacts with the graph support that is currently disabled for sparse. Regression: CPU Next unless you would rather I stop: engine-level end-to-end (real Ouro-1.4B, concurrency 1/8/32) to confirm the per-layer gain survives scheduling, and updating the PR description, which still describes the branch as blocked. |
Runs the real StageEngine with asynchronous self-speculation and compares dense against sparse decode at full and reduced retention: throughput, peak memory and per-request token agreement, across concurrency 1/8/32. The result is a correctness pass and a performance diagnosis. Full retention reproduces dense tokens exactly at every concurrency (1/1, 8/8, 32/32), so the merged O_ref contract holds through scheduling, sampling and delivery. Throughput is 5-19x worse on the tiny model and the gap widens with concurrency, which shows the remaining cost is per-layer Python fixed cost scaling with rows, not per-key attention work: with two layers and 48-token contexts there is nothing to amortise it against. A serving answer needs a trained Ouro-1.4B checkpoint, which this host does not have.
|
Engine-level e2e is in ( Correctness: cleanFull retention reproduces the dense engine's tokens exactly at every concurrency: 1/1, 8/8, 32/32. That means the merged Throughput on tiny Ouro (2 layers, async self-speculation, prompt 32, max_tokens 16, median of 3)
5–19x slower, and the gap widens with concurrency. That is the useful part of the result: the per-layer sweep says 1.4–2.0x at 32 rows, but the engine says the cost grows with rows while the per-key work does not. With two layers and 48-token contexts there is nothing to amortise the fixed cost against, and dense batching keeps getting better while the sparse path adds per-row Python work. So the remaining work is Python overhead, not attention mathConcretely, per layer: Two ways out, both larger than anything in this branch: precompute the per-layer plan and chunk indices once per submission instead of per layer (the runner already knows every field), or capture the per-layer Python sequence in a CUDA graph — which currently conflicts with sparse being excluded from graph capture. What I am not claimingThese are tiny-model numbers on a randomly initialised 2-layer architecture. They cannot answer whether 24 layers at 8k–32k amortise the fixed cost enough to win; the per-layer table suggests it is close (1.74x at 32k with 24 layers), but that needs a trained Ouro-1.4B checkpoint and this host has none. I would rather say that than extrapolate. Regression unchanged: CPU 969 passed / 614 skipped, GPU 303 passed / 1 skipped / 13 deselected. |
|
Thanks for Trained Ouro-1.4B through your e2e harnessI ran the structure of
So in eager mode, neither sparse path beats dense at these sizes, even with the real model. That is consistent with your tiny-model result. On this host, eager decoding is itself dominated by host time. Dense with CUDA graphs is 322 tok/s at 1k × 8, against 59.7 in eager. That points at the second way out you mention. With CUDA graphs, the same harness on the same GH200 puts the fused path at parity with dense: 0.92x at 1k × 1, 1.00x at 1k × 8, 0.98x at 8k × 1, and 1.02x at 8k × 4. The wall time includes prefill, which is the same for both arms. GH200 is the hardest case for sparse; the kernel breakdown below shows why. CUDA graphs are where it pays offCapturing the per-layer sequence in CUDA graphs works on a local branch. Selection counts are decided on the device, so there is no
The gain shrinks on a GH200: block speculation is +7.7% at B=2 × 32k, and a token pipeline loses 1.3% at 14.6k. An nsys breakdown of the pipeline case (106 decode steps, graphs):
At batch 1 with short contexts, it is parity. Merging with your headMy eager-scope fused path merges cleanly onto
One numerical note: on sm_90, the split capture rounds FP32 outputs up to 1.8e-7 from the dense kernel. BF16 is bit for bit, and the FP32 test now allows one ulp. If it helps, I can open the eager part against your branch as one commit, and the graph / P/D / token-pipeline part as a stacked follow-up. |
|
@kunkunblueberry I think this PR finally has good result, and can be opened for review after such a long iteration. |
|
Please resolve the conflicts |
Token-id prompts drive a trained model into a fixed point, so every implementation agrees and the comparison says nothing. This encodes real text with the checkpoint's tokenizer and reports throughput, peak memory and per-request token agreement against the dense engine. Measured on Ouro-1.4B (24 layers, 16 KV heads, 200-token prompt, 16 tokens out, median of 2, A100): concurrency 1 dense 4.77 tok/s, sparse r=1.0 3.74, r=0.1 3.87; concurrency 8 dense 42.75, r=1.0 17.27, r=0.1 22.01. Full retention reproduces dense tokens exactly at both concurrencies.
|
Downloaded the trained checkpoint and re-ran the e2e on it. Setup
Results
This is much better than the tiny model, and it confirms the shape of the remaining cost: at concurrency 1 the feature is 20% slower than dense, not 5x, because 24 layers amortise the per-layer fixed cost. At concurrency 8 it is 2–2.5x slower, because that cost scales with rows while the dense Triton kernel keeps batching better. That is the same diagnosis as before, now measured on the model that matters. Two findings worth recording1. Token agreement is a weak quality signal, and I nearly reported a false pass. The first run used synthetic token-id prompts and the model collapsed to a fixed point (token 0 sixteen times), so every implementation "matched". After switching to real text I instrumented the retained probability mass:
So at 2% retention the correction discards 17% of the attention mass on average (44% at the worst row) and the 16 greedy tokens are still identical. The sparse path is definitely executing ( 2. A corrupted checkpoint looks like a broken implementation. My first download assembled to the right byte count but the wrong sha256, and the loader accepted it (it validates keys and shapes, not hashes). The model then emitted token 0 forever, which reads exactly like a broken attention path. Worth adding a hash check to the load path, or at least documenting that shape validation is not integrity validation. What I am still not claimingPeak memory at concurrency 8 is noisy across configurations (dense 18.60 vs r=1.0 14.24 GiB) because the KV pool and the reference pool are sized differently per run, so I would not read a memory conclusion from that column yet. Partial-retention task quality is still unmeasured: the mass table above is a mechanism check, not an accuracy result. |
Adds --model-path, --cache-blocks and --sparse-context so the same harness can run against Ouro-1.4B. The KV pool has to be sized to the workload: 4096 blocks at 24 layers, 16 heads and 128 dims is 51 GiB of K and V, which OOMs; 768 blocks fits a 32-request, 200-token burst.
# Conflicts: # README.md
|
@0z5a conflicts resolved: I reproduced the CI steps locally first (ruff 0.13.2 On your data — thank you, this is the piece I could not get. My e2e topped out at 0.78x at concurrency 1 and 0.40x at 8 on the A100, and I had no way to test the graph hypothesis from here. Your numbers make the ordering unambiguous:
That also settles the argument about what this PR is: the per-layer work is done, and the remaining 2.5x at high concurrency is host time that graphs remove. Yes please open the eager part ( Two notes on your conflicts description: I want the store to end up exactly as you describe — one helper, whole-plan device slot ids when they are contiguous, device arange for contiguous rows, pinned copy otherwise — that is strictly better than my One finding from my side that affects how we report quality: I instrumented the retained probability mass, and token agreement is a weak signal. At r=0.02 on Ouro-1.4B (10 of 200 keys) the correction discards 17% of attention mass on average and 44% on the worst row, yet the 16 greedy tokens are still bit-identical to dense, at every concurrency. Any quality gate that compares short generations will pass trivially. The mass table is in my previous comment; if your graph campaign can log retained mass per step, the two data sets would line up. |
|
Opened as one commit against your branch: kunkunblueberry#2 ( On retained mass: agreed that short-generation token agreement says little. I'll log the retained mass per step in the graph runs, so the two data sets can be lined up. |
SparseDecodeAttentionConfig(implementation="fused") routes uniform capture or consume submissions to Triton kernels with no per-layer host readbacks; the default batched path and its oracle contract are unchanged, and it keeps the fused consume from 4802af6. - Capture is the dense Triton split kernel that also keeps the scores, so O_ref matches the dense output (bit for bit in BF16; within one ulp in FP32 on sm_90). It builds the first radix histogram per part; selection is a 4 x 8-bit radix select over 16 parts per (row, head), each pass closed in a fixed order by the last part to finish (integer tickets), writing positions only. The retained contribution is gathered over the compact list. Backends other than Triton keep their call plus one serial selection kernel. - Consume is split-K over the retained ranks with a fixed-order merge. Its per-stream scratch follows rows, heads and the padded head dimension. - Retained counts are decided on the device; a batched consume after a fused capture reads them back once. - The store writes its metadata without host-scalar copies: one row-slot helper (the whole plan reuses its device slot ids, contiguous rows use a device arange, other rows a pinned copy), index_fill_ retirement, a readback-free masked accept that drops the round's host counts, and index storage that grows on demand within the logical bucket (the engine budgets the growth peak). - Reused launches call the compiled Triton 3.8 kernel directly, skipping its Python dispatch. The key function is generated per argument-type layout and mirrors Triton's specialization; the tests assert on every reuse that it is the kernel Triton would pick. Measured with benchmarks/sparse_decode_context_sweep.py (32 rows, r=0.1, RTX PRO 6000): four-loop sparse/dense 0.66-0.73x at >= 8k, 1.1x at 2k, 1.6-1.7x at 512. Eager GSM8K-87 at batch 1 (RTX 5090): 1.064x dense wall time, 61 vs 59 correct. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
A speculative ticket releases its leases from the host's current stream, but the store clears slot metadata on its owner stream. The slot index was copied and allocated on the current stream and freed when release returned, so a later allocation on that stream could overwrite it before the owner stream ran index_fill_, clearing other slots or indexing out of range. With CUDA graphs the host runs far enough ahead for this to surface as an index out-of-bounds device assert. The index is now built inside the owner stream. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Thanks for marking this ready, and for getting CI green on
Once it is merged and CI is green, I'll approve with the scope as it stands:
|
Brings in #2: - kernels/sparse_decode.py: opt-in fused capture/consume kernels, the path both measurement campaigns used; - attention_reference.py: build the retirement index on the store's owner stream, which removes a use-after-free that only showed up under CUDA graphs.
|
Merged: Both commits are in, and I read them rather than just taking the merge:
Verified locally on the merged head, in CI's own terms: The description now documents the Ready for your approval. After it lands I will take the stacked graph / P-D / token-pipeline follow-up, and I will label it GPU-dependent with the GH200 numbers next to the L40S ones rather than quoting the best case alone. |
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. |
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 Plan
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, 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.benchmarks/sparse_decode_context_sweep.py: per-layer dense-baseline sweep and memory breakdown.benchmarks/sparse_decode_engine_e2e.py: engine-level throughput, peak memory and token agreement across concurrency.vllm-rlt Version / Commit:
88c8f96(upstream/main, merged here) + this branch.Test Result
A100 (SM80), torch 2.13.0+cu130:
pytest tests/ -q: 981 passed, 652 skipped, 0 failed.-k "not flash_attn_4". The deselections are allflash_attn_4(FlashAttention-4 is not supported on SM80).O_refcontract.ruff check/ruff format --checkclean on all touched files.Not claimed: a serving speedup, TTFT under concurrency, or partial-retention task quality. The feature stays
enabled=falseby default.