perf: keep homogeneous batches on the vectorized sampling path

[Vissirexa]

What

Every request was inserted into BatchGenerator with its own per-row sampler, so GenerationBatch._step always took the any(self.samplers) branch: per-row logprobs slice + sampler call + mx.concatenate on every decode step — even when all rows share identical sampling params (the overwhelmingly common local case). With omlx's mx.compile-free sampler that is ~8–10 extra tiny Metal kernel launches per row per token.

_create_batch_generator now records the sampling-param key of the fallback sampler it installs, and both insert sites pass a None row sampler when the request's params match. All-None rows keep mlx-lm's vectorized fallback_sampler(logprobs) path — one sampling op for the whole batch.

Safety

• Rows with differing params keep their own sampler exactly as before (heterogeneous batches unchanged).
• mlx-lm's GenerationBatch.filter skips reindexing an all-None samplers list (any() is False) — the sampler sibling of the known stale-filter alignment issue (fix: normalise logits_processors row slots dropped to None by batch merge jundot/omlx#1799/thinking_budget is silently ignored on /v1/completions jundot/omlx#1825). The existing per-step uid-registry realign already rebuilds sampler rows positionally; a new test pins the remove-then-extend case.
• The VLM-MTP router still receives the real sampler object (it reads .temp/.top_p for draft acceptance).

Bench

gemma-4-e4b-it-qat-OptiQ-4bit, temp=0.7 top_p=0.9, 256 decode tokens, median of 3 (old = per-row samplers forced via monkeypatch in the same process):

batch	per-row (old)	elided (new)	Δ
1	107.7 tok/s	109.9 tok/s	+2.1%
4	315.1 tok/s	322.3 tok/s	+2.3%

Testing

• New TestHomogeneousRowSamplerElision in tests/test_scheduler_logits_processors.py: elide/keep decisions, param-key insensitivity to non-sampler fields, realign-after-remove-then-extend, source-level guards on both insert sites.
• pytest tests/ -m "not slow and not integration" green (except the two pre-existing test_glm_mtp_patch.py tolerance failures, identical on main).

Part of the inference-speed work tracked in docs/perf-analysis/inference-speed-analysis.md (finding F2 stage 1).

🤖 Generated with Claude Code

https://claude.ai/code/session_01VFknQkdZoPFqsH6vtNS11q