[MRV2][Spec Decode] Avoid rejection sampler OOM by chunking (#48630)

Signed-off-by: mgoin <[email protected]>
Signed-off-by: Michael Goin <[email protected]>
Co-authored-by: Nick Hill <[email protected]>

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

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

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

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"]

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

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], (

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",

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:

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

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(

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

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(