[ROCm][FEAT] Integrate AITER tgemm.

[vllmellm]

Purpose

This PR integrates aiter tgemm kernel that support bfloat16, float16, and fp8 data types, resulting in improved model performance.

This PR also introduces _aiter_ops.py as proposed in the RFC here. The aiter_ops namespace provides several key benefits:

• Centralized kernel registration: Ensures that kernels from the aiter package are properly registered

• Environment availability checks: Encapsulates aiter support detection and environment compatibility validation

• Reduced code duplication: Eliminates the need for duplicate helper functions across different vLLM modules

This implementation establishes the foundation for future refactoring efforts, where existing kernels throughout the vLLM repository will be migrated to use this unified approach for better maintainability and consistency.

This PR uses 7aa65b6 commit from aiter repo.

Benchmark Results

meta-llama/Llama-3.3-70B-Instruct tp2

Metric	With AITER tgemm	With AITER tuned tgemm	Without AITER tgemm
Request Throughput (req/s)	1.25	1.24	1.25
Output Token Thpt (tok/s)	696.23	706.41	701.09
Total Token Thpt (tok/s)	1939.97	1944.30	1953.18
Mean TTFT (ms)	457.92	461.52	448.17
Median TTFT (ms)	290.64	289.67	289.39
P99 TTFT (ms)	5261.99	4960.19	5408.76
Mean TPOT (ms)	44.80	44.71	45.90
Median TPOT (ms)	44.52	44.10	44.13
P99 TPOT (ms)	63.25	70.00	66.27
Mean ITL (ms)	44.43	43.81	44.10
Median ITL (ms)	32.39	31.99	32.11
P99 ITL (ms)	260.39	259.57	259.04

amd/Llama-3.3-70B-Instruct-FP8-KV tp2

Metric	With AITER tgemm	With AITER tuned tgemm	Without AITER tgemm
Request Throughput (req/s)	1.77	1.78	1.80
Output Token Thpt (tok/s)	994.62	980.54	986.99
Total Token Thpt (tok/s)	2763.94	2752.32	2779.85
Mean TTFT (ms)	350.12	344.54	337.96
Median TTFT (ms)	215.54	215.35	215.35
P99 TTFT (ms)	3683.01	3631.55	3827.46
Mean TPOT (ms)	32.93	33.62	32.22
Median TPOT (ms)	31.05	31.53	31.41
P99 TPOT (ms)	50.99	58.38	56.30
Mean ITL (ms)	31.03	31.44	31.24
Median ITL (ms)	21.97	22.28	21.96
P99 ITL (ms)	195.56	194.72	195.24

benchmark setting

python vllm/benchmarks/benchmark_serving.py --backend vllm --model "$model_name" --dataset-name random --num-prompts 1000 --max-concurrency 32 --random-input-len 1000 --random-output-len 1000

AITER tgemm tuning guide

1. Run vllm serve command with AITER_TUNE_GEMM=1 environment flag:
   Example:

   `VLLM_USE_V1=1 \
    VLLM_ROCM_USE_AITER=1 \
    AITER_TUNE_GEMM=1 \
    vllm serve meta-llama/Llama-3.3-70B-Instruct \
    --tensor-parallel-size 2 \
    --disable-log-requests \
    --trust-remote-code \
    --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE", "cudagraph_capture_sizes": [1,2,4,8,16,24,32]}'`
   

   The above command will record the requested shapes based on cudagraph_capture_sizes into aiter/configs/untuned_gemm.csv in the directory where you have installed/cloned the aiter package/repo.

2. Run python3 gradlib/gradlib/gemm_tuner.py --tuned_file aiter/configs/tuned_gemm.csv --input_file aiter/configs/untuned_gemm.csv in the directory where aiter package exists.

for more instruction, follow the documentation here.

Test Plan

Test models that are afftected by this change, using lm_eval on gsm8k dataset.

environment setting

Step 1: run vllm serve

VLLM_USE_V1=1 VLLM_ROCM_USE_AITER=1 vllm serve $MODEL_NAME --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE", "cudagraph_capture_sizes": [1,2,4,8,16,24,32]}' -tp 2 --trust-remote-code --swap-space 16 --distributed-executor-backend mp

Step 2: run lm_eval

lm_eval --model local-completions --tasks gsm8k --model_args model=$MODEL_NAME,base_url=http://localhost:8000/v1/completions --trust_remote_code --num_fewshot 5 --batch_size 256

Test Results

meta-llama/Llama-3.3-70B-Instruct tp2

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9356	±	0.0068
		strict-match	5	exact_match	↑	0.9083	±	0.0080

amd/Llama-3.3-70B-Instruct-FP8-KV tp2

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9378	±	0.0067
		strict-match	5	exact_match	↑	0.8984	±	0.0083

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[mergify]

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

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