@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0162596603

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 3, 14:04 UTC

[Model Runner V2] FP32 gumbel sampling. (#41775)

Signed-off-by: PatchouliTaisa <patchychen@tencent.com> Co-authored-by: PatchouliTaisa <patchychen@tencent.com>

PatchyTIS committed May 16, 2026 at 00:20 UTC 0162596603c296d8ec5021806bf362debafa6b17

8 files changed +50 -7

vllm/config/model.py

@@ -233,6 +233,10 @@ class ModelConfig:

233 Processed means the values after applying all processors, including

234 temperature and top_k/top_p.

235 """

236 + use_fp64_gumbel: bool = False

237 + """Whether to use FP64 (instead of FP32) for the Gumbel noise used by the

238 + sampler. FP64 reduces the chance of ties in Gumbel-max sampling at the cost

239 + of significantly lower kernel throughput on most GPUs."""

240 disable_sliding_window: bool = False

241 """Whether to disable sliding window. If True, we will disable the sliding

242 window functionality of the model, capping to sliding window size. If the

@@ -365,6 +369,7 @@ class ModelConfig:

369 "spec_target_max_model_len",

370 "enforce_eager",

371 "logprobs_mode",

372 + "use_fp64_gumbel",

373 "disable_cascade_attn",

374 "skip_tokenizer_init",

375 "served_model_name",

vllm/engine/arg_utils.py

@@ -519,6 +519,7 @@ class EngineArgs:

519 max_num_seqs: int | None = None

520 max_logprobs: int = ModelConfig.max_logprobs

521 logprobs_mode: LogprobsMode = ModelConfig.logprobs_mode

522 + use_fp64_gumbel: bool = ModelConfig.use_fp64_gumbel

523 disable_log_stats: bool = False

524 aggregate_engine_logging: bool = False

525 revision: str | None = ModelConfig.revision

@@ -804,6 +805,7 @@ class EngineArgs:

805 )

806 model_group.add_argument("--max-logprobs", **model_kwargs["max_logprobs"])

807 model_group.add_argument("--logprobs-mode", **model_kwargs["logprobs_mode"])

808 + model_group.add_argument("--use-fp64-gumbel", **model_kwargs["use_fp64_gumbel"])

809 model_group.add_argument(

810 "--disable-sliding-window", **model_kwargs["disable_sliding_window"]

811 )

@@ -1542,6 +1544,7 @@ class EngineArgs:

1544 enable_return_routed_experts=self.enable_return_routed_experts,

1545 max_logprobs=self.max_logprobs,

1546 logprobs_mode=self.logprobs_mode,

1547 + use_fp64_gumbel=self.use_fp64_gumbel,

1548 disable_sliding_window=self.disable_sliding_window,

1549 disable_cascade_attn=self.disable_cascade_attn,

1550 skip_tokenizer_init=self.skip_tokenizer_init,

vllm/v1/worker/gpu/model_runner.py

@@ -218,6 +218,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

218 req_states=self.req_states,

219 logprobs_mode=self.model_config.logprobs_mode,

220 num_speculative_tokens=self.num_speculative_steps + 1,

221 + use_fp64_gumbel=self.model_config.use_fp64_gumbel,

222 )

223 if self.speculative_config is not None:

224 self.rejection_sampler = RejectionSampler(

vllm/v1/worker/gpu/sample/gumbel.py

+28 -6

@@ -2,7 +2,18 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import torch

5 -from vllm.triton_utils import tl, triton

5 +from vllm.triton_utils import HAS_TRITON, tl, triton

6 +

7 +# Smallest positive normal fp32 value. Used to clamp the uniform draw so that

8 +# `log(u)` cannot produce -inf (and thus `-log(-log(u))` stays finite).

9 +#

10 +# Triton requires globals accessed from `@triton.jit` functions to be wrapped

11 +# in `tl.constexpr(...)`. We can only do that when Triton is actually

12 +# available — on the CPU worker path `tl` is a placeholder whose `constexpr`

13 +# attribute is `None`, and `tl.constexpr(...)` would crash at import time.

14 +_FP32_TINY = (

15 + tl.constexpr(float.fromhex("0x1p-126")) if HAS_TRITON else float.fromhex("0x1p-126")

16 +)

19 @triton.jit

@@ -79,6 +90,7 @@ def gumbel_block_argmax(

90 processed_logits_col_ptr,

91 vocab_size,

92 APPLY_TEMPERATURE: tl.constexpr,

93 + USE_FP64: tl.constexpr,

94 ):

95 req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx)

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

