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

README.md

@@ -16,8 +16,8 @@ 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
+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 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.
@@ -32,7 +32,7 @@ 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**.
 
@@ -71,7 +71,7 @@ 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`
@@ -122,106 +122,115 @@ 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 attention. Both run on the same GPU,
+on two cards: an RTX 3090 Ti (sm_86) and an RTX 5090 Laptop GPU (sm_120). 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 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
+fp32-renormalized output, against the original model's embedding of the same function. Peak
+memory is the encode's peak CUDA allocation above the loaded weights, 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
+### Agreement with the original model
 
 | 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).
+| `eager` | default | 0.99902 / 0.99991 | 0.99904 / 0.99991 |
+| `eager` | flattened | 0.99916 / 0.99991 | 0.99904 / 0.99991 |
+| `torch_varlen` | default | 0.99896 / 0.99993 | 0.99894 / 0.99993 |
+| `torch_varlen` | flattened | 0.99896 / 0.99993 | 0.99894 / 0.99993 |
+| `flash_attn` | default | 0.99841 / 0.99993 | 0.99914 / 0.99993 |
+| `flash_attn` | flattened | 0.99841 / 0.99993 | 0.99914 / 0.99993 |
+| `torch_varlen`, compiled | default | 0.99960 / 0.99995 | 0.99936 / 0.99995 |
+| `torch_varlen`, compiled | flattened | 0.99960 / 0.99995 | 0.99936 / 0.99995 |
+| `flash_attn`, compiled | default | 0.99946 / 0.99995 | 0.99952 / 0.99995 |
+| `flash_attn`, compiled | flattened | 0.99946 / 0.99995 | 0.99952 / 0.99995 |
+
+`eager` at batch 4, the others at batch 256. The mean cosine is 0.99991 or higher everywhere; the
+minimum is the one function, of 16,384, that moves most under bf16 weights and arithmetic.
+Compiled rows sit closer to the original than uncompiled ones: inductor computes each fused run
+of element-wise work in fp32 and rounds to bf16 once, instead of after every operation. `auto`,
+the default with no override, picked `torch_varlen` on both cards.
 
 ### Speed and memory
 
-| load | batch size | peak memory, eager | peak memory, compiled | 3090 Ti eager | 5090 Laptop eager | 3090 Ti compiled | 5090 Laptop compiled |
+| model / load | batch size | peak memory | peak memory, compiled | 3090 Ti | 5090 Laptop | 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 |
+| `nomic-ai/CodeRankEmbed` (fp32) | 4 | 8.4 GiB | – | 68.9 s | 77.8 s | – | – |
+| default | 32 | 2.5 GiB | 1.6 GiB | 16.3 s | 14.6 s | 14.0 s | 12.4 s |
+| default | 256 | 10.9 GiB | 7.1 GiB | 15.1 s | 15.4 s | 12.8 s | 12.6 s |
+| default | 1,024 | out of memory | 16.3 GiB | out of memory | out of memory | 13.8 s | 13.8 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
+| flattened | 32 | 1.4 GiB | 0.9 GiB | 15.5 s | 13.8 s | 13.2 s | 12.5 s |
+| flattened | 256 | 6.9 GiB | 4.5 GiB | 13.6 s | 14.5 s | 11.4 s | 12.2 s |
+| flattened | 1,024 | 17.0 GiB | 11.1 GiB | 13.4 s | 14.9 s | 10.9 s | 12.4 s |
+| flattened | 2,048 | out of memory | 16.6 GiB | out of memory | out of memory | 10.8 s | 12.8 s |
+
+This repo's rows are `torch_varlen`; `flash_attn` runs within 1% of it on the 3090 Ti and within
+6% on the 5090 Laptop. 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. 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.
+- **Against the original**, the same 16,384 functions encode 5.1–5.6× faster uncompiled and
+  6.4× faster compiled (10.8 s on the 3090 Ti, 12.2 s on the 5090 Laptop, against 68.9 s and
+  77.8 s), in a third of the memory at batch 32.
+- **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
+  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.
+- **Flattening** cuts peak memory by 36–42% against the default load and wall time by 5–10%.
+- **Compiling** is 1.10–1.23× faster at the same batch size and cuts peak memory by about a
+  third. It compiles one graph, with no graph breaks, once per process (the first encode takes
+  about 4–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.
 
 ### 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.
+tokens after truncation to 8,192 tokens), this repo 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. The original runs at
+batch 1 and 2, the most its padded fp32 attention fits at 8,192 tokens. 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 |
+| model / 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.
+| `nomic-ai/CodeRankEmbed` (fp32) | sorted | 1 | 6.8 GiB | 26.08 s | 33.05 s |
+| `nomic-ai/CodeRankEmbed` (fp32) | sorted | 2 | 12.9 GiB | 27.13 s | 37.84 s |
+| `nomic-ai/CodeRankEmbed` (fp32) | file | 1 | 6.7 GiB | 26.79 s | 34.54 s |
+| `nomic-ai/CodeRankEmbed` (fp32) | file | 2 | 12.9 GiB | 43.24 s | 58.98 s |
+| default | sorted | 4 | 1.3 GiB | 5.54 s | 5.46 s |
+| default | sorted | 32 | 6.1 GiB | 4.78 s | 5.17 s |
+| default | sorted | 64 | 9.1 GiB | 4.79 s | 5.25 s |
+| default | sorted | 353 (largest) | 18.2 GiB | 5.73 s | 5.83 s |
+| default | file | 4 | 0.6 GiB | 6.40 s | 5.80 s |
+| default | file | 32 | 1.6 GiB | 6.34 s | 6.03 s |
+| default | file | 64 | 2.8 GiB | 6.41 s | 6.16 s |
+| flattened | sorted | 4 | 0.8 GiB | 5.47 s | 5.10 s |
+| flattened | sorted | 32 | 3.8 GiB | 4.54 s | 4.97 s |
+| flattened | sorted | 64 | 6.1 GiB | 4.48 s | 5.01 s |
+| flattened | sorted | 628 (largest) | 19.2 GiB | 4.59 s | 5.25 s |
+| flattened | file | 4 | 0.6 GiB | 5.77 s | 5.37 s |
+| flattened | file | 32 | 1.6 GiB | 4.72 s | 4.97 s |
+| flattened | file | 64 | 2.7 GiB | 4.51 s | 5.04 s |
+
+Against the original, the 768 files agree at a mean cosine of 0.99991 and a minimum of 0.99881 to
+0.99941, depending on card and tier. This repo is 5.8× (3090 Ti) and 6.7× (5090 Laptop) faster
+than the original's best time. 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 1.0000 under the flattened load, against 0.99987 under the default load.
 
 ## What changed vs the source repo
 
@@ -233,10 +242,10 @@ load.
    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
+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
+   `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 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
@@ -269,8 +278,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},
 }
 ```