We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Mistral/Mixtral generate nonsense past 4096 tokens in prompt
If the prompt contains more than 4k tokens, the model will begin generating nonsense. This seems to be true for both Mistral and Mixtral. I launched the engine at both 32k and 8k max model lengths for testing. Tested with the latest vllm release.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `vllm/worker/model_runner.py`
Content:
```
1 from typing import Dict, List, Optional, Tuple
2 
3 import torch
4 
5 from vllm.config import ModelConfig, ParallelConfig, SchedulerConfig
6 from vllm.logger import init_logger
7 from vllm.model_executor import get_model, InputMetadata, SamplingMetadata
8 from vllm.sampling_params import SamplingParams, SamplingType
9 from vllm.sequence import SamplerOutput, SequenceData, SequenceGroupMetadata
10 
11 logger = init_logger(__name__)
12 
13 _PAD_SLOT_ID = -1
14 
15 
16 class ModelRunner:
17 
18     def __init__(
19         self,
20         model_config: ModelConfig,
21         parallel_config: ParallelConfig,
22         scheduler_config: SchedulerConfig,
23     ):
24         self.model_config = model_config
25         self.parallel_config = parallel_config
26         self.scheduler_config = scheduler_config
27 
28         # model_config can be None in tests/samplers/test_sampler.py.
29         # FIXME(woosuk): This is a hack to make the tests work. Refactor this.
30         self.sliding_window = (model_config.get_sliding_window()
31                                if model_config is not None else None)
32         self.model = None
33         self.block_size = None  # Set after initial profiling.
34 
35     def load_model(self) -> None:
36         self.model = get_model(self.model_config)
37 
38     def set_block_size(self, block_size: int) -> None:
39         self.block_size = block_size
40 
41     def _prepare_prompt(
42         self,
43         seq_group_metadata_list: List[SequenceGroupMetadata],
44     ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata]:
45         assert len(seq_group_metadata_list) > 0
46         input_tokens: List[List[int]] = []
47         input_positions: List[List[int]] = []
48         slot_mapping: List[List[int]] = []
49 
50         prompt_lens: List[int] = []
51         for seq_group_metadata in seq_group_metadata_list:
52             assert seq_group_metadata.is_prompt
53             seq_ids = list(seq_group_metadata.seq_data.keys())
54             assert len(seq_ids) == 1
55             seq_id = seq_ids[0]
56 
57             seq_data = seq_group_metadata.seq_data[seq_id]
58             prompt_tokens = seq_data.get_token_ids()
59             prompt_len = len(prompt_tokens)
60             prompt_lens.append(prompt_len)
61 
62             input_tokens.append(prompt_tokens)
63             # NOTE(woosuk): Here we assume that the first token in the prompt
64             # is always the first token in the sequence.
65             input_positions.append(list(range(prompt_len)))
66 
67             if seq_group_metadata.block_tables is None:
68                 # During memory profiling, the block tables are not initialized
69                 # yet. In this case, we just use a dummy slot mapping.
70                 slot_mapping.append([_PAD_SLOT_ID] * prompt_len)
71                 continue
72 
73             # Compute the slot mapping.
74             slot_mapping.append([])
75             block_table = seq_group_metadata.block_tables[seq_id]
76             # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,
77             # where start_idx is max(0, prompt_len - sliding_window).
78             # For example, if the prompt len is 10, sliding window is 8, and
79             # block size is 4, the first two tokens are masked and the slot
80             # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].
81             start_idx = 0
82             if self.sliding_window is not None:
83                 start_idx = max(0, prompt_len - self.sliding_window)
84             for i in range(prompt_len):
85                 if i < start_idx:
86                     slot_mapping[-1].append(_PAD_SLOT_ID)
87                     continue
88 
89                 block_number = block_table[i // self.block_size]
90                 block_offset = i % self.block_size
91                 slot = block_number * self.block_size + block_offset
92                 slot_mapping[-1].append(slot)
93 
94         max_prompt_len = max(prompt_lens)
95         input_tokens = _make_tensor_with_pad(input_tokens,
96                                              max_prompt_len,
97                                              pad=0,
98                                              dtype=torch.long)
99         input_positions = _make_tensor_with_pad(input_positions,
100                                                 max_prompt_len,
101                                                 pad=0,
102                                                 dtype=torch.long)
103         slot_mapping = _make_tensor_with_pad(slot_mapping,
104                                              max_prompt_len,
105                                              pad=_PAD_SLOT_ID,
106                                              dtype=torch.long)
107 
108         input_metadata = InputMetadata(
109             prompt_lens=prompt_lens,
110             slot_mapping=slot_mapping,
111             max_context_len=None,
112             context_lens=None,
113             block_tables=None,
114         )
115         return input_tokens, input_positions, input_metadata
116 
117     def _prepare_decode(
118         self,
119         seq_group_metadata_list: List[SequenceGroupMetadata],
120     ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata]:
121         assert len(seq_group_metadata_list) > 0
122         input_tokens: List[List[int]] = []
123         input_positions: List[List[int]] = []
124         slot_mapping: List[List[int]] = []
125         context_lens: List[int] = []
126         block_tables: List[List[int]] = []
127 
128         for seq_group_metadata in seq_group_metadata_list:
129             assert not seq_group_metadata.is_prompt
130 
131             seq_ids = list(seq_group_metadata.seq_data.keys())
132             for seq_id in seq_ids:
133                 seq_data = seq_group_metadata.seq_data[seq_id]
134                 generation_token = seq_data.get_last_token_id()
135                 input_tokens.append([generation_token])
136 
137                 context_len = seq_data.get_len()
138                 if self.sliding_window is not None:
139                     context_len = min(context_len, self.sliding_window)
140                 context_lens.append(context_len)
141 
142                 position = context_len - 1
143                 input_positions.append([position])
144 
145                 block_table = seq_group_metadata.block_tables[seq_id]
146                 block_number = block_table[position // self.block_size]
147                 block_offset = position % self.block_size
148                 slot = block_number * self.block_size + block_offset
149                 slot_mapping.append([slot])
150 
151                 if self.sliding_window is not None:
152                     sliding_window_blocks = (self.sliding_window //
153                                              self.block_size)
154                     block_table = block_table[-sliding_window_blocks:]
155                 block_tables.append(block_table)
156 
157         input_tokens = _make_tensor_with_pad(input_tokens,
158                                              max_len=1,
159                                              pad=0,
160                                              dtype=torch.long)
161         input_positions = _make_tensor_with_pad(input_positions,
162                                                 max_len=1,
163                                                 pad=0,
164                                                 dtype=torch.long)
165         slot_mapping = _make_tensor_with_pad(slot_mapping,
166                                              max_len=1,
167                                              pad=_PAD_SLOT_ID,
168                                              dtype=torch.long)
169         max_context_len = max(context_lens)
170         context_lens = torch.tensor(context_lens,
171                                     dtype=torch.int,
172                                     device="cuda")
173         max_block_table_len = max([len(t) for t in block_tables])
174         block_tables = _make_tensor_with_pad(block_tables,
175                                              max_len=max_block_table_len,
176                                              pad=0,
177                                              dtype=torch.int)
178 
179         input_metadata = InputMetadata(
180             prompt_lens=[],
181             slot_mapping=slot_mapping,
182             max_context_len=max_context_len,
183             context_lens=context_lens,
184             block_tables=block_tables,
185         )
186         return input_tokens, input_positions, input_metadata
187 
188     def _prepare_sample(
189         self,
190         seq_group_metadata_list: List[SequenceGroupMetadata],
191         prompt_lens: List[int],
192     ) -> SamplingMetadata:
193         seq_groups: List[Tuple[List[int], SamplingParams]] = []
194         selected_token_indices: List[int] = []
195         selected_token_start_idx = 0
196         categorized_sample_indices = {t: [] for t in SamplingType}
197         categorized_sample_indices_start_idx = 0
198 
199         max_prompt_len = max(prompt_lens) if prompt_lens else 1
200         for i, seq_group_metadata in enumerate(seq_group_metadata_list):
201             seq_ids = list(seq_group_metadata.seq_data.keys())
202             sampling_params = seq_group_metadata.sampling_params
203             seq_groups.append((seq_ids, sampling_params))
204 
205             if seq_group_metadata.is_prompt:
206                 assert len(seq_ids) == 1
207                 prompt_len = prompt_lens[i]
208                 if sampling_params.prompt_logprobs is not None:
209                     # NOTE: prompt token positions do not need sample, skip
210                     categorized_sample_indices_start_idx += prompt_len - 1
211 
212                 categorized_sample_indices[
213                     sampling_params.sampling_type].append(
214                         categorized_sample_indices_start_idx)
215                 categorized_sample_indices_start_idx += 1
216 
217                 if sampling_params.prompt_logprobs is not None:
218                     selected_token_indices.extend(
219                         range(selected_token_start_idx,
220                               selected_token_start_idx + prompt_len - 1))
221                 selected_token_indices.append(selected_token_start_idx +
222                                               prompt_len - 1)
223                 selected_token_start_idx += max_prompt_len
224             else:
225                 num_seqs = len(seq_ids)
226                 selected_token_indices.extend(
227                     range(selected_token_start_idx,
228                           selected_token_start_idx + num_seqs))
229                 selected_token_start_idx += num_seqs
230 
231                 categorized_sample_indices[
232                     sampling_params.sampling_type].extend(
233                         range(categorized_sample_indices_start_idx,
234                               categorized_sample_indices_start_idx + num_seqs))
235                 categorized_sample_indices_start_idx += num_seqs
236 
237         selected_token_indices = torch.tensor(selected_token_indices,
238                                               dtype=torch.long,
239                                               device="cuda")
240         categorized_sample_indices = {
241             t: torch.tensor(seq_ids, dtype=torch.int, device="cuda")
242             for t, seq_ids in categorized_sample_indices.items()
243         }
244 
245         seq_data: Dict[int, SequenceData] = {}
246         for seq_group_metadata in seq_group_metadata_list:
247             seq_data.update(seq_group_metadata.seq_data)
248 
249         sampling_metadata = SamplingMetadata(
250             seq_groups=seq_groups,
251             seq_data=seq_data,
252             prompt_lens=prompt_lens,
253             selected_token_indices=selected_token_indices,
254             categorized_sample_indices=categorized_sample_indices,
255         )
256         return sampling_metadata
257 
258     @torch.inference_mode()
259     def execute_model(
260         self,
261         seq_group_metadata_list: List[SequenceGroupMetadata],
262         kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],
263         cache_events: Optional[List[torch.cuda.Event]] = None,
264     ) -> SamplerOutput:
265         # NOTE: We assume that all sequences in the group are all prompts or
266         # all decodes.
267         # Prepare input tensors.
268         is_prompt = seq_group_metadata_list[0].is_prompt
269         if is_prompt:
270             inputs = self._prepare_prompt(seq_group_metadata_list)
271             input_tokens, input_positions, input_metadata = inputs
272         else:
273             inputs = self._prepare_decode(seq_group_metadata_list)
274             input_tokens, input_positions, input_metadata = inputs
275         sampling_metadata = self._prepare_sample(seq_group_metadata_list,
276                                                  input_metadata.prompt_lens)
277 
278         # Execute the model.
279         hidden_states = self.model(
280             input_ids=input_tokens,
281             positions=input_positions,
282             kv_caches=kv_caches,
283             input_metadata=input_metadata,
284             cache_events=cache_events,
285         )
286 
287         # Sample the next token.
288         output = self.model.sample(
289             hidden_states=hidden_states,
290             sampling_metadata=sampling_metadata,
291         )
292         return output
293 
294     @torch.inference_mode()
295     def profile_run(self) -> None:
296         # Enable top-k sampling to reflect the accurate memory usage.
297         vocab_size = self.model_config.get_vocab_size()
298         sampling_params = SamplingParams(top_p=0.99, top_k=vocab_size - 1)
299         max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens
300         max_num_seqs = self.scheduler_config.max_num_seqs
301 
302         # Profile memory usage with max_num_sequences sequences and the total
303         # number of tokens equal to max_num_batched_tokens.
304         seqs: List[SequenceGroupMetadata] = []
305         for group_id in range(max_num_seqs):
306             seq_len = (max_num_batched_tokens // max_num_seqs +
307                        (group_id < max_num_batched_tokens % max_num_seqs))
308             seq_data = SequenceData([0] * seq_len)
309             seq = SequenceGroupMetadata(
310                 request_id=str(group_id),
311                 is_prompt=True,
312                 seq_data={group_id: seq_data},
313                 sampling_params=sampling_params,
314                 block_tables=None,
315             )
316             seqs.append(seq)
317 
318         # Run the model with the dummy inputs.
319         num_layers = self.model_config.get_num_layers(self.parallel_config)
320         kv_caches = [(None, None)] * num_layers
321         self.execute_model(seqs, kv_caches)
322         return
323 
324 
325 def _pad_to_max(x: List[int], max_len: int, pad: int) -> List[int]:
326     assert len(x) <= max_len
327     return x + [pad] * (max_len - len(x))
328 
329 
330 def _make_tensor_with_pad(
331     x: List[List[int]],
332     max_len: int,
333     pad: int,
334     dtype: torch.dtype,
335 ) -> torch.Tensor:
336     padded_x = [_pad_to_max(x_i, max_len, pad) for x_i in x]
337     return torch.tensor(padded_x, dtype=dtype, device="cuda")
338
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -134,14 +134,14 @@
                 generation_token = seq_data.get_last_token_id()
                 input_tokens.append([generation_token])
 
-                context_len = seq_data.get_len()
-                if self.sliding_window is not None:
-                    context_len = min(context_len, self.sliding_window)
-                context_lens.append(context_len)
-
-                position = context_len - 1
+                seq_len = seq_data.get_len()
+                position = seq_len - 1
                 input_positions.append([position])
 
+                context_len = seq_len if self.sliding_window is None else min(
+                    seq_len, self.sliding_window)
+                context_lens.append(context_len)
+
                 block_table = seq_group_metadata.block_tables[seq_id]
                 block_number = block_table[position // self.block_size]
                 block_offset = position % self.block_size

{"golden_diff": "diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py\n--- a/vllm/worker/model_runner.py\n+++ b/vllm/worker/model_runner.py\n@@ -134,14 +134,14 @@\n                 generation_token = seq_data.get_last_token_id()\n                 input_tokens.append([generation_token])\n \n-                context_len = seq_data.get_len()\n-                if self.sliding_window is not None:\n-                    context_len = min(context_len, self.sliding_window)\n-                context_lens.append(context_len)\n-\n-                position = context_len - 1\n+                seq_len = seq_data.get_len()\n+                position = seq_len - 1\n                 input_positions.append([position])\n \n+                context_len = seq_len if self.sliding_window is None else min(\n+                    seq_len, self.sliding_window)\n+                context_lens.append(context_len)\n+\n                 block_table = seq_group_metadata.block_tables[seq_id]\n                 block_number = block_table[position // self.block_size]\n                 block_offset = position % self.block_size\n", "issue": "[BUG] Mistral/Mixtral generate nonsense past 4096 tokens in prompt\nIf the prompt contains more than 4k tokens, the model will begin generating nonsense. This seems to be true for both Mistral and Mixtral. I launched the engine at both 32k and 8k max model lengths for testing. Tested with the latest vllm release.\n", "before_files": [{"content": "from typing import Dict, List, Optional, Tuple\n\nimport torch\n\nfrom vllm.config import ModelConfig, ParallelConfig, SchedulerConfig\nfrom vllm.logger import init_logger\nfrom vllm.model_executor import get_model, InputMetadata, SamplingMetadata\nfrom vllm.sampling_params import SamplingParams, SamplingType\nfrom vllm.sequence import SamplerOutput, SequenceData, SequenceGroupMetadata\n\nlogger = init_logger(__name__)\n\n_PAD_SLOT_ID = -1\n\n\nclass ModelRunner:\n\n    def __init__(\n        self,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n        scheduler_config: SchedulerConfig,\n    ):\n        self.model_config = model_config\n        self.parallel_config = parallel_config\n        self.scheduler_config = scheduler_config\n\n        # model_config can be None in tests/samplers/test_sampler.py.\n        # FIXME(woosuk): This is a hack to make the tests work. Refactor this.\n        self.sliding_window = (model_config.get_sliding_window()\n                               if model_config is not None else None)\n        self.model = None\n        self.block_size = None  # Set after initial profiling.\n\n    def load_model(self) -> None:\n        self.model = get_model(self.model_config)\n\n    def set_block_size(self, block_size: int) -> None:\n        self.block_size = block_size\n\n    def _prepare_prompt(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n    ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata]:\n        assert len(seq_group_metadata_list) > 0\n        input_tokens: List[List[int]] = []\n        input_positions: List[List[int]] = []\n        slot_mapping: List[List[int]] = []\n\n        prompt_lens: List[int] = []\n        for seq_group_metadata in seq_group_metadata_list:\n            assert seq_group_metadata.is_prompt\n            seq_ids = list(seq_group_metadata.seq_data.keys())\n            assert len(seq_ids) == 1\n            seq_id = seq_ids[0]\n\n            seq_data = seq_group_metadata.seq_data[seq_id]\n            prompt_tokens = seq_data.get_token_ids()\n            prompt_len = len(prompt_tokens)\n            prompt_lens.append(prompt_len)\n\n            input_tokens.append(prompt_tokens)\n            # NOTE(woosuk): Here we assume that the first token in the prompt\n            # is always the first token in the sequence.\n            input_positions.append(list(range(prompt_len)))\n\n            if seq_group_metadata.block_tables is None:\n                # During memory profiling, the block tables are not initialized\n                # yet. In this case, we just use a dummy slot mapping.\n                slot_mapping.append([_PAD_SLOT_ID] * prompt_len)\n                continue\n\n            # Compute the slot mapping.\n            slot_mapping.append([])\n            block_table = seq_group_metadata.block_tables[seq_id]\n            # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,\n            # where start_idx is max(0, prompt_len - sliding_window).\n            # For example, if the prompt len is 10, sliding window is 8, and\n            # block size is 4, the first two tokens are masked and the slot\n            # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].\n            start_idx = 0\n            if self.sliding_window is not None:\n                start_idx = max(0, prompt_len - self.sliding_window)\n            for i in range(prompt_len):\n                if i < start_idx:\n                    slot_mapping[-1].append(_PAD_SLOT_ID)\n                    continue\n\n                block_number = block_table[i // self.block_size]\n                block_offset = i % self.block_size\n                slot = block_number * self.block_size + block_offset\n                slot_mapping[-1].append(slot)\n\n        max_prompt_len = max(prompt_lens)\n        input_tokens = _make_tensor_with_pad(input_tokens,\n                                             max_prompt_len,\n                                             pad=0,\n                                             dtype=torch.long)\n        input_positions = _make_tensor_with_pad(input_positions,\n                                                max_prompt_len,\n                                                pad=0,\n                                                dtype=torch.long)\n        slot_mapping = _make_tensor_with_pad(slot_mapping,\n                                             max_prompt_len,\n                                             pad=_PAD_SLOT_ID,\n                                             dtype=torch.long)\n\n        input_metadata = InputMetadata(\n            prompt_lens=prompt_lens,\n            slot_mapping=slot_mapping,\n            max_context_len=None,\n            context_lens=None,\n            block_tables=None,\n        )\n        return input_tokens, input_positions, input_metadata\n\n    def _prepare_decode(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n    ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata]:\n        assert len(seq_group_metadata_list) > 0\n        input_tokens: List[List[int]] = []\n        input_positions: List[List[int]] = []\n        slot_mapping: List[List[int]] = []\n        context_lens: List[int] = []\n        block_tables: List[List[int]] = []\n\n        for seq_group_metadata in seq_group_metadata_list:\n            assert not seq_group_metadata.is_prompt\n\n            seq_ids = list(seq_group_metadata.seq_data.keys())\n            for seq_id in seq_ids:\n                seq_data = seq_group_metadata.seq_data[seq_id]\n                generation_token = seq_data.get_last_token_id()\n                input_tokens.append([generation_token])\n\n                context_len = seq_data.get_len()\n                if self.sliding_window is not None:\n                    context_len = min(context_len, self.sliding_window)\n                context_lens.append(context_len)\n\n                position = context_len - 1\n                input_positions.append([position])\n\n                block_table = seq_group_metadata.block_tables[seq_id]\n                block_number = block_table[position // self.block_size]\n                block_offset = position % self.block_size\n                slot = block_number * self.block_size + block_offset\n                slot_mapping.append([slot])\n\n                if self.sliding_window is not None:\n                    sliding_window_blocks = (self.sliding_window //\n                                             self.block_size)\n                    block_table = block_table[-sliding_window_blocks:]\n                block_tables.append(block_table)\n\n        input_tokens = _make_tensor_with_pad(input_tokens,\n                                             max_len=1,\n                                             pad=0,\n                                             dtype=torch.long)\n        input_positions = _make_tensor_with_pad(input_positions,\n                                                max_len=1,\n                                                pad=0,\n                                                dtype=torch.long)\n        slot_mapping = _make_tensor_with_pad(slot_mapping,\n                                             max_len=1,\n                                             pad=_PAD_SLOT_ID,\n                                             dtype=torch.long)\n        max_context_len = max(context_lens)\n        context_lens = torch.tensor(context_lens,\n                                    dtype=torch.int,\n                                    device=\"cuda\")\n        max_block_table_len = max([len(t) for t in block_tables])\n        block_tables = _make_tensor_with_pad(block_tables,\n                                             max_len=max_block_table_len,\n                                             pad=0,\n                                             dtype=torch.int)\n\n        input_metadata = InputMetadata(\n            prompt_lens=[],\n            slot_mapping=slot_mapping,\n            max_context_len=max_context_len,\n            context_lens=context_lens,\n            block_tables=block_tables,\n        )\n        return input_tokens, input_positions, input_metadata\n\n    def _prepare_sample(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n        prompt_lens: List[int],\n    ) -> SamplingMetadata:\n        seq_groups: List[Tuple[List[int], SamplingParams]] = []\n        selected_token_indices: List[int] = []\n        selected_token_start_idx = 0\n        categorized_sample_indices = {t: [] for t in SamplingType}\n        categorized_sample_indices_start_idx = 0\n\n        max_prompt_len = max(prompt_lens) if prompt_lens else 1\n        for i, seq_group_metadata in enumerate(seq_group_metadata_list):\n            seq_ids = list(seq_group_metadata.seq_data.keys())\n            sampling_params = seq_group_metadata.sampling_params\n            seq_groups.append((seq_ids, sampling_params))\n\n            if seq_group_metadata.is_prompt:\n                assert len(seq_ids) == 1\n                prompt_len = prompt_lens[i]\n                if sampling_params.prompt_logprobs is not None:\n                    # NOTE: prompt token positions do not need sample, skip\n                    categorized_sample_indices_start_idx += prompt_len - 1\n\n                categorized_sample_indices[\n                    sampling_params.sampling_type].append(\n                        categorized_sample_indices_start_idx)\n                categorized_sample_indices_start_idx += 1\n\n                if sampling_params.prompt_logprobs is not None:\n                    selected_token_indices.extend(\n                        range(selected_token_start_idx,\n                              selected_token_start_idx + prompt_len - 1))\n                selected_token_indices.append(selected_token_start_idx +\n                                              prompt_len - 1)\n                selected_token_start_idx += max_prompt_len\n            else:\n                num_seqs = len(seq_ids)\n                selected_token_indices.extend(\n                    range(selected_token_start_idx,\n                          selected_token_start_idx + num_seqs))\n                selected_token_start_idx += num_seqs\n\n                categorized_sample_indices[\n                    sampling_params.sampling_type].extend(\n                        range(categorized_sample_indices_start_idx,\n                              categorized_sample_indices_start_idx + num_seqs))\n                categorized_sample_indices_start_idx += num_seqs\n\n        selected_token_indices = torch.tensor(selected_token_indices,\n                                              dtype=torch.long,\n                                              device=\"cuda\")\n        categorized_sample_indices = {\n            t: torch.tensor(seq_ids, dtype=torch.int, device=\"cuda\")\n            for t, seq_ids in categorized_sample_indices.items()\n        }\n\n        seq_data: Dict[int, SequenceData] = {}\n        for seq_group_metadata in seq_group_metadata_list:\n            seq_data.update(seq_group_metadata.seq_data)\n\n        sampling_metadata = SamplingMetadata(\n            seq_groups=seq_groups,\n            seq_data=seq_data,\n            prompt_lens=prompt_lens,\n            selected_token_indices=selected_token_indices,\n            categorized_sample_indices=categorized_sample_indices,\n        )\n        return sampling_metadata\n\n    @torch.inference_mode()\n    def execute_model(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n        kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],\n        cache_events: Optional[List[torch.cuda.Event]] = None,\n    ) -> SamplerOutput:\n        # NOTE: We assume that all sequences in the group are all prompts or\n        # all decodes.\n        # Prepare input tensors.\n        is_prompt = seq_group_metadata_list[0].is_prompt\n        if is_prompt:\n            inputs = self._prepare_prompt(seq_group_metadata_list)\n            input_tokens, input_positions, input_metadata = inputs\n        else:\n            inputs = self._prepare_decode(seq_group_metadata_list)\n            input_tokens, input_positions, input_metadata = inputs\n        sampling_metadata = self._prepare_sample(seq_group_metadata_list,\n                                                 input_metadata.prompt_lens)\n\n        # Execute the model.\n        hidden_states = self.model(\n            input_ids=input_tokens,\n            positions=input_positions,\n            kv_caches=kv_caches,\n            input_metadata=input_metadata,\n            cache_events=cache_events,\n        )\n\n        # Sample the next token.\n        output = self.model.sample(\n            hidden_states=hidden_states,\n            sampling_metadata=sampling_metadata,\n        )\n        return output\n\n    @torch.inference_mode()\n    def profile_run(self) -> None:\n        # Enable top-k sampling to reflect the accurate memory usage.\n        vocab_size = self.model_config.get_vocab_size()\n        sampling_params = SamplingParams(top_p=0.99, top_k=vocab_size - 1)\n        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens\n        max_num_seqs = self.scheduler_config.max_num_seqs\n\n        # Profile memory usage with max_num_sequences sequences and the total\n        # number of tokens equal to max_num_batched_tokens.\n        seqs: List[SequenceGroupMetadata] = []\n        for group_id in range(max_num_seqs):\n            seq_len = (max_num_batched_tokens // max_num_seqs +\n                       (group_id < max_num_batched_tokens % max_num_seqs))\n            seq_data = SequenceData([0] * seq_len)\n            seq = SequenceGroupMetadata(\n                request_id=str(group_id),\n                is_prompt=True,\n                seq_data={group_id: seq_data},\n                sampling_params=sampling_params,\n                block_tables=None,\n            )\n            seqs.append(seq)\n\n        # Run the model with the dummy inputs.\n        num_layers = self.model_config.get_num_layers(self.parallel_config)\n        kv_caches = [(None, None)] * num_layers\n        self.execute_model(seqs, kv_caches)\n        return\n\n\ndef _pad_to_max(x: List[int], max_len: int, pad: int) -> List[int]:\n    assert len(x) <= max_len\n    return x + [pad] * (max_len - len(x))\n\n\ndef _make_tensor_with_pad(\n    x: List[List[int]],\n    max_len: int,\n    pad: int,\n    dtype: torch.dtype,\n) -> torch.Tensor:\n    padded_x = [_pad_to_max(x_i, max_len, pad) for x_i in x]\n    return torch.tensor(padded_x, dtype=dtype, device=\"cuda\")\n", "path": "vllm/worker/model_runner.py"}], "after_files": [{"content": "from typing import Dict, List, Optional, Tuple\n\nimport torch\n\nfrom vllm.config import ModelConfig, ParallelConfig, SchedulerConfig\nfrom vllm.logger import init_logger\nfrom vllm.model_executor import get_model, InputMetadata, SamplingMetadata\nfrom vllm.sampling_params import SamplingParams, SamplingType\nfrom vllm.sequence import SamplerOutput, SequenceData, SequenceGroupMetadata\n\nlogger = init_logger(__name__)\n\n_PAD_SLOT_ID = -1\n\n\nclass ModelRunner:\n\n    def __init__(\n        self,\n        model_config: ModelConfig,\n        parallel_config: ParallelConfig,\n        scheduler_config: SchedulerConfig,\n    ):\n        self.model_config = model_config\n        self.parallel_config = parallel_config\n        self.scheduler_config = scheduler_config\n\n        # model_config can be None in tests/samplers/test_sampler.py.\n        # FIXME(woosuk): This is a hack to make the tests work. Refactor this.\n        self.sliding_window = (model_config.get_sliding_window()\n                               if model_config is not None else None)\n        self.model = None\n        self.block_size = None  # Set after initial profiling.\n\n    def load_model(self) -> None:\n        self.model = get_model(self.model_config)\n\n    def set_block_size(self, block_size: int) -> None:\n        self.block_size = block_size\n\n    def _prepare_prompt(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n    ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata]:\n        assert len(seq_group_metadata_list) > 0\n        input_tokens: List[List[int]] = []\n        input_positions: List[List[int]] = []\n        slot_mapping: List[List[int]] = []\n\n        prompt_lens: List[int] = []\n        for seq_group_metadata in seq_group_metadata_list:\n            assert seq_group_metadata.is_prompt\n            seq_ids = list(seq_group_metadata.seq_data.keys())\n            assert len(seq_ids) == 1\n            seq_id = seq_ids[0]\n\n            seq_data = seq_group_metadata.seq_data[seq_id]\n            prompt_tokens = seq_data.get_token_ids()\n            prompt_len = len(prompt_tokens)\n            prompt_lens.append(prompt_len)\n\n            input_tokens.append(prompt_tokens)\n            # NOTE(woosuk): Here we assume that the first token in the prompt\n            # is always the first token in the sequence.\n            input_positions.append(list(range(prompt_len)))\n\n            if seq_group_metadata.block_tables is None:\n                # During memory profiling, the block tables are not initialized\n                # yet. In this case, we just use a dummy slot mapping.\n                slot_mapping.append([_PAD_SLOT_ID] * prompt_len)\n                continue\n\n            # Compute the slot mapping.\n            slot_mapping.append([])\n            block_table = seq_group_metadata.block_tables[seq_id]\n            # Mask the [0, start_idx) tokens of the prompt with _PAD_SLOT_ID,\n            # where start_idx is max(0, prompt_len - sliding_window).\n            # For example, if the prompt len is 10, sliding window is 8, and\n            # block size is 4, the first two tokens are masked and the slot\n            # mapping will be [-1, -1, 2, 3, 4, 5, 6, 7, 0, 1].\n            start_idx = 0\n            if self.sliding_window is not None:\n                start_idx = max(0, prompt_len - self.sliding_window)\n            for i in range(prompt_len):\n                if i < start_idx:\n                    slot_mapping[-1].append(_PAD_SLOT_ID)\n                    continue\n\n                block_number = block_table[i // self.block_size]\n                block_offset = i % self.block_size\n                slot = block_number * self.block_size + block_offset\n                slot_mapping[-1].append(slot)\n\n        max_prompt_len = max(prompt_lens)\n        input_tokens = _make_tensor_with_pad(input_tokens,\n                                             max_prompt_len,\n                                             pad=0,\n                                             dtype=torch.long)\n        input_positions = _make_tensor_with_pad(input_positions,\n                                                max_prompt_len,\n                                                pad=0,\n                                                dtype=torch.long)\n        slot_mapping = _make_tensor_with_pad(slot_mapping,\n                                             max_prompt_len,\n                                             pad=_PAD_SLOT_ID,\n                                             dtype=torch.long)\n\n        input_metadata = InputMetadata(\n            prompt_lens=prompt_lens,\n            slot_mapping=slot_mapping,\n            max_context_len=None,\n            context_lens=None,\n            block_tables=None,\n        )\n        return input_tokens, input_positions, input_metadata\n\n    def _prepare_decode(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n    ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata]:\n        assert len(seq_group_metadata_list) > 0\n        input_tokens: List[List[int]] = []\n        input_positions: List[List[int]] = []\n        slot_mapping: List[List[int]] = []\n        context_lens: List[int] = []\n        block_tables: List[List[int]] = []\n\n        for seq_group_metadata in seq_group_metadata_list:\n            assert not seq_group_metadata.is_prompt\n\n            seq_ids = list(seq_group_metadata.seq_data.keys())\n            for seq_id in seq_ids:\n                seq_data = seq_group_metadata.seq_data[seq_id]\n                generation_token = seq_data.get_last_token_id()\n                input_tokens.append([generation_token])\n\n                seq_len = seq_data.get_len()\n                position = seq_len - 1\n                input_positions.append([position])\n\n                context_len = seq_len if self.sliding_window is None else min(\n                    seq_len, self.sliding_window)\n                context_lens.append(context_len)\n\n                block_table = seq_group_metadata.block_tables[seq_id]\n                block_number = block_table[position // self.block_size]\n                block_offset = position % self.block_size\n                slot = block_number * self.block_size + block_offset\n                slot_mapping.append([slot])\n\n                if self.sliding_window is not None:\n                    sliding_window_blocks = (self.sliding_window //\n                                             self.block_size)\n                    block_table = block_table[-sliding_window_blocks:]\n                block_tables.append(block_table)\n\n        input_tokens = _make_tensor_with_pad(input_tokens,\n                                             max_len=1,\n                                             pad=0,\n                                             dtype=torch.long)\n        input_positions = _make_tensor_with_pad(input_positions,\n                                                max_len=1,\n                                                pad=0,\n                                                dtype=torch.long)\n        slot_mapping = _make_tensor_with_pad(slot_mapping,\n                                             max_len=1,\n                                             pad=_PAD_SLOT_ID,\n                                             dtype=torch.long)\n        max_context_len = max(context_lens)\n        context_lens = torch.tensor(context_lens,\n                                    dtype=torch.int,\n                                    device=\"cuda\")\n        max_block_table_len = max([len(t) for t in block_tables])\n        block_tables = _make_tensor_with_pad(block_tables,\n                                             max_len=max_block_table_len,\n                                             pad=0,\n                                             dtype=torch.int)\n\n        input_metadata = InputMetadata(\n            prompt_lens=[],\n            slot_mapping=slot_mapping,\n            max_context_len=max_context_len,\n            context_lens=context_lens,\n            block_tables=block_tables,\n        )\n        return input_tokens, input_positions, input_metadata\n\n    def _prepare_sample(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n        prompt_lens: List[int],\n    ) -> SamplingMetadata:\n        seq_groups: List[Tuple[List[int], SamplingParams]] = []\n        selected_token_indices: List[int] = []\n        selected_token_start_idx = 0\n        categorized_sample_indices = {t: [] for t in SamplingType}\n        categorized_sample_indices_start_idx = 0\n\n        max_prompt_len = max(prompt_lens) if prompt_lens else 1\n        for i, seq_group_metadata in enumerate(seq_group_metadata_list):\n            seq_ids = list(seq_group_metadata.seq_data.keys())\n            sampling_params = seq_group_metadata.sampling_params\n            seq_groups.append((seq_ids, sampling_params))\n\n            if seq_group_metadata.is_prompt:\n                assert len(seq_ids) == 1\n                prompt_len = prompt_lens[i]\n                if sampling_params.prompt_logprobs is not None:\n                    # NOTE: prompt token positions do not need sample, skip\n                    categorized_sample_indices_start_idx += prompt_len - 1\n\n                categorized_sample_indices[\n                    sampling_params.sampling_type].append(\n                        categorized_sample_indices_start_idx)\n                categorized_sample_indices_start_idx += 1\n\n                if sampling_params.prompt_logprobs is not None:\n                    selected_token_indices.extend(\n                        range(selected_token_start_idx,\n                              selected_token_start_idx + prompt_len - 1))\n                selected_token_indices.append(selected_token_start_idx +\n                                              prompt_len - 1)\n                selected_token_start_idx += max_prompt_len\n            else:\n                num_seqs = len(seq_ids)\n                selected_token_indices.extend(\n                    range(selected_token_start_idx,\n                          selected_token_start_idx + num_seqs))\n                selected_token_start_idx += num_seqs\n\n                categorized_sample_indices[\n                    sampling_params.sampling_type].extend(\n                        range(categorized_sample_indices_start_idx,\n                              categorized_sample_indices_start_idx + num_seqs))\n                categorized_sample_indices_start_idx += num_seqs\n\n        selected_token_indices = torch.tensor(selected_token_indices,\n                                              dtype=torch.long,\n                                              device=\"cuda\")\n        categorized_sample_indices = {\n            t: torch.tensor(seq_ids, dtype=torch.int, device=\"cuda\")\n            for t, seq_ids in categorized_sample_indices.items()\n        }\n\n        seq_data: Dict[int, SequenceData] = {}\n        for seq_group_metadata in seq_group_metadata_list:\n            seq_data.update(seq_group_metadata.seq_data)\n\n        sampling_metadata = SamplingMetadata(\n            seq_groups=seq_groups,\n            seq_data=seq_data,\n            prompt_lens=prompt_lens,\n            selected_token_indices=selected_token_indices,\n            categorized_sample_indices=categorized_sample_indices,\n        )\n        return sampling_metadata\n\n    @torch.inference_mode()\n    def execute_model(\n        self,\n        seq_group_metadata_list: List[SequenceGroupMetadata],\n        kv_caches: List[Tuple[torch.Tensor, torch.Tensor]],\n        cache_events: Optional[List[torch.cuda.Event]] = None,\n    ) -> SamplerOutput:\n        # NOTE: We assume that all sequences in the group are all prompts or\n        # all decodes.\n        # Prepare input tensors.\n        is_prompt = seq_group_metadata_list[0].is_prompt\n        if is_prompt:\n            inputs = self._prepare_prompt(seq_group_metadata_list)\n            input_tokens, input_positions, input_metadata = inputs\n        else:\n            inputs = self._prepare_decode(seq_group_metadata_list)\n            input_tokens, input_positions, input_metadata = inputs\n        sampling_metadata = self._prepare_sample(seq_group_metadata_list,\n                                                 input_metadata.prompt_lens)\n\n        # Execute the model.\n        hidden_states = self.model(\n            input_ids=input_tokens,\n            positions=input_positions,\n            kv_caches=kv_caches,\n            input_metadata=input_metadata,\n            cache_events=cache_events,\n        )\n\n        # Sample the next token.\n        output = self.model.sample(\n            hidden_states=hidden_states,\n            sampling_metadata=sampling_metadata,\n        )\n        return output\n\n    @torch.inference_mode()\n    def profile_run(self) -> None:\n        # Enable top-k sampling to reflect the accurate memory usage.\n        vocab_size = self.model_config.get_vocab_size()\n        sampling_params = SamplingParams(top_p=0.99, top_k=vocab_size - 1)\n        max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens\n        max_num_seqs = self.scheduler_config.max_num_seqs\n\n        # Profile memory usage with max_num_sequences sequences and the total\n        # number of tokens equal to max_num_batched_tokens.\n        seqs: List[SequenceGroupMetadata] = []\n        for group_id in range(max_num_seqs):\n            seq_len = (max_num_batched_tokens // max_num_seqs +\n                       (group_id < max_num_batched_tokens % max_num_seqs))\n            seq_data = SequenceData([0] * seq_len)\n            seq = SequenceGroupMetadata(\n                request_id=str(group_id),\n                is_prompt=True,\n                seq_data={group_id: seq_data},\n                sampling_params=sampling_params,\n                block_tables=None,\n            )\n            seqs.append(seq)\n\n        # Run the model with the dummy inputs.\n        num_layers = self.model_config.get_num_layers(self.parallel_config)\n        kv_caches = [(None, None)] * num_layers\n        self.execute_model(seqs, kv_caches)\n        return\n\n\ndef _pad_to_max(x: List[int], max_len: int, pad: int) -> List[int]:\n    assert len(x) <= max_len\n    return x + [pad] * (max_len - len(x))\n\n\ndef _make_tensor_with_pad(\n    x: List[List[int]],\n    max_len: int,\n    pad: int,\n    dtype: torch.dtype,\n) -> torch.Tensor:\n    padded_x = [_pad_to_max(x_i, max_len, pad) for x_i in x]\n    return torch.tensor(padded_x, dtype=dtype, device=\"cuda\")\n", "path": "vllm/worker/model_runner.py"}]}