@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 1134545b6f

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

Co-authored-by: vllm-agent <vllm-agent@users.noreply.github.com>

vllm-agent committed Jul 21, 2026 at 01:36 UTC 1134545b6f0064e5fa1459455d18b05d8ffaae1f

5 files changed +25 -85

tests/v1/sample/test_topk_topp_sampler.py

-38

@@ -5,13 +5,11 @@ 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

8 from vllm.platforms import current_platform

9 from vllm.triton_utils import HAS_TRITON

10 from vllm.utils.torch_utils import set_random_seed

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

12 apply_top_k_top_p_pytorch,

14 - flashinfer_sample,

13 random_sample,

14 )

15 from vllm.v1.sample.sampler import Sampler

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

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

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

1045 )

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

+4 -6

@@ -392,9 +392,8 @@ 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. 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)

395 + # Apply top-p.

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

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

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

399 # at least one

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

500 probs, k, deterministic=True

501 )

502 else:

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

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

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

504 next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits(

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

505 + logits, k, p, deterministic=True

506 )

507

508 return next_token_ids.view(-1)

vllm/v1/sample/ops/topk_topp_triton.py

+13 -19

@@ -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 - ).to(tl.float32)

134 + )

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 - ).to(tl.float32)

167 + )

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 - ).to(tl.float32)

308 + )

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 - ).to(tl.float32)

460 + )

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 - ).to(tl.float32)

603 + )

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 - ).to(tl.float32)

629 + )

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 - ).to(tl.float32)

663 + )

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 - ).to(tl.float32)

757 + )

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 - ).to(tl.float32)

838 + )

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 in (torch.float32, torch.bfloat16, torch.float16)

881 + assert logits.dtype == torch.float32

882 batch_size, vocab_size = logits.shape

883 topk_enabled = k is not None

884 topp_enabled = p is not None

@@ -911,9 +911,7 @@ 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 = torch.empty(

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

916 - )

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

915 _TRITON_BUFFER_CACHE[buf_key] = buffer

916 if buffer.shape[0] > NUM_PROGRAMS:

917 buffer = buffer[:NUM_PROGRAMS]

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

919 # Cache lookup table entries on each device.

920 tables = _TRITON_TABLE_CACHE.get(logits.device)

921 if tables is None:

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 - )

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 _TRITON_TABLE_CACHE[logits.device] = (

925 normal_cdf_to_sigma_table,

926 percentile_to_std_table,

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

+1 -3

@@ -218,9 +218,7 @@ 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(

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

223 - ).to(tl.float32)

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

222 logits += bias

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

224

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

+7 -19

@@ -83,7 +83,11 @@ 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 - return_logprobs = self.returns_logprobs(idx_mapping_np)

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

92 sampled, processed_logits = self.sample(

93 logits,

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

102 if return_logprobs:

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

104 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,13 +142,6 @@ 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 -

145 def apply_sampling_params(

146 self,

147 logits: torch.Tensor,

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

152 expanded_local_pos: torch.Tensor,

153 skip_top_k_top_p: bool = False,

154 ) -> torch.Tensor:

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()

155 + # Copy logits to a new FP32 tensor.

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

157

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

159 self.logit_bias_state.apply_logit_bias(