@tdoublep can you help to review?

The following are benchmark results obtained on an NVIDIA H100 GPU.

vllm bench latency --model meta-llama/Llama-3.1-8B-Instruct --batch-size 1 --input-len 500 --output-len <variable>

vllm bench latency --model meta-llama/Llama-3.1-8B-Instruct --batch-size <variable> --input-len 500 --output-len 100

vllm serve meta-llama/Llama-3.1-8B-Instruct --disable-log-requests

vllm bench serve --model meta-llama/Llama-3.1-8B-Instruct --dataset-name sharegpt \
                 --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json

This PR:

============ Serving Benchmark Result ============
Successful requests:                     984
Failed requests:                         16
Benchmark duration (s):                  19.56
Total input tokens:                      211284
Total generated tokens:                  194325
Request throughput (req/s):              50.32
Output token throughput (tok/s):         9936.97
Peak output token throughput (tok/s):    20576.00
Peak concurrent requests:                984.00
Total Token throughput (tok/s):          20741.15
---------------Time to First Token----------------
Mean TTFT (ms):                          3336.41
Median TTFT (ms):                        3269.26
P99 TTFT (ms):                           6205.97
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          79.15
Median TPOT (ms):                        44.26
P99 TPOT (ms):                           210.99
---------------Inter-token Latency----------------
Mean ITL (ms):                           35.44
Median ITL (ms):                         23.97
P99 ITL (ms):                            215.13
==================================================

Current upstream:

============ Serving Benchmark Result ============
Successful requests:                     984
Failed requests:                         16
Benchmark duration (s):                  19.16
Total input tokens:                      209644
Total generated tokens:                  194099
Request throughput (req/s):              51.36
Output token throughput (tok/s):         10131.23
Peak output token throughput (tok/s):    20148.00
Peak concurrent requests:                984.00
Total Token throughput (tok/s):          21073.84
---------------Time to First Token----------------
Mean TTFT (ms):                          3368.81
Median TTFT (ms):                        3345.50
P99 TTFT (ms):                           6208.56
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          78.77
Median TPOT (ms):                        43.73
P99 TPOT (ms):                           211.84
---------------Inter-token Latency----------------
Mean ITL (ms):                           35.38
Median ITL (ms):                         24.00
P99 ITL (ms):                            213.75
==================================================

For these experiments on an NVIDIA H100 GPU, no clear performance differences were observed.

Here is a table summarizing the benefit of the following changes on MI300:

1. No dim mask: Removing the dim. mask trick
2. No exp2: Replacing exp2 with exp
3. No q cg: Removing cache modifier from the q load
4. No tile mask: Removing the tile mask trick

I can do (1) and (3) as their effect is minimal if you prefer to minimize the diff.

What do you think @tdoublep?

Num KV heads: 8
Num Q heads: 64
Dtype: bf16

@cagrikymk Sorry for slow response on this one. Thank you for the detailed experiments!! If I understand correctly, lower is better in this table right?

@tdoublep Oh, I forgot to include some important details.

These numbers are the speedup against the current code in the PR.

Each column undoes a change and shows how much the perf. is affected by it.

So, if the number is low, it means it is needed for good perf., hence, as you said, lower is better as that indicates the importance of keeping that optimization.

With this PR, Triton attention gets faster but is still slower than using VLLM_ATTENTION_BACKEND=ROCM_AITER_UNIFIED_ATTN on gfx1151:

# This PR
Preprocessing 34171.30 tokens/s; TTFT 116 ms
Decode 44.23 tokens/s

After adding the following changes on top

• select_2d_config, select_3d_config: Use TILE_SIZE = block_size on RDNA (from [ROCm] Restore 16-wide fast path in Triton unified attention #30582)
• select_3d_config: MIN_SEGMENTS = 16 if TILE_SIZE <= 16 else 8 on RDNA
• use_2d_kernel: Force cu_mult = 4 on RDNA
• Kernel fast path (in both 2D and 3D kernels) from [ROCm] Restore 16-wide fast path in Triton unified attention #30582: When TILE_SIZE == BLOCK_SIZE:
  - Direct block table lookup with j instead of seq_offset // BLOCK_SIZE
  - Direct offs_t indexing instead of seq_offset % BLOCK_SIZE
  it's on par with VLLM_ATTENTION_BACKEND=ROCM_AITER_UNIFIED_ATTN:

# After the changes above
Preprocessing 57191.28 tokens/s; TTFT 69 ms
Decode 51.30 tokens/s

Test case is

vllm serve Qwen/Qwen3-4B-AWQ --max-model-len 4096 --gpu-memory-utilization 0.38 --max-num-batched-tokens 128
vllm bench serve --model Qwen/Qwen3-4B-AWQ --num-warmups 5 --random-input-len 3968 --random-output-len 128 --num-prompts 1 --max-concurrency 1

@cagrikymk Is this PR stale?