@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 27930df9c2

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

Signed-off-by: Giancarlo Delfin <gdelfin@inferact.ai> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Giancarlo Delfin committed Aug 6, 2026 at 16:32 UTC 27930df9c2bd14047be35ff2a986ca72fc65631a

2 files changed +229 -74

tests/v1/spec_decode/test_rejection_sampler_utils.py

+122

@@ -367,6 +367,128 @@ def test_placeholder_draft_token_rejected():

367 assert (recovered >= 0).all() and (recovered < VOCAB_SIZE).all()

368

369

370 +@pytest.mark.parametrize("has_draft_logits", [True, False])

371 +@pytest.mark.parametrize("num_placeholders", [1, 2])

372 +def test_block_verification_placeholder_truncates_block(

373 + has_draft_logits: bool, num_placeholders: int

374 +):

375 + """Trailing placeholder (-1) draft ids end the block early. The last real

376 + draft token must then be verified against the closed-form threshold rather

377 + than a residual mass derived from the placeholder position, whose logits

378 + describe a token that was never drafted.

379 + """

380 + torch.manual_seed(42)

381 + device = "cuda"

382 + num_trials = 20 * VOCAB_SIZE

383 + K = 3

384 + temperature = 1.0

385 +

386 + target_logits_1d = torch.randn(VOCAB_SIZE, device=device) / temperature

387 + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device)

388 +

389 + inputs = _build_rejection_sample_inputs(

390 + target_logits_1d,

391 + draft_logits_1d,

392 + K,

393 + temperature=temperature,

394 + num_trials=num_trials,

395 + )

396 + if not has_draft_logits:

397 + inputs["draft_logits"] = None

398 + inputs["draft_sampled"].view(num_trials, K + 1)[:, K + 1 - num_placeholders :] = -1

399 +

400 + sampled, num_sampled = rejection_sample(

401 + **inputs, num_speculative_steps=K, use_block_verification=True

402 + )

403 +

404 + num_real = K - num_placeholders

405 + assert (num_sampled <= num_real + 1).all(), (

406 + "Accepted a draft token at or after a placeholder."

407 + )

408 + target_probs = torch.softmax(target_logits_1d, dim=0)

409 + for pos in range(num_real + 1):

410 + accepted_mask = num_sampled >= pos + 1

411 + _assert_distribution_match(

412 + sampled[accepted_mask, pos], target_probs, device, label=f"position {pos}"

413 + )

414 +

415 +

416 +def test_greedy_placeholder_emits_target_argmax():

417 + """Greedy sampling skips resampling and relies on the rejection kernel

418 + storing the target argmax at the rejected position. A placeholder must not

419 + bypass that store, or the output slot is returned uninitialized.

420 + """

421 + torch.manual_seed(0)

422 + device = "cuda"

423 + num_trials = 512

424 + K = 3

425 +

426 + target_logits_1d = torch.randn(VOCAB_SIZE, device=device)

427 + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device)

428 +

429 + inputs = _build_rejection_sample_inputs(

430 + target_logits_1d,

431 + draft_logits_1d,

432 + K,

433 + temperature=0.0,

434 + num_trials=num_trials,

435 + )

436 + target_argmax = int(target_logits_1d.argmax())

437 + draft_sampled = inputs["draft_sampled"].view(num_trials, K + 1)

438 + # Accept the first draft so that the rejection lands on the placeholder.

439 + draft_sampled[:, 1] = target_argmax

440 + draft_sampled[:, 2:] = -1

441 +

442 + sampled, num_sampled = rejection_sample(**inputs, num_speculative_steps=K)

443 +

444 + assert torch.equal(num_sampled, torch.full_like(num_sampled, 2))

445 + steps = torch.arange(K + 1, device=device).unsqueeze(0)

446 + emitted = sampled[steps < num_sampled.unsqueeze(1)]

447 + assert (emitted == target_argmax).all(), (

448 + "Greedy sampling emitted a token that is not the target argmax"

449 + )

450 +

451 +

452 +@pytest.mark.parametrize("use_block_verification", [False, True])

453 +def test_placeholder_blocks_later_draft_tokens(use_block_verification: bool):

454 + """A placeholder is not necessarily the final draft. Nothing at or after

455 + one may be accepted, even when a valid draft follows it.

456 + """

457 + torch.manual_seed(0)

458 + device = "cuda"

459 + num_trials = 4 * VOCAB_SIZE

460 + K = 3

461 + temperature = 1.0

462 +

