Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
positive666 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ede6d5d
Fuse target temperature in rejection sampler
positive666 Aug 20, 2026
bceeaff
Warm fused rejection sampler kernels
positive666 Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
188 changes: 188 additions & 0 deletions benchmarks/kernels/benchmark_rejection_sampler_temperature.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,188 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from functools import partial

import torch

from vllm.triton_utils import triton
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.v1.worker.gpu.sample.gumbel import apply_temperature
from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import (
rejection_sample,
)


def make_inputs(
batch_size: int,
num_speculative_steps: int,
vocab_size: int,
temperature: float,
dtype: torch.dtype,
with_draft_logits: bool,
) -> dict[str, torch.Tensor]:
device = "cuda"
num_logits = batch_size * (num_speculative_steps + 1)
target_logits = torch.randn(num_logits, vocab_size, dtype=dtype, device=device)
draft_sampled = torch.randint(
vocab_size,
(batch_size, num_speculative_steps + 1),
dtype=torch.int64,
device=device,
)
draft_sampled[:, 0] = 0
idx_mapping = torch.arange(batch_size, dtype=torch.int32, device=device)
expanded_idx_mapping = idx_mapping.repeat_interleave(num_speculative_steps + 1)
expanded_local_pos = torch.arange(
num_speculative_steps + 1, dtype=torch.int32, device=device
).repeat(batch_size)

inputs = {
"target_logits": target_logits,
"draft_sampled": draft_sampled.flatten(),
"cu_num_logits": torch.arange(batch_size + 1, dtype=torch.int32, device=device)
* (num_speculative_steps + 1),
"pos": torch.arange(num_logits, dtype=torch.int32, device=device),
"idx_mapping": idx_mapping,
"expanded_idx_mapping": expanded_idx_mapping,
"expanded_local_pos": expanded_local_pos,
"temperature": torch.full(
(batch_size,), temperature, dtype=torch.float32, device=device
),
"seed": torch.arange(batch_size, dtype=torch.int64, device=device),
}
if with_draft_logits:
inputs["draft_logits"] = torch.randn(
batch_size,
num_speculative_steps,
vocab_size,
dtype=dtype,
device=device,
)
return inputs


def run_baseline(
inputs: dict[str, torch.Tensor], num_speculative_steps: int
) -> tuple[torch.Tensor, torch.Tensor]:
processed_logits = torch.empty_like(
inputs["target_logits"], dtype=torch.float32
).copy_(inputs["target_logits"])
apply_temperature(
processed_logits,
inputs["expanded_idx_mapping"],
inputs["temperature"],
)
return rejection_sample(
target_logits=processed_logits,
draft_logits=inputs.get("draft_logits"),
num_speculative_steps=num_speculative_steps,
**{
key: value
for key, value in inputs.items()
if key not in ("target_logits", "draft_logits")
},
)


def run_fused(
inputs: dict[str, torch.Tensor], num_speculative_steps: int
) -> tuple[torch.Tensor, torch.Tensor]:
return rejection_sample(
target_logits=inputs["target_logits"],
draft_logits=inputs.get("draft_logits"),
num_speculative_steps=num_speculative_steps,
apply_target_temperature=True,
**{
key: value
for key, value in inputs.items()
if key not in ("target_logits", "draft_logits")
},
)


def assert_outputs_equal(
baseline: tuple[torch.Tensor, torch.Tensor],
fused: tuple[torch.Tensor, torch.Tensor],
num_speculative_steps: int,
) -> None:
baseline_sampled, baseline_num_sampled = baseline
fused_sampled, fused_num_sampled = fused
torch.testing.assert_close(fused_num_sampled, baseline_num_sampled, rtol=0, atol=0)
steps = torch.arange(
num_speculative_steps + 1, device=baseline_sampled.device
).unsqueeze(0)
valid = steps < baseline_num_sampled.unsqueeze(1)
torch.testing.assert_close(
fused_sampled[valid], baseline_sampled[valid], rtol=0, atol=0
)


def measure_peak_memory(callable_) -> int:
torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
allocated = torch.cuda.memory_allocated()
output = callable_()
torch.cuda.synchronize()
peak = torch.cuda.max_memory_allocated() - allocated
del output
return peak


def main(args) -> None:
dtype = getattr(torch, args.dtype)
print(
"batch rows baseline_ms fused_ms speedup "
"baseline_peak_MiB fused_peak_MiB saved_MiB"
)
for batch_size in args.batch_sizes:
inputs = make_inputs(
batch_size,
args.num_speculative_steps,
args.vocab_size,
args.temperature,
dtype,
args.with_draft_logits,
)
baseline_call = partial(run_baseline, inputs, args.num_speculative_steps)
fused_call = partial(run_fused, inputs, args.num_speculative_steps)

