[Kernel] [V1] Improved performance for V1 Triton (ROCm) backend

[tdoublep]

In this PR we target performance improvement for the V1 ROCmAttentionBackend (which should be renamed TritonAttentionBackend once this PR is merged).

cc @SageMoore

Performance

All of the below results are for meta-llama/Llama-3.1-8B-Instruct on an NVIDIA H100 GPU.

The ROCmAttention backend on main (we have to hack it a bit to make this happen on NVIDIA) currently produces the following results for the serving benchmark:

$ python benchmarks/benchmark_serving.py \
    --model meta-llama/Llama-3.1-8B-Instruct \
    --dataset-name sharegpt \
    --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json
============ Serving Benchmark Result ============
Successful requests:                     1000      
Benchmark duration (s):                  107.33    
Total input tokens:                      215196    
Total generated tokens:                  197284    
Request throughput (req/s):              9.32      
Output token throughput (tok/s):         1838.05   
Total Token throughput (tok/s):          3842.98   
---------------Time to First Token----------------
Mean TTFT (ms):                          10264.30  
Median TTFT (ms):                        8952.08   
P99 TTFT (ms):                           24269.34  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          408.87    
Median TPOT (ms):                        261.25    
P99 TPOT (ms):                           1208.50   
---------------Inter-token Latency----------------
Mean ITL (ms):                           210.67    
Median ITL (ms):                         172.46    
P99 ITL (ms):                            1308.82   
==================================================

Whereas, using the default FlashAttention backend produces:

============ Serving Benchmark Result ============
Successful requests:                     1000      
Benchmark duration (s):                  20.47     
Total input tokens:                      215196    
Total generated tokens:                  198001    
Request throughput (req/s):              48.86     
Output token throughput (tok/s):         9673.74   
Total Token throughput (tok/s):          20187.58  
---------------Time to First Token----------------
Mean TTFT (ms):                          3473.32   
Median TTFT (ms):                        3405.53   
P99 TTFT (ms):                           6302.31   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          82.06     
Median TPOT (ms):                        47.72     
P99 TPOT (ms):                           212.20    
---------------Inter-token Latency----------------
Mean ITL (ms):                           36.34     
Median ITL (ms):                         24.48     
P99 ITL (ms):                            216.40    
==================================================

So, the Triton backend is really under-performing on this hardware: the total throughput is more than 5x worse than the FlashAttention baseline.

Whereas, using the optimized ROCmAttention from this branch produces:

============ Serving Benchmark Result ============
Successful requests:                     1000      
Benchmark duration (s):                  29.38     
Total input tokens:                      215196    
Total generated tokens:                  197281    
Request throughput (req/s):              34.03     
Output token throughput (tok/s):         6713.84   
Total Token throughput (tok/s):          14037.37  
---------------Time to First Token----------------
Mean TTFT (ms):                          3711.10   
Median TTFT (ms):                        3781.32   
P99 TTFT (ms):                           6787.05   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          100.91    
Median TPOT (ms):                        66.34     
P99 TPOT (ms):                           249.37    
---------------Inter-token Latency----------------
Mean ITL (ms):                           53.15     
Median ITL (ms):                         42.07     
P99 ITL (ms):                            259.42    
==================================================

