Vectorize batched sampling for rows that share a sampler

[pierre427]

Vectorize batched sampling for rows that share a sampler

Problem

GenerationBatch._step has two sampling paths: a vectorized one that
runs the sampler once over the whole [batch, vocab] logprob tensor,
and a per-row fallback that loops in Python, calling the sampler B
separate times on [1, vocab] slices and concatenating the results.
The vectorized path only triggers when no row carries a custom
sampler — but the server attaches a freshly-constructed sampler closure
to every request (server.py builds one per request via
_make_sampler). So in real server operation the vectorized path is
unreachable: every decode step, for every generated token, pays B
Python-level sampler dispatches plus a concatenate. At
--decode-concurrency 32 that is 32 dispatches per token, all doing
identical math for the overwhelmingly common case of requests that
share sampling settings.

Change

Two small pieces that compose:

1. A batch_groupable contract on samplers. make_sampler's
   outputs now carry an attribute declaring themselves row-independent:
   applying them to k stacked rows at once is semantically identical
   to k separate calls. Temperature scaling, top-p, min-p, and top-k
   all have this property (per-row transforms over the vocab axis).
   XTC does not — it draws one scalar random gate, which would
   become correlated across requests if grouped — so XTC-bearing
   samplers are never marked. Arbitrary third-party callables are also
   unmarked and keep the exact historical one-call-per-row contract
   (a stateful callable, or one drawing shared randomness, cannot be
   silently regrouped); the feature is invisible to them.

2. Identity grouping in the decode step. _step groups rows by
   sampler object; each marked group runs one vectorized call over
   its rows' logprobs. To make groups actually form in the server,
   _make_sampler is memoized on the sampling parameters plus the
   actual special-token id values (small clear-on-overflow cache), so
   concurrent requests with identical settings share one sampler
   object and collapse into a single [B, vocab] call. Mixed-config
   batches degrade gracefully to one call per unique config — never
   worse than the previous per-row loop.

No public API changes. Samplers produced by make_sampler already
vectorize over a leading batch axis (the fallback path has always
required this of the default sampler).

Correctness

• test_batch_shared_sampler_matches_fallback: rows sharing one
  make_sampler object produce bit-identical tokens to a
  fallback-path (sampler=) run under a fixed seed — the grouped
  single-group call consumes RNG identically to the fallback path.
• test_batch_shared_stateful_sampler_keeps_per_row_calls: an unmarked
  stateful callable is still called exactly once per row, with per-row
  widths asserted.
• test_xtc_sampler_not_batch_groupable: XTC samplers are excluded
  from grouping (their scalar random gate is not row-independent).
• test_batch_shared_nonvectorizing_sampler: a shared constant closure
  keeps per-row semantics.
• All existing sampler/logits-processor/stop-matcher batch tests pass
  unchanged. (test_batch_matches_single, test_batch_sliding_window,
  test_many_batches, and the test_batch_continued_generation* trio
  fail identically on pristine main in my environment — pre-existing
  numeric baselines, unrelated to this change.)

Measured

Apple M5 Max (128 GB), mlx 0.32.0.dev; a harness driving
BatchGenerator directly; counterbalanced 4-rep medians; generation
t/s from the engine's own BatchStats (prompt time excluded by the
scheduler itself, not by wall-clock subtraction). "per-row" = distinct
sampler objects per request (current server behavior); "grouped" = one
shared object (server behavior with the memoization).

Baseline control (pristine main): a sampler object shared across rows
gains nothing without the grouping engine — memoized/per-row =
1.001x (B=8), 0.997x (B=32) — confirming any delta below is the engine
change alone.

Qwen3-0.6B-4bit (150k vocab), 128-token prompts, 64 generated:

B	per-row t/s	grouped t/s	speedup	engine-step delta
8	1746	1893	1.084x	-0.35 ms
32	2156	2475	1.148x	-1.75 ms

The grouped path matches the vectorized fallback path within noise at
every batch size — it recovers all recoverable overhead. The saved cost
is a roughly fixed 1-2 ms of Python dispatch + concatenate per decode
step, so the relative win shrinks as per-step forward time grows:

• 4096-token prompts (same model): 1.038x / 1.039x at B=8/32 — the
  dilution is mechanical (attention time grows with context; the saved
  dispatch does not).
• Qwen3-Coder-30B-A3B-8bit: within-tree grouping ratios 0.993-1.021x —
  below this session's measurement floor. No 30B gain is claimed;
  the win is concentrated where batching is used most aggressively
  (small/mid models at high concurrency), which is stated as the
  honest scope.

It is a strict improvement in all measured configurations: the grouped
path never does more work than the per-row loop.

