From 1134545b6f0064e5fa1459455d18b05d8ffaae1f Mon Sep 17 00:00:00 2001 From: vllm-agent Date: Tue, 21 Jul 2026 01:36:45 -0700 Subject: [PATCH] Revert "[Sampler] Stop upcasting logits to fp32 in apply_sampling_params" (#48641) (#49033) Co-authored-by: vllm-agent --- tests/v1/sample/test_topk_topp_sampler.py | 38 ----------------------- vllm/v1/sample/ops/topk_topp_sampler.py | 10 +++--- vllm/v1/sample/ops/topk_topp_triton.py | 32 ++++++++----------- vllm/v1/worker/gpu/sample/logit_bias.py | 4 +-- vllm/v1/worker/gpu/sample/sampler.py | 26 +++++----------- 5 files changed, 25 insertions(+), 85 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 439ce0ea34..8a3d313f1d 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -5,13 +5,11 @@ import torch from torch import Generator from tests.utils import large_gpu_mark -from vllm.model_executor.layers.vocab_parallel_embedding import pad_vocab_size from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON from vllm.utils.torch_utils import set_random_seed from vllm.v1.sample.ops.topk_topp_sampler import ( apply_top_k_top_p_pytorch, - flashinfer_sample, random_sample, ) from vllm.v1.sample.sampler import Sampler @@ -1045,39 +1043,3 @@ class TestFlashInferDistributionMatch: f"{label}: distribution differs from theoretical: " f"chi2={chi2:.2f} p_value={p_value:.2e} alpha={self.ALPHA}" ) - - -@pytest.mark.skipif( - not FLASHINFER_TOPK_TOPP_SUPPORTED, - reason="FlashInfer top-k/top-p sampler is not available on this platform.", -) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("k, p", [(20, 0.95), (20, None), (None, 0.95)]) -def test_flashinfer_sample_padded_vocab( - dtype: torch.dtype, k: int | None, p: float | None -): - """flashinfer_sample must accept the logits the sampler actually hands it. - - compute_logits slices the padding off the vocab, so for a vocab that isn't a - multiple of 64 (e.g. opt's 50272) the logits are a strided view in the model - dtype, while FlashInfer requires contiguous fp32. - """ - torch.set_default_device(DEVICE_TYPE) - batch_size = 8 - org_vocab_size = 50272 - padded_vocab_size = pad_vocab_size(org_vocab_size) - assert padded_vocab_size != org_vocab_size - - logits = torch.randn(batch_size, padded_vocab_size, dtype=dtype)[ - ..., :org_vocab_size - ] - # A single row stays contiguous despite the padded stride, hence batch_size > 1. - assert not logits.is_contiguous() - - token_ids = flashinfer_sample( - logits, - torch.full((batch_size,), k, dtype=torch.int32) if k is not None else None, - torch.full((batch_size,), p, dtype=torch.float32) if p is not None else None, - ) - assert token_ids.shape == (batch_size,) - assert torch.all((token_ids >= 0) & (token_ids < org_vocab_size)) diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 0f3dc79b98..69b35830ad 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -392,9 +392,8 @@ def apply_top_k_top_p_pytorch( logits_sort.masked_fill_(top_k_mask, -float("inf")) if p is not None: - # Apply top-p. The cumsum below runs over the whole vocab, so accumulating - # in a low-precision dtype makes the nucleus undershoot p. - probs_sort = logits_sort.softmax(dim=-1, dtype=torch.float32) + # Apply top-p. + probs_sort = logits_sort.softmax(dim=-1) probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort) top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1) # at least one @@ -501,10 +500,9 @@ def flashinfer_sample( probs, k, deterministic=True ) else: - # Both top-k and top-p. FlashInfer requires contiguous fp32 logits; the - # branches above get that from softmax(). + # Both top-k and top-p. next_token_ids = flashinfer.sampling.top_k_top_p_sampling_from_logits( - logits.float().contiguous(), k, p, deterministic=True + logits, k, p, deterministic=True ) return next_token_ids.view(-1) diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 1b9058ba16..c284ff6187 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -131,7 +131,7 @@ def _topk_topp_kernel( mask_n = offs < VOCAB_SIZE logits_blk0 = tl.load( LOGITS_ROW + offs, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) # Exclude -inf values (e.g. from grammar bitmasks) from # statistics to avoid NaN in pivot computation. finite_mask = (logits_blk0 > -float("inf")) & mask_n @@ -164,7 +164,7 @@ def _topk_topp_kernel( mask_n = offs_n < VOCAB_SIZE logits_blk = tl.load( LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) max_logit = tl.maximum(max_logit, tl.max(logits_blk)) # Exclude -inf from min to keep binary search bounds @@ -305,7 +305,7 @@ def _topk_topp_kernel( mask_n = offs_n < VOCAB_SIZE logits_blk2 = tl.load( LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) above_0 = logits_blk2 > k_pivot_0 above_1 = logits_blk2 > k_pivot_1 @@ -457,7 +457,7 @@ def _topk_topp_kernel( LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf"), - ).to(tl.float32) + ) outlier_mask = (probs_blk > min_logit) & mask_n @@ -600,7 +600,7 @@ def _topk_topp_kernel( mask_n = offs < VOCAB_SIZE logits_blk0 = tl.load( LOGITS_ROW + offs, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) # Exclude -inf values (e.g. from grammar bitmasks) from # statistics to avoid NaN in pivot computation. finite_mask = (logits_blk0 > -float("inf")) & mask_n @@ -626,7 +626,7 @@ def _topk_topp_kernel( mask_n = offs_n < VOCAB_SIZE logits_blk = tl.load( LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) max_logit = tl.maximum(max_logit, tl.max(logits_blk)) # Exclude -inf from min to keep binary search bounds # finite (avoids NaN pivots). @@ -660,7 +660,7 @@ def _topk_topp_kernel( probs_blk = tl.load( LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) probs_blk = tl.exp(probs_blk - max_sample) probs_blk = probs_blk / sum_exp_logits @@ -754,7 +754,7 @@ def _topk_topp_kernel( probs_blk = tl.load( LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) probs_blk = tl.exp(probs_blk - max_sample) probs_blk = probs_blk / sum_exp_logits tl.store(BUFFER_ROW + offs_n, probs_blk, mask=mask_n) @@ -835,7 +835,7 @@ def _topk_topp_kernel( mask_n = offs_n < VOCAB_SIZE logits_blk = tl.load( LOGITS_ROW + offs_n, mask=mask_n, other=-float("inf") - ).to(tl.float32) + ) keep_mask = (logits_blk > final_pivot) & mask_n # Duplicate logit handling @@ -878,7 +878,7 @@ def apply_top_k_top_p_triton( The masked logits tensor. It may or may not be modified in-place. """ assert logits.ndim == 2 - assert logits.dtype in (torch.float32, torch.bfloat16, torch.float16) + assert logits.dtype == torch.float32 batch_size, vocab_size = logits.shape topk_enabled = k is not None topp_enabled = p is not None @@ -911,9 +911,7 @@ def apply_top_k_top_p_triton( buffer = _TRITON_BUFFER_CACHE.get(buf_key) if buffer is None or buffer.shape[0] < NUM_PROGRAMS: size = min(next_power_of_2(NUM_PROGRAMS), num_sm) - buffer = torch.empty( - (size, vocab_size), dtype=torch.float32, device=logits.device - ) + buffer = logits.new_empty((size, vocab_size)) _TRITON_BUFFER_CACHE[buf_key] = buffer if buffer.shape[0] > NUM_PROGRAMS: buffer = buffer[:NUM_PROGRAMS] @@ -921,12 +919,8 @@ def apply_top_k_top_p_triton( # Cache lookup table entries on each device. tables = _TRITON_TABLE_CACHE.get(logits.device) if tables is None: - normal_cdf_to_sigma_table = torch.tensor( - _NORMAL_CDF_TO_SIGMA_TABLE, dtype=torch.float32, device=logits.device - ) - percentile_to_std_table = torch.tensor( - _PERCENTILE_TO_STD_TABLE, dtype=torch.float32, device=logits.device - ) + normal_cdf_to_sigma_table = logits.new_tensor(_NORMAL_CDF_TO_SIGMA_TABLE) + percentile_to_std_table = logits.new_tensor(_PERCENTILE_TO_STD_TABLE) _TRITON_TABLE_CACHE[logits.device] = ( normal_cdf_to_sigma_table, percentile_to_std_table, diff --git a/vllm/v1/worker/gpu/sample/logit_bias.py b/vllm/v1/worker/gpu/sample/logit_bias.py index 1812251d47..6c95ed7aac 100644 --- a/vllm/v1/worker/gpu/sample/logit_bias.py +++ b/vllm/v1/worker/gpu/sample/logit_bias.py @@ -218,9 +218,7 @@ def _bias_kernel( mask=mask, ) bias = tl.load(bias_ptr + req_state_idx * bias_stride + block, mask=mask) - logits = tl.load( - logits_ptr + token_idx * logits_stride + token_ids, mask=mask - ).to(tl.float32) + logits = tl.load(logits_ptr + token_idx * logits_stride + token_ids, mask=mask) logits += bias tl.store(logits_ptr + token_idx * logits_stride + token_ids, logits, mask=mask) diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index d9d6754110..f0a83c92ef 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -83,7 +83,11 @@ class Sampler: # that num_nans is computed before applying penalties and temperature. num_nans = get_num_nans(logits) if self.compute_nans else None - return_logprobs = self.returns_logprobs(idx_mapping_np) + max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) + max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids( + idx_mapping_np + ) + return_logprobs = max_num_logprobs != NO_LOGPROBS or max_per_req_token_ids > 0 sampled, processed_logits = self.sample( logits, @@ -98,10 +102,6 @@ class Sampler: if return_logprobs: if self.logprobs_mode in ("processed_logprobs", "processed_logits"): logits = processed_logits - max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np) - max_per_req_token_ids = self.logprob_token_ids_state.max_num_token_ids( - idx_mapping_np - ) expanded_logits = logits.shape[0] != idx_mapping_np.shape[0] cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None num_logprobs = max_num_logprobs if max_num_logprobs != NO_LOGPROBS else 0 @@ -142,13 +142,6 @@ class Sampler: ) return sampler_output - def returns_logprobs(self, idx_mapping_np: np.ndarray) -> bool: - """Whether any request in the batch produces logprobs this step.""" - return ( - self.sampling_states.max_num_logprobs(idx_mapping_np) != NO_LOGPROBS - or self.logprob_token_ids_state.max_num_token_ids(idx_mapping_np) > 0 - ) - def apply_sampling_params( self, logits: torch.Tensor, @@ -159,13 +152,8 @@ class Sampler: expanded_local_pos: torch.Tensor, skip_top_k_top_p: bool = False, ) -> torch.Tensor: - # The ops below upcast to fp32 internally, so the input dtype is kept and - # mutated in place. Only raw_logprobs reads the unmodified logits - # afterward, so copy just for that case. - if self.logprobs_mode.startswith("raw_") and self.returns_logprobs( - idx_mapping_np - ): - logits = logits.clone() + # Copy logits to a new FP32 tensor. + logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits) # Apply logit bias (e.g., allowed_token_ids, min_tokens) in place. self.logit_bias_state.apply_logit_bias(