Softmax-free GPU sampling: Gumbel-max categorical sampling kernel

[inureyes]

Background

mlxcel samples on GPU through MLX graph ops. The fused sampler fused_sample (src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp, around lines 4525-4597) runs: greedy argmax when temperature==0 || top_k==1; otherwise temperature scale, top-k via argpartition, top-p filter, min-p compiled kernel, and finally mlx::core::random::categorical. The Rust orchestration is sample_token_optimized_core (src/lib/mlxcel-core/src/sampling.rs, around line 297), which first applies logit-level pre-steps (token bias, repetition/DRY/frequency/presence penalties, XTC) and then calls the fused sampler. Batched sampling loops per row (batched_sample, sampling.rs around line 414).

Two cost sources are structural: categorical internally needs a normalization pass over the vocabulary, and argpartition is a partition-sort-class op over the vocabulary (32K-152K entries), both per token per row.

The Gumbel-max trick removes both. Add i.i.d. Gumbel noise g = -log(-log(u)) to logits / temperature and take the argmax; the result is an exact sample from the softmax categorical distribution. No softmax, no sort, one index-carrying max-reduction pass. Use a counter-based (Philox-style) RNG so each (row, element) gets an independent stream derived from a seed and offset, which keeps the kernel stateless and reproducible.

This issue implements Gumbel-max for the no-filter sampling path. It deliberately comes before the rejection-based top-k/top-p work (#901) so the kernel plumbing (JIT strings for both backends, per-row RNG, batch-wide launch, parity harness) is validated on the simplest case first.

Task

1. Kernel (Metal + CUDA JIT strings, following the paged_attention.cpp dual-source pattern; a new file src/lib/mlx-cpp/turbo/sampling.cpp is a reasonable home): input logits [B, V] (f32 or bf16/f16 upcast), temperature, per-row RNG (seed, offset); compute logits[i]/temp + gumbel(philox(seed, row, i, offset)) and reduce to the argmax index per row. Batch-wide grid: all rows in one launch (grid.y = B), removing the per-row Rust loop for this path.
2. RNG semantics: derive per-row Philox seeds deterministically from mlxcel's existing seed state so that a fixed seed reproduces the same token stream across runs on the same backend. Document that sampled streams will differ from the current categorical-based streams at equal seeds (RNG change), and add a CHANGELOG note.
3. Wiring: in fused_sample, use the Gumbel-max kernel when no top-k/top-p/min-p/XTC filter is active and temperature > 0. Greedy stays on argmax. All logit-level pre-steps (penalties, bias, XTC) remain untouched upstream. Env kill switch MLXCEL_SAMPLING_GUMBEL=0 restores categorical.
4. Batched path: extend batched_sample to call the batch-wide kernel once when every row qualifies (uniform no-filter config), falling back to the per-row loop otherwise.

Performance validation (mandatory)

• Add a sampling microbenchmark (under benchmarks/ or as an example binary) measuring the sampling step alone: vocab {32K, 64K, 152K} x batch {1, 4, 8}, comparing baseline (categorical path) vs Gumbel-max, on Metal and CUDA.
• Required outcomes: sampling-step time reduced at every point in the matrix (expected >= 2x at batch >= 4); end-to-end decode tok/s with temperature sampling on one small model (e.g. Qwen3.5-0.8B-4bit) shows no regression at batch 1 and measurable improvement at batch 4.
• Record in docs/benchmark_results/gumbel-sampling-<hw>-<date>.md.

Regression guard (mandatory)

• Statistical correctness: on synthetic logits (several shapes: peaked, flat, bimodal, with -inf masked entries), draw >= 1e6 samples and compare empirical frequencies against exact softmax probabilities with a chi-square test; document the acceptance threshold in the test. Must also verify temperature scaling behaves identically to the reference (same test at temp 0.5 and 1.5).
• Masked tokens (logit == -inf from bias/XTC pre-steps) must never be sampled (explicit test).
• Determinism: same seed, same backend => same token stream (test).
• Greedy (temperature==0) token streams byte-identical to baseline.
• Full sampling unit tests and cargo test -p mlxcel-core pass; kill switch restores baseline behavior; per-request sampling options (top-k/top-p/min-p/XTC set) still route through the existing path unchanged.

References

• src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp (fused_sample)
• src/lib/mlxcel-core/src/sampling.rs
• src/lib/mlx-cpp/turbo/paged_attention.cpp (dual JIT-string pattern to imitate)

Line numbers are indicative as of current main; search for the named symbols.

---

Part of epic #909.