Repository navigation
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.
…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.
…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.
…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.
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose and integration
High-concurrency follow-up for #157, based on its
048b9d7c79a5711857ce3287812502ad8f274ad0head. This is a stacked, cherry-pick handoff for @kunkunblueberry; please do not merge the cumulative parent stack into main independently. The incremental change is onlyvllm_rlt/attention/sparse.pyandtests/test_sparse_decode_batched.py.The scalar reference path launches/gathers/synchronizes per row, which makes B=32/64 asynchronous sparse decode very slow. Group rows by mode and exact context/retained length, bound gather scratch, and read consume counts once per layer. Preserve the scalar implementation as an oracle. Keep the exact per-row cuBLAS reduction shapes: folding rows into the matmul batch produced real BF16 differences and changed generated tokens in our first prototype, so that prototype is not an accepted speedup.
Cherry-pick these three commits in order, or squash this range after review:
The middle commit corrects the first commit's numerical batching issue; consume the complete range. The final commit prepares views once and uses per-row
bmmwith the same shapes and strides as the scalar einsum. The parent's memory-bounded retirement fix is already in the base and is not claimed as our change.Required human contributor answers
This follow-up reuses the human contributor's scope and four recorded answers in #157's contributor-intent record, grounded in the #156 claim. It serves that shared implementation rather than starting a separate RFC cell.
PR necessity: the maintainer explicitly requested a local-branch PR for the #157 author to cherry-pick. This is that delivery vehicle. The observed problem is documented in the high-concurrency Request changes review. We reproduced the poor B=32/64 throughput on GH200 with a real Ouro-1.4B engine. The parent's
048b9d7fixes the separately reported allocation blow-up; this PR targets remaining reference-path overhead and guards against new numerical drift.Test plan and current status
Candidate:
398333dd08fcb41e60d7a616846fb9d5738fa7d4.Comparison baseline:
048b9d7c79a5711857ce3287812502ad8f274ad0.pre-commit run --all-files: Ruff lint, Ruff format and complete CPU regression hook PASS, on the own GH200_1 runtime withCUDA_VISIBLE_DEVICES="".90120f2: targeted GPU suites 135 passed, 1 skipped (the existing two-GPU reservation case); real captured-input replay is FP32 bitwise exact against scalar for both capture and consume.90120f2real-model B=32: full retention completed with identical tokens/depths to fresh scalar baseline, 18.590 vs 8.430 output tok/s; first r=0.1 observation 33.875 vs 9.462 tok/s. These are single interim observations, not final-head performance qualification.Frozen real-model configuration
One GH200 (97,871 MiB HBM), BF16 Ouro-1.4B; model weight SHA256
58872a72616c736595b8b7662079c5b12c5a162ec16eae94f21c348dfa9885af; PyTorch 2.14.0+cu130, Triton 3.8, OMP threads 4. NativeStageEngine, Triton attention, async scheduling, default multistream, CUDA graphs disabled; 3 draft tokens, 2 draft loops, 4 target loops, block mode. B=32/64, 256 input tokens, 16 measured output tokens, separate 4-token warmup, temperature 0, ignore EOS, fixed 4 loops. Cache block size 16 and 2368/4672 blocks; max batched tokens 4096, prefill chunk 256, admission scan 128. Each case uses a fresh process. Prompts are deterministic token IDs: request i uses[1 + (j * 137 + i * 53) % 49000 for j in range(256)].Sparse
key_retention0.1/1,key_bucket=None,max_context_tokens=None(the original 65,536 model context bound). We did not shrink the reference capacity to 512 to obtain the comparison. Timing includes enqueue, prefill, decode, output delivery, drain and CUDA synchronization. Output token IDs, exit depths, completed requests, actual concurrency, reference slots and KV drain are checked.These are bounded short-burst engine measurements, not HTTP serving, steady-state performance or model-quality evaluation. The original full-retention-versus-dense greedy discrepancy reported on #157 is not claimed fixed. Same-algorithm output equality does not establish sparse quality parity with dense.