@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 425c4eafb0

[Sampler] Stop upcasting logits to fp32 in apply_sampling_params (#48641)

Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Michael Goin committed Jul 17, 2026 at 21:15 UTC 425c4eafb064f3804538a93584ec07c196d14cf9

5 files changed +85 -25

tests/v1/sample/test_topk_topp_sampler.py

+38

@@ -5,11 +5,13 @@ import torch

5 from torch import Generator

7 from tests.utils import large_gpu_mark

8 +from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size

9 from vllm.platforms import current_platform

10 from vllm.triton_utils import HAS_TRITON

11 from vllm.utils.torch_utils import set_random_seed

12 from vllm.v1.sample.ops.topk_topp_sampler import (

13 apply_top_k_top_p_pytorch,

14 + flashinfer_sample,

15 random_sample,

16 )

17 from vllm.v1.sample.sampler import Sampler

@@ -1043,3 +1045,39 @@ class TestFlashInferDistributionMatch:

1045 f"{label}: distribution differs from theoretical: "

1046 f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}"

1047 )

1048 +

1049 +

1050 +@pytest.mark.skipif(

1051 + not FLASHINFER_TOPK_TOPP_SUPPORTED,

1052 + reason="FlashInfer top-k/top-p sampler is not available on this platform.",

1053 +)

1054 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32])

1055 +@pytest.mark.parametrize("k, p", [(20, 0.95), (20, None), (None, 0.95)])

1056 +def test_flashinfer_sample_padded_vocab(

1057 + dtype: torch.dtype, k: int | None, p: float | None

1058 +):

1059 + """flashinfer_sample must accept the logits the sampler actually hands it.

1060 +

1061 + compute_logits slices the padding off the vocab, so for a vocab that isn't a

1062 + multiple of 64 (e.g. opt's 50272) the logits are a strided view in the model

1063 + dtype, while FlashInfer requires contiguous fp32.

1064 + """

1065 + torch.set_default_device(DEVICE_TYPE)

1066 + batch_size = 8

1067 + org_vocab_size = 50272

1068 + padded_vocab_size = pad_vocab_size(org_vocab_size)

1069 + assert padded_vocab_size != org_vocab_size

1070 +

1071 + logits = torch.randn(batch_size, padded_vocab_size, dtype=dtype)[

1072 + ..., :org_vocab_size

1073 + ]

1074 + # A single row stays contiguous despite the padded stride, hence batch_size > 1.

1075 + assert not logits.is_contiguous()

1076 +

1077 + token_ids = flashinfer_sample(

1078 + logits,

1079 + torch.full((batch_size,), k, dtype=torch.int32) if k is not None else None,

1080 + torch.full((batch_size,), p, dtype=torch.float32) if p is not None else None,

1081 + )

1082 + assert token_ids.shape == (batch_size,)

1083 + assert torch.all((token_ids >= 0) & (token_ids < org_vocab_size))

vllm/v1/sample/ops/topk_topp_sampler.py

+6 -4

@@ -392,8 +392,9 @@ def apply_top_k_top_p_pytorch(

392 logits_sort.masked_fill_(top_k_mask, -float("inf"))

393

394 if p is not None:

395 - # Apply top-p.

396 - probs_sort = logits_sort.softmax(dim=-1)

395 + # Apply top-p. The cumsum below runs over the whole vocab, so accumulating

396 + # in a low-precision dtype makes the nucleus undershoot p.

397 + probs_sort = logits_sort.softmax(dim=-1, dtype=torch.float32)

398 probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort)

399 top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)

400 # at least one

@@ -500,9 +501,10 @@ def flashinfer_sample(

501 probs, k, deterministic=True

502 )

503 else:

503 - # Both top-k and top-p.

504 + # Both top-k and top-p. FlashInfer requires contiguous fp32 logits; the

505 + # branches above get that from softmax().

506 next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(

505 - logits, k, p, deterministic=True

507 + logits.float().contiguous(), k, p, deterministic=True

508 )

