[ROCm][WIP]: Fused aiter rope kvcache mla

[Rohan138]

Fuse RoPE+cache+cat ops for MLA e.g. DeepSeek, Kimi on ROCm + AITER similar to #33443.

Requires:

• [ROCm]: Enable customop and rope+kvcache fusion for AITER RoPE #35180
• feat(mla): add default do_kv_cache_update for MLA #33658
• [Performance] Extract kv update ops from MLA attention backends #34627

Purpose

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 introduces a fusion for RoPE, KV cache update, and concatenation operations for MLA on ROCm with AITER, which is a significant performance optimization. The changes include adding a new custom op for rotary embedding, refactoring the MLA attention layer to separate KV cache updates, and updating tests and configurations. The implementation looks mostly correct, but I've found a potential logic bug in the configuration and a minor inconsistency in function signatures that should be addressed.

[gemini-code-assist]

There seems to be a logic inversion here. The docstring on line 143 states that this fusion should be enabled when using 'AITER Triton GEMMs', but the code on line 149 checks for not rocm_aiter_ops.is_triton_gemm_enabled(). This appears to be a bug that would incorrectly disable the fusion when Triton GEMMs are in use. Should this be rocm_aiter_ops.is_triton_gemm_enabled() to match the docstring and the previous logic?

Suggested change

	rocm_aiter_ops.is_rmsnorm_enabled()
	and not rocm_aiter_ops.is_triton_gemm_enabled()
	rocm_aiter_ops.is_rmsnorm_enabled()
	and rocm_aiter_ops.is_triton_gemm_enabled()

[mergify]

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

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

[mergify]

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

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