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 fc3fd6c19..646591909 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 @@ -60,6 +60,9 @@ if _is_cuda: if _use_aiter: + from aiter.ops.fused_qk_norm_rope_cache_quant import ( + fused_qk_rmsnorm as fused_qk_rmsnorm_bf16, + ) from aiter.ops.triton.batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant import ( batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant, ) @@ -160,6 +163,15 @@ class DeepseekMLAForwardMixin: output_unquantized_inp1=False, ) + elif _use_aiter: + q, k_nope = fused_qk_rmsnorm_bf16( + q, + self.q_a_layernorm.weight, + self.q_a_layernorm.variance_epsilon, + k_nope, + self.kv_a_layernorm.weight, + self.kv_a_layernorm.variance_epsilon, + ) else: q = self.q_a_layernorm(q) k_nope = self.kv_a_layernorm(k_nope)