[V1] LoRA - Add triton kernels for V1

[varun-sundar-rabindranath]

Add shrink and expand triton kernels for V1.

Why do we need a new set of kernels:

• V0 sorts/groups requests based on LoRA ID. The SGMV kernels take advantage of this and groups the compute within thread blocks.
• V1 doesn't group requests based on LoRA ID. The new set of kernels have information about which input tokens map to which LoRA ID and they use this information to load the appropriate input tokens. The rest of the matmul is very similar to the SGMV kernels.

Kernel Code Change:
The new kernels re-use a lot of the code from the existing SGMV kernels. The main changes are,

1. Kernel Launch Grid formulation (this was required so the kernels are CUDAGraph compatible. Note that SGMV kernels are not)
2. Loading of the input tokens (A matrix) for the matmul.
   All other kernel code is the same as the existing SGMV kernels. I refactored the code so it can be reused.

benchmark serving numbers :

server command : VLLM_USE_V1="1" vllm serve meta-llama/Llama-2-7b-hf --max-loras 4 --max-lora-rank 8 --enable-lora --lora-modules lora1=yard1/llama-2-7b-sql-lora-test lora2=yard1/llama-2-7b-sql-lora-test lora3=yard1/llama-2-7b-sql-lora-test lora4=yard1/llama-2-7b-sql-lora-test --no-enable-prefix-caching

benchmark command : python3 benchmarks/benchmark_serving.py --model meta-llama/Llama-2-7b-hf --dataset-name sharegpt --dataset-path ./ShareGPT_V3_unfiltered_cleaned_split.json --num-prompts 500 --request-rate inf --lora-modules lora1 lora2 lora3 lora4

V1 LoRA - This PR:

============ Serving Benchmark Result ============
Successful requests:                     500       
Benchmark duration (s):                  68.81     
Total input tokens:                      117316    
Total generated tokens:                  110942    
Request throughput (req/s):              7.27      
Output token throughput (tok/s):         1612.20   
Total Token throughput (tok/s):          3317.02   
---------------Time to First Token----------------
Mean TTFT (ms):                          7103.84   
Median TTFT (ms):                        7136.82   
P99 TTFT (ms):                           14040.12  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          126.93    
Median TPOT (ms):                        94.96     
P99 TPOT (ms):                           231.31    
---------------Inter-token Latency----------------
Mean ITL (ms):                           76.47     
Median ITL (ms):                         58.02     
P99 ITL (ms):                            236.28    
==================================================

V1 LoRA - Main:

============ Serving Benchmark Result ============
Successful requests:                     500       
Benchmark duration (s):                  117.84    
Total input tokens:                      117316    
Total generated tokens:                  110942    
Request throughput (req/s):              4.24      
Output token throughput (tok/s):         941.44    
Total Token throughput (tok/s):          1936.96   
---------------Time to First Token----------------
Mean TTFT (ms):                          10277.45  
Median TTFT (ms):                        9370.56   
P99 TTFT (ms):                           22882.99  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          259.87    
Median TPOT (ms):                        236.82    
P99 TPOT (ms):                           445.65    
---------------Inter-token Latency----------------
Mean ITL (ms):                           173.39    
Median ITL (ms):                         164.76    
P99 ITL (ms):                            459.73    
==================================================

Kernel micro benchmark:

Please find kernel microbenchmark here - https://docs.google.com/spreadsheets/d/1b_8KsDGdiSGWlHODMszug_-do7OlSPlSzoV84VkmVPc/edit?usp=sharing (sheet : "V1 : Dont Sort Tokens By LoRA ")

Note : The V0 SGMV and BGMV kernels are not tuned. But the V1 kernels are tuned with triton auto-tuner. Therefore the discrepancy between the V1 and SGMV/BGMV kernels could be partially explained by the tuning.
The SGMV kernel depends heavily on the input being sorted. V1 kernels aren't affected as much.

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

Requesting reviews on this file. The utilities in this file deal with loading the stored triton configs.
cc @tlrmchlsmth @mgoin

[jeejeelee]

QQ: Is there a significant gain to stored tuned config?

[varun-sundar-rabindranath]

Not significant - but I do see some gains. On average I see about 200 tokens/s more throughput when using the tuned configs.

[jeejeelee]

TBH, storing these configs is a bit crazy, I'm not sure if this is the right direction.

[varun-sundar-rabindranath]

It appears, it is common to have these configs for triton kernels - the fused_moe kernels also have such configs here https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/fused_moe/configs

Arguments for storing configs:

• The default/static config isn't going to be optimal for all cases.
• Generating these configs is easy. With new GPUs and Triton version updates we can simply run the triton autotuner and drop in these configs without having to do any code change.

I did a tuned-kernels vs untuned-kernels micro benchmark run -
https://docs.google.com/spreadsheets/d/1f87TKWuwJXVK2-8YHGBoVbo6ueEWGpLPPUIkgYs6dMo/edit?usp=sharing .
TLDR the tuned shrink kernels are better than the untuned versions in most cases. For expand kernels the tuned version is much better at low batch size regimes

E2E performance:

• The 200 tokens / s number I shared is on the A100 GPU. The E2E performance is limited in this case. But I believe with CUDA Graphs enabled, the tuned kernels will have a bigger impact.
• Also, I haven't checked the E2E performance on H100 GPU. it might be better.

What do you think ?

[varun-sundar-rabindranath]

Spoke with @jeejeelee IRL - The plan is to remove the triton configs from this PR and introduce them in a separate PR so we can reason about them separately.

@jeejeelee I have removed the configs from this PR. can you take another pass at the PR when you find the time please ! Thanks 🙏

[varun-sundar-rabindranath]

Is there a better way to call the base class method ? (Note that this class inherits from multiple classes. )

[jeejeelee]

Maybe super().update_metadata is better

[jeejeelee]

All the LoRA tests have failed again

[varun-sundar-rabindranath]

All the LoRA tests have failed again

Looking into this now 👍

[varun-sundar-rabindranath]

Update : I enabled tests in tests/lora/test_layers.py for V1. The tests work locally but OOM's on the CI - I am tracking this down.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @varun-sundar-rabindranath.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[jeejeelee]

It seems these modifications have significantly increased the time consumption for lora testing

[varun-sundar-rabindranath]

It seems these modifications have significantly increased the time consumption for lora testing

Yes. This PR adds the v1_kernel tests in test_punica_ops.py and enables test_layers.py to run for V1 also. I believe most of it is coming from the test_layers.py that now runs for both V1 and V0 (effectively doubling its run time) - Ill see what we can do here.

[Edit]
@jeejeelee

Update : Reduced the tests in commits a18d273 and ba94947
The times are now,

a maximum of 7 minute increase. Do you think we should prune further ?

[jeejeelee]

Thank you very much. Let's continue moving forward in the direction we discussed on Slack.