@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7fe7800fa4
Fork 0

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 1, 21:15 UTC

[BUG] Fix FP64 Gumbel precision coverage (#43150)

Signed-off-by: tianyu-z <zhangtianyupro@gmail.com> Signed-off-by: Tianyu Zhang <53099276+tianyu-z@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Tianyu Zhang committed Jun 5, 2026 at 04:04 UTC 7fe7800fa4ce9709e9e3361f6b44b1ccf28c775f
11 files changed +391 -21
tests/v1/sample/test_rejection_sampler.py
+64 -1
@@ -544,13 +544,15 @@ def native_sample_recovered_tokens(
544 target_probs: torch.Tensor, # [num_tokens, vocab_size]
545 sampling_metadata: SamplingMetadata,
546 device: torch.device,
547 + use_fp64_gumbel: bool = False,
548 ) -> torch.Tensor:
549 batch_size = len(num_draft_tokens)
550 vocab_size = target_probs.shape[-1]
551 + q_dtype = torch.float64 if use_fp64_gumbel else torch.float32
552
553 q = torch.empty(
554 (batch_size, vocab_size),
553 - dtype=torch.float32,
555 + dtype=q_dtype,
556 device=device,
557 )
558 q.exponential_()
@@ -935,6 +937,67 @@ def test_sample_recovered_tokens(
937 assert torch.equal(recovered_token_ids, ref_recovered_token_ids)
938
939
940 +def test_sample_recovered_tokens_uses_fp64_exponential_race_when_requested():
941 + batch_size = 2
942 + vocab_size = 64
943 + max_spec_len = 2
944 + num_tokens = batch_size * max_spec_len
945 +
946 + draft_probs = torch.rand(
947 + num_tokens,
948 + vocab_size,
949 + dtype=torch.float32,
950 + device=DEVICE_TYPE,
951 + )
952 + draft_probs = F.softmax(draft_probs, dim=-1)
953 + target_probs = torch.rand(
954 + num_tokens,
955 + vocab_size,
956 + dtype=torch.float32,
957 + device=DEVICE_TYPE,
958 + )
959 + target_probs = F.softmax(target_probs, dim=-1)
960 + draft_token_ids = torch.multinomial(draft_probs, num_samples=1).to(torch.int32)
961 +
962 + generators = {
963 + i: torch.Generator(device=DEVICE_TYPE).manual_seed(i) for i in range(batch_size)
964 + }
965 + sampling_metadata = create_sampling_metadata(
966 + all_greedy=False,
967 + temperature=torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE),
968 + generators=generators,
969 + )
970 + spec_decode_metadata = create_spec_decode_metadata(
971 + draft_token_ids.reshape(batch_size, max_spec_len).tolist(),
972 + target_probs.log(),
973 + )
974 +
975 + expected = native_sample_recovered_tokens(
976 + max_spec_len,
977 + spec_decode_metadata.num_draft_tokens,
978 + spec_decode_metadata.cu_num_draft_tokens,
979 + draft_token_ids,
980 + draft_probs,
981 + target_probs,
982 + sampling_metadata,
983 + device=torch.device(DEVICE_TYPE),
984 + use_fp64_gumbel=True,
985 + )
986 + actual = sample_recovered_tokens(
987 + max_spec_len,
988 + spec_decode_metadata.num_draft_tokens,
989 + spec_decode_metadata.cu_num_draft_tokens,
990 + draft_token_ids,
991 + draft_probs,
992 + target_probs,
993 + sampling_metadata,
994 + device=torch.device(DEVICE_TYPE),
995 + use_fp64_gumbel=True,
996 + )
997 +
998 + assert torch.equal(actual, expected)
999 +
1000 +
1001 ########################### Tests for Synthetic Rejection Sampling #########
1002
1003
tests/v1/sample/test_topk_topp_sampler.py
+39 -1
@@ -6,7 +6,12 @@ from torch import Generator
6
7 from vllm.platforms import current_platform
8 from vllm.triton_utils import HAS_TRITON
9 -from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p_pytorch
9 +from vllm.utils.torch_utils import set_random_seed
10 +from vllm.v1.sample.ops.topk_topp_sampler import (
11 + apply_top_k_top_p_pytorch,
12 + random_sample,
13 +)
14 +from vllm.v1.sample.sampler import Sampler
15
16 DEVICE_TYPE = current_platform.device_type
17
@@ -38,6 +43,10 @@ def _flashinfer_topk_topp_supported() -> bool:
43 FLASHINFER_TOPK_TOPP_SUPPORTED = _flashinfer_topk_topp_supported()
44
45
46 +def _seed_default_generator(seed: int) -> None:
47 + set_random_seed(seed)
48 +
49 +
50 @pytest.fixture(autouse=True)
51 def reset_default_device():
52 """
@@ -49,6 +58,35 @@ def reset_default_device():
58 torch.set_default_device(original_device)
59
60
61 +def test_sampler_threads_fp64_gumbel_to_topk_topp_sampler():
62 + sampler = Sampler(use_fp64_gumbel=True)
63 +
64 + assert sampler.topk_topp_sampler.use_fp64_gumbel
65 +
66 +
67 +def test_random_sample_uses_fp64_exponential_race_when_requested():
68 + torch.set_default_device(DEVICE_TYPE)
69 + probs = torch.tensor(
70 + [
71 + [0.70, 0.20, 0.10],
72 + [0.05, 0.15, 0.80],
73 + [0.25, 0.25, 0.50],
74 + ],
75 + dtype=torch.float32,
76 + device=DEVICE_TYPE,
77 + )
78 +
79 + _seed_default_generator(12345)
80 + q = torch.empty(probs.shape, dtype=torch.float64, device=probs.device)
81 + q.exponential_()
82 + expected = q.reciprocal_().mul_(probs).argmax(dim=-1).view(-1)
83 +
84 + _seed_default_generator(12345)
85 + actual = random_sample(probs.clone(), {}, use_fp64_gumbel=True)
86 +
87 + assert torch.equal(actual, expected)
88 +
89 +
90 def test_topk_impl_equivalence():
91 torch.set_default_device(DEVICE_TYPE)
92 generator = Generator(device=DEVICE_TYPE).manual_seed(33)
tests/v1/spec_decode/test_eagle.py
+2 -1
@@ -1034,7 +1034,8 @@ def test_propose_stores_probabilistic_draft_probs(monkeypatch):
1034 proposer.model = model_mock
1035 proposer._draft_attn_layer_names = {"layer.0"}
1036
1037 - def fake_compute_probs(logits, sampling_metadata):
1037 + def fake_compute_probs(logits, sampling_metadata, use_fp64_gumbel):
1038 + assert use_fp64_gumbel == proposer.use_fp64_gumbel
1039 probs = torch.softmax(logits, dim=-1)
1040 return probs.argmax(dim=-1), probs
1041
tests/v1/spec_decode/test_llm_base_proposer_sampling.py new
+70
@@ -0,0 +1,70 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.platforms import current_platform
7 +from vllm.utils.torch_utils import set_random_seed
8 +from vllm.v1.sample.logits_processor import LogitsProcessors
9 +from vllm.v1.sample.metadata import SamplingMetadata
10 +from vllm.v1.spec_decode.llm_base_proposer import (
11 + compute_probs_and_sample_next_token,
12 +)
13 +
14 +DEVICE_TYPE = current_platform.device_type
15 +
16 +
17 +def _seed_default_generator(seed: int) -> None:
18 + set_random_seed(seed)
19 +
20 +
21 +def _make_sampling_metadata(batch_size: int) -> SamplingMetadata:
22 + return SamplingMetadata(
23 + temperature=torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE),
24 + all_greedy=False,
25 + all_random=True,
26 + top_p=None,
27 + top_k=None,
28 + generators={},
29 + max_num_logprobs=None,
30 + no_penalties=True,
31 + prompt_token_ids=None,
32 + frequency_penalties=torch.empty(0, device=DEVICE_TYPE),
33 + presence_penalties=torch.empty(0, device=DEVICE_TYPE),
34 + repetition_penalties=torch.empty(0, device=DEVICE_TYPE),
35 + output_token_ids=[[] for _ in range(batch_size)],
36 + spec_token_ids=[[] for _ in range(batch_size)],
37 + allowed_token_ids_mask=None,
38 + bad_words_token_ids={},
39 + logitsprocs=LogitsProcessors(),
40 + )
41 +
42 +
43 +def test_compute_probs_and_sample_next_token_uses_fp64_exponential_race():
44 + batch_size = 4
45 + vocab_size = 32
46 + generator = torch.Generator(device=DEVICE_TYPE).manual_seed(11)
47 + logits = torch.randn(
48 + batch_size,
49 + vocab_size,
50 + dtype=torch.float32,
51 + device=DEVICE_TYPE,
52 + generator=generator,
53 + )
54 + metadata = _make_sampling_metadata(batch_size)
55 +
56 + _seed_default_generator(12345)
57 + probs = logits.softmax(dim=-1, dtype=torch.float32)
58 + q = torch.empty(probs.shape, dtype=torch.float64, device=probs.device)
59 + q.exponential_()
60 + expected_ids = q.reciprocal_().mul_(probs).argmax(dim=-1).view(-1)
61 +
62 + _seed_default_generator(12345)
63 + actual_ids, actual_probs = compute_probs_and_sample_next_token(
64 + logits.clone(),
65 + metadata,
66 + use_fp64_gumbel=True,
67 + )
68 +
69 + assert torch.equal(actual_ids, expected_ids)
70 + assert torch.allclose(actual_probs, probs)
tools/gumbel_precision/prove_exponential_race_precision.py new
+141
@@ -0,0 +1,141 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""CUDA proof for fp32 exponential-race tail truncation.
4 +
5 +This script is intentionally not a unit test. It is a reproducible, GPU-only
6 +statistical proof for the hidden Gumbel-max idiom:
7 +
8 + q.exponential_()
9 + sample = (probs / q).argmax()
10 +
11 +For q ~ Exp(1), this is equivalent to argmax(log(probs) + Gumbel). On CUDA,
12 +fp32 exponential samples inherit a 24-bit uniform lower-tail cutoff, so very
13 +small q values are impossible. The many-tail experiment below chooses a case
14 +where a correct sampler should select a low-probability tail token dozens of
15 +times, while fp32 q cannot select one.
16 +"""
17 +
18 +from __future__ import annotations
19 +
20 +import argparse
21 +import math
22 +import time
23 +
24 +import torch
25 +
26 +
27 +def _seed(seed: int) -> None:
28 + torch.manual_seed(seed)
29 +
30 +
31 +def measure_exponential_lower_tail(
32 + *,
33 + device: torch.device,
34 + samples: int,
35 + chunk_size: int,
36 + seed: int,
37 +) -> None:
38 + threshold = 2.0**-24
39 + print(f"lower-tail threshold: {threshold:.18e}")
40 + for dtype in (torch.float32, torch.float64):
41 + _seed(seed)
42 + count_below = 0
43 + min_q = float("inf")
44 + max_q = 0.0
45 + start = time.perf_counter()
46 + remaining = samples
47 + while remaining > 0:
48 + n = min(chunk_size, remaining)
49 + q = torch.empty((n,), dtype=dtype, device=device)
50 + q.exponential_()
51 + count_below += int((q < threshold).sum().item())
52 + min_q = min(min_q, float(q.min().item()))
53 + max_q = max(max_q, float(q.max().item()))
54 + remaining -= n
55 + torch.accelerator.synchronize()
56 + elapsed = time.perf_counter() - start
57 + print(
58 + f"{dtype}: samples={samples} count(q < 2^-24)={count_below} "
59 + f"min={min_q:.18e} max={max_q:.6f} elapsed={elapsed:.2f}s"
60 + )
61 +
62 +
63 +def run_many_tail_race(
64 + *,
65 + device: torch.device,
66 + trials: int,
67 + num_tail_tokens: int,
68 + gap: float,
69 + chunk_trials: int,
70 + seed: int,
71 +) -> None:
72 + p_tail = math.exp(-gap)
73 + expected_tail_hits = (
74 + trials * (num_tail_tokens * p_tail) / (1.0 + num_tail_tokens * p_tail)
75 + )
76 + print(
77 + "many-tail race: "
78 + f"trials={trials} num_tail_tokens={num_tail_tokens} gap={gap} "
79 + f"expected_tail_hits={expected_tail_hits:.4f}"
80 + )
81 +
82 + for dtype in (torch.float32, torch.float64):
83 + _seed(seed)
84 + hits = 0
85 + p0 = torch.tensor(1.0, dtype=dtype, device=device)
86 + pt = torch.tensor(p_tail, dtype=dtype, device=device)
87 + start = time.perf_counter()
88 + remaining = trials
89 + while remaining > 0:
90 + batch = min(chunk_trials, remaining)
91 + q0 = torch.empty((batch,), dtype=dtype, device=device)
92 + q0.exponential_()
93 + qt = torch.empty((batch, num_tail_tokens), dtype=dtype, device=device)
94 + qt.exponential_()
95 + head_score = p0 / q0
96 + tail_score = (pt / qt).amax(dim=-1)
97 + hits += int((tail_score > head_score).sum().item())
98 + remaining -= batch
99 + torch.accelerator.synchronize()
100 + elapsed = time.perf_counter() - start
101 + print(f"{dtype}: tail_hits={hits} elapsed={elapsed:.2f}s")
102 +
103 +
104 +def main() -> None:
105 + parser = argparse.ArgumentParser()
106 + parser.add_argument("--lower-tail-samples", type=int, default=200_000_000)
107 + parser.add_argument("--lower-tail-chunk-size", type=int, default=10_000_000)
108 + parser.add_argument("--race-trials", type=int, default=100_000)
109 + parser.add_argument("--race-tail-tokens", type=int, default=262_144)
110 + parser.add_argument("--race-gap", type=float, default=20.5)
111 + parser.add_argument("--race-chunk-trials", type=int, default=64)
112 + parser.add_argument("--seed", type=int, default=2026)
113 + args = parser.parse_args()
114 +
115 + if not torch.accelerator.is_available():
116 + raise RuntimeError("CUDA is required for this proof.")
117 +
118 + device = torch.accelerator.current_accelerator()
119 + if device.type != "cuda":
120 + raise RuntimeError("CUDA is required for this proof.")
121 +
122 + print(f"torch={torch.__version__} cuda={torch.version.cuda}")
123 + print(f"device={device}")
124 + measure_exponential_lower_tail(
125 + device=device,
126 + samples=args.lower_tail_samples,
127 + chunk_size=args.lower_tail_chunk_size,
128 + seed=args.seed,
129 + )
130 + run_many_tail_race(
131 + device=device,
132 + trials=args.race_trials,
133 + num_tail_tokens=args.race_tail_tokens,
134 + gap=args.race_gap,
135 + chunk_trials=args.race_chunk_trials,
136 + seed=args.seed,
137 + )
138 +
139 +
140 +if __name__ == "__main__":
141 + main()
vllm/config/model.py
+4 -3
@@ -235,9 +235,10 @@ class ModelConfig:
235 temperature and top_k/top_p.
236 """
237 use_fp64_gumbel: bool = False
238 - """Whether to use FP64 (instead of FP32) for the Gumbel noise used by the
239 - sampler. FP64 reduces the chance of ties in Gumbel-max sampling at the cost
240 - of significantly lower kernel throughput on most GPUs."""
238 + """Whether to use FP64 (instead of FP32) random noise for Gumbel-max and
239 + equivalent exponential-race sampling. FP64 preserves lower-tail sampling
240 + events that fp32 uniform/exponential draws can truncate, at the cost of
241 + significantly lower throughput on most GPUs."""
242 disable_sliding_window: bool = False
243 """Whether to disable sliding window. If True, we will disable the sliding
244 window functionality of the model, capping to sliding window size. If the
vllm/v1/sample/ops/topk_topp_sampler.py
+36 -7
@@ -75,9 +75,14 @@ class TopKTopPSampler(nn.Module):
75 Implementations may update the logits tensor in-place.
76 """
77
78 - def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs") -> None:
78 + def __init__(
79 + self,
80 + logprobs_mode: LogprobsMode = "raw_logprobs",
81 + use_fp64_gumbel: bool = False,
82 + ) -> None:
83 super().__init__()
84 self.logprobs_mode = logprobs_mode
85 + self.use_fp64_gumbel = use_fp64_gumbel
86 if current_platform.is_cuda():
87 # FlashInfer doesn't expose post-top-k/top-p logits/logprobs,
88 # so it can't be used when the configured mode requires them.
@@ -142,7 +147,10 @@ class TopKTopPSampler(nn.Module):
147 elif self.logprobs_mode == "processed_logprobs":
148 logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
149 probs = logits.softmax(dim=-1, dtype=torch.float32)
145 - return random_sample(probs, generators), logits_to_return
150 + return (
151 + random_sample(probs, generators, self.use_fp64_gumbel),
152 + logits_to_return,
153 + )
154
155 def forward_cuda(
156 self,
@@ -163,6 +171,8 @@ class TopKTopPSampler(nn.Module):
171 "PyTorch-native implementation."
172 )
173 return self.forward_native(logits, generators, k, p)
174 + if self.use_fp64_gumbel:
175 + return self.forward_native(logits, generators, k, p)
176 assert self.logprobs_mode not in ("processed_logits", "processed_logprobs"), (
177 "FlashInfer does not support returning logits/logprobs"
178 )
@@ -190,16 +200,16 @@ class TopKTopPSampler(nn.Module):
200 elif self.logprobs_mode == "processed_logprobs":
201 logits_to_return = logits.log_softmax(dim=-1, dtype=torch.float32)
202
193 - if len(generators) != logits.shape[0]:
203 + if len(generators) != logits.shape[0] and not self.use_fp64_gumbel:
204 return compiled_random_sample(logits), logits_to_return
205
206 probs = logits.softmax(dim=-1, dtype=torch.float32)
197 - q = torch.empty_like(probs)
207 + q = empty_exponential_noise_like(probs, self.use_fp64_gumbel)
208 q.exponential_()
209 for i, generator in generators.items():
210 q[i].exponential_(generator=generator)
211
202 - return probs.div_(q).argmax(dim=-1).view(-1), logits_to_return
212 + return sample_with_exponential_noise(probs, q), logits_to_return
213
214 def forward_hip(
215 self,
@@ -216,6 +226,8 @@ class TopKTopPSampler(nn.Module):
226 "falling back to PyTorch-native."
227 )
228 return self.forward_native(logits, generators, k, p)
229 + if self.use_fp64_gumbel:
230 + return self.forward_native(logits, generators, k, p)
231 assert self.logprobs_mode not in (
232 "processed_logits",
233 "processed_logprobs",
@@ -404,16 +416,33 @@ def apply_top_k_only(logits: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
416 return logits.masked_fill_(logits < top_k_mask, -float("inf"))
417
418
419 +def empty_exponential_noise_like(
420 + probs: torch.Tensor, use_fp64_gumbel: bool
421 +) -> torch.Tensor:
422 + dtype = torch.float64 if use_fp64_gumbel else probs.dtype
423 + return torch.empty(probs.shape, dtype=dtype, device=probs.device)
424 +
425 +
426 +def sample_with_exponential_noise(probs: torch.Tensor, q: torch.Tensor) -> torch.Tensor:
427 + if q.dtype == probs.dtype:
428 + scores = probs.div_(q)
429 + else:
430 + scores = q.reciprocal_()
431 + scores.mul_(probs)
432 + return scores.argmax(dim=-1).view(-1)
433 +
434 +
435 def random_sample(
436 probs: torch.Tensor,
437 generators: dict[int, torch.Generator],
438 + use_fp64_gumbel: bool = False,
439 ) -> torch.Tensor:
440 """Randomly sample from the probabilities.
441
442 We use this function instead of torch.multinomial because torch.multinomial
443 causes CPU-GPU synchronization.
444 """
416 - q = torch.empty_like(probs)
445 + q = empty_exponential_noise_like(probs, use_fp64_gumbel)
446 # NOTE(woosuk): To batch-process the requests without their own seeds,
447 # which is the common case, we first assume that every request does
448 # not have its own seed. Then, we overwrite the values for the requests
@@ -425,7 +454,7 @@ def random_sample(
454 # one by one. Optimize this.
455 for i, generator in generators.items():
456 q[i].exponential_(generator=generator)
428 - return probs.div_(q).argmax(dim=-1).view(-1)
457 + return sample_with_exponential_noise(probs, q)
458
459
460 def flashinfer_sample(
vllm/v1/sample/rejection_sampler.py
+13 -2
@@ -65,6 +65,7 @@ class RejectionSampler(nn.Module):
65 ):
66 super().__init__()
67 self.sampler = sampler
68 + self.use_fp64_gumbel = getattr(sampler, "use_fp64_gumbel", False)
69 logprobs_mode = self.sampler.logprobs_mode
70 self.is_processed_logprobs_mode = logprobs_mode.startswith("processed")
71 self.is_logits_logprobs_mode = logprobs_mode.endswith("logits")
@@ -176,6 +177,7 @@ class RejectionSampler(nn.Module):
177 sampling_metadata,
178 synthetic_mode=self.synthetic_mode,
179 synthetic_conditional_rates=self.synthetic_conditional_rates,
180 + use_fp64_gumbel=self.use_fp64_gumbel,
181 )
182
183 logprobs_tensors = None
@@ -406,6 +408,7 @@ def rejection_sample(
408 sampling_metadata: SamplingMetadata,
409 synthetic_mode: bool = False,
410 synthetic_conditional_rates: torch.Tensor | None = None,
411 + use_fp64_gumbel: bool = False,
412 ) -> torch.Tensor:
413 assert draft_token_ids.ndim == 1
414 assert draft_probs is None or draft_probs.ndim == 2
@@ -480,6 +483,7 @@ def rejection_sample(
483 target_probs,
484 sampling_metadata,
485 device,
486 + use_fp64_gumbel,
487 )
488
489 # Rejection sampling for random sampling requests.
@@ -669,13 +673,15 @@ def sample_recovered_tokens(
673 target_probs: torch.Tensor,
674 sampling_metadata: SamplingMetadata,
675 device: torch.device,
676 + use_fp64_gumbel: bool = False,
677 ) -> torch.Tensor:
678 # NOTE(woosuk): Create only one distribution for each request.
679 batch_size = len(num_draft_tokens)
680 vocab_size = target_probs.shape[-1]
681 + q_dtype = torch.float64 if use_fp64_gumbel else torch.float32
682 q = torch.empty(
683 (batch_size, vocab_size),
678 - dtype=torch.float32,
684 + dtype=q_dtype,
685 device=device,
686 )
687 q.exponential_()
@@ -699,6 +705,7 @@ def sample_recovered_tokens(
705 vocab_size,
706 BLOCK_SIZE,
707 NO_DRAFT_PROBS=draft_probs is None,
708 + USE_FP64_GUMBEL=use_fp64_gumbel,
709 )
710 return recovered_token_ids
711
@@ -861,6 +868,7 @@ def sample_recovered_tokens_kernel(
868 vocab_size,
869 BLOCK_SIZE: tl.constexpr,
870 NO_DRAFT_PROBS: tl.constexpr,
871 + USE_FP64_GUMBEL: tl.constexpr,
872 ):
873 req_idx = tl.program_id(0)
874 start_idx = 0 if req_idx == 0 else tl.load(cu_num_draft_tokens_ptr + req_idx - 1)
@@ -877,7 +885,10 @@ def sample_recovered_tokens_kernel(
885 if NO_DRAFT_PROBS:
886 draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
887
880 - max_val = float("-inf")
888 + if USE_FP64_GUMBEL:
889 + max_val = tl.full((), float("-inf"), tl.float64)
890 + else:
891 + max_val = tl.full((), float("-inf"), tl.float32)
892 recovered_id = 0
893 for v in range(0, vocab_size, BLOCK_SIZE):
894 vocab_offset = v + tl.arange(0, BLOCK_SIZE)
vllm/v1/sample/sampler.py
+7 -2
@@ -58,11 +58,16 @@ class Sampler(nn.Module):
58 9. Return the final `SamplerOutput`.
59 """
60
61 - def __init__(self, logprobs_mode: LogprobsMode = "raw_logprobs"):
61 + def __init__(
62 + self,
63 + logprobs_mode: LogprobsMode = "raw_logprobs",
64 + use_fp64_gumbel: bool = False,
65 + ):
66 super().__init__()
63 - self.topk_topp_sampler = TopKTopPSampler(logprobs_mode)
67 + self.topk_topp_sampler = TopKTopPSampler(logprobs_mode, use_fp64_gumbel)
68 self.pin_memory = is_pin_memory_available()
69 self.logprobs_mode = logprobs_mode
70 + self.use_fp64_gumbel = use_fp64_gumbel
71
72 def forward(
73 self,
vllm/v1/spec_decode/llm_base_proposer.py
+11 -3
@@ -32,6 +32,10 @@ from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
32 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher
33 from vllm.v1.kv_cache_interface import KVCacheConfig, UniformTypeKVCacheSpecs
34 from vllm.v1.sample.metadata import SamplingMetadata
35 +from vllm.v1.sample.ops.topk_topp_sampler import (
36 + empty_exponential_noise_like,
37 + sample_with_exponential_noise,
38 +)
39 from vllm.v1.sample.sampler import _SAMPLING_EPS
40 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
41 from vllm.v1.spec_decode.utils import (
@@ -113,6 +117,7 @@ class SpecDecodeBaseProposer:
117 self.use_local_argmax_reduction: bool = (
118 self.speculative_config.use_local_argmax_reduction
119 )
120 + self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel
121
122 self.max_batch_size = vllm_config.scheduler_config.max_num_seqs
123 self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
@@ -409,7 +414,9 @@ class SpecDecodeBaseProposer:
414 return logits.argmax(dim=-1), None
415 if sampling_metadata.all_greedy:
416 return logits.argmax(dim=-1), None
412 - return compute_probs_and_sample_next_token(logits, sampling_metadata)
417 + return compute_probs_and_sample_next_token(
418 + logits, sampling_metadata, self.use_fp64_gumbel
419 + )
420
421 def _sample_draft_tokens(
422 self,
@@ -1656,6 +1663,7 @@ class SpecDecodeBaseProposer:
1663 def compute_probs_and_sample_next_token(
1664 logits: torch.Tensor,
1665 sampling_metadata: SamplingMetadata,
1666 + use_fp64_gumbel: bool = False,
1667 ) -> tuple[torch.Tensor, torch.Tensor]:
1668 if sampling_metadata.all_greedy:
1669 # For greedy requests, draft_probs is not used in rejection sampling.
@@ -1682,11 +1690,11 @@ def compute_probs_and_sample_next_token(
1690 # of the generated tokens after rejection sampling.
1691
1692 # TODO(woosuk): Consider seeds.
1685 - q = torch.empty_like(probs)
1693 + q = empty_exponential_noise_like(probs, use_fp64_gumbel)
1694 q.exponential_()
1695 # NOTE(woosuk): We shouldn't use `probs.div_(q)` because the draft_probs
1696 # will be used later for rejection sampling.
1689 - next_token_ids = probs.div(q).argmax(dim=-1).view(-1)
1697 + next_token_ids = sample_with_exponential_noise(probs.clone(), q)
1698 if not sampling_metadata.all_random:
1699 greedy_token_ids = probs.argmax(dim=-1)
1700 next_token_ids = torch.where(is_greedy, greedy_token_ids, next_token_ids)
vllm/v1/worker/gpu_model_runner.py
+4 -1
@@ -504,7 +504,10 @@ class GPUModelRunner(
504 self.use_async_scheduling = self.scheduler_config.async_scheduling
505
506 # Sampler
507 - self.sampler = Sampler(logprobs_mode=self.model_config.logprobs_mode)
507 + self.sampler = Sampler(
508 + logprobs_mode=self.model_config.logprobs_mode,
509 + use_fp64_gumbel=self.model_config.use_fp64_gumbel,
510 + )
511
512 self.eplb_state: EplbState | None = None
513 self._moe_model: MixtureOfExperts | None = None