[Bugfix][TPU] Fix pad slot id (vllm-project#5977)

opendatahub-io · Jul 1, 2024 · 9ea7506 · 9ea7506
1 parent 28773bf
commit 9ea7506
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/vllm/worker/tpu_model_runner.py b/vllm/worker/tpu_model_runner.py
@@ -19,7 +19,7 @@
 
 logger = init_logger(__name__)
 
-_PAD_SLOT_ID = 0  # FIXME(woosuk)
+_PAD_SLOT_ID = -1  # NOTE(woosuk): In PyTorch XLA, index -1 is ignored.
 # FIXME(woosuk): Temporarily disabled top-p sampling since it's too slow.
 _ENABLE_TOP_P = False
 # FIXME(woosuk): A temporary hack to support `n > 1`.