vllm-pack
vllm/vllm/v1/worker/gpu_model_runner.py · ef9339e · L1923–2064 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
1923 def _prepare_inputs(
1924 self,
1925 scheduler_output: "SchedulerOutput",
1926 num_scheduled_tokens: np.ndarray,
1927 ) -> tuple[
1928 torch.Tensor,
1929 SpecDecodeMetadata | None,
1930 int,
1931 ]:
1932 """Returns:
1933 tuple[logits_indices, spec_decode_metadata, max_num_sampled_tokens]
1934
1935 """
1936 total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
1937 assert total_num_scheduled_tokens > 0
1938 num_reqs = self.input_batch.num_reqs
1939 assert num_reqs > 0
1940
1941 # OPTIMIZATION: Start copying the block table first.
1942 # This way, we can overlap the copy with the following CPU operations.
1943 self.input_batch.block_table.commit_block_table(num_reqs)
1944
1945 # Get request indices.
1946 # E.g., [2, 5, 3] -> [0, 0, 1, 1, 1, 1, 1, 2, 2, 2]
1947 req_indices = np.repeat(self.arange_np[:num_reqs], num_scheduled_tokens)
1948
1949 # cu_num_tokens: [2, 5, 3] -> [2, 7, 10]
1950 # self.query_pos.np[:10]: [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
1951 cu_num_tokens = self._get_cumsum_and_arange(
1952 num_scheduled_tokens, self.query_pos.np
1953 )
1954
1955 # Get positions.
1956 positions_np = (
1957 self.input_batch.num_computed_tokens_cpu[req_indices]
1958 + self.query_pos.np[: cu_num_tokens[-1]]
1959 )
1960
1961 # Calculate M-RoPE positions.
1962 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
1963 if self.uses_mrope:
1964 self._calc_mrope_positions(scheduler_output)
1965
1966 # Get token indices.
1967 # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
1968 # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
1969 # where M is the max_model_len.
1970 token_indices = (
1971 positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
1972 )
1973 token_indices_tensor = torch.from_numpy(token_indices)
1974
1975 # NOTE(woosuk): We use torch.index_select instead of np.take here
1976 # because torch.index_select is much faster than np.take for large
1977 # tensors.
1978 torch.index_select(
1979 self.input_batch.token_ids_cpu_tensor.flatten(),
1980 0,
1981 token_indices_tensor,
1982 out=self.input_ids.cpu[:total_num_scheduled_tokens],
1983 )
1984 if self.enable_prompt_embeds:
1985 is_token_ids = self.input_batch.is_token_ids_tensor.flatten()
1986 torch.index_select(
1987 is_token_ids,
1988 0,
1989 token_indices_tensor,
1990 out=self.is_token_ids.cpu[:total_num_scheduled_tokens],
1991 )
1992
1993 # Because we did not pre-allocate a massive prompt_embeds CPU tensor on
1994 # the InputBatch, we need to fill in the prompt embeds into the expected
1995 # spots in the GpuModelRunner's pre-allocated prompt_embeds tensor.
1996 if self.input_batch.req_prompt_embeds:
1997 output_idx = 0
1998 for req_idx in range(num_reqs):
1999 num_sched = num_scheduled_tokens[req_idx]
2000
2001 # Skip if this request doesn't have embeddings
2002 if req_idx not in self.input_batch.req_prompt_embeds:
2003 output_idx += num_sched
2004 continue
2005
2006 # Skip if no tokens scheduled
2007 if num_sched <= 0:
2008 output_idx += num_sched
2009 continue
2010
2011 req_embeds = self.input_batch.req_prompt_embeds[req_idx]
2012 start_pos = self.input_batch.num_computed_tokens_cpu[req_idx]
2013
2014 # Skip if trying to read beyond available embeddings
2015 if start_pos >= req_embeds.shape[0]:
2016 output_idx += num_sched
2017 continue
2018
2019 # Copy available embeddings
2020 end_pos = start_pos + num_sched
2021 actual_end = min(end_pos, req_embeds.shape[0])
2022 actual_num_sched = actual_end - start_pos
2023
2024 if actual_num_sched > 0:
2025 self.inputs_embeds.cpu[
2026 output_idx : output_idx + actual_num_sched
2027 ].copy_(req_embeds[start_pos:actual_end])
2028
2029 output_idx += num_sched
2030
2031 # Prepare the attention metadata.
2032 self.query_start_loc.np[0] = 0
2033 self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens
2034 # Note: pad query_start_loc to be non-decreasing, as kernels
2035 # like FlashAttention requires that
2036 self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1])
2037 self.query_start_loc.copy_to_gpu()
2038 query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]
2039
2040 # Compute optimistic seq_lens (assumes all draft tokens from previous
2041 # iteration accepted). Store in optimistic_seq_lens_cpu for use by
2042 # _build_attention_metadata (max_seq_len) and discard_request_mask.
2043 # seq_lens (GPU) will be computed later using the same optimistic values.
2044 torch.add(
2045 self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs],
2046 torch.from_numpy(num_scheduled_tokens),
2047 out=self.optimistic_seq_lens_cpu[:num_reqs],
2048 )
2049 self.optimistic_seq_lens_cpu[num_reqs:].fill_(0)
2050
2051 # Build prev_positions mapping: current pos -> prev pos (-1 if new).
2052 # Used for gathering from previous iteration's GPU tensors.
2053 prev_req_id_to_index = self.input_batch.prev_req_id_to_index
2054 self._compute_prev_positions(num_reqs)
2055
2056 num_tokens = [self.requests[r].num_tokens for r in self.input_batch.req_ids]
2057 num_tokens_np = np.array(num_tokens, dtype=np.int32)
2058
2059 # Record which requests should not be sampled,
2060 # so that we could clear the sampled tokens before returning
2061 self.discard_request_mask.np[:num_reqs] = (
2062 self.optimistic_seq_lens_cpu[:num_reqs].numpy() < num_tokens_np
2063 )
2064 self.discard_request_mask.copy_to_gpu(num_reqs)