fix(server): row-level per-request sampling in the continuous batch (supersedes #1353)

[mmmugh]

Supersedes #1353. That PR added top_k/min_p to the single shared batch sampler; @lucasnewman closed it because a continuous batch needs row-level sampling when requests have different generation parameters. This does that: every row in the batch is sampled with its own temperature / top_p / top_k / min_p / seed.

The framing

mlx-vlm's GenerationBatch was forked from mlx-lm's, and mlx-lm already carries per-row samplers (self.samplers: List[Callable], threaded through extend/filter). The fork collapsed that list into one shared sampler built from the first request — which is the bug. On main today, GenerationBatch.extend()/filter() never touch self.sampler, so a merging batch's sampler is dropped ("request #1 wins"); _make_sampler drops top_k/min_p; and greedy_sampling is a batch-wide bool from request #1. This PR restores per-row sampling.

Approach

• A frozen SamplingConfig(temperature, top_p, top_k, min_p, seed) carried as a per-row List[SamplingConfig] through insert → extend → filter, mirroring the existing logits_processors list exactly.
• A row-aware _PositionedTargetSampler(configs) that builds [B] mx.arrays and applies temperature → top_k → top_p → min_p in one sorted-space pass, with fused greedy where(temp < eps, argmax, sampled). top_k is a rank-mask (exactly-k), because apply_top_k's argpartition(kth=int) can't take a per-row k. The vendored sample_utils.py is untouched.
• The two duplicated _PositionedTargetSampler copies (which had drifted on seed handling) are unified.
• Per-request seeds stay stateless: _position_seed(seed, position) → mx.random.key, so a seeded request is reproducible and invariant to batch composition (vLLM/SGLang semantics: same seed = same stream). Per-row differentiation comes from each row's seed.
• Out-of-domain top_k/min_p are rejected at admission (422) via Pydantic field constraints.

The proof

The load-bearing test is a strengthened test_positioned_target_sampler_is_batch_grouping_invariant: two rows with distinct row_ids and distinct temperature/top_p/top_k/min_p/seed, asserting the batched draw equals each row drawn alone. That property is what #1353 lacked, and it fails on main today. Plus regression tests for the batch lifecycle (extend/filter alignment, greedy reconciliation, first-token per-row), MTP staggered completion, and speculative no-drop. 499 tests green.

Live verification

Verified end-to-end against a real mlx_vlm.server (Qwen3-VL-8B-Instruct-8bit, MLX 0.32.0, Apple Silicon), on the continuous-batching path.

Check	Result	Detail
Greedy (temp=0) deterministic solo	✅	identical across runs
top_k=1 @ temp=2.0 tracks greedy (top_k binds)	✅	shares a 305-char prefix with the greedy answer
Unconstrained temp=2.0 diverges from greedy (incoherent)	✅	shares only 17 chars with greedy
Stochastic temp=1.3, top_k=60, min_p=0.02 genuinely samples	✅	≠ greedy
Seeded stochastic config reproducible (same seed twice, solo)	✅	identical
Co-batched top_k=1 row stays greedy-coherent every time	✅	prefix-match with greedy across 5 co-batched pairs: [120,120,120,120,120]
Co-batched temp=2.0 row stays incoherent every time	✅	prefix-match with greedy across the same pairs: [17,17,17,17,17]
The two co-batched rows never collapse to one mode	✅	greedy row ≠ hot row in every pair

The discriminator (the bar #1353 could not pass): a top_k=1 (greedy) request and an unconstrained temperature=2.0 (incoherent) request were co-batched. The greedy row tracked the greedy answer every time while the hot row stayed wild every time, in the same batches — they retained distinct per-config behavior. Under the old shared-sampler behavior the whole batch uses one parameter set, so the greedy row would have gone wild whenever the hot request opened the batch; it never did.

Why behavioral, not byte-identity: the batched model forward on this GPU is floating-point non-deterministic across batch composition — two identical greedy requests, co-batched, share a long common prefix then diverge (reduction order depends on batch position). That is a pre-existing model property on the untouched greedy path, unrelated to sampling. So byte-exact per-row isolation is proven by the deterministic unit test test_positioned_target_sampler_is_batch_grouping_invariant (identical logprobs in → identical token out); the live test proves param adherence, which is robust to that FP noise.

Scope

• This PR: plain AR + MTP — full per-row. MTP rides the same BatchGenerator; its walks call sample_target on the full batch, so per-row composes (one gate keeps positioned-sampler MTP batches at full width, since _mtp_rounds_batch otherwise compacts mid-round).
• dflash/eagle3 (_run_speculative): interim explicit degradation, full per-row in a follow-up PR. A heterogeneous batch is serialized-by-config (never the silent last-wins), with a TODO(PR2) marker. Full per-row there is tractable (no accept-math change — verification is exact-match-to-target-sample, not q/p rejection) but has its own sampler_is_greedy surface + test matrix, so it's a separate reviewable PR.

An open question for you

The plumbing (per-row config through insert/extend/filter) is the same whether the sampler applies params vectorized (this PR) or as mlx-lm's per-row loop-of-closures. I went vectorized because it preserves the sample_target grouping-invariance that mlx-vlm's positioned sampler adds over mlx-lm's global-RNG closures (and speculative/mtp.py duck-types on sample_target). Measured ~1.6–3.1× faster than the loop at B=4–16, though that's a couple percent of a decode step at small B — happy to switch to the loop shape if you'd prefer parity with mlx-lm; it drops into the same seam.

Known limitations (follow-ups)

• The model-dependent top_k >= vocab_size check raises HTTPException(422) but some pre-existing broad except Exception handlers in openai.py/anthropic.py downgrade it to 400/500. No crash either way (the sampler is robust to out-of-domain params); it's a status-code nicety I left rather than reworking unrelated error handling.
• _run_speculative's serialize-by-config can defer a lone off-config request under sustained majority-config load (interim path; resolved by the PR-2 per-row work).

Credits

Restores the per-row samplers pattern from mlx-lm. Vectorized-sampler shape (rank-mask top_k, fused greedy, per-row [B] params) follows vLLM, SGLang, and TGI. The vendored-in-tree sample_utils/BatchKVCache this builds on came from @lucasnewman's #1494.

[Lazarus-931]

hey @mmmugh , the CI fails b/c of pre-commit, run pre-commit run --all-files and you should be good to go

[lucasnewman]

Is this just a drive-by fix? I don't object to it, but it seems only tangentially related.

[mmmugh]

You're right—removed it. It was tangential, and batched_row_sample already degrades an
out-of-range top_k safely (it just acts as "disabled"), so it wasn't guarding against a crash
anyway. Kept the top_k/min_p/temperature Pydantic field constraints, since those validate the
new per-row params directly. (Pushed in the latest commit.)

[lucasnewman]

This is a non-trivial change from a performance perspective, since we won't be dropping finished sequences. Can you benchmark batched inference before and after?

[mmmugh]

Good call—I reworked it so we do keep dropping finished sequences. Instead of disabling
mid-round compaction for positioned samplers, the sampler now compacts its per-row configs in
lockstep with the batch via _PositionedTargetSampler.select(keep_slots), so active_idx shrinks
as before and the configs stay aligned to the shrunk logprobs. The regression test now asserts
compaction happens (cache.filter_calls == [[1]]) rather than being suppressed, and reverting the
one sampler.select(...) line reproduces the width-mismatch it guards against.

Benchmarked it end-to-end. I split Qwen/Qwen3.5-4B into a native qwen3_5_mtp drafter
(python -m mlx_vlm.speculative.drafters.qwen3_5_mtp.split) and served with
--draft-kind mtp --draft-block-size 4 on an M-series Mac Studio. Workload: 8 concurrent chat
requests with max_tokens spread 48→256 so rows finish at staggered rounds and mid-round compaction
fires repeatedly, temperature=0.7 so the batch runs the vectorized per-row sampler. I instrumented
_mtp_rounds_batch to confirm the change is actually exercised—it runs at B=8 with the positioned
sampler, and compaction fires on every drop (n_active 8→7→…→1, select() each time). 3-trial
median aggregate decode throughput (per-run variance <1%):

	shared sampling config	distinct per-row config
main (c9e27b0)	135.6 tok/s	134.6 tok/s
this branch	135.2 tok/s	119.1 tok/s

So the compaction rework itself is free: with a shared config the branch matches main (135.2 vs
135.6, within noise) while dropping finished rows exactly like main (verified firing 8→1).

Being upfront about the other cell—when every row carries a different sampling config, per-row
MTP decode runs ~12% slower than the uniform batch (119 vs 135). It's the same _mtp_rounds_batch
path (not the serialize_speculative_batch_by_config gate, which is dflash/eagle3-only): the cost of
actually honoring distinct seeds/params across a speculative batch, which main doesn't do at all.
It only appears when configs genuinely differ; shared-config batches are unaffected. I can chase that
down in the per-row-speculative follow-up if it's worth it. (Batched AR per-row overhead, for
reference, was ~1%: branch 272 vs main 275 tok/s.)

[lucasnewman]

@mmmugh Thanks, see the comments for some questions. Please make sure to run pre-commit run --all to clear the formatter so we can run tests.

[mmmugh]

Thanks @Lazarus-931—good catch. Ran pre-commit run --all (black/isort/autoflake) and pushed the formatting fix in 21d3002; it's on the current head. 🙏

[lucasnewman]

@mmmugh Can you sign your commits and force push so we can clear the merge requirements?

[mmmugh]

Done—signed all 8 commits and force-pushed. The branch head is now a343cd6 and GitHub shows every commit as Verified, so that should clear the signed-commits requirement. 🙏

[mmmugh]

@lucasnewman before I finish the merge, I'd like your steer on one integration point, since main has moved a lot in the sampler area since I opened this (the top-nσ / p-less / typical-p modes, plus _PositionedTargetSampler itself).

This PR generalizes _PositionedTargetSampler to carry a per-row SamplingConfig, which meant replacing the batch generator's per-request sampler= slot with per-row sampling_configs threaded through insert/extend/filter. Your new modes route through that same sampler= slot (_make_sampler(args)) and aren't expressible in the current SamplingConfig (temperature/top_p/top_k/min_p/seed)—so they collide in server/generation.py. (The other two conflicts are trivial: schema field constraints and a co-located test.)

How would you like the two to coexist?

Extend SamplingConfig + the per-row sampler to honor top_n_sigma/p_less/typical_p per-row (architecturally consistent, a bit more surface).
Keep the new modes on a batch-level sampler path alongside per-row (applied batch-global, not per-row).
Happy to implement whichever fits where you're taking this.

[lucasnewman]

@lucasnewman before I finish the merge, I'd like your steer on one integration point, since main has moved a lot in the sampler area since I opened this (the top-nσ / p-less / typical-p modes, plus _PositionedTargetSampler itself).

This PR generalizes _PositionedTargetSampler to carry a per-row SamplingConfig, which meant replacing the batch generator's per-request sampler= slot with per-row sampling_configs threaded through insert/extend/filter. Your new modes route through that same sampler= slot (_make_sampler(args)) and aren't expressible in the current SamplingConfig (temperature/top_p/top_k/min_p/seed)—so they collide in server/generation.py. (The other two conflicts are trivial: schema field constraints and a co-located test.)

How would you like the two to coexist?

Extend SamplingConfig + the per-row sampler to honor top_n_sigma/p_less/typical_p per-row (architecturally consistent, a bit more surface). Keep the new modes on a batch-level sampler path alongside per-row (applied batch-global, not per-row). Happy to implement whichever fits where you're taking this.

@Lazarus-931 Can you take a look?

[Lazarus-931]

hi @mmmugh, yes great point, I think we'll merge this first with extended SamplingConfig, we'll edit those prs need be. Feel free to make changes here however for that!

[mmmugh]

Merged current main; the branch is no longer conflicting. Per @Lazarus-931's call above, SamplingConfig now carries top_n_sigma / p_less / typical_p, honored per row: _new_modes_keep() builds the per-row keep mask with each mode gated per row, so a row that doesn't use one is untouched. One open question for you at the end.

Conflicts: schemas.py keeps this PR's field constraints plus main's three new fields; both adjacent test_speculative.py tests survive; _make_sampler is deleted (per-row admission leaves it with no callers, and make_sampler is no longer imported, so keeping it would NameError). A second merge picks up Mage-Flow and the black 26.3.1 bump, so the formatter check here runs the same version CI will.

Two bugs a pre-push audit turned up, both in main's own features, fixed and regression-tested:

• APC. _apc_pick_for destructures the queued-sequence tuple as 6 elements; insert() now appends a 7th (the config). Separate hunks, so git merged both sides cleanly—a semantic-only conflict that would have raised ValueError on every APC-enabled prefill. The new test feeds real insert() output to the real consumer so the two can't drift; one upstream test that hand-built a 6-tuple is updated to match.
• float16. top_n_sigma took max/std in the raw dtype, and mx.std over an fp16 logprob vector overflows to inf at any real vocab size, silently no-opping the mode. Upstream's _top_n_sigma casts to fp32 unconditionally and the deleted _make_sampler routed through it. Now fp32: 0 mismatches vs apply_top_n_sigma at V=8192/32000 in fp16, bf16 and fp32.

Four smaller ones in the new code: upstream's "drop if x < t" phrasing, so a single -inf logprob (grammar masking, logit_bias) can't NaN the threshold and pin a row to argmax; modes chained ahead of the rank filters per make_sampler's order, or typical_p can empty a Qwen-default top_k=20 window and collapse the row to greedy; the extra mask work skipped unless a row opts in; typical_p range-gated, and a non-bool p_less coerced rather than raising in the GPU thread and tearing down the batch.

Your call. Combined with top_p I renormalize over the survivors instead of reproducing make_sampler's chain. apply_top_p compares cumulative mass against 1 - top_p on a distribution that no longer sums to 1 once a mode has masked it, so it can reject every token: top_n_sigma=1.0 with top_p=0.9 leaves a row of all -inf upstream that then gets sampled from. Looks like a genuine upstream bug worth its own look, but happy to match upstream exactly here instead. With no mode active, behaviour is unchanged and still matches apply_top_p / apply_min_p exactly.

Masks are checked token-for-token against sample_utils.apply_*, plus per-row activation, grouping invariance, fp16 parity, NaN tolerance, typical_p+top_k and server wiring; every fix is mutation-tested (8/8 fail when reverted). 1826 passed, 4 skipped, plus one unrelated pre-existing failure (test_qwen_target_verify_quantized_linear_matches_singleton_batch_path, an exact bf16 quantized-matmul equality check) that fails identically on pristine origin/main in a clean worktree, so not caused by this PR. pre-commit clean under 26.3.1.

@lucasnewman your approve predates the merge and these fixes. The composition order is the part worth re-reading.

[mmmugh]

Rebased onto current main—back to mergeable, commits signed, suite green.

Whenever one of you has a moment: CI has never run on this branch (fork PRs need a maintainer to approve the workflow), so it's sitting at "blocked" on a check that hasn't been triggered rather than on anything failing. A run would tell us if it's actually good on your infra.

Still open from the last update, whenever you get to it:

• top_p combined with one of the new modes: I renormalize over the survivors rather than reproducing make_sampler's chain, because that chain can reject every token once a mode has masked mass away (top_n_sigma=1.0 with top_p=0.9 leaves an all -inf row). Happy to match upstream exactly instead if you'd rather keep them identical.
• @lucasnewman your approve predates the merge and the fixes that came with it, so it'll need a fresh look.

No rush—just flagging that it's ready and only needs a CI trigger to move.