Model update

blocks_jvlm.py

@@ -1,7 +1,6 @@
 # Copyright 2025 Jina AI. All rights reserved.
 
 from abc import ABCMeta, abstractmethod
-from contextlib import nullcontext
 from copy import deepcopy
 from functools import wraps
 from math import prod, sqrt
@@ -712,7 +711,6 @@ def eager_attention_forward(
     dropout: float = 0.0,
     **_,
 ):
-    assert isinstance(module.num_key_value_groups, int)
     key_states = repeat_kv(key, module.num_key_value_groups)
     value_states = repeat_kv(value, module.num_key_value_groups)
 
@@ -1239,7 +1237,13 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
             if config.pooling_type == ImagePooling2DType.attention_2wide:
                 pooling_input_size *= 2
 
-            attn_implementation, _ = self._resolve_attn_pooling(attn_implementation)
+            # Flash Attention can cause Inf grads in the attention pooling layer
+            # because of very large batch sizes. Setting this to sdpa does not cost us
+            # much since sequence lengths in the case of attention pooling are very
+            # small
+            attn_implementation = attn_implementation or 'eager'
+            if attn_implementation.startswith('flash'):
+                attn_implementation = 'sdpa'
             self.pooling = MHSDPA(
                 config.attn_pooling_config,
                 hidden_size=pooling_input_size,
@@ -1280,23 +1284,6 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
         self.projector_dropout = Dropout(config.projector_dropout)
         self.feature_dropout = Dropout(config.feature_dropout)
 
-    @staticmethod
-    def _resolve_attn_pooling(attn_implementation: Optional[str] = None):
-        """
-        Flash Attention can cause Inf grads in the attention pooling layer because of
-        very large batch sizes. Setting this to sdpa does not cost us much since
-        sequence lengths in the case of attention pooling are tiny
-        """
-        attn_runtime_ctx = nullcontext()
-        if (
-            attn_implementation is not None
-            and attn_implementation.startswith('flash')
-        ):
-            attn_implementation = 'sdpa'
-            attn_runtime_ctx = sdpa_kernel(backends=[SDPBackend.MATH])
-
-        return attn_implementation, attn_runtime_ctx
-
     def forward(
         self,
         image_features: torch.Tensor,
@@ -1361,10 +1348,22 @@ class VisionLanguageConnector(GradientCheckpointingLayer):
             image_features = image_features.contiguous()
             if self.pooling_type == ImagePooling2DType.attention_meanq:
                 query = image_features.mean(-2, keepdim=True)
-                attn_implementation, attn_runtime_ctx = self._resolve_attn_pooling(
-                    attn_implementation
-                )
-                with attn_runtime_ctx:
+                # Flash Attention can cause Inf grads in the attention pooling layer
+                # because of very large batch sizes. Setting this to sdpa does not cost
+                # us much since sequence lengths in the case of attention pooling are
+                # very small
+                attn_implementation = attn_implementation or 'eager'
+                if attn_implementation.startswith('flash'):
+                    attn_implementation = 'sdpa'
+                if attn_implementation == 'sdpa':
+                    with sdpa_kernel(backends=[SDPBackend.MATH]):
+                        image_features, _ = self.pooling(
+                            xq=query,
+                            xk=image_features,
+                            attn_implementation='sdpa',
+                            **kwargs,
+                        )
+                else:
                     image_features, _ = self.pooling(
                         xq=query,
                         xk=image_features,