Update attention.py

InternLM · Aug 13, 2024 · bc41ee8 · bc41ee8
1 parent bd28606
commit bc41ee8
Showing 1 changed file with 0 additions and 2 deletions.
diff --git a/internlm/model/ops/attention.py b/internlm/model/ops/attention.py
@@ -376,7 +376,6 @@ def _npu_varlen_kvpacked_attn(
 ):
     # TODO: support npu native varlen flash attention
     k, v = kv.unbind(dim=2)
-    # k, v = k.squeeze(dim=2), v.squeeze(dim=2)
     return _npu_varlen_qkvsplited_attn(
         q,
         k,
@@ -393,7 +392,6 @@ def _npu_varlen_kvpacked_attn(
 
 def _npu_fixedlen_kvpacked_attn(q: torch.Tensor, kv: torch.Tensor, dropout_p: float, softmax_scale=None, causal=False):
     k, v = kv.unbind(dim=2)
-    # k, v = k.squeeze(dim=2), v.squeeze(dim=2)
     return _npu_fixedlen_qkvsplited_attn(q, k, v, dropout_p, softmax_scale, causal)