Epic: Fused paged decode, sorting-free sampling, and serving performance techniques

[inureyes]

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.