@@ -103,16 +115,21 @@ def gumbel_block_argmax(

115 mask=mask,

116 )

117

106 - logits = logits.to(tl.float64)

118 + # fp32 is the default reduction dtype; fp64 is ~1/32–1/64x the throughput

119 + # on H100/Ada/Blackwell and empirically indistinguishable for Gumbel-max.

120 + if USE_FP64:

121 + logits = logits.to(tl.float64)

122 if temp != 0.0:

123 # Calculate the seed for gumbel noise.

124 seed = tl.load(seeds_ptr + req_state_idx)

125 pos = tl.load(pos_ptr + token_idx)

126 gumbel_seed = tl.randint(seed, pos)

127

113 - # tl.rand returns fp32, so build a true fp64 uniform from 64 random

114 - # bits before applying the double-log transform.

115 - u = tl_rand64(gumbel_seed, block, includes_zero=False)

128 + if USE_FP64:

129 + u = tl_rand64(gumbel_seed, block, includes_zero=False)

130 + else:

131 + u = tl.rand(gumbel_seed, block)

132 + u = tl.maximum(u, _FP32_TINY)

133 gumbel_noise = -tl.log(-tl.log(u))

134

135 # Apply gumbel noise.

@@ -140,6 +157,7 @@ def _gumbel_sample_kernel(

157 vocab_size,

158 BLOCK_SIZE: tl.constexpr,

159 APPLY_TEMPERATURE: tl.constexpr,

160 + USE_FP64: tl.constexpr,

161 ):

162 token_idx = tl.program_id(0)

163 block_idx = tl.program_id(1)

@@ -166,6 +184,7 @@ def _gumbel_sample_kernel(

184 processed_logits_col_ptr,

185 vocab_size,

186 APPLY_TEMPERATURE=APPLY_TEMPERATURE,

187 + USE_FP64=USE_FP64,

188 )

189 token_id = block_idx * BLOCK_SIZE + idx

190 tl.store(local_argmax_ptr + token_idx * local_argmax_stride + block_idx, token_id)

@@ -181,12 +200,14 @@ def gumbel_sample(

200 apply_temperature: bool,

201 output_processed_logits: torch.Tensor | None = None,

202 output_processed_logits_col: torch.Tensor | None = None,

203 + use_fp64: bool = False,

204 ) -> torch.Tensor:

205 num_tokens, vocab_size = logits.shape

206 BLOCK_SIZE = 1024

207 num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)

208 local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)

189 - local_max = logits.new_empty(num_tokens, num_blocks, dtype=torch.float64)

209 + local_max_dtype = torch.float64 if use_fp64 else torch.float32

210 + local_max = logits.new_empty(num_tokens, num_blocks, dtype=local_max_dtype)

211 _gumbel_sample_kernel[(num_tokens, num_blocks)](

212 local_argmax,

213 local_argmax.stride(0),

@@ -204,6 +225,7 @@ def gumbel_sample(

225 vocab_size,

226 BLOCK_SIZE=BLOCK_SIZE,

227 APPLY_TEMPERATURE=apply_temperature,

228 + USE_FP64=use_fp64,

229 )

230 # NOTE(woosuk): Use int64 for later indexing.

231 max_block_idx = local_max.argmax(dim=-1, keepdim=True)

vllm/v1/worker/gpu/sample/sampler.py

@@ -31,11 +31,13 @@ class Sampler:

31 req_states: RequestState,

32 logprobs_mode: LogprobsMode = "raw_logprobs",

33 num_speculative_tokens: int = 1,

34 + use_fp64_gumbel: bool = False,

35 ):