509

510 return next_token_ids.view(-1)

vllm/v1/sample/ops/topk_topp_triton.py

+19 -13

@@ -131,7 +131,7 @@ def _topk_topp_kernel(

131 mask_n = offs < VOCAB_SIZE

132 logits_blk0 = tl.load(

133 LOGITS_ROW + offs, mask=mask_n, other=-float("inf")

134 - )

134 + ).to(tl.float32)

135 # Exclude -inf values (e.g. from grammar bitmasks) from

136 # statistics to avoid NaN in pivot computation.

137 finite_mask = (logits_blk0 > -float("inf")) & mask_n

@@ -164,7 +164,7 @@ def _topk_topp_kernel(

164 mask_n = offs_n < VOCAB_SIZE

165 logits_blk = tl.load(

166 LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")

167 - )

167 + ).to(tl.float32)

168

169 max_logit = tl.maximum(max_logit, tl.max(logits_blk))

170 # Exclude -inf from min to keep binary search bounds

@@ -305,7 +305,7 @@ def _topk_topp_kernel(

305 mask_n = offs_n < VOCAB_SIZE

306 logits_blk2 = tl.load(

307 LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")

308 - )

308 + ).to(tl.float32)

309

310 above_0 = logits_blk2 > k_pivot_0

311 above_1 = logits_blk2 > k_pivot_1

@@ -457,7 +457,7 @@ def _topk_topp_kernel(

457 LOGITS_ROW + offs_n,

458 mask=mask_n,

459 other=-float("inf"),

460 - )

460 + ).to(tl.float32)

461

462 outlier_mask = (probs_blk > min_logit) & mask_n

463

@@ -600,7 +600,7 @@ def _topk_topp_kernel(

600 mask_n = offs < VOCAB_SIZE

601 logits_blk0 = tl.load(

602 LOGITS_ROW + offs, mask=mask_n, other=-float("inf")

603 - )

603 + ).to(tl.float32)

604 # Exclude -inf values (e.g. from grammar bitmasks) from

605 # statistics to avoid NaN in pivot computation.

606 finite_mask = (logits_blk0 > -float("inf")) & mask_n

@@ -626,7 +626,7 @@ def _topk_topp_kernel(

626 mask_n = offs_n < VOCAB_SIZE

627 logits_blk = tl.load(

628 LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")

629 - )

629 + ).to(tl.float32)

630 max_logit = tl.maximum(max_logit, tl.max(logits_blk))

631 # Exclude -inf from min to keep binary search bounds

632 # finite (avoids NaN pivots).

@@ -660,7 +660,7 @@ def _topk_topp_kernel(

660

661 probs_blk = tl.load(

662 LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")

663 - )

663 + ).to(tl.float32)

664 probs_blk = tl.exp(probs_blk - max_sample)

665 probs_blk = probs_blk / sum_exp_logits

666

@@ -754,7 +754,7 @@ def _topk_topp_kernel(

754

755 probs_blk = tl.load(

756 LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")

757 - )

757 + ).to(tl.float32)

758 probs_blk = tl.exp(probs_blk - max_sample)

759 probs_blk = probs_blk / sum_exp_logits

760 tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n)

@@ -835,7 +835,7 @@ def _topk_topp_kernel(

835 mask_n = offs_n < VOCAB_SIZE

836 logits_blk = tl.load(

837 LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf")

838 - )

838 + ).to(tl.float32)

839 keep_mask = (logits_blk > final_pivot) & mask_n

840

841 # Duplicate logit handling