Thus, we see a 3.65x performance boost in terms of throughput, and performance is only 30% off the FlashAttention baseline. We feel like with smart auto-tuning of the kernels (#13305 already goes in this direction), that the performance gap can be closed further.

We are in the process of evaluating the performance on AMD GPUs.

Correctness

The new kernel produces correct results.

Using FlashAttention on H100 we see:

$ export VLLM_V1_USE_TRITON_BACKEND=0
$ VLLM_USE_V1=1 lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --tasks gsm8k --num_fewshot 5 --batch_size auto --limit 500
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.792|±  |0.0182|
|     |       |strict-match    |     5|exact_match|↑  |0.770|±  |0.0188|

Using TritonAttention on H100 we see:

$ export VLLM_V1_USE_TRITON_BACKEND=1
$ VLLM_USE_V1=1 lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --tasks gsm8k --num_fewshot 5 --batch_size auto --limit 500
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.798|±  |0.0180|
|     |       |strict-match    |     5|exact_match|↑  |0.780|±  |0.0185|

How is this performance improvement achieved?

We have noticed that the prefix_prefill kernel is not well optimized for sequences that are in decode phase. While it may be possible to optimize the kernel to make it better for decode, it is quite limiting to try to combine everything within a single Triton kernel. It implies we have to always use the same grid for prefill, chunked prefills and decodes, and also we have to take great care to choose the tile sizes and auto-tuning settings to work well across all cases. Instead, we think it is better to handle prefill/chunked prefill separately from decodes. We have implemented chunked_prefill_paged_decode in ibm_triton_lib so that the prefix_prefill kernel is used for the sequences in the batch that have query length > 1, and we used an optimized paged decoding Triton kernel for all other sequences. As you can see, no change to vLLM V1 metadata is required.

What is ibm_triton_lib?

To run the code on this branch one needs to install the following dependency:

pip install "git+https://github.com/foundation-model-stack/vllm-triton-backend.git#egg=ibm-triton-lib&subdirectory=ibm-triton-lib"

This is, hopefully, just a temporary place for these kernels. If preferred by the community, we would be happy to contribute these new kernels to vLLM repo. Alternatively, we feel like there might be some benefit to having a Triton-equivalent to FlashAttention (e.g., an external project that implements optimized Triton kernels). We feel this could un-burden vLLM from performing extensive AOT auto-tuning that is required to get the most out of these kernels (see @bringlein's talk from Ray Summit). Rather, we could do this in an external project and bundle the optimized configs inside a wheel that vLLM could easily consume. We discussed this a bit already on this RFC. If there is interested from the wider community, we could think about having this as a vLLM ecosystem project. Please note that currently the Triton kernels being used here have not been auto-tuned, so there is definitely more performance improvements to be had here.

Update: this PR is now self-contained and does not have a dependency on ibm_triton_lib

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

🚀

[SageMoore]

Thanks for the contribution! I'm totally fine with the dependency if it's easily installable via the requirements-rocm-build.txt file. I do think it would be nice to have this kernel integrated into vllm so that any performance optimizations made to this kernel in the future are visible to the vllm community.

[bringlein]

I did some benchmarks for different number of concurrent users on A100 and H100 for this PR, using the benchmarks/benchmark_serving.py as described by @tdoublep above:

[tdoublep]

I do think it would be nice to have this kernel integrated into vllm so that any performance optimizations made to this kernel in the future are visible to the vllm community.

@SageMoore We've added the kernel to this PR. This is now a standalone contribution with no dependency on ibm_triton_lib.

[maleksan85]

Hi Thomas,

Can we do joined effort and:

1. take the kernel from my PR: [Kernel][ROCM] Upstream prefix prefill speed up for vLLM V1 #13305 plus return for ctx len == 1 (because your _fwd_kernel seems the same as original except return for context len 1). Use mine only starting from triton 3.2
2. take your kernel kernel_paged_attention_2d
3. put them into path of vllm upstream: /vllm/attention/ops/prefix_prefill.py

With this joined thing above I got:

VLLM_ENABLE_V1_MULTIPROCESSING=0 VLLM_USE_V1=1 lm_eval --model vllm --model_args pretrained=/data/models/Llama-3.1-8B-Instruct --tasks gsm8k --num_fewshot 5 --batch_size auto --limit 500

2025-03-04:19:53:53,697 INFO [evaluation_tracker.py:269] Output path not provided, skipping saving results aggregated
vllm (pretrained=/data/models/Llama-3.1-8B-Instruct), gen_kwargs: (None), limit: 500.0, num_fewshot: 5, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.790	±	0.0182
		strict-match	5	exact_match	↑	0.774	±	0.0187

and:

python vllm/benchmarks/benchmark_serving.py --model /data/models/Llama-3.1-8B-Instruct --dataset-name sharegpt --dataset-path datasets/ShareGPT/ShareGPT_V3_unfiltered_cleaned_split.json

============ Serving Benchmark Result ============
Successful requests:                     1000
Benchmark duration (s):                  23.68
Total input tokens:                      215196
Total generated tokens:                  198343
Request throughput (req/s):              42.22
Output token throughput (tok/s):         8374.29
Total Token throughput (tok/s):          17460.14
---------------Time to First Token----------------
Mean TTFT (ms):                          1311.16
Median TTFT (ms):                        1366.36
P99 TTFT (ms):                           1808.26
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          45.81
Median TPOT (ms):                        44.80
P99 TPOT (ms):                           72.73
---------------Inter-token Latency----------------
Mean ITL (ms):                           38.17
Median ITL (ms):                         39.32
P99 ITL (ms):                            81.78
==================================================

cc @SageMoore

[tdoublep]

@maleksan85 awesome! Yes, there is no reason why these two improvements can't be combined. I had tried it yesterday but didn't see additive gains like that. I will try it again using the latest commits.

[maleksan85]

prefix_prefill.py.txt
I tried benchmarks above on latest main + prefix_prefill.py (from attach) + your file triton_paged_decode_attention_2d.py https://github.com/foundation-model-stack/vllm-triton-backend/blob/40b2fa560f5c5ae00724960852f426a3eb99d71e/ibm-triton-lib/ibm_triton_lib/kernels/triton_paged_decode_attention_2d.py#L73 copied into same folder as prefix_prefill.py.

[SageMoore]

This PR generally looks good to me. I only have some small nits. Thanks so much for the contribution! The performance improvements are very exciting.

[SageMoore]

Nit: Could you explain a bit why all of the strides are constexpr? My understanding is that triton will generate a new kernel each time one of these constexprs changes, which could be non-ideal.

[bringlein]

Yes, that's right, triton will generate a new binary if a constant changes. However, having strides as constant enables better optimization of memory accesses (and therefore, automatic tiling, coalescing, etc.). Our experience shows this could be quite significant, esp. for small requests (we've seen latency reduction up to 2x for this 2d kernel in the past for very small batches).
Additionally, we don't expect that num_query_heads or block_table_stride changes within a running inference server (although it will change based on the vllm configuration or model).

[tdoublep]

Thanks for the review @SageMoore @maleksan85.

We have addressed all the comments.

[SageMoore]

LGTM

[tlrmchlsmth]

A couple of things:

1. any environment variables should be handled invllm/envs.py
2. We should try to avoid proliferation of environment variables in general when possible. I know you're working on a triton backend in general - is there any community consensus for how this should be integrated yet?

[tdoublep]

1. Sure, we actually just added this env variable as a way to test the ROCm/Triton backend on NVIDIA hardware. We were kind of hoping that this PR might be merged in the meantime which would give the user a way to select the ROCm/Triton backend without introducing a new env variable.

2. We have developed a Triton backend for vLLM V0 which was open-sourced as an out-of-tree plugin here. Since V1 already has a Triton-only backend (e.g., what is currently called ROCmBackend) I think the V0 backend code we have developed is now somewhat redundant, however the kernels that we have developed can be re-used for V1 (with some important changes) which is essentially what this PR is doing). We are happy to contribute these kernels to vLLM repo (as suggested by @SageMoore above) so the community can track the progress and contribute further optimizations.

