Repository navigation
Sparse decode attention: per-token reference lifecycle for async self-speculation - #157
kunkunblueberry wants to merge 43 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.
…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.
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
Two additional findings from the sparse/reference review.
bjf-frz
left a comment
bjf-frz
left a comment
There was a problem hiding this comment.
Clarify the helper's lifecycle responsibility.
Review comments on the reference store, the runner and the plan types. **reserve is exception-safe.** It popped slots from the pool before `_index()` and `torch.arange()` could fail, so an allocation error leaked the slots and the next request failed with "no retired attention reference slots". It now peeks, allocates, and only then removes them from the pool. test_failed_reservation_keeps_the_pool_intact injects the failure and checks that a later request still reserves. **Write-only state removed.** `valid`, `uses`, `position` and `accepted` were written by the store and by three kernels and never read back: availability is the retained count and ownership is the lease. The tensors, their writes, the `bind_positions` API and the corresponding kernel arguments (`VALID` in `_capture_kernel` and `_select_pass_kernel`, `USES` in `_consume_merge_kernel`) are gone, which also drops 12 bytes per slot from the accounted payload. `capture` and `consume` remain as liveness assertions. **Naming and types.** `AttentionReferenceStore` is `SparseDecodeReferenceStore`. `SparseDecodePlan.owners` and `.store` carry real types through a `TYPE_CHECKING` import, and the unused `key_bucket` field is gone. **Runner.** `_generations` and `SpeculativeRequestState.generation` are removed; the value had no production reader left. The class-level `reference_store` / `sparse_decode` defaults that existed for test doubles are gone, and the one double that skips `__init__` now sets both fields explicitly. Regression: CPU 981 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 217 passed / 4 skipped; ruff 0.13.2 clean.
|
All six are in 1.
|
Two review findings that are not style. **Fragmented slots were read as a span (P1).** The batched path accepted a chunk when `first + len == last + 1`, which `[5, 3, 7]` satisfies, and then used `slice(5, 8)`. CONSUME read slot 6 — another request's reference — and CAPTURE overwrote it while leaving slot 3 unwritten. `slot_span` now checks every slot against `first + offset` and returns None otherwise, so a fragmented chunk takes the indexed path. Both call sites use it. `test_fragmented_consume_reads_each_rows_own_reference` and `test_fragmented_capture_writes_only_the_slots_it_reserved` build the exact fragmentation from the report (release 7, then 3, then 5; reserve three rows gives `[5, 3, 7]`) with slots 4 and 6 held by unrelated live leases. Verified to fail against the endpoint check: slot 3 keeps its sentinel, which is the reported symptom. **Kernel scratch outlived its engine (P1).** `_WORKSPACE` and `_CONSUME_WORKSPACE` were module-level dicts keyed by (device, stream) that nothing released, so a closed engine's tensors stayed reachable — the capture score buffer alone is rows x heads x capacity floats. They move to `kernels/sparse_workspace.py` (Triton-free, so CPU-only installs can release them) and `SparseDecodeReferenceStore.close()` releases its device after the caller has drained; `StageEngine.close()` calls it after its own drain. **Cleanup.** `_reference_metadata` is `_prepare_sparse_reference` on both the runner and the speculative runner, since it reserves and releases leases rather than building metadata. `SparseDecodePlan.reference_depth` is gone: it returned a module constant and had no production reader. Regression: CPU 986 passed / 650 skipped; GPU (four sparse modules, async speculative, lifecycle) 222 passed / 4 skipped; ruff 0.13.2 clean.
|
Both P1s and the cleanup items are in P1: fragmented slots were read as a span — confirmed, and worse than a read
def slot_span(slots: Sequence[int]) -> slice | None:
first = slots[0]
for offset, slot in enumerate(slots):
if slot != first + offset:
return None
return slice(first, first + len(slots))Both call sites use it (CONSUME read, CAPTURE write). The regression reproduces your exact fragmentation.
I verified both fail against the endpoint check by temporarily restoring it: Scope: the fragmentation regression runs on CPU. I did not build a GPU-specific fragmentation case, so I am not claiming GPU coverage beyond the existing suite. P1:
|
|
ROCm compatibility check: I merged #171 (ROCm platform support) into this branch as This is compatibility evidence, not a performance or accuracy claim: the branch's own numbers still show sparse slower than dense on this device. Branch: https://github.com/0z5a/vllm-rlt/tree/rocm-sparse |
Priority scheduling cuts high-priority completion 7.80s to 2.22s; incremental KV gives +3.8%; sparse decode InferMatrixLab#157 passes 137 tests on the merged ROCm branch.
RFC InferMatrixLab#156 sections 3.1, 6.1 and 6.2. One host-side description of which rows of a recurrent loop execute, for token-sparse prefill, sparse decode attention and WCPB budget accounting: - TokenSparsePrefillConfig (section 7.1): off by default, nested retention with two dense loops, selection window aligned with the prefill chunk when unset. - SelectionWindows: absolute window geometry, max(1, ceil(f * len)) budgets and execution_rows(), exact for whole windows and bounded by min(overlap, budget) for fragments. - select_active(): FP32 top-k nested in the previous set, ties broken by the smaller position, no forced final token. - PrefillSelections and ActiveRows: runner-owned active sets per request, window and depth; a per-loop view with generated-token rows first, then active prompt rows, original positions, inactive rows for KV inheritance and take() for per-row plans such as InferMatrixLab#157's modes. A loop with no inactive row keeps the identity order. Nothing is wired into the scheduler, runner or model yet; the feature stays off and the dense path is unchanged. Signed-off-by: MinhaoLi0318 <lminhao039@gmail.com>
Priority scheduling cuts high-priority completion 7.80s to 2.22s; incremental KV gives +3.8%; sparse decode InferMatrixLab#157 passes 137 tests on the merged ROCm branch.
|
@0z5a thank you for running it on AMD — that is the one platform I could not reach, and "no sparse-kernel change was needed" is the useful part: it means the kernel only ever depended on #171's platform selection and Triton's I have recorded it in the description under Test setup, worded exactly as you framed it: compatibility only, not a performance or accuracy claim — 137 passed / 3 skipped on one MI350X (gfx950, ROCm 7.2, torch 2.14.1+rocm7.2, Triton 3.5.1), with the sentence that the branch is still slower than dense there and that the CUDA numbers do not transfer. If you would rather it read differently, say so and I will change it. Two other things happened while you were testing: The branch is rebased on All review threads are resolved, including the two P1s from the last round — fragmented slots ( One thing I want to be explicit about for the record, since your ROCm branch is where the sparse kernels would next be touched: #171 is still open, so this PR does not depend on it and does not claim ROCm support. When #171 lands, a rebase should be enough — your run is the evidence for that, and it belongs with #171 rather than here. |
|
Hi @kunkunblueberry , I previously volunteered to help with prefix-cache validation for sparse decode attention in #156 (comment). I’ve now opened kunkunblueberry#4 against this feature branch. The PR adds regression tests for prefix-cache hits and reuse, using cold baselines, and checks that generated-token sparse references remain request-local. It also covers concurrent requests, cancellation and slot reuse, and includes a fix for the folded capture histogram layout. The PR description has the test results, including failures reproduced on the base branch. Could you please review the PR at your convenience? Thanks! |
|
What it covers, all on the synchronous path: prefix-cache hits compared against a cache-disabled cold baseline of the same workload (prompt lengths 63/257/1025, 2–4 concurrent requests), request-local reference ownership when two requests share a prefix, cancellation between capture and consume with physical slot reuse, and prefix pages verified byte-identical after the run. The harness self-tests — One line of production code changes. Verified on the merged head (A100, torch 2.13.0+cu130, Triton 3.7.1): CPU 999 passed / 725 skipped, GPU across seven modules 256 passed / 4 skipped, ruff 0.13.2 clean. One thing I want on the record rather than buried: @zignaix-t's run at the declared floor (torch 2.5.1 / Triton 3.1.0) still shows five failures — four non-split fused fallback cases and one async tensor-transfer compile failure — none of which this fix addresses, and none of which reproduce on 3.7.1. So |
@zignaix-t followed up on the two questions from the prefix-cache review with controlled comparisons on the same A40-8Q, which change the record: - Hoisting the tl.where alone is NOT enough on Triton 3.1.0. Six folded-capture regression cases scored 3/6 for the original code, 3/6 with hoist-only and 6/6 with the vote reorder; all three variants pass on 3.7.1. The kernel comment now carries those numbers, the divisibility requirement on BLOCK_T, and the fact that the reorder's cost has not been measured. - The five remaining fused failures are Triton-version-associated, not introduced by the histogram fix: they fail on 3.1.0 and pass on 3.7.1 with the same GPU, driver and torch, for both code variants (56 executions, no skips). Intermediate versions are untested, so the minimum working Triton is unknown. We therefore document the limitation instead of raising the declared floor to an untested version. docs/runtime.md now says the fused implementation is only verified on 3.7.1 and requires a newer Triton than the declared triton>=3.1, while batched stays the default and keeps the bitwise contract. The PR description carries the same note. Regression: CPU 999 passed / 725 skipped; folded-capture GPU cases 6 passed; ruff 0.13.2 clean.
3135432 to
8fc2427
Compare
|
A correction and an update to my previous comment, both from @zignaix-t's follow-up measurements ( Correction. I wrote that the five failures on Triton 3.1.0 were "none of which this fix addresses". That was literally true but left open the wrong reading — that the histogram change might have introduced them. They reran all five with the original production code and with the fix, keeping GPU, driver and torch 2.5.1 fixed and switching only Triton in isolated environments: 3.1.0 fails with both variants, 3.7.1 passes with both (56 executions, no skips). So they are Triton-version-associated and predate the fix. Intermediate versions were not tested, and the torch 2.5.1 / Triton 3.7.1 diagnostic pairing is not officially supported. Both open questions from the prefix-cache review are now settled.
Regression on this head: CPU 999 passed / 725 skipped, folded-capture GPU cases 6 passed, ruff 0.13.2 clean. |
|
Description audit: the body was carrying two generations of text, and the stale half contradicted the measured half. Removed the leading block ( Corrected four stale statements in the sections I kept, each checked against the tree at
Added an explicit open-items list instead of leaving it to be inferred: partial-retention task quality unvalidated, The measured tables, the four contributor-intent answers and the claim boundaries are unchanged. One stale reference I did not fix, because it is a code change and this head is under review:
@0z5a your approval is at
|
`SpeculativeRoundTicket` has been `AsyncSpeculativeOutput` since 0ce8ca4 renamed it, but the documentation table for the async speculative path still named the old class, so that row pointed at something that does not exist.
The batched capture and consume kernels load the physical page id as int32 from the block table and then multiply it by the physical block stride. With the layer-interleaved KV layout the element offset crosses 2**31 well inside a supported cache size, so the product wrapped and addressed memory outside the allocation. On an H800 that surfaced as an illegal memory access in the shared selected-value kernel. This is @Levius-Fubuki's fix from #3 (12e8049), applied unmodified to the two kernels, so his reproduction and regression test carry over as the evidence for it: `test_sparse_score_and_value_kernels_address_large_kv_offsets` builds a BF16 cache whose physical block 2 begins at element 2**31 and asserts the capture and consume kernels read from it. The fused kernels in `kernels/sparse_decode.py` already promote their block ids, which is why the capability change in #3 is separate from this fix: only the batched path was mis-addressing, and it is the default.
|
Two commits landed on the head since the description cleanup. Both are small and neither changes reviewed semantics:
Why the second one is here instead of a follow-up. It is a defect in the default Verification on @0z5a the approval request above now targets |
Status: opt-in and correctness-complete; not a serving speedup on its own eager path
This branch is no longer blocked on anyone else's work. It carries the per-token reference lifecycle (§5), the numerical contract that makes the result checkable, and both fused kernels. It does not beat dense attention end to end on its own eager path, and that is a scoped result rather than an open question:
kretained keys.Defaults stay
enabled=false, andimplementation="batched"keeps the bitwise contract.Open after this PR, stated so it is not read as done:
implementation="fused"is verified on Triton 3.7.1 only (see the known limitation indocs/runtime.md). The minimum working Triton above the declaredtriton>=3.1floor is unknown, and CI installs the latest torch, so it cannot see this..github/workflows/cpu-tests.yml). Every GPU number in this description comes from the runs cited below, not from CI.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, a new allocation identity after same-position replacement, reuse ordered after the last reader/writer) are the contract every later kernel change must respect; this branch enforces that identity with the lease object rather than a counter.
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—SparseDecodeReferenceStore, the runner-owned per-token slot pool (§5):empty → source-ready → reusable → retired, one slot per live token, each slot bound to the lease that reserved it so a plan from a previous owner is rejected after same-position replacement or cross-request reuse, 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 when the round retires. Outstanding rounds hold disjoint slots.config.py/entrypoints/stage_engine.py— thesparse_decode_attentiongroup (enabled, default false;key_retention;key_bucketas index-buffer capacity only;implementation) with explicit capability checks and reference-state memory accounted before admission.Numerical contract: dense and capture rows keep the selected backend's output, the actually rounded recurrence input is stored as
O_ref, and full-retention consume rows use the same dense backend as the control. The fused capture pass computes the same FP32 dot product as the oracle with a different reduction order: max score delta 1.1e-5 on a magnitude of 65, retained mass delta 9e-8, attention output delta 2e-8, identical selected sets on 511/512 row-head pairs.batched_paged_attentionwithoutfused_scoresstill reproduces the scalar oracle bit for bit.Explicitly out of scope: synchronous speculation, token pipelines, CUDA graph capture with sparse enabled, other matrix rows.
Test setup
Hardware and software. A100 (SM80), torch 2.13.0+cu130, Triton attention, ruff 0.13.2 (the version CI pins). GPU suites run with
--run-gpu -k "not flash_attn_4"; the deselected cases areflash_attn_4, which does not support SM80.Model.
ByteDance/Ouro-1.4Bat revision574fa66c, sha256 verified against the LFS digest (58872a72…): 24 layers, 16 attention heads, 16 KV heads, 65,536 context, 4 UT steps, 1.435B parameters, BF16. Properties that do not depend on trained weights are checked on the test-only tiny Ouro (2 layers, 128 context).Regression tests kept in this PR.
tests/test_sparse_decode_attention.py— plan and arithmetic contracts (RFC §3.2 worked example = 23.5, full retention ≡ dense, ceil retention rule, capacity), store lifecycle (one slot per live token, owner-checked release, acceptance-driven invalidation at accepted counts 0/1/2, root survival, stale-plan rejection after a new lease and after cross-request slot reuse, disjoint outstanding rounds, metadata-only retirement allocation bound, byte accounting, plan carriage), end-to-end decode through the real model + paged KV + attention, and engine-level full-retention parity.tests/test_sparse_decode_batched.py(from @0z5a) — the batched path must be scalar-exact.tests/test_sparse_decode_sync.py(from @Levius-Fubuki) — ordinary decode and synchronous speculation at draft depths 1/2.tests/test_sparse_decode_fused.py(from @0z5a) — the opt-in fused kernels.tests/test_sparse_decode_prefix_cuda.pyand the prefix cases intest_sparse_decode_sync.py(from @zignaix-t, PR Q1: qualify BF16 numerics and validate last-exited KV independently #4) — prefix-cache reuse under synchronous sparse decode: cache hits checked against a cache-disabled cold baseline of the same workload, request-local reference ownership, concurrent requests, cancellation between capture and consume with physical slot reuse, and prefix pages verified byte-identical afterwards. The harness self-tests:test_sparse_trace_rejects_missing_duplicate_and_foreign_stepscorrupts the trace five ways and requires the checker to reject each.vllm-rlt Version / Commit:
e0a5515(upstream/main, merged here) + this branch.Fused-path Triton limitation.
implementation="fused"is verified on Triton 3.7.1 but not at the declared floor: on torch 2.5.1 / Triton 3.1.0 the folded capture histogram wrote the wrong selected indices (fixed here) and five other fused cases still fail. All five pass on the same hardware with Triton 3.7.1, with or without the histogram fix; intermediate versions are untested, so the minimum working Triton is unknown. CI installs latest torch and cannot see this. See the limitation note indocs/runtime.md.ROCm compatibility (compatibility only, not a performance or accuracy claim). @0z5a merged #171 (ROCm platform support) into a scratch branch and ran the four sparse modules on one MI350X (gfx950, ROCm 7.2, torch 2.14.1+rocm7.2, Triton 3.5.1): 137 passed, 3 skipped. No sparse-kernel change was needed beyond #171's platform selection and Triton
ieeedot precision. This says the kernels are portable and nothing more: on that device the branch is still slower than dense, and the CUDA numbers above do not transfer.Test results
Regression on the current head. CPU
OMP_NUM_THREADS=1 pytest -q: 999 passed, 725 skipped, 0 failed. GPU (async speculative, async lifecycle, five sparse modules including the fused path and prefix-cache reuse): 256 passed, 4 skipped.ruff check .andruff format --check .clean under 0.13.2.Numerical contract. The fused capture computes the same FP32 dot product as the scalar oracle with a different reduction order: max score delta 1.1e-5 on a magnitude of 65, retained mass delta 9e-8, attention output delta 2e-8, identical selected sets on 511/512 row-head pairs at 8k context. Full retention reproduces the dense engine's tokens bit for bit at concurrency 1/8/32 on the tiny model, and 5904/5904 tokens across all exit depths on the real 1.4B model (@Levius-Fubuki's run for the merged
O_refcontract). On SM90 the split capture rounds FP32 to within 1.8e-7 of the dense kernel while BF16 stays bit-exact, which is why the FP32 test allows one ulp.Per-layer cost (32 rows, 24 layers, 16 heads, head_dim 128, BF16, r=0.1, median of 10, dense = Triton):
Engine end to end (Ouro-1.4B, 200-token real-text prompt from the checkpoint's tokenizer, 16 tokens out, median of 2, eager, A100):
Eager decoding on this host is host-bound, so these numbers describe the batched path, not the algorithm's ceiling. With CUDA Graphs the fused path reaches parity on GH200 (0.92x at 1k × 1, 1.00x at 1k × 8, 0.98x at 8k × 1, 1.02x at 8k × 4) and gains on lower-bandwidth GPUs (+9.3% and +9.6% on RTX PRO 6000 at 32k and 16k, +7.3% on L40S at 32k, +13.5% for a token pipeline at 22.5k). That work is a stacked follow-up, not this PR.
Retained mass — why token agreement is not the quality signal. On Ouro-1.4B, r=0.02 retains 10 of 200 keys and discards 17% of the attention mass on average (44% on the worst row), yet the 16 greedy tokens stay bit-identical to dense at every concurrency. A quality gate that compares short generations will pass trivially; retained mass, logit margins or task accuracy over a real cohort are the measurements that discriminate.
Evidence attachments. The numerical diagnosis and regression report is kept as a release attachment. The benchmark harnesses that produced the tables above were removed from this PR at review request; their setup is described here and their results are the numbers above.
Not claimed: a serving speedup on this PR's own path, TTFT under concurrency, or partial-retention task quality. The feature stays
enabled=falseby default, andimplementation="batched"keeps the bitwise contract.