@@ -878,7 +878,7 @@ def apply_top_k_top_p_triton(

878 The masked logits tensor. It may or may not be modified in-place.

879 """

880 assert logits.ndim == 2

881 - assert logits.dtype == torch.float32

881 + assert logits.dtype in (torch.float32, torch.bfloat16, torch.float16)

882 batch_size, vocab_size = logits.shape

883 topk_enabled = k is not None

884 topp_enabled = p is not None

@@ -911,7 +911,9 @@ def apply_top_k_top_p_triton(

911 buffer = _TRITON_BUFFER_CACHE.get(buf_key)

912 if buffer is None or buffer.shape[0] < NUM_PROGRAMS:

913 size = min(next_power_of_2(NUM_PROGRAMS), num_sm)

914 - buffer = logits.new_empty((size, vocab_size))

914 + buffer = torch.empty(

915 + (size, vocab_size), dtype=torch.float32, device=logits.device

916 + )

917 _TRITON_BUFFER_CACHE[buf_key] = buffer

918 if buffer.shape[0] > NUM_PROGRAMS:

919 buffer = buffer[:NUM_PROGRAMS]

@@ -919,8 +921,12 @@ def apply_top_k_top_p_triton(

921 # Cache lookup table entries on each device.

922 tables = _TRITON_TABLE_CACHE.get(logits.device)

923 if tables is None:

922 - normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE)

923 - percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE)

924 + normal_cdf_to_sigma_table = torch.tensor(

925 + _NORMAL_CDF_TO_SIGMA_TABLE, dtype=torch.float32, device=logits.device

926 + )

927 + percentile_to_std_table = torch.tensor(

928 + _PERCENTILE_TO_STD_TABLE, dtype=torch.float32, device=logits.device

929 + )

930 _TRITON_TABLE_CACHE[logits.device] = (

931 normal_cdf_to_sigma_table,

932 percentile_to_std_table,

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

+3 -1

@@ -218,7 +218,9 @@ def _bias_kernel(

218 mask=mask,

219 )

220 bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask)

221 - logits = tl.load(logits_ptr + token_idx * logits_stride + token_ids, mask=mask)

221 + logits = tl.load(

222 + logits_ptr + token_idx * logits_stride + token_ids, mask=mask

223 + ).to(tl.float32)

224 logits += bias

225 tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask)

226

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

+19 -7

@@ -83,11 +83,7 @@ class Sampler:

83 # that num_nans is computed before applying penalties and temperature.

84 num_nans = get_num_nans(logits) if self.compute_nans else None

86 - max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)

87 - max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(

88 - idx_mapping_np

89 - )

90 - return_logprobs = max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0

86 + return_logprobs = self.returns_logprobs(idx_mapping_np)

88 sampled, processed_logits = self.sample(

89 logits,

@@ -102,6 +98,10 @@ class Sampler:

98 if return_logprobs:

99 if self.logprobs_mode in ("processed_logprobs", "processed_logits"):

100 logits = processed_logits

101 + max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)

102 + max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids(

103 + idx_mapping_np

104 + )

105 expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]

106 cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None

107 num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0

@@ -142,6 +142,13 @@ class Sampler:

142 )

143 return sampler_output

144

145 + def returns_logprobs(self, idx_mapping_np: np.ndarray) -> bool:

146 + """Whether any request in the batch produces logprobs this step."""

147 + return (

148 + self.sampling_states.max_num_logprobs(idx_mapping_np) != NO_LOGPROBS

149 + or self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np) > 0

150 + )

151 +

152 def apply_sampling_params(

153 self,

154 logits: torch.Tensor,

@@ -152,8 +159,13 @@ class Sampler:

159 expanded_local_pos: torch.Tensor,

160 skip_top_k_top_p: bool = False,

161 ) -> torch.Tensor:

155 - # Copy logits to a new FP32 tensor.

156 - logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits)

162 + # The ops below upcast to fp32 internally, so the input dtype is kept and

163 + # mutated in place. Only raw_logprobs reads the unmodified logits

164 + # afterward, so copy just for that case.

165 + if self.logprobs_mode.startswith("raw_") and self.returns_logprobs(

166 + idx_mapping_np

167 + ):

168 + logits = logits.clone()

169

170 # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place.

171 self.logit_bias_state.apply_logit_bias(