Only register rocm_aiter_ops if aiter is found

[mgoin]

Purpose

In #24490 there was an addition of a new import for rocm.py

Without this change, you can see the rocm import warnings on nvidia devices by default

WARNING 11-10 22:31:06 [rocm.py:39] Failed to import from amdsmi with ModuleNotFoundError("No module named 'amdsmi'")
WARNING 11-10 22:31:06 [rocm.py:50] Failed to import from vllm._rocm_C with ModuleNotFoundError("No module named 'vllm._rocm_C'")

Thanks to @tjtanaa, he pointed out there was an oversight and the ops class should act as a no-op when it is on non-ROCm platform and aiter is not installed.

I found the culprit using

import traceback
print(traceback.print_stack())

Output:

python benchmarks/kernels/bench_block_fp8_gemm.py
Printing stack trace
  File "/home/mgoin/code/vllm/benchmarks/kernels/bench_block_fp8_gemm.py", line 6, in <module>
    from vllm.model_executor.layers.quantization.utils.fp8_utils import (
  File "<frozen importlib._bootstrap>", line 1360, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1331, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 935, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 999, in exec_module
  File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
  File "/home/mgoin/code/vllm/vllm/model_executor/layers/quantization/utils/fp8_utils.py", line 15, in <module>
    from vllm._aiter_ops import rocm_aiter_ops
  File "<frozen importlib._bootstrap>", line 1360, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1331, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 935, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 999, in exec_module
  File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
  File "/home/mgoin/code/vllm/vllm/_aiter_ops.py", line 941, in <module>
    rocm_aiter_ops.register_ops_once()
  File "/home/mgoin/code/vllm/vllm/_aiter_ops.py", line 35, in wrapper
    from vllm.platforms.rocm import on_gfx9
  File "<frozen importlib._bootstrap>", line 1360, in _find_and_load
  File "<frozen importlib._bootstrap>", line 1331, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 935, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 999, in exec_module
  File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
  File "/home/mgoin/code/vllm/vllm/platforms/rocm.py", line 27, in <module>
    print(traceback.print_stack())

Test Plan

Test Result

---

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.

[gemini-code-assist]

Code Review

This pull request correctly addresses an issue where importing rocm_aiter_ops at the module level in fp8_utils.py caused unnecessary warnings on non-ROCm platforms. By moving the import into the _run_aiter method, it becomes a lazy import, ensuring that ROCm-specific code is only loaded when it is actually used. This change is clean, effective, and has no significant performance impact. The fix is well-implemented and resolves the reported problem.

[yewentao256]

LGTM, thanks for the work!

[tjtanaa]

LGTM. Thank you for the fix.