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.
…tion Brings in the basic/synchronous cell from @Levius-Fubuki (PR #1 in this fork), based on 048b9d7: - ordinary decode and synchronous block speculation at draft depths 1/2, per-input reference ownership, rejection/replacement, cancellation, capacity checks and stream-safe retirement; - numerical contract fix: dense/capture outputs keep the selected backend and the actually rounded recurrence input is stored as O_ref; retention 1.0 uses the same dense backend as the control, which removes the full-retention token divergence traced to Triton-vs-Torch BF16 rounding.
…ched path The cherry-picked batching commits replaced sparse_paged_attention with a dispatcher and dropped the dense_attention handoff, so _attend_prepared would have raised TypeError and the numerical contract would have been lost. Both implementations now honour it: - the dispatcher forwards dense_attention to either implementation; - the batched path uses the backend output for dense rows, stores the actually rounded recurrence input as O_ref for capture rows, and returns the backend output directly for full-retention consume rows, while keeping selector probabilities and the retained contribution in the Torch FP32 oracle. benchmarks/sparse_decode_batched_ab.py isolates the batching gain from that contract by timing both implementations on the same plan and store.
Measures the four-loop decode cost against the dense Triton kernel (the acceptance baseline, not the scalar sparse reference) across 256-32768 tokens, plus the three memory components: KV storage, the persistent reference pool and the transient capture/consume workspace. The result is the diagnosis the review asked for: capture = dense + selector, and the selector costs about 25x per key relative to the fused kernel, so no context in the sweep is a net win.
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.
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).
… 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.
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.
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.
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
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>
Brings in kunkunblueberry#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.
…ecode On top of the eager fused path: - Async speculative CUDA graphs keep one entry per bucket and sparse mode. Their slot ids form a static plan refilled before each replay, with a never-bound padding slot for padded rows, and the graph owns its capture workspace with the index pool allocated whole. The engine accepts sparse with graphs only for async block speculation, the fused implementation and r < 1, and budgets the graph workspace and selection scratch. - P/D passes the sparse decode configuration to the decode stage only. Measured end to end with Ouro-1.4B, async self-speculation K=3, 512 greedy tokens, interleaved dense/sparse engines: B=2 x 32k +9.3% (r=0.1) / +6.1% (r=0.2), B=4 x 16k +9.6% / +7.2% (RTX PRO 6000); B=2 x 14.6k +3% to +8% (RTX 5090); batch 1 parity. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Ordinary decode builds its reference plans from host-known depths and positions, and every reference read, write and retirement is ordered on the core stream, so submitting a round before collecting the previous one needs no change to the lifecycle; the pool stays one slot per scheduled request. The runner now publishes each capture row's retained count when it builds the plan, since the row's causal length is known there. A batched consume after a fused capture then sizes its gather without reading the device count back. Tests: asynchronous and synchronous ordinary decode produce the same tokens (batched on CPU; batched and fused on CUDA, where every consume is also checked to be sized from a published count equal to the device count); the pool never exceeds one live reference per scheduled request, including abort and id reuse. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Token pipelines advance, shift and restart their rows on the GPU, so the host knows neither a row's loop nor which reference slot it owns. They now carry both on the device: - DeviceSparseDecodePlan holds one int8 mode and one slot id per row as device tensors, plus the owning leases, which the store checks live on the host. Only the fused kernels read it. - The fused kernels take a per-row mode pointer for such plans: the split and selection kernels compute dense and capture rows (consume rows do no dense work, dense rows keep no scores), the consume kernels compute consume rows, and each writes only its own rows. Dense and capture rows stay the dense kernel's output; uniform host plans compile without the mode branches. - Each pipeline request reserves one slot per pipeline row for its lifetime. The slots form a device permutation moved with the rows on every emission, so a row consumes the slot its loop-1 capture wrote; the request's release returns them, and buffer discard drops them from surviving states. - Each recurrent offset derives its modes from the rows' loops (1 captures, later loops consume, inactive rows stay dense). With CUDA Graphs, pipeline runners capture device-plan entries whose slot ids and modes are refilled before replay. - The engine accepts token pipelines with sparse decode only for asynchronous scheduling, the fused implementation and key_retention < 1. GPU tests (to run on CUDA): one mixed-mode launch equals the uniform launches row for row and slot for slot; pipeline sparse tokens equal block sparse tokens (dense pipeline equals dense block as the precondition), with and without graphs and for draft depths 1 and 2; slots follow requests through abort and id reuse; pipelines reject the batched implementation and full retention. docs/runtime.md lists the supported sparse modes. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
On an L40S (sm_89) the mixed-mode kernel variant rounds dense and capture rows up to one ulp away from the dense kernel (1.8e-7), while consume rows and every store field stay bit-identical; the RTX 5090 matched exactly. The mixed-mode test now requires exact consume rows and store fields, and dense and capture rows within 1e-6; uniform launches keep their bit-for-bit test. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
Pipeline rows sit at multiples of draft_loops, so whether any row can be in its capture loop at a given offset is known on the host: with the default two draft loops, offset 0 runs loops 0 and 2 and never captures. Device plans now carry that as may_capture; when it is false the fused path launches the dense split, the first pass (which merges the dense rows) and the consume kernels, four launches instead of eight, and graphs key such steps separately. Outputs and store fields are identical (tested). Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
A token suspended after its reference loop must consume that same reference when it resumes. Suspension used to drop it, so the remaining loops fell back to dense attention; sparse decode therefore rejected preemption outright. Suspension now copies the token's reference slot (retained indices, counts, mass, output, contribution and the per-slot bookkeeping, plus the published host counts) into storage the saved state owns, resident or on the CPU with the offloaded KV, including a later resident-to-offload conversion. Resumption binds a fresh slot and generation and copies it back on the owner stream. Preemption is lifted for ordinary decode, synchronous and asynchronous. Speculative rows hold per-round references with no snapshot, so speculative decoding with sparse attention still rejects preemption explicitly. Tests: a suspension right after the capture loop (resident, offload, and resident then offload) resumes with a sparse consume and the same tokens as an uninterrupted run; KV-pressure preemption matches the unpreempted run, synchronous and asynchronous, with references carried across suspensions; abort while suspended frees everything and the id is reusable. CUDA variant for the batched and fused kernels. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
A suspended token pipeline keeps several live rows, each past or before its capture loop, and each row reads the reference slot the device permutation gives it. The pipeline snapshot now carries the request's whole reference lease and, per row, which lease slot it reads; restoration binds a fresh lease and generation, copies the slots back and rebuilds the permutation, so every row resumes with the reference it captured. Live rows must form a prefix, as the restored pipeline puts them back in rows 0..k-1. Restoring a reference now reserves its slots on the owner stream as well: the slot ids are copied on the current stream and must be on the stream that indexes with them. Preemption is lifted for token pipelines; block speculation with sparse attention still rejects it. Test (CUDA): a priority arrival suspends a sparse pipeline with rows past their capture loop, with and without graphs, resident and offloaded; it resumes to the same tokens as an uninterrupted run and releases every slot. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
…test The schedule is deterministic and a random tiny model rejects most drafts, so a single suspension landed at the same point of the pipeline cycle in every parametrization, with no row past its capture loop; the test stopped at its precondition before checking any output. It now suspends repeatedly, one more submission later each time, until a suspension holds such a row, checks the saved permutation at every suspension, and then compares the output with an uninterrupted run. Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
0z5a
marked this pull request as ready for review
November 2, 2026 15:55
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.
Stacked on #157. Please review only the nine commits after
43d9cc9. Until #157 merges, the diff also shows #157's own commits.This PR adds six execution features to the opt-in fused sparse decode path. Each feature has its own A/B end-to-end measurement below. The speedup depends on the GPU; it is not universal.
Every A/B below compares sparse with dense in the same execution mode.
Common setup and claim rule
1. CUDA graphs for sparse async self-speculation (
1d98a84)What changes
DeviceRecurrentGraphskeeps one entry per (bucket, sparse mode).A/B: dense vs sparse, both with graphs (block speculation, decode throughput):
b0a73bc)The 5090 2 × 14.6k cell was measured on
b0a73bc. The other cells come from earlier revisions of the same kernels.2. Sparse decode on the P/D decode stage (
1d98a84,6f5ebdd)What changes:
pd/engine.pysends the sparse configuration to the decode stage only; prefill stays dense.A/B: dense P/D vs sparse on the decode stage. 14.6k context per request, graphs on the decode stage, throughput over the window in which every request decodes, median of two rounds per engine:
3. Sparse decode with asynchronous ordinary decode (
bc0622f)What changes
A/B: three arms, async, no graphs. 2 × RTX 5090, 2 × 14.6k, decode-window throughput:
maindenseE equals M within 0.1%, so the branch does not slow dense down.
4. Sparse decode in asynchronous token pipelines (
9cb7407,1aba42a,b0a73bc)What changes (
SpeculativeConfig(mode="pipeline"))implementation="fused"and r < 1.A/B: dense pipeline vs sparse pipeline, both with graphs.
b0a73bc)b0a73bc)Against dense block speculation
Why GH200 is different. An nsys breakdown of the 14.6k, two-draft-loop case:
5. Preemption and recovery for ordinary sparse decode (
40056e8)What changes
A/B: M / E / S as in §3, async, preemption enabled. KV fixed at 7,424 blocks, 2 × 14.6k, metric = output tokens / round wall time (prefill included). Every round in every arm had exactly one suspension and one resume. 2 × RTX 5090:
Sparse supports preemption here without a slowdown. The round includes the two 14.6k prefills and an 11.5 GB KV offload and restore, so there is no throughput gain to claim.
6. Preemption and recovery for sparse token pipelines (
0f0569c,44737aa)What changes
store.restorenow reserves on the owner stream. Before, slot ids were copied on the current stream but used on the owner stream.A/B: three arms, async + priority scheduling + preemption, 1 draft loop.
maindense pipeline + graphsWorkload:
Pre-registered primary metric: tokens emitted after the injection, divided by the seconds from the injection to the end of the round. This excludes the prefill that all arms share. One GH200 per cell, measured on
44737aa, median of two rounds per engine:main(−1.6% to +1.5%), and +7.8% to +46.8% against eager dense.High concurrency on GH200 (final head
44737aa)This is not a separate feature. It checks how the token-pipeline gain (§4) holds as the batch grows while the total context stays about 60k tokens.
Setup
maindense (989921a), this PR dense, this PR sparse (r=0.1). All three use the asynchronous token pipeline with one draft loop and CUDA graphs. Prefix caching is off.max_num_seqs = B. Each request gets its own non-overlapping prompt text. The KV cache size is fixed per cell and equal across arms.main, densemainEach cell lists both engines, in tok/s.
mainin every cell, so there is no regression.mainengines ran 2–6% faster than the earlier ones. The mirrored order cancels a linear drift. At 32 requests, even the slower early engines are faster than both sparse engines.main) and +10–14% at 7.5k and above. At 1.9k × 32 requests it loses about 7%. Sparse decode therefore stays opt-in.Quality
Short generations that agree with dense say little about quality. The graph runs therefore logged the retained attention mass of every fresh capture: one value per (slot, layer, head), with retired slots excluded. Measured at 2 × 14.6k with graphs, 512 tokens:
GSM8K, 87 questions (lm_eval strict, batch 1, about 1k context, graphs):
For scale, a dense pipeline and dense block speculation give identical outputs on only 21 of the 87 questions.
Tests
ruff checkis clean, and the touched files need no formatting changes.pytest -qon this head gives 945 passed, 679 skipped (local CPU environment).44737aa(GH200; the same environment also ranmain989921a):--run-gpusuite excepttests/test_pd.py: 1408 passed, 40 skipped. All sparse decode tests pass (172), including the four pipeline preemption CUDA cases, and so do all Triton cases oftest_async_speculative_graphs.main. 171 come from this environment: native CUDA kernels not built (114) and FA3/FA4 not installed (57). The other 9 fail the same way onmain.tests/test_pd.pyon 2 GPUs: 45 passed, 3 skipped (the 4-GPU cases).main: 39 passed, 1 skipped.tests/test_pd.pycases. They passed on 4× RTX 5090, including the 2P2D case, before the preemption commits.