Repository navigation
Epic: Fused paged decode, sorting-free sampling, and serving performance techniques #909
Description
Overview
This epic tracks a coordinated set of serving-performance improvements across attention, sampling, speculative decoding, and kernel tuning. Each sub-issue is written to be self-contained: it carries the problem statement with measured evidence, the algorithm to implement, the mlxcel integration points, and mandatory performance-validation and regression gates.
The headline motivator: since v0.4 the server defaults to batched decode (--parallel 4), and docs/adr/0001-paged-attention-gather-vs-fused-kernel.md measured the production gather-then-SDPA decode path at ~48% overhead for batch 4 at just 1024 tokens of context (2x-3x past 4096). The ADR's own trigger for building the fused path has fired. Around that anchor, the epic adds sorting-free GPU sampling (the vocabulary-sized sort/partition work in the per-token sampling path is avoidable), distribution-preserving speculative acceptance, shared-prefix compute deduplication, sparsity via page indirection, small decode-loop fusions, and a shape-bucketed autotuner.
Sub-issues and dependencies
| # | Issue | Priority | Blocked by |
|---|---|---|---|
| #898 | Fused paged-attention decode v2: CSR page table, cross-CTA split-KV, merge kernels | P0 | (root) |
| #899 | Route production batched paged decode through the fused v2 kernel | P0 | #898 |
| #900 | Softmax-free GPU sampling: Gumbel-max categorical sampling kernel | P0 | (root) |
| #901 | Sorting-free top-k/top-p/min-p sampling with dual-pivot rejection | P1 | #900 |
| #902 | Distribution-preserving speculative acceptance (chain speculative sampling) | P1 | #900 |
| #903 | Cascade attention: compute shared prompt prefixes once per decode batch | P1 | #898, #899 |
| #904 | Fused sparse-attention decode via page indirection (DSA, block-sparse) | P2 | #899 |
| #905 | Fused residual-add RMSNorm and fused RoPE + KV-append kernels | P2 | (root) |
| #906 | Shape-bucketed kernel autotuner and cold-L2 benchmark methodology | P2 | (root) |
| #907 | MLA matrix-absorbed decode with compressed-latent KV cache | P3 | (root; reuses #898 merge infra at Stage 2) |
| #908 | Mixed prefill/decode step execution: design spike and prototype | P3 | #899 |
Wave structure implied by the DAG: wave 0 = #898, #900, #905, #906 (independent roots; #907 Stage 1 may also start here); wave 1 = #899, #901, #902; wave 2 = #903, #904, #908.
Epic-level completion criteria
- Every sub-issue is closed only after its feature is integrated into the actual serving/CLI logic (not a standalone module), its mandatory performance validation ran on real hardware with numbers recorded under
docs/benchmark_results/, and its regression gates passed (token parity where specified, full relevant test suites, kill-switch behavior). - Deliverable at epic close: the production batched/paged decode path runs the fused v2 kernel by default; sampling no longer sorts the vocabulary on the hot path; speculative decoding preserves the target distribution at temperature > 0; and a benchmark report comparing epic-start vs epic-end on the standard serving scenarios (single-client decode, 4-client batched decode, long-context 16K/32K, shared-prefix batch) exists under
docs/benchmark_results/. - No regression anywhere: epic-end numbers for single-client decode and prefill on the standard models are within noise (3%) of epic-start; any intentional behavior change (RNG streams, sampling tie-breaking) is documented in the CHANGELOG.
Working notes for implementers
- Line numbers in sub-issues are indicative as of current
main; search for the named symbols. - Backend expectations per sub-issue: Metal and CUDA are both in scope for kernels unless the issue says otherwise; performance validation runs on at least one Apple Silicon machine, plus GB10 for CUDA-affecting changes when hardware is available.
- Every benchmark result lands under
docs/benchmark_results/following the existing report format, with before/after numbers on the same hardware and the exact commands recorded.
Activity
[-]Epic: Adopt FlashInfer-derived attention, sampling, and tuning techniques[/-][+]Epic: Fused paged decode, sorting-free sampling, and serving performance techniques[/+]test(autotune): stop pinning the rep-scaling test to an exact ceiling (…
test(core): fail loudly when another mlxcel-core test binary shares t…
inureyes commentedon Aug 3, 2026
Epic #909 complete
All eleven sub-issues are closed and their PRs merged.
Benchmark reports for each are under
docs/benchmark_results/, and the mixed-step decision is recorded asdocs/adr/0005-mixed-prefill-decode-step-execution.md.What the measurements actually said
Four sub-issues shipped with their premise corrected rather than their feature built, and four shipped default-off or narrowly routed. That is the honest headline, and it is worth stating plainly because the epic's framing assumed eleven wins.
The recurring pattern: an op-level speedup does not survive contact with the call site. #899's fused decode measured 2.78x on the kernel and 1.15x once routed through production batched decode, where the discriminator turned out to be batch size, not total token count. #901's dual-pivot sampling measured 1.17x on the op and 0.58x, a regression, at the real sampler. #906's autotuner looked like a 1.17x median across 11 of 12 shape cells until a variance guard was added, after which only ctx=16384 reproduced at 1.15x to 1.23x.
Three separate benchmarks in this epic compared a code path against itself and reported the resulting noise as a win. Each was caught by a counter that can only move in one arm, which is now the pattern the harnesses enforce:
examples/sparse_paged_decode_bench.rsandscripts/bench_mixed_step_admission.pyboth refuse to print a table when attribution fails.Integration verification
The full workspace suite was run once on the final merged tree, in a single process with nothing else on the machine, since #1008 makes concurrent runs meaningless:
Two targets failed,
granite4_vision_parityandhunyuan_vl_parity, both of them the pre-existing flaky tests tracked in #997, which was filed from an unrelated branch before this epic finished. Nothing else in the workspace failed, and #964 passed in this run.Defects found along the way
fix(server): a chunked prefill is starved until the decode batch drains #1011 (fixed in fix(server): bound a parked chunked prefill's wait with a fairness grant #1013): a chunked prefill was starved until the decode batch drained, so an admitted request's TTFT was unbounded under continuous load. Found by the Mixed prefill/decode step execution: design spike and prototype #908 spike, which set out to fix the opposite symptom. Measured: the prefill held at chunk 1 of 19 for 20 seconds while decode advanced 138 to 394 steps.
The cause was structural.
decide_tickis a pure function ofTickState, andTickStatecarried no history, so alternation was not merely absent, it was inexpressible. The fix is a grant counter (--prefill-grant-interval, default 16) returned fromdecide_tickalongside the choice, so the one application site cannot take a decision without its successor counter. Bound: one chunk perN + 1ticks. Measured TTFT 50.5 s to 23.8 s, mean stream ITL 1.60x during the admission window only, p50 unmoved.It produced one more corrected premise, and this one is about measurement itself: the p95 figure both fix(server): a chunked prefill is starved until the decode batch drains #1011 and ADR 0005 asked for is the wrong number to tune on. Grants land at
1/(N+1), so p95 catches them below interval 19 and misses them above, relocating the same thirteen hiccups to p99 without shrinking any of them.--prefill-chunk-sizeis the lever that shrinks a hiccup rather than moving it, which is where ADR 0005 arrives from the opposite direction. The benchmark record says this instead of quietly reporting the flattering percentile.fix(test): mlxcel-core release test binary intermittently hangs forever, blocking later runs via the cargo lock #1008 (closed): concurrent
mlxcel-corerelease test binaries contend for Metal command queues. 7 of 12 concurrent suites aborted and the survivors reported 3 to 4 false failures.--test-threadscaps do not help, because the contention is between processes. Now fails loudly by name viagpu_exclusivity_tests.rs(test(core): fail loudly when another mlxcel-core test binary shares the GPU #1009).fix(speculative): draft KV cache rewind trims two entries too many, starving the drafter #994 (closed): draft KV cache rewind trimmed two entries too many.
fix(tests): fused MoE GeGLU parity threshold is tighter than the gather_qmm reference's own jitter #964 (open): the fused MoE GeGLU parity threshold is tighter than the kernel can satisfy. Pre-existing, unrelated to this epic's changes, surfaced by running the suite repeatedly.
Durable changes beyond the features
chore(release): prepare v0.5.0-beta.1