From 2ea844ec81f42a7f96bd8a110391bcde5baccc16 Mon Sep 17 00:00:00 2001
From: kk <43161300+kkHuang-amd@users.noreply.github.com>
Date: Mon, 15 Dec 2025 17:33:33 +0800
Subject: [PATCH] Fused two elementwise kernels for k_nope and k_pe concat
 (#14862)

1 TC failure to check, but irrelevant to this code change
---
 python/sglang/srt/models/deepseek_v2.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py
index 1f5ce486a..403bea8bc 100644
--- a/python/sglang/srt/models/deepseek_v2.py
+++ b/python/sglang/srt/models/deepseek_v2.py
@@ -2698,6 +2698,9 @@ class DeepseekV2AttentionMLA(nn.Module):
                 attn_dtype = k_nope.dtype
             k = k_nope.new_empty(*k_shape, dtype=attn_dtype)
             concat_and_cast_mha_k_triton(k, k_nope, k_pe)
+        elif _is_hip and self.current_attention_backend == "aiter":
+            k = k_nope.new_empty(*k_shape)
+            concat_and_cast_mha_k_triton(k, k_nope, k_pe)
         else:
             k = k_nope.new_empty(*k_shape)
             k[..., : self.qk_nope_head_dim] = k_nope