perf(server): batched decode sampler fast path ([B, vocab] -> [B] fused sampling) #325

Summary

Add a [B, vocab] -> [B] fused batched sampler fast path to the batched decode loop so that B active sequences are sampled with a single eval/sync point, instead of B per-row slice + sample + host scalar-extract round trips.

Source: idea.md "1. Batched Decode Sampler Fast Path" (P0, expected payoff medium to high for small/fast models and B>1 serving, risk medium).

Current state

• BatchScheduler::execute_batched_decode() runs a single batched model forward for all active sequences, then slices one row at a time, calls sample_token_optimized(), evaluates the token array, extracts item_i32, and optionally computes per-row logprobs (src/server/batch/scheduler.rs, in the execute_batched_decode body near lines 4137 / 4213 / 4254).
• The standalone batched_sample() helper also loops by row and slices logits before calling the single-row sampler (src/lib/mlxcel-core/src/sampling.rs:309).
• sample_token_optimized() is already fused for a single row (src/lib/mlxcel-core/src/sampling.rs:223), so the missing piece is a true [B, vocab] -> [B] fused sampler.

Proposal

• Add a scheduler fast path that triggers when ALL active rows share a compatible sampling config AND none of them needs a structured-output mask, a thinking-budget override, or a per-token logprobs payload.
• Dispatch a new batched fused sampler over the full [B, vocab] logits and return B token ids after one eval/sync point.
• Keep the existing per-row path as the correctness fallback for structured output, row-specific logprobs, token-bias observability, and mixed sampling configs.

Scope

• New batched fused sampler in src/lib/mlxcel-core/src/sampling.rs (greedy first, then temperature/top-k/top-p in batched form where feasible).
• Fast-path gate plus dispatch in execute_batched_decode().
• Preserve the per-row fallback exactly.

Validation

• Server benchmark with B=2/4/8 on qwen3-0.6b-4bit, qwen3-4b-4bit, and a small Llama/Gemma dense model.
• Compare total decode tok/s, per-token latency, and greedy output parity at temperature=0.
• Include one fallback test with structured output or logprobs enabled to confirm behavior is unchanged.
• Suggested first experiment (idea.md): measure server decode overhead on qwen3-0.6b-4bit at B=1/2/4/8, no logprobs then logprobs enabled, to size the opportunity before changing kernels.

Acceptance criteria

• Batched fused sampler returns [B] token ids with one eval/sync per decode step for the compatible-config case.
• Greedy (temperature=0) output is byte-identical to the current per-row path, or within the documented jitter class.
• Fast path is gated off whenever any row needs structured output, a thinking-budget override, or per-token logprobs; the per-row fallback is unchanged for those.
• INTEGRATION: wired into the live execute_batched_decode decode loop, not an orphaned helper.
• Decode tok/s and per-token latency recorded for B=2/4/8; no regression at B=1.
• Unit/integration coverage for the batched sampler and the fallback gate.

Risk and priority

idea.md priority P0; risk medium. Main upside is on fast small models and B>1 serving, where per-row sync and host scalar extraction become a visible fraction of decode.

References

• idea.md "1. Batched Decode Sampler Fast Path".
• Code: src/server/batch/scheduler.rs (execute_batched_decode), src/lib/mlxcel-core/src/sampling.rs (sample_token_optimized, batched_sample, compute_logprobs).