vllm-metadata
vllm/vllm/v1/worker/gpu_model_runner.py · ef9339e · L2102–2196 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
2102 # Update num_computed_tokens on GPU. In async spec decode,
2103 # CPU values are optimistic (all drafts accepted). The kernel
2104 # corrects on GPU using the previous step's
2105 # valid_sampled_token_count_gpu. Otherwise, just copy from CPU.
2106 if (
2107 self.use_async_spec_decode
2108 and self.valid_sampled_token_count_gpu is not None
2109 and prev_req_id_to_index
2110 ):
2111 self.prev_positions.copy_to_gpu(num_reqs)
2112 self.prev_num_draft_tokens.copy_to_gpu()
2113 cpu_values = self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs].to(
2114 device=self.device, non_blocking=True
2115 )
2116 update_num_computed_tokens_for_batch_change(
2117 self.num_computed_tokens,
2118 self.num_accepted_tokens.gpu[:num_reqs],
2119 self.prev_positions.gpu[:num_reqs],
2120 self.valid_sampled_token_count_gpu,
2121 self.prev_num_draft_tokens.gpu,
2122 cpu_values,
2123 )
2124 else:
2125 self.num_computed_tokens[:num_reqs].copy_(
2126 self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],
2127 non_blocking=True,
2128 )
2129
2130 self.req_indices.np[:total_num_scheduled_tokens] = req_indices
2131 self.req_indices.copy_to_gpu(total_num_scheduled_tokens)
2132 req_indices_gpu = self.req_indices.gpu[:total_num_scheduled_tokens]
2133
2134 self.query_pos.copy_to_gpu(total_num_scheduled_tokens)
2135 self.num_scheduled_tokens.np[:num_reqs] = num_scheduled_tokens
2136 self.num_scheduled_tokens.copy_to_gpu(num_reqs)
2137 num_scheduled_tokens_gpu = self.num_scheduled_tokens.gpu[:num_reqs]
2138 self.positions[:total_num_scheduled_tokens] = (
2139 self.num_computed_tokens[req_indices_gpu].to(torch.int64)
2140 + self.query_pos.gpu[:total_num_scheduled_tokens]
2141 )
2142 self.seq_lens[:num_reqs] = (
2143 self.num_computed_tokens[:num_reqs] + num_scheduled_tokens_gpu
2144 )
2145 self.seq_lens[num_reqs:].fill_(0)
2146
2147 self.input_batch.block_table.compute_slot_mapping(
2148 num_reqs,
2149 self.query_start_loc.gpu[: num_reqs + 1],
2150 self.positions[:total_num_scheduled_tokens],
2151 )
2152
2153 # Copy the tensors to the GPU.
2154 self._prepare_input_ids(
2155 scheduler_output,
2156 num_reqs,
2157 total_num_scheduled_tokens,
2158 cu_num_tokens,
2159 )
2160
2161 if self.uses_mrope:
2162 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
2163 # Copy one row at a time. mrope_positions is allocated as
2164 # [num_dims, max_num_tokens + 1] with a dummy trailing column to keep it
2165 # non-contiguous for torch.compile, so cpu[:, :N] is a strided view.
2166 # copy_() cannot express a strided source as a single
2167 # cudaMemcpyAsync, so it first gathers into a contiguous *pageable*
2168 # temporary, and a pageable H2D ignores non_blocking=True and
2169 # synchronizes the stream before the transfer starts. Each row is
2170 # contiguous within the pinned allocation, so per-row copies stay on
2171 # the pinned path and are genuinely asynchronous.
2172 for row in range(self.mrope_positions.gpu.shape[0]):
2173 self.mrope_positions.gpu[row, :total_num_scheduled_tokens].copy_(
2174 self.mrope_positions.cpu[row, :total_num_scheduled_tokens],
2175 non_blocking=True,
2176 )
2177 if self.use_async_spec_decode and self.uses_mrope:
2178 drift = self.num_computed_tokens[req_indices_gpu].to(
2179 torch.int64
2180 ) - async_tensor_h2d(
2181 self.input_batch.num_computed_tokens_cpu_tensor[req_indices],
2182 device=self.device,
2183 dtype=torch.int64,
2184 )
2185 self.mrope_positions.gpu[:, :total_num_scheduled_tokens] += drift
2186
2187 use_spec_decode = self.speculative_config is not None
2188 if not use_spec_decode:
2189 # NOTE(woosuk): Due to chunked prefills, the batch may contain
2190 # partial requests. While we should not sample any token
2191 # from these partial requests, we do so for simplicity.
2192 # We will ignore the sampled tokens from the partial requests.
2193 # TODO: Support prompt logprobs.
2194 logits_indices = query_start_loc[1:] - 1
2195 spec_decode_metadata = None
2196 num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)