[Bugfix][Hardware][AMD] Fix hardcoded device in AITER MLA and Fused MOE

[c0de128]

Summary

Fix hardcoded device="cuda" in AITER MLA sparse attention and Fused MOE initialization code. This ensures tensors are created on the correct device in multi-GPU setups and improves ROCm compatibility.

Changes

1. vllm/attention/ops/rocm_aiter_mla_sparse.py

Replace 5 instances of device="cuda" with device=q.device or device=q_fp8.device:

• Lines 46, 49: Use q.device in fp8_mqa_logits_torch()
• Lines 127, 135: Use q.device in fp8_paged_mqa_logits_torch()
• Line 194: Use q_fp8.device in rocm_fp8_paged_mqa_logits()

2. vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

Add device parameter to init_aiter_topK_meta_data() function:

• Replace 3 instances of device="cuda" with device=device parameter
• Default value device="cuda" maintains backward compatibility

3. vllm/model_executor/layers/fused_moe/layer.py

Update call site to pass current_platform.device_type to the init function.

Test Plan

• Code inspection verified correct device propagation
• Linting passed (ruff format, ruff check)
• The fix follows existing vLLM patterns for device handling

cc @hongxiayang @tjtanaa

---

Note

Ensures tensors are allocated on the correct device (CUDA/ROCm) instead of being hardcoded to CUDA, improving multi-GPU correctness and ROCm compatibility.

• In vllm/v1/attention/ops/rocm_aiter_mla_sparse.py, uses q.device/q_fp8.device for torch.arange, tensor fills, and logits buffers.
• In vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py, adds device param to init_aiter_topK_meta_data() and switches tensor allocations to device=device.
• In vllm/model_executor/layers/fused_moe/layer.py, passes current_platform.device_type to init_aiter_topK_meta_data() for shared-expert TopK buffers.

^{Written by Cursor Bugbot for commit d87dadc. This will update automatically on new commits. Configure here.}

[c0de128]

/ci-run

[gemini-code-assist]

Code Review

This pull request correctly addresses the issue of hardcoded device="cuda" in AITER MLA and Fused MOE components. The changes are well-implemented, replacing the hardcoded device with dynamic device information from existing tensors or function parameters. This is a crucial fix for improving ROCm compatibility and ensuring correctness in multi-GPU setups. The approach of adding a device parameter with a backward-compatible default is sound. Overall, this is a solid bugfix that improves the robustness and portability of the code.

[c0de128]

Hi @tjtanaa, AMD CI passed (#2374). This fixes hardcoded device parameters in AITER MLA and Fused MOE code. Would appreciate a review when you have time. Thanks!

[mergify]

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

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

[c0de128]

/buildkite run

[c0de128]

Closing this PR to reduce maintainer review burden. The fix is available in this branch if needed in the future. Thank you for your time!