36 if logprobs_mode not in ("processed_logprobs", "raw_logprobs"):

37 raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")

38 self.logprobs_mode = logprobs_mode

39 self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.

40 + self.use_fp64_gumbel = use_fp64_gumbel

42 self.sampling_states = SamplingStates(max_num_reqs, vocab_size)

43 self.penalties_state = PenaltiesState(req_states)

@@ -192,5 +194,6 @@ class Sampler:

194 self.sampling_states.seeds.gpu,

195 pos,

196 apply_temperature=False,

197 + use_fp64=self.use_fp64_gumbel,

198 )

199 return sampled, processed_logits

vllm/v1/worker/gpu/spec_decode/eagle/speculator.py

@@ -63,6 +63,7 @@ class EagleSpeculator:

63 self.hidden_size = self.hidden_size * hc_mult

64 self.vocab_size = self.draft_model_config.get_vocab_size()

65 self.dtype = vllm_config.model_config.dtype

66 + self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel

68 # DP configuration

69 self.dp_size = vllm_config.parallel_config.data_parallel_size

@@ -247,6 +248,7 @@ class EagleSpeculator:

248 apply_temperature=True,

249 output_processed_logits=draft_logits,

250 output_processed_logits_col=draft_step,

251 + use_fp64=self.use_fp64_gumbel,

252 )

253 else:

254 return logits.argmax(dim=-1)

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

@@ -125,6 +125,7 @@ class RejectionSampler:

125 self.sampler.sampling_states.seeds.gpu,

126 self.num_speculative_steps,

127 self.synthetic_conditional_rates,

128 + use_fp64=self.sampler.use_fp64_gumbel,

129 )

130 logprobs_tensors = self._get_logprobs_tensors(

131 input_batch,

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

+7 -1

@@ -339,6 +339,7 @@ def _resample_kernel(

339 vocab_size,

340 BLOCK_SIZE: tl.constexpr,

341 HAS_DRAFT_LOGITS: tl.constexpr,

342 + USE_FP64: tl.constexpr,

343 ):

344 req_idx = tl.program_id(0)

345 resample_idx = tl.load(rejected_step_ptr + req_idx)

@@ -415,6 +416,7 @@ def _resample_kernel(

416 None, # processed_logits_col_ptr

417 vocab_size,

418 APPLY_TEMPERATURE=False,

419 + USE_FP64=USE_FP64,

420 )

421 token_id = block_idx * BLOCK_SIZE + idx

422 tl.store(

@@ -512,6 +514,7 @@ def rejection_sample(

514 num_speculative_steps: int,

515 # [num_speculative_steps]

516 synthetic_conditional_rates: torch.Tensor | None = None,

517 + use_fp64: bool = False,

518 ) -> tuple[torch.Tensor, torch.Tensor]:

519 num_reqs = cu_num_logits.shape[0] - 1

520 num_logits, vocab_size = target_logits.shape

@@ -620,7 +623,9 @@ def rejection_sample(

623 num_reqs, resample_num_blocks, dtype=torch.int64

624 )

625 resampled_local_max = target_logits.new_empty(

623 - num_reqs, resample_num_blocks, dtype=torch.float64

626 + num_reqs,

627 + resample_num_blocks,

628 + dtype=torch.float64 if use_fp64 else torch.float32,

629 )

630 _resample_kernel[(num_reqs, resample_num_blocks)](

631 resampled_local_argmax,

@@ -644,6 +649,7 @@ def rejection_sample(

649 vocab_size,

650 BLOCK_SIZE=RESAMPLE_BLOCK_SIZE,

651 HAS_DRAFT_LOGITS=has_draft_logits,

652 + USE_FP64=use_fp64,

653 )

654

655 # Insert the resampled tokens into the output sampled.