[inureyes]

Summary

Fixes the "last-seed-wins" determinism bug in batched prefill so a seeded request's first token no longer depends on the seeds of the other rows it happens to be batched with. The change is a one-line reseed of the global MLX RNG at the exact first-token sample point, plus a real-model regression test.

Root cause

seed_rng_if_needed seeds the process-global MLX RNG (ffi::random_seed) and is called from begin_prefill. In run_padded_batched_prefill, every row in a cohort runs begin_prefill up front (one loop) before any row runs finish_prefill (a later per-row loop that samples the first token). The first-token sampler is fused_sample, which takes only (logits, temperature, top_k, top_p, min_p) and draws from the global RNG with no per-call key. So by the time row 0 samples, the global RNG holds the LAST cohort row's seed, and every seeded row's first token is governed by whichever seed was applied last ("last-seed-wins"). PR #346 (#332 cohort splitting) widened the exposure: a cold row that previously fell back to sequential prefill now runs batched inside a BatchedCold cohort and newly inherits this behavior.

Fix

Insert seed_rng_if_needed(&seq.sampling) in finish_prefill, immediately before the first-token sample_token_optimized call, so the global RNG is reseeded to THIS row's own seed at the exact point it samples. begin_prefill's existing seeding is intentionally left in place (it is harmless once overwritten here, and is still reached from the disaggregated handoff path). finish_prefill is the single per-row first-token sampling chokepoint for every prefill path (batched cohort, sequential, full, chunked-final), so this one insertion fixes both the pre-existing all-cold case and the #332 cohort-split case holistically rather than patching each path.

Greedy is a no-op

Greedy / temperature == 0 / top_k == 1 rows take the fused_sample argmax path, which consumes no RNG, so the reseed has no effect on them. The regression test pins this.

Scope: decode is out of scope (honest AC #1 re-framing)

This fix makes a seeded row's FIRST token independent of its siblings' seeds, and greedy output is fully identical regardless of batch composition. It does NOT make a seeded request's full multi-token output byte-identical between "run alone" (B=1) and "batched" (B>1). Two pre-existing, already-documented effects remain and are not this seed bug: (1) padded batched prefill is not bitwise-identical to single-sequence prefill on Metal (the #203 / #325 / #326 forward-pass jitter class), and (2) the batched fused DECODE path (execute_batched_decode -> batched_fused_sample) samples all rows in ONE fused_sample draw over [B, vocab] from the shared global RNG, so a seeded row's decode trajectory depends on batch composition (the documented batched-vs-B=1 jitter class noted in sampling.rs). Full per-sequence multi-token determinism across batch sizes would require threading a per-row RNG key into fused_sample (an FFI / MLX-side change), which is disproportionate for this priority:low bug and is left as a follow-up. AC #1 is therefore satisfied at the first-token / batched-prefill scope of this bug; the residual is the existing jitter class, not last-seed-wins.

Regression test (invariant choice)

src/server/batch/scheduler_seed_determinism_tests.rs (real qwen3-0.6b-4bit, #[ignore], soft-skips when the model dir is absent), mirroring the scheduler_cohort_parity_tests.rs harness. It deliberately does NOT compare a batched (B>1) row against a single-prefill (B=1) row, which would conflate the forward-pass jitter class. Instead it holds the cohort COMPOSITION constant (same prompts in the same positions, so the measured row's logits are byte-identical) and varies ONLY seed values, isolating the RNG variable. The measured row is always at index 0 (sampled first) and the seeded sibling at index 1 (seeded last), which is exactly where "last-seed-wins" would leak. Checks:

• Sibling-seed invariance (load-bearing, bug-discriminating): the seeded stochastic row's first token is identical across distinct sibling seeds. This fails under "last-seed-wins" and holds exactly under the fix.
• Own-seed dependence (positive control): the same row's first token takes more than one distinct value across its own seeds, proving the sampler genuinely consumes RNG here (so the invariance check is not vacuous) and that the reseed reaches the row.
• Greedy invariance (AC test: v0.0.27 live verification (post-OSS-release) #2): a temperature == 0 row's first token is independent of any sibling's seed.
• Cohort-split coverage (AC chore: complete OSS release rollout #1 cohort-split clause / AC chore: add OSS security automation and community health files #3): a seeded row's first token in a perf(server): split batched prefill into compatible cohorts instead of all-or-nothing fallback #332 cohort-split mixed window ([cold A, cold B, adopted C], which the planner splits into BatchedCold {A, B} + Sequential {C}) matches the same row in an all-cold batched window of the same cold composition ([cold A, cold B]).

Test plan

• cargo check --lib --tests --features metal,accelerate
• cargo clippy --lib --tests --features metal,accelerate -- -D warnings
• cargo test -p mlxcel --lib --release --features metal,accelerate batched_prefill_seed_determinism -- --ignored --nocapture (observed: sibling seeds [1,2,3] -> first token [4792, 4792, 4792]; own seeds [10..60] -> 4 distinct first tokens; greedy [334, 334, 334]; cohort-split all-cold == mixed == 334)
• cargo fmt --check

Closes #347

[inureyes]

Implementation Review Summary

Intent

Reseed the global MLX RNG to each row's own seed at the batched-prefill first-token sample point, eliminating "last-seed-wins" for seeded sampling in a batched cohort (#347).

Findings

No CRITICAL, HIGH, or MEDIUM correctness findings. The one-line reseed and its regression test are correct as written; no changes were made by this review.

Verification

• Reseed placement correct. seed_rng_if_needed(&seq.sampling) (scheduler.rs:3924) sits after the structured-output mask (3885-3909) and immediately before sample_token_optimized (3925). Nothing between consumes RNG: sample_token_optimized_core applies only deterministic token-bias and penalty graph ops before fused_sample, and the lone eval in the token-bias branch is an argmax, not a draw.
• finish_prefill is the single first-token chokepoint for every classic prefill path: batched cohort (3418), full (3577), chunked first-chunk-terminal (3731), chunked continuation-final (3868). Sequential-cohort rows reuse the full/chunked finish sites. The speculative MTP/DFlash burst is a separate, pre-existing pipeline that bypasses finish_prefill, but it is structurally immune to last-seed-wins: a batched burst window only admits seed-identical rows (sampling_config_eq compares seed) and a B=1 burst is single-seed, so there is no distinct "last seed" to leak.
• Regression test discriminates the bug without the f16 batched-vs-B=1 jitter confound. It holds cohort composition constant and varies only sibling seeds, measuring row 0 (sampled first) with the seeded sibling at row 1 (seeded last), which is exactly where last-seed-wins leaks. The positive control (own-seed variation yielding two or more distinct tokens) keeps the invariant non-vacuous. Two cold rows of different lengths do batch into one BatchedCold cohort on qwen3 (can_pad == true), so the test genuinely exercises the begin-all-then-finish-all ordering.
• Greedy unaffected. seed_rng_if_needed is gated on seed.is_some() and the argmax path consumes no RNG; pinned by the greedy-invariance assertion.
• AC chore: complete OSS release rollout #1 reframing is honest. First-token / batched-prefill determinism is achieved. The residual run-alone (B=1) vs batched (B>1) multi-token difference is the pre-existing forward-pass jitter class plus the batched-decode shared-RNG draw, both correctly identified and deferred to the FFI key change.
• cargo check --lib --tests --features metal,accelerate passes. No placeholder code, conventions followed (test mirrors scheduler_cohort_parity_tests.rs), Closes #347 present.

Verdict: correct and complete for its stated scope.