[TPU][V1][Bugfix] Fix chunked prefill with padding

[NickLucche]

Currently if you run:

VLLM_USE_V1=1 vllm serve llava-hf/llava-1.5-7b-hf --max-model-len 2512 --max-num-seqs 16 --max-num-batched-tokens 128 --chat-template examples/template_llava.jinja --enforce-eager &

# query 
python examples/online_serving/openai_chat_completion_client_for_multimodal.py


[ ... ]
ERROR 03-18 13:12:42 [core.py:340]   File "/home/nick/work/vllm/vllm/v1/worker/tpu_model_runner.py", line 564, in execute_model
ERROR 03-18 13:12:42 [core.py:340]     inputs_embeds = self.model.get_input_embeddings(
ERROR 03-18 13:12:42 [core.py:340]                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-18 13:12:42 [core.py:340]   File "/home/nick/work/vllm/vllm/v1/worker/tpu_model_runner.py", line 862, in get_input_embeddings
ERROR 03-18 13:12:42 [core.py:340]     return self.model.get_input_embeddings(*args, **kwargs)
ERROR 03-18 13:12:42 [core.py:340]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-18 13:12:42 [core.py:340]   File "/home/nick/work/vllm/vllm/model_executor/models/llava.py", line 779, in get_input_embeddings
ERROR 03-18 13:12:42 [core.py:340]     inputs_embeds = merge_multimodal_embeddings(
ERROR 03-18 13:12:42 [core.py:340]                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-18 13:12:42 [core.py:340]   File "/home/nick/work/vllm/vllm/model_executor/models/utils.py", line 455, in merge_multimodal_embeddings
ERROR 03-18 13:12:42 [core.py:340]     return _merge_multimodal_embeddings(
ERROR 03-18 13:12:42 [core.py:340]            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ERROR 03-18 13:12:42 [core.py:340]   File "/home/nick/work/vllm/vllm/model_executor/models/utils.py", line 371, in _merge_multimodal_embeddings
ERROR 03-18 13:12:42 [core.py:340]     raise ValueError(
ERROR 03-18 13:12:42 [core.py:340] ValueError: Attempted to assign 69 = 69 multimodal tokens to 113 placeholders
ERROR 03-18 13:12:42 [core.py:340] 
CRITICAL 03-18 13:12:42 [core_client.py:269] Got fatal signal from worker processes, shutting down. See stack trace above for root cause issue.
[1]    99083 killed     VLLM_USE_V1=1 vllm serve llava-hf/llava-1.5-7b-hf --max-model-len 2512  16   

I tried to summarize the issue in this (poorly drawn) diagram:

This is most noticeable with multimodal because spurious image tokens lingering on a re-used tensor will cause the server to crash, but the same happens with text-only.
Basically when we use CP in combination with padding, we must treat the last/terminal chunk carefully, otherwise it will contain values from previous iterations as the buffer (input_ids_cpu in this case) is re-used.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[DarkLight1337]

Please fix the pre-commit failures