[tdoublep]

Perhaps we should simply remove the changes to platforms/cuda.py from this PR? It would mean that we can only run this backend on AMD for now, but that is the current behaviour on main anyway. Adding the option to use the Triton backend on NVIDIA can be addressed in a separate PR (e.g. #14071)

[tlrmchlsmth]

All of that sounds good and makes sense to me.

For now, let's remove the changes platforms/cuda.py for this PR and iterate on it future ones

[tdoublep]

done

[SageMoore]

I ran the same benchmark_serving experiments that you posted above on an MI300x machine and the results are quite good.

branch: main
Command: python benchmarks/benchmark_serving.py \
			   --model meta-llama/Llama-3.1-8B-Instruct \
			   --dataset-name sharegpt \
			   --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
			   --port 4000
Results:
============ Serving Benchmark Result ============
Successful requests:                     1000
Benchmark duration (s):                  136.33
Total input tokens:                      215196
Total generated tokens:                  197051
Request throughput (req/s):              7.34
Output token throughput (tok/s):         1445.39
Total Token throughput (tok/s):          3023.89
---------------Time to First Token----------------
Mean TTFT (ms):                          23988.34
Median TTFT (ms):                        18652.88
P99 TTFT (ms):                           62421.28
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          358.72
Median TPOT (ms):                        325.15
P99 TPOT (ms):                           722.79
---------------Inter-token Latency----------------
Mean ITL (ms):                           249.32
Median ITL (ms):                         171.76
P99 ITL (ms):                            826.05
==================================================


branch: tpa-v1-clean
Command: python benchmarks/benchmark_serving.py \
			   --model meta-llama/Llama-3.1-8B-Instruct \
			   --dataset-name sharegpt \
			   --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
			   --port 4000
Results:
============ Serving Benchmark Result ============
Successful requests:                     1000
Benchmark duration (s):                  36.09
Total input tokens:                      215196
Total generated tokens:                  197037
Request throughput (req/s):              27.71
Output token throughput (tok/s):         5459.77
Total Token throughput (tok/s):          11422.72
---------------Time to First Token----------------
Mean TTFT (ms):                          6215.30
Median TTFT (ms):                        5872.31
P99 TTFT (ms):                           12840.71
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          74.53
Median TPOT (ms):                        73.89
P99 TPOT (ms):                           122.55
---------------Inter-token Latency----------------
Mean ITL (ms):                           55.83
Median ITL (ms):                         47.13
P99 ITL (ms):                            132.60
==================================================

[tlrmchlsmth]

should this be torch.float8_e4m3fnuz in some cases on RoCM?

[tdoublep]

Not entirely sure, but this logic is taking direct from what is already happening here.

[tlrmchlsmth]

FYI @SageMoore is going to double-check accuracy on an MI300 system

[SageMoore]

VLLM_USE_V1=1 lm_eval --model vllm --model_args pretrained=meta-llama/Llama-3.1-8B-Instruct --tasks gsm8k --num_fewshot 5 --batch_size auto --limit 500

vllm (pretrained=meta-llama/Llama-3.1-8B-Instruct), gen_kwargs: (None), limit: 500.0, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.780|±  |0.0185|
|     |       |strict-match    |     5|exact_match|↑  |0.762|±  |0.0191|

Looks reasonable

[tlrmchlsmth]

Thanks @SageMoore. And since this is just KV cache compression, the fp8_e4m3fn on RoCM issue is not worrying me like it did when it first caught my eye

[tlrmchlsmth]

This looks great to me, thanks for the contribution!

[tlrmchlsmth]

Thanks @SageMoore. And since this is just KV cache compression, the fp8_e4m3fn on RoCM issue is not worrying me like it did when it first caught my eye

[tdoublep]

PR failure is unrelated. I checked and that test is actually failing on main (see here for a quick fix).

[bringlein]

just do document: This PR enables similar performance improvement also on MI250 (4.1x)