Padding-free encoder, flattening opt-in, torch.compile recipe; card measured on public corpora

README.md

@@ -65,6 +65,17 @@ file ships all three paths itself, so no runtime patching or post-load hooks are
   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.
+- **Padding-free on the varlen tiers.** The model unpads the batch once, right after the input
+  ids, runs the embeddings and every encoder layer on the real tokens only, and pads back once at
+  the end. `last_hidden_state` keeps its `[B, S, 768]` shape, with zeros at padding positions, so
+  pooling is unchanged.
+- **Flattened batches (opt-in).** Loading with
+  `model_kwargs={"attn_implementation": "flash_attention_2"}` lets sentence-transformers send each
+  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
+  the `eager` tier the flattened batch is re-padded and runs the padded path.
 - **`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
@@ -93,73 +104,188 @@ q = model.encode(queries, normalize_embeddings=True)
 d = model.encode(codes,   normalize_embeddings=True)
 ```
 
+With `flash_attn` installed, opt in to flattened batches:
+
+```python
+model = SentenceTransformer(
+    "handwoven8588/CodeRankEmbed-flash-attn",
+    trust_remote_code=True,
+    model_kwargs={"attn_implementation": "flash_attention_2"},
+)
+```
+
+On a varlen tier the encoder layers run on packed tokens only, so they compile cleanly. For about
+1.2× the throughput and a third less peak memory (see [torch.compile](#torchcompile) below):
+
+```python
+import torch
+
+model[0].auto_model.encoder = torch.compile(model[0].auto_model.encoder, dynamic=True)
+```
+
+`dynamic=True` matters: the number of packed tokens changes with every batch.
+
 ## Parity & performance
 
 The weights are the original CodeRankEmbed weights (bf16-cast), so embeddings match the fp32
 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 the first 64 `.py` files,
-by path, of the PenStock inference library's source tree; 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) | 32 | `torch_varlen` | 0.999941 | 0.999856 | 600 MiB | 0.197 s |
-| RTX 3090 Ti (sm_86) | 32 | `flash_attn` | 1.000000 | 1.000000 | 600 MiB | 0.197 s |
-| RTX 3090 Ti (sm_86) | 32 | `eager` | 0.999912 | 0.999856 | 807 MiB | 0.262 s |
-| RTX 3090 Ti (sm_86) | 64 | `torch_varlen` | 0.999941 | 0.999856 | 1201 MiB | 0.223 s |
-| RTX 3090 Ti (sm_86) | 64 | `flash_attn` | 1.000000 | 1.000000 | 1201 MiB | 0.218 s |
-| RTX 3090 Ti (sm_86) | 64 | `eager` | 0.999910 | 0.999856 | 1614 MiB | 0.302 s |
-| RTX 5090 Laptop GPU (sm_120) | 32 | `torch_varlen` | 0.999944 | 0.999919 | 600 MiB | 0.177 s |
-| RTX 5090 Laptop GPU (sm_120) | 32 | `flash_attn` | 1.000000 | 1.000000 | 600 MiB | 0.175 s |
-| RTX 5090 Laptop GPU (sm_120) | 32 | `eager` | 0.999912 | 0.999870 | 807 MiB | 0.249 s |
-| RTX 5090 Laptop GPU (sm_120) | 64 | `torch_varlen` | 0.999944 | 0.999919 | 1201 MiB | 0.202 s |
-| RTX 5090 Laptop GPU (sm_120) | 64 | `flash_attn` | 1.000000 | 1.000000 | 1201 MiB | 0.203 s |
-| RTX 5090 Laptop GPU (sm_120) | 64 | `eager` | 0.999913 | 0.999834 | 1614 MiB | 0.298 s |
+**Protocol.** 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. They are encoded with sentence-transformers in one `encode()`
+call per row at the stated `batch_size=` (64 to 4,096 batches per row), under the default load
+and under the flattened opt-in. Each row encodes once untimed, then again for the measurement.
+Cosine similarity is taken on fp32-renormalized output against one encode by the **pre-dispatch
+revision of this repo (`e361c6f`) on the same GPU at batch 32**, which runs the original
+`flash_attn` path on padded batches. `torch_varlen` is a different, torch-native kernel and
+`eager` a different algorithm, so their cosines measure those differences. `flash_attn` is the
+same attention kernel; where it sits below 1.000000, the rest of the network runs on packed
+tokens or on differently composed batches, so its matrix shapes, and with them its bf16
+rounding, differ. `auto`, the default with no override, kept `torch_varlen` on both cards and
+matches its rows. **Encode peak** is the encode's peak CUDA allocation above what was allocated
+before it (the loaded model's ~266 MiB of bf16 weights). Wall time covers all 16,384 functions.
+`eager` runs at batch 4 only: padded attention needs `batch × heads × seq²`, which at this
+corpus's 4,746-token maximum is about 17 GB of attention scores at batch 32. For comparison, the
+pre-dispatch revision at batch 32 peaks at 4,456 MiB and takes 23.8 s (3090 Ti) and 21.7 s
+(5090 Laptop).
+
+
+| GPU | load | batch size | tier | mean cos | min cos | encode peak | wall |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| RTX 3090 Ti (sm_86) | default | 4 | `eager` | 0.999923 | 0.999404 | 4,321 MiB | 41.791 s |
+| RTX 3090 Ti (sm_86) | default | 32 | `torch_varlen` | 0.999948 | 0.999403 | 2,511 MiB | 16.327 s |
+| RTX 3090 Ti (sm_86) | default | 32 | `flash_attn` | 0.999974 | 0.999814 | 2,511 MiB | 16.330 s |
+| RTX 3090 Ti (sm_86) | default | 256 | `torch_varlen` | 0.999949 | 0.999403 | 11,135 MiB | 15.027 s |
+| RTX 3090 Ti (sm_86) | default | 256 | `flash_attn` | 0.999979 | 0.999776 | 11,136 MiB | 15.056 s |
+| RTX 3090 Ti (sm_86) | default | 1,024 | `torch_varlen` | out of memory | | | |
+| RTX 3090 Ti (sm_86) | default | 1,024 | `flash_attn` | out of memory | | | |
+| RTX 3090 Ti (sm_86) | default | 2,048 | `torch_varlen` | out of memory | | | |
+| RTX 3090 Ti (sm_86) | default | 2,048 | `flash_attn` | out of memory | | | |
+| RTX 3090 Ti (sm_86) | flattened | 4 | `eager` | 0.999923 | 0.999404 | 4,329 MiB | 51.793 s |
+| RTX 3090 Ti (sm_86) | flattened | 32 | `torch_varlen` | 0.999948 | 0.999403 | 1,468 MiB | 15.529 s |
+| RTX 3090 Ti (sm_86) | flattened | 32 | `flash_attn` | 0.999967 | 0.999776 | 1,466 MiB | 15.496 s |
+| RTX 3090 Ti (sm_86) | flattened | 256 | `torch_varlen` | 0.999949 | 0.999403 | 7,073 MiB | 13.622 s |
+| RTX 3090 Ti (sm_86) | flattened | 256 | `flash_attn` | 0.999978 | 0.999776 | 7,072 MiB | 13.641 s |
+| RTX 3090 Ti (sm_86) | flattened | 1,024 | `torch_varlen` | 0.999949 | 0.999403 | 17,411 MiB | 13.371 s |
+| RTX 3090 Ti (sm_86) | flattened | 1,024 | `flash_attn` | 0.999978 | 0.999776 | 17,412 MiB | 13.478 s |
+| RTX 3090 Ti (sm_86) | flattened | 2,048 | `torch_varlen` | out of memory | | | |
+| RTX 3090 Ti (sm_86) | flattened | 2,048 | `flash_attn` | out of memory | | | |
+| RTX 5090 Laptop GPU (sm_120) | default | 4 | `eager` | 0.999925 | 0.999721 | 4,321 MiB | 33.589 s |
+| RTX 5090 Laptop GPU (sm_120) | default | 32 | `torch_varlen` | 0.999949 | 0.999748 | 2,511 MiB | 14.391 s |
+| RTX 5090 Laptop GPU (sm_120) | default | 32 | `flash_attn` | 0.999976 | 0.999766 | 2,511 MiB | 14.423 s |
+| RTX 5090 Laptop GPU (sm_120) | default | 256 | `torch_varlen` | 0.999951 | 0.999748 | 11,134 MiB | 15.171 s |
+| RTX 5090 Laptop GPU (sm_120) | default | 256 | `flash_attn` | 0.999988 | 0.999811 | 11,136 MiB | 15.155 s |
+| RTX 5090 Laptop GPU (sm_120) | default | 1,024 | `torch_varlen` | out of memory | | | |
+| RTX 5090 Laptop GPU (sm_120) | default | 1,024 | `flash_attn` | out of memory | | | |
+| RTX 5090 Laptop GPU (sm_120) | default | 2,048 | `torch_varlen` | out of memory | | | |
+| RTX 5090 Laptop GPU (sm_120) | default | 2,048 | `flash_attn` | out of memory | | | |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 4 | `eager` | 0.999924 | 0.998160 | 4,329 MiB | 44.613 s |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 32 | `torch_varlen` | 0.999950 | 0.999748 | 1,468 MiB | 13.699 s |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 32 | `flash_attn` | 0.999985 | 0.999811 | 1,466 MiB | 13.758 s |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 256 | `torch_varlen` | 0.999951 | 0.999748 | 7,072 MiB | 14.402 s |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 256 | `flash_attn` | 0.999988 | 0.999811 | 7,072 MiB | 14.478 s |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 1,024 | `torch_varlen` | 0.999951 | 0.999748 | 17,412 MiB | 14.739 s |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 1,024 | `flash_attn` | 0.999988 | 0.999811 | 17,412 MiB | 14.783 s |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 2,048 | `torch_varlen` | out of memory | | | |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 2,048 | `flash_attn` | out of memory | | | |
 
 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` (`2.8.3.post1` on the
-3090 Ti), all measured loading this repo at `1954de4`.
+3090 Ti).
 
-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.
+Memory tracks the real tokens in the heaviest batch, not the number of functions: both
+loads run the encoder padding-free, and each out-of-memory row failed allocating one 3,072-wide
+MLP activation over exactly that batch's tokens. sentence-transformers length-sorts the corpus,
+so under the default load the first batch of 1,024 gathers the longest functions, 869,948 tokens,
+which does not fit on 24 GB. When it flattens, it interleaves longest with shortest to even out
+batches, so its heaviest batch of 1,024 holds 593,716 tokens and fits at 17.4 GB; at 2,048 even
+that reaches 889,626 tokens and does not. Up to that limit, flattening cuts encode memory by
+36–42% and wall time by 5–9% against the default load. The interleaving
+also changes which texts share a padded `eager` batch, which is why the flattened `eager` row's
+minimum differs on the 5090 Laptop: its lowest text is a 2,166-token function padded next to a
+2,195-token one (the same text alone scores 0.999936).
 
 **Real documents at the model defaults.** 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), encoded with sentence-transformers at this model's default
-`max_seq_length` of 8192 and one `encode()` call per row, so the files are length-sorted before
-batching. Wall time covers the whole corpus; peak VRAM is the peak CUDA allocation, weights
-included. The largest batch row is the largest batch size that fit on each card (the 3090 Ti was
-sharing about 1 GiB with other processes).
-
-| GPU | batch size | tier | wall (768 files) | peak VRAM |
-| --- | --- | --- | --- | --- |
-| RTX 3090 Ti (sm_86) | 4 | `torch_varlen` | 6.57 s | 1,295 MiB |
-| RTX 3090 Ti (sm_86) | 4 | `flash_attn` | 6.58 s | 1,296 MiB |
-| RTX 3090 Ti (sm_86) | 4 | `eager` | 14.18 s | 12,961 MiB |
-| RTX 3090 Ti (sm_86) | 32 | `torch_varlen` | 6.37 s | 7,973 MiB |
-| RTX 3090 Ti (sm_86) | 64 | `torch_varlen` | 7.24 s | 15,659 MiB |
-| RTX 3090 Ti (sm_86) | 71 | `torch_varlen` | 7.60 s | 17,340 MiB |
-| RTX 5090 Laptop GPU (sm_120) | 4 | `torch_varlen` | 5.54 s | 1,319 MiB |
-| RTX 5090 Laptop GPU (sm_120) | 4 | `flash_attn` | 5.57 s | 1,319 MiB |
-| RTX 5090 Laptop GPU (sm_120) | 4 | `eager` | 16.27 s | 12,983 MiB |
-| RTX 5090 Laptop GPU (sm_120) | 32 | `torch_varlen` | 6.36 s | 7,995 MiB |
-| RTX 5090 Laptop GPU (sm_120) | 64 | `torch_varlen` | 7.48 s | 15,681 MiB |
-| RTX 5090 Laptop GPU (sm_120) | 73 | `torch_varlen` | 7.83 s | 17,843 MiB |
-
-At batch 4 `eager` takes 2.2× (3090 Ti) to 2.9× (5090 Laptop) as long as the varlen tiers and needs
-about ten times the memory. At batch 32 with the same files in random order, `eager` ran out of
-memory on 21 of the 24 batches (3090 Ti). Both varlen tiers are also gated downstream at min
-cosine > 0.997 against the fp32 `nomic-ai/CodeRankEmbed` reference.
+`max_seq_length` of 8192, under the default load and under the flattened opt-in. **Sorted** is
+one `encode()` call over the whole corpus, so sentence-transformers length-sorts the files before
+batching; **file** is one `encode()` call per batch of consecutive files, so each batch mixes short
+and long files. Wall time covers the whole corpus; peak VRAM is the peak CUDA allocation, weights
+included. The largest sorted batch on each card is the largest batch size that fit (the 5090
+Laptop was sharing about 0.4 GiB with another process).
+
+| GPU | load | order | batch size | tier | wall (768 files) | peak VRAM |
+| --- | --- | --- | --- | --- | --- | --- |
+| RTX 3090 Ti (sm_86) | default | sorted | 4 | `torch_varlen` | 5.61 s | 1,296 MiB |
+| RTX 3090 Ti (sm_86) | default | sorted | 32 | `torch_varlen` | 4.79 s | 6,186 MiB |
+| RTX 3090 Ti (sm_86) | default | sorted | 64 | `torch_varlen` | 4.80 s | 9,253 MiB |
+| RTX 3090 Ti (sm_86) | default | sorted | 353 | `torch_varlen` | 5.77 s | 18,637 MiB |
+| RTX 3090 Ti (sm_86) | default | sorted | 32 | `flash_attn` | 4.79 s | 6,187 MiB |
+| RTX 3090 Ti (sm_86) | default | file | 4 | `torch_varlen` | 6.49 s | 636 MiB |
+| RTX 3090 Ti (sm_86) | default | file | 32 | `torch_varlen` | 6.39 s | 1,614 MiB |
+| RTX 3090 Ti (sm_86) | default | file | 64 | `torch_varlen` | 6.44 s | 2,801 MiB |
+| RTX 3090 Ti (sm_86) | flattened | sorted | 4 | `torch_varlen` | 5.51 s | 792 MiB |
+| RTX 3090 Ti (sm_86) | flattened | sorted | 32 | `torch_varlen` | 4.54 s | 3,843 MiB |
+| RTX 3090 Ti (sm_86) | flattened | sorted | 64 | `torch_varlen` | 4.47 s | 6,237 MiB |
+| RTX 3090 Ti (sm_86) | flattened | sorted | 628 | `torch_varlen` | 4.60 s | 19,658 MiB |
+| RTX 3090 Ti (sm_86) | flattened | sorted | 32 | `flash_attn` | 4.54 s | 3,844 MiB |
+| RTX 3090 Ti (sm_86) | flattened | file | 4 | `torch_varlen` | 5.83 s | 634 MiB |
+| RTX 3090 Ti (sm_86) | flattened | file | 32 | `torch_varlen` | 4.73 s | 1,607 MiB |
+| RTX 3090 Ti (sm_86) | flattened | file | 64 | `torch_varlen` | 4.53 s | 2,787 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | sorted | 4 | `torch_varlen` | 4.75 s | 1,319 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | sorted | 32 | `torch_varlen` | 4.72 s | 6,209 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | sorted | 64 | `torch_varlen` | 4.84 s | 9,276 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | sorted | 353 | `torch_varlen` | 5.98 s | 18,661 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | sorted | 32 | `flash_attn` | 4.71 s | 6,210 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | file | 4 | `torch_varlen` | 5.64 s | 660 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | file | 32 | `torch_varlen` | 5.70 s | 1,636 MiB |
+| RTX 5090 Laptop GPU (sm_120) | default | file | 64 | `torch_varlen` | 5.84 s | 2,823 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | sorted | 4 | `torch_varlen` | 4.56 s | 816 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | sorted | 32 | `torch_varlen` | 4.56 s | 3,866 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | sorted | 64 | `torch_varlen` | 4.63 s | 6,260 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | sorted | 628 | `torch_varlen` | 4.75 s | 19,681 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | sorted | 32 | `flash_attn` | 4.57 s | 3,867 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | file | 4 | `torch_varlen` | 4.75 s | 658 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | file | 32 | `torch_varlen` | 4.70 s | 1,630 MiB |
+| RTX 5090 Laptop GPU (sm_120) | flattened | file | 64 | `torch_varlen` | 4.66 s | 2,811 MiB |
+
+Flattening keeps peak memory well below the default load at every sorted batch size and fits a
+batch of 628 on 24 GB, where the default load fits 353. It also makes a text's embedding
+independent of the other texts in its batch: between the file-order and sorted runs, the minimum
+cosine over the 768 files is 0.9999998 on both cards and both varlen tiers under the flattened
+load, against as low as 0.99987 under the default load.
+
+### torch.compile
+
+The parity table's 16,384 functions on `torch_varlen`, with the encoder layers eager and compiled
+(`torch.compile(model[0].auto_model.encoder, dynamic=True)`). Cosine is against the same
+pre-dispatch encode as the parity table.
+
+| GPU | load | batch size | eager | compiled | compiled min cos (mean) |
+| --- | --- | --- | --- | --- | --- |
+| RTX 3090 Ti (sm_86) | default | 32 | 16.38 s, 2,511 MiB | 14.05 s, 1,635 MiB | 0.998156 (0.999943) |
+| RTX 3090 Ti (sm_86) | default | 256 | 15.06 s, 11,135 MiB | 12.74 s, 7,249 MiB | 0.998156 (0.999943) |
+| RTX 3090 Ti (sm_86) | default | 1,024 | out of memory | 13.81 s, 16,701 MiB | 0.998156 (0.999943) |
+| RTX 3090 Ti (sm_86) | flattened | 32 | 15.58 s, 1,468 MiB | 13.16 s, 956 MiB | 0.998156 (0.999943) |
+| RTX 3090 Ti (sm_86) | flattened | 256 | 13.67 s, 7,073 MiB | 11.37 s, 4,602 MiB | 0.998156 (0.999943) |
+| RTX 3090 Ti (sm_86) | flattened | 1,024 | 13.40 s, 17,411 MiB | 10.90 s, 11,325 MiB | 0.998156 (0.999943) |
+| RTX 5090 Laptop GPU (sm_120) | default | 32 | 14.33 s, 2,511 MiB | 12.19 s, 1,637 MiB | 0.998657 (0.999944) |
+| RTX 5090 Laptop GPU (sm_120) | default | 256 | 15.08 s, 11,134 MiB | 12.30 s, 7,249 MiB | 0.998657 (0.999944) |
+| RTX 5090 Laptop GPU (sm_120) | default | 1,024 | out of memory | 13.39 s, 16,701 MiB | 0.999250 (0.999944) |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 32 | 13.67 s, 1,468 MiB | 11.72 s, 956 MiB | 0.999250 (0.999944) |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 256 | 14.35 s, 7,072 MiB | 11.66 s, 4,601 MiB | 0.999250 (0.999944) |
+| RTX 5090 Laptop GPU (sm_120) | flattened | 1,024 | 14.67 s, 17,412 MiB | 11.80 s, 11,326 MiB | 0.999250 (0.999944) |
+
+Compiling makes the encode 1.17–1.24× faster and cuts its peak memory by about a third, because
+inductor fuses the element-wise work around each attention call; the default load's batch of 1,024
+then fits. It compiles one graph with no graph breaks, once per process (the first encode took about
+4 s longer on the 5090 Laptop and 10 s on the 3090 Ti), 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. The mean cosine is unchanged; a few texts move by bf16 reassociation,
+to a minimum of 0.998156 against 0.999403 eager on the 3090 Ti.
 
 ## What changed vs the source repo
 
@@ -178,19 +304,25 @@ cosine > 0.997 against the fp32 `nomic-ai/CodeRankEmbed` reference.
    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). 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.
+   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
    pinned load fetches the weights from the pinned commit too. It also accepts transformers'
-   `dtype=` alongside the older `torch_dtype=`.
+   `dtype=` alongside the older `torch_dtype=`, and returns the model in eval mode, as
+   transformers' own `from_pretrained` does.
+5. **Flattened input**: `NomicBertModel.forward` also accepts sentence-transformers' flattened
+   batches (`input_ids [1, ΣL]` with per-sequence `position_ids` and `cu_seq_lens_q`), and the
+   model declares flash-attention support so that sentence-transformers sends them when loaded
+   with `attn_implementation="flash_attention_2"` (see Behavior).
 
 ## License & attribution
 

modeling_hf_nomic_bert.py

@@ -12,7 +12,7 @@ import os
 import re
 from collections import OrderedDict
 from functools import partial
-from typing import List, Optional, Tuple, Union
+from typing import NamedTuple, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -236,6 +236,21 @@ def _pad(x_u: torch.Tensor, indices: torch.Tensor, B: int, S: int) -> torch.Tens
     return out.view(B, S, *x_u.shape[1:])
 
 
+class _Packed(NamedTuple):
+    """Metadata for hidden states kept packed ``[N, D]`` across the encoder.
+
+    ``cu_seqlens`` / ``max_seqlen`` feed the varlen kernel; ``positions`` (int64,
+    ``[N]``) is each token's column in the padded batch, the position RoPE uses on
+    the dense path; ``rope_seqlen`` is the padded width ``S``, the length the dense
+    path sizes the RoPE cache to.
+    """
+
+    cu_seqlens: torch.Tensor
+    max_seqlen: int
+    positions: torch.Tensor
+    rope_seqlen: int
+
+
 def select_attention_impl(device: torch.device) -> str:
     """The attention tier for ``device`` under the current ``NOMIC_BERT_ATTN_IMPL``.
 
@@ -541,6 +556,10 @@ class NomicBertPreTrainedModel(PreTrainedModel):
     supports_gradient_checkpointing = True
     _no_split_modules = ["Block"]
     _skip_keys_device_placement = "past_key_values"
+    # Pre-packed input (cu_seq_lens_q / position_ids through forward's **kwargs) runs on
+    # either varlen tier, so the model can accept a flattened batch.
+    _supports_attention_backend = True
+    _supports_flash_attn = True
 
     def __init__(self, config, *inputs, **kwargs):
         super().__init__(config)
@@ -598,6 +617,9 @@ class NomicBertPreTrainedModel(PreTrainedModel):
         if rotary_scaling_factor:
             config.rotary_scaling_factor = rotary_scaling_factor
 
+        if kwargs.get("attn_implementation") == "flash_attention_2":
+            # Opt-in: advertises flash attention so sentence-transformers flattens batches.
+            config._attn_implementation = "flash_attention_2"
         if config.n_positions <= 0 and config.rotary_emb_fraction > 0:
             config.n_positions = 2048
         if num_labels:
@@ -663,6 +685,9 @@ class NomicBertPreTrainedModel(PreTrainedModel):
             logger.warning(load_return)
         else:
             logger.debug(load_return)
+        # transformers' native from_pretrained returns the model in eval mode; this
+        # override must too, or a direct forward runs the config's dropout.
+        model.eval()
         return model
 
     def _set_gradient_checkpointing(self, module, value=False):
@@ -929,6 +954,29 @@ class NomicBertRotaryEmbedding(nn.Module):
         return torch.stack((q_rot, k_rot, qkv[:, :, 2]), dim=2)
 
 
+def _rotary_packed(rotary: "NomicBertRotaryEmbedding", qkv_u: torch.Tensor, packed: _Packed) -> torch.Tensor:
+    """RoPE on packed ``[N, 3, H, D]``: each token rotated by its own column position.
+
+    Same arithmetic as ``NomicBertRotaryEmbedding.forward`` on the dense
+    ``[B, S, 3, H, D]`` tensor, so every real token gets bit-identical q/k. The caller
+    sizes the cache first (``NomicBertModel._size_rotary_caches``), outside the encoder:
+    the update branches on the module's cached length, which ``torch.compile`` guards as
+    a static int, so updating it here would recompile a compiled encoder on every growth.
+    """
+    pattern = "... d -> ... 1 (2 d)" if not rotary.interleaved else "... d -> ... 1 (d 2)"
+    cos = repeat(rotary._cos_cached[packed.positions], pattern)
+    sin = repeat(rotary._sin_cached[packed.positions], pattern)
+    ro_dim = cos.shape[-1]
+
+    def rot(x):
+        return torch.cat(
+            [x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], rotary.interleaved) * sin, x[..., ro_dim:]],
+            dim=-1,
+        )
+
+    return torch.stack((rot(qkv_u[:, 0]), rot(qkv_u[:, 1]), qkv_u[:, 2]), dim=1)
+
+
 class NomicBertDynamicNTKRotaryEmbedding(NomicBertRotaryEmbedding):
     def __init__(self, rotary_scaling_factor, max_position_embeddings, **kwargs):
         super().__init__(**kwargs)
@@ -1069,8 +1117,12 @@ class NomicBertAttention(nn.Module):
         is_padded_inputs: Optional[bool] = True,
         cu_seqlens: Optional[torch.Tensor] = None,
         max_seq_len: Optional[int] = None,
+        packed: Optional[_Packed] = None,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 
+        if packed is not None:
+            return self._forward_packed(hidden_states, packed)
+
         has_layer_past = past_key_value is not None
 
         if has_layer_past:
@@ -1172,6 +1224,23 @@ class NomicBertAttention(nn.Module):
 
         return attn_output
 
+    def _forward_packed(self, hidden_states: torch.Tensor, packed: _Packed) -> torch.Tensor:
+        """Varlen attention on packed ``[N, D]`` hidden states; returns ``[N, D]``."""
+        qkv = self.Wqkv(hidden_states)
+        qkv = rearrange(qkv, "n (three h d) -> n three h d", three=3, d=self.head_dim)
+        if self.rotary_emb_dim > 0:
+            qkv = _rotary_packed(self.rotary_emb, qkv, packed)
+        orig_dtype = qkv.dtype
+        if orig_dtype not in (torch.float16, torch.bfloat16):
+            qkv = qkv.to(torch.bfloat16)
+        cu, max_s = packed.cu_seqlens, packed.max_seqlen
+        if select_attention_impl(hidden_states.device) == "torch_varlen":
+            q, k, v = qkv.unbind(1)
+            out = _torch_varlen_attn(q, k, v, cu, cu, max_s, max_s)
+        else:
+            out = flash_attn_varlen_qkvpacked_func(qkv, cu, max_s, dropout_p=0.0, causal=False)
+        return self.out_proj(rearrange(out.to(orig_dtype), "n h d -> n (h d)"))
+
 
 class NomicBertBlock(nn.Module):
     def __init__(
@@ -1225,6 +1294,7 @@ class NomicBertBlock(nn.Module):
         use_cache: Optional[bool] = False,
         cu_seqlens: Optional[torch.Tensor] = None,
         max_seq_len: Optional[int] = None,
+        packed: Optional[_Packed] = None,
     ):
         r"""Pass the input through the encoder layer.
         Args:
@@ -1244,6 +1314,7 @@ class NomicBertBlock(nn.Module):
                 is_padded_inputs=is_padded_inputs,
                 cu_seqlens=cu_seqlens,
                 max_seq_len=max_seq_len,
+                packed=packed,
             )
 
             dropped = self.dropout2(hidden_states)
@@ -1260,6 +1331,7 @@ class NomicBertBlock(nn.Module):
                 is_padded_inputs=is_padded_inputs,
                 cu_seqlens=cu_seqlens,
                 max_seq_len=max_seq_len,
+                packed=packed,
             )
             hidden_states = self.norm1((self.dropout1(attn_outputs) + hidden_states).to(dtype=self.norm1.weight.dtype))
             mlp_out = self.mlp(hidden_states)
@@ -1287,6 +1359,7 @@ class NomicBertEncoder(nn.Module):
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         is_padded_inputs: Optional[bool] = True,
+        packed: Optional[_Packed] = None,
     ):
         """If subset_mask is not None, we only want output for the subset of the sequence.
         This means that we only compute the last layer output for these tokens.
@@ -1314,6 +1387,7 @@ class NomicBertEncoder(nn.Module):
                     None,
                     None,
                     is_padded_inputs,
+                    packed=packed,
                     # if you freeze ANY layers, you need `use_reentrant=False`
                     # https://github.com/huggingface/transformers/issues/21381
                     # https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/7
@@ -1331,6 +1405,7 @@ class NomicBertEncoder(nn.Module):
                     is_padded_inputs,
                     output_attentions,
                     use_cache,
+                    packed=packed,
                 )
         return hidden_states
 
@@ -1427,16 +1502,26 @@ class NomicBertModel(NomicBertPreTrainedModel):
         token_type_ids=None,
         position_ids=None,
         return_dict=None,
+        cu_seq_lens_q=None,
+        max_length_q=None,
+        **kwargs,
     ):
         if token_type_ids is None:
             token_type_ids = torch.zeros_like(input_ids)
+        if cu_seq_lens_q is not None:
+            self.attention_impl = select_attention_impl(input_ids.device)
+            _log_attention_impl_once(self.attention_impl, input_ids.device)
+            return self._forward_prepacked(input_ids, token_type_ids, position_ids, cu_seq_lens_q, max_length_q)
+
+        impl = select_attention_impl(input_ids.device)
+        self.attention_impl = impl
+        _log_attention_impl_once(impl, input_ids.device)
+        if impl != "eager" and (attention_mask is None or attention_mask.ndim == 2):
+            return self._forward_packed(input_ids, attention_mask, token_type_ids, position_ids)
+
         hidden_states = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)
         hidden_states = self.emb_ln(hidden_states)
         hidden_states = self.emb_drop(hidden_states)
-
-        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.
@@ -1454,6 +1539,94 @@ class NomicBertModel(NomicBertPreTrainedModel):
             pooler_output=pooled_output,
         )
 
+    def _size_rotary_caches(self, seqlen, device):
+        """Grow each layer's RoPE cos/sin cache to ``seqlen`` before the packed encoder runs
+        (see ``_rotary_packed``), in the dtype its q/k projection produces."""
+        for layer in self.encoder.layers:
+            attn = layer.attn
+            if attn.rotary_emb_dim > 0:
+                attn.rotary_emb._update_cos_sin_cache(seqlen, device=device, dtype=attn.Wqkv.weight.dtype)
+
+    def _forward_packed(self, input_ids, attention_mask, token_type_ids, position_ids):
+        """Varlen tiers: unpad once, run embeddings and every layer on real tokens only.
+
+        ``last_hidden_state`` comes back padded ``[B, S, D]`` with zeros at padding
+        positions, so masked pooling downstream is unchanged.
+        """
+        B, S = input_ids.shape
+        mask = (
+            torch.ones(B, S, device=input_ids.device, dtype=torch.bool)
+            if attention_mask is None
+            else attention_mask.to(torch.bool)
+        )
+        seqlens = mask.sum(-1, dtype=torch.int32)
+        max_s = int(seqlens.max())
+        if max_s == 0:
+            raise RuntimeError(
+                "NomicBertModel.forward: every sequence is empty (max_s == 0); "
+                "pre-filter empty inputs before encoding."
+            )
+        indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten()
+        cu = F.pad(torch.cumsum(seqlens, 0, dtype=torch.int32), (1, 0))
+        columns = indices % S
+        pos_u = position_ids.flatten()[indices] if position_ids is not None else columns
+        hidden_states = self.embeddings(
+            input_ids.flatten()[indices][None],
+            position_ids=pos_u[None],
+            token_type_ids=token_type_ids.flatten()[indices][None],
+        )[0]
+        hidden_states = self.emb_drop(self.emb_ln(hidden_states))
+        packed = _Packed(cu_seqlens=cu, max_seqlen=max_s, positions=columns, rope_seqlen=S)
+        self._size_rotary_caches(packed.rope_seqlen, hidden_states.device)
+        hidden_states = self.encoder(hidden_states, packed=packed)
+        sequence_output = _pad(hidden_states, indices, B, S)
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+        )
+
+    def _forward_prepacked(self, input_ids, token_type_ids, position_ids, cu_seq_lens_q, max_length_q):
+        """Input already packed by the caller: ``input_ids [1, ΣL]`` with per-sequence
+        ``position_ids`` and ``cu_seq_lens_q`` boundaries. Returns ``[1, ΣL, D]``."""
+        if self.attention_impl == "eager":
+            return self._forward_prepacked_eager(input_ids, token_type_ids, cu_seq_lens_q)
+        if position_ids is None:
+            raise ValueError("NomicBertModel.forward: pre-packed input needs per-sequence position_ids")
+        positions = position_ids.reshape(-1).long()
+        max_s = int(max_length_q) if max_length_q is not None else int((cu_seq_lens_q[1:] - cu_seq_lens_q[:-1]).max())
+        hidden_states = self.embeddings(input_ids, position_ids=position_ids, token_type_ids=token_type_ids)[0]
+        hidden_states = self.emb_drop(self.emb_ln(hidden_states))
+        packed = _Packed(cu_seqlens=cu_seq_lens_q.to(torch.int32), max_seqlen=max_s, positions=positions,
+                         rope_seqlen=max_s)
+        self._size_rotary_caches(packed.rope_seqlen, hidden_states.device)
+        sequence_output = self.encoder(hidden_states, packed=packed)[None]
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+        )
+
+    def _forward_prepacked_eager(self, input_ids, token_type_ids, cu_seq_lens_q):
+        """Pre-packed input on the eager tier (forced, or no varlen kernel on this device):
+        re-pad to ``[B, S]`` (right padding, so each token keeps its per-sequence
+        position), run the padded eager forward, and gather the real tokens back to
+        ``[1, ΣL, D]``, the shape the caller's pooling expects."""
+        cu = cu_seq_lens_q.long()
+        seqlens = cu[1:] - cu[:-1]
+        B, S = seqlens.shape[0], int(seqlens.max())
+        mask = torch.arange(S, device=input_ids.device)[None, :] < seqlens[:, None]
+        indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten()
+        ids = _pad(input_ids.reshape(-1), indices, B, S)
+        types = _pad(token_type_ids.reshape(-1), indices, B, S)
+        out = self.forward(ids, attention_mask=mask.long(), token_type_ids=types)
+        sequence_output = out.last_hidden_state.flatten(0, 1)[indices][None]
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+        )
+
 
 class NomicBertForPreTraining(NomicBertPreTrainedModel):
     _tied_weights_keys = ["predictions.decoder.bias", "cls.predictions.decoder.weight"]