iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model Runner V2] Fix -1 placeholder draft token ids in rejection sam… (#50939)

Browse Source 
Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
Giancarlo Delfin
2026-08-06 16:32:00 -07:00
committed by GitHub
co-authored by Nick Hill
parent 9bca7d840d
commit 27930df9c2
2 changed files with 229 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/spec_decode/test_rejection_sampler_utils.py
+122
View File
@@ -367,6 +367,128 @@ def test_placeholder_draft_token_rejected():
assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all()
@pytest.mark.parametrize("has_draft_logits", [True, False])
@pytest.mark.parametrize("num_placeholders", [1, 2])
def test_block_verification_placeholder_truncates_block(
has_draft_logits: bool, num_placeholders: int
):
"""Trailing placeholder (-1) draft ids end the block early. The last real
draft token must then be verified against the closed-form threshold rather
than a residual mass derived from the placeholder position, whose logits
describe a token that was never drafted.
"""
torch.manual_seed(42)
device = "cuda"
num_trials = 20 * VOCAB_SIZE
K = 3
temperature = 1.0
target_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature
draft_logits_1d = torch.randn(VOCAB_SIZE, device=device)
inputs = _build_rejection_sample_inputs(
target_logits_1d,
draft_logits_1d,
K,
temperature=temperature,
num_trials=num_trials,
)
if not has_draft_logits:
inputs["draft_logits"] = None
inputs["draft_sampled"].view(num_trials, K + 1)[:, K + 1 - num_placeholders :] = -1
sampled, num_sampled = rejection_sample(
**inputs, num_speculative_steps=K, use_block_verification=True
)
num_real = K - num_placeholders
assert (num_sampled <= num_real + 1).all(), (
"Accepted a draft token at or after a placeholder."
)
target_probs = torch.softmax(target_logits_1d, dim=0)
for pos in range(num_real + 1):
accepted_mask = num_sampled >= pos + 1
_assert_distribution_match(
sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"
)
def test_greedy_placeholder_emits_target_argmax():
"""Greedy sampling skips resampling and relies on the rejection kernel
storing the target argmax at the rejected position. A placeholder must not
bypass that store, or the output slot is returned uninitialized.
"""
torch.manual_seed(0)
device = "cuda"
num_trials = 512
K = 3
target_logits_1d = torch.randn(VOCAB_SIZE, device=device)
draft_logits_1d = torch.randn(VOCAB_SIZE, device=device)
inputs = _build_rejection_sample_inputs(
target_logits_1d,
draft_logits_1d,
K,
temperature=0.0,
num_trials=num_trials,
)
target_argmax = int(target_logits_1d.argmax())
draft_sampled = inputs["draft_sampled"].view(num_trials, K + 1)
# Accept the first draft so that the rejection lands on the placeholder.
draft_sampled[:, 1] = target_argmax
draft_sampled[:, 2:] = -1
sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)
assert torch.equal(num_sampled, torch.full_like(num_sampled, 2))
steps = torch.arange(K + 1, device=device).unsqueeze(0)
emitted = sampled[steps < num_sampled.unsqueeze(1)]
assert (emitted == target_argmax).all(), (
"Greedy sampling emitted a token that is not the target argmax"
)
@pytest.mark.parametrize("use_block_verification", [False, True])
def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool):
"""A placeholder is not necessarily the final draft. Nothing at or after
one may be accepted, even when a valid draft follows it.
"""
torch.manual_seed(0)
device = "cuda"
num_trials = 4 * VOCAB_SIZE
K = 3
temperature = 1.0
# An identical draft is accepted with probability ~1, so any token after
# the placeholder would be accepted unless it is explicitly blocked. The
# draft logits are stored pre-temperature, so pass the unscaled base.
draft_logits_1d = torch.randn(VOCAB_SIZE, device=device)
target_logits_1d = draft_logits_1d / temperature
inputs = _build_rejection_sample_inputs(
target_logits_1d,
draft_logits_1d,
K,
temperature=temperature,
num_trials=num_trials,
)
inputs["draft_sampled"].view(num_trials, K + 1)[:, 2] = -1
_, num_sampled = rejection_sample(
**inputs,
num_speculative_steps=K,
use_block_verification=use_block_verification,
)
# Only the first draft is verifiable, plus one resampled token.
assert (num_sampled <= 2).all(), "Accepted a draft token past a placeholder."
assert (num_sampled == 2).float().mean().item() > 0.9, (
"The first draft was rarely accepted; the test is not exercising "
"acceptance past the placeholder."
)
@pytest.mark.parametrize(
"num_speculative_steps,temperature",
[
vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py
+107 -74
View File
@@ -55,7 +55,7 @@ def _compute_global_residual_mass(
target_local_max_stride,
target_local_sumexp_ptr,
target_local_sumexp_stride,
draft_sampled_ptr,
draft_token,
logit_idx,
vocab_num_blocks,
PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,
@@ -74,7 +74,6 @@ def _compute_global_residual_mass(
# One-hot draft. M_s is a point mass at this draft token
# so the residual mass reduces to the closed form:
# p * (1 - M_b(draft_token)).
draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
target_lse = _compute_global_logsumexp(
target_local_max_ptr,
target_local_max_stride,
@@ -347,35 +346,40 @@ def _compute_cumulative_log_p_kernel(
if temp == 0.0:
return
log_p = 0.0
log_p = tl.zeros((), tl.float32)
for step in range(num_draft_tokens):
logit_idx = start_idx + step
draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
target_logprob, draft_logprob, _, _ = _compute_global_logprobs_and_logsumexp(
draft_token,
True, # mask
logit_idx,
req_state_idx,
step,
temp,
target_logits_ptr,
target_logits_stride,
target_local_max_ptr,
target_local_max_stride,
target_local_sumexp_ptr,
target_local_sumexp_stride,
draft_logits_ptr,
draft_logits_stride_0,
draft_logits_stride_1,
draft_local_max_ptr,
draft_local_max_stride,
draft_local_sumexp_ptr,
draft_local_sumexp_stride,
vocab_num_blocks,
PADDED_VOCAB_NUM_BLOCKS,
HAS_DRAFT_LOGITS,
)
log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)
# -1 placeholder tokens can never be accepted. Skip their reductions
# and carry the last valid cumulative value.
if draft_token >= 0:
target_logprob, draft_logprob, _, _ = (
_compute_global_logprobs_and_logsumexp(
draft_token,
True, # mask
logit_idx,
req_state_idx,
step,
temp,
target_logits_ptr,
target_logits_stride,
target_local_max_ptr,
target_local_max_stride,
target_local_sumexp_ptr,
target_local_sumexp_stride,
draft_logits_ptr,
draft_logits_stride_0,
draft_logits_stride_1,
draft_local_max_ptr,
draft_local_max_stride,
draft_local_sumexp_ptr,
draft_local_sumexp_stride,
vocab_num_blocks,
PADDED_VOCAB_NUM_BLOCKS,
HAS_DRAFT_LOGITS,
)
)
log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)
tl.store(cumulative_log_p_ptr + logit_idx, log_p)
@@ -406,6 +410,8 @@ def _compute_local_residual_mass_kernel(
draft_local_sumexp_ptr,
draft_local_sumexp_stride,
# [num_logits]
draft_sampled_ptr,
# [num_logits]
expanded_idx_mapping_ptr,
# [num_logits]
expanded_local_pos_ptr,
@@ -425,6 +431,11 @@ def _compute_local_residual_mass_kernel(
# first and last (bonus) positions aren't needed for this computation.
return
if tl.load(draft_sampled_ptr + logit_idx + 1) < 0:
# -1 placeholder token. The rejection kernel treats the preceding token
# as the end of the block, so this position's residual mass is unused.
return
req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64)
temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)
if temp == 0.0:
@@ -539,41 +550,23 @@ def _rejection_kernel(
accepted_length = tl.zeros((), tl.int64)
target_lse = 0.0
draft_lse = 0.0
accepted = True
verifying = True
for i in range(num_draft_tokens):
logit_idx = start_idx + i
draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)
pos = tl.load(pos_ptr + logit_idx)
u = tl_rand32(seed, pos, includes_zero=False)
if USE_BLOCK_VERIFICATION and not is_greedy:
# Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444
prefix_joint_ratio = tl.exp(
tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32)
)
if i < num_draft_tokens - 1:
residual_mass = _compute_global_residual_mass(
local_residual_mass_ptr,
local_residual_mass_stride,
prefix_joint_ratio,
target_logits_ptr,
target_logits_stride,
target_local_max_ptr,
target_local_max_stride,
target_local_sumexp_ptr,
target_local_sumexp_stride,
draft_sampled_ptr,
logit_idx + 1,
vocab_num_blocks,
PADDED_VOCAB_NUM_BLOCKS,
HAS_DRAFT_LOGITS,
)
denom = residual_mass + 1.0 - prefix_joint_ratio
h = tl.where(denom > 0.0, residual_mass / denom, 1.0)
else:
h = prefix_joint_ratio
accepted_length = tl.where(u <= h, i + 1, accepted_length)
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
elif accepted:
# -1 is used for placeholder draft token ids that should be rejected.
is_valid_draft = draft_sampled >= 0
# Avoid possible OOB ptr access.
draft_sampled = tl.maximum(0, draft_sampled)
if not is_greedy:
# A -1 placeholder ends verification. Greedy is excluded because it
# stores the target argmax upon first rejection, so it rejects the
# placeholder via `accepted` instead.
verifying &= is_valid_draft
if verifying:
pos = tl.load(pos_ptr + logit_idx)
u = tl_rand32(seed, pos, includes_zero=False)
if is_greedy:
# Greedy sampling. Accept IFF draft matches target argmax.
# NOTE: Target argmax is stored directly so that resampling
@@ -589,20 +582,51 @@ def _rejection_kernel(
)
if SYNTHETIC_MODE:
rate = tl.load(synthetic_conditional_rates_ptr + i)
# -1 is used for padded draft token ids that should be rejected.
accepted &= (u < rate) & (draft_sampled >= 0)
accepted = u < rate
else:
accepted &= target_argmax == draft_sampled
accepted = target_argmax == draft_sampled
accepted &= is_valid_draft
verifying = accepted
accepted_length += accepted
tl.store(
sampled_ptr + req_idx * sampled_stride + i,
draft_sampled if accepted else target_argmax,
)
elif USE_BLOCK_VERIFICATION:
# Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444
prefix_joint_ratio = tl.exp(
tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32)
)
next_draft_token = tl.load(
draft_sampled_ptr + logit_idx + 2,
mask=i < num_draft_tokens - 1,
other=-1,
).to(tl.int64)
if next_draft_token >= 0:
residual_mass = _compute_global_residual_mass(
local_residual_mass_ptr,
local_residual_mass_stride,
prefix_joint_ratio,
target_logits_ptr,
target_logits_stride,
target_local_max_ptr,
target_local_max_stride,
target_local_sumexp_ptr,
target_local_sumexp_stride,
next_draft_token,
logit_idx + 1,
vocab_num_blocks,
PADDED_VOCAB_NUM_BLOCKS,
HAS_DRAFT_LOGITS,
)
denom = residual_mass + 1.0 - prefix_joint_ratio
h = tl.where(denom > 0.0, residual_mass / denom, 1.0)
else:
h = prefix_joint_ratio
accepted_length = tl.where(u <= h, i + 1, accepted_length)
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
else:
# Speculative decoding (Leviathan et al., 2023): https://arxiv.org/abs/2211.17192
# -1 is used for padded draft token ids that should be rejected.
is_valid_draft = draft_sampled >= 0
# Avoid possible OOB ptr access.
draft_sampled = tl.maximum(0, draft_sampled)
target_logprob, draft_logprob, target_lse, draft_lse = (
_compute_global_logprobs_and_logsumexp(
draft_sampled,
@@ -631,14 +655,15 @@ def _rejection_kernel(
)
if SYNTHETIC_MODE:
rate = tl.load(synthetic_conditional_rates_ptr + i)
accepted &= u < rate
accepted = u < rate
else:
# Probability ratio test: p(x) > u * q(x)
# Equivalent log form: log_p(x) > log(u) + log_q(x)
accepted &= target_logprob > tl.log(u) + draft_logprob
accepted &= is_valid_draft
accepted = target_logprob > tl.log(u) + draft_logprob
verifying = accepted
accepted_length += accepted
tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)
accepted_length += accepted
tl.store(rejected_steps_ptr + req_idx, accepted_length)
if USE_BLOCK_VERIFICATION and not is_greedy and accepted_length < num_draft_tokens:
# Compute the target and draft log exponential sums for the
@@ -722,6 +747,13 @@ def _resample_kernel(
# the target argmax is already in the sampled tensor.
return
rejected_draft_token = tl.load(
draft_sampled_ptr + resample_token_idx + 1,
mask=not is_bonus,
other=0,
)
is_valid_rejected_draft = rejected_draft_token >= 0
block_idx = tl.program_id(1)
block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = block < vocab_size
@@ -732,8 +764,9 @@ def _resample_kernel(
).to(tl.float32)
# Compute the residual logits to resample the rejected token from.
if is_bonus:
# Bonus token (no rejections). Directly use the target logits.
if is_bonus or not is_valid_rejected_draft:
# Bonus token (no rejections) or -1 placeholder token. In either case,
# directly use the target logits.
residual_logits = target_logits
elif HAS_DRAFT_LOGITS:
# draft_logits is stored pre-temperature, so apply scale first.
@@ -783,7 +816,6 @@ def _resample_kernel(
# p_tau * M_b(x) / Z otherwise
# Therefore p_tau is a constant that cancels under normalization,
# and does not need to be applied.
rejected_draft_token = tl.load(draft_sampled_ptr + resample_token_idx + 1)
residual_logits = tl.where(
block != rejected_draft_token,
target_logits,
@@ -1036,6 +1068,7 @@ def rejection_sample(
draft_local_max.stride(0),
draft_local_sumexp,
draft_local_sumexp.stride(0),
draft_sampled,
expanded_idx_mapping,
expanded_local_pos,
temperature,