[Feature] [ROCm]: AITER Kernel Integration

[tjtanaa]

🚀 The feature, motivation and pitch

This is an issue that tracks PRs related to AITER https://github.com/ROCm/aiter .

AITER is AMD’s centralized repository that support various of high performance AI operators for AI workloads acceleration, where a good, unified place for all the customer operator-level requests, which can match different customers' needs. Developers can focus on operators, and let the customers integrate this op collection into their own private/public/whatever framework.

Note: This issue tracker description has been reorganized.

Based on AITER commit: 5a77249

The kernels from #14007 has been broken down into the following PRs for ease of review:

• AITER Linear ([FEAT] [ROCm]: Support AITER Linear #14916)
• AITER RMS Norm ([FEAT] [ROCm]: Add AITER RMS Norm (Layer Norm) Feature #14959)
• AITER Fused MoE + Block Scaled Fused MoE ([FEAT][ROCm] Integrate Fused MoE Kernels from AITER #14967)
• AITER Block Scaled A8W8 GEMM ([FEAT] [ROCm]: Add AITER Block-Scaled GEMM Feature #14968)
• AITER Paged Attention ([FEAT][ROCm] Integrate Paged Attention Kernel from AITER #15001)
• AITER INT8 a8w8 GEMM kernel ([FEAT] [ROCm] Add AITER int8 scaled gemm kernel #15433)
• AITER MLA ([FEAT][ROCm]: Support AITER MLA #15893)
• AITER Tkw1 for Llama4 FP8 ([ROCm] Add aiter tkw1 kernel for Llama4 fp8 #16727) ([ROCm] (Deprecated) Enable AITER Tkw1 kernel #16418)
• AITER CK_MoE for Llama4 BF16 ([ROCM] enable aiter fused moe kernel for llama4 bf16 checkpoints #16674)
• Enable AITER Fused MoE in V1 Engine ([FEAT] [ROCm]: AITER Fused MOE V1 Support #16752) To be merged after
  • AITER Tkw1 ([ROCm] Add aiter tkw1 kernel for Llama4 fp8 #16727)
  • AITER CK_MoE for Llama4 ([ROCM] enable aiter fused moe kernel for llama4 bf16 checkpoints #16674)
• AITER 2Stage CK MoE [FEAT] [ROCm]: Add AITER CK 2 Stages MoE support #17110
• AITER MLA V1 ([FEAT][ROCm]: Support AITER MLA on V1 Engine #17523)
• AITER biased group topk ([FEAT] [ROCm] [V1]: Add AITER biased group topk for DeepSeekV3 #17955)

Enhancement::

• Restrict Fused MoE based on Model that are actually using the kernel [Misc][ROCm] Restrict Aiter moe to specific models. #16435
• [BugFix] [ROCm]: Bugfix and handle addition case of input for rocm_aiter_rms_norm #17857

Bugfix

• [BUG] [ROCm] [MLA] Fix variable name bug due to change in variable name in PR #17483 #17961

Based on AITER commit: c1debd87ce0391aa27438d9e07e76e4fea7c4b70

• Fix MLA Backend v0 due to AITER API change in newer version ([BugFix][AMD] Compatible patch for latest AITER(05/07/2025) #17864)
  • It has been reverted (Revert "[BugFix][AMD] Compatible patch for latest AITER(05/07/2025)" #17910) as it introduced new properties that causes pre-commit to fail. The bug fix PR is ([BugFix][AMD] Compatible patch for AITER lib after 04/20 #17912)
• Use AITER fused moe external API ([FEAT] [ROCm] Upgrade AITER Fused MoE kernels. #18271)
• [FEAT][ROCm] Upgrade AITER MLA v1 backend #18338
• [FEAT][ROCm] Add AITER grouped topk for DeepSeekV2 #18825
• Enable full context length of DeepSeekV3 [ROCm] Remove unnecessary assertion of max_model_len in ROCM_AITER_MLA attention backend. #18938

Based on AITER commit: 648764942e552a8bb5fe16026703716a81f05374

• AITER MHA V1 ([Hardware][AMD] integrate aiter chunked prefill into vllm #18596) ([Hardware][AMD] integrate aiter into vllm #17710)
• Patch for new AITER commit ([ROCm] [AITER] [Bugfix] Patch for AITER commit 648764942e552a8bb5fe16026703716a81f05374 #18990)

Archived on 2025-05-14

The kernels from #14007 has been broken down into the following PRs for ease of review:

• AITER Linear ([FEAT] [ROCm]: Support AITER Linear #14916)
• AITER RMS Norm ([FEAT] [ROCm]: Add AITER RMS Norm (Layer Norm) Feature #14959)
• AITER Fused MoE + Block Scaled Fused MoE ([FEAT][ROCm] Integrate Fused MoE Kernels from AITER #14967)
• AITER Block Scaled A8W8 GEMM ([FEAT] [ROCm]: Add AITER Block-Scaled GEMM Feature #14968)
• AITER Paged Attention ([FEAT][ROCm] Integrate Paged Attention Kernel from AITER #15001)
• AITER INT8 a8w8 GEMM kernel ([FEAT] [ROCm] Add AITER int8 scaled gemm kernel #15433)
• AITER MLA ([FEAT][ROCm]: Support AITER MLA #15893)
• AITER Tkw1 for Llama4 FP8 ([ROCm] Add aiter tkw1 kernel for Llama4 fp8 #16727) ([ROCm] (Deprecated) Enable AITER Tkw1 kernel #16418)
• AITER CK_MoE for Llama4 BF16 ([ROCM] enable aiter fused moe kernel for llama4 bf16 checkpoints #16674)
• Enable AITER Fused MoE in V1 Engine ([FEAT] [ROCm]: AITER Fused MOE V1 Support #16752) To be merged after
  • AITER Tkw1 ([ROCm] Add aiter tkw1 kernel for Llama4 fp8 #16727)
  • AITER CK_MoE for Llama4 ([ROCM] enable aiter fused moe kernel for llama4 bf16 checkpoints #16674)
• AITER 2Stage CK MoE [FEAT] [ROCm]: Add AITER CK 2 Stages MoE support #17110
• AITER MLA V1 ([FEAT][ROCm]: Support AITER MLA on V1 Engine #17523)
• Fix MLA Backend v0 due to AITER API change in newer version ([BugFix][AMD] Compatible patch for latest AITER(05/07/2025) #17864)
  • It has been reverted (Revert "[BugFix][AMD] Compatible patch for latest AITER(05/07/2025)" #17910) as it introduced new properties that causes pre-commit to fail. The bug fix PR is ([BugFix][AMD] Compatible patch for AITER lib after 04/20 #17912)
• AITER MHA V1 ([Hardware][AMD] integrate aiter into vllm #17710)
• AITER biased group topk ([FEAT] [ROCm] [V1]: Add AITER biased group topk for DeepSeekV3 #17955)

Enhancement::

• Restrict Fused MoE based on Model that are actually using the kernel [Misc][ROCm] Restrict Aiter moe to specific models. #16435
• [BugFix] [ROCm]: Bugfix and handle addition case of input for rocm_aiter_rms_norm #17857

Bugfix

• [BUG] [ROCm] [MLA] Fix variable name bug due to change in variable name in PR #17483 #17961

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.