[Kernel] Fuse temperature + softmax in sampling for decode speedup

[Godmook]

Motivation

The standard sampling path in sampler.py executes temperature scaling and softmax as two separate CUDA kernels:

logits.div_(sampling_info.temperatures)       # Kernel 1
logits[:] = torch.softmax(logits, dim=-1)     # Kernel 2

This is called on every decode step and incurs:

• 2 kernel launches (~10us overhead each)
• 6 global memory passes over the full [batch_size, vocab_size] tensor (4 reads + 2 writes)

For large-vocab models (Llama 3 128K, Qwen 152K), this becomes a meaningful bottleneck in the decode latency budget.

Modifications

New file: python/sglang/srt/layers/fused_sampling.py

Triton kernel that fuses div_(temperature) + softmax into a single kernel with two variants:

• Single-pass kernel (vocab ≤ 32768): Loads entire vocab into registers once. Computes max, exp, sum, normalize all in-register. Only 1 read + 1 write = 2 memory passes.
• Multi-pass kernel (vocab > 32768): 2-pass online softmax with @triton.autotune over (BLOCK_SIZE, num_warps, num_stages). 2 reads + 1 write = 3 memory passes.

Both variants apply logits / temperature during the load, avoiding a separate division pass.

Modified: python/sglang/srt/layers/sampler.py (3 lines changed)

• Import fused_temperature_softmax_inplace
• Replace the div_ + softmax two-line sequence with the fused kernel in the standard sampling path
• Original implementation kept as fallback for non-CUDA platforms

New file: test/registered/sampling/test_fused_temperature_softmax.py

14 correctness tests comparing fused output against the reference (div_ + softmax) across:

• Batch sizes: 1, 2, 16, 64, 128, 512
• Vocab sizes: 1024, 32000, 128256
• Temperatures: 0.01 to 100.0
• Dtypes: bf16, fp16, fp32
• Edge cases: empty batch, temperature=1.0

New file: benchmark/kernels/bench_fused_temperature_softmax.py

Microbenchmark using torch.cuda.Event timing (200 iters, 50 warmup).

Accuracy Tests

All 14 tests pass — fused output matches reference within tolerance (atol=1e-4, rtol=1e-3 for bf16):

test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_basic PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_large_vocab PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_batch_sizes PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_temperature_one PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_very_low_temperature PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_very_high_temperature PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_fp16_input PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_fp32_input PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_mixed_temperatures PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_empty_batch PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_inplace_basic PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_inplace_bf16 PASSED
test_fused_temperature_softmax.py::TestFusedTemperatureSoftmax::test_inplace_large_vocab PASSED

13 passed

Benchmarking and Profiling

Measured on NVIDIA A40 (SM 8.6), bf16, torch.cuda.Event timing, 200 iters after 50 warmup:

bs	vocab	original (us)	fused inplace (us)	speedup
1	32000	25.5	23.4	1.09x
1	128256	43.0	33.8	1.27x
32	32000	28.5	21.3	1.34x
32	128256	150.5	46.7	3.22x
128	32000	140.9	34.7	4.06x
128	128256	612.3	177.1	3.46x
512	32000	541.7	121.4	4.46x
512	128256	2416.8	694.8	3.48x

Speedup across all configurations: 1.09x–4.46x. Improvement comes from:

• Eliminating 1 kernel launch (~10us)
• Reducing memory passes from 6 to 2 (single-pass) or 3 (multi-pass)
• Autotune selecting optimal (BLOCK_SIZE, num_warps, num_stages) per vocab size

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.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[DarkSharpness]

LGTM. I just wonder how long it would take to compile & autotune this triton kernel? Also, the autotune must be done in warm-up stage and should not introduce extra runtime overhead.

[Godmook]

LGTM. I just wonder how long it would take to compile & autotune this triton kernel? Also, the autotune must be done in warm-up stage and should not introduce extra runtime overhead.

Measured on A100 (SM 8.0).

1. Compile + Autotune overhead (first call per vocab_size)

vocab_size	first call
32K	0.003s
128K	1.447s

Autotuning is cached by key=["vocab_size"], so changing batch_size within the same vocab_size does not trigger re-autotuning.

---

2. Runtime latency after warmup (200 iters, 50 warmup, bf16)

bs	vocab	cached (ms)
1	32K	0.051
32	32K	0.050
128	32K	0.050
512	32K	0.112
1	128K	0.063
32	128K	0.064
128	128K	0.118
512	128K	0.462

Once autotuned, runtime overhead is negligible (0.05–0.46ms).

---

3. Autotune overhead per vocab_size x dtype

vocab	dtype	autotune (s)
32K	bf16	0.003
32K	fp16	0.122
32K	fp32	0.120
128K	bf16	1.448
128K	fp16	3.686
128K	fp32	3.667

Each dtype triggers a separate autotune. fp16/fp32 takes ~30x longer than bf16.

---

4. Disk cache hit (first call after runtime restart)

vocab	1st call (s)
32K	0.003
128K	1.455

Even with Triton's ~/.triton/cache present, the same autotune cost is incurred after every process restart. The disk cache reuses compiled PTX/cubin but re-runs the autotune config search on every new process.

[Godmook]

I added autotune.

Changes

