Three-tier attention: torch varlen, flash_attn, eager (inline mask)

README.md

@@ -16,33 +16,57 @@ language:
 # CodeRankEmbed-flash-attn
 
 A **bf16 quantization of [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed)**
-with **flash-attention built into a custom `modeling_hf_nomic_bert.py` shipped in this repo.** It is
-not a finetune — the weights are the original CodeRankEmbed weights cast to bf16 (no further
-training) — and attention runs via `flash_attn` varlen instead of the original eager `O(seq²)` path.
+with a **three-tier attention dispatch built into a custom `modeling_hf_nomic_bert.py` shipped in
+this repo.** It is not a finetune — the weights are the original CodeRankEmbed weights cast to bf16
+(no further training). Two of the three tiers replace the original eager `O(seq²)` attention with an
+`O(N)` unpadded path; the third keeps the original eager algorithm as the correctness reference and
+universal fallback.
 
 ## Why
 
 `nomic-ai/CodeRankEmbed` loads through `trust_remote_code`, and its attention path is **eager
 only** — activation memory grows as `batch × heads × seq²`, which OOMs at large batches even though
-the model is only 137M params. `flash_attn`'s varlen path computes the same attention in `O(N)`
-memory by packing unpadded sequences, so the large batches that OOM the eager path run comfortably —
-with **parity embeddings** (no quality change). This repo ships that flash path in the modeling file
-itself, so no runtime patching or post-load hooks are needed.
+the model is only 137M params. This repo adds two attention paths that compute the same attention in
+`O(N)` memory by packing unpadded sequences, so the large batches that OOM the eager path run
+comfortably — with **parity embeddings** (no quality change):
+
+- `torch_varlen` — `torch.nn.attention.varlen.varlen_attn`, shipped in torch itself (no extra
+  package), available from torch **2.10.0** onward.
+- `flash_attn` — the original `flash_attn` varlen-packed kernel. Kept as a fallback for older torch
+  builds that don't yet have `torch.nn.attention.varlen` but do have `flash_attn` installed. This
+  package is **optional**.
+
+Both run the same FA2-family kernel and are gated by the same GPU-capability check. The modeling
+file ships all three paths itself, so no runtime patching or post-load hooks are needed.
 
 ## Behavior
 
-- **Loads bf16 by default.** `flash_attn` requires half precision and the model runs bf16 in any
-  real serving setup, so the weights are stored bf16 and `config.json` declares `torch_dtype:
-  bfloat16`. The upstream custom `from_pretrained` silently dropped `torch_dtype` and always loaded
-  fp32; the copy in this repo honors it, so the model loads bf16 natively, like any normal HF
-  model. Pass `torch_dtype=torch.float32` to load fp32 (note: the stored weights are bf16-precision,
-  so this only widens the dtype, not the precision).
-- **CUDA + `flash_attn`** → flash-varlen path (fast, low VRAM). Attention tensors are cast to bf16
-  internally as a safety net; with the default bf16 load this is a no-op.
-- **CPU, or no `flash_attn`** → the original eager attention algorithm runs unchanged. Because the
-  model is bf16, eager runs in bf16 here — numerically equivalent to the flash path, just without
-  its memory and throughput wins. The model loads and encodes on any host; the forward selects the
-  path automatically (`_FLASH_AVAILABLE and hidden_states.is_cuda`).
+- **Three-tier attention, chosen automatically per device** (override with
+  `NOMIC_BERT_ATTN_IMPL=torch_varlen|flash_attn|eager`):
+  1. **`torch_varlen`** — CUDA, compute capability sm_80+ (Ampere or newer), torch **≥ 2.10.0**. No
+     third-party kernel needed.
+  2. **`flash_attn`** — CUDA, compute capability sm_80+, the `flash_attn` package importable, on a
+     torch that doesn't yet ship `torch.nn.attention.varlen` (typically an older torch). This
+     dependency is **optional**.
+  3. **`eager`** — everything else: CPU, pre-Ampere GPUs, or neither of the above available. The
+     original padded attention algorithm, unchanged, runs on any host.
+
+  `auto` (the default) prefers `torch_varlen`, then `flash_attn`, then `eager`. A forced override
+  (e.g. `NOMIC_BERT_ATTN_IMPL=torch_varlen`) **raises** `RuntimeError` if that tier's precondition
+  doesn't hold — a forced tier never falls back silently. An unrecognized override raises
+  `ValueError`.
+- **See which tier engaged**: `model[0].auto_model.attention_impl` after a forward pass, or the
+  one-time `NomicBert attention impl=... device=... capability=... torch=... flash_attn=...
+  override=...` INFO log line (one line per distinct `(impl, device)`).
+- **Loads bf16 by default.** `flash_attn` and `torch_varlen` both require half precision and the
+  model runs bf16 in any real serving setup, so the weights are stored bf16 and `config.json`
+  declares `torch_dtype: bfloat16`. The upstream custom `from_pretrained` silently dropped
+  `torch_dtype` and always loaded fp32; the copy in this repo honors it, so the model loads bf16
+  natively, like any normal HF model. Pass `torch_dtype=torch.float32` to load fp32 (note: the
+  stored weights are bf16-precision, so this only widens the dtype, not the precision).
+- **`eager` runs in bf16 too** (because the stored weights are bf16), numerically equivalent to the
+  varlen tiers, just without their memory and throughput wins. The model loads and encodes on any
+  host regardless of which tier engages.
 
 ## Usage
 
