[Performance][MRV2] Cache logits-processing request state (#52329)
Signed-off-by: Rui Rui4 Cheng <chengrui4@lenovo.com> Co-authored-by: Rui Rui4 Cheng <chengrui4@lenovo.com>
This commit is contained in:
Cheng Rui
co-authored by
Rui Rui4 Cheng
Rui Rui4 Cheng
parent
502af5ed00
commit
6664d397bf
@@ -0,0 +1,90 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytest.importorskip("triton")
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required for sampler flag tests", allow_module_level=True)
|
||||
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.worker.gpu.sample.sampler import Sampler
|
||||
from vllm.v1.worker.gpu.states import RequestState
|
||||
|
||||
DEVICE = torch.device("cuda")
|
||||
VOCAB_SIZE = 128
|
||||
|
||||
|
||||
class MockReasoningConfig:
|
||||
reasoning_start_token_ids = [90]
|
||||
reasoning_end_token_ids = [91]
|
||||
natural_reasoning_end_token_ids = [91]
|
||||
|
||||
|
||||
def _make_sampler() -> Sampler:
|
||||
req_states = RequestState(
|
||||
max_num_reqs=4,
|
||||
max_model_len=64,
|
||||
max_num_batched_tokens=16,
|
||||
num_speculative_steps=1,
|
||||
vocab_size=VOCAB_SIZE,
|
||||
device=DEVICE,
|
||||
)
|
||||
return Sampler(
|
||||
max_num_reqs=4,
|
||||
vocab_size=VOCAB_SIZE,
|
||||
device=DEVICE,
|
||||
req_states=req_states,
|
||||
reasoning_config=MockReasoningConfig(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("sampling_params", "expected"),
|
||||
[
|
||||
pytest.param(SamplingParams(), False, id="defaults"),
|
||||
pytest.param(SamplingParams(temperature=0.0), False, id="greedy"),
|
||||
pytest.param(
|
||||
SamplingParams(thinking_token_budget=3), True, id="thinking-budget"
|
||||
),
|
||||
pytest.param(SamplingParams(logit_bias={1: 1.0}), True, id="logit-bias"),
|
||||
pytest.param(SamplingParams(frequency_penalty=0.1), True, id="penalty"),
|
||||
pytest.param(SamplingParams(_bad_words_token_ids=[[1]]), True, id="bad-words"),
|
||||
pytest.param(SamplingParams(temperature=0.7), True, id="temperature"),
|
||||
pytest.param(SamplingParams(min_p=0.1), True, id="min-p"),
|
||||
pytest.param(SamplingParams(top_k=10), True, id="top-k"),
|
||||
pytest.param(SamplingParams(top_p=0.9), True, id="top-p"),
|
||||
pytest.param(
|
||||
SamplingParams.for_sampler_warmup(), True, id="all-logits-processors"
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_logits_processing_cache_matches_request_features(
|
||||
sampling_params: SamplingParams, expected: bool
|
||||
):
|
||||
sampler = _make_sampler()
|
||||
sampler.add_request(3, prompt_len=1, sampling_params=sampling_params)
|
||||
|
||||
assert sampler.needs_logits_processing[3] == expected
|
||||
|
||||
|
||||
def test_logits_processing_cache_is_overwritten_when_slot_is_reused():
|
||||
sampler = _make_sampler()
|
||||
sampler.add_request(3, 1, SamplingParams.for_sampler_warmup())
|
||||
sampler.add_request(3, 1, SamplingParams())
|
||||
|
||||
assert not sampler.needs_logits_processing[3]
|
||||
|
||||
|
||||
def test_logits_processing_cache_only_checks_active_requests():
|
||||
sampler = _make_sampler()
|
||||
sampler.add_request(0, 1, SamplingParams(temperature=0.0))
|
||||
sampler.add_request(2, 1, SamplingParams.for_sampler_warmup())
|
||||
|
||||
sampling_only = np.array([0], dtype=np.int32)
|
||||
with_processing = np.array([0, 2], dtype=np.int32)
|
||||
|
||||
assert not np.any(sampler.needs_logits_processing[sampling_only])
|
||||
assert np.any(sampler.needs_logits_processing[with_processing])
|
||||
@@ -12,6 +12,7 @@ if not torch.cuda.is_available():
|
||||
)
|
||||
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.worker.gpu.sample.sampler import Sampler
|
||||
from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
|
||||
from vllm.v1.worker.gpu.states import RequestState
|
||||
|
||||
@@ -185,6 +186,44 @@ def test_v2_thinking_budget_ignores_plain_request():
|
||||
assert torch.all(out == 0)
|
||||
|
||||
|
||||
def test_v2_greedy_sampling_applies_thinking_budget():
|
||||
"""Greedy-only requests must not bypass thinking-budget processing."""
|
||||
req_states = _make_req_states([1, START, 10, 11, 12], prompt_len=1)
|
||||
sampler = Sampler(
|
||||
max_num_reqs=4,
|
||||
vocab_size=VOCAB_SIZE,
|
||||
device=DEVICE,
|
||||
req_states=req_states,
|
||||
reasoning_config=MockReasoningConfig(),
|
||||
)
|
||||
sampler.add_request(
|
||||
req_idx=3,
|
||||
prompt_len=1,
|
||||
sampling_params=SamplingParams(
|
||||
temperature=0.0,
|
||||
thinking_token_budget=3,
|
||||
),
|
||||
)
|
||||
sampler.apply_staged_writes()
|
||||
|
||||
idx_mapping = torch.tensor([3], dtype=torch.int32, device=DEVICE)
|
||||
idx_mapping_np = idx_mapping.cpu().numpy()
|
||||
expanded_idx_mapping = idx_mapping.clone()
|
||||
input_ids = torch.tensor([12], dtype=torch.int32, device=DEVICE)
|
||||
logits = torch.zeros((1, VOCAB_SIZE), device=DEVICE)
|
||||
out = sampler.apply_sampling_params(
|
||||
logits,
|
||||
expanded_idx_mapping,
|
||||
idx_mapping,
|
||||
idx_mapping_np,
|
||||
torch.tensor([4], dtype=torch.int32, device=DEVICE),
|
||||
input_ids,
|
||||
torch.tensor([0], dtype=torch.int32, device=DEVICE),
|
||||
)
|
||||
|
||||
assert out[0, END].item() == pytest.approx(1.0e9)
|
||||
|
||||
|
||||
def test_v2_thinking_budget_latest_prefill_end_disables_forcing():
|
||||
req_states = _make_req_states(
|
||||
[1, START, 10, 11, 12, END, 13],
|
||||
|
||||
@@ -53,6 +53,7 @@ class Sampler:
|
||||
self.bad_words_state = BadWordsState(req_states)
|
||||
self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
|
||||
self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config)
|
||||
self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool)
|
||||
self.num_speculative_tokens = num_speculative_tokens
|
||||
self.return_sampling_mask = return_sampling_mask
|
||||
self.use_flashinfer = (
|
||||
@@ -69,6 +70,22 @@ class Sampler:
|
||||
self.logprob_token_ids_state.add_request(req_idx, sampling_params)
|
||||
self.thinking_budget_state.add_request(req_idx, sampling_params)
|
||||
|
||||
states = self.sampling_states
|
||||
temperature = states.temperature.np[req_idx]
|
||||
self.needs_logits_processing[req_idx] = (
|
||||
self.logit_bias_state.use_logit_bias[req_idx]
|
||||
or self.penalties_state.use_penalty[req_idx]
|
||||
or self.bad_words_state.num_bad_words.np[req_idx] > 0
|
||||
or (
|
||||
self.thinking_budget_state.enabled
|
||||
and self.thinking_budget_state.use_thinking_budget[req_idx]
|
||||
)
|
||||
or (temperature != 0.0 and temperature != 1.0)
|
||||
or states.min_p.np[req_idx] != 0.0
|
||||
or states.top_k.np[req_idx] != states.vocab_size
|
||||
or states.top_p.np[req_idx] != 1.0
|
||||
)
|
||||
|
||||
def apply_staged_writes(self) -> None:
|
||||
self.sampling_states.apply_staged_writes()
|
||||
self.penalties_state.apply_staged_writes()
|
||||
@@ -172,7 +189,7 @@ class Sampler:
|
||||
expanded_local_pos: torch.Tensor,
|
||||
skip_top_k_top_p: bool = False,
|
||||
) -> torch.Tensor:
|
||||
if not self._requires_logits_processing(idx_mapping_np):
|
||||
if not np.any(self.needs_logits_processing[idx_mapping_np]):
|
||||
return logits
|
||||
|
||||
# Copy logits to a new FP32 tensor.
|
||||
@@ -228,24 +245,6 @@ class Sampler:
|
||||
logits, expanded_idx_mapping, idx_mapping_np
|
||||
)
|
||||
|
||||
def _requires_logits_processing(self, idx_mapping_np: np.ndarray) -> bool:
|
||||
if np.any(self.logit_bias_state.use_logit_bias[idx_mapping_np]):
|
||||
return True
|
||||
if np.any(self.penalties_state.use_penalty[idx_mapping_np]):
|
||||
return True
|
||||
if np.any(self.bad_words_state.num_bad_words.np[idx_mapping_np] > 0):
|
||||
return True
|
||||
|
||||
states = self.sampling_states
|
||||
temperatures = states.temperature.np[idx_mapping_np]
|
||||
if np.any((temperatures != 0.0) & (temperatures != 1.0)):
|
||||
return True
|
||||
if np.any(states.min_p.np[idx_mapping_np] != 0.0):
|
||||
return True
|
||||
if np.any(states.top_k.np[idx_mapping_np] != states.vocab_size):
|
||||
return True
|
||||
return bool(np.any(states.top_p.np[idx_mapping_np] != 1.0))
|
||||
|
||||
def sample(
|
||||
self,
|
||||
logits: torch.Tensor,
|
||||
|
||||
Reference in New Issue
Block a user