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

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/encode-time.png filter=lfs diff=lfs merge=lfs -text
+assets/encode-memory.png filter=lfs diff=lfs merge=lfs -text
+assets/parity.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -60,13 +60,19 @@ file ships all three paths itself, so no runtime patching or post-load hooks are
   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)`). 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.
+- **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
+  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.
 - **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
@@ -93,73 +99,129 @@ 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 Speed and memory 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 |
-
-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`.
-
-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.
-
-**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.
+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.
+
+**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).
+
+### 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[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.
 
 ## What changed vs the source repo
 
@@ -178,19 +240,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"]