e535d90deb · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/v1/worker/gpu/async_utils.py

+1 -5

@@ -4,11 +4,7 @@

4 import numpy as np

5 import torch

6

7 -from vllm.v1.outputs import (

8 - AsyncModelRunnerOutput,

9 - LogprobsTensors,

10 - ModelRunnerOutput,

11 -)

7 +from vllm.v1.outputs import AsyncModelRunnerOutput, LogprobsTensors, ModelRunnerOutput

8 from vllm.v1.worker.gpu.sample.output import SamplerOutput

9

10

vllm/v1/worker/gpu/attn_utils.py

+7 -17

@@ -32,9 +32,7 @@ def get_kv_cache_spec(vllm_config: VllmConfig) -> dict[str, KVCacheSpec]:

32

33

34 def init_attn_backend(

35 - kv_cache_config: KVCacheConfig,

36 - vllm_config: VllmConfig,

37 - device: torch.device,

35 + kv_cache_config: KVCacheConfig, vllm_config: VllmConfig, device: torch.device

36 ):

37 attn_backends: dict[str, type[AttentionBackend]] = {}

38 attn_metadata_builders: list[AttentionMetadataBuilder] = []

@@ -50,10 +48,7 @@ def init_attn_backend(

48 attn_backends[layer_name] = attn_backend

49

50 attn_metadata_builder = attn_backend.get_builder_cls()(

53 - kv_cache_group_spec.kv_cache_spec,

54 - layer_names,

55 - vllm_config,

56 - device,

51 + kv_cache_group_spec.kv_cache_spec, layer_names, vllm_config, device

52 )

53 attn_metadata_builders.append(attn_metadata_builder) # type: ignore

54

@@ -65,10 +60,7 @@ def init_attn_backend(

60 return attn_backends, attn_metadata_builders

61

62

68 -def _allocate_kv_cache(

69 - kv_cache_config: KVCacheConfig,

70 - device: torch.device,

71 -):

63 +def _allocate_kv_cache(kv_cache_config: KVCacheConfig, device: torch.device):

64 kv_cache_raw_tensors: dict[str, torch.Tensor] = {}

65 for kv_cache_tensor in kv_cache_config.kv_cache_tensors:

66 tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=device)

@@ -141,12 +133,11 @@ def init_kv_cache(

133

134

135 def build_slot_mappings_by_layer(

144 - slot_mappings: torch.Tensor,

145 - kv_cache_config: KVCacheConfig,

136 + slot_mappings: torch.Tensor, kv_cache_config: KVCacheConfig

137 ) -> dict[str, torch.Tensor]:

138 slot_mappings_by_layer: dict[str, torch.Tensor] = {}

148 - for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):

149 - slot_mapping = slot_mappings[i]

139 + kv_cache_groups = kv_cache_config.kv_cache_groups

140 + for slot_mapping, kv_cache_group in zip(slot_mappings, kv_cache_groups):

141 for layer_name in kv_cache_group.layer_names:

142 slot_mappings_by_layer[layer_name] = slot_mapping

143 return slot_mappings_by_layer

@@ -188,8 +179,7 @@ def build_attn_metadata(

179

180 attn_metadata_builder = attn_metadata_builders[i]

181 metadata = attn_metadata_builder.build(

191 - common_prefix_len=0,

192 - common_attn_metadata=common_attn_metadata,

182 + common_prefix_len=0, common_attn_metadata=common_attn_metadata

183 )

184 for layer_name in kv_cache_spec.layer_names:

185 attn_metadata[layer_name] = metadata

vllm/v1/worker/gpu/block_table.py

+2 -5

@@ -71,9 +71,7 @@ class BlockTables:

71 def _make_ptr_tensor(self, x: Iterable[torch.Tensor]) -> torch.Tensor:

72 # NOTE(woosuk): Use uint64 instead of int64 to cover all possible addresses.

73 return torch.tensor(

74 - [t.data_ptr() for t in x],

75 - dtype=torch.uint64,

76 - device=self.device,

74 + [t.data_ptr() for t in x], dtype=torch.uint64, device=self.device

75 )

76

77 def append_block_ids(

@@ -96,8 +94,7 @@ class BlockTables:

94 self.num_blocks.copy_to_uva()

95

96 def gather_block_tables(

99 - self,

100 - idx_mapping: torch.Tensor,

97 + self, idx_mapping: torch.Tensor

98 ) -> tuple[torch.Tensor, ...]:

99 num_reqs = idx_mapping.shape[0]

100 _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs)](

vllm/v1/worker/gpu/buffer_utils.py

+9 -20

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

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from collections.abc import Iterable, Sequence

4 +from functools import partial

5

6 import numpy as np

7 import torch

@@ -81,10 +82,7 @@ class UvaBufferPool:

82

83 class UvaBackedTensor:

84 def __init__(

84 - self,

85 - size: int | Sequence[int],

86 - dtype: torch.dtype,

87 - max_concurrency: int = 2,

85 + self, size: int | Sequence[int], dtype: torch.dtype, max_concurrency: int = 2

86 ):

87 self.dtype = dtype

88 self.max_concurrency = max_concurrency

@@ -135,25 +133,16 @@ class StagedWriteTensor:

133 self._staged_write_contents: list[int | float] = []

134 self._staged_write_cu_lens: list[int] = []

135

138 - self.write_indices = UvaBufferPool(

139 - self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency

140 - )

141 - self.write_starts = UvaBufferPool(

142 - self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency

143 - )

136 + new_buffer = partial(UvaBufferPool, max_concurrency=max_concurrency)

137 +

138 + self.write_indices = new_buffer(self.num_rows, dtype=torch.int32)

139 + self.write_starts = new_buffer(self.num_rows, dtype=torch.int32)

140 init_size = next_power_of_2(self.num_rows)

145 - self.write_contents = UvaBufferPool(

146 - init_size, dtype=dtype, max_concurrency=max_concurrency

147 - )

148 - self.write_cu_lens = UvaBufferPool(

149 - self.num_rows, dtype=torch.int32, max_concurrency=max_concurrency

150 - )

141 + self.write_contents = new_buffer(init_size, dtype=dtype)

142 + self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32)

143

144 def stage_write(

153 - self,

154 - index: int,

155 - start: int,

156 - x: Iterable[int] | Iterable[float],

145 + self, index: int, start: int, x: Iterable[int] | Iterable[float]

146 ) -> None:

147 assert index >= 0

148 assert start >= 0

vllm/v1/worker/gpu/cudagraph_utils.py

+2 -11

@@ -24,12 +24,7 @@ from vllm.v1.worker.gpu.input_batch import InputBuffers

24

25

26 class CudaGraphManager:

27 - def __init__(

28 - self,

29 - vllm_config: VllmConfig,

30 - uses_mrope: bool,

31 - device: torch.device,

32 - ):

27 + def __init__(self, vllm_config: VllmConfig, uses_mrope: bool, device: torch.device):

28 self.vllm_config = vllm_config

29 self.scheduler_config = vllm_config.scheduler_config

30 self.uses_mrope = uses_mrope

@@ -41,11 +36,7 @@ class CudaGraphManager:

36 self.dp_size = vllm_config.parallel_config.data_parallel_size

37 self.compilation_config = vllm_config.compilation_config

38 assert self.compilation_config is not None

44 - self.cudagraph_mode: CUDAGraphMode

45 - if self.compilation_config.cudagraph_mode is None:

46 - self.cudagraph_mode = CUDAGraphMode.NONE

47 - else:

48 - self.cudagraph_mode = self.compilation_config.cudagraph_mode

39 + self.cudagraph_mode = self.compilation_config.cudagraph_mode

40 self.cudagraph_sizes = get_cudagraph_sizes(

41 self.compilation_config.cudagraph_capture_sizes,

42 self.max_num_reqs,

vllm/v1/worker/gpu/dp_utils.py

+2 -8

@@ -13,10 +13,7 @@ def make_num_tokens_across_dp(dp_size: int, num_tokens: int) -> torch.Tensor | N

13

14

15 def get_batch_metadata_across_dp(

16 - num_tokens: int,

17 - cudagraph_size: int,

18 - dp_size: int,

19 - dp_rank: int,

16 + num_tokens: int, cudagraph_size: int, dp_size: int, dp_rank: int

17 ) -> tuple[torch.Tensor, torch.Tensor]:

18 assert dp_size > 1

19 # Use CPU group to avoid CPU-GPU synchronization.

@@ -29,10 +26,7 @@ def get_batch_metadata_across_dp(

26

27

28 def get_cudagraph_and_dp_padding(

32 - num_tokens: int,

33 - cudagraph_size: int | None,

34 - dp_size: int,

35 - dp_rank: int,

29 + num_tokens: int, cudagraph_size: int | None, dp_size: int, dp_rank: int

30 ) -> tuple[bool, int, torch.Tensor | None]:

31 if dp_size == 1:

32 if cudagraph_size is not None:

vllm/v1/worker/gpu/kv_connector.py

+4 -4

@@ -65,10 +65,10 @@ class ActiveKVConnector(KVConnector):

65

66 if scheduler_output.preempted_req_ids:

67 self.kv_connector.handle_preemptions(scheduler_output.preempted_req_ids)

68 - assert scheduler_output.kv_connector_metadata is not None

69 - self.kv_connector.bind_connector_metadata(

70 - scheduler_output.kv_connector_metadata

71 - )

68 + kv_connector_metadata = scheduler_output.kv_connector_metadata

69 + assert kv_connector_metadata is not None

70 + self.kv_connector.bind_connector_metadata(kv_connector_metadata)

71 +

72 # TODO: sort out KV Connectors' use of forward_context

73 if is_forward_context_available():

74 self.kv_connector.start_load_kv(get_forward_context())

vllm/v1/worker/gpu/lora_utils.py

+2 -5

@@ -15,10 +15,7 @@ class LoraState:

15 self.lora_requests: dict[str, LoRARequest] = {}

16

17 def add_request(

18 - self,

19 - req_id: str,

20 - req_index: int,

21 - lora_request: LoRARequest | None,

18 + self, req_id: str, req_index: int, lora_request: LoRARequest | None

19 ) -> None:

20 if lora_request is not None:

21 self.lora_requests[req_id] = lora_request

@@ -41,7 +38,7 @@ class LoraState:

38

39 active_lora_requests: set[LoRARequest] = set()

40 for req_id in req_ids:

44 - lora_request = self.lora_requests.get(req_id, None)

41 + lora_request = self.lora_requests.get(req_id)

42 if lora_request is not None:

43 active_lora_requests.add(lora_request)

44 return prompt_lora_mapping, token_lora_mapping, active_lora_requests

vllm/v1/worker/gpu/mm/encoder_runner.py

+8 -23

@@ -23,10 +23,7 @@ class EncoderRunner:

23 self.device = device

24

25 self.inputs_embeds = torch.zeros(

26 - max_num_tokens,

27 - hidden_size,

28 - dtype=dtype,

29 - device=device,

26 + max_num_tokens, hidden_size, dtype=dtype, device=device

27 )

28 self.req_id_to_mm_features: dict[str, list[MultiModalFeatureSpec]] = {}

29 self.encoder_cache: dict[str, torch.Tensor] = {}

@@ -57,8 +54,7 @@ class EncoderRunner:

54 self.req_id_to_mm_features.pop(req_id, None)

55

56 def prepare_mm_inputs(

60 - self,

61 - scheduled_encoder_inputs: dict[str, list[int]],

57 + self, scheduled_encoder_inputs: dict[str, list[int]]

58 ) -> tuple[list[str], list[tuple[str, MultiModalKwargsItem]]]:

59 mm_hashes: list[str] = []

60 mm_kwargs: list[tuple[str, MultiModalKwargsItem]] = []

@@ -85,20 +81,16 @@ class EncoderRunner:

81

82 encoder_outputs: list[torch.Tensor] = []

83 for modality, num_items, mm_kwargs_group in group_mm_kwargs_by_modality(

88 - mm_kwargs,

89 - device=self.device,

90 - pin_memory=False,

84 + mm_kwargs, device=self.device, pin_memory=False

85 ):

86 curr_group_outputs = model.embed_multimodal(**mm_kwargs_group)

87 sanity_check_mm_encoder_outputs(

94 - curr_group_outputs,

95 - expected_num_items=num_items,

88 + curr_group_outputs, expected_num_items=num_items

89 )

90 encoder_outputs.extend(curr_group_outputs)

91

92 # Cache the encoder outputs by mm_hash

100 - for mm_hash, output in zip(mm_hashes, encoder_outputs):

101 - self.encoder_cache[mm_hash] = output

93 + self.encoder_cache.update(zip(mm_hashes, encoder_outputs))

94 return encoder_outputs

95

96 def gather_mm_embeddings(

@@ -115,9 +107,7 @@ class EncoderRunner:

107 if all_decode:

108 # All decode requests, so no need to gather any embeddings.

109 return [], torch.zeros(

118 - total_num_scheduled_tokens,

119 - dtype=torch.bool,

120 - device=self.device,

110 + total_num_scheduled_tokens, dtype=torch.bool, device=self.device

111 )

112

113 query_start = computed_prefill_lens.tolist()

@@ -125,10 +115,7 @@ class EncoderRunner:

115

116 mm_embeds: list[torch.Tensor] = []

117 is_mm_embed = torch.zeros(

128 - total_num_scheduled_tokens,

129 - dtype=torch.bool,

130 - device="cpu",

131 - pin_memory=True,

118 + total_num_scheduled_tokens, dtype=torch.bool, device="cpu", pin_memory=True

119 )

120 for i, req_id in enumerate(req_ids):

121 if not is_prefilling[i]:

@@ -189,9 +176,7 @@ class EncoderRunner:

176 is_mm_embed: torch.Tensor,

177 ) -> torch.Tensor:

178 x = model.embed_input_ids(

192 - input_ids,

193 - multimodal_embeddings=mm_embeds,

194 - is_multimodal=is_mm_embed,

179 + input_ids, multimodal_embeddings=mm_embeds, is_multimodal=is_mm_embed

180 )

181 # Copy to the pre-allocated buffer for CUDA graphs.

182 self.inputs_embeds[: x.shape[0]] = x

vllm/v1/worker/gpu/mm/mrope_utils.py

+1 -4

@@ -51,10 +51,7 @@ class MRopeState:

51 mm_features: list,

52 ) -> None:

53 prefill_mrope_positions, prefill_mrope_delta = (

54 - mrope_model.get_mrope_input_positions(

55 - prefill_token_ids,

56 - mm_features,

57 - )

54 + mrope_model.get_mrope_input_positions(prefill_token_ids, mm_features)

55 )

56 for i in range(3):

57 pos = prefill_mrope_positions[i].tolist()

vllm/v1/worker/gpu/model_runner.py

+16 -23

@@ -339,10 +339,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):

339 gc.collect()

340

341 def reset_mm_cache(self) -> None:

342 - self.encoder_runner.reset_mm_cache()

342 + if self.supports_mm_inputs:

343 + self.encoder_runner.reset_mm_cache()

344

345 def reset_encoder_cache(self) -> None:

345 - self.encoder_runner.reset_encoder_cache()

346 + if self.supports_mm_inputs:

347 + self.encoder_runner.reset_encoder_cache()

348

349 def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:

350 # SP is not supported yet.

@@ -402,10 +404,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

404

405 def finish_requests(self, scheduler_output: SchedulerOutput) -> None:

406 finished_req_ids = scheduler_output.finished_req_ids

405 - if scheduler_output.preempted_req_ids:

406 - finished_req_ids = finished_req_ids.union(

407 - scheduler_output.preempted_req_ids

408 - )

407 + preempted_req_ids = scheduler_output.preempted_req_ids

408 + if preempted_req_ids:

409 + finished_req_ids = finished_req_ids.union(preempted_req_ids)

410 for req_id in finished_req_ids:

411 self.req_states.remove_request(req_id)

412 if self.supports_mm_inputs:

@@ -477,28 +478,21 @@ class GPUModelRunner(LoRAModelRunnerMixin):

478 )

