[BUG] [ROCm] Fix import bug on ROCm

[tjtanaa]

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.

Purpose

Fix the bug introduce by PR #21069
following the bugfix PR #21036

Bug Log

1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/executor/executor_base.py", line 53, in __init__
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     self._init_executor()
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/executor/uniproc_executor.py", line 49, in _init_executor
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     self.collective_rpc("load_model")
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/executor/uniproc_executor.py", line 58, in collective_rpc
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     answer = run_method(self.driver_worker, method, args, kwargs)
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/utils/__init__.py", line 2948, in run_method
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     return func(*args, **kwargs)
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/v1/worker/gpu_worker.py", line 212, in load_model
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     self.model_runner.load_model(eep_scale_up=eep_scale_up)
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/v1/worker/gpu_model_runner.py", line 1972, in load_model
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     self.drafter.load_model(self.model)
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/v1/spec_decode/eagle.py", line 389, in load_model
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     from vllm.compilation.backends import set_model_tag
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/compilation/backends.py", line 27, in <module>
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     from .pass_manager import PostGradPassManager
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/compilation/pass_manager.py", line 12, in <module>
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     from .collective_fusion import AllReduceFusionPass, AsyncTPPass
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/app/fix-specdec-rocm/vllm/compilation/collective_fusion.py", line 41, in <module>
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]   File "/usr/local/lib/python3.10/dist-packages/torch/_ops.py", line 1267, in __getattr__
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667]     raise AttributeError(
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m ERROR 08-01 09:45:20 [core.py:667] AttributeError: '_OpNamespace' '_C' object has no attribute 'scaled_fp4_quant'
^[[1;36m(EngineCore_0 pid=162345)^[[0;0m Process EngineCore_0:

Test Plan

Test Result

(Optional) Documentation Update

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[gemini-code-assist]

Code Review

This pull request aims to fix an import bug on ROCm by restricting some CUDA-specific imports. While the change correctly identifies and restricts these imports, it introduces a critical runtime error. The code that uses the imported modules is not conditionally guarded, which will lead to a NameError on ROCm platforms. I've provided a critical comment with details on how to resolve this.

[gemini-code-assist]

This change correctly restricts CUDA-specific imports to the CUDA platform, which should fix the import error on ROCm. However, this introduces a new problem. The passes imported here (FusionPass, AllReduceFusionPass, etc.) are used later in this file within PostGradPassManager.configure based on config.compilation_config.pass_config flags.

On a ROCm platform, this if condition will be false, the imports will be skipped, but the code that uses these passes will still execute if the corresponding config flags are enabled. This will lead to a NameError at runtime.

To fix this, the usage of these passes within PostGradPassManager.configure must also be guarded by if current_platform.is_cuda():.

[ProExpertProg]

@tjtanaa this was a correct suggestion by Gemini, please pay attention next time!

[ProExpertProg]

As @gshtras pointed out in Slack this breaks custom passes on ROCm. We should only disable the pass that uses the scaled_fp4_quant function

[tjtanaa]

@ProExpertProg Thank you for pointing out. So, I was following the wrong fix.
@gshtras Thank you very much for your explanation.