[Kernel][ROCM] Upstream prefix prefill speed up for vLLM V1

[maleksan85]

Speed up prefix prefill with vLLM V1 on AMG GPUs

Improvements:

1. Vectorization in the context loop (most complex one as k cache shape is very specific)
2. Refactoring for online softmax computation
3. Refactoring to the kernel so autotune might select the best configs per shape
4. Plus adding new spectrum of unrolling/staging in autotuner

More details on triton kernel tunning: https://rocm.docs.amd.com/en/docs-6.1.1/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html

see last comments

[maleksan85]

HIP_VISIBLE_DEVICES=6 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-04-08:18:10:02,846 INFO [lm_eval.loggers.evaluation_tracker:272] 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.808	±	0.0176
		strict-match	5	exact_match	↑	0.782	±	0.0185

[maleksan85]

python3 benchmarks/benchmark_serving.py
--backend vllm
--model /data/models/Llama-3.1-70B-Instruct
--dataset-name random
--random-input-len 10000
--random-output-len 100
--num-prompts 300
--seed 42
--ignore-eos
--percentile-metrics "ttft,tpot,itl,e2el"

PR (like 20% gain)

============ Serving Benchmark Result ============
Successful requests:                     300
Benchmark duration (s):                  409.78
Total input tokens:                      3000000
Total generated tokens:                  30000
Request throughput (req/s):              0.73
Output token throughput (tok/s):         73.21
Total Token throughput (tok/s):          7394.28
---------------Time to First Token----------------
Mean TTFT (ms):                          205042.73
Median TTFT (ms):                        203406.19
P99 TTFT (ms):                           400609.81
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          1610.15
Median TPOT (ms):                        2027.83
P99 TPOT (ms):                           2239.19
---------------Inter-token Latency----------------
Mean ITL (ms):                           1610.15
Median ITL (ms):                         80.56
P99 ITL (ms):                            5252.32
----------------End-to-end Latency----------------
Mean E2EL (ms):                          364447.21
Median E2EL (ms):                        404161.34
P99 E2EL (ms):                           409588.24
==================================================

Upstream

============ Serving Benchmark Result ============
Successful requests:                     300
Benchmark duration (s):                  498.15
Total input tokens:                      3000000
Total generated tokens:                  30000
Request throughput (req/s):              0.60
Output token throughput (tok/s):         60.22
Total Token throughput (tok/s):          6082.51
---------------Time to First Token----------------
Mean TTFT (ms):                          249095.71
Median TTFT (ms):                        248711.87
P99 TTFT (ms):                           488484.85
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          1957.47
Median TPOT (ms):                        2462.50
P99 TPOT (ms):                           2732.60
---------------Inter-token Latency----------------
Mean ITL (ms):                           1957.47
Median ITL (ms):                         80.32
P99 ITL (ms):                            8005.81
----------------End-to-end Latency----------------
Mean E2EL (ms):                          442885.68
Median E2EL (ms):                        492500.58
P99 E2EL (ms):                           497952.19
==================================================

[maleksan85]

python3 benchmarks/benchmark_serving.py
--backend vllm
--model /data/models/Llama-3.1-70B-Instruct
--dataset-name random
--random-input-len 5000
--random-output-len 100
--num-prompts 500
--seed 42
--ignore-eos
--percentile-metrics "ttft,tpot,itl,e2el"

PR (10% gain)

============ Serving Benchmark Result ============
Successful requests:                     500
Benchmark duration (s):                  319.37
Total input tokens:                      2500000
Total generated tokens:                  50000
Request throughput (req/s):              1.57
Output token throughput (tok/s):         156.56
Total Token throughput (tok/s):          7984.50
---------------Time to First Token----------------
Mean TTFT (ms):                          155485.39
Median TTFT (ms):                        149836.40
P99 TTFT (ms):                           310684.27
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          1219.18
Median TPOT (ms):                        1556.81
P99 TPOT (ms):                           1629.28
---------------Inter-token Latency----------------
Mean ITL (ms):                           1219.18
Median ITL (ms):                         77.67
P99 ITL (ms):                            4265.61
----------------End-to-end Latency----------------
Mean E2EL (ms):                          276184.44
Median E2EL (ms):                        310784.82
P99 E2EL (ms):                           319205.24
==================================================

Upstream

============ Serving Benchmark Result ============
Successful requests:                     500
Benchmark duration (s):                  355.99
Total input tokens:                      2500000
Total generated tokens:                  50000
Request throughput (req/s):              1.40
Output token throughput (tok/s):         140.45
Total Token throughput (tok/s):          7163.04
---------------Time to First Token----------------
Mean TTFT (ms):                          172121.19
Median TTFT (ms):                        162339.60
P99 TTFT (ms):                           349045.74
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          1369.76
Median TPOT (ms):                        1699.35
P99 TPOT (ms):                           1892.04
---------------Inter-token Latency----------------
Mean ITL (ms):                           1369.76
Median ITL (ms):                         78.00
P99 ITL (ms):                            6167.44
----------------End-to-end Latency----------------
Mean E2EL (ms):                          307727.51
Median E2EL (ms):                        349138.54
P99 E2EL (ms):                           355831.83
==================================================

[SageMoore]

Looks reasonable. Just a few nits. Thanks for all of the hard work making this kernel faster.