[ROCm] Enable Triton ScaledMM fallback + kernel selection fix

[shivampr]

Purpose

This PR fixes #14397 : triton_scaled_mm never used on ROCm.

On ROCm, TritonScaledMMLinearKernel was never selected because:

• Only AITriton path was registered for ROCm
• CUTLASS path was skipped on ROCm
• Triton fallback was missing from the dispatch list
• Kernel selection didn't check is_supported()
• As a result, ROCm always fell back to slow torch.scaled_mm

This PR:

• Enables Triton fallback for ROCm when AITriton is not available
• Adds Triton fallback to CUDA after CUTLASS
• Implements kernel support checks via is_supported()
• Adds minimal integration test (no GPU required)

Test Plan (what I actually ran)

Minimal smoke test (ROCm-mocked, no native ops):

# From repo root
PYTHONPATH="$(pwd)" python3 mini_tests/select_triton_rocm.py

What the script does (in-process mocks):

• Forces VLLM_TARGET_DEVICE=rocm
• Mocks so ROCm platform is detected
• Stubs vllm._custom_ops (no compiled extensions needed)
• Uses torch CPU wheel only
• Selects kernel via choose_scaled_mm_linear_kernel(...)
• Asserts TritonScaledMMLinearKernel is chosen

Test Results

• ROCm mocked + AITriton disabled → Triton fallback selected: PASS
  Output includes:

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

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

[gemini-code-assist]

The get_min_capability method returns a hardcoded value of 75, which corresponds to a CUDA compute capability (Turing architecture). This is confusing and potentially incorrect when running on ROCm, as it might allow this kernel to be selected on older ROCm hardware that doesn't fully support Triton, leading to runtime errors. To make the check more robust and clear, this method should be platform-aware, returning the appropriate minimum capability for both CUDA and ROCm.

    def get_min_capability(cls) -> int:
        if current_platform.is_rocm():
            # gfx90a (MI200) is the first ROCm GPU with widespread Triton support.
            return 90
        # Volta support is limited; Turing (SM 7.5) is a safer baseline for CUDA.
        return 75

[shivampr]

Ignoring this as I did not touch this code, I can take this up in followup PR.

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