Support for vLLM 0.6.3

.devcontainer/Dockerfile

@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
+FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
 RUN apt-get update && apt-get install -y libgl1 libglib2.0-0 ffmpeg locales
 
 # Set the locale

aana/deployments/hqq_text_generation_deployment.py

@@ -118,7 +118,7 @@ async def apply_config(self, config: dict[str, Any]):
                 self.dtype = Dtype.BFLOAT16
 
         if config_obj.quantize_on_fly:
-            self.model_kwargs["device_map"] = "cpu"
+            self.model_kwargs.pop("device_map", None)
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_id, torch_dtype=self.dtype.to_torch(), **self.model_kwargs
             )

aana/deployments/vllm_deployment.py

@@ -256,9 +256,9 @@ async def generate_stream(  # noqa: C901
             else:
                 inputs = TokensPrompt(prompt_token_ids=prompt_token_ids)
             results_generator = self.engine.generate(
+                inputs,
                 sampling_params=sampling_params_vllm,
                 request_id=request_id,
-                inputs=inputs,
             )
 
             num_returned = 0