Pin weights to the requested revision; harden the varlen probe; card: all tiers measured on one GPU with stated batch sizes

README.md

@@ -54,13 +54,19 @@ file ships all three paths itself, so no runtime patching or post-load hooks are
   `auto` (the default) prefers `torch_varlen`, then `flash_attn`, then `eager`. Before accepting a
   varlen tier, `auto` runs one tiny kernel probe per device (a capability check alone can't tell
   whether the installed build has a kernel for the GPU's architecture, e.g. on ROCm); if the probe
-  raises, it logs one WARNING naming the tier and the error and steps down to the next tier. A
-  forced override (e.g. `NOMIC_BERT_ATTN_IMPL=torch_varlen`) is not probed and **raises**
+  raises a `RuntimeError` (how a missing or unsupported kernel fails), it logs one WARNING naming
+  the tier and the error and steps down to the next tier. An out-of-memory error, any other
+  exception type, or a CUDA error left over from earlier work propagates instead of demoting the
+  tier, so a transient failure never pins the device to a slower tier. A forced override (e.g. `NOMIC_BERT_ATTN_IMPL=torch_varlen`) is not probed and **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)`).
+  override=...` INFO log line (one line per distinct `(impl, device)`). When `flash_attn` is
+  unavailable the line names why (`flash_attn=absent(<error>)`); an installed `flash_attn` that
+  fails to import also logs a WARNING.
+- **`revision=` pins everything.** Loading with `revision=<commit>` fetches the code, config,
+  tokenizer **and weights** from that commit.
 - **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
@@ -90,34 +96,44 @@ 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. 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 |
+original to within bf16 precision. The table below measures every tier **on the same GPU**, so
+`eager` is the O(seq²) baseline the varlen tiers are compared against.
+
+**Protocol.** 64 real Python source snippets (the first 40 lines of each file; 428 tokens mean,
+638 max), encoded with sentence-transformers at an explicit **batch size of 32 and of 64** (the
+`batch_size=` passed to `encode`; with 64 snippets, 64 is one forward). Each tier encodes once
+untimed, then again for the measurement. Cosine similarity is taken on fp32-renormalized output
+against the **pre-dispatch revision of this repo on the same GPU** (its `flash_attn` path).
+`flash_attn` runs the same computation as before the dispatch, so its cosine is a parity check.
+`torch_varlen` is a different, torch-native kernel, and `eager` a different algorithm, so theirs
+are the real signal. `auto` is the default with no override: the probe keeps `torch_varlen`.
+**Encode peak** is the encode's peak CUDA allocation above the loaded model's own ~266 MiB of
+bf16 weights. Wall time is one encode of all 64 snippets.
+
+| GPU | batch size | tier | mean cos | min cos | encode peak | wall |
 | --- | --- | --- | --- | --- | --- | --- |
