diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index bf9bea80bf..982582ffae 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -412,3 +412,65 @@ def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int) f"Block verification mean accepted length {mean_block:.4f} is worse " f"than standard {mean_standard:.4f}." ) + + +@pytest.mark.parametrize("has_draft_logits", [True, False]) +def test_chunked_requests_match_full_batch(has_draft_logits: bool): + torch.manual_seed(7) + device = "cuda" + num_reqs = 5 + num_speculative_steps = 3 + vocab_size = 257 + + target_logits = torch.randn(vocab_size, device=device) + draft_logits = torch.randn(vocab_size, device=device) + inputs = _build_rejection_sample_inputs( + target_logits, + draft_logits, + num_speculative_steps, + temperature=0.6, + num_trials=num_reqs, + ) + padded_target_logits = torch.empty( + inputs["target_logits"].shape[0], vocab_size + 3, device=device + ) + padded_target_logits[:, :vocab_size].copy_(inputs["target_logits"]) + inputs["target_logits"] = padded_target_logits[:, :vocab_size] + assert inputs["target_logits"].stride(-1) == 1 + assert not inputs["target_logits"].is_contiguous() + if not has_draft_logits: + inputs["draft_logits"] = None + + sampled, num_sampled = rejection_sample( + **inputs, num_speculative_steps=num_speculative_steps + ) + + sampled_chunks = [] + num_sampled_chunks = [] + for start, end in ((0, 2), (2, 5)): + lo = start * (num_speculative_steps + 1) + hi = end * (num_speculative_steps + 1) + chunk_inputs = dict(inputs) + for name in ( + "target_logits", + "draft_sampled", + "pos", + "expanded_idx_mapping", + "expanded_local_pos", + ): + chunk_inputs[name] = inputs[name][lo:hi] + chunk_inputs["cu_num_logits"] = inputs["cu_num_logits"][start : end + 1] - lo + chunk_inputs["idx_mapping"] = inputs["idx_mapping"][start:end] + + chunk_sampled, chunk_num_sampled = rejection_sample( + **chunk_inputs, num_speculative_steps=num_speculative_steps + ) + sampled_chunks.append(chunk_sampled) + num_sampled_chunks.append(chunk_num_sampled) + + chunked_sampled = torch.cat(sampled_chunks) + chunked_num_sampled = torch.cat(num_sampled_chunks) + assert torch.equal(chunked_num_sampled, num_sampled) + steps = torch.arange(num_speculative_steps + 1, device=device) + valid = steps.unsqueeze(0) < num_sampled.unsqueeze(1) + assert torch.equal(chunked_sampled[valid], sampled[valid]) diff --git a/tests/v1/test_outputs.py b/tests/v1/test_outputs.py index 89d551e344..4696eefa29 100644 --- a/tests/v1/test_outputs.py +++ b/tests/v1/test_outputs.py @@ -2,7 +2,32 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest import TestCase -from vllm.v1.outputs import LogprobsLists +import torch + +from vllm.v1.outputs import LogprobsLists, LogprobsTensors + + +def test_logprobs_tensors_cat(): + first = LogprobsTensors( + torch.tensor([[1, 2]]), + torch.tensor([[0.1, 0.2]]), + torch.tensor([1]), + ) + second = LogprobsTensors( + torch.tensor([[3, 4]]), + torch.tensor([[0.3, 0.4]]), + torch.tensor([2]), + ) + + result = LogprobsTensors.cat([first, second], [0, 1, 2]) + + assert result.logprob_token_ids.tolist() == [[1, 2], [3, 4]] + assert result.logprobs.tolist() == ( + first.logprobs.tolist() + second.logprobs.tolist() + ) + assert result.selected_token_ranks.tolist() == [1, 2] + assert result.cu_num_generated_tokens == [0, 1, 2] + assert LogprobsTensors.cat([first]) is first class TestLogprobsLists(TestCase): diff --git a/tests/v1/worker/test_gpu_rejection_sampler_chunking.py b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py new file mode 100644 index 0000000000..22cb71e67b --- /dev/null +++ b/tests/v1/worker/test_gpu_rejection_sampler_chunking.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import MethodType, SimpleNamespace +from typing import get_args + +import numpy as np +import pytest +import torch + +from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode +from vllm.platforms import current_platform +from vllm.v1.worker.gpu.spec_decode.rejection_sampler import ( + RejectionSampler, + _iter_request_chunks, +) + + +def test_iter_request_chunks_preserves_request_boundaries(): + cu_num_logits = np.array([0, 3, 4, 11, 13], dtype=np.int32) + + assert list(_iter_request_chunks(cu_num_logits, max_chunk_logits=5)) == [ + (0, 2), + (2, 3), + (3, 4), + ] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA") +@pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode)) +def test_chunked_scores_match_full_batch(logprobs_mode: str): + device = torch.device("cuda") + cu_num_logits_np = np.array([0, 3, 4, 8, 10], dtype=np.int32) + num_logits_per_req = np.diff(cu_num_logits_np) + idx_mapping_np = np.array([7, 2, 9, 1], dtype=np.int32) + input_batch = SimpleNamespace( + num_reqs=4, + cu_num_logits_np=cu_num_logits_np, + cu_num_logits=torch.from_numpy(cu_num_logits_np).to(device), + idx_mapping_np=idx_mapping_np, + idx_mapping=torch.from_numpy(idx_mapping_np).to(device), + expanded_idx_mapping=torch.from_numpy( + np.repeat(idx_mapping_np, num_logits_per_req) + ).to(device), + expanded_local_pos=torch.from_numpy( + np.concatenate( + [np.arange(count, dtype=np.int32) for count in num_logits_per_req] + ) + ).to(device), + ) + rejection_sampler = object.__new__(RejectionSampler) + rejection_sampler.sampler = SimpleNamespace(logprobs_mode=logprobs_mode) + rejection_sampler.num_speculative_steps = 3 + + def fake_verify( + self, + logits, + _draft_logits, + _draft_sampled, + _pos, + cu_num_logits, + idx_mapping, + *_mappings, + ): + num_sampled = torch.diff(cu_num_logits).to(torch.int32) + sampled = ( + idx_mapping.to(torch.int64).unsqueeze(1) + torch.arange(4, device=device) + ) % logits.shape[1] + return logits.float() + 1, sampled, num_sampled + + rejection_sampler._verify = MethodType(fake_verify, rejection_sampler) + logits = torch.arange(170, dtype=torch.float32, device=device).view(10, 17) + + sampled, num_sampled, chunked_logprobs = rejection_sampler._verify_in_chunks( + logits, + input_batch, + draft_logits=None, + draft_sampled=torch.arange(10, device=device), + pos=torch.arange(10, device=device), + max_chunk_logits=5, + max_num_logprobs=2, + ) + score_logits = logits + 1 if logprobs_mode in PROCESSED_LOGPROBS_MODES else logits + full_logprobs = rejection_sampler._get_logprobs_tensors( + sampled, + num_sampled, + score_logits, + input_batch.cu_num_logits, + input_batch.cu_num_logits_np, + max_num_logprobs=2, + ) + + assert sampled[:, 0].tolist() == idx_mapping_np.tolist() + assert num_sampled.tolist() == num_logits_per_req.tolist() + assert chunked_logprobs is not None + assert full_logprobs is not None + assert torch.equal( + chunked_logprobs.logprob_token_ids, + full_logprobs.logprob_token_ids, + ) + assert torch.equal(chunked_logprobs.logprobs, full_logprobs.logprobs) + assert torch.equal( + chunked_logprobs.selected_token_ranks, + full_logprobs.selected_token_ranks, + ) + assert ( + chunked_logprobs.cu_num_generated_tokens + == full_logprobs.cu_num_generated_tokens + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index 6b032ae762..d9d2f57dc4 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -90,6 +90,10 @@ ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"] LogprobsMode = Literal[ "raw_logits", "raw_logprobs", "processed_logits", "processed_logprobs" ] +PROCESSED_LOGPROBS_MODES: tuple[LogprobsMode, ...] = ( + "processed_logits", + "processed_logprobs", +) HfOverrides = dict[str, Any] | Callable[[PretrainedConfig], PretrainedConfig] ModelImpl = Literal["auto", "vllm", "transformers", "terratorch"] LayerBlockType = Literal["attention", "linear_attention", "mamba"] diff --git a/vllm/v1/outputs.py b/vllm/v1/outputs.py index 9f13ad939f..abb6a2be0d 100644 --- a/vllm/v1/outputs.py +++ b/vllm/v1/outputs.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod +from collections.abc import Sequence from copy import copy from dataclasses import dataclass, field from typing import TYPE_CHECKING, NamedTuple, TypeAlias @@ -90,6 +91,32 @@ class LogprobsTensors(NamedTuple): self.selected_token_ranks[mask], ) + @staticmethod + def cat( + tensors: Sequence["LogprobsTensors"], + cu_num_generated_tokens: list[int] | None = None, + ) -> "LogprobsTensors": + """Concatenate flattened logprob tensors.""" + assert tensors + assert cu_num_generated_tokens is not None or all( + tensor.cu_num_generated_tokens is None for tensor in tensors + ) + if len(tensors) == 1: + tensor = tensors[0] + if cu_num_generated_tokens is None: + return tensor + return tensor._replace(cu_num_generated_tokens=cu_num_generated_tokens) + return LogprobsTensors( + logprob_token_ids=torch.cat( + [tensor.logprob_token_ids for tensor in tensors] + ), + logprobs=torch.cat([tensor.logprobs for tensor in tensors]), + selected_token_ranks=torch.cat( + [tensor.selected_token_ranks for tensor in tensors] + ), + cu_num_generated_tokens=cu_num_generated_tokens, + ) + @staticmethod def empty_cpu( num_positions: int, num_tokens_per_position: int diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 69b35830ad..f608a0583d 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -7,7 +7,7 @@ import torch.nn as nn from vllm import envs from vllm._aiter_ops import rocm_aiter_ops -from vllm.config.model import LogprobsMode +from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode from vllm.logger import init_logger from vllm.platforms import CpuArchEnum, current_platform from vllm.triton_utils import HAS_TRITON @@ -87,7 +87,7 @@ class TopKTopPSampler(nn.Module): # FlashInfer doesn't expose post-top-k/top-p logits/logprobs, # so it can't be used when the configured mode requires them. can_use_flashinfer = ( - logprobs_mode not in ("processed_logits", "processed_logprobs") + logprobs_mode not in PROCESSED_LOGPROBS_MODES and flashinfer_sampler_supported() ) self.forward = ( @@ -108,7 +108,7 @@ class TopKTopPSampler(nn.Module): else: self.forward = self.forward_native elif ( - logprobs_mode not in ("processed_logits", "processed_logprobs") + logprobs_mode not in PROCESSED_LOGPROBS_MODES and rocm_aiter_ops.is_enabled() ): self.aiter_ops = None @@ -165,7 +165,7 @@ class TopKTopPSampler(nn.Module): return self.forward_native(logits, generators, k, p) if self.use_fp64_gumbel: return self.forward_native(logits, generators, k, p) - assert self.logprobs_mode not in ("processed_logits", "processed_logprobs"), ( + assert self.logprobs_mode not in PROCESSED_LOGPROBS_MODES, ( "FlashInfer does not support returning logits/logprobs" ) # flashinfer sampling functions expect contiguous logits. @@ -236,10 +236,9 @@ class TopKTopPSampler(nn.Module): return self.forward_native(logits, generators, k, p) if self.use_fp64_gumbel: return self.forward_native(logits, generators, k, p) - assert self.logprobs_mode not in ( - "processed_logits", - "processed_logprobs", - ), "aiter sampler does not support returning logits/logprobs." + assert self.logprobs_mode not in PROCESSED_LOGPROBS_MODES, ( + "aiter sampler does not support returning logits/logprobs." + ) if self.aiter_ops is None and not self._init_aiter_ops(): return self.forward_native(logits, generators, k, p) return self.aiter_sample(logits, k, p, generators), None @@ -300,10 +299,7 @@ class TopKTopPSampler(nn.Module): logits.shape[0], dtype=torch.int64, device=logits.device ) logits_to_return = None - if ( - self.logprobs_mode == "processed_logits" - or self.logprobs_mode == "processed_logprobs" - ): + if self.logprobs_mode in PROCESSED_LOGPROBS_MODES: logits_to_return = torch.empty_like(logits) assert len(generators) != logits.shape[0], ( diff --git a/vllm/v1/sample/rejection_sampler.py b/vllm/v1/sample/rejection_sampler.py index 0f56b8a4a2..a1baf0a808 100644 --- a/vllm/v1/sample/rejection_sampler.py +++ b/vllm/v1/sample/rejection_sampler.py @@ -10,6 +10,7 @@ from typing import TYPE_CHECKING import torch import torch.nn as nn +from vllm.config.model import PROCESSED_LOGPROBS_MODES from vllm.logger import init_logger from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsLists, LogprobsTensors, SamplerOutput @@ -67,10 +68,7 @@ class RejectionSampler(nn.Module): self.sampler = sampler self.use_fp64_gumbel = getattr(sampler, "use_fp64_gumbel", False) logprobs_mode = self.sampler.logprobs_mode - self.is_processed_logprobs_mode = logprobs_mode in ( - "processed_logprobs", - "processed_logits", - ) + self.is_processed_logprobs_mode = logprobs_mode in PROCESSED_LOGPROBS_MODES self.is_logits_logprobs_mode = logprobs_mode in ( "raw_logits", "processed_logits", diff --git a/vllm/v1/worker/gpu/sample/prompt_logprob.py b/vllm/v1/worker/gpu/sample/prompt_logprob.py index 4d4cc24482..50baf39ebb 100644 --- a/vllm/v1/worker/gpu/sample/prompt_logprob.py +++ b/vllm/v1/worker/gpu/sample/prompt_logprob.py @@ -132,15 +132,7 @@ class PromptLogprobsWorker: if prompt_logprobs_list: # Merge the in-progress logprobs. - logprobs = LogprobsTensors( - logprob_token_ids=torch.cat( - [x.logprob_token_ids for x in prompt_logprobs_list] - ), - logprobs=torch.cat([x.logprobs for x in prompt_logprobs_list]), - selected_token_ranks=torch.cat( - [x.selected_token_ranks for x in prompt_logprobs_list] - ), - ) + logprobs = LogprobsTensors.cat(prompt_logprobs_list) prompt_logprobs_list.clear() if logprobs is None: diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index f0a83c92ef..e34e2acf37 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -5,7 +5,7 @@ import numpy as np import torch import vllm.envs as envs -from vllm.config.model import LogprobsMode +from vllm.config.model import PROCESSED_LOGPROBS_MODES, LogprobsMode from vllm.sampling_params import SamplingParams from vllm.v1.sample.ops.topk_topp_sampler import ( apply_top_k_top_p, @@ -100,7 +100,7 @@ class Sampler: ) if return_logprobs: - if self.logprobs_mode in ("processed_logprobs", "processed_logits"): + if self.logprobs_mode in PROCESSED_LOGPROBS_MODES: logits = processed_logits expanded_logits = logits.shape[0] != idx_mapping_np.shape[0] cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None @@ -221,10 +221,7 @@ class Sampler: # any greedy requests or per-request seeds, or if post-processed # logprobs need to be returned for any requests. (top_k is None and top_p is None) - or ( - return_logprobs - and self.logprobs_mode in ("processed_logprobs", "processed_logits") - ) + or (return_logprobs and self.logprobs_mode in PROCESSED_LOGPROBS_MODES) or self.sampling_states.any_greedy(idx_mapping_np) or self.sampling_states.any_explicit_seed(idx_mapping_np) ) diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index 4753d28174..63d479efa0 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -1,8 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator + +import numpy as np import torch from vllm.config import SpeculativeConfig +from vllm.config.model import PROCESSED_LOGPROBS_MODES from vllm.triton_utils import tl, triton from vllm.v1.outputs import LogprobsTensors from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates @@ -19,6 +23,29 @@ from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( rejection_sample, ) +# Cap on the FP32 target-logits buffer materialized by apply_sampling_params. +# TODO(mgoin): Chunking is a workaround. The rejection kernels already upcast +# per vocab block on load and apply ops like temperature and gumbel, so folding +# sampling-param application into those kernels would remove this buffer and +# its traffic entirely. +MAX_CHUNK_BYTES = 2**30 # 1GB +_FP32_BYTES = 4 + + +def _iter_request_chunks( + cu_num_logits: np.ndarray, max_chunk_logits: int +) -> Iterator[tuple[int, int]]: + """Yield maximally packed request ranges without splitting requests.""" + assert max_chunk_logits > 0 + num_reqs = cu_num_logits.size - 1 + start = 0 + while start < num_reqs: + max_logit = int(cu_num_logits[start]) + max_chunk_logits + end = int(np.searchsorted(cu_num_logits, max_logit, side="right") - 1) + end = min(num_reqs, max(start + 1, end)) + yield start, end + start = end + @triton.jit def _flatten_sampled_kernel( @@ -66,18 +93,17 @@ class RejectionSampler: def _get_logprobs_tensors( self, - input_batch: InputBatch, sampled: torch.Tensor, num_sampled: torch.Tensor, logits: torch.Tensor, + cu_num_logits: torch.Tensor, + cu_num_logits_np: np.ndarray, + max_num_logprobs: int, ) -> LogprobsTensors | None: - max_num_logprobs = self.sampler.sampling_states.max_num_logprobs( - input_batch.idx_mapping_np - ) if max_num_logprobs == NO_LOGPROBS: return None - num_reqs = input_batch.cu_num_logits.shape[0] - 1 + num_reqs = cu_num_logits.shape[0] - 1 num_logits = logits.shape[0] flat_sampled = torch.zeros( num_logits, dtype=sampled.dtype, device=sampled.device @@ -87,19 +113,122 @@ class RejectionSampler: sampled, sampled.stride(0), num_sampled, - input_batch.cu_num_logits, + cu_num_logits, num_warps=1, ) - expanded_logits = num_logits != input_batch.idx_mapping.shape[0] + expanded_logits = num_logits != num_reqs return compute_topk_scores( logits, max_num_logprobs, flat_sampled, - input_batch.cu_num_logits_np.tolist() if expanded_logits else None, + cu_num_logits_np.tolist() if expanded_logits else None, logits_mode=self.sampler.logprobs_mode in ("raw_logits", "processed_logits"), ) + def _verify( + self, + logits: torch.Tensor, + draft_logits: torch.Tensor | None, + draft_sampled: torch.Tensor, + pos: torch.Tensor, + cu_num_logits: torch.Tensor, + idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + expanded_idx_mapping: torch.Tensor, + expanded_local_pos: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + processed_logits = self.sampler.apply_sampling_params( + logits, + expanded_idx_mapping, + idx_mapping_np, + pos, + draft_sampled, + expanded_local_pos, + ) + sampled, num_sampled = rejection_sample( + processed_logits, + draft_logits, + draft_sampled, + cu_num_logits, + pos, + idx_mapping, + expanded_idx_mapping, + expanded_local_pos, + self.sampler.sampling_states.temperature.gpu, + self.sampler.sampling_states.seeds.gpu, + self.num_speculative_steps, + self.synthetic_conditional_rates, + use_fp64=self.sampler.use_fp64_gumbel, + use_block_verification=self.use_block_verification, + ) + return processed_logits, sampled, num_sampled + + def _verify_in_chunks( + self, + logits: torch.Tensor, + input_batch: InputBatch, + draft_logits: torch.Tensor | None, + draft_sampled: torch.Tensor, + pos: torch.Tensor, + max_chunk_logits: int, + max_num_logprobs: int, + ) -> tuple[torch.Tensor, torch.Tensor, LogprobsTensors | None]: + cu_num_logits_np = input_batch.cu_num_logits_np + use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES + sampled_chunks: list[torch.Tensor] = [] + num_sampled_chunks: list[torch.Tensor] = [] + logprobs_chunks: list[LogprobsTensors] = [] + + for start, end in _iter_request_chunks(cu_num_logits_np, max_chunk_logits): + lo = int(cu_num_logits_np[start]) + hi = int(cu_num_logits_np[end]) + chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo + chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] - lo + # draft_logits uses persistent request-state indices and stays global. + processed_logits, sampled, num_sampled = self._verify( + logits[lo:hi], + draft_logits, + draft_sampled[lo:hi], + pos[lo:hi], + chunk_cu_num_logits, + input_batch.idx_mapping[start:end], + input_batch.idx_mapping_np[start:end], + input_batch.expanded_idx_mapping[lo:hi], + input_batch.expanded_local_pos[lo:hi], + ) + chunk_logprobs = self._get_logprobs_tensors( + sampled, + num_sampled, + processed_logits if use_processed_logits else logits[lo:hi], + chunk_cu_num_logits, + chunk_cu_num_logits_np, + max_num_logprobs, + ) + if chunk_logprobs is not None: + logprobs_chunks.append(chunk_logprobs) + del processed_logits + sampled_chunks.append(sampled) + num_sampled_chunks.append(num_sampled) + + if len(sampled_chunks) == 1: + logprobs_tensors = logprobs_chunks[0] if logprobs_chunks else None + return sampled_chunks[0], num_sampled_chunks[0], logprobs_tensors + + logprobs_tensors = None + if logprobs_chunks: + expanded_logits = logits.shape[0] != input_batch.num_reqs + logprobs_tensors = LogprobsTensors.cat( + logprobs_chunks, + cu_num_generated_tokens=( + cu_num_logits_np.tolist() if expanded_logits else None + ), + ) + + sampled = torch.cat(sampled_chunks) + num_sampled = torch.cat(num_sampled_chunks) + return sampled, num_sampled, logprobs_tensors + def __call__( self, logits: torch.Tensor, @@ -112,37 +241,19 @@ class RejectionSampler: draft_sampled = input_batch.input_ids[input_batch.logits_indices] pos = input_batch.positions[input_batch.logits_indices] - processed_logits = self.sampler.apply_sampling_params( - logits, - input_batch.expanded_idx_mapping, - input_batch.idx_mapping_np, - pos, - draft_sampled, - input_batch.expanded_local_pos, + + max_num_logprobs = self.sampler.sampling_states.max_num_logprobs( + input_batch.idx_mapping_np ) - sampled, num_sampled = rejection_sample( - processed_logits, + max_chunk_logits = max(1, MAX_CHUNK_BYTES // (logits.shape[1] * _FP32_BYTES)) + sampled, num_sampled, logprobs_tensors = self._verify_in_chunks( + logits, + input_batch, draft_logits, draft_sampled, - input_batch.cu_num_logits, pos, - input_batch.idx_mapping, - input_batch.expanded_idx_mapping, - input_batch.expanded_local_pos, - self.sampler.sampling_states.temperature.gpu, - self.sampler.sampling_states.seeds.gpu, - self.num_speculative_steps, - self.synthetic_conditional_rates, - use_fp64=self.sampler.use_fp64_gumbel, - use_block_verification=self.use_block_verification, - ) - logprobs_tensors = self._get_logprobs_tensors( - input_batch, - sampled, - num_sampled, - processed_logits - if self.sampler.logprobs_mode in ("processed_logprobs", "processed_logits") - else logits, + max_chunk_logits, + max_num_logprobs, ) num_sampled, num_rejected = get_num_sampled_and_rejected( diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index f700ea846c..0c59a1c65e 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -888,6 +888,10 @@ def rejection_sample( use_fp64: bool = False, use_block_verification: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: + assert target_logits.ndim == 2 and target_logits.stride(-1) == 1 + assert draft_logits is None or ( + draft_logits.ndim == 3 and draft_logits.stride(-1) == 1 + ) num_reqs = cu_num_logits.shape[0] - 1 num_logits, vocab_size = target_logits.shape draft_logits_stride_0 = 0 diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e8804721f8..f7646682e7 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -37,6 +37,7 @@ from vllm.config import ( update_config, ) from vllm.config.cache import CacheConfig +from vllm.config.model import PROCESSED_LOGPROBS_MODES from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.eplb.eplb_state import EplbState from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group @@ -6215,10 +6216,7 @@ class GPUModelRunner( # memory during profile_run. # No .clone() of logits: warmup output is discarded, so any in-place # mutation by forward_native does not affect correctness. - if self.sampler.logprobs_mode not in ( - "processed_logits", - "processed_logprobs", - ): + if self.sampler.logprobs_mode not in PROCESSED_LOGPROBS_MODES: self.sampler( logits=logits, sampling_metadata=replace(