fix: TPU multimodal kwargs and outlines installation in TPU docker

Dockerfile.tpu

@@ -4,9 +4,7 @@ ARG BASE_IMAGE="us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:night
 FROM $BASE_IMAGE
 
 WORKDIR /workspace
-COPY . /workspace/aphrodite-engine
 
-ENV APHRODITE_TARGET_DEVICE="tpu"
 # Install aiohttp separately to avoid build errors.
 RUN pip install aiohttp
 # Install the TPU and Pallas dependencies.
@@ -14,6 +12,13 @@ RUN pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases
 RUN pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
 
 # Build Aphrodite.
+COPY . /workspace/aphrodite-engine
+ENV APHRODITE_TARGET_DEVICE="tpu"
 RUN cd /workspace/aphrodite-engine && python setup.py develop
 
+# Re-install outlines to avoid dependency errors.
+# The outlines version must follow requirements-common.txt.
+RUN pip uninstall outlines -y
+RUN pip install "outlines>=0.0.43"
+
 CMD ["/bin/bash"]

aphrodite/task_handler/tpu_model_runner.py

@@ -153,7 +153,7 @@ class TPUModelRunner:
         # Dummy run.
         num_samples = _MAX_NUM_SAMPLES if is_prompt else 1
         self.model(token_ids, position_ids, kv_caches, attn_metadata,
-                   input_lens, t, p, num_samples)
+                   input_lens, None, t, p, num_samples)
 
     def warmup_model(
         self,