[ROCm][Aiter] Add triton fp8 bmm kernel for mla

[divakar-amd]

Purpose

Replace torch.bmm with aiter batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant kernel for MLA

Test Plan

• Correctness check
• Perf check

Test Result

Correctness Result

Generated Outputs:
------------------------------------------------------------
Prompt:    'Hello, my name is'
Output:    ' Christian Munoz and\nthis is my blog where I cover different\nIT topics'
------------------------------------------------------------
Prompt:    'The president of the United States is'
Output:    ' in charge of which branch of government? A. judicial B. legislative C.'
------------------------------------------------------------
Prompt:    'The capital of France is'
Output:    ' Paris. Paris is located along the Seine River in the north-central part of the'
------------------------------------------------------------
Prompt:    'The future of AI is'
Output:    " a fascinating and rapidly evolving field. Here's a glimpse into some key areas shaping"
------------------------------------------------------------

Performance test on DeepSeek-R1 with full-cudagraph capture mode

REQUEST_RATES=(1 5 7 9)
TOTAL_SECONDS=20 
TP=8
OUTPUT_LEN=128

DATASET_PATH="ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json"
PYTHON_BENCH_SCRIPT="benchmarks/benchmark_serving.py"

   VLLM_USE_V1=1 vllm serve $MODEL \
        --port 8004 \
        --tensor-parallel-size $TP \
        --max-num-seqs 256 \
        --no-enable-prefix-caching \
        --swap-space 16 \
        --disable-log-requests \
        --disable-uvicorn-access-log \
        --block-size 1 \
        -O '{"full_cuda_graph":true}' 

        python3 $PYTHON_BENCH_SCRIPT \
            --model $MODEL \
            --percentile-metrics ttft,tpot,itl,e2el \
            --dataset-path $DATASET_PATH \
            --request-rate $REQUEST_RATE \
            --num-prompts $(($TOTAL_SECONDS * $REQUEST_RATE)) \
            --ignore-eos \
            --port 8004 \
            --sharegpt-output-len $OUTPUT_LEN

(Optional) Documentation Update

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

🚀

[gemini-code-assist]

Code Review

This pull request introduces a Triton FP8 BMM kernel for MLA on ROCm to enhance performance. The changes include adding an environment variable to toggle this feature and updating the MLA attention implementation to leverage the new kernel for matrix multiplications. The implementation also involves quantizing weights to FP8 and includes a warmup loop for the Triton kernel. My review has identified two high-severity issues: the use of a hardcoded FP8 dtype instead of a platform-specific one, and a hardcoded warmup range for the Triton kernel that is insufficient for default configurations, potentially leading to performance degradation.

[mergify]

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

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

[divakar-amd]

Re-created another PR with some updates: #23264