baseline_output = baseline_call()
fused_output = fused_call()
torch.cuda.synchronize()
assert_outputs_equal(baseline_output, fused_output, args.num_speculative_steps)
del baseline_output, fused_output

baseline_ms = triton.testing.do_bench(
baseline_call, warmup=args.warmup_ms, rep=args.rep_ms
)
fused_ms = triton.testing.do_bench(
fused_call, warmup=args.warmup_ms, rep=args.rep_ms
)
baseline_peak = measure_peak_memory(baseline_call)
fused_peak = measure_peak_memory(fused_call)
mib = 1024**2
print(
f"{batch_size:5d} {batch_size * (args.num_speculative_steps + 1):4d} "
f"{baseline_ms:11.3f} {fused_ms:8.3f} "
f"{baseline_ms / fused_ms:7.3f}x "
f"{baseline_peak / mib:17.1f} {fused_peak / mib:14.1f} "
f"{(baseline_peak - fused_peak) / mib:9.1f}"
)


if __name__ == "__main__":
parser = FlexibleArgumentParser(
description="Benchmark in-kernel target temperature for rejection sampling."
)
parser.add_argument("--batch-sizes", type=int, nargs="+", default=[1, 8, 32, 128])
parser.add_argument("--num-speculative-steps", type=int, default=4)
parser.add_argument("--vocab-size", type=int, default=151936)
parser.add_argument("--temperature", type=float, default=0.6)
parser.add_argument(
"--dtype", choices=["float32", "float16", "bfloat16"], default="bfloat16"
)
parser.add_argument("--with-draft-logits", action="store_true")
parser.add_argument("--warmup-ms", type=int, default=100)
parser.add_argument("--rep-ms", type=int, default=500)
main(parser.parse_args())
160 changes: 149 additions & 11 deletions tests/v1/spec_decode/test_rejection_sampler_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,17 +139,22 @@ def _assert_distribution_match(


@pytest.mark.parametrize(
"num_speculative_steps,temperature",
"num_speculative_steps,temperature,apply_target_temperature",
[
(1, 0.6),
(3, 0.6),
(1, 1.0),
(3, 1.0),
(1, 0.6, False),
(3, 0.6, False),
(1, 1.0, False),
(3, 1.0, False),
(1, 0.6, True),
(3, 0.6, True),
],
)
@pytest.mark.parametrize("draft_logits_dtype", [torch.float32, torch.bfloat16])
def test_stochastic_rejection_sample(
num_speculative_steps: int, temperature: float, draft_logits_dtype: torch.dtype
num_speculative_steps: int,
temperature: float,
apply_target_temperature: bool,
draft_logits_dtype: torch.dtype,
):
"""
Verify that rejection sampling produces the target distribution.
Expand All @@ -167,13 +172,14 @@ def test_stochastic_rejection_sample(
device = "cuda"
num_trials = 10 * VOCAB_SIZE

target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
raw_target_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32)
target_logits_1d = raw_target_logits_1d
draft_logits_1d = torch.randn(VOCAB_SIZE, device=device, dtype=torch.float32).to(
draft_logits_dtype
)

if temperature > 0:
target_logits_1d /= temperature
if temperature > 0 and not apply_target_temperature:
target_logits_1d = target_logits_1d / temperature

inputs = _build_rejection_sample_inputs(
target_logits_1d,
Expand All @@ -184,17 +190,149 @@ def test_stochastic_rejection_sample(
)

sampled, num_sampled = rejection_sample(
**inputs, num_speculative_steps=num_speculative_steps
**inputs,
num_speculative_steps=num_speculative_steps,
apply_target_temperature=apply_target_temperature,
)

target_probs = torch.softmax(target_logits_1d, dim=0)
processed_target_logits_1d = raw_target_logits_1d
if temperature > 0:
processed_target_logits_1d = processed_target_logits_1d / temperature
target_probs = torch.softmax(processed_target_logits_1d, dim=0)
for pos in range(num_speculative_steps + 1):
accepted_mask = num_sampled >= pos + 1
_assert_distribution_match(
sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"
)


@pytest.mark.parametrize(
("target_dtype", "vocab_size", "num_trials"),
[
(torch.float32, VOCAB_SIZE, 64),
(torch.float16, VOCAB_SIZE, 64),
(torch.bfloat16, VOCAB_SIZE, 64),
(torch.bfloat16, 8193, 4),
],
)
@pytest.mark.parametrize("has_draft_logits", [True, False])
@pytest.mark.parametrize("verification_mode", ["standard", "block", "synthetic"])
def test_in_kernel_target_temperature_matches_preprocessing(
target_dtype: torch.dtype,
vocab_size: int,
num_trials: int,
has_draft_logits: bool,
verification_mode: str,
):
torch.manual_seed(42)
device = "cuda"
num_speculative_steps = 3
temperature = 0.6

raw_target_logits_1d = torch.randn(
vocab_size, device=device, dtype=torch.float32
).to(target_dtype)
processed_target_logits_1d = raw_target_logits_1d.float() / temperature
draft_logits_1d = torch.randn(vocab_size, device=device, dtype=torch.bfloat16)
processed_inputs = _build_rejection_sample_inputs(
processed_target_logits_1d,
draft_logits_1d,
num_speculative_steps,
temperature=temperature,
num_trials=num_trials,
)
raw_inputs = dict(processed_inputs)
raw_inputs["target_logits"] = (
raw_target_logits_1d.unsqueeze(0)
.expand(num_trials * (num_speculative_steps + 1), -1)
.contiguous()
)
if not has_draft_logits:
processed_inputs["draft_logits"] = None
raw_inputs["draft_logits"] = None

verification_kwargs = {}
if verification_mode == "block":
verification_kwargs["use_block_verification"] = True
elif verification_mode == "synthetic":
verification_kwargs["synthetic_conditional_rates"] = torch.full(
(num_speculative_steps,), 0.5, dtype=torch.float32, device=device
)

expected = rejection_sample(
**processed_inputs,
num_speculative_steps=num_speculative_steps,
**verification_kwargs,
)
actual = rejection_sample(
**raw_inputs,
num_speculative_steps=num_speculative_steps,
apply_target_temperature=True,
**verification_kwargs,
)

assert torch.equal(actual[1], expected[1])
steps = torch.arange(num_speculative_steps + 1, device=device).unsqueeze(0)
valid = steps < expected[1].unsqueeze(1)
assert torch.equal(actual[0][valid], expected[0][valid])


def test_in_kernel_target_temperature_supports_mixed_requests():
torch.manual_seed(42)
device = "cuda"
num_trials = 6
num_speculative_steps = 2
num_logits = num_trials * (num_speculative_steps + 1)
temperatures = torch.tensor(
[0.0, 0.6, 1.0, 0.8, 0.0, 1.0], dtype=torch.float32, device=device
)
idx_mapping = torch.arange(num_trials, dtype=torch.int32, device=device)
expanded_idx_mapping = idx_mapping.repeat_interleave(num_speculative_steps + 1)
expanded_local_pos = torch.arange(
num_speculative_steps + 1, dtype=torch.int32, device=device
).repeat(num_trials)
row_temperatures = temperatures[expanded_idx_mapping.long()]

raw_target_logits = torch.randn(
num_logits, VOCAB_SIZE, dtype=torch.bfloat16, device=device
)
processed_target_logits = raw_target_logits.float()
scaled_rows = (row_temperatures != 0.0) & (row_temperatures != 1.0)
processed_target_logits[scaled_rows] /= row_temperatures[scaled_rows, None]
draft_sampled = torch.randint(
VOCAB_SIZE, (num_logits,), dtype=torch.int64, device=device
)
cu_num_logits = torch.arange(num_trials + 1, dtype=torch.int32, device=device) * (
num_speculative_steps + 1
)
pos = torch.arange(num_logits, dtype=torch.int32, device=device)
seed = torch.arange(num_trials, dtype=torch.int64, device=device)
common_inputs = dict(
draft_logits=None,
draft_sampled=draft_sampled,
cu_num_logits=cu_num_logits,
pos=pos,
idx_mapping=idx_mapping,
expanded_idx_mapping=expanded_idx_mapping,
expanded_local_pos=expanded_local_pos,
temperature=temperatures,
seed=seed,
num_speculative_steps=num_speculative_steps,
)

expected = rejection_sample(target_logits=processed_target_logits, **common_inputs)
actual = rejection_sample(
target_logits=raw_target_logits,
apply_target_temperature=True,
**common_inputs,
)

assert torch.equal(actual[1], expected[1])
steps = torch.arange(num_speculative_steps + 1, device=device).unsqueeze(0)
valid = steps < expected[1].unsqueeze(1)
assert torch.equal(actual[0][valid], expected[0][valid])


@pytest.mark.parametrize("num_speculative_steps", [1, 3])
def test_greedy_rejection_sample(num_speculative_steps: int):
"""
Expand Down
Loading
Loading