Card: every comparison against nomic-ai/CodeRankEmbed; plain wording; citation fixed

README.md

@@ -16,23 +16,23 @@ language:
 # CodeRankEmbed-flash-attn
 
 A **bf16 quantization of [`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed)**
-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.
+with a custom `modeling_hf_nomic_bert.py`, shipped in this repo, that **picks one of three
+attention implementations for the hardware it runs on.** It is not a finetune — the weights are the
+original CodeRankEmbed weights cast to bf16 (no further training). Two of the three replace the
+original's padded `O(seq²)` attention with an `O(N)` unpadded path; the third keeps the original's
+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. 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):
+only** — activation memory grows as `batch × heads × seq²`, which runs out of memory at large
+batches even though 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 the
+eager path cannot fit 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
+- `flash_attn` — the `flash_attn` package's varlen 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**.
 
@@ -41,7 +41,7 @@ file ships all three paths itself, so no runtime patching or post-load hooks are
 
 ## Behavior
 
-- **Three-tier attention, chosen automatically per device** (override with
+- **Three attention tiers, 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.
@@ -49,7 +49,7 @@ file ships all three paths itself, so no runtime patching or post-load hooks are
      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.
+     original's padded attention algorithm, unchanged, runs on any host.
 
   `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
@@ -57,9 +57,10 @@ file ships all three paths itself, so no runtime patching or post-load hooks are
   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`.
+  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. The
   modeling file also logs the tier, and why, once at INFO.
 - **Padding-free on the varlen tiers.** The model unpads the batch once, right after the input
@@ -71,11 +72,11 @@ file ships all three paths itself, so no runtime patching or post-load hooks are
   batch flattened (`[1, ΣL]` token ids with sequence boundaries, no padding at all), and the model
   returns one `[1, ΣL, 768]` output that sentence-transformers pools per sequence. The opt-in needs
   the `flash_attn` package installed, because transformers checks for it at load and raises
-  otherwise, even though the attention itself still runs on whichever tier the dispatch picks. On
+  otherwise, even though the attention itself still runs on whichever tier the model picks. On
   the `eager` tier the flattened batch is re-padded and runs the padded path.
 - **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
+  declares `torch_dtype: bfloat16`. The original's 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).
@@ -122,106 +123,92 @@ model[0].auto_model.encoder = torch.compile(model[0].auto_model.encoder, dynamic
 
 ## Parity & performance
 
-The weights are the original CodeRankEmbed weights (bf16-cast), so embeddings match the fp32
-original to within bf16 precision. Every number below is measured on two cards, an RTX 3090 Ti
-(sm_86) and an RTX 5090 Laptop GPU (sm_120), against the **pre-dispatch revision of this repo
-(`e361c6f`) on the same GPU**, which runs the original `flash_attn` path on padded batches.
+Every number below compares this repo with the original model,
+[`nomic-ai/CodeRankEmbed`](https://huggingface.co/nomic-ai/CodeRankEmbed) (revision `3c4b608`),
+run as published: fp32 weights, its own modeling file, padded eager attention. Both run on the
+same GPU, an RTX 3090 Ti (sm_86). The original's modeling file calls
+`get_extended_attention_mask`, which transformers has deprecated and no longer ships by 5.17, so
+it was measured on transformers 5.11; this repo does not call it and loads on 5.17 too.
 
 **Corpus.** The first 16,384 `document` fields (real Python functions) of
 [`lightonai/cornstack`](https://huggingface.co/datasets/lightonai/cornstack)'s Python split, in
 shard order (`train-00000-of-00682.parquet`, then `train-00001-…`): 2,684,878 tokens, 7 to 4,746
 per function, 164 mean, 92 median. Each row is one `encode()` call over all of them at the stated
 `batch_size=`, after one untimed warm-up of the same call. Cosine similarity is per function, on
-fp32-renormalized output, against one batch-32 encode by the pre-dispatch revision. Peak memory is
-the encode's peak CUDA allocation above the loaded weights (~266 MiB), the same on both cards to
-within 32 MiB. `torch 2.12.1+cu130`, `transformers 5.11.0`, `sentence-transformers 6.1.0`,
-`flash-attn 2.8.3` (`2.8.3.post1` on the 3090 Ti).
-
-### Agreement with the pre-dispatch revision
-
-| tier | load | 3090 Ti min / mean | 5090 Laptop min / mean |
-| --- | --- | --- | --- |
-| `eager` | default | 0.99940 / 0.99992 | 0.99972 / 0.99992 |
-| `eager` | flattened | 0.99940 / 0.99992 | 0.99816 / 0.99992 |
-| `torch_varlen` | default | 0.99940 / 0.99995 | 0.99975 / 0.99995 |
-| `torch_varlen` | flattened | 0.99940 / 0.99995 | 0.99975 / 0.99995 |
-| `flash_attn` | default | 0.99978 / 0.99998 | 0.99981 / 0.99999 |
-| `flash_attn` | flattened | 0.99978 / 0.99998 | 0.99981 / 0.99999 |
-| `torch_varlen`, compiled | default | 0.99816 / 0.99994 | 0.99866 / 0.99994 |
-| `torch_varlen`, compiled | flattened | 0.99816 / 0.99994 | 0.99925 / 0.99994 |
-| `flash_attn`, compiled | default | 0.99818 / 0.99994 | 0.99876 / 0.99994 |
-| `flash_attn`, compiled | flattened | 0.99818 / 0.99994 | 0.99876 / 0.99994 |
-
-`eager` at batch 4, the others at batch 256. `torch_varlen` is a different, torch-native kernel and
-`eager` a different algorithm; `flash_attn` is the same kernel, and where it sits below 1 the rest
-of the network is running on packed tokens, so its matrix shapes and bf16 rounding differ.
-Compiling reassociates bf16 arithmetic, which leaves the mean unchanged and moves a few functions
-further. `auto`, the default with no override, kept `torch_varlen` on both cards. The flattened
-`eager` minimum on the 5090 Laptop is one 2,166-token function padded next to a 2,195-token one:
-when it flattens, sentence-transformers interleaves long and short texts, which changes who shares
-a padded `eager` batch (the same function alone scores 0.99994).
+fp32-renormalized output, against the original's embedding of the same function. Peak memory is
+the encode's peak CUDA allocation above the loaded weights. **Load** is how the model was loaded:
+`default` (as in Usage) or `flattened` (the opt-in above). `torch 2.12.1+cu130`,
+`transformers 5.11.0`, `sentence-transformers 6.1.0`, `flash-attn 2.8.3.post1`.
+
+### Agreement with the original model
+
+| model | attention | compiled | load | batch size | min cosine | mean cosine |
+| --- | --- | --- | --- | --- | --- | --- |
+| `nomic-ai/CodeRankEmbed` | `eager` | – | default | 4 | 1 (the reference) | 1 (the reference) |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `eager` | – | default | 4 | 0.99902 | 0.99991 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `eager` | – | flattened | 4 | 0.99916 | 0.99991 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 256 | 0.99896 | 0.99993 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 256 | 0.99896 | 0.99993 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | – | default | 256 | 0.99841 | 0.99993 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | – | flattened | 256 | 0.99841 | 0.99993 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 256 | 0.99960 | 0.99995 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 256 | 0.99960 | 0.99995 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | ✓ | default | 256 | 0.99946 | 0.99995 |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `flash_attn` | ✓ | flattened | 256 | 0.99946 | 0.99995 |
+
+The original is the reference every other row is measured against, so its cosine is 1 by
+definition, not a measurement. The mean cosine is 0.99991 or higher in every configuration; the
+minimum is the one function, of 16,384, that moves most under bf16 weights and arithmetic. The
+batch size does not change these figures (a varlen tier gives the same five decimals at batch 32
+and 256), so the largest batch that ran in every load is shown; `eager` runs at batch 4 (see
+below). Compiled rows sit closer to the original than uncompiled ones.
 
 ### Speed and memory
 
-| load | batch size | peak memory, eager | peak memory, compiled | 3090 Ti eager | 5090 Laptop eager | 3090 Ti compiled | 5090 Laptop compiled |
-| --- | --- | --- | --- | --- | --- | --- | --- |
-| pre-dispatch revision | 32 | 4.4 GiB | – | 23.8 s | 21.4 s | – | – |
-| default | 32 | 2.5 GiB | 1.6 GiB | 16.4 s | 14.3 s | 14.0 s | 12.2 s |
-| default | 256 | 10.9 GiB | 7.1 GiB | 15.1 s | 15.1 s | 12.7 s | 12.3 s |
-| default | 1,024 | out of memory | 16.3 GiB | out of memory | out of memory | 13.8 s | 13.4 s |
-| default | 2,048 | out of memory | out of memory | out of memory | out of memory | out of memory | out of memory |
-| flattened | 32 | 1.4 GiB | 0.9 GiB | 15.6 s | 13.7 s | 13.2 s | 11.7 s |
-| flattened | 256 | 6.9 GiB | 4.5 GiB | 13.7 s | 14.4 s | 11.4 s | 11.7 s |
-| flattened | 1,024 | 17.0 GiB | 11.1 GiB | 13.4 s | 14.7 s | 10.9 s | 11.8 s |
-| flattened | 2,048 | out of memory | 16.6 GiB | out of memory | out of memory | 10.7 s | 11.4 s |
-
-`torch_varlen`; `flash_attn` runs within 1% of it. Compiled means
+| model | attention | compiled | load | batch size | peak memory | time |
+| --- | --- | --- | --- | --- | --- | --- |
+| `nomic-ai/CodeRankEmbed` | `eager` | – | default | 4 | 8.4 GiB | 68.9 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 32 | 2.5 GiB | 16.3 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 256 | 10.9 GiB | 15.1 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 1,024 | out of memory | out of memory |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | default | 2,048 | out of memory | out of memory |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 32 | 1.4 GiB | 15.5 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 256 | 6.9 GiB | 13.6 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 1,024 | 17.0 GiB | 13.4 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | – | flattened | 2,048 | out of memory | out of memory |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 32 | 1.6 GiB | 14.0 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 256 | 7.1 GiB | 12.8 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 1,024 | 16.3 GiB | 13.8 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | default | 2,048 | out of memory | out of memory |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 32 | 0.9 GiB | 13.2 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 256 | 4.5 GiB | 11.4 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 1,024 | 11.1 GiB | 10.9 s |
+| `handwoven8588/CodeRankEmbed-flash-attn` | `torch_varlen` | ✓ | flattened | 2,048 | 16.6 GiB | 10.8 s |
+
+This repo's rows are `torch_varlen`; `flash_attn` runs within 1% of it in both time and memory
+at every cell. The original runs at batch 4: its padded fp32 attention needs
+`batch × heads × seq²` memory, 8.4 GiB at batch 4 for this corpus's longest function. The same
+algorithm in bf16 (this repo's `eager` tier, batch 4) takes 41.2 s in 4.2 GiB; the rest of the
+gain comes from the unpadded path and the batches it allows. Compiled means
 `model[0].auto_model.encoder = torch.compile(model[0].auto_model.encoder, dynamic=True)`.
 
-- **Memory follows the real tokens in the heaviest batch, not the number of functions.** Both
-  loads run the encoder padding-free, and every out-of-memory cell failed allocating one 3,072-wide
-  MLP activation over that batch's tokens. sentence-transformers length-sorts the corpus, so the
-  default load's first batch of 1,024 holds the longest functions, 869,948 tokens. Flattened, it
-  interleaves longest with shortest, so its heaviest batch of 1,024 holds 593,716 tokens.
-- **Flattening** cuts peak memory by 36–42% against the default load and wall time by up to 9%.
-- **Compiling** is 1.17–1.24× faster at the same batch size and cuts peak memory by about a third,
-  because inductor fuses the element-wise work around each attention call. It compiles one graph,
-  with no graph breaks, once per process (the first encode takes 4–10 s longer), and does not
-  recompile as batch lengths change; the one further compile comes once a batch passes about
-  699,000 tokens, where inductor switches to 64-bit indexing.
-- **Together**, flattened and compiled at batch 2,048 encode the corpus in 10.7 s (3090 Ti) and
-  11.4 s (5090 Laptop), 2.2× and 1.9× the pre-dispatch revision. `eager` at batch 4 takes 34–52 s.
-
-### Real documents
-
-768 Python files sampled with a fixed seed from one shard of
-[`HuggingFaceTB/stack-edu`](https://huggingface.co/datasets/HuggingFaceTB/stack-edu) (730,431
-tokens after truncation to the model's 8192 `max_seq_length`), on `torch_varlen`, uncompiled.
-**Sorted** is one `encode()` call over all 768, so sentence-transformers length-sorts them;
-**file** is one call per batch of consecutive files, so each batch mixes short and long files.
-Peak memory here includes the weights; **largest** is the largest batch that fit on either card.
-
-| load | order | batch size | peak memory | 3090 Ti | 5090 Laptop |
-| --- | --- | --- | --- | --- | --- |
-| default | sorted | 4 | 1.3 GiB | 5.61 s | 4.75 s |
-| default | sorted | 32 | 6.1 GiB | 4.79 s | 4.72 s |
-| default | sorted | 64 | 9.1 GiB | 4.80 s | 4.84 s |
-| default | sorted | 353 (largest) | 18.2 GiB | 5.77 s | 5.98 s |
-| default | file | 4 | 0.6 GiB | 6.49 s | 5.64 s |
-| default | file | 32 | 1.6 GiB | 6.39 s | 5.70 s |
-| default | file | 64 | 2.8 GiB | 6.44 s | 5.84 s |
-| flattened | sorted | 4 | 0.8 GiB | 5.51 s | 4.56 s |
-| flattened | sorted | 32 | 3.8 GiB | 4.54 s | 4.56 s |
-| flattened | sorted | 64 | 6.1 GiB | 4.47 s | 4.63 s |
-| flattened | sorted | 628 (largest) | 19.2 GiB | 4.60 s | 4.75 s |
-| flattened | file | 4 | 0.6 GiB | 5.83 s | 4.75 s |
-| flattened | file | 32 | 1.6 GiB | 4.73 s | 4.70 s |
-| flattened | file | 64 | 2.7 GiB | 4.53 s | 4.66 s |
-
-Flattening also makes a file's embedding independent of the other files in its batch: between
-the file-order and sorted runs, the minimum cosine over the 768 files is 0.9999998 under the
-flattened load on both cards and both varlen tiers, against as low as 0.99987 under the default
-load.
+- **Against the original**, the same 16,384 functions encode 5.1× faster uncompiled (13.4 s,
+  flattened at batch 1,024) and 6.4× faster compiled (10.8 s, flattened at batch 2,048), against
+  68.9 s. At batch 32 this repo needs 2.5 GiB (default load) or 1.4 GiB (flattened), against the
+  original's 8.4 GiB at batch 4.
+- **Memory follows the real tokens in the heaviest batch, not the number of functions.** This
+  repo runs the encoder on real tokens only, and every out-of-memory cell failed allocating one
+  activation over that batch's tokens: the 3,072-wide MLP hidden state, or, compiled, the
+  2,304-wide QKV projection. sentence-transformers length-sorts the corpus, so the default load's
+  first batch of 1,024 holds the longest functions, about 870,000 tokens; flattened, it pairs the
+  longest with the shortest, so its heaviest batch of 1,024 holds about 594,000.
+- **Flattening** cuts peak memory by 32–42% against the default load at the same batch size, and
+  wall time by 5–21% (more at larger batches).
+- **Compiling** is 1.17–1.23× faster at the same batch size and cuts peak memory by 35%. It
+  compiles one graph, with no graph breaks, once per process (the first encode takes about 5 s
+  longer), and does not recompile as batch lengths change; the only further compiles come once a
+  batch passes about 699,000 and 932,000 tokens, where inductor switches to 64-bit indexing for
+  the 3,072-wide MLP activation and the 2,304-wide QKV projection.
 
 ## What changed vs the source repo
 
@@ -229,29 +216,30 @@ load.
    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
+2. **`from_pretrained` dtype fix**: the original's 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. **Three-tier attention dispatch**: `NomicBertAttention.forward` now selects one of three
-   attention implementations at call time — `torch_varlen` (torch's own
+3. **Three attention tiers**: `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 algorithm, numerically
-   unchanged, and the default off CUDA sm_80+; it now converts a raw 2-D `[B, S]` mask to the
-   additive form before adding it, for the case where a `device_map` split hands it a mask a varlen
-   tier upstream never passes). On both varlen tiers `NomicBertModel.forward` unpads the batch once
-   with torch-native `_unpad`/`_pad` helpers (a replacement for `flash_attn.bert_padding`), keeps
-   the hidden states packed `[ΣL, 768]` through the embeddings and every layer, and pads back once at
-   the end. Rotary embeddings rotate each packed token by its position within its own sequence,
-   the same rotation it gets in the padded batch. On the eager tier `NomicBertModel.forward` builds
-   the additive attention mask inline instead of calling the (now-removed-upstream)
-   `get_extended_attention_mask` helper. 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
+   `flash_attn` (the `flash_attn` package's varlen kernel, kept as a fallback for older torch
+   builds that have it installed), and `eager` (the original's padded attention algorithm,
+   numerically unchanged, and the default off CUDA sm_80+; it now converts a raw 2-D `[B, S]` mask
+   to the additive form before adding it, for the case where a `device_map` split hands it a mask
+   a varlen tier upstream never passes). On both varlen tiers `NomicBertModel.forward` unpads the
+   batch once with torch-native `_unpad`/`_pad` helpers (a replacement for
+   `flash_attn.bert_padding`), keeps the hidden states packed `[ΣL, 768]` through the embeddings
+   and every layer, and pads back once at the end. Rotary embeddings rotate each packed token by
+   its position within its own sequence, the same rotation it gets in the padded batch. On the
+   eager tier `NomicBertModel.forward` builds the additive attention mask inline instead of calling
+   the (now-removed-upstream) `get_extended_attention_mask` helper. 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 original's 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=`, and returns the model in eval mode, as
    transformers' own `from_pretrained` does.
@@ -269,8 +257,12 @@ derives from Tri Dao's BERT implementation, and `CodeRankEmbed` was trained by t
 
 ```bibtex
 @misc{suresh2025cornstackhighqualitycontrastivedata,
-  title  = {CoRNStack: High-Quality Contrastive Data for Text and Code Retrieval},
-  author = {Suresh, K N Q and Wang, Xiang and Khan, Saqib and others},
-  year   = {2025},
+      title={CoRNStack: High-Quality Contrastive Data for Better Code Retrieval and Reranking},
+      author={Tarun Suresh and Revanth Gangi Reddy and Yifei Xu and Zach Nussbaum and Andriy Mulyar and Brandon Duderstadt and Heng Ji},
+      year={2025},
+      eprint={2412.01007},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2412.01007},
 }
 ```