fix(server): batched-prefill seed determinism (last-seed-wins for seeded sampling in a batched cohort) #347

Problem / Background

In the batched-prefill path, a per-request RNG seed is applied to the global MLX RNG, and the first-token sampler then consumes that global RNG with no per-call seed. Because a batched cohort runs every row's begin_prefill (seeding) before any row's finish_prefill (sampling), the last seeded row's seed governs first-token sampling for the entire cohort ("last-seed-wins").

Verified call chain (this repo):

• seed_rng_if_needed (src/lib/mlxcel-core/src/generation_policy.rs:26-30) seeds the global MLX RNG via ffi::random_seed(seed), and only when sampling.seed is Some.
• begin_prefill (src/server/batch/scheduler.rs:385-390) calls seed_rng_if_needed(&seq.sampling) at line 388.
• finish_prefill (src/server/batch/scheduler.rs:3765) samples the first token via sample_token_optimized(...) at line 3802.
• sample_token_optimized -> ffi::fused_sample (src/lib/mlxcel-core/src/sampling.rs:359-365) takes only (logits, temperature, top_k, top_p, min_p). There is no per-call seed argument; it draws from the global RNG.
• Ordering: run_padded_batched_prefill (src/server/batch/scheduler.rs:3163) calls begin_prefill for all rows up front (lines 3182-3186), runs one batched forward pass, then loops per row calling finish_prefill (around line 3309). So every row's seed is applied before the first sample is drawn, and the last row's seed is the live global RNG state when row 0 is sampled.

Pre-existing, widened by #332

This is a pre-existing property of the all-cold batched-prefill path, which already ordered all begin_prefill before all finish_prefill. PR #346 (#332 cohort splitting) widens exposure: a cold row that previously fell back to sequential prefill (because its collected window held an incompatible sibling) now runs batched inside a BatchedCold cohort, so it newly inherits last-seed-wins. The in-code comment at src/server/batch/scheduler.rs:3117-3128 documents exactly this "cold row now runs batched" behavior change.

Impact (scoped honestly)

• Output tokens remain valid; there is no corruption.
• Greedy / temperature == 0 requests are unaffected: the argmax path consumes no RNG (the FusedSampleParams field docs in sampling.rs note that temperature == 0.0 and top_k == 1 select the greedy argmax path).
• Affected only when a request supplies an explicit seed and temperature > 0 and is batched with other seeded rows.
• Net: a per-request seed-reproducibility gap, not a crash or corruption. Hence priority:low.

Suggested fix direction (not to implement here)

• Per-row reseed immediately before that row's sampling in finish_prefill; or
• A per-sequence RNG stream (instead of the shared global RNG) for seeded rows; or
• (cheap stopgap) exclude seeded non-greedy rows from batched cohorts in the prefill classifier (the PrefillRow classification in execute_batched_prefill / plan_prefill_cohorts).

The fix must be holistic: it has to cover the pre-existing all-cold case, not only the #332 mixed-window case.

Acceptance Criteria

• A seeded request (seed set, temperature > 0) produces identical output whether run alone, in an all-cold batched cohort, or in a cohort-split mixed window.
• Greedy / temperature == 0 requests are unchanged.
• A regression test pins per-request seed reproducibility under batched prefill, covering both the all-cold and the cohort-split cases.

Technical Considerations

• The seed is applied to a process-global RNG (ffi::random_seed), so any solution that keeps the global RNG must guarantee the seed is live at the exact point each row is sampled (reseed inside finish_prefill, not in begin_prefill).
• Watch the interaction with the batched fused decode fast path (config_supports_fused_batch / batched_fused_sample): seeded non-greedy rows already need care there, and a per-sequence RNG stream would have to extend to decode steps for full end-to-end reproducibility, not just the first token.

Related: PR #346, issue #332.