Fix XTC sampling threshold to be per-row for batched logits

[sohumt123]

Problem

apply_xtc in mlx_lm/sample_utils.py computes its masking threshold with a global .min() over all axes:

mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()

Every other filter in the sampler chain (apply_top_p, apply_top_k, apply_min_p) reduces along axis=-1, so samplers produced by make_sampler are expected to vectorize over a leading batch axis — and they are called that way on the full [B, vocab] tensor via batch_generate (GenerationBatch._step -> self.fallback_sampler(logprobs)).

With B >= 2, the smallest above-threshold probability of any row in the batch becomes the mask threshold for every row. Example with xtc_threshold=0.2:

• Row A: probs [0.6, 0.25, 0.1, 0.05] — above-threshold min is 0.25, so 0.6 is correctly masked.
• Row B: probs [0.9, 0.05, 0.03, 0.02] — only 0.9 is above threshold, so XTC should mask nothing.

Batched together, row A's 0.25 becomes the global threshold and row B's top token (p=0.9) is masked to -inf, forcing row B to sample from its near-zero tail. Row B processed alone is unaffected — so batching silently changes (and corrupts) sampling results.

Fix

Reduce along the last axis with keepdims so the threshold is computed per distribution:

mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min(
    axis=-1, keepdims=True
)

This matches the axis convention used by the other filters and is a no-op for the single-row case.

Testing

• Added a regression test in tests/test_sample_utils.py asserting that each row of a batched apply_xtc call equals the same row processed alone (fails on main, passes with this change).
• python -m unittest tests.test_sample_utils -v — 7 tests, all pass.
• black --check clean.

Note: PR #1540 excludes XTC from its vectorized batching because the single scalar mx.random.uniform(0, 1) gate is shared batch-wide; that gate is orthogonal to this fix and left untouched here, but the threshold bug above affects the existing batch_generate fallback path today regardless.

[michalk8]

LGTM, thanks a lot!