-| 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.
+| RTX 3090 Ti (sm_86) | 32 | `torch_varlen` | 0.999941 | 0.999856 | 600 MiB | 0.232 s |
+| RTX 3090 Ti (sm_86) | 32 | `flash_attn` | 1.000000 | 1.000000 | 600 MiB | 0.232 s |
+| RTX 3090 Ti (sm_86) | 32 | `eager` | 0.999912 | 0.999856 | 807 MiB | 0.284 s |
+| RTX 3090 Ti (sm_86) | 64 | `torch_varlen` | 0.999941 | 0.999856 | 1201 MiB | 0.248 s |
+| RTX 3090 Ti (sm_86) | 64 | `flash_attn` | 1.000000 | 1.000000 | 1201 MiB | 0.250 s |
+| RTX 3090 Ti (sm_86) | 64 | `eager` | 0.999910 | 0.999856 | 1614 MiB | 0.316 s |
+| RTX 5090 Laptop GPU (sm_120) | 32 | `torch_varlen` | 0.999944 | 0.999919 | 600 MiB | 0.234 s |
+| RTX 5090 Laptop GPU (sm_120) | 32 | `flash_attn` | 1.000000 | 1.000000 | 600 MiB | 0.236 s |
+| RTX 5090 Laptop GPU (sm_120) | 32 | `eager` | 0.999912 | 0.999870 | 807 MiB | 0.318 s |
+| RTX 5090 Laptop GPU (sm_120) | 64 | `torch_varlen` | 0.999944 | 0.999919 | 1201 MiB | 0.272 s |
+| RTX 5090 Laptop GPU (sm_120) | 64 | `flash_attn` | 1.000000 | 1.000000 | 1201 MiB | 0.271 s |
+| RTX 5090 Laptop GPU (sm_120) | 64 | `eager` | 0.999913 | 0.999834 | 1614 MiB | 0.378 s |
+
+Cosines rounded to 6 decimal places, VRAM to the nearest MiB. `torch 2.12.1+cu130`,
+`transformers 5.11.0`, `sentence-transformers 6.1.0`, `flash-attn 2.8.3`.
+
+Eager's cost grows with `batch × heads × seq²`: at these short snippets it needs about a third more
+encode memory than the varlen tiers. At longer inputs the gap stops being a percentage. At a
+batch size of **256** inputs of ~300–2048 tokens, `torch_varlen` stays under 20 GB peak VRAM,
+while eager's attention scores alone (`256 × 12 heads × 2048² × 2 bytes`) exceed a 24 GB card.
+Both varlen tiers are also gated downstream at min cosine > 0.997 against the fp32
+`nomic-ai/CodeRankEmbed` reference.
 
 ## What changed vs the source repo
 
@@ -144,6 +160,11 @@ at batch size 256.
    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.
+4. **`from_pretrained` forwards the Hub revision**: the upstream custom `from_pretrained` fetched
+   the weight file without the caller's `revision` (and cache/token options), so a pinned load
+   still took the weights from the default branch. The copy here passes them through, so a
+   pinned load fetches the weights from the pinned commit too. It also accepts transformers'
+   `dtype=` alongside the older `torch_dtype=`.
 
 ## License & attribution
 

modeling_hf_nomic_bert.py

@@ -33,13 +33,13 @@ 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
+#   3. eager        — the original padded algorithm, 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). In auto mode a tiny kernel probe then confirms the chosen varlen tier
 # actually runs on the device, stepping down a tier if it raises (_select_cached).
-# When neither is available (CPU-only hosts, or builds without them) tier 3 runs
-# unchanged, so this model loads and encodes everywhere.
+# When neither is available (CPU-only hosts, or builds without them) tier 3 runs,
+# 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
 
@@ -47,13 +47,25 @@ try:  # pragma: no cover - import guard, exercised by environment
 except ImportError:  # pragma: no cover
     _torch_varlen_attn = None  # type: ignore[assignment]
     _TORCH_VARLEN_AVAILABLE = False
+# Why flash_attn is unavailable, or None. A missing package and a broken install
+# (e.g. an extension built against another torch) both leave the tier unavailable,
+# but only a broken install is worth a WARNING; the INFO tier line and a forced
+# flash_attn error both name the cause.
+_FLASH_IMPORT_ERROR: Optional[str] = None
 try:  # pragma: no cover - import guard, exercised by environment
     from flash_attn import flash_attn_varlen_qkvpacked_func
 
     _FLASH_AVAILABLE = True
-except Exception:  # pragma: no cover
+except Exception as _exc:  # pragma: no cover
     flash_attn_varlen_qkvpacked_func = None  # type: ignore[assignment]
     _FLASH_AVAILABLE = False
+    _FLASH_IMPORT_ERROR = f"{type(_exc).__name__}: {_exc}"
+    if not isinstance(_exc, ModuleNotFoundError) or _exc.name != "flash_attn":
+        logging.getLogger(__name__).warning(
+            "NomicBert attention: flash_attn is installed but failed to import (%s); "
+            "the flash_attn tier is unavailable",
+            _FLASH_IMPORT_ERROR,
+        )
 
 ATTENTION_IMPLS = ("torch_varlen", "flash_attn", "eager")
 ATTN_IMPL_ENV = "NOMIC_BERT_ATTN_IMPL"  # auto (default) | torch_varlen | flash_attn | eager
@@ -91,7 +103,10 @@ def resolve_attention_impl(
             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")
+        raise RuntimeError(
+            f"{ATTN_IMPL_ENV}=flash_attn requested but flash_attn is not importable"
+            f" ({_FLASH_IMPORT_ERROR})"
+        )
     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):
@@ -114,7 +129,12 @@ def _probe_varlen_kernel(impl: str, device: torch.device) -> None:
     only a proxy: a torch build can ship ``torch.nn.attention.varlen`` without an
     FA kernel for the device's arch, and on ROCm ``device.type`` is ``"cuda"`` and
     the capability is the gfx major.
+
+    The device is synchronized BEFORE the kernel call, outside anything that would
+    read a failure as "this tier has no kernel": an asynchronous CUDA error left by
+    earlier work surfaces there as itself instead of being blamed on the probe.
     """
+    torch.cuda.synchronize(device)
     qkv = torch.randn(2, 3, 1, 64, device=device, dtype=torch.bfloat16)
     cu = torch.tensor([0, 2], device=device, dtype=torch.int32)
     if impl == "torch_varlen":
@@ -125,23 +145,39 @@ def _probe_varlen_kernel(impl: str, device: torch.device) -> None:
     torch.cuda.synchronize(device)
 
 
+@functools.lru_cache(maxsize=None)
+def _is_oom(exc: BaseException) -> bool:
+    """True for an out-of-memory failure, allocator-level or driver-level.
+
+    ``torch.cuda.OutOfMemoryError`` is the caching allocator's; a driver-level
+    allocation failure inside a kernel surfaces as a plain ``RuntimeError`` (or
+    ``torch.AcceleratorError``) whose message says ``out of memory``.
+    """
+    return isinstance(exc, torch.cuda.OutOfMemoryError) or "out of memory" in str(exc).lower()
+
+
 @functools.lru_cache(maxsize=None)
 def _select_cached(dev_type, dev_index, override):
     """Resolve the tier for one device, probing the kernel in auto mode.
 
     A forced tier is returned exactly as ``resolve_attention_impl`` decides it (or
     raises), with no probe: it never falls back. In ``auto`` mode a varlen tier is
-    accepted only if ``_probe_varlen_kernel`` runs; if the probe raises, one
-    WARNING names the rejected tier and the exception, that tier is marked
+    accepted only if ``_probe_varlen_kernel`` runs; if the probe raises a
+    ``RuntimeError`` (which covers ``NotImplementedError`` and
+    ``torch.AcceleratorError``: the ways a missing or unsupported kernel fails),
+    one WARNING names the rejected tier and the exception, that tier is marked
     unavailable, and the pure selector picks again (``torch_varlen`` ->
     ``flash_attn`` -> ``eager``). The probe is the single authority on whether a
     varlen kernel runs, on every backend: ROCm gets no separate predicate, it is
     probed like CUDA. The result, including a fallback, is memoized per
