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