[BugFix] Fix the missing value for argument 'no_lora_flag_cpu' error when running benchmark_lora.py

[classicsong]

Purpose

FIx the bug #25013 by adding the argument no_lora_flag_cpu when preparing the input arguments in as_lora_shrink_kwargs and as_lora_expand_kwargs.

Test Plan

python3 benchmarks/kernels/benchmark_lora.py model_bench --models meta-llama/Llama-3-8b  --arg-pool-size 32 --batch-sizes 32 --dtype torch.float16  --lora-ranks 16 --num-loras 1 4 --op-types lora_shrink lora_expand --seq-lengths 16 --sort-by-lora-id 1 --cuda-graph-nops 32 

Test Result

Output after the bug fix

Namespace(cmd='model_bench', models=['meta-llama/Llama-3-8b'], tp_sizes=[1], lora_ranks=[16], dtype=torch.float16, arg_pool_size=32, cuda_graph_nops=32, num_loras=[1, 4], num_active_loras=None, sort_by_lora_id=[True], op_types=[<OpType.LORA_SHRINK: 1>, <OpType.LORA_EXPAND: 2>], seq_lengths=[16], batch_sizes=[32], expand_fn_add_inputs=[True, False], output_directory=None, test_correctness=False, func=<function run_model_bench at 0x734bc084df80>)
Model bench :
 Hidden Sizes {6144, 4096, 28672} LoRA Ranks [16]
Benchmarking 32 invocations inside a CUDA Graph
[-------------------------------------------------------------------------------------------------------------------------------- lora-torch.float16 | cugraph 32 ops ---------------------------------------------------------------------------------------------------------------------------------]
                                                                                                                 |  single-lora roofline using torch.mm (f16xf16=>f16)  |  LORA_SHRINK() (f16xf16=>f32)  |  LORA_EXPAND(add_inputs=True) (f32xf16=>f16)  |  LORA_EXPAND(add_inputs=False) (f32xf16=>f16)
1 threads: ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
      {"bs": 32, "sl": 16, "m": 512, "k": 6144, "n": 16, "num_loras": 1, "sort_by_lora": true, "num_slices": 1}  |                        228.4                         |             602.9              |                                               |                                              
      {"bs": 32, "sl": 16, "m": 512, "k": 6144, "n": 16, "num_loras": 1, "sort_by_lora": true, "num_slices": 2}  |                        233.8                         |             862.0              |                                               |                                              
      {"bs": 32, "sl": 16, "m": 512, "k": 6144, "n": 16, "num_loras": 1, "sort_by_lora": true, "num_slices": 3}  |                        278.2                         |             844.9              |                                               |                                              
      {"bs": 32, "sl": 16, "m": 512, "k": 16, "n": 6144, "num_loras": 1, "sort_by_lora": true, "num_slices": 1}  |                        168.7                         |                                |                      275.8                    |                     183.5                    
      {"bs": 32, "sl": 16, "m": 512, "k": 16, "n": 6144, "num_loras": 1, "sort_by_lora": true, "num_slices": 2}  |                        242.0                         |                                |                      753.5                    |                     496.3                    
      {"bs": 32, "sl": 16, "m": 512, "k": 16, "n": 6144, "num_loras": 1, "sort_by_lora": true, "num_slices": 3}  |                        410.7                         |                                |                     1108.8                    |                     723.2                    

Times are in microseconds (us).

Note on Comparison with torch.mm : The torch.mm numbers are benchmark numbers of a simple matmul emulating the single lora case. It is provided as a roofline for comparing our LoRA Kernel implementations. It is expected that the LoRA kernels will be slower than torch.mm in cases where num_loras is big. But for small num_loras the goal should be to match the torch.mm numbers.

...

Output before the bug fix

Benchmarking 32 invocations inside a CUDA Graph
Traceback (most recent call last):
  File "/home/ec2-user/vllm/benchmarks/kernels/benchmark_lora.py", line 1065, in <module>
    args.func(args)
  File "/home/ec2-user/vllm/benchmarks/kernels/benchmark_lora.py", line 918, in run_model_bench
    run(args, bench_contexts)
  File "/home/ec2-user/vllm/benchmarks/kernels/benchmark_lora.py", line 793, in run
    bench_optype(
  File "/home/ec2-user/vllm/benchmarks/kernels/benchmark_lora.py", line 642, in bench_optype
    op_type.bench_fn()(**kwargs)
  File "/home/ec2-user/.local/lib/python3.9/site-packages/torch/_ops.py", line 1243, in __call__
    return self._op(*args, **kwargs)
RuntimeError: vllm::lora_shrink() is missing value for argument 'no_lora_flag_cpu'. Declaration: vllm::lora_shrink(Tensor inputs, Tensor[] lora_a_weights, Tensor(a2!) output_tensor, Tensor token_lora_mapping, Tensor token_indices_sorted_by_lora_ids, Tensor num_tokens_per_lora, Tensor lora_token_start_loc, Tensor lora_ids, Tensor no_lora_flag_cpu, float scaling) -> ()

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[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 ask your reviewers to trigger select CI tests on top of fastcheck CI.

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.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

🚀

[gemini-code-assist]

Code Review

This pull request fixes a bug where the no_lora_flag_cpu argument was missing when running benchmark_lora.py. The fix correctly adds this argument. My review focuses on improving the implementation by removing redundant code. I've identified duplicated and unnecessary logic for calculating the no_lora_flag_cpu flag in two different methods and suggested a more efficient and maintainable approach by reusing an existing pre-computed value.

[jeejeelee]

It looks like there was a PR previously: #23774

[classicsong]

It looks like there was a PR previously: #23774

Can we merge that PR or this one? Either is fine.

[jeejeelee]

Let's merge the other one instead, mainly because it's earlier. Thank you for your contribution.

[jeejeelee]

Fixed by #23774