1. python/sglang/srt/layers/fused_sampling.py

Added warmup_fused_temperature_softmax() at the end of the file.
Runs both fused_temperature_softmax and fused_temperature_softmax_inplace with a dummy input to trigger JIT compile and autotune before the first real request.

2. python/sglang/srt/model_executor/model_runner.py

Added _warmup_fused_sampling() call inside kernel_warmup(), after _flashinfer_autotune().

---

Boot flow

Server start
  → ModelRunner.__init__()
    → kernel_warmup()
      → _flashinfer_autotune()        # existing
      → _warmup_fused_sampling()      # new
        → warmup_fused_temperature_softmax(vocab_size)
          - vocab ≤ 32K : single-pass JIT only  (~3ms)
          - vocab > 32K : multi-pass autotune   (~1–4s)
    → init_device_graphs()
  → Ready to serve (no latency spike on first request)

---

Autotune is cached by key=["vocab_size"]. Warming up once with the model's vocab_size covers all batch sizes at runtime — no re-autotuning occurs during serving. Runtime overhead after warmup is 0.

Also, I reinforced test codes and refactor some minor changes. @DarkSharpness Could you check it? Thanks!

[DarkSharpness]

/tag-and-rerun-ci

[Godmook]

/rerun-failed-ci

[DarkSharpness]

/rerun-failed-ci

[DarkSharpness]

/rerun-failed-ci

[Godmook]

/rerun-failed-ci

[Godmook]

/rerun-failed-ci

[Godmook]

/rerun-failed-ci

[Godmook]

Sampling Kernel: Correctness Fix + Hybrid Dispatch

Problem

The original fused Triton kernel used a numerically different computation order
from PyTorch, causing structured output (grammar-constrained decoding) to fail.

Three sources of divergence:

	Original kernel	PyTorch path
Temperature scaling	x * (1/temp)	logits.div_(temp)
Softmax accumulation	Online softmax (with correction term)	3-pass: max → sum → normalize
Normalization	exp(x - max - log(sum))	exp(x - max) / sum

---

Fix

Rewrote the kernel to follow the same mathematical operation order as PyTorch
(x / temp → exp(x - max) / sum). Also removed the grammar fallback branch in
sampler.py — the fused kernel now always runs.

The only remaining difference is sub-ULP reduction ordering inside tl.max /
tl.sum vs PyTorch's CUDA kernels, which is extremely unlikely to affect
sampling results in practice.

---

Performance Impact

Aligning with PyTorch's algorithm required moving from 2-pass to 3-pass
reduction (one extra read pass), which regressed small-batch latency slightly.

Before (original kernel, numerically divergent from PyTorch):

bs	vocab	triton_ip (µs)	flashinfer (µs)	fi/tri
1	32K	63.0	61.5	0.98x
1	128K	77.2	62.0	0.80x
32	32K	61.3	64.2	1.05x
32	128K	77.1	86.1	1.12x
128	32K	62.5	74.1	1.19x
128	128K	118.8	308.2	2.59x
512	32K	110.8	259.5	2.34x
512	128K	463.6	1183.2	2.55x

After fix (numerically aligned with PyTorch, grammar decoding works):

bs	vocab	triton_ip (µs)	flashinfer (µs)	fi/tri
1	32K	50.6	43.5	0.70x
1	128K	70.3	44.2	0.56x
32	32K	50.2	44.0	0.72x
32	128K	76.7	81.0	0.95x
128	32K	50.4	71.0	1.15x
128	128K	176.6	308.7	1.74x
512	32K	113.5	259.7	1.49x
512	128K	609.1	1184.2	1.81x

Small-batch latency improved (~63µs → ~50µs at bs=1/32K) but triton_ip now
falls behind flashinfer at small batches. Large-batch advantage also narrowed
(bs=512/128K: 2.55x → 1.81x over flashinfer).

---

Hybrid Dispatch

To recover small-batch performance without sacrificing correctness, added a
dispatch layer:

• bs < 128 → torch.div_ + torch.softmax (faster for small batches,
  mathematically identical)
• bs >= 128 → fused Triton kernel (2–3x faster than flashinfer at large
  batches)

Both paths perform the same operations (x / temp → softmax) so grammar-
constrained decoding works correctly on either path.

Memory access pattern comparison:

Path	Reads	Writes
PyTorch (div_ + softmax)	4	2
Fused Triton (3-pass)	3	1

@DarkSharpness Do you think it is good approach or not?

[DarkSharpness]

/ temp → softmax) so grammar- constrained decoding works correctly on either path.

Memory access pattern comparison:

For now, the baseline is naive torch implementation. As long as your implementation is faster than torch native (div + softmax), it will be fine. We can try flashinfer solution in later PRs as follow-up.

[Godmook]

@DarkSharpness I think XPU test and Xeon test has some problems now. I try to fix XPU CI tests in another PR now. So I'll try to do rerun and if there is no problem for other CIs, I think it is good now. I'll do rerun ci after 2 hours because of time limit...

[BBuf]

Can we move this import to the top of the file?

[Godmook]

fused_sampling pulls in Triton at import time. we keep this import inside if is_cuda(): with the other CUDA-only deps (flashinfer, sgl_kernel) so sampler stays importable on non-CUDA builds.

Thanks for help me :)