[ROCm] Enable Triton ScaledMM fallback + kernel selection fix

[shivampr]

Purpose

Fixes #14397 — triton_scaled_mm was never used on ROCm due to missing dispatch and checks.
This PR:

• Enables Triton fallback for ROCm when AITriton is unavailable

• Adds Triton fallback after CUTLASS on CUDA

• Implements is_supported() checks for kernel selection

• Adds a lightweight integration test validating ROCm dispatch logic

---

Test Plan

1. Mocked test (no GPU)

python3 mini_tests/select_triton_rocm.py

Result

Selected kernel: TritonScaledMMLinearKernel
OK: TritonScaledMMLinearKernel chosen on ROCm fallback.

2. MI300X (ROCm 7.0, vLLM built from this PR)

(a) Triton kernel functional test

max_abs_err≈2.5e-01, max_rel_err≈3.9e-03

(b) OpenAI-compatible API test

python3 -m vllm.entrypoints.openai.api_server \
  --model Qwen/Qwen2.5-0.5B-Instruct --dtype bfloat16 --host 0.0.0.0 --port 8000

Then:

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model":"Qwen/Qwen2.5-0.5B-Instruct","messages":[{"role":"user","content":"Say hi from MI300X."}]}'

Response

"Hello! How can I assist you today?"

Confirms successful end-to-end inference on ROCm.

[mergify]

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

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

[gemini-code-assist]

Code Review

This pull request addresses an issue where triton_scaled_mm was not being used on ROCm by fixing the kernel selection logic. It correctly adds TritonScaledMMLinearKernel as a fallback for both ROCm and CUDA, and introduces an is_supported check to ensure kernels are compatible with the current platform. The changes are accompanied by a new integration test to verify the fix.

My review focuses on improving the robustness of the kernel selection. I've suggested making the get_min_capability check in the Triton kernel platform-aware to prevent it from being selected on unsupported ROCm hardware. Additionally, I've pointed out a confusing try-except block in the new test file that should be simplified for clarity and to avoid masking potential errors.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

[ProExpertProg]

Is this ready for review again?

[shivampr]

@ProExpertProg yes!
Sorry will ping you directly from now on if its review ready.

[ProExpertProg]

Just one note about is_supported

[ProExpertProg]

@shivampr could you report lm-eval results for a model that uses int8 Triton scaled mm to check this works?

[shivampr]

@ProExpertProg

I verified the ROCm int8 Triton ScaledMM path using lm-eval with vLLM.

Environment:
- Base image: RunPod ROCm vLLM 0.9.2 / ROCm 7.0
- Backend: `VLLM_TARGET_DEVICE=rocm`
- Aiter disabled so we exercise Triton fallback:
  ```bash
  export VLLM_TARGET_DEVICE=rocm
  export VLLM_DISABLED_KERNELS="AiterScaledMMLinearKernel"

Command :

lm_eval \
  --model vllm \
  --model_args pretrained="RedHatAI/Meta-Llama-3.1-8B-Instruct-quantized.w8a8",dtype=auto,tensor_parallel_size=1,max_model_len=4096,gpu_memory_utilization=0.9,trust_remote_code=true \
  --tasks gsm8k \
  --num_fewshot 5 \
  --limit 50 \
  --batch_size auto

Form logs :

INFO ... Automatically detected platform rocm.
INFO ... Using TritonScaledMMLinearKernel for CompressedTensorsW8A8Int8

lm-eval results :

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.70	±	0.0655
		strict-match	5	exact_match	↑	0.64	±	0.0686

[mergify]

Hi @shivampr, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint