From 8589b92a891b6142cf5f068cde0f87c812a568b5 Mon Sep 17 00:00:00 2001 From: Alan Kao Date: Tue, 21 Apr 2026 17:35:09 +0800 Subject: [PATCH] [AMD] Fused qk rmsnorm bf16 for amd/Kimi-K2.5-MXFP4 (#23186) Co-authored-by: Claude Sonnet 4.6 --- .../attention_forward_methods/forward_mla.py | 12 ++++++++++++ 1 file changed, 12 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 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)