From 839f7f2696f98db0af375b8ced62af4ebbe07efb Mon Sep 17 00:00:00 2001 From: Matti Varjokallio Date: Wed, 13 May 2026 10:27:09 +0300 Subject: [PATCH] [AMD] Add _skip_rope_for_aiter_fused_mla method and check to avoid double rotating with gfx950 and Aiter backend (#24148) --- .../attention_forward_methods/forward_mla.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index cd41e62bd..88f31c14b 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -352,10 +352,12 @@ class DeepseekMLAForwardMixin: q_nope_out = q_nope_out.transpose(0, 1) skip_rope_for_nsa_tilelang_fused = self._skip_rope_for_nsa_tilelang_fused() + skip_rope_for_aiter_fused_mla = self._skip_rope_for_aiter_fused_mla() if ( self.rotary_emb is not None and (not self._fuse_rope_for_trtllm_mla(forward_batch)) and (not skip_rope_for_nsa_tilelang_fused) + and (not skip_rope_for_aiter_fused_mla) and (not _use_aiter or not _is_gfx95_supported or self.use_nsa) ): q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe) @@ -664,3 +666,15 @@ class DeepseekMLAForwardMixin: or server_args.nsa_prefill_backend == "tilelang" ) ) + + def _skip_rope_for_aiter_fused_mla(self: DeepseekV2AttentionMLA) -> bool: + """ + Skip rope in prepare and let the fused kernel in forward_absorb_core handle it, + when running aiter-backend MLA on gfx95 (i.e., the `else` branch in forward_absorb_core + that calls fused_qk_rope_cat_and_cache_mla). + """ + return ( + _use_aiter_gfx95 + and self.current_attention_backend + not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS + )