A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 00:00 UTC
[Kernel] Optimize sample_recovered_tokens_kernel (#34974)
Signed-off-by: Xin Yang <xyangx@amazon.com>
Xin Yang committed
Feb 20, 2026 at 19:59 UTC
7a5adad48026d130348064ae7d41072ff999d1bf
2 files changed
+178
-33
tests/v1/sample/test_rejection_sampler.py
+126
-1
@@ -11,7 +11,11 @@ from tests.v1.sample.utils import create_allowed_token_ids
11
from vllm.platforms import current_platform
12
from vllm.v1.sample.logits_processor import LogitsProcessors
13
from vllm.v1.sample.metadata import SamplingMetadata
14
-from vllm.v1.sample.rejection_sampler import PLACEHOLDER_TOKEN_ID, RejectionSampler
14
+from vllm.v1.sample.rejection_sampler import (
15
+ PLACEHOLDER_TOKEN_ID,
16
+ RejectionSampler,
17
+ sample_recovered_tokens,
18
+)
19
from vllm.v1.sample.sampler import Sampler, SamplerOutput
20
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
21
@@ -518,6 +522,70 @@ def estimate_rejection_sampling_pdf(
522
return hist.hist
523
524
525
+def native_sample_recovered_tokens(
526
+ max_spec_len: int,
527
+ num_draft_tokens: list[int],
528
+ cu_num_draft_tokens: torch.Tensor, # [batch_size]
529
+ draft_token_ids: torch.Tensor, # [num_tokens]
530
+ draft_probs: torch.Tensor | None, # [num_tokens, vocab_size]
531
+ target_probs: torch.Tensor, # [num_tokens, vocab_size]
532
+ sampling_metadata: SamplingMetadata,
533
+ device: torch.device,
534
+) -> torch.Tensor:
535
+ batch_size = len(num_draft_tokens)
536
+ vocab_size = target_probs.shape[-1]
537
+
538
+ q = torch.empty(
539
+ (batch_size, vocab_size),
540
+ dtype=torch.float32,
541
+ device=device,
542
+ )
543
+ q.exponential_()
544
+
545
+ states = {
546
+ i: generator.get_state()
547
+ for i, generator in sampling_metadata.generators.items()
548
+ }
549
+ for i, generator in sampling_metadata.generators.items():
550
+ # Do not generate random numbers for requests with no draft tokens.
551
+ # This can be important for reproducibility.
552
+ if num_draft_tokens[i] > 0:
553
+ q[i].exponential_(generator=generator)
554
+
555
+ # In order to generate the same exponential later, reset the CUDA RNG
556
+ # state because RNG state advances after each call.
557
+ generator.set_state(states[i])
558
+
559
+ inv_q = q.reciprocal()
560
+
561
+ out = torch.empty_like(draft_token_ids)
562
+
563
+ for req_idx in range(batch_size):
564
+ start_idx = 0 if req_idx == 0 else int(cu_num_draft_tokens[req_idx - 1].item())
565
+ end_idx = int(cu_num_draft_tokens[req_idx].item())
566
+ num_tokens = end_idx - start_idx
567
+
568
+ for pos in range(max_spec_len):
569
+ if pos >= num_tokens:
570
+ continue
571
+ token_idx = start_idx + pos
572
+
573
+ if draft_probs is None:
574
+ # prob is target_probs[token_idx] except draft_token_id is zeroed
575
+ prob = target_probs[token_idx].clone()
576
+ draft_token_id = draft_token_ids[token_idx]
577
+ prob[draft_token_id] = 0.0
578
+ else:
579
+ prob = (target_probs[token_idx] - draft_probs[token_idx]).clamp_min_(
580
+ 0.0
581
+ )
582
+
583
+ score = prob * inv_q[req_idx]
584
+ recovered_id = torch.argmax(score, dim=-1)
585
+ out[token_idx] = recovered_id
586
+ return out
587
+
588
+
589
def _test_masked_logits(
590
rejection_sampler,
591
batch_size: int,
@@ -778,3 +846,60 @@ def test_allowed_token_ids(rejection_sampler):
846
device=logits.device,
847
)
848
assert torch.equal(output.sampled_token_ids, expected)
849
+
850
+
851
+@pytest.mark.parametrize("batch_size", [1, 100])
852
+@pytest.mark.parametrize("vocab_size", [100, 8192, 10000])
853
+@pytest.mark.parametrize("max_spec_len", [1, 3])
854
+@pytest.mark.parametrize("no_draft_probs", [True, False])
855
+def test_sample_recovered_tokens(
856
+ batch_size: int, vocab_size: int, max_spec_len: int, no_draft_probs: bool
857
+):
858
+ num_tokens = batch_size * max_spec_len
859
+
860
+ # Create random draft probabilities.
861
+ draft_probs = torch.rand(num_tokens, vocab_size, dtype=torch.float32, device=DEVICE)
862
+ draft_probs = F.softmax(draft_probs, dim=-1)
863
+
864
+ # Create random target probabilities.
865
+ target_logits = torch.rand(
866
+ num_tokens, vocab_size, dtype=torch.float32, device=DEVICE
867
+ )
868
+ target_probs = F.softmax(target_logits, dim=-1)
869
+
870
+ # Randomly sample draft token ids from draft probs
871
+ draft_token_ids = torch.multinomial(draft_probs, num_samples=1).to(torch.int32)
872
+
873
+ temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE)
874
+ generators = {
875
+ i: torch.Generator(device=DEVICE).manual_seed(i) for i in range(batch_size)
876
+ }
877
+ sampling_metadata = create_sampling_metadata(
878
+ all_greedy=False, temperature=temperature, generators=generators
879
+ )
880
+
881
+ spec_decode_metadata = create_spec_decode_metadata(
882
+ draft_token_ids.reshape(batch_size, max_spec_len).tolist(), target_logits
883
+ )
884
+
885
+ ref_recovered_token_ids = native_sample_recovered_tokens(
886
+ max_spec_len,
887
+ spec_decode_metadata.num_draft_tokens,
888
+ spec_decode_metadata.cu_num_draft_tokens,
889
+ draft_token_ids,
890
+ None if no_draft_probs else draft_probs,
891
+ target_probs,
892
+ sampling_metadata,
893
+ device=DEVICE,
894
+ )
895
+ recovered_token_ids = sample_recovered_tokens(
896
+ max_spec_len,
897
+ spec_decode_metadata.num_draft_tokens,
898
+ spec_decode_metadata.cu_num_draft_tokens,
899
+ draft_token_ids,
900
+ None if no_draft_probs else draft_probs,
901
+ target_probs,
902
+ sampling_metadata,
903
+ device=DEVICE,
904
+ )
905
+ assert torch.equal(recovered_token_ids, ref_recovered_token_ids)
vllm/v1/sample/rejection_sampler.py
+52
-32
@@ -623,16 +623,19 @@ def sample_recovered_tokens(
623
if num_draft_tokens[i] > 0:
624
q[i].exponential_(generator=generator)
625
626
+ inv_q = q.reciprocal()
627
+
628
recovered_token_ids = torch.empty_like(draft_token_ids)
629
+ BLOCK_SIZE = 8192
630
sample_recovered_tokens_kernel[(batch_size, max_spec_len)](
631
recovered_token_ids,
632
cu_num_draft_tokens,
633
draft_token_ids,
634
draft_probs,
635
target_probs,
633
- q,
636
+ inv_q,
637
vocab_size,
635
- triton.next_power_of_2(vocab_size),
638
+ BLOCK_SIZE,
639
NO_DRAFT_PROBS=draft_probs is None,
640
)
641
return recovered_token_ids
@@ -776,9 +779,9 @@ def sample_recovered_tokens_kernel(
779
draft_token_ids_ptr, # [num_tokens]
780
draft_probs_ptr, # [num_tokens, vocab_size] or None
781
target_probs_ptr, # [num_tokens, vocab_size]
779
- q_ptr, # [batch_size, vocab_size]
782
+ inv_q_ptr, # [batch_size, vocab_size]
783
vocab_size,
781
- PADDED_VOCAB_SIZE: tl.constexpr,
784
+ BLOCK_SIZE: tl.constexpr,
785
NO_DRAFT_PROBS: tl.constexpr,
786
):
787
req_idx = tl.program_id(0)
@@ -791,33 +794,50 @@ def sample_recovered_tokens_kernel(
794
if pos >= num_draft_tokens:
795
return
796
794
- vocab_offset = tl.arange(0, PADDED_VOCAB_SIZE)
797
+ token_idx = start_idx + pos
798
+
799
if NO_DRAFT_PROBS:
796
- draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
797
- prob = tl.load(
798
- target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset,
799
- mask=((vocab_offset < vocab_size) & (vocab_offset != draft_token_id)),
800
- other=0,
801
- )
802
- else:
803
- draft_prob = tl.load(
804
- draft_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset,
805
- mask=vocab_offset < vocab_size,
806
- other=0,
807
- )
808
- target_prob = tl.load(
809
- target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset,
810
- mask=vocab_offset < vocab_size,
811
- other=0,
800
+ draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
801
+
802
+ max_val = float("-inf")
803
+ recovered_id = 0
804
+ for v in range(0, vocab_size, BLOCK_SIZE):
805
+ vocab_offset = v + tl.arange(0, BLOCK_SIZE)
806
+ vocab_mask = vocab_offset < vocab_size
807
+
808
+ if NO_DRAFT_PROBS:
809
+ prob = tl.load(
810
+ target_probs_ptr + token_idx * vocab_size + vocab_offset,
811
+ mask=(vocab_mask & (vocab_offset != draft_token_id)),
812
+ other=0.0,
813
+ )
814
+ else:
815
+ draft_prob = tl.load(
816
+ draft_probs_ptr + token_idx * vocab_size + vocab_offset,
817
+ mask=vocab_mask,
818
+ other=0.0,
819
+ )
820
+ target_prob = tl.load(
821
+ target_probs_ptr + token_idx * vocab_size + vocab_offset,
822
+ mask=vocab_mask,
823
+ other=0.0,
824
+ )
825
+ prob = tl.maximum(target_prob - draft_prob, 0.0)
826
+ # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because
827
+ # `tl.argmax` will select the maximum value.
828
+
829
+ inv_q = tl.load(
830
+ inv_q_ptr + req_idx * vocab_size + vocab_offset,
831
+ mask=vocab_mask,
832
+ other=0.0,
833
)
813
- prob = tl.maximum(target_prob - draft_prob, 0)
814
- # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because
815
- # `tl.argmax` will select the maximum value.
816
-
817
- q = tl.load(
818
- q_ptr + req_idx * vocab_size + vocab_offset,
819
- mask=vocab_offset < vocab_size,
820
- other=float("-inf"),
821
- )
822
- recovered_id = tl.argmax(prob / q, axis=-1)
823
- tl.store(output_token_ids_ptr + start_idx + pos, recovered_id)
834
+
835
+ # Local tile reduction
836
+ score = prob * inv_q
837
+ local_max, local_id = tl.max(score, axis=0, return_indices=True)
838
+
839
+ if local_max > max_val:
840
+ max_val = local_max
841
+ recovered_id = v + local_id
842
+
843
+ tl.store(output_token_ids_ptr + token_idx, recovered_id)