[ROCm] Add AITER RoPE + KV cache fusion for MLA prefill and decode

[khairulkabir1661]

Summary

This PR adds AITER fused kernel support for DeepSeek MLA attention on AMD ROCm, implementing RoPE + KV cache fusion for both prefill and decode paths. This builds on top of #38299 (RMSNorm + FP8 quantization fusion).

Changes

Core Functionality

• AITER fused decode kernel: Fuses RoPE application and KV cache writes for decode tokens using AMD's AITER library
• Prefill RoPE + KV cache fusion: Separate fusion path for prefill tokens
• Mixed batch handling: Correctly handles batches containing both prefill and decode tokens

Implementation Details

• vllm/model_executor/layers/attention/mla_attention.py:
  • Add _run_aiter_fused_decode() for fused decode path
  • Add prefill fusion in forward_impl()
  • Update unified_mla_kv_cache_update() to skip KV writes when using fused paths
  • Add custom ops for torch.compile/CUDA graph compatibility
• vllm/model_executor/layers/mla.py:
  • Pass RoPE caches and modules to MLAAttention
  • Skip RoPE in mla.py when using fused path (applied in custom op instead)
• vllm/envs.py:
  • Add VLLM_ROCM_USE_AITER flag to enable AITER kernels (default: enabled on ROCm)

Code Quality

• Clean up verbose comments throughout mla_attention.py
• Remove unused parameters and debug logging
• Simplify docstrings and inline comments

Testing

Tested on AMD MI300X with DeepSeek-V3:

• ✅ Mixed batches (prefill + decode)
• ✅ Pure prefill batches
• ✅ Pure decode batches
• ✅ CUDA graph mode
• ✅ Eager mode

Performance

Expected improvements on AMD MI300X:

• Reduced memory bandwidth usage (fused RoPE + KV cache write)
• Better kernel launch overhead (fewer separate operations)

Dependencies

• Requires: [ROCm] Add AITER RMSNorm+FP8 quantization fusion for MLA #38299 (RMSNorm + FP8 quantization fusion)
• AITER library must be installed on ROCm systems

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

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

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 includes a wide range of updates, including hardware support enhancements (AMD, RISC-V, XPU), build system improvements (Docker, CMake), new features like adaptive concurrency search and weight transfer backends, and refactoring of entrypoint tests. A critical issue was identified in csrc/cumem_allocator.cpp where a ROCm-specific workaround is being executed on CUDA platforms, which should be guarded by an #ifdef USE_ROCM block.

[gemini-code-assist]

This block of code is a workaround for a ROCm driver issue, as stated in the comment. However, it is not guarded by #ifdef USE_ROCM and will therefore also be executed on CUDA platforms. This is unnecessary for CUDA and could potentially introduce performance overhead or instability, as it relies on driver-specific behavior of cuMemAddressFree and cuMemAddressReserve. This workaround should be wrapped in an #ifdef USE_ROCM block to ensure it only applies to ROCm builds.

#ifdef USE_ROCM
  // ROCm workaround: hipMemRelease does not return physical VRAM to the
  // free pool while the virtual-address reservation is still held.
  // Cycling cuMemAddressFree → cuMemAddressReserve (at the same address)
  // forces the driver to actually release the physical pages while keeping
  // the same VA available for a later create_and_map.
  if (first_error == no_error) {
    first_error = cuMemAddressFree(d_mem, size);
    if (first_error == no_error) {
      CUdeviceptr d_mem_new = 0;
      first_error = cuMemAddressReserve(&d_mem_new, size, 0, d_mem, 0);
      if (first_error == no_error && d_mem_new != d_mem) {
        cuMemAddressFree(d_mem_new, size);
        snprintf(error_msg, sizeof(error_msg),
                 "ROCm: VA re-reserve got %p instead of %p", (void*)d_mem_new,
                 (void*)d_mem);
        error_code = CUresult(1);
        std::cerr << error_msg << std::endl;
        return;
      }
    }
  }
#endif

[mergify]

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

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