@@ -63,35 +87,58 @@ d = model.encode(codes,   normalize_embeddings=True)
 ## Parity & performance
 
 The weights are the original CodeRankEmbed weights (bf16-cast), so embeddings match the fp32
-original to within bf16 precision. Measured on an RTX 3090 Ti, `flash_attn` 2.8.3, CLS pooling,
-L2-normalized, batch size 64:
-
-| metric | `nomic-ai/CodeRankEmbed` (fp32, eager) | this repo (bf16, flash-varlen) |
-| --- | --- | --- |
-| cosine vs fp32 reference | 1.000000 | **0.9986** |
-| peak VRAM (bs=64) | 6.7 GB | **2.1 GB** (≈3.3× less) |
-| throughput (bs=64) | 52,000 tok/s | **162,000 tok/s** (≈3.1× faster) |
-
-Same 512-snippet corpus (~236k tokens, 20–900 tokens each), batch size 64. The flash-varlen path
-also scales to far larger batches (bs=256 fits in ~7.6 GB) where eager O(seq²) blows up. Parity
-(>0.997) reproduced on both an RTX 4060 and the 3090 Ti.
+original to within bf16 precision. The table below re-measures all three tiers after adding the
+dispatch: each tier's output on a 64-code-snippet corpus (fp32-renormalized), compared by cosine
+similarity against this same repo's **pre-dispatch output on the same device** (`flash_attn` on
+CUDA, `eager` on CPU — `flash_attn` can't run on CPU). `eager` and `flash_attn` each run the same
+computation as before the dispatch was added, so their cosines are parity checks (round to
+1.000000 at 6 decimal places; true values are ≥ 0.9999997). `torch_varlen` is a different,
+torch-native kernel, and its cosine is the real signal.
+
+| GPU | tier | device | transformers | mean cos | min cos | peak VRAM |
+| --- | --- | --- | --- | --- | --- | --- |
+| RTX 3090 Ti (sm_86) | `flash_attn` | cuda | 5.16.1 | 1.000000 | 1.000000 | 1145 MiB |
+| RTX 3090 Ti (sm_86) | `torch_varlen` | cuda | 5.16.1 | 0.999941 | 0.999856 | 875 MiB |
+| RTX 3090 Ti (sm_86) | `eager` | cpu | 5.16.1 | 1.000000 | 1.000000 | – (CPU) |
+| RTX 5090 Laptop GPU (sm_120) | `flash_attn` | cuda | 5.11.0 | 1.000000 | 1.000000 | 1169 MiB |
+| RTX 5090 Laptop GPU (sm_120) | `torch_varlen` | cuda | 5.11.0 | 0.999944 | 0.999919 | 899 MiB |
+| RTX 5090 Laptop GPU (sm_120) | `eager` | cpu | 5.11.0 | 1.000000 | 1.000000 | – (CPU) |
+| CPU only | `eager` | cpu | 5.17.0 | 1.000000 | 1.000000 | – (CPU) |
+
+Cosines rounded to 6 decimal places; peak VRAM rounded to the nearest MiB. `eager` always runs on
+CPU in this protocol (it is the universal fallback tier); the GPU named in the first column is the
+host each row's measurement ran on, not the device `eager` used on that row. `torch` was
+`2.12.1+cu130` for every row except the standalone CPU-only row (`torch 2.14.0+cpu`), which used a
+separate `transformers==5.17.0` install to check the eager path against a newer transformers.
+
+Separately, both varlen tiers (`torch_varlen`, `flash_attn`) are gated in this repo's downstream
+test suite at min cosine > 0.997 against the fp32 `nomic-ai/CodeRankEmbed` reference (not
+re-measured here — see the table above for this repo's own numbers), and stay under 20 GB peak VRAM
+at batch size 256.
 
 ## What changed vs the source repo
 
-1. **Weights**: fp32 → bf16. `flash_attn` only accepts half precision and the model runs bf16 in
-   any real serving configuration, so the weights are stored bf16 and (via the load fix below)
-   arrive bf16 — which is simply how this model is used, and removes the need for a post-load dtype
-   cast. Parity-neutral (cosine 0.9986 vs the fp32 original); the smaller download is incidental,
-   not the reason.
+1. **Weights**: fp32 → bf16. `flash_attn` and `torch_varlen` only accept half precision and the
+   model runs bf16 in any real serving configuration, so the weights are stored bf16 and (via the
+   load fix below) arrive bf16 — which is simply how this model is used, and removes the need for a
+   post-load dtype cast. Parity-neutral; the smaller download is incidental, not the reason.
 2. **`from_pretrained` dtype fix**: the upstream custom `from_pretrained` instantiated the model
    fp32 and `load_state_dict`-ed the checkpoint into fp32 params, **ignoring `torch_dtype`**. The
    copy here adds the standard transformers dtype resolution (explicit arg → `config.torch_dtype` →
    checkpoint dtype) so the model loads in its declared dtype.
-3. **Flash-varlen forward**: `NomicBertAttention.forward` gains an `unpad →
-   flash_attn_varlen_qkvpacked_func(causal=False) → repad` branch (CUDA + flash_attn); the original
-   eager block is kept as the fallback. `NomicBertModel.forward` skips
-   `get_extended_attention_mask` on the flash path. Rotary embeddings are applied to the dense
-   `[B, S, 3, H, D]` tensor **before** unpadding — the correctness keystone.
+3. **Three-tier attention dispatch**: `NomicBertAttention.forward` now selects one of three
+   attention implementations at call time — `torch_varlen` (torch's own
+   `torch.nn.attention.varlen.varlen_attn`, no third-party kernel, needs torch ≥ 2.10.0),
+   `flash_attn` (the original flash-attn varlen-packed kernel, kept as a fallback for older torch
+   builds that have it installed), and `eager` (the original padded attention, unchanged, and the
+   default off CUDA sm_80+). Both varlen tiers unpad the input with torch-native `_unpad`/`_pad`
+   helpers (a replacement for `flash_attn.bert_padding`) before the packed kernel call, then repad
+   the output; `NomicBertModel.forward` builds the additive attention mask inline instead of calling
+   the (now-removed-upstream) `get_extended_attention_mask` helper. Rotary embeddings are applied to
+   the dense `[B, S, 3, H, D]` tensor **before** unpadding — the correctness keystone: applying RoPE
+   after unpadding would hand each packed position the wrong sequence's rotation. Set
+   `NOMIC_BERT_ATTN_IMPL=torch_varlen|flash_attn|eager` to force a tier (raises if it can't engage);
+   `model[0].auto_model.attention_impl` and the one-time INFO log line report which tier engaged.
 
 ## License & attribution
 

modeling_hf_nomic_bert.py

@@ -6,6 +6,8 @@
 import logging
 
 # Inspired by https://github.com/huggingface/transformers/blob/main/src/transformers/models/bert/modeling_bert.py
+import functools
+import importlib.metadata
 import os
 import re
 from collections import OrderedDict
@@ -26,27 +28,179 @@ from transformers.models.bert.modeling_bert import (
 from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME, WEIGHTS_INDEX_NAME, WEIGHTS_NAME
 from transformers.utils.hub import cached_file, get_checkpoint_shard_files
 
-# Flash-attention varlen path (optional). When ``flash_attn`` is importable and
-# the tensors live on CUDA, attention runs as unpadded varlen packing — O(N)
-# activation memory instead of the eager path's O(seq^2). When flash_attn is
-# unavailable (CPU-only hosts, or builds without it) the original eager path
-# runs byte-for-byte unchanged, so this model loads and encodes everywhere.
+from .configuration_hf_nomic_bert import NomicBertConfig
+
+# Three-tier attention dispatch (auto by default, override via NOMIC_BERT_ATTN_IMPL):
+#   1. torch_varlen — torch.nn.attention.varlen, no third-party kernel
+#   2. flash_attn   — the flash-attn varlen-packed kernel (optional dependency)
+#   3. eager        — the original padded path, byte-for-byte, runs everywhere
+# Both tier 1 and tier 2 run the same FA2-family kernel and refuse pre-Ampere
+# hardware, so a single capability predicate gates both (resolve_attention_impl,
+# below). When neither is available (CPU-only hosts, or builds without them)
+# tier 3 runs unchanged, so this model loads and encodes everywhere.
+try:  # pragma: no cover - import guard, exercised by environment
+    from torch.nn.attention.varlen import varlen_attn as _torch_varlen_attn
+
+    _TORCH_VARLEN_AVAILABLE = True
+except ImportError:  # pragma: no cover
+    _torch_varlen_attn = None  # type: ignore[assignment]
+    _TORCH_VARLEN_AVAILABLE = False
 try:  # pragma: no cover - import guard, exercised by environment
     from flash_attn import flash_attn_varlen_qkvpacked_func
-    from flash_attn.bert_padding import pad_input, unpad_input
 
     _FLASH_AVAILABLE = True
 except Exception:  # pragma: no cover
     flash_attn_varlen_qkvpacked_func = None  # type: ignore[assignment]
-    pad_input = None  # type: ignore[assignment]
-    unpad_input = None  # type: ignore[assignment]
     _FLASH_AVAILABLE = False
 
-from .configuration_hf_nomic_bert import NomicBertConfig
+ATTENTION_IMPLS = ("torch_varlen", "flash_attn", "eager")
+ATTN_IMPL_ENV = "NOMIC_BERT_ATTN_IMPL"  # auto (default) | torch_varlen | flash_attn | eager
+
+
+def resolve_attention_impl(
+    *,
+    override: str,
+    torch_varlen_available: bool,
+    flash_available: bool,
+    is_cuda: bool,
+    capability: Optional[Tuple[int, int]],
+) -> str:
+    """Pick an attention tier from an override, availability flags, and device info.
+
+    ``varlen_ok`` gates both ``torch_varlen`` and ``flash_attn``: both run the FA2
+    kernel family, which refuses pre-Ampere hardware. ``override == "auto"`` prefers
+    ``torch_varlen``, then ``flash_attn``, then ``eager``. A named override returns
+    that tier if its precondition holds, else raises ``RuntimeError`` naming the
+    tier and the failed precondition — a forced tier never falls back silently.
+    Any other override raises ``ValueError`` listing ``ATTENTION_IMPLS``.
+    """
+    if override != "auto" and override not in ATTENTION_IMPLS:
+        raise ValueError(f"{ATTN_IMPL_ENV}={override!r}: expected 'auto' or one of {ATTENTION_IMPLS}")
+    varlen_ok = is_cuda and capability is not None and tuple(capability) >= (8, 0)
+    if override == "auto":
+        if torch_varlen_available and varlen_ok:
+            return "torch_varlen"
+        return "flash_attn" if flash_available and varlen_ok else "eager"
+    if override == "eager":
+        return "eager"
+    if override == "torch_varlen" and not torch_varlen_available:
+        raise RuntimeError(
+            f"{ATTN_IMPL_ENV}=torch_varlen requested but torch.nn.attention.varlen "
+            f"is unavailable (torch {torch.__version__}; needs >= 2.10)"
+        )
+    if override == "flash_attn" and not flash_available:
+        raise RuntimeError(f"{ATTN_IMPL_ENV}=flash_attn requested but flash_attn is not importable")
+    if not is_cuda:
+        raise RuntimeError(f"{ATTN_IMPL_ENV}={override} requested but device is cpu (not CUDA)")
+    if capability is None or tuple(capability) < (8, 0):
+        raise RuntimeError(
+            f"{ATTN_IMPL_ENV}={override} requested but compute capability "
+            f"{tuple(capability) if capability else None} < (8, 0)"
+        )
+    return override
+
+
+@functools.lru_cache(maxsize=None)
+def _select_cached(dev_type, dev_index, override):
+    is_cuda = dev_type == "cuda"
+    cap = torch.cuda.get_device_capability(torch.device(dev_type, dev_index)) if is_cuda else None
+    return resolve_attention_impl(
+        override=override,
+        torch_varlen_available=_TORCH_VARLEN_AVAILABLE,
+        flash_available=_FLASH_AVAILABLE,
+        is_cuda=is_cuda,
+        capability=cap,
+    )
+
+
+def _unpad(x: torch.Tensor, mask: torch.Tensor):
+    """Flatten ``[B,S,...]`` to the packed ``[N,...]`` of true positions in ``mask``.
+
+    Torch-native replacement for flash-attn's ``bert_padding.unpad_input``. Returns
+    ``(x_u, indices, cu_seqlens, max_seqlen)``: ``indices`` (int64) locate each
+    packed row in the flattened ``B*S``; ``cu_seqlens`` (int32, length ``B+1``) is
+    the exclusive prefix sum of per-row valid lengths, as the varlen kernels need.
+    """
+    seqlens = mask.sum(-1, dtype=torch.int32)
+    indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten()
+    max_s = int(seqlens.max())
+    cu = F.pad(torch.cumsum(seqlens, 0, dtype=torch.int32), (1, 0))
+    return x.flatten(0, 1)[indices], indices, cu, max_s
+
+
+def _pad(x_u: torch.Tensor, indices: torch.Tensor, B: int, S: int) -> torch.Tensor:
+    """Scatter the packed ``[N,...]`` back to ``[B,S,...]`` (inverse of ``_unpad``).
+
+    Torch-native replacement for flash-attn's ``bert_padding.pad_input``. Padding
+    positions are zero-filled.
+    """
+    out = torch.zeros(B * S, *x_u.shape[1:], device=x_u.device, dtype=x_u.dtype)
+    out[indices] = x_u
+    return out.view(B, S, *x_u.shape[1:])
+
+
+def select_attention_impl(device: torch.device) -> str:
+    """Bind ``resolve_attention_impl`` to a device, the env override and a memo.
+
+    Reads ``NOMIC_BERT_ATTN_IMPL`` on every call (so a forced-unavailable tier
+    keeps raising rather than sticking to a memoized exception — ``lru_cache``
+    never caches a raised call) and memoizes the resolved tier per
+    ``(device.type, device.index, override)``.
+    """
+    override = os.environ.get(ATTN_IMPL_ENV) or "auto"
+    return _select_cached(device.type, device.index, override)
+
 
 logger = logging.getLogger(__name__)
 
 
+def _inline_extended_mask(attention_mask: Optional[torch.Tensor], dtype: torch.dtype) -> Optional[torch.Tensor]:
+    """Build the additive ``[B, 1, 1, S]`` mask the eager tier consumes.
+
+    Replaces the ``PreTrainedModel`` helper, which newer transformers removed,
+    with the same arithmetic inline, so this file has no dependency on that
+    helper. ``None`` stays ``None``.
+    """
+    if attention_mask is None:
+        return None
+    return (1.0 - attention_mask[:, None, None, :].to(dtype)) * torch.finfo(dtype).min
+
+
+_LOGGED_IMPLS: set = set()
+
+
+def _log_attention_impl_once(impl: str, device: torch.device) -> None:
+    """Emit one INFO line per ``(impl, device)`` naming why this tier engaged.
+
+    INFO, not WARNING, on every tier: eager on CPU is documented behaviour, not a
+    degradation. The line names the reason (capability, torch version,
+    flash_attn presence, the active override) so an operator reading a worker
+    log can tell why a faster tier was or wasn't selected.
+    """
+    key = (impl, str(device))
+    if key in _LOGGED_IMPLS:
+        return
+    _LOGGED_IMPLS.add(key)
+    is_cuda = device.type == "cuda"
+    cap = torch.cuda.get_device_capability(device) if is_cuda else None
+    if _FLASH_AVAILABLE:
+        try:
+            fa = importlib.metadata.version("flash-attn")
+        except importlib.metadata.PackageNotFoundError:
+            fa = "unknown"
+    else:
+        fa = "absent"
+    logger.info(
+        "NomicBert attention impl=%s device=%s capability=%s torch=%s flash_attn=%s override=%s",
+        impl,
+        device,
+        cap,
+        torch.__version__,
+        fa,
+        os.environ.get(ATTN_IMPL_ENV) or "auto",
+    )
+
+
 # adapted from flash attention, added safe serialization option for hf models
 def state_dict_from_pretrained(model_name, safe_serialization=False, device=None, dtype=None):
     # If not fp32, then we don't want to load directly to the GPU
@@ -813,13 +967,14 @@ class NomicBertAttention(nn.Module):
             if self.rotary_head_dim:
                 qkv = rearrange(qkv, "b h three s d -> b s three h d")
 
-        if _FLASH_AVAILABLE and hidden_states.is_cuda:
-            # --- Flash-attention varlen path ---
-            # flash_attn derives per-sequence lengths from ``cu_seqlens`` (built by
-            # unpad_input), so it consumes the raw bool [B, S] mask — NOT the
-            # additive [B, 1, 1, S] mask that get_extended_attention_mask makes.
-            # NomicBertModel.forward skips get_extended_attention_mask when flash is
-            # active and passes the bool mask straight through.
+        impl = select_attention_impl(hidden_states.device)
+        if impl != "eager":
+            # --- torch_varlen / flash_attn varlen paths ---
+            # Both consume the raw bool [B, S] mask (per-sequence lengths derived
+            # from cu_seqlens) — NOT the additive [B, 1, 1, S] mask that
+            # NomicBertModel.forward builds inline (_inline_extended_mask) for the
+            # eager tier. NomicBertModel.forward calls the same select_attention_impl
+            # and, on a varlen tier, passes this raw bool mask straight through.
             #
             # Correctness keystone: RoPE was already applied above, on the dense
             # [B, S, 3, H, D] tensor, BEFORE unpadding — so per-sequence positions
@@ -827,17 +982,15 @@ class NomicBertAttention(nn.Module):
             # absolute positions of sequence 1's tail and silently drift embeddings.
             B, S = qkv.shape[0], qkv.shape[1]
 
-            # flash_attn requires fp16/bf16 inputs. sentence-transformers loads the
-            # weights as fp32 by default (it silently drops config.torch_dtype), so
-            # if the model is fp32 we cast qkv to bf16 for the flash call and cast
-            # the result back to the model dtype before out_proj. This makes the
-            # model work regardless of how it was loaded — no external cast needed.
+            # Both varlen kernels require fp16/bf16 inputs. sentence-transformers
+            # loads the weights as fp32 by default (it silently drops
+            # config.torch_dtype), so if the model is fp32 we cast qkv to bf16 for
+            # the kernel call and cast the result back to the model dtype before
+            # out_proj. This makes the model work regardless of how it was loaded —
+            # no external cast needed.
             orig_dtype = qkv.dtype
-            flash_dtype = (
-                orig_dtype if orig_dtype in (torch.float16, torch.bfloat16) else torch.bfloat16
-            )
-            if flash_dtype != orig_dtype:
-                qkv = qkv.to(flash_dtype)
+            if orig_dtype not in (torch.float16, torch.bfloat16):
+                qkv = qkv.to(torch.bfloat16)
 
             if attention_mask is None:
                 bool_mask = torch.ones(B, S, device=qkv.device, dtype=torch.bool)
@@ -847,7 +1000,7 @@ class NomicBertAttention(nn.Module):
             else:
                 bool_mask = attention_mask.to(torch.bool)
 
-            qkv_u, indices, cu_s, max_s, *_ = unpad_input(qkv, bool_mask)
+            qkv_u, indices, cu_s, max_s = _unpad(qkv, bool_mask)
             if max_s == 0:
                 raise RuntimeError(
                     "NomicBertAttention.forward: every sequence is empty (max_s == 0); "
@@ -855,15 +1008,19 @@ class NomicBertAttention(nn.Module):
                 )
             # causal=False: bidirectional encoder. Default softmax_scale is
             # 1/sqrt(head_dim), numerically identical to / self.norm_factor above.
-            out_u = flash_attn_varlen_qkvpacked_func(
-                qkv_u, cu_s, max_s, dropout_p=0.0, causal=False
-            )
-            attn_output = pad_input(out_u, indices, B, S)  # [B, S, H, D]
+            if impl == "torch_varlen":
+                q, k, v = qkv_u.unbind(1)
+                out_u = _torch_varlen_attn(q, k, v, cu_s, cu_s, max_s, max_s)
+            else:
+                out_u = flash_attn_varlen_qkvpacked_func(
+                    qkv_u, cu_s, max_s, dropout_p=0.0, causal=False
+                )
+            attn_output = _pad(out_u, indices, B, S)  # [B, S, H, D]
             attn_output = attn_output.to(orig_dtype)  # restore model dtype for out_proj
             attn_output = rearrange(attn_output, "... h d -> ... (h d)")
             return self.out_proj(attn_output)
 
-        # --- Eager path (fallback: CPU, or flash_attn unavailable) ---
+        # --- Eager path (no varlen tier engaged: CPU, pre-Ampere, forced, or no kernel) ---
         # Byte-for-byte identical to nomic-ai/CodeRankEmbed upstream.
         query, key, value = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]
 
@@ -1110,6 +1267,7 @@ class NomicBertPreTrainingHeads(nn.Module):
 class NomicBertModel(NomicBertPreTrainedModel):
     def __init__(self, config: GPT2Config, add_pooling_layer=True):
         super().__init__(config)
+        self.attention_impl = None
         self.pad_vocab_size_multiple = getattr(config, "pad_vocab_size_multiple", 1)
         if config.vocab_size % self.pad_vocab_size_multiple != 0:
             config.vocab_size += self.pad_vocab_size_multiple - (config.vocab_size % self.pad_vocab_size_multiple)
@@ -1146,20 +1304,18 @@ class NomicBertModel(NomicBertPreTrainedModel):
         hidden_states = self.emb_ln(hidden_states)
         hidden_states = self.emb_drop(hidden_states)
 
-        if _FLASH_AVAILABLE and hidden_states.is_cuda:
-            # Flash-varlen path: pass the raw bool [B, S] mask straight through.
-            # Do NOT call get_extended_attention_mask here — it produces an
-            # additive [B, 1, 1, S] mask that flash_attn (which derives lengths
-            # from cu_seqlens via unpad_input) cannot consume.
-            sequence_output = self.encoder(
-                hidden_states, attention_mask=attention_mask, return_dict=return_dict
-            )
-        else:
-            # Eager path: original additive [B, 1, 1, S] extended mask.
-            attention_mask = self.get_extended_attention_mask(attention_mask, input_ids.shape)
-            sequence_output = self.encoder(
-                hidden_states, attention_mask=attention_mask, return_dict=return_dict
-            )
+        impl = select_attention_impl(hidden_states.device)
+        self.attention_impl = impl
+        _log_attention_impl_once(impl, hidden_states.device)
+        if impl == "eager":
+            # No varlen tier engaged (CPU, pre-Ampere, forced, or no kernel):
+            # build the additive [B, 1, 1, S] mask inline for the eager block.
+            attention_mask = _inline_extended_mask(attention_mask, hidden_states.dtype)
+        # else: torch_varlen / flash_attn consume the raw bool [B, S] mask as-is —
+        # NomicBertAttention.forward derives lengths from it via cu_seqlens.
+        sequence_output = self.encoder(
+            hidden_states, attention_mask=attention_mask, return_dict=return_dict
+        )
 
         pooled_output = self.pooler(sequence_output) if self.pooler is not None else None