@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 30172b4947

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 23:30 UTC

[V1] Optimize handling of sampling metadata and req_ids list (#13244)

Signed-off-by: Nick Hill <nhill@redhat.com>

Nick Hill committed Feb 18, 2025 at 12:15 UTC 30172b4947c52890b808c6da3a6c7580f55cbb74

15 files changed +255 -298

tests/v1/sample/test_rejection_sampler.py

+2 -7

@@ -26,17 +26,13 @@ def create_logits_tensor(token_ids: List[int],

26 def create_sampling_metadata(spec_tokens: List[List[int]]) -> SamplingMetadata:

27 batch_size = len(spec_tokens)

28 return SamplingMetadata(

29 - temperature=0.0,

29 + temperature=torch.tensor([]),

30 all_greedy=True,

31 all_random=False,

32 - rejection_sampling=True,

32 spec_token_ids=spec_tokens,

33 top_p=None,

34 top_k=None,

36 - no_top_p=False,

37 - no_top_k=False,

35 min_p=torch.empty(batch_size, ),

39 - no_min_p=True,

36 generators={},

37 max_num_logprobs=0,

38 no_penalties=False,

@@ -45,8 +41,7 @@ def create_sampling_metadata(spec_tokens: List[List[int]]) -> SamplingMetadata:

41 presence_penalties=torch.tensor([]),

42 repetition_penalties=torch.tensor([]),

43 output_token_ids=[],

48 - min_tokens=[],

49 - stop_token_ids=[],

44 + min_tokens={},

45 logit_bias=[None] * batch_size,

46 )

tests/v1/sample/test_sampler.py

+18 -26

@@ -77,25 +77,20 @@ def _create_default_sampling_metadata(

77 temperature=torch.full((batch_size, ), 0.0),

78 all_greedy=True,

79 all_random=False,

80 - rejection_sampling=False,

81 - top_p=torch.empty(batch_size, ),

82 - top_k=torch.empty(batch_size, ),

83 - no_top_p=True,

84 - no_top_k=True,

85 - min_p=torch.empty(batch_size, ),

86 - no_min_p=True,

80 + top_p=None,

81 + top_k=None,

82 + min_p=None,

83 generators={},

84 max_num_logprobs=0,

85 prompt_token_ids=_create_prompt_tokens_tensor(prompt_token_ids,

86 vocab_size, device),

87 output_token_ids=output_token_ids,

92 - spec_token_ids=[],

88 + spec_token_ids=None,

89 frequency_penalties=_create_penalty_tensor(batch_size, 0.0, device),

90 presence_penalties=_create_penalty_tensor(batch_size, 0.0, device),

91 repetition_penalties=_create_penalty_tensor(batch_size, 1.0, device),

92 no_penalties=True,

97 - min_tokens=[],

98 - stop_token_ids=[],

93 + min_tokens={},

94 logit_bias=[None] * batch_size,

95 )

96 return fake_sampling_metadata

@@ -104,10 +99,10 @@ def _create_default_sampling_metadata(

99 def _generate_min_token_penalties_and_stop_tokens(

100 num_output_tokens: int, batch_size: int, vocab_size: int,

101 batch_indices_for_min_token_penalty: List[int]

107 -) -> Tuple[List[int], List[Set[int]]]:

102 +) -> Dict[int, Tuple[int, Set[int]]]:

103 """

109 - Generates and returns a list of minimum token penalties (`min_tokens`)

110 - and a corresponding list of stop token IDs (`stop_token_ids`) for each

104 + Generates and returns a dict of minimum token penalties and

105 + corresponding stop token IDs (`min_tokens`, `stop_token_ids`) for each

106 batch.

107

108 If a batch index is included in `batch_indices_for_min_token_penalty`,

@@ -115,22 +110,19 @@ def _generate_min_token_penalties_and_stop_tokens(

110 and a random set of stop token IDs is created. Otherwise, a lower

111 `min_tokens` value is assigned, and the stop token IDs set is empty.

112 """

118 - stop_token_ids: List[Set[int]] = []

119 - min_tokens: List[int] = []

113 + min_tokens: Dict[int, Tuple[int, Set[int]]] = {}

114 for index in range(batch_size):

115 if index in batch_indices_for_min_token_penalty:

122 - min_tokens.append(

116 + min_tokens[index] = (

117 np.random.randint(num_output_tokens + 1,

124 - 2 * num_output_tokens))

125 - stop_token_ids.append(

118 + 2 * num_output_tokens),

119 set(

120 np.random.randint(0, vocab_size - 1)

121 for _ in range(np.random.randint(0, vocab_size))))

129 -

122 else:

131 - min_tokens.append(np.random.randint(0, num_output_tokens))

132 - stop_token_ids.append(set())

133 - return (min_tokens, stop_token_ids)

123 + min_tokens[index] = (np.random.randint(0,

124 + num_output_tokens), set())

125 + return min_tokens

126

127

128 def _create_weighted_output_token_list(

@@ -165,7 +157,7 @@ def _create_weighted_output_token_list(

157 output_token_ids_for_batch.extend(

158 [token_id for _ in range(index + 1)])

159 output_token_ids.append(output_token_ids_for_batch)

168 - return (output_token_ids, sorted_token_ids_in_output)

160 + return output_token_ids, sorted_token_ids_in_output

161

162

163 @pytest.mark.parametrize("device", CUDA_DEVICES)

@@ -182,17 +174,17 @@ def test_sampler_min_tokens_penalty(device: str, batch_size: int):

174 NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))

175 batch_indices_for_min_token_penalty = np.random.randint(

176 0, batch_size - 1, size=np.random.randint(0, batch_size)).tolist()

185 - min_tokens, stop_token_ids = _generate_min_token_penalties_and_stop_tokens(

177 + min_tokens = _generate_min_token_penalties_and_stop_tokens(

178 NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE,

179 batch_indices_for_min_token_penalty)

180 sampling_metadata.min_tokens = min_tokens

189 - sampling_metadata.stop_token_ids = stop_token_ids

181 sampler = Sampler()

182 logits = sampler.apply_penalties(fake_logits, sampling_metadata)

183 logits = logits.cpu()

184 for batch_idx in range(batch_size):

185 for token_id in range(VOCAB_SIZE):

195 - if token_id in stop_token_ids[batch_idx]:

186 + _, stop_token_ids = min_tokens.get(batch_idx, (0, set()))

187 + if token_id in stop_token_ids:

188 assert logits[batch_idx][token_id] == -float("inf")

189 else:

190 assert logits[batch_idx][token_id] != -float("inf")

tests/v1/worker/test_gpu_input_batch.py

+21 -26

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Dict, List, Set, Tuple

3 +from typing import Dict, List, Optional, Set, Tuple

5 import numpy as np

6 import pytest

@@ -41,7 +41,7 @@ def _remove_requests(

41 for index in req_indices_to_remove:

42 input_batch.remove_request(reqs[index].req_id)

43 req_ids_to_remove.add(reqs[index].req_id)

44 - return (req_ids_to_remove, req_indices_to_remove_list)

44 + return req_ids_to_remove, req_indices_to_remove_list

47 def _construct_expected_sampling_metadata(

@@ -64,8 +64,7 @@ def _construct_expected_sampling_metadata(

64 top_p = [0.0 for _ in range(num_reqs)]

65 min_p = [0.0 for _ in range(num_reqs)]

66 temperature = [0.0 for _ in range(num_reqs)]

67 - stop_token_ids: List[Set[int]] = [set() for _ in range(num_reqs)]

68 - min_tokens = [0 for _ in range(num_reqs)]

67 + min_tokens = {}

68 logit_bias = [None] * num_reqs

69 for req in reqs:

70 if req.req_id not in req_ids_retained:

@@ -83,22 +82,21 @@ def _construct_expected_sampling_metadata(

82 top_p[index_in_input_batch] = req.sampling_params.top_p

83 min_p[index_in_input_batch] = req.sampling_params.min_p

84 temperature[index_in_input_batch] = req.sampling_params.temperature

86 - stop_token_ids[

87 - index_in_input_batch] = req.sampling_params.all_stop_token_ids

88 - min_tokens[index_in_input_batch] = req.sampling_params.min_tokens

85 + min_tokens[index_in_input_batch] = (

86 + req.sampling_params.min_tokens,

87 + req.sampling_params.all_stop_token_ids)

88 logit_bias[index_in_input_batch] = req.sampling_params.logit_bias

89 return SamplingMetadata(

90 temperature=torch.tensor(temperature, dtype=torch.float,

91 device=device),

92 all_greedy=False,

93 all_random=True,

95 - rejection_sampling=False,

96 - top_p=torch.tensor(top_p, dtype=torch.float, device=device),

97 - top_k=torch.tensor(top_k, dtype=torch.int, device=device),

98 - no_top_p=all(x == 1.0 for x in top_p),

99 - no_top_k=all(x == 0 for x in top_k),

100 - min_p=torch.tensor(min_p, dtype=torch.float, device=device),

101 - no_min_p=all(x == 0.0 for x in min_p),

94 + top_p=None if all(x == 1.0 for x in top_p) else torch.tensor(

95 + top_p, dtype=torch.float, device=device),

96 + top_k=None if all(x == 0 for x in top_k) else torch.tensor(

97 + top_k, dtype=torch.int, device=device),

98 + min_p=None if all(x == 0.0 for x in min_p) else torch.tensor(

99 + min_p, dtype=torch.float, device=device),

100 generators={},

101 max_num_logprobs=0,

102 prompt_token_ids=make_tensor_with_pad(

@@ -117,9 +115,8 @@ def _construct_expected_sampling_metadata(

115 dtype=torch.float,

116 device=device),

117 output_token_ids=output_token_ids,

120 - spec_token_ids=[],

118 + spec_token_ids=None,

119 min_tokens=min_tokens,

122 - stop_token_ids=stop_token_ids,

120 no_penalties=(all(x == 0 for x in presence_penalties)

121 and all(x == 0 for x in frequency_penalties)

122 and all(x == 1 for x in repetition_penalties)),

@@ -206,8 +203,7 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

203 input_batch.condense(req_indices_to_remove)

204

205 # Generate the sampling metadata

209 - sampling_metadata = input_batch.make_sampling_metadata(

210 - req_id_output_token_ids, req_id_to_spec_token_ids={}, skip_copy=False)

206 + sampling_metadata = input_batch._make_sampling_metadata()

207

208 # Create expected output.

209 expected_sampling_metadata = _construct_expected_sampling_metadata(

@@ -216,13 +212,16 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

212 input_batch.req_id_to_index,

213 device=torch.device(device))

214

215 + def same(t1: Optional[torch.Tensor], t2: Optional[torch.Tensor]) -> bool:

216 + return (t1 is None

217 + and t2 is None) or (t1 is not None and t2 is not None

218 + and torch.allclose(t1, t2))

219 +

220 # Assert the actual and expected output.

221 assert torch.allclose(expected_sampling_metadata.temperature,

222 sampling_metadata.temperature)

222 - assert torch.allclose(expected_sampling_metadata.top_p,

223 - sampling_metadata.top_p)

224 - assert torch.allclose(expected_sampling_metadata.top_k,

225 - sampling_metadata.top_k)

223 + assert same(expected_sampling_metadata.top_p, sampling_metadata.top_p)

224 + assert same(expected_sampling_metadata.top_k, sampling_metadata.top_k)

225 assert torch.allclose(

226 expected_sampling_metadata.frequency_penalties,

227 sampling_metadata.frequency_penalties,

@@ -240,10 +239,6 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

239 assert (expected_sampling_metadata.output_token_ids ==

240 sampling_metadata.output_token_ids)

241 assert expected_sampling_metadata.min_tokens == sampling_metadata.min_tokens

243 - assert expected_sampling_metadata.stop_token_ids == \

244 - sampling_metadata.stop_token_ids

242 assert expected_sampling_metadata.no_penalties == \

243 sampling_metadata.no_penalties

247 - assert expected_sampling_metadata.no_top_p == sampling_metadata.no_top_p

248 - assert expected_sampling_metadata.no_top_k == sampling_metadata.no_top_k

244 assert expected_sampling_metadata.logit_bias == sampling_metadata.logit_bias

tests/v1/worker/test_gpu_model_runner.py

+22 -11

@@ -5,6 +5,7 @@ from vllm.config import CacheConfig, ModelConfig, SchedulerConfig, VllmConfig

5 from vllm.sampling_params import SamplingParams

6 from vllm.v1.core.scheduler_output import (CachedRequestData, NewRequestData,

7 SchedulerOutput)

8 +from vllm.v1.sample.metadata import SamplingMetadata

9 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

@@ -82,14 +83,21 @@ def _is_req_added(model_runner, req_id: str) -> bool:

83 return req_id in model_runner.requests

86 +def _is_sampling_metadata_changed(model_runner,

87 + sampling_metadata_before: SamplingMetadata):

88 + return model_runner.input_batch.sampling_metadata is not (

89 + sampling_metadata_before)

90 +

91 +

92 def test_update_states_new_request(model_runner):

93 req_id = "req_0"

95 # new req

96 scheduler_output = _schedule_new_request(req_id)

91 - batch_changed = model_runner._update_states(scheduler_output)

92 - assert batch_changed is True

98 + metadata_before = model_runner.input_batch.sampling_metadata

99 + model_runner._update_states(scheduler_output)

100 + assert _is_sampling_metadata_changed(model_runner, metadata_before)

101 assert _is_req_added(model_runner, req_id)

102 assert _is_req_scheduled(model_runner, req_id)

103

@@ -117,8 +125,9 @@ def test_update_states_request_finished(model_runner):

125 free_encoder_input_ids=[],

126 )

127

120 - batch_changed = model_runner._update_states(scheduler_output)

121 - assert batch_changed is True

128 + metadata_before = model_runner.input_batch.sampling_metadata

129 + model_runner._update_states(scheduler_output)

130 + assert _is_sampling_metadata_changed(model_runner, metadata_before)

131 assert not _is_req_added(model_runner, req_id)

132 assert not _is_req_scheduled(model_runner, req_id)

133

@@ -142,7 +151,7 @@ def test_update_states_request_resumed(model_runner):

151 scheduled_spec_decode_tokens={},

152 scheduled_encoder_inputs={},

153 num_common_prefix_blocks=0,

145 - finished_req_ids={},

154 + finished_req_ids=set(),

155 free_encoder_input_ids=[],

156 )

157

@@ -171,8 +180,9 @@ def test_update_states_request_resumed(model_runner):

180 free_encoder_input_ids=[],

181 )

182

174 - batch_changed = model_runner._update_states(scheduler_output)

175 - assert batch_changed is True

183 + metadata_before = model_runner.input_batch.sampling_metadata

184 + model_runner._update_states(scheduler_output)

185 + assert _is_sampling_metadata_changed(model_runner, metadata_before)

186 assert _is_req_added(model_runner, req_id)

187 assert _is_req_scheduled(model_runner, req_id)

188

@@ -200,8 +210,9 @@ def test_update_states_no_changes(model_runner):

210 free_encoder_input_ids=[],

211 )

212

203 - batch_changed = model_runner._update_states(scheduler_output)

204 - assert batch_changed is False

213 + metadata_before = model_runner.input_batch.sampling_metadata

214 + model_runner._update_states(scheduler_output)

215 + assert not _is_sampling_metadata_changed(model_runner, metadata_before)

216 assert _is_req_added(model_runner, req_id)

217 assert _is_req_scheduled(model_runner, req_id)

218

@@ -233,8 +244,8 @@ def test_update_states_request_unscheduled(model_runner):

244 free_encoder_input_ids=[],

245 )

246

236 - batch_changed = model_runner._update_states(scheduler_output)

237 - assert batch_changed is True

247 + metadata_before = model_runner._update_states(scheduler_output)

248 + assert _is_sampling_metadata_changed(model_runner, metadata_before)

249

250 assert _is_req_added(model_runner, req_ids[0])

251 assert _is_req_scheduled(model_runner, req_ids[0])

vllm/model_executor/layers/utils.py

+3 -3

@@ -45,7 +45,7 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,

45 vocab_size, num_seqs)

46 output_bin_counts, output_mask = get_token_bin_counts_and_mask(

47 output_tokens_tensor, vocab_size, num_seqs)

48 - repetition_penalties = repetition_penalties.unsqueeze_(dim=1).repeat(

48 + repetition_penalties = repetition_penalties.unsqueeze(dim=1).repeat(

49 1, vocab_size)

50 logits[logits > 0] /= torch.where(prompt_mask | output_mask,

51 repetition_penalties, 1.0)[logits > 0]

@@ -53,6 +53,6 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,

53 repetition_penalties, 1.0)[logits <= 0]

54 # We follow the definition in OpenAI API.

55 # Refer to https://platform.openai.com/docs/api-reference/parameter-details

56 - logits -= frequency_penalties.unsqueeze_(dim=1) * output_bin_counts

57 - logits -= presence_penalties.unsqueeze_(dim=1) * output_mask

56 + logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts

57 + logits -= presence_penalties.unsqueeze(dim=1) * output_mask

58 return logits

vllm/v1/core/scheduler.py

+4 -2

@@ -195,8 +195,10 @@ class Scheduler:

195 request.num_computed_tokens -

196 request.num_tokens)

197 if num_scheduled_spec_tokens > 0:

198 + # Trim spec_token_ids list to num_scheduled_spec_tokens.

199 + del request.spec_token_ids[num_scheduled_spec_tokens:]

200 scheduled_spec_decode_tokens[request.request_id] = (

199 - request.spec_token_ids[:num_scheduled_spec_tokens])

201 + request.spec_token_ids)

202

203 # Encoder-related.

204 if encoder_inputs_to_schedule:

@@ -567,7 +569,7 @@ class Scheduler:

569 outputs.append(

570 EngineCoreOutput(

571 request_id=req_id,

570 - new_token_ids=new_token_ids or [],

572 + new_token_ids=new_token_ids,

573 finish_reason=request.get_finished_reason(),

574 new_logprobs=new_logprobs,

575 new_prompt_logprobs_tensors=prompt_logprobs_tensors,

vllm/v1/sample/metadata.py

+10 -11

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from dataclasses import dataclass

4 -from typing import Dict, List, Optional, Set

4 +from typing import Dict, List, Optional, Set, Tuple

6 import torch

@@ -12,15 +12,13 @@ class SamplingMetadata:

12 temperature: torch.Tensor

13 all_greedy: bool

14 all_random: bool

15 - rejection_sampling: bool

16 - spec_token_ids: List[List[int]]

18 - top_p: torch.Tensor

19 - top_k: torch.Tensor

20 - no_top_p: bool

21 - no_top_k: bool

22 - min_p: torch.Tensor

23 - no_min_p: bool

16 + # None when there are no speculated tokens.

17 + spec_token_ids: Optional[List[List[int]]]

18 +

19 + top_p: Optional[torch.Tensor]

20 + top_k: Optional[torch.Tensor]

21 + min_p: Optional[torch.Tensor]

23 generators: Dict[int, torch.Generator]

@@ -34,7 +32,8 @@ class SamplingMetadata:

32 repetition_penalties: torch.Tensor

34 output_token_ids: List[List[int]]

37 - min_tokens: List[int]

38 - stop_token_ids: List[Set[int]]

35 +

36 + # req_index -> (min_tokens, stop_token_ids)

37 + min_tokens: Dict[int, Tuple[int, Set[int]]]

39 logit_bias: List[Optional[Dict[int, float]]]

vllm/v1/sample/ops/penalties.py

+6 -7

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import List, Set, Tuple

3 +from typing import Dict, List, Set, Tuple

5 import torch

@@ -8,18 +8,17 @@ from vllm.model_executor.layers.utils import apply_penalties

8 from vllm.utils import is_pin_memory_available, make_tensor_with_pad

11 -def apply_min_token_penalties(logits: torch.Tensor,

12 - output_token_ids: List[List[int]],

13 - stop_token_ids: List[Set[int]],

14 - min_tokens: List[int]) -> None:

11 +def apply_min_token_penalties(

12 + logits: torch.Tensor, output_token_ids: List[List[int]],

13 + min_tokens: Dict[int, Tuple[int, Set[int]]]) -> None:

14 """

15 Applies minimum token penalty by setting the logits of the stop tokens

16 to -inf.

17 """

18 min_tokens_logits_to_penalize: List[Tuple[int, int]] = []

20 - for index, min_token in enumerate(min_tokens):

19 + for index, (min_token, stop_token_ids) in min_tokens.items():

20 if len(output_token_ids[index]) < min_token:

22 - for stop_token_id in stop_token_ids[index]:

21 + for stop_token_id in stop_token_ids:

22 min_tokens_logits_to_penalize.append((index, stop_token_id))

23 if min_tokens_logits_to_penalize:

24 logits[tuple(zip(*min_tokens_logits_to_penalize))] = -float("inf")

vllm/v1/sample/ops/topk_topp_sampler.py

+20 -28

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Dict

3 +from typing import Dict, Optional

5 import torch

6 import torch.nn as nn

@@ -55,13 +55,11 @@ class TopKTopPSampler(nn.Module):

55 self,

56 logits: torch.Tensor,

57 generators: Dict[int, torch.Generator],

58 - no_top_k: bool,

59 - k: torch.Tensor,

60 - no_top_p: bool,

61 - p: torch.Tensor,

58 + k: Optional[torch.Tensor],

59 + p: Optional[torch.Tensor],

60 ) -> torch.Tensor:

61 """PyTorch-native implementation of top-k and top-p sampling."""

64 - logits = apply_top_k_top_p(logits, no_top_k, k, no_top_p, p)

62 + logits = apply_top_k_top_p(logits, k, p)

63 probs = logits.softmax(dim=-1, dtype=torch.float32)

64 return random_sample(probs, generators)

@@ -69,37 +67,33 @@ class TopKTopPSampler(nn.Module):

67 self,

68 logits: torch.Tensor,

69 generators: Dict[int, torch.Generator],

72 - no_top_k: bool,

73 - k: torch.Tensor,

74 - no_top_p: bool,

75 - p: torch.Tensor,

70 + k: Optional[torch.Tensor],

71 + p: Optional[torch.Tensor],

72 ) -> torch.Tensor:

73 """More optimized implementation for top-k and top-p sampling."""

74 probs = logits.softmax(dim=-1, dtype=torch.float32)

79 - if no_top_k and no_top_p:

75 + if k is None and p is None:

76 # We prefer `random_sample` over `flashinfer_sample` when sorting is

77 # not needed. This is because `random_sample` does not require

78 # CPU-GPU synchronization while `flashinfer_sample` does.

79 return random_sample(probs, generators)

84 - return flashinfer_sample(probs, no_top_k, k, no_top_p, p, generators)

80 + return flashinfer_sample(probs, k, p, generators)

83 def apply_top_k_top_p(

84 logits: torch.Tensor,

89 - no_top_k: bool,

90 - k: torch.Tensor,

91 - no_top_p: bool,

92 - p: torch.Tensor,

85 + k: Optional[torch.Tensor],

86 + p: Optional[torch.Tensor],

87 ) -> torch.Tensor:

88 """Apply top-k and top-p masks to the logits.

90 This function sorts the logits tensor, which can be slow for large batches.

91 """

98 - if no_top_k and no_top_p:

92 + if k is None and p is None:

93 return logits

94 logits_sort, logits_idx = logits.sort(dim=-1, descending=False)

102 - if not no_top_k:

96 + if k is not None:

97 # Apply top-k.

98 top_k_mask = logits_sort.size(1) - k.to(torch.long)

99 # Get all the top_k values.

@@ -107,7 +101,7 @@ def apply_top_k_top_p(

101 top_k_mask = logits_sort < top_k_mask

102 logits_sort.masked_fill_(top_k_mask, -float("inf"))

103

110 - if not no_top_p:

104 + if p is not None:

105 # Apply top-p.

106 probs_sort = logits_sort.softmax(dim=-1)

107 probs_sum = probs_sort.cumsum(dim=-1)

@@ -147,10 +141,8 @@ def random_sample(

141

142 def flashinfer_sample(

143 probs: torch.Tensor,

150 - no_top_k: bool,

151 - k: torch.Tensor,

152 - no_top_p: bool,

153 - p: torch.Tensor,

144 + k: Optional[torch.Tensor],

145 + p: Optional[torch.Tensor],

146 generators: Dict[int, torch.Generator],

147 ) -> torch.Tensor:

148 """Sample from the probabilities using FlashInfer.

@@ -167,7 +159,7 @@ def flashinfer_sample(

159 does not. Call this function at the end of the forward pass to minimize

160 the synchronization overhead.

161 """

170 - assert not (no_top_k and no_top_p)

162 + assert not (k is None and p is None)

163 max_top_k_round = 32

164 batch_size = probs.shape[0]

165 uniform_samples = torch.empty((max_top_k_round, batch_size),

@@ -178,11 +170,11 @@ def flashinfer_sample(

170 for i, generator in generators.items():

171 uniform_samples[:, i].uniform_(generator=generator)

172

181 - if no_top_k:

173 + if k is None:

174 # Top-p only.

175 next_token_ids, success = flashinfer.sampling.top_p_sampling_from_probs(

176 probs, uniform_samples, p, deterministic=True)

185 - elif no_top_p:

177 + elif p is None:

178 # Top-k only.

179 next_token_ids, success = flashinfer.sampling.top_k_sampling_from_probs(

180 probs, uniform_samples, k, deterministic=True)

@@ -194,9 +186,9 @@ def flashinfer_sample(

186

187 # NOTE: CPU-GPU synchronization happens here.

188 if not success.all():

197 - if not no_top_k:

189 + if k is not None:

190 probs = flashinfer.sampling.top_k_renorm_prob(probs, k)

199 - if not no_top_p:

191 + if p is not None:

192 probs = flashinfer.sampling.top_p_renorm_prob(probs, p)

193 next_token_ids = flashinfer.sampling.sampling_from_probs(

194 probs, uniform_samples[0], deterministic=True)

vllm/v1/sample/rejection_sampler.py

@@ -68,6 +68,7 @@ class RejectionSampler(nn.Module):

68 # NOTE: The following input preparationg can be moved

69 # to the model runner with a persistent manner for better

70 # performance.

71 + assert sampling_metadata.spec_token_ids is not None

72 spec_token_ids = sampling_metadata.spec_token_ids

73 max_spec_len = max(len(s) for s in spec_token_ids)

74 batch_size = len(spec_token_ids)

@@ -119,6 +120,7 @@ class RejectionSampler(nn.Module):

120 logits: torch.Tensor,

121 sampling_metadata: SamplingMetadata,

122 ) -> SamplerOutput:

123 + assert sampling_metadata.spec_token_ids is not None

124 spec_lens = [len(x) for x in sampling_metadata.spec_token_ids]

125 # Add 1 to include the 'bonus' token.

126 sample_lens = [x + 1 for x in spec_lens]

vllm/v1/sample/sampler.py

+6 -7

@@ -26,7 +26,7 @@ class Sampler(nn.Module):

26 logits: torch.Tensor,

27 sampling_metadata: SamplingMetadata,

28 ) -> SamplerOutput:

29 - if sampling_metadata.rejection_sampling:

29 + if sampling_metadata.spec_token_ids:

30 if sampling_metadata.max_num_logprobs:

31 raise NotImplementedError(

32 "Rejection sampling does not support logprobs.")

@@ -104,16 +104,14 @@ class Sampler(nn.Module):

104 logits = self.apply_temperature(logits, sampling_metadata.temperature)

105

106 # Apply min_p.

107 - if not sampling_metadata.no_min_p:

107 + if sampling_metadata.min_p is not None:

108 logits = self.apply_min_p(logits, sampling_metadata.min_p)

109

110 # Apply top_k and/or top_p.

111 random_sampled = self.topk_topp_sampler(

112 logits,

113 sampling_metadata.generators,

114 - sampling_metadata.no_top_k,

114 sampling_metadata.top_k,

116 - sampling_metadata.no_top_p,

115 sampling_metadata.top_p,

116 )

117

@@ -179,9 +177,10 @@ class Sampler(nn.Module):

177 logits: torch.Tensor,

178 sampling_metadata: SamplingMetadata,

179 ) -> torch.Tensor:

182 - apply_min_token_penalties(logits, sampling_metadata.output_token_ids,

183 - sampling_metadata.stop_token_ids,

184 - sampling_metadata.min_tokens)

180 + if sampling_metadata.min_tokens:

181 + apply_min_token_penalties(logits,

182 + sampling_metadata.output_token_ids,

183 + sampling_metadata.min_tokens)

184 if not sampling_metadata.no_penalties:

185 assert sampling_metadata.prompt_token_ids is not None

186 logits = apply_all_penalties(

vllm/v1/utils.py

+11

@@ -188,3 +188,14 @@ def bind_kv_cache(

188 for layer_name, kv_cache in kv_caches.items():

189 # NOTE: Use list because of v0 PP virtual engine.

190 forward_context[layer_name].kv_cache = [kv_cache]

191 +

192 +

193 +def copy_slice(from_tensor: torch.Tensor, to_tensor: torch.Tensor,

194 + length: int) -> None:

195 + """

196 + Copy the first length elements of a tensor into another tensor in a

197 + non-blocking manner.

198 +

199 + Used to copy pinned CPU tensor data to pre-allocated GPU tensors.

200 + """

201 + to_tensor[:length].copy_(from_tensor[:length], non_blocking=True)

vllm/v1/worker/gpu_input_batch.py

+102 -111

@@ -1,9 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -

2 # Datastructures defining an input batch

4 from dataclasses import dataclass

6 -from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple

5 +from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, cast

7 import numpy as np

8 import torch

@@ -12,6 +11,7 @@ from vllm.lora.request import LoRARequest

11 from vllm.multimodal import MultiModalKwargs

12 from vllm.sampling_params import SamplingParams, SamplingType

13 from vllm.v1.sample.metadata import SamplingMetadata

14 +from vllm.v1.utils import copy_slice

15 from vllm.v1.worker.block_table import BlockTable

17 _SAMPLING_EPS = 1e-5

@@ -63,7 +63,7 @@ class InputBatch:

63 self.pin_memory = pin_memory

64 self.vocab_size = vocab_size

66 - self.req_ids: List[Optional[str]] = [None] * max_num_reqs

66 + self._req_ids: List[Optional[str]] = []

67 self.req_id_to_index: Dict[str, int] = {}

69 # TODO(woosuk): This buffer could be too large if max_model_len is big.

@@ -171,11 +171,8 @@ class InputBatch:

171 self.repetition_penalties_cpu_tensor.numpy()

172 self.repetition_penalties_reqs: Set[str] = set()

173

174 - self.min_tokens: List[int] = [0] * max_num_reqs

175 - self.stop_token_ids: List[Set[int]] = [

176 - set() for _ in range(max_num_reqs)

177 - ]

178 - self.prompt_token_ids: Optional[torch.Tensor] = None

174 + # req_index -> (min_tokens, stop_token_ids)

175 + self.min_tokens: Dict[int, Tuple[int, Set[int]]] = {}

176

177 # lora related

178 self.request_lora_mapping = np.zeros((self.max_num_reqs, ),

@@ -196,6 +193,17 @@ class InputBatch:

193 self.logit_bias: List[Optional[Dict[int,

194 float]]] = [None] * max_num_reqs

195

196 + self.req_output_token_ids: List[Optional[List[int]]] = []

197 +

198 + # This is updated each time the batch constituents change.

199 + self.sampling_metadata = self._make_sampling_metadata()

200 +

201 + @property

202 + def req_ids(self) -> List[str]:

203 + # None elements should only be present transiently

204 + # while performing state updates to the batch.

205 + return cast(List[str], self._req_ids)

206 +

207 def add_request(

208 self,

209 request: "CachedRequestState",

@@ -206,7 +214,13 @@ class InputBatch:

214 assert req_index < self.max_num_reqs

215

216 req_id = request.req_id

209 - self.req_ids[req_index] = req_id

217 + if req_index == len(self._req_ids):

218 + self._req_ids.append(req_id)

219 + self.req_output_token_ids.append(request.output_token_ids)

220 + else:

221 + self._req_ids[req_index] = req_id

222 + self.req_output_token_ids[req_index] = request.output_token_ids

223 +

224 self.req_id_to_index[req_id] = req_index

225

226 # Copy the prompt token ids and output token ids.

@@ -255,8 +269,9 @@ class InputBatch:

269 req_index] = sampling_params.repetition_penalty

270 if sampling_params.repetition_penalty != 1.0:

271 self.repetition_penalties_reqs.add(req_id)

258 - self.min_tokens[req_index] = sampling_params.min_tokens

259 - self.stop_token_ids[req_index] = sampling_params.all_stop_token_ids

272 + if sampling_params.min_tokens:

273 + self.min_tokens[req_index] = (sampling_params.min_tokens,

274 + sampling_params.all_stop_token_ids)

275

276 # NOTE(woosuk): self.generators should not include the requests that

277 # do not have their own generator.

@@ -284,16 +299,20 @@ class InputBatch:

299 self.request_lora_mapping[req_index] = 0

300

301 def remove_request(self, req_id: str) -> Optional[int]:

302 + """This method must always be followed by a call to condense()."""

303 +

304 req_index = self.req_id_to_index.pop(req_id, None)

305 if req_index is None:

306 return None

290 - self.req_ids[req_index] = None

307 + self._req_ids[req_index] = None

308 + self.req_output_token_ids[req_index] = None

309

310 self.greedy_reqs.discard(req_id)

311 self.random_reqs.discard(req_id)

312 self.top_p_reqs.discard(req_id)

313 self.top_k_reqs.discard(req_id)

314 self.min_p_reqs.discard(req_id)

315 + self.min_tokens.pop(req_index, None)

316 self.frequency_penalties_reqs.discard(req_id)

317 self.presence_penalties_reqs.discard(req_id)

318 self.repetition_penalties_reqs.discard(req_id)

@@ -313,33 +332,17 @@ class InputBatch:

332 self.logit_bias[req_index] = None

333 return req_index

334

316 - def clear(self) -> None:

317 - self.req_ids = [None] * self.max_num_reqs

318 - self.req_id_to_index.clear()

319 - self.greedy_reqs.clear()

320 - self.random_reqs.clear()

321 - self.top_p_reqs.clear()

322 - self.top_k_reqs.clear()

323 - self.min_p_reqs.clear()

324 - self.frequency_penalties_reqs.clear()

325 - self.presence_penalties_reqs.clear()

326 - self.repetition_penalties_reqs.clear()

327 - self.generators.clear()

328 - self.num_logprobs.clear()

329 - self.num_prompt_logprobs.clear()

330 - self.request_lora_mapping.fill(0)

331 - self.lora_id_to_lora_request.clear()

332 - self.lora_id_to_request_ids.clear()

333 - self.logit_bias = [None] * self.max_num_reqs

334 -

335 def condense(self, empty_req_indices: List[int]) -> None:

336 - if self.num_reqs == 0:

336 + num_reqs = self.num_reqs

337 + if num_reqs == 0:

338 # The batched states are empty.

339 + self._req_ids.clear()

340 + self.req_output_token_ids.clear()

341 return

342

343 # NOTE(woosuk): This function assumes that the empty_req_indices

344 # is sorted in descending order.

342 - last_req_index = self.num_reqs + len(empty_req_indices) - 1

345 + last_req_index = num_reqs + len(empty_req_indices) - 1

346 while empty_req_indices:

347 # Find the largest non-empty index.

348 while last_req_index in empty_req_indices:

@@ -351,10 +354,13 @@ class InputBatch:

354 break

355

356 # Swap the states.

354 - req_id = self.req_ids[last_req_index]

357 + req_id = self._req_ids[last_req_index]

358 + output_token_ids = self.req_output_token_ids[last_req_index]

359 assert req_id is not None

356 - self.req_ids[empty_index] = req_id

357 - self.req_ids[last_req_index] = None

360 + self._req_ids[empty_index] = req_id

361 + self._req_ids[last_req_index] = None

362 + self.req_output_token_ids[empty_index] = output_token_ids

363 + self.req_output_token_ids[last_req_index] = None

364 self.req_id_to_index[req_id] = empty_index

365

366 num_tokens = self.num_tokens[last_req_index]

@@ -379,13 +385,14 @@ class InputBatch:

385 self.repetition_penalties_cpu[

386 empty_index] = self.repetition_penalties_cpu[last_req_index]

387 self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index]

382 - self.min_tokens[empty_index] = self.min_tokens[last_req_index]

383 - self.stop_token_ids[empty_index] = self.stop_token_ids[

384 - last_req_index]

388 generator = self.generators.pop(last_req_index, None)

389 if generator is not None:

390 self.generators[empty_index] = generator

391

392 + min_token = self.min_tokens.pop(last_req_index, None)

393 + if min_token is not None:

394 + self.min_tokens[empty_index] = min_token

395 +

396 self.request_lora_mapping[empty_index] = self.request_lora_mapping[

397 last_req_index]

398

@@ -394,87 +401,71 @@ class InputBatch:

401 # Decrement last_req_index since it is now empty.

402 last_req_index -= 1

403

397 - def make_sampling_metadata(

398 - self,

399 - req_id_output_token_ids: Dict[str, List[int]],

400 - req_id_to_spec_token_ids: Dict[str, List[int]],

401 - skip_copy: bool = False,

402 - ) -> SamplingMetadata:

403 - if not skip_copy:

404 - self.temperature[:self.num_reqs].copy_(

405 - self.temperature_cpu_tensor[:self.num_reqs], non_blocking=True)

406 - self.top_p[:self.num_reqs].copy_(

407 - self.top_p_cpu_tensor[:self.num_reqs], non_blocking=True)

408 - self.top_k[:self.num_reqs].copy_(

409 - self.top_k_cpu_tensor[:self.num_reqs], non_blocking=True)

410 - self.min_p[:self.num_reqs].copy_(

411 - self.min_p_cpu_tensor[:self.num_reqs], non_blocking=True)

412 - if not self.no_penalties:

413 - # Since syncing these tensors is expensive only copy them

414 - # if necessary i.e. if there are requests which require

415 - # penalties to be applied during sampling.

416 - self.frequency_penalties[:self.num_reqs].copy_(

417 - self.frequency_penalties_cpu_tensor[:self.num_reqs],

418 - non_blocking=True,

419 - )

420 - self.presence_penalties[:self.num_reqs].copy_(

421 - self.presence_penalties_cpu_tensor[:self.num_reqs],

422 - non_blocking=True,

423 - )

424 - self.repetition_penalties[:self.num_reqs].copy_(

425 - self.repetition_penalties_cpu_tensor[:self.num_reqs],

426 - non_blocking=True,

427 - )

428 - # The prompt tokens are used only for applying penalties during

429 - # the sampling process. Hence copy these tensors only when

430 - # there are requests which need penalties to be applied.

431 - self.prompt_token_ids = self._make_prompt_token_ids_tensor()

432 -

433 - output_token_ids: List[List[int]] = []

434 - spec_token_ids: List[List[int]] = []

435 - rejection_sampling = False

436 - for req_id in self.req_ids[:self.num_reqs]:

437 - assert req_id is not None

438 - # Currently we create a tensor for output_token_ids from scratch

439 - # at each step. However, for the penalties computation what we

440 - # need is stats about the token ids present in the output. This

441 - # stats can be maintained incrementally instead of computing it

442 - # from scratch at each step.

443 - # TODO - Replace this with incremental update to output token

444 - # statistics.

445 - output_token_ids.append(req_id_output_token_ids[req_id])

446 - req_spec_token_ids = req_id_to_spec_token_ids.get(req_id, [])

447 - spec_token_ids.append(req_spec_token_ids)

448 - if req_spec_token_ids:

449 - # If any of the requests require speculative decoding, set the

450 - # flag to True.

451 - rejection_sampling = True

404 + # Trim lists to the batch size.

405 + del self._req_ids[self.num_reqs:]

406 + del self.req_output_token_ids[self.num_reqs:]

407 +

408 + def refresh_sampling_metadata(self):

409 + self.sampling_metadata = self._make_sampling_metadata()

410 +

411 + def _make_sampling_metadata(self) -> SamplingMetadata:

412 + num_reqs = self.num_reqs

413 + copy_slice(self.temperature_cpu_tensor, self.temperature, num_reqs)

414 + if not self.no_top_p:

415 + copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs)

416 + if not self.no_top_k:

417 + copy_slice(self.top_k_cpu_tensor, self.top_k, num_reqs)

418 + if not self.no_min_p:

419 + copy_slice(self.min_p_cpu_tensor, self.min_p, num_reqs)

420 +

421 + if not self.no_penalties:

422 + # Since syncing these tensors is expensive only copy them

423 + # if necessary i.e. if there are requests which require

424 + # penalties to be applied during sampling.

425 + copy_slice(self.frequency_penalties_cpu_tensor,

426 + self.frequency_penalties, num_reqs)

427 + copy_slice(self.presence_penalties_cpu_tensor,

428 + self.presence_penalties, num_reqs)

429 + copy_slice(self.repetition_penalties_cpu_tensor,

430 + self.repetition_penalties, num_reqs)

431 +

432 + # The prompt tokens are used only for applying penalties during

433 + # the sampling process. Hence copy these tensors only when

434 + # there are requests which need penalties to be applied.

435 + prompt_token_ids = self._make_prompt_token_ids_tensor()

436 + else:

437 + prompt_token_ids = None

438

439 return SamplingMetadata(

454 - temperature=self.temperature[:self.num_reqs],

440 + temperature=self.temperature[:num_reqs],

441 all_greedy=self.all_greedy,

442 all_random=self.all_random,

457 - rejection_sampling=rejection_sampling,

458 - top_p=self.top_p[:self.num_reqs],

459 - top_k=self.top_k[:self.num_reqs],

460 - min_p=self.min_p[:self.num_reqs],

461 - no_min_p=self.no_min_p,

462 - no_top_p=self.no_top_p,

463 - no_top_k=self.no_top_k,

443 + top_p=None if self.no_top_p else self.top_p[:num_reqs],

444 + top_k=None if self.no_top_k else self.top_k[:num_reqs],

445 + min_p=None if self.no_min_p else self.min_p[:num_reqs],

446 generators=self.generators,

447 max_num_logprobs=self.max_num_logprobs,

466 - prompt_token_ids=self.prompt_token_ids,

467 - frequency_penalties=self.frequency_penalties[:self.num_reqs],

468 - presence_penalties=self.presence_penalties[:self.num_reqs],

469 - repetition_penalties=self.repetition_penalties[:self.num_reqs],

470 - output_token_ids=output_token_ids,

471 - spec_token_ids=spec_token_ids,

472 - min_tokens=self.min_tokens[:self.num_reqs],

473 - stop_token_ids=self.stop_token_ids[:self.num_reqs],

448 + prompt_token_ids=prompt_token_ids,

449 + frequency_penalties=self.frequency_penalties[:num_reqs],

450 + presence_penalties=self.presence_penalties[:num_reqs],

451 + repetition_penalties=self.repetition_penalties[:num_reqs],

452 + output_token_ids=cast(List[List[int]], self.req_output_token_ids),

453 + spec_token_ids=None,

454 + min_tokens=self.min_tokens,

455 no_penalties=self.no_penalties,

475 - logit_bias=self.logit_bias[:self.num_reqs],

456 + logit_bias=self.logit_bias[:num_reqs],

457 )

458

459 + def get_sampling_metadata(

460 + self,

461 + req_id_to_spec_token_ids: Dict[str, List[int]],

462 + ) -> SamplingMetadata:

463 + # Set the new spec token ids in the cached sampling metadata.

464 + self.sampling_metadata.spec_token_ids = [

465 + req_id_to_spec_token_ids.get(req_id, []) for req_id in self.req_ids

466 + ] if req_id_to_spec_token_ids else None

467 + return self.sampling_metadata

468 +

469 def _make_prompt_token_ids_tensor(self) -> torch.Tensor:

470 max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()

471 prompt_token_ids_cpu_tensor = torch.empty(

vllm/v1/worker/gpu_model_runner.py

+28 -57

@@ -31,7 +31,6 @@ from vllm.v1.engine.mm_input_cache import MMInputCacheClient

31 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

32 KVCacheSpec)

33 from vllm.v1.outputs import LogprobsTensors, ModelRunnerOutput

34 -from vllm.v1.sample.metadata import SamplingMetadata

34 from vllm.v1.sample.rejection_sampler import INVALID_TOKEN_ID

35 from vllm.v1.spec_decode.ngram_proposer import NgramProposer

36 from vllm.v1.utils import bind_kv_cache

@@ -224,16 +223,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

223 pin_memory=self.pin_memory)

224 self.seq_lens_np = self.seq_lens_cpu.numpy()

225

227 - def _update_states(self, scheduler_output: "SchedulerOutput") -> bool:

226 + def _update_states(self, scheduler_output: "SchedulerOutput") -> None:

227 """Update the cached states and the persistent batch with the scheduler

228 output.

229

230 The updated states are used by the `_prepare_inputs` function to create

231 the input GPU tensors for the model.

232

234 - Returns:

235 - True if there is a new/resumed/paused/finished request in the batch.

236 - If False, we can skip copying SamplingMetadata to the GPU.

233 + The SamplingMetadata is updated and copied to the GPU if there is a

234 + new/resumed/paused/finished request in the batch.

235 """

236 # Remove finished requests from the cached states.

237 for req_id in scheduler_output.finished_req_ids:

@@ -344,9 +342,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

342 num_new_tokens = (num_computed_tokens +

343 len(req_data.new_token_ids) -

344 req_state.num_tokens)

347 - new_token_ids = (req_data.new_token_ids[-num_new_tokens:]

348 - if num_new_tokens > 0 else [])

349 - req_state.output_token_ids.extend(new_token_ids)

345 + if num_new_tokens == 1:

346 + # Avoid slicing list in most common case.

347 + req_state.output_token_ids.append(req_data.new_token_ids[-1])

348 + elif num_new_tokens > 0:

349 + req_state.output_token_ids.extend(

350 + req_data.new_token_ids[-num_new_tokens:])

351 # Update the block IDs.

352 if not req_data.resumed_from_preemption:

353 # Append the new blocks to the existing block IDs.

@@ -380,7 +381,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

381 self.input_batch.num_tokens_no_spec[req_index] = end_token_index

382 # Add spec_token_ids to token_ids_cpu.

383 spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(

383 - req_id, [])

384 + req_id, ())

385 if spec_token_ids:

386 start_index = end_token_index

387 end_token_index += len(spec_token_ids)

@@ -410,7 +411,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

411 if removed_req_indices:

412 self.input_batch.condense(removed_req_indices)

413

413 - return batch_changed

414 + if batch_changed:

415 + self.input_batch.refresh_sampling_metadata()

416

417 def _prepare_inputs(

418 self,

@@ -429,8 +431,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

431 # TODO: The Python loop can be slow. Optimize.

432 num_scheduled_tokens = np.empty(num_reqs, dtype=np.int32)

433 max_num_scheduled_tokens = 0

432 - for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

433 - assert req_id is not None

434 + for i, req_id in enumerate(self.input_batch.req_ids):

435 num_tokens = scheduler_output.num_scheduled_tokens[req_id]

436 num_scheduled_tokens[i] = num_tokens

437 max_num_scheduled_tokens = max(max_num_scheduled_tokens,

@@ -669,10 +670,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

670

671 def _calc_mrope_positions(self, scheduler_output: "SchedulerOutput"):

672 mrope_pos_ptr = 0

672 - num_reqs = self.input_batch.num_reqs

673 - for index, req_id in enumerate(self.input_batch.req_ids[:num_reqs]):

674 - assert req_id is not None

675 -

673 + for index, req_id in enumerate(self.input_batch.req_ids):

674 req = self.requests[req_id]

675 assert req.mrope_positions is not None

676

@@ -726,12 +724,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):

724 self,

725 scheduler_output: "SchedulerOutput",

726 cu_num_tokens: np.ndarray,

729 - ) -> Tuple[torch.Tensor, torch.Tensor]:

727 + ) -> torch.Tensor:

728 # Get the number of spec decode tokens for each request.

729 num_reqs = self.input_batch.num_reqs

730 num_spec_decode_tokens = np.empty(num_reqs, dtype=np.int32)

733 - for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

734 - assert req_id is not None

731 + for i, req_id in enumerate(self.input_batch.req_ids):

732 num_spec_decode_tokens[i] = len(

733 scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))

734

@@ -769,22 +766,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

766 return torch.from_numpy(spec_decode_logits_indices).to(

767 self.device, non_blocking=True)

768

772 - def _prepare_sampling(

773 - self,

774 - batch_changed: bool,

775 - req_to_spec_token_ids: Dict[str, List[int]],

776 - ) -> SamplingMetadata:

777 - # Create the sampling metadata.

778 - req_id_output_token_ids: Dict[str, List[int]] = \

779 - {req_id: req.output_token_ids \

780 - for req_id, req in self.requests.items()}

781 -

782 - sampling_metadata = self.input_batch.make_sampling_metadata(

783 - req_id_output_token_ids,

784 - req_to_spec_token_ids,

785 - skip_copy=not batch_changed)

786 - return sampling_metadata

787 -

769 def _execute_encoder(self, scheduler_output: "SchedulerOutput"):

770 scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs

771 if not scheduled_encoder_inputs:

@@ -838,9 +819,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

819 scheduler_output: "SchedulerOutput",

820 ) -> List[torch.Tensor]:

