iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: vllm-agent <vllm-agent@users.noreply.github.com>
This commit is contained in:
vllm-agent
2026-07-21 09:36:45 +01:00
committed by GitHub
co-authored by vllm-agent
parent 3e0c887511
commit 1134545b6f
5 changed files with 25 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/sample/test_topk_topp_sampler.py
-38
View File
@@ -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))
vllm/v1/sample/ops/topk_topp_sampler.py
+4 -6
View File
@@ -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)
vllm/v1/sample/ops/topk_topp_triton.py
+13 -19
View File
@@ -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,
vllm/v1/worker/gpu/sample/logit_bias.py
+1 -3
View File
@@ -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)
vllm/v1/worker/gpu/sample/sampler.py
+7 -19
View File
@@ -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(