Add FA2-capable modeling file but disable flash_attention_2 via config (incompatible with PyLate query expansion; ~10 ndcg@10 loss)

README.md

@@ -67,3 +67,7 @@ reranked = rank.rerank(
 ```
 
 `trust_remote_code=True` is required: it loads `modeling_lfm2_bidirectional.py`, which replaces the causal attention mask and short-conv padding with their non-causal equivalents. Without it the model will silently use causal attention and produce poor embeddings.
+
+### Note on FlashAttention-2
+
+`attn_implementation="flash_attention_2"` is intentionally rejected at load time (`disable_flash_attention` in `config.json`). ColBERT query expansion pads queries with mask tokens that have `attention_mask=0` but are still scored in MaxSim; FA2's unpadding cannot compute embeddings for those positions, which severely degrades retrieval quality (~10 ndcg@10 points on NanoBEIR). Use the default `sdpa` (or `eager`).

config.json

@@ -60,5 +60,6 @@
   "vocab_size": 64402,
   "auto_map": {
     "AutoModel": "modeling_lfm2_bidirectional.Lfm2BidirectionalModel"
-  }
+  },
+  "disable_flash_attention": true
 }

modeling_lfm2_bidirectional.py

@@ -6,6 +6,17 @@ Wired into the HF repo via `auto_map` in config.json so that
     SentenceTransformer(repo, trust_remote_code=True)
 
 both return a model with the encoder-style patches already applied.
+
+Supports `attn_implementation` in {"eager", "sdpa", "flash_attention_2"}:
+eager/sdpa consume a 4D additive pad-only mask and reproduce the exact
+training-time behavior; flash_attention_2 receives the 2D padding mask (or
+None) and runs the kernel non-causally via `Lfm2Attention.is_causal = False`,
+yielding outputs equivalent to the unpadded forward.
+
+Repos may set `"disable_flash_attention": true` in config.json to reject
+flash_attention_2 at load time (used for ColBERT, where PyLate query expansion
+tokens — attention_mask=0 but scored in MaxSim — are incompatible with FA2
+unpadding and severely degrade retrieval quality).
 """
 
 from typing import Optional
@@ -14,6 +25,7 @@ import torch
 import torch.nn.functional as F
 from transformers.models.lfm2 import modeling_lfm2 as _lfm2_mod
 from transformers.models.lfm2.modeling_lfm2 import (
+    Lfm2Attention,
     Lfm2Model,
     Lfm2ShortConv,
     apply_mask_to_padding_states,
@@ -27,7 +39,14 @@ def _bidirectional_mask(
     cache_position: Optional[torch.LongTensor],
     past_key_values,
     position_ids: Optional[torch.LongTensor],
-) -> torch.Tensor:
+) -> Optional[torch.Tensor]:
+    if config._attn_implementation == "flash_attention_2":
+        # FA2 only uses the 2D padding mask to unpad sequences; causality is
+        # controlled by `Lfm2Attention.is_causal` (set to False below).
+        if attention_mask is not None and not attention_mask.all():
+            return attention_mask
+        return None
+
     device = input_embeds.device
     dtype = input_embeds.dtype
     bsz, q_len = input_embeds.shape[:2]
@@ -52,6 +71,10 @@ def _noncausal_shortconv_forward(
     cache_position=None,
     attention_mask: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
+    # eager/sdpa pass the 4D additive mask, on which this is a no-op — matching
+    # the behavior the checkpoints were trained with. FA2 passes the 2D pad
+    # mask, so pads are zeroed before the conv (closest match to the unpadded
+    # forward, since FA2 cannot reproduce the padded sdpa pad-state evolution).
     x = apply_mask_to_padding_states(hidden_states, attention_mask)
 
     BCx = self.in_proj(x).transpose(-1, -2)
@@ -98,5 +121,18 @@ class Lfm2BidirectionalModel(Lfm2Model):
     """LFM2 patched for encoder-style use: full bidirectional attention + non-causal short-conv."""
 
     def __init__(self, config):
+        if (
+            getattr(config, "_attn_implementation", None) == "flash_attention_2"
+            and getattr(config, "disable_flash_attention", False)
+        ):
+            raise ValueError(
+                "flash_attention_2 is disabled for this model: query expansion "
+                "tokens (attention_mask=0 but scored in MaxSim) are incompatible "
+                "with FA2 unpadding and severely degrade retrieval quality. "
+                "Load with attn_implementation='sdpa' (default) or 'eager'."
+            )
         _install_patches()
         super().__init__(config)
+        for module in self.modules():
+            if isinstance(module, Lfm2Attention):
+                module.is_causal = False