[perf] Sample next tokens from temperature-scaled logits without materializing full-vocab probs

[positive666]

Motivation

On the default CUDA (flashinfer) sampling path, every non-greedy decode step runs an explicit full-vocab softmax plus an in-place copy-back (logits[:] = torch.softmax(logits, dim=-1)), then samples from the materialized probs. The filter-free (temperature-only) case additionally uses torch.multinomial, whose cost explodes with batch size.

The fused flashinfer samplers can consume temperature-scaled logits directly, and the repo already ships an alignment test (python/sglang/kernels/aot/tests/test_sampling.py) showing top_k_top_p_sampling_from_logits draws RNG-for-RNG identical tokens to softmax + top_k_top_p_sampling_from_probs (both joint and top_k_first orders). This is the output-side sibling of the input-logprob fast path from #31958.

Modifications

• Sampler._should_sample_from_scaled_logits: the fast path engages only when nothing downstream needs the probs tensor — it is skipped for probs-derived logprobs (return_logprob), RL sampling-mask capture, per-request sampling_seed, deterministic inference mode, min_p, and non-flashinfer backends. All those paths are byte-for-byte unchanged.
• Sampler._sample_from_scaled_logits: top_k_top_p_sampling_from_logits(..., filter_apply_order="joint") for the top-k/top-p case, sampling_from_logits for the filter-free case.
• environ.py: SGLANG_DISABLE_SAMPLING_FROM_LOGITS (EnvBool(False)) escape hatch, placed next to its input-side sibling SGLANG_ENABLE_FAST_INPUT_LOGPROBS.
• New unit test test/registered/unit/sampling/test_sampler_from_logits.py (gate truth table, bitwise parity vs the probs path under a shared RNG state, distribution parity vs torch.multinomial).

Note: #35966 restructures the same region of Sampler.forward for sampling-mask capture. The two changes are logically independent (the mask path stays on the probs pipeline and is excluded by this gate); whichever lands second needs a trivial rebase.

Accuracy Tests

• Unit tests 5/5 on L40S, including sampler-level bitwise parity: with the same RNG state, the fused top-k/top-p path returns tokens identical to softmax + top_k_top_p_sampling_from_probs (bs=64, vocab 32k).
• Filter-free path: empirical distribution over 200k draws, L1(new, true)=0.0072 vs L1(torch.multinomial, true)=0.0085 (vocab 64) — at least as faithful as the current path.
• E2E A/B (env flag on/off, same build): greedy outputs byte-identical; return_logprob and seeded requests fall back to the existing pipeline and return unchanged results.

Speed Tests and Profiling

Kernel-level (L40S, vocab=151936, fp32 logits; old = softmax + copy-back + sample-from-probs, per step):

bs	config	old (us)	new (us)	saved
1–32	temperature-only	~168	~61	~107 us (2.8x)
128	temperature-only	988	61	927 us (16.3x)
128	top_p=0.95, top_k=ALL	577	382	195 us
128	top_p=0.95, top_k=50	675	485	190 us
32	top_p=0.95, top_k=ALL	108	101	7 us

E2E (qwen3-0.6b, 1x L40S, bench_serving random 256/256, concurrency 128, ABAB design across 4 server restarts):

sampling	baseline median TPOT	fast median TPOT	delta
temperature=1.0	20.52 / 20.94 ms	19.89 / 20.22 / 20.22 ms	~3–4% faster, output throughput +3.5–6%
+ top_p=0.95	20.30–20.78 ms (4 reps)	20.14–20.97 ms (4 reps)	within restart noise (kernel saving ~1% of step)

The temperature-only saving is a fixed per-step cost (vocab- and batch-dependent, model-size independent), so the relative gain is largest for small/fast models and large running batches.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #32657339125
Latest PR Test (Extra): ❌ Run #32657339077
Latest PR Test (AMD ROCm 7.2): ❌ Run #32657339116