Notes / possible follow-ups

• Per-row heterogeneous params still take one call per unique config.
  A batched sampler taking per-row parameter arrays (temp: [B], ...)
  would collapse those too; left out to keep this change small — happy
  to explore if there's interest. The same follow-up is the natural
  place to thread per-row RNG keys, which would let seeded requests
  (currently excluded from batching by args.seed) batch
  deterministically.
• Logits processors have the same per-row loop shape; not touched here.

🤖 Generated with Claude Code

[pierre427]

Pushed 56c033f: fixed a latent crash in PromptProcessingBatch.extend. It planted [None] logits_processors placeholders while filter() uses [[]], so a mixed batch ([None, [proc]]) crashes _step at the for processor in self.logits_processors[e] line with TypeError: 'NoneType' object is not iterable.

This is the same None-vs-[] invariant that #1513 just fixed for the server loop. Used the list-comprehension form [[] for _ in …] rather than [[]] * n so the placeholders don't alias one shared list. Left the sampler None placeholders intact — those legitimately fall back and are never iterated.

[pierre427]

While hardening the batched-sampler path I hit a second lane-alignment issue in the same family as the extend [] invariant this PR already fixes, but in filter.

GenerationBatch.filter gates the re-index on any(self.samplers) / any(self.logits_processors). When a per-lane list is non-empty but every entry is falsy (all-None samplers, or all-[] processors — reachable via the batch API, or when the all-None placeholder from extend reaches a GenerationBatch), any(...) is False, the list is not re-indexed, and it stays longer than uids. The next extend() then appends at the wrong offset and silently binds lanes to the wrong request — greedy / wrong-sampler output with no error.

Gating on truthiness of the list itself re-indexes whenever per-lane info exists while still leaving a genuinely empty [] untouched. Proposed diff (happy to push it onto this branch):

-        if any(self.samplers):
+        # Re-index whenever a per-lane list exists (len == old uids), even if
+        # every entry is falsy (all-None samplers / all-[] processors). any()
+        # skips the all-falsy case, leaving the list longer than uids so a later
+        # extend() binds lanes to the wrong request (silent greedy/wrong-sampler).
+        if self.samplers:
             self.samplers = [self.samplers[idx] for idx in keep]
-        if any(self.logits_processors):
+        if self.logits_processors:
             self.logits_processors = [self.logits_processors[idx] for idx in keep]

Plus the sibling aliasing bug in PromptProcessingBatch.filter's else-branch ([[]] * n is n references to one list object, so an in-place append on one lane mutates every lane):

-            self.logits_processors = [[]] * len(keep)
+            self.logits_processors = [[] for _ in keep]

Repro: build a GenerationBatch of B≥2 with all-None per-lane samplers; filter(keep) a strict subset (any(samplers) is False → not re-indexed); extend() a request carrying a real sampler (appended at the wrong offset); _step then reads a shifted lane and silently samples with the wrong sampler. No exception is raised — the output is just wrong.

#1513 already covers the extend None→[] crash, so there's no overlap with that PR.

[pierre427]

Pushed the two filter fixes to this branch in 94c841e — GenerationBatch.filter now gates on if self.samplers: / if self.logits_processors: (re-indexes even when every lane entry is falsy), and PromptProcessingBatch.filter's else-branch uses [[] for _ in keep] instead of the aliasing [[]] * n.

[pierre427]

Revalidated 2026-07-22 on the current tree: grouped/memoized vs per-row 1.05 / 1.07 / 1.07 / 1.10 / 1.12x at B=2/4/8/16/32 (Qwen3-0.6B-4bit, 256-token prompts, 128 generated, 4-rep medians) — consistent with the original table; the state-budget admission layer downstream of this was also re-exercised under a deliberately tight budget (flood 8 -> 4 active + 4 queued, clean second-wave drain, no over-admission).

[pierre427]

Rebased onto current main (254d153) and added one commit.

The rebase picks up #1176/#1575/#1372 (XTC special tokens, per-row thresholds, the 0.1 default), which this PR's sampler memoization now composes with: _uncached_make_sampler carries #1176's xtc_special_tokens computation, and the new commit fixes the memo key to the same plural tokenizer.eos_token_ids contract — it previously read the singular attribute, which the tokenizer contract doesn't guarantee (the TestMakeSampler fake added in #1176 exercises exactly that, and now passes against the memoized path).

Re-verified locally on mlx-community/Qwen1.5-0.5B-Chat-4bit: the five batching tests this PR adds, plus tests/test_server.py (34) and tests/test_sample_utils.py, all green; black clean.

🤖 Generated with Claude Code