[Perf][Spec Decode] Fuse target temperature in rejection sampler

[positive666]

Summary

Fuse target-temperature scaling into the MRV2 rejection-sampling kernels for temperature-only requests. This avoids materializing an FP32 target-logits buffer and removes chunking on this path.

Requests using other logits processors or processed logprobs continue to use the existing fallback. Startup warmup now compiles the fused-temperature variants to avoid first-request JIT compilation.

Results

L40S, BF16 target logits, vocabulary 151936, speculative K=4:

• One-hot draft rejection-sampling microbenchmark: 1.57x-6.89x for batch sizes 1, 8, 32, and 128.
• Full draft logits microbenchmark: 1.49x-4.56x.
• Batch 128 peak temporary memory: 371.4 MiB to 0.5 MiB (370.9 MiB saved).

Two five-round Qwen3.5-9B end-to-end comparisons (ngram K=4, 128 prompts, 64 output tokens) remained within run-to-run noise: -0.66% and +0.99% throughput.

A fresh Triton-cache run confirmed that the fused rejection-sampler kernels do not compile during inference after startup warmup.

Tests

.venv/bin/python -m pytest -q \
  tests/v1/spec_decode/test_rejection_sampler_utils.py \
  tests/v1/worker/test_gpu_rejection_sampler_chunking.py \
  tests/v1/worker/test_gpu_sampler_flags.py \
  tests/v1/worker/test_gpu_rejection_sampler_i64.py \
  tests/v1/worker/test_gpu_warmup_blocks.py
# 125 passed

uv run --no-sync ruff check <9 changed Python files>
uv run --no-sync ruff format --check <9 changed Python files>

Correctness coverage includes FP32/FP16/BF16 target logits, one-hot/full draft logits, standard/block/synthetic verification, mixed temperatures, multi-block vocabularies, and fallback paths.

Benchmark commands:

.venv/bin/python benchmarks/kernels/benchmark_rejection_sampler_temperature.py \
  --batch-sizes 1 8 32 128 --warmup-ms 100 --rep-ms 300
.venv/bin/python benchmarks/kernels/benchmark_rejection_sampler_temperature.py \
  --batch-sizes 1 8 32 128 --warmup-ms 100 --rep-ms 300 --with-draft-logits

Duplicate check

No equivalent open PR was found. #48928 targets FlashInfer AIR top-p sampling, while this fast path excludes top-p. #45369 changes the legacy sampler path.

AI assistance

AI assistance was used for implementation and validation. The submitter reviewed every changed line and ran the tests above.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.