-    ``(device, override)``, so the probe runs once per device per process. A
-    ``torch.cuda.OutOfMemoryError`` from the probe is re-raised immediately
-    instead, with no WARNING and no tier marked unavailable; ``lru_cache`` never
-    memoizes a raised call, so a transient OOM never permanently pins the device
-    to a lower tier.
+    ``(device, override)``, so the probe runs once per device per process.
+
+    Three failures propagate instead, with no WARNING and no tier marked
+    unavailable, and ``lru_cache`` never memoizes a raised call, so none of them
+    pins the device to a lower tier: an out-of-memory error (``_is_oom``: transient,
+    not a missing kernel), an exception that is not a ``RuntimeError`` (a
+    ``TypeError`` from a changed kernel signature is a bug to fix, not a tier to
+    skip), and an error from the pre-probe synchronize (earlier work's fault).
     """
     is_cuda = dev_type == "cuda"
     device = torch.device(dev_type, dev_index)
@@ -160,9 +196,9 @@ def _select_cached(dev_type, dev_index, override):
         try:
             _probe_varlen_kernel(impl, device)
             return impl
-        except torch.cuda.OutOfMemoryError:
-            raise
-        except Exception as exc:
+        except RuntimeError as exc:
+            if _is_oom(exc):
+                raise
             logger.warning(
                 "NomicBert attention: auto rejected impl=%s on device=%s (probe raised %s: %s); "
                 "trying the next tier",
@@ -201,12 +237,11 @@ def _pad(x_u: torch.Tensor, indices: torch.Tensor, B: int, S: int) -> torch.Tens
 
 
 def select_attention_impl(device: torch.device) -> str:
-    """Bind ``resolve_attention_impl`` to a device, the env override and a memo.
+    """The attention tier for ``device`` under the current ``NOMIC_BERT_ATTN_IMPL``.
 
-    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)``.
+    Reads the override on every call, so changing the env var between forwards
+    takes effect, and delegates to ``_select_cached``, which resolves (and in
+    auto mode probes) once per ``(device.type, device.index, override)``.
     """
     override = os.environ.get(ATTN_IMPL_ENV) or "auto"
     return _select_cached(device.type, device.index, override)
@@ -215,9 +250,10 @@ def select_attention_impl(device: torch.device) -> str:
 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``.
+    ``attention_mask`` is the ``{0,1}`` ``[B, S]`` tokenizer mask: kept positions
+    become ``0.0`` and padding becomes ``torch.finfo(dtype).min``, the arithmetic
+    of transformers' extended-attention-mask helper, inline so this file does not
+    depend on that helper. ``None`` stays ``None``.
     """
     if attention_mask is None:
         return None
@@ -247,7 +283,7 @@ def _log_attention_impl_once(impl: str, device: torch.device) -> None:
         except importlib.metadata.PackageNotFoundError:
             fa = "unknown"
     else:
-        fa = "absent"
+        fa = f"absent({_FLASH_IMPORT_ERROR})"
     logger.info(
         "NomicBert attention impl=%s device=%s capability=%s torch=%s flash_attn=%s override=%s",
         impl,
@@ -260,7 +296,7 @@ def _log_attention_impl_once(impl: str, device: torch.device) -> None:
 
 
 # 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):
+def state_dict_from_pretrained(model_name, safe_serialization=False, device=None, dtype=None, **hub_kwargs):
     # If not fp32, then we don't want to load directly to the GPU
     mapped_device = "cpu" if dtype not in [torch.float32, None] else device
     is_sharded = False
@@ -288,10 +324,14 @@ def state_dict_from_pretrained(model_name, safe_serialization=False, device=None
         load_safe = True
     else:  # Try loading from HF hub instead of from local files
         weight_name = WEIGHTS_NAME if not safe_serialization else SAFE_WEIGHTS_NAME
-        resolved_archive_file = cached_file(model_name, weight_name, _raise_exceptions_for_missing_entries=False)
+        resolved_archive_file = cached_file(
+            model_name, weight_name, _raise_exceptions_for_missing_entries=False, **hub_kwargs
+        )
         if resolved_archive_file is None:
             weight_index = WEIGHTS_INDEX_NAME if not safe_serialization else SAFE_WEIGHTS_INDEX_NAME
-            resolved_archive_file = cached_file(model_name, weight_index, _raise_exceptions_for_missing_entries=False)
+            resolved_archive_file = cached_file(
+                model_name, weight_index, _raise_exceptions_for_missing_entries=False, **hub_kwargs
+            )
             if resolved_archive_file is not None:
                 is_sharded = True
 
@@ -308,7 +348,9 @@ def state_dict_from_pretrained(model_name, safe_serialization=False, device=None
     if is_sharded:
         # resolved_archive_file becomes a list of files that point to the different
         # checkpoint shards in this case.
-        resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(model_name, resolved_archive_file)
+        resolved_archive_file, sharded_metadata = get_checkpoint_shard_files(
+            model_name, resolved_archive_file, **hub_kwargs
+        )
         state_dict = {}
         for sharded_file in resolved_archive_file:
             state_dict.update(loader(sharded_file))
@@ -528,9 +570,26 @@ class NomicBertPreTrainedModel(PreTrainedModel):
             *inputs, **kwargs: additional input for the specific NomicBert class
                 (ex: num_labels for NomicBertForSequenceClassification)
         """
+        # The Hub-resolution kwargs transformers' Auto classes forward (revision above
+        # all) reach every file this method downloads, so a pinned revision pins the
+        # weights, not just the code and config.
+        hub_kwargs = {
+            k: kwargs[k]
+            for k in (
+                "revision",
+                "cache_dir",
+                "token",
+                "local_files_only",
+                "force_download",
+                "proxies",
+                "subfolder",
+                "_commit_hash",
+            )
+            if kwargs.get(k) is not None
+        }
         # Instantiate model.
         if config is None:
-            config = cls.config_class.from_pretrained(model_name)
+            config = cls.config_class.from_pretrained(model_name, **hub_kwargs)
         remove_cls = cls != NomicBertForPreTraining
         remove_bert_prefix = cls != NomicBertForPreTraining
         ignore_mismatched_shapes = kwargs.pop("ignore_mismatched_sizes", False)
@@ -569,7 +628,9 @@ class NomicBertPreTrainedModel(PreTrainedModel):
             load_return = model.load_state_dict(state_dict, strict=False)
         else:
             # TODO: can probably check config class and see if we need to remap from a bert model
-            state_dict = state_dict_from_pretrained(model_name, safe_serialization=kwargs.get("safe_serialization", False))
+            state_dict = state_dict_from_pretrained(
+                model_name, safe_serialization=kwargs.get("safe_serialization", False), **hub_kwargs
+            )
             state_dict = remap_bert_state_dict(
                 state_dict,
                 config,
@@ -581,21 +642,27 @@ class NomicBertPreTrainedModel(PreTrainedModel):
                 state_dict = filter_shapes(state_dict, model)
 
             load_return = model.load_state_dict(state_dict, strict=True)
-        # Honor torch_dtype like transformers' native from_pretrained does. Our custom
-        # override above bypassed it (it instantiates fp32 and load_state_dict upcasts
-        # the checkpoint into the fp32 params, ignoring torch_dtype). Resolve explicit
-        # kwarg > config.torch_dtype > checkpoint dtype, then cast the model.
+        # Honor the requested dtype like transformers' native from_pretrained does.
+        # This custom override instantiates fp32 and load_state_dict upcasts the
+        # checkpoint into the fp32 params, so the dtype is applied here. Resolve
+        # explicit kwarg (`dtype`, or the older name `torch_dtype`) > config >
+        # checkpoint dtype, then cast the model.
         import torch as _torch
-        _td = kwargs.get("torch_dtype")
+        _td = kwargs.get("dtype")
         if _td is None:
-            _td = getattr(config, "torch_dtype", None)
+            _td = kwargs.get("torch_dtype")
+        if _td is None:
+            _td = getattr(config, "dtype", None) or getattr(config, "torch_dtype", None)
         if _td == "auto" or _td is None:
             _td = next(iter(state_dict.values())).dtype
         if isinstance(_td, str):
             _td = getattr(_torch, _td)
         if _td is not None:
             model = model.to(_td)
-        logger.warning(load_return)
+        if load_return.missing_keys or load_return.unexpected_keys:
+            logger.warning(load_return)
+        else:
+            logger.debug(load_return)
         return model
 
     def _set_gradient_checkpointing(self, module, value=False):
@@ -1028,11 +1095,11 @@ class NomicBertAttention(nn.Module):
         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
+            # Both consume the tokenizer's {0,1} [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.
+            # and, on a varlen tier, passes the {0,1} 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
@@ -1040,12 +1107,10 @@ class NomicBertAttention(nn.Module):
             # absolute positions of sequence 1's tail and silently drift embeddings.
             B, S = qkv.shape[0], qkv.shape[1]
 
-            # 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.
+            # Both varlen kernels require fp16/bf16 inputs. The model loads bf16 by
+            # default, but a caller can ask for fp32 (dtype=torch.float32), so an
+            # fp32 qkv is cast to bf16 for the kernel call and the result is cast
+            # back to the model dtype before out_proj.
             orig_dtype = qkv.dtype
             if orig_dtype not in (torch.float16, torch.bfloat16):
                 qkv = qkv.to(torch.bfloat16)
@@ -1376,7 +1441,7 @@ class NomicBertModel(NomicBertPreTrainedModel):
             # 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 —
+        # else: torch_varlen / flash_attn consume the {0,1} [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