463 + # An identical draft is accepted with probability ~1, so any token after

464 + # the placeholder would be accepted unless it is explicitly blocked. The

465 + # draft logits are stored pre-temperature, so pass the unscaled base.

466 + draft_logits_1d = torch.randn(VOCAB_SIZE, device=device)

467 + target_logits_1d = draft_logits_1d / temperature

468 +

469 + inputs = _build_rejection_sample_inputs(

470 + target_logits_1d,

471 + draft_logits_1d,

472 + K,

473 + temperature=temperature,

474 + num_trials=num_trials,

475 + )

476 + inputs["draft_sampled"].view(num_trials, K + 1)[:, 2] = -1

477 +

478 + _, num_sampled = rejection_sample(

479 + **inputs,

480 + num_speculative_steps=K,

481 + use_block_verification=use_block_verification,

482 + )

483 +

484 + # Only the first draft is verifiable, plus one resampled token.

485 + assert (num_sampled <= 2).all(), "Accepted a draft token past a placeholder."

486 + assert (num_sampled == 2).float().mean().item() > 0.9, (

487 + "The first draft was rarely accepted; the test is not exercising "

488 + "acceptance past the placeholder."

489 + )

490 +

491 +

492 @pytest.mark.parametrize(

493 "num_speculative_steps,temperature",

494 [

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

+107 -74

@@ -55,7 +55,7 @@ def _compute_global_residual_mass(

55 target_local_max_stride,

56 target_local_sumexp_ptr,

57 target_local_sumexp_stride,

58 - draft_sampled_ptr,

58 + draft_token,

59 logit_idx,

60 vocab_num_blocks,

61 PADDED_VOCAB_NUM_BLOCKS: tl.constexpr,

@@ -74,7 +74,6 @@ def _compute_global_residual_mass(

74 # One-hot draft. M_s is a point mass at this draft token

75 # so the residual mass reduces to the closed form:

76 # p * (1 - M_b(draft_token)).

77 - draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)

77 target_lse = _compute_global_logsumexp(

78 target_local_max_ptr,

79 target_local_max_stride,

@@ -347,35 +346,40 @@ def _compute_cumulative_log_p_kernel(

346 if temp == 0.0:

347 return

348

350 - log_p = 0.0

349 + log_p = tl.zeros((), tl.float32)

350 for step in range(num_draft_tokens):

351 logit_idx = start_idx + step

352 draft_token = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)

354 - target_logprob, draft_logprob, _, _ = _compute_global_logprobs_and_logsumexp(

355 - draft_token,

356 - True, # mask

357 - logit_idx,

358 - req_state_idx,

359 - step,

360 - temp,

361 - target_logits_ptr,

362 - target_logits_stride,

363 - target_local_max_ptr,

364 - target_local_max_stride,

365 - target_local_sumexp_ptr,

366 - target_local_sumexp_stride,

367 - draft_logits_ptr,

368 - draft_logits_stride_0,

369 - draft_logits_stride_1,

370 - draft_local_max_ptr,

371 - draft_local_max_stride,

372 - draft_local_sumexp_ptr,

373 - draft_local_sumexp_stride,

374 - vocab_num_blocks,

375 - PADDED_VOCAB_NUM_BLOCKS,

376 - HAS_DRAFT_LOGITS,

377 - )

378 - log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)

353 + # -1 placeholder tokens can never be accepted. Skip their reductions

354 + # and carry the last valid cumulative value.

355 + if draft_token >= 0:

356 + target_logprob, draft_logprob, _, _ = (

357 + _compute_global_logprobs_and_logsumexp(

358 + draft_token,

359 + True, # mask

360 + logit_idx,

361 + req_state_idx,

362 + step,

363 + temp,

364 + target_logits_ptr,

365 + target_logits_stride,

366 + target_local_max_ptr,

367 + target_local_max_stride,

368 + target_local_sumexp_ptr,

369 + target_local_sumexp_stride,

370 + draft_logits_ptr,

371 + draft_logits_stride_0,

372 + draft_logits_stride_1,

373 + draft_local_max_ptr,

374 + draft_local_max_stride,

375 + draft_local_sumexp_ptr,

376 + draft_local_sumexp_stride,

377 + vocab_num_blocks,

378 + PADDED_VOCAB_NUM_BLOCKS,

379 + HAS_DRAFT_LOGITS,

380 + )

381 + )

382 + log_p = tl.minimum(log_p + (target_logprob - draft_logprob), 0.0)

383 tl.store(cumulative_log_p_ptr + logit_idx, log_p)

384

385

@@ -406,6 +410,8 @@ def _compute_local_residual_mass_kernel(

410 draft_local_sumexp_ptr,

411 draft_local_sumexp_stride,

412 # [num_logits]

413 + draft_sampled_ptr,

414 + # [num_logits]

415 expanded_idx_mapping_ptr,

416 # [num_logits]

417 expanded_local_pos_ptr,

@@ -425,6 +431,11 @@ def _compute_local_residual_mass_kernel(

431 # first and last (bonus) positions aren't needed for this computation.

432 return

433

434 + if tl.load(draft_sampled_ptr + logit_idx + 1) < 0:

435 + # -1 placeholder token. The rejection kernel treats the preceding token

436 + # as the end of the block, so this position's residual mass is unused.

437 + return

438 +

439 req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64)

440 temp = tl.load(temp_ptr + req_state_idx).to(tl.float32)

441 if temp == 0.0:

@@ -539,41 +550,23 @@ def _rejection_kernel(

550 accepted_length = tl.zeros((), tl.int64)

551 target_lse = 0.0

552 draft_lse = 0.0

542 - accepted = True

553 + verifying = True

554 for i in range(num_draft_tokens):

555 logit_idx = start_idx + i

556 draft_sampled = tl.load(draft_sampled_ptr + logit_idx + 1).to(tl.int64)

546 - pos = tl.load(pos_ptr + logit_idx)

547 - u = tl_rand32(seed, pos, includes_zero=False)

548 - if USE_BLOCK_VERIFICATION and not is_greedy:

549 - # Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444

550 - prefix_joint_ratio = tl.exp(

551 - tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32)

552 - )

553 - if i < num_draft_tokens - 1:

554 - residual_mass = _compute_global_residual_mass(

555 - local_residual_mass_ptr,

556 - local_residual_mass_stride,

557 - prefix_joint_ratio,

558 - target_logits_ptr,

559 - target_logits_stride,

560 - target_local_max_ptr,

561 - target_local_max_stride,

562 - target_local_sumexp_ptr,

563 - target_local_sumexp_stride,

564 - draft_sampled_ptr,

565 - logit_idx + 1,

566 - vocab_num_blocks,

567 - PADDED_VOCAB_NUM_BLOCKS,

568 - HAS_DRAFT_LOGITS,

569 - )

570 - denom = residual_mass + 1.0 - prefix_joint_ratio

571 - h = tl.where(denom > 0.0, residual_mass / denom, 1.0)

572 - else:

573 - h = prefix_joint_ratio

574 - accepted_length = tl.where(u <= h, i + 1, accepted_length)

575 - tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)

576 - elif accepted:

557 + # -1 is used for placeholder draft token ids that should be rejected.

558 + is_valid_draft = draft_sampled >= 0

559 + # Avoid possible OOB ptr access.

560 + draft_sampled = tl.maximum(0, draft_sampled)

561 + if not is_greedy:

562 + # A -1 placeholder ends verification. Greedy is excluded because it

563 + # stores the target argmax upon first rejection, so it rejects the

564 + # placeholder via `accepted` instead.

565 + verifying &= is_valid_draft

566 +

567 + if verifying:

568 + pos = tl.load(pos_ptr + logit_idx)

569 + u = tl_rand32(seed, pos, includes_zero=False)

570 if is_greedy:

571 # Greedy sampling. Accept IFF draft matches target argmax.

572 # NOTE: Target argmax is stored directly so that resampling

@@ -589,20 +582,51 @@ def _rejection_kernel(

582 )

583 if SYNTHETIC_MODE:

584 rate = tl.load(synthetic_conditional_rates_ptr + i)

592 - # -1 is used for padded draft token ids that should be rejected.

593 - accepted &= (u < rate) & (draft_sampled >= 0)

585 + accepted = u < rate

586 else:

595 - accepted &= target_argmax == draft_sampled

587 + accepted = target_argmax == draft_sampled

588 + accepted &= is_valid_draft

589 + verifying = accepted

590 + accepted_length += accepted

591 tl.store(

592 sampled_ptr + req_idx * sampled_stride + i,

593 draft_sampled if accepted else target_argmax,

594 )

595 + elif USE_BLOCK_VERIFICATION:

596 + # Block verification (Sun et al., 2024): https://arxiv.org/abs/2403.10444

597 + prefix_joint_ratio = tl.exp(

598 + tl.load(cumulative_log_p_ptr + logit_idx).to(tl.float32)

599 + )

600 + next_draft_token = tl.load(

601 + draft_sampled_ptr + logit_idx + 2,

602 + mask=i < num_draft_tokens - 1,

603 + other=-1,

604 + ).to(tl.int64)

605 + if next_draft_token >= 0:

606 + residual_mass = _compute_global_residual_mass(

607 + local_residual_mass_ptr,

608 + local_residual_mass_stride,

609 + prefix_joint_ratio,

610 + target_logits_ptr,

611 + target_logits_stride,

612 + target_local_max_ptr,

613 + target_local_max_stride,

614 + target_local_sumexp_ptr,

615 + target_local_sumexp_stride,

616 + next_draft_token,

617 + logit_idx + 1,

618 + vocab_num_blocks,

619 + PADDED_VOCAB_NUM_BLOCKS,

620 + HAS_DRAFT_LOGITS,

621 + )

622 + denom = residual_mass + 1.0 - prefix_joint_ratio

623 + h = tl.where(denom > 0.0, residual_mass / denom, 1.0)

624 + else:

625 + h = prefix_joint_ratio

626 + accepted_length = tl.where(u <= h, i + 1, accepted_length)

627 + tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)

628 else:

629 # Speculative decoding (Leviathan et al., 2023): https://arxiv.org/abs/2211.17192

602 - # -1 is used for padded draft token ids that should be rejected.

603 - is_valid_draft = draft_sampled >= 0

604 - # Avoid possible OOB ptr access.

605 - draft_sampled = tl.maximum(0, draft_sampled)

630 target_logprob, draft_logprob, target_lse, draft_lse = (

631 _compute_global_logprobs_and_logsumexp(

632 draft_sampled,

@@ -631,14 +655,15 @@ def _rejection_kernel(

655 )

656 if SYNTHETIC_MODE:

657 rate = tl.load(synthetic_conditional_rates_ptr + i)

634 - accepted &= u < rate

658 + accepted = u < rate

659 else:

660 # Probability ratio test: p(x) > u * q(x)

661 # Equivalent log form: log_p(x) > log(u) + log_q(x)

638 - accepted &= target_logprob > tl.log(u) + draft_logprob

639 - accepted &= is_valid_draft

662 + accepted = target_logprob > tl.log(u) + draft_logprob

663 + verifying = accepted

664 + accepted_length += accepted

665 tl.store(sampled_ptr + req_idx * sampled_stride + i, draft_sampled)

641 - accepted_length += accepted

666 +

667 tl.store(rejected_steps_ptr + req_idx, accepted_length)

668 if USE_BLOCK_VERIFICATION and not is_greedy and accepted_length < num_draft_tokens:

669 # Compute the target and draft log exponential sums for the

@@ -722,6 +747,13 @@ def _resample_kernel(

747 # the target argmax is already in the sampled tensor.

748 return

749

750 + rejected_draft_token = tl.load(

751 + draft_sampled_ptr + resample_token_idx + 1,

752 + mask=not is_bonus,

753 + other=0,

754 + )

755 + is_valid_rejected_draft = rejected_draft_token >= 0

756 +

757 block_idx = tl.program_id(1)

758 block = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

759 mask = block < vocab_size

@@ -732,8 +764,9 @@ def _resample_kernel(

764 ).to(tl.float32)

765

766 # Compute the residual logits to resample the rejected token from.

735 - if is_bonus:

736 - # Bonus token (no rejections). Directly use the target logits.

767 + if is_bonus or not is_valid_rejected_draft:

768 + # Bonus token (no rejections) or -1 placeholder token. In either case,

769 + # directly use the target logits.

770 residual_logits = target_logits

771 elif HAS_DRAFT_LOGITS:

772 # draft_logits is stored pre-temperature, so apply scale first.

@@ -783,7 +816,6 @@ def _resample_kernel(

816 # p_tau * M_b(x) / Z otherwise

817 # Therefore p_tau is a constant that cancels under normalization,

818 # and does not need to be applied.

786 - rejected_draft_token = tl.load(draft_sampled_ptr + resample_token_idx + 1)

819 residual_logits = tl.where(

820 block != rejected_draft_token,

821 target_logits,

@@ -1036,6 +1068,7 @@ def rejection_sample(

1068 draft_local_max.stride(0),

1069 draft_local_sumexp,

1070 draft_local_sumexp.stride(0),

1071 + draft_sampled,

1072 expanded_idx_mapping,

1073 expanded_local_pos,

1074 temperature,