821 encoder_outputs: List[torch.Tensor] = []

841 - num_reqs = self.input_batch.num_reqs

842 - for req_id in self.input_batch.req_ids[:num_reqs]:

843 - assert req_id is not None

822 + for req_id in self.input_batch.req_ids:

823 num_scheduled_tokens = scheduler_output.num_scheduled_tokens[

824 req_id]

825 req_state = self.requests[req_id]

@@ -882,7 +861,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

861 scheduler_output: "SchedulerOutput",

862 intermediate_tensors: Optional[IntermediateTensors] = None,

863 ) -> Union[ModelRunnerOutput, torch.Tensor]:

885 - batch_changed = self._update_states(scheduler_output)

864 + self._update_states(scheduler_output)

865

866 if self.is_multimodal_model:

867 # Run the multimodal encoder if any.

@@ -964,8 +943,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):

943 logits = self.model.compute_logits(sample_hidden_states, None)

944

945 # Sample the next token and get logprobs if needed.

967 - sampling_metadata = self._prepare_sampling(

968 - batch_changed, scheduler_output.scheduled_spec_decode_tokens)

946 + sampling_metadata = self.input_batch.get_sampling_metadata(

947 + scheduler_output.scheduled_spec_decode_tokens)

948 sampler_output = self.model.sample(

949 logits=logits,

950 sampling_metadata=sampling_metadata,

@@ -973,14 +952,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

952

953 # TODO(woosuk): The following loop can be slow since it iterates over

954 # the requests one by one. Optimize.

976 - num_reqs = self.input_batch.num_reqs

977 - req_ids: List[str] = []

978 - # Because `input_batch.req_ids` is a list of length `max_num_reqs`,

979 - # we need to stop at `num_reqs`.

980 - # FIXME(woosuk): This is hacky. Refactor.

981 - for i, req_id in zip(range(num_reqs), self.input_batch.req_ids):

982 - assert req_id is not None

983 - req_ids.append(req_id)

955 + for i, req_id in enumerate(self.input_batch.req_ids):

956 req_state = self.requests[req_id]

957 seq_len = (req_state.num_computed_tokens +

958 scheduler_output.num_scheduled_tokens[req_id])

@@ -1027,7 +999,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

999 valid_sampled_token_ids)

1000

1001 model_runner_output = ModelRunnerOutput(

1030 - req_ids=req_ids,

1002 + req_ids=self.input_batch.req_ids,

1003 req_id_to_index=self.input_batch.req_id_to_index,

1004 sampled_token_ids=valid_sampled_token_ids,

1005 spec_token_ids=spec_token_ids,

@@ -1041,19 +1013,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1013 sampled_token_ids: List[List[int]],

1014 ) -> List[List[int]]:

1015 # TODO(woosuk): Optimize.

1044 - num_reqs = len(sampled_token_ids)

1016 draft_token_ids: List[List[int]] = []

1046 - for i in range(num_reqs):

1047 - if len(sampled_token_ids[i]) == 0:

1017 + for i, sampled_ids in enumerate(sampled_token_ids):

1018 + num_sampled_ids = len(sampled_ids)

1019 + if not num_sampled_ids:

1020 # Skip speculative decoding.

1021 draft_token_ids.append([])

1022 continue

1023

1024 # Add sampled_token_ids to token_ids_cpu.

1025 start_idx = self.input_batch.num_tokens_no_spec[i]

1054 - end_idx = start_idx + len(sampled_token_ids[i])

1055 - self.input_batch.token_ids_cpu[

1056 - i, start_idx:end_idx] = sampled_token_ids[i]

1026 + end_idx = start_idx + num_sampled_ids

1027 + self.input_batch.token_ids_cpu[i, start_idx:end_idx] = sampled_ids

1028 drafter_output = self.drafter.propose(

1029 self.input_batch.token_ids_cpu[i, :end_idx],

1030 self.speculative_config.ngram_prompt_lookup_min,

@@ -1204,7 +1175,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1175 # multiplying the list, to avoid Dynamo from treating them as

1176 # tensor aliasing.

1177 dummy_kv_caches = [

1207 - torch.tensor([], dtype=torch.float32, device=self.device)

1178 + torch.tensor((), dtype=torch.float32, device=self.device)

1179 for _ in range(self.num_attn_layers)

1180 ]

1181

vllm/v1/worker/tpu_model_runner.py

-2

@@ -1048,8 +1048,6 @@ def swap_positions(b: InputBatch, id_1, id_2):

1048

1049 b.min_tokens[id_1], b.min_tokens[id_2] = b.min_tokens[id_2], b.min_tokens[

1050 id_1]

1051 - b.stop_token_ids[id_1], b.stop_token_ids[id_2] = b.stop_token_ids[

1052 - id_2], b.stop_token_ids[id_1]

1051

1052 gen_1 = b.generators.pop(id_1, None)

1053 gen_2 = b.generators.pop(id_2, None)