479

480 def prepare_inputs(

480 - self,

481 - scheduler_output: SchedulerOutput,

482 - num_tokens_after_padding: int,

481 + self, scheduler_output: SchedulerOutput, num_tokens_after_padding: int

482 ) -> InputBatch:

483 num_tokens = scheduler_output.total_num_scheduled_tokens

484 assert num_tokens > 0

486 - num_reqs = len(scheduler_output.num_scheduled_tokens)

485 + num_tokens_per_req = scheduler_output.num_scheduled_tokens

486 + num_reqs = len(num_tokens_per_req)

487

488 # Decode first, then prefill.

489 # batch_idx -> req_id

490 - req_ids = sorted(

491 - scheduler_output.num_scheduled_tokens.keys(),

492 - key=lambda k: scheduler_output.num_scheduled_tokens[k],

493 - )

494 - num_scheduled_tokens = np.array(

495 - [scheduler_output.num_scheduled_tokens[i] for i in req_ids], dtype=np.int32

496 - )

490 + req_ids = sorted(num_tokens_per_req, key=num_tokens_per_req.get) # type: ignore[arg-type]

491 + numtoks_iter = map(num_tokens_per_req.get, req_ids)

492 + num_scheduled_tokens = np.fromiter(numtoks_iter, dtype=np.int32, count=num_reqs)

