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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user