[Sampler] Add opt-in vocab-sharded sampling (USE_VOCAB_SHARDED_SAMPLING)

[gutianyu-google]

Summary

sample() all-gathers the (batch, vocab) logits over the tensor-parallel axis on every decode step ("Unshard the logits explicitly…", with a TODO asking whether the vocab dimension could stay sharded) and then runs top-k / top-p / categorical on the replicated array, on every device. For a 151,936-token vocab that is a ~78 MB f32 gather per DP rank per step, followed by the 32–33 vocab-wide reduction sweeps of topk_mask and topp_mask, all redundant across the TP devices. In an xprof profile of a Qwen3-0.6B GRPO rollout (DP16 × TP8, TPU v7x) this was ~30% of the decode step's device time.

This PR adds USE_VOCAB_SHARDED_SAMPLING (default off). Each vocab shard keeps its top 64 candidates and only those (values and global ids, ~0.4 MB) are gathered; top-k, top-p and the categorical draw run on the merged candidates.

Exactness. The kept set is the same as the replicated path's:

• the union of the per-shard top-64 contains the global top-k for k <= 64;
• ties at the k-th value are kept, like topk_mask;
• the top-p cutoff is the smallest prefix of the sorted probabilities with mass >= p, like topp_mask;
• candidates outside the top-k set are masked, so the softmax normalizes over exactly the same set (up to float32 summation order);
• greedy rows take the argmax with the lowest id among ties, like jnp.argmax.

The sampled tokens follow the same distribution as before but are not bitwise identical, since the random draw is over the candidates rather than the full vocab. Different batch shards fold their shard index into the key so they draw independent noise.

Fallback. Rows are exact when greedy, or when 0 < top_k <= 64 and top_p > 0. A batch containing any other sampling row takes the replicated path through a lax.cond; both branches leave the returned logits vocab-sharded, so a runtime fallback never adds the gather to the hot path. Padded request slots carry DEFAULT_SAMPLING_PARAMS (temperature −1, top_k 0), i.e. greedy, so they never force the fallback (in an earlier version of this change a guard on top_k > 0 alone was silently tripped by the padded rows on every step).

Constraints. The returned logits are the raw, still-sharded logits, so TpuPlatform.check_and_update_config rejects the env var together with logprobs_mode="processed_*". Sampled-token logprobs (compute_and_gather_logprobs) still run on the replicated logits; sharding that path is a natural follow-up.

Measured on TPU v7x GRPO rollouts (2048 sequences/step, generation cap 8192, top_k 50, top_p 1.0): −5.8 s/step at DP16 × TP8 (123.2 s → 118.0 s) and −9.5 s/step at DP32 × TP4 (99.7 s → 90.2 s), with reward and completion-length distributions unchanged across 20-step runs.

Tests

tests/layers/jax/sample/test_sampling.py::TestVocabShardedSampling (mesh shards the vocab over every local device; run with --xla_force_host_platform_device_count=4 to exercise the multi-shard merge on CPU):

• _can_sample_vocab_sharded: padded slots and greedy rows never force the fallback; a real row without a top-k filter or above the candidate budget does.
• greedy rows match jnp.argmax, including the lowest-id tie rule;
• sampled tokens always lie in the top-k / top-p set of _apply_sampling_transforms;
• over 8192 identical rows the empirical distribution matches the replicated path's probabilities within binomial noise, with zero mass outside the kept set;
• a batch with a top_k=2000 row produces exactly the replicated path's tokens (same key);
• env off keeps the replicated path.

Also: tests/test_envs.py (default off), tests/platforms/test_tpu_platform.py (processed_* logprobs modes rejected, raw_logprobs accepted).