493

498 - idx_mapping_list = [

499 - self.req_states.req_id_to_index[req_id] for req_id in req_ids

500 - ]

501 - idx_mapping_np = np.array(idx_mapping_list, dtype=np.int32)

494 + idx_mapping_iter = map(self.req_states.req_id_to_index.get, req_ids)

495 + idx_mapping_np = np.fromiter(idx_mapping_iter, dtype=np.int32, count=num_reqs)

496 idx_mapping = async_copy_to_gpu(idx_mapping_np, device=self.device)

497

498 # Get the number of draft tokens for each request.

@@ -889,8 +883,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

883

884 @torch.inference_mode()

885 def sample_tokens(

892 - self,

893 - grammar_output: GrammarOutput | None,

886 + self, grammar_output: GrammarOutput | None

887 ) -> AsyncOutput | ModelRunnerOutput:

888 assert self.execute_model_state is not None

889 hidden_states, input_batch, kv_connector_output = self.execute_model_state

vllm/v1/worker/gpu/sample/logit_bias.py

+2 -9

@@ -13,11 +13,7 @@ MAX_NUM_STOP_TOKEN_IDS = 128

13

14

15 class LogitBiasState:

16 - def __init__(

17 - self,

18 - max_num_reqs: int,

19 - device: torch.device,

20 - ):

16 + def __init__(self, max_num_reqs: int, device: torch.device):

17 self.max_num_reqs = max_num_reqs

18

19 # Allowed token IDs.

@@ -54,10 +50,7 @@ class LogitBiasState:

50 self.use_logit_bias = np.zeros(max_num_reqs, dtype=bool)

51

52 def add_request(

57 - self,

58 - req_idx: int,

59 - prompt_len: int,

60 - sampling_params: SamplingParams,

53 + self, req_idx: int, prompt_len: int, sampling_params: SamplingParams

54 ) -> None:

55 # Using any logit bias.

56 use_logit_bias = False

vllm/v1/worker/gpu/sample/logprob.py

+7 -21

@@ -73,19 +73,12 @@ def _ranks_kernel(

73

74

75 def compute_token_logprobs(

76 - logits: torch.Tensor,

77 - token_ids: torch.Tensor,

76 + logits: torch.Tensor, token_ids: torch.Tensor

77 ) -> torch.Tensor:

79 - batch_size = logits.shape[0]

80 - vocab_size = logits.shape[1]

78 + batch_size, vocab_size = logits.shape

79 token_ids = token_ids.to(torch.int64)

80 num_logprobs = token_ids.shape[1]

83 - logprobs = torch.empty(

84 - batch_size,

85 - num_logprobs,

86 - dtype=torch.float32,

87 - device=logits.device,

88 - )

81 + logprobs = logits.new_empty((batch_size, num_logprobs), dtype=torch.float32)

82 _topk_log_softmax_kernel[(batch_size,)](

83 logprobs,

84 logits,

@@ -107,23 +100,16 @@ def compute_topk_logprobs(

100 ) -> LogprobsTensors:

101 assert num_logprobs >= 0

102 batch_size, vocab_size = logits.shape

110 - if num_logprobs == 0:

111 - logprob_token_ids = sampled_token_ids.unsqueeze(-1)

112 - else:

103 + logprob_token_ids = sampled_token_ids.unsqueeze(-1)

104 + if num_logprobs > 0:

105 topk_indices = torch.topk(logits, num_logprobs, dim=-1).indices

114 - logprob_token_ids = torch.cat(

115 - (sampled_token_ids.unsqueeze(-1), topk_indices), dim=1

116 - )

106 + logprob_token_ids = torch.cat((logprob_token_ids, topk_indices), dim=1)

107

108 # NOTE(woosuk): Here, to save GPU memory, we do not materialize the full

109 # logprobs tensor. Instead, we only compute and return the logprobs of

110 # the topk + 1 tokens.

111 logprobs = compute_token_logprobs(logits, logprob_token_ids)

122 - token_ranks = torch.empty(

123 - batch_size,

124 - dtype=torch.int64,

125 - device=logits.device,

126 - )

112 + token_ranks = torch.empty(batch_size, dtype=torch.int64, device=logits.device)

113 _ranks_kernel[(batch_size,)](

114 token_ranks,

115 logits,

vllm/v1/worker/gpu/sample/min_p.py

+1 -3

@@ -42,9 +42,7 @@ def _min_p_kernel(

42

43

44 def apply_min_p(

45 - logits: torch.Tensor,

46 - idx_mapping: torch.Tensor,

47 - min_p: torch.Tensor,

45 + logits: torch.Tensor, idx_mapping: torch.Tensor, min_p: torch.Tensor

46 ) -> None:

47 num_reqs, vocab_size = logits.shape

48 BLOCK_SIZE = 1024

vllm/v1/worker/gpu/sample/prompt_logprob.py

+1 -3

@@ -23,11 +23,9 @@ class PromptLogprobsWorker:

23 def add_request(self, req_id: str, req_idx: int, sampling_params: SamplingParams):

24 # For now, only support prompt logprobs for the prompt tokens (not top-k).

25 uses_prompt_logprobs = sampling_params.prompt_logprobs is not None

26 + self.uses_prompt_logprobs[req_idx] = uses_prompt_logprobs

27 if uses_prompt_logprobs:

27 - self.uses_prompt_logprobs[req_idx] = True

28 self.in_progress_prompt_logprobs[req_id] = []

29 - else:

30 - self.uses_prompt_logprobs[req_idx] = False

29

30 def remove_request(self, req_id: str) -> None:

31 self.in_progress_prompt_logprobs.pop(req_id, None)

vllm/v1/worker/gpu/sample/sampler.py

+4 -10

@@ -26,7 +26,7 @@ class Sampler:

26 device: torch.device,

27 logprobs_mode: LogprobsMode = "raw_logprobs",

28 ):

29 - if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:

29 + if logprobs_mode not in ("processed_logprobs", "raw_logprobs"):

30 raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")

31 self.logprobs_mode = logprobs_mode

32 self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.

@@ -36,10 +36,7 @@ class Sampler:

36 self.logit_bias_state = LogitBiasState(max_num_reqs, device)

37

38 def add_request(

39 - self,

40 - req_idx: int,

41 - prompt_len: int,

42 - sampling_params: SamplingParams,

39 + self, req_idx: int, prompt_len: int, sampling_params: SamplingParams

40 ) -> None:

41 self.sampling_states.add_request(req_idx, sampling_params)

42 self.penalties_state.add_request(req_idx, sampling_params)

@@ -74,11 +71,8 @@ class Sampler:

71

72 max_num_logprobs = self.sampling_states.max_num_logprobs(idx_mapping_np)

73 if max_num_logprobs != NO_LOGPROBS:

77 - logits = (

78 - processed_logits

79 - if self.logprobs_mode == "processed_logprobs"

80 - else logits

81 - )

74 + if self.logprobs_mode == "processed_logprobs":

75 + logits = processed_logits

76 expanded_logits = logits.shape[0] != idx_mapping_np.shape[0]

77 cu_num_logits = cu_num_logits_np.tolist() if expanded_logits else None

78 logprobs_tensors = compute_topk_logprobs(

vllm/v1/worker/gpu/sample/states.py

+6 -9

@@ -35,22 +35,19 @@ class SamplingStates:

35 def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:

36 self.temperature.np[req_idx] = sampling_params.temperature

37 self.top_p.np[req_idx] = sampling_params.top_p

38 - if 0 < sampling_params.top_k < self.vocab_size:

39 - top_k = sampling_params.top_k

40 - else:

38 + top_k = sampling_params.top_k

39 + if top_k <= 0 or top_k > self.vocab_size:

40 top_k = self.vocab_size

41 self.top_k.np[req_idx] = top_k

42 self.min_p.np[req_idx] = sampling_params.min_p

43

45 - if sampling_params.seed is not None:

46 - seed = sampling_params.seed

47 - else:

44 + seed = sampling_params.seed

45 + if seed is None:

46 seed = np.random.randint(_NP_INT64_MIN, _NP_INT64_MAX)

47 self.seeds.np[req_idx] = seed

48

51 - if sampling_params.logprobs is not None:

52 - num_logprobs = sampling_params.logprobs

53 - else:

49 + num_logprobs = sampling_params.logprobs

50 + if num_logprobs is None:

51 num_logprobs = NO_LOGPROBS

52 self.num_logprobs[req_idx] = num_logprobs

53

vllm/v1/worker/gpu/spec_decode/__init__.py

+1 -4

@@ -5,10 +5,7 @@ import torch

5 from vllm.config import VllmConfig

6

7

8 -def init_speculator(

9 - vllm_config: VllmConfig,

10 - device: torch.device,

11 -):

8 +def init_speculator(vllm_config: VllmConfig, device: torch.device):

9 speculative_config = vllm_config.speculative_config

10 assert speculative_config is not None

11 if speculative_config.use_eagle():

vllm/v1/worker/gpu/spec_decode/eagle.py

+4 -15

@@ -54,26 +54,15 @@ class EagleSpeculator:

54 device=device,

55 )

56 self.hidden_states = torch.zeros(

57 - self.max_num_tokens,

58 - self.hidden_size,

59 - dtype=self.dtype,

60 - device=device,

57 + self.max_num_tokens, self.hidden_size, dtype=self.dtype, device=device

58 )

59 self.idx_mapping = torch.zeros(

63 - self.max_num_reqs,

64 - dtype=torch.int32,

65 - device=device,

60 + self.max_num_reqs, dtype=torch.int32, device=device

61 )

62 self.temperature = torch.zeros(

68 - self.max_num_reqs,

69 - dtype=torch.float32,

70 - device=device,

71 - )

72 - self.seeds = torch.zeros(

73 - self.max_num_reqs,

74 - dtype=torch.int64,

75 - device=device,

63 + self.max_num_reqs, dtype=torch.float32, device=device

64 )

65 + self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device)

66 self.draft_tokens = torch.zeros(

67 self.max_num_reqs,

68 self.num_speculative_steps,

vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py

+5 -15

@@ -19,11 +19,7 @@ from vllm.v1.worker.gpu.input_batch import InputBuffers

19

20

21 class EagleCudaGraphManager:

22 - def __init__(

23 - self,

24 - vllm_config: VllmConfig,

25 - device: torch.device,

26 - ):

22 + def __init__(self, vllm_config: VllmConfig, device: torch.device):

23 self.vllm_config = vllm_config

24 self.scheduler_config = vllm_config.scheduler_config

25 self.device = device

@@ -35,16 +31,10 @@ class EagleCudaGraphManager:

31 self.compilation_config = vllm_config.compilation_config

32 assert self.compilation_config is not None

33

38 - cudagraph_mode: CUDAGraphMode

39 - if self.compilation_config.cudagraph_mode is None:

40 - cudagraph_mode = CUDAGraphMode.NONE

41 - else:

42 - cudagraph_mode = self.compilation_config.cudagraph_mode

43 - if cudagraph_mode == CUDAGraphMode.FULL:

44 - # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.

45 - cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY

46 -

47 - self.cudagraph_mode = cudagraph_mode

34 + self.cudagraph_mode = self.compilation_config.cudagraph_mode

35 + if self.cudagraph_mode == CUDAGraphMode.FULL:

36 + # NOTE(woosuk): For Eagle, we only use CUDA graphs for decode.

37 + self.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY

38

39 self.cudagraph_sizes = get_cudagraph_sizes(

40 self.compilation_config.cudagraph_capture_sizes,

vllm/v1/worker/gpu/structured_outputs.py

+1 -6

@@ -10,12 +10,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch

10

11

12 class StructuredOutputsWorker:

13 - def __init__(

14 - self,

15 - max_num_logits: int,

16 - vocab_size: int,

17 - device: torch.device,

18 - ):

13 + def __init__(self, max_num_logits: int, vocab_size: int, device: torch.device):

14 self.logits_indices = torch.zeros(

15 max_num_logits, dtype=torch.int32, device=device

16 )

[ModelRunner V2] Misc minor simplifications and optimizations (#33467)