AITER backend extension and workload optimizations

[HaiShaw]

Co-author: @kkHuang-amd

Motivation

• DeepSeek optimization
• 1 simple flag: SGLANG_USE_AITER

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Hello @HaiShaw, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

Summary of Changes

Hello team,

Gemini here, providing a summary for this pull request. This PR introduces significant changes related to integrating and optimizing the AITER backend, particularly focusing on MLA (Multi-Layer Attention) architecture and general workload optimizations on HIP (AMD) devices. The core intent is to enable and leverage AITER kernels for improved performance, especially for models like DeepSeek v2, controlled by a new environment variable SGLANG_USE_AITER.

The changes involve updating documentation, modifying attention and layer normalization layers to conditionally use AITER kernels, adapting MoE and FP8 quantization layers for AITER compatibility, and updating the model runner and DeepSeek v2 model implementation to dispatch to the AITER backend when enabled. MLA-specific logic and data structures have been added to the AITER attention backend to support this architecture.

Highlights

• AITER Backend Integration: Introduces a new environment variable SGLANG_USE_AITER to enable the AITER backend for HIP devices, replacing the old SGLANG_AITER_MOE flag.
• MLA Attention Support: Adds specific logic and data structures within the AITER attention backend (aiter_backend.py) to support the MLA attention architecture, including new metadata fields and an MLA-specific index updater.
• AITER Kernel Usage: Integrates AITER kernels for attention (flash_attn_varlen_func, mla_decode_fwd) and RMS normalization (rmsnorm2d_fwd, rmsnorm2d_fwd_with_add) when the AITER backend is enabled.
• MoE and Quantization Updates: Updates MoE and FP8 quantization layers to respect the new SGLANG_USE_AITER flag and utilize AITER-compatible functions (gemm_a8w8_blockscale_CK, ck_moe_2stages, shuffle_weight) on HIP devices.
• Model and Runner Configuration: Configures the model runner to select the AITER attention backend for MLA architecture on HIP devices and updates the DeepSeek v2 model implementation to work with the AITER backend, including handling MLA-specific requirements like kv_b_proj.

Changelog

Click here to see the changelog

• docs/references/environment_variables.md
  • Renamed environment variable SGLANG_AITER_MOE to SGLANG_USE_AITER and updated its description.
• python/sglang/srt/layers/attention/aiter_backend.py
  • Imported new AITER attention and MLA decode functions.
  • Added MLA-specific fields (qo_indptr, kv_last_page_len, max_extend_len, max_prefix_extend_len) to ForwardMetadata.
  • Added use_mla flag based on model config.
  • Conditionally initialized AiterMlaIndicesUpdaterPrefill and workspace_buffer based on use_mla.
  • Modified init_forward_metadata and init_forward_metadata_capture_cuda_graph to handle MLA-specific metadata initialization.
  • Updated forward_extend and forward_decode to dispatch to AITER MLA kernels or existing AITER kernels based on use_mla.
  • Added AiterMlaIndicesUpdaterPrefill class for MLA-specific index updates.
• python/sglang/srt/layers/layernorm.py
  • Introduced _use_aiter flag.
  • Conditionally imported AITER RMS norm kernels.
  • Added forward_aiter method using AITER kernels.
  • Updated the main forward method to dispatch to forward_aiter when enabled.
• python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
  • Replaced SGLANG_AITER_MOE with SGLANG_USE_AITER in conditional checks.
• python/sglang/srt/layers/moe/fused_moe_triton/layer.py
  • Introduced _use_aiter flag.
  • Conditionally imported AITER MoE functions.
  • Replaced checks for _is_hip and get_bool_env_var("SGLANG_AITER_MOE") or use_aiter_moe with _use_aiter in various methods (process_weights_after_loading, forward_cuda, process_weights_hip_int4, process_weights_hip_scale_padding, maybe_apply_hip_fused_experts).
• python/sglang/srt/layers/quantization/fp8.py
  • Renamed internal flags use_hip_int4 and use_aiter_moe.
  • Defined _use_aiter based on the new environment variable.
  • Conditionally imported gemm_a8w8_blockscale_CK from AITER.
  • Replaced old flag checks with new ones.
  • Updated aiter_w8a8_block_fp8_linear to use the new AITER GEMM function.
• python/sglang/srt/model_executor/model_runner.py
  • Added logic to select 'aiter' as the attention backend for MLA architecture on HIP devices.
  • Included 'aiter' in the list of supported MLA attention backends.
• python/sglang/srt/models/deepseek_v2.py
  • Introduced _use_aiter flag.
  • Conditionally imported AITER rope function.
  • Initialized self.attn_mha.kv_b_proj to None in __init__.
  • Modified _dispatch_mla_subtype to handle 'aiter' backend selection based on forward mode.
  • Ensured self.attn_mha.kv_b_proj is set in forward_prepare.
• scripts/amd_ci_exec.sh
  • Updated the environment variable used in AMD CI scripts from SGLANG_AITER_MOE to SGLANG_USE_AITER.
• test/srt/test_nightly_gsm8k_eval_amd.py
  • Updated the environment variable used in tests from SGLANG_AITER_MOE to SGLANG_USE_AITER.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

---

AITER's kernels, swift and bright,
Bring MLA speed, day and night.
With flags set true,
Optimized for you,
DeepSeek runs fast, bathed in light.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces significant extensions to the AITER backend, including optimizations for Multi-Latent Attention (MLA) primarily targeting DeepSeek models. The changes involve new AITER kernel integrations, modifications to metadata handling, and conditional logic for MLA vs. non-MLA paths. The environment variable SGLANG_AITER_MOE has also been aptly renamed to the more generic SGLANG_USE_AITER.

Overall, the changes are aimed at improving performance and extending backend capabilities. However, a critical issue regarding uninitialized variables has been identified that needs to be addressed. There's also a point regarding in-place modification of a shared buffer that would be good to clarify.

Summary of Findings

• Uninitialized Variables in CUDA Graph Metadata: In aiter_backend.py, within init_forward_metadata_capture_cuda_graph, if self.use_mla is false, qo_indptr, kv_last_page_len, and max_extend_len are used without prior initialization, which will cause a runtime error. (Severity: Critical)
• In-place Modification of Backend Buffer: In AiterMlaIndicesUpdaterPrefill.update_single_wrapper (aiter_backend.py), self.attn_backend.kv_indptr is modified in-place. This needs confirmation that it's safe and doesn't cause issues with buffer reuse. (Severity: Medium)
• Environment Variable Rename: The environment variable SGLANG_AITER_MOE was consistently renamed to SGLANG_USE_AITER across documentation, scripts, and tests, which improves clarity and generality.
• AITER Backend and MLA Integration: The PR successfully integrates AITER backend enhancements, including specific paths for Multi-Latent Attention (MLA), particularly for DeepSeek models. This involves new kernel usages and conditional logic in attention and layernorm layers.

Merge Readiness

This pull request makes substantial progress in extending the AITER backend and optimizing for MLA. However, there is a critical issue identified concerning uninitialized variables in python/sglang/srt/layers/attention/aiter_backend.py (within init_forward_metadata_capture_cuda_graph) that must be addressed before merging, as it will likely cause runtime errors. Additionally, there's a medium-severity question about an in-place buffer modification that would be good to clarify to ensure long-term stability.

I recommend that these changes, especially the critical issue, be addressed before this PR is merged. As an AI, I am not authorized to approve pull requests; please ensure further review and approval from other maintainers.

[HaiShaw]

cc @kkHuang-amd

[HaiShaw]

@saienduri - note the flag change to SGLANG_USE_AITER

[zhyncs]

@HaiShaw May you fix the conflicts? Thanks!