[Kernel] Unified Triton kernel that doesn't distinguish between prefill + decode

[tdoublep]

In this PR we add:

• A new Triton kernel (triton_unified_attention) that works like flash_attn_varlen_func and can handle arbitrary query length. The kernel does GQA "packing" along the query dimension to ensure the Tensor cores are well used.
• Added a new unit test that is based on the unit tests for flash_attn_varlen_func
• Updated the V1 Triton attention backend to use this kernel. Note that the memory layout for the key cache also changes to match what the Flash attention backend is doing.

Best performance is obtained when using the jit cache decorator from triton_dejavu. In this code I'm using the jit cache decorator from triton_dejavu package but if #16606 is merged we could use it directly from vLLM.

Note that the unit tests don't currently pass when I enable the jit cache, but they all pass if it disabled. This is because we are testing different combinations of numbers of heads etc, which we assume to be constant in the decorator. We probably need to think of a good testing strategy for kernels with this decorator (cc @bringlein).

Initial benchmarking

Here are some initial benchmarking results on H100 for llama3.1-8b using:

python benchmark_serving.py \
    --model meta-llama/Llama-3.1-8B-Instruct  \
    --dataset-name sharegpt \
    --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json

Note that with these changes, the Triton backend significantly outperforms FlashAttention backend on an H100 GPU for this workloads.

Correctness

$ 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.800|±  |0.0179|
|     |       |strict-match    |     5|exact_match|↑  |0.784|±  |0.0184|

@bringlein the correctness check only looks good if I use branch tpa-grid-copy from triton-dejavu, if I use main it fails. you should be able to reproduce on H100.

Further benchmarking

I've run the following scenario across both H100 and MI300x using different backends from main, as well as using the changes from this PR:

MODEL=mistralai/Mistral-Small-24B-Instruct-2501
REQUEST_RATES=(1 5 7 9)
TOTAL_SECONDS=120

for REQUEST_RATE in "${REQUEST_RATES[@]}";
do
    NUM_PROMPTS=$(($TOTAL_SECONDS * $REQUEST_RATE))

    echo ""
    echo "===== RUNNING $MODEL FOR $NUM_PROMPTS PROMPTS WITH $REQUEST_RATE QPS ====="
    echo ""

    python3 benchmark_serving.py \
        --model $MODEL \
        --dataset-name random \
        --request-rate $REQUEST_RATE \
        --random-input-len 1000 \
        --random-output-len 100 \
        --num-prompts $NUM_PROMPTS \
        --ignore-eos --seed $REQUEST_RATE
done

Here are the results:

Main takeaways:

• TritonBackend from this PR is out-performing FlashAttentionBackend on H100
• The changes from this PR significantly improve the performance on MI300x at high QPS
• For some reason, using compile seems to make the results worse on MI300x (something that @SageMoore had mentioned to me previously). Do we have an issue to track that somewhere?

cc @robertgshaw2-redhat @tlrmchlsmth @SageMoore @bringlein @jvlunteren @LucasWilkinson

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

🚀

[tdoublep]

Re: the weirdness with torch compile on MI300x, I followed the suggestion of @robertgshaw2-redhat and re-ran everything inside the latest rocm/vllm-dev:nightly image that uses newer version of torch (2.7.0a0+git295f2ed). I now see better results without needing to use --enforce-eager on MI300x. I also ran the baseline V0 performance on MI300x using the ROCmFlashAttentionBackend.

[LucasWilkinson]

Nice work! Overall looks pretty good! Left a few comments

[LucasWilkinson]

can we add support for non-casual attention too? could be a future PR, but its useful for cascade attention and MLA

[tdoublep]

sure, I don't think that would be too hard