Fused two elementwise kernels for k_nope and k_pe concat (#14862)

1 TC failure to check, but irrelevant to this code change
This commit is contained in:
kk
2025-12-15 17:33:33 +08:00
committed by GitHub
parent 1e2d753804
commit 2ea844ec81

View File

@@ -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