[Bugfix] LoRA V0 - Fix case where max_num_seqs is between cudagraph capture sizes

[varun-sundar-rabindranath]

FIX #15269

Repro command :

VLLM_USE_V1=0 python3 benchmarks/benchmark_throughput.py --model meta-llama/Llama-2-7b-hf --backend vllm --dataset ./ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 200 --max-loras 4 --max-lora-rank 8 --lora-path yard1/llama-2-7b-sql-lora-test --enable-lora --max-num-seqs 50

Bug:

...
  File "/home/shadeform/precog/.venv/lib/python3.12/site-packages/vllm/lora/punica_wrapper/punica_gpu.py", line 66, in update_metadata
    self.prompt_mapping_meta.prepare_tensors(self.sampler_indices)
  File "/home/shadeform/precog/.venv/lib/python3.12/site-packages/vllm/lora/ops/triton_ops/lora_kernel_metadata.py", line 76, in prepare_tensors
    self.token_lora_mapping[:num_tokens].copy_(token_lora_mapping,J
RuntimeError: The size of tensor a (50) must match the size of tensor b (56) at non-singleton dimension 0

PR that introduced the bug: #14685

Affects : V0 engine. V1 engine is fine.

Fix / Why ? :
In V0, during cudagraph capture, the cuda graph capture size could be greater than the max_num_seqs setting. In #14685 , we assume that max_num_seqs will always be respected. This assumption is true for V1, but not for V0.
The line where the error occurs deals with the LogitsProcessor. Before #14685, _sampler_indices at

vllm/vllm/lora/punica_wrapper/punica_base.py

Line 136 in cfbb8c9

self._sampler_indices = torch.empty(max_num_batched_tokens,

was used in the place of token_lora_mapping . Before #14685 we seem to have handled this issue by just allocating a buffer as big as max_num_batched_tokens. We use the same fix here.

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

🚀

[varun-sundar-rabindranath]

Please take a look @jeejeelee. Thanks ! 🙌

[jeejeelee]

After the test modifications, overall LGTM