fix lora and prompt adapter

vllm-project · Jul 22, 2024 · f386f69 · f386f69
1 parent 4d4eacc
commit f386f69
Showing 1 changed file with 17 additions and 9 deletions.
diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
@@ -530,12 +530,12 @@ def build(self) -> ModelInputForGPU:
             lora_requests = set(r for data in self.inter_data_list
                                 for r in data.lora_requests)
             lora_index_mapping = flatten_2d_lists([
-                inter_data.lora_index_mapping
+                flatten_2d_lists(inter_data.lora_index_mapping)
                 for inter_data in self.inter_data_list
             ])
             lora_index_mapping.extend([0] * cuda_graph_pad_size)
             lora_prompt_mapping = flatten_2d_lists([
-                inter_data.lora_prompt_mapping
+                flatten_2d_lists(inter_data.lora_prompt_mapping)
                 for inter_data in self.inter_data_list
             ])
             lora_mapping = LoRAMapping(
@@ -546,14 +546,22 @@ def build(self) -> ModelInputForGPU:
         # Prompt adapter data.
         prompt_adapter_requests: Set[PromptAdapterRequest] = set()
         prompt_adapter_mapping = None
-        if self.enable_prompt_adapter and self.inter_data_list[0].is_prompt:
-            assert len(self.inter_data_list) == 1
-            if self.inter_data_list[0].prompt_adapter_request is not None:
-                prompt_adapter_requests = set(
-                    [self.inter_data_list[0].prompt_adapter_request])
+        if self.enable_prompt_adapter:
+            prompt_adapter_requests = set(
+                data.prompt_adapter_request for data in self.inter_data_list
+                if data.prompt_adapter_request is not None)
+            prompt_adapter_index_mapping = flatten_2d_lists([
+                inter_data.prompt_adapter_index_mapping
+                for inter_data in self.inter_data_list
+            ])
+            prompt_adapter_index_mapping.extend([0] * cuda_graph_pad_size)
+            prompt_adapter_prompt_mapping = flatten_2d_lists([
+                inter_data.prompt_adapter_prompt_mapping
+                for inter_data in self.inter_data_list
+            ])
             prompt_adapter_mapping = PromptAdapterMapping(
-                self.inter_data_list[0].prompt_adapter_index_mapping,
-                self.inter_data_list[0].prompt_adapter_prompt_mapping,
+                prompt_adapter_index_mapping,
+                prompt_adapter_prompt_mapping,
             )
 
         # Multi-modal data.