perf(qwen4): overlap streamed PLE decode without changing sampled AR (stacked on #475)

[davidtai]

This AR change is stacked on MTPLX #475, whose tested parent head is cd6023ecf5b78bfd75733fae0b0ad8004161e20b. This upstream PR targets main, so GitHub also displays the parent changes while #475 remains open. Review the focused AR child PR or its AR-only compare for the actual new diff. After #475 merges, this PR can be rebased on upstream main to show only the AR change.

The problem

The existing sampled-AR pipeline overlaps host graph construction by making the next token part of the lazy MLX graph. That removes a host bubble, but it also replaces the request's established sampler:

• normal AR calls _sample_from_logits, applies MTPLX's existing filtering order, and draws from the request's numpy.random.Generator;
• the pipelined path called _mx_lazy_sample and drew with mx.random.categorical.

The same request and seed therefore produced a different continuation. A regression test catches the divergence at token 3 under temperature 1 / top-p 0.95 sampling. This is a correctness failure even though both streams are individually deterministic.

The pipeline also required a resident 32 GB n-gram table. On the 128 GB Qwen3.8 serving configuration, the PLE table is streamed from its SafeTensors sidecar, so constructing a forward on an unknown token would otherwise force a host synchronization.

The resulting schedule

This PR keeps the original sampler and overlaps only work that is independent of its result:

GPU:   forward(token k) ------------------------- logits k
Host:              build forward(deferred token k+1)
CPU:                                             sample logits k with the original RNG
CPU:                                             fill token + exact packed PLE rows
GPU:                                             submit prebuilt forward(k+1)

For every decode step:

1. Allocate a fresh MLX-owned int64[1] token leaf initialized to zero.
2. Build the next forward graph against that leaf while the current forward is executing.
3. Materialize the current logits through _sample_from_logits with the same NumPy generator used by ordinary AR.
4. Fill the token leaf with that sampled token.
5. Compute the 16 exact PLE row IDs from the concrete token and prior history, copy the existing packed Q4 weights/scales/biases into fresh MLX-owned leaves, and advance the PLE history.
6. Submit the already-built consumer graph.

The model still dequantizes the same packed sidecar bytes through the ordinary MLX operations. There is no alternate arithmetic, token prediction, speculative rollback, or cache repair.

Construction-time contract

The optimized lane resolves its invariants once when it is installed or activated:

• the native extension ABI must expose both deferred-row and deferred-token factories;
• all PLE routes are selected before model state is changed;
• token creation and each layer's PLE flush/discard callable are prebound;
• enabled decode calls those bound functions directly, without per-token metadata walks, environment reads, eligibility checks, proof counters, or stock fallback branches;
• a sidecar read/fill error poisons the route, preserves the original exception, and prevents the next activation until the model is reloaded.

The existing request-level MTPLX_AR_PIPELINE eligibility gate remains unchanged. If the model cannot install the complete route, it declines before entering the pipeline and uses the explicit classic AR route.

The public stats envelope now includes ar_pipeline_active and ar_pipeline_variant, allowing a benchmark to prove that variant=streamed actually ran.

Correctness gates

The focused test battery covers:

• a consumer graph built before token/PLE fill reading the exact filled bytes at evaluation;
• independent fresh native leaves across steps;
• streamed PLE history and exact packed-plane publication;
• activation-bound routing and failure poisoning;
• same-seed equality between classic and pipelined temperature-1 AR;
• public engagement telemetry.

Results:

• 24 CPU protocol/static tests passed;
• 16 guarded MLX/native tests passed after rebuilding the extension from the final commit;
• the production service was restored to the same model after every guarded run.

The full-model parity gate used the unchanged #475 head as control and final commit 775697bb40a8a19c8ecd06cea73b32cd3ae5a8a0 as candidate:

Setting	Value
Model	Qwen3.8 Flash-Next MTPLX Optimized-Speed
Input	16,384 tokens, cold cache bypass
Output	128 tokens
Generation	target AR; MTPLX_AR_PIPELINE=0 on #475 control, =1 on candidate
Sampling	temperature 1, top-p 0.95, top-k 20
Reasoning	xhigh
Seeds	20260829, 20260830, 20260831
Thermal	max fans; each sample started at or below 40 C

All three pairs matched on request-body hash, prompt hash, completion count, finish reason, visible-text hash, and reasoning-text hash. Every candidate response reported ar_pipeline_active=true and ar_pipeline_variant=streamed.

Seed	#475 control tok/s	Candidate tok/s	Output parity	Streamed route
20260829	48.81	55.78	yes	yes
20260830	49.83	56.65	yes	yes
20260831	48.97	55.81	yes	yes
median	48.97	55.81	3/3	3/3

That is a 13.97% median decode-throughput gain with the same 85.08 GB peak-memory reading.

Sustained A/B/A result

The final performance confirmation used the same artifact, prompt, sampler, reasoning mode, seeds, cache policy, fan policy, and 40 C start gate, with a 1,024-token output cap. It ran control A, candidate B, then control A2. Seeds 20260829 and 20260830 reached 1,024 tokens; seed 20260831 emitted the same stop at token 666 in all three arms.

Arm	Median decode tok/s	Median prefill tok/s	Median TTFT	Median wall time	Peak memory
#475 control A	46.28	1,630.21	10.267 s	32.371 s	85.080 GB
Candidate B	52.46	1,630.99	10.282 s	29.717 s	85.080 GB
#475 control A2	46.18	1,631.52	10.270 s	32.411 s	85.080 GB

Across the six surrounding control samples, the candidate improves median decode throughput by 13.47% and reduces median end-to-end wall time by 8.26%. Prefill, TTFT, and peak memory are unchanged within run noise. All A/B/A output hashes match for every seed, and all candidate records prove the streamed route was active.

Scope

This changes only the Qwen4/Qwen3.8 sampled target-AR lane layered on #475's native cached-PLE extension. MTP verification, draft sampling, greedy AR, constrained decoding, penalty-bearing sampling, prefill arithmetic, and default-off behavior are unchanged.

[davidtai]

It should be noted that there are several fixes that need to go in and this is just one.