Update modeling/lance/qwen2_navit.py

modeling/lance/qwen2_navit.py

@@ -499,34 +499,34 @@ class PackedAttentionMoT(Qwen2Attention):
         cu_seqlens_k = torch.nn.functional.pad(torch.cumsum(key_values_lens, dim=0), (1, 0))
 
         if FLASH_ATTN_AVAILABLE:
-    packed_attn_output = flash_attn_varlen_func(
-        q=packed_query_states,
-        k=merged_key_states,
-        v=merged_value_states,
-        cu_seqlens_q=cu_seqlens_q.to(torch.int32),
-        cu_seqlens_k=cu_seqlens_k.to(torch.int32),
-        max_seqlen_q=max(query_lens).item(),
-        max_seqlen_k=max(key_values_lens).item(),
-        causal=is_causal,
-    )
-else:
-    q = packed_query_states.transpose(0, 1).unsqueeze(0)
-    k = merged_key_states.transpose(0, 1).unsqueeze(0)
-    v = merged_value_states.transpose(0, 1).unsqueeze(0)
-
-    packed_attn_output = scaled_dot_product_attention(
-        q,
-        k,
-        v,
-        is_causal=is_causal,
-    )
-
-    packed_attn_output = (
-        packed_attn_output
-        .squeeze(0)
-        .transpose(0, 1)
-        .contiguous()
-    )
+            packed_attn_output = flash_attn_varlen_func(
+                q=packed_query_states,
+                k=merged_key_states,
+                v=merged_value_states,
+                cu_seqlens_q=cu_seqlens_q.to(torch.int32),
+                cu_seqlens_k=cu_seqlens_k.to(torch.int32),
+                max_seqlen_q=max(query_lens).item(),
+                max_seqlen_k=max(key_values_lens).item(),
+                causal=is_causal,
+            )
+        else:
+            q = packed_query_states.transpose(0, 1).unsqueeze(0)
+            k = merged_key_states.transpose(0, 1).unsqueeze(0)
+            v = merged_value_states.transpose(0, 1).unsqueeze(0)
+
+            with sdpa_kernel(backends=[SDPBackend.EFFICIENT_ATTENTION]):
+                packed_attn_output = scaled_dot_product_attention(
+                    q,
+                    k,
+                    v,
+                    is_causal=is_causal,
+                )
+
+            packed_attn_output = (
+                packed_attn_output.squeeze(0)
+                .transpose(0, 1)
+                .contiguous()
+            )
         packed_attn_output = packed_attn_output.reshape(-1, self.hidden_size)
         if mode == "und":
             packed_attn_output = self.o_proj(packed_attn_output)