[TPU][V1] Fix padding recompilation when max-num-batched-tokens is not even

[NickLucche]

max-num-batched-tokens values that are not power of 2s (or simply not even when using bucketing) can silently cause recompiliations.
This is due to the fact in such cases both bucketed and exponential padding will return a maximal padding value that will be > max-num-batched-tokens.
Assume:

# 317
self.max_num_tokens = scheduler_config.max_num_batched_tokens
self.num_tokens_paddings = _get_token_paddings(..)
print(self.num_tokens_paddings[-1]) # 512 for exponential, 320 for bucketed

In turn, auxiliary data structures will be instantiated with the original uneven value (say 317)

self.input_ids_cpu = torch.zeros(self.max_num_tokens, # 317
                                 dtype=torch.int32,
                                 device="cpu")

Causing the following

# Padded value=512, which is > 317
padded_total_num_scheduled_tokens = _get_padded_token_len(
    self.num_tokens_paddings, total_num_scheduled_tokens)
...
# Select 512 positions in 317 position array will yield the full 317 positions
self.input_ids = self.input_ids_cpu[:
                                    padded_total_num_scheduled_tokens].to(
                                        self.device)
# Causing `input_ids` to run with a potentially uncompiled size of 317

To verify:

VLLM_XLA_CACHE_PATH= VLLM_XLA_CHECK_RECOMPILATION=1 VLLM_USE_V1=1 vllm serve Qwen/Qwen2.5-1.5B-Instruct \
 --port 8004 \
 --gpu-memory-utilization 0.95 \
 --max-num-seqs 8 \
 --max-num-batched-tokens 93 \
 --tensor-parallel-size 1 \
 --max-model-len 256
 
 
 # Some long request
 curl http://localhost:8004/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen2.5-1.5B-Instruct",
    "messages": [
      {
        "role": "system",
        "content": "You are a helpful assistant."
      },
      {
        "role": "user",
        "content": "Imagine a lone samurai wandering across feudal Japan, passing through misty mountains, quiet villages, and fields touched by the morning dew. He reflects on the impermanence of life, the fleeting nature of honor, and the silent strength required to walk a solitary path. His sword is sheathed not just in steel, but in wisdom. His heart carries both the burden of war and the beauty of cherry blossoms falling in the spring breeze. He meets no one, yet learns from every shadow, every whisper of wind, and every broken piece of pottery he finds along the way. Please write a short poem that captures the spirit of this samurai journey and his internal reflections. The tone should be serene, slightly melancholic, yet filled with a sense of dignity and resolve."
      }
    ],
    "temperature": 0.4,
    "min_p": 0.8,
    "max_tokens": 48,
    "stream": false
  }'

We can either have the last padding element be always exactly equal to max-num-batched-tokens or adjust for its "actual" padded value computed by _get_token_paddings (a power of 2 or "bucketed" power of 2).

This PR implements the latter because I ultimately think having these sizes be aligned can help with memory accesses here and there.
I am open to discuss better approaches if you have suggestions.

[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.

🚀

[yaochengji]

I think we should also add a test when padding_gap == 0?

[NickLucche]

I did in a previous PR it's below. That is for exponential padding.

[NickLucche]

Thanks for reviewing @yaochengji .

[yaochengji]

LGTM, thanks for the fix!