CodeRankEmbed-flash-attn

A bf16 quantization of nomic-ai/CodeRankEmbed 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 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 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.

Both run the same FA2-family kernel and are gated by the same GPU-capability check. The modeling file ships all three paths itself, so no runtime patching or post-load hooks are needed.

Behavior

• 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.
  2. flash_attn — CUDA, compute capability sm_80+, the flash_attn package importable, on a 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'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 whether the installed build has a kernel for the GPU's architecture, e.g. on ROCm); if the probe raises a RuntimeError (how a missing or unsupported kernel fails), it logs one WARNING naming the tier and the error and steps down to the next tier. An out-of-memory error, any other exception type, or a CUDA error left over from earlier work propagates instead of demoting the tier, so a transient failure never pins the device to a slower tier. A forced override (e.g. NOMIC_BERT_ATTN_IMPL=torch_varlen) is not probed and raises RuntimeError if that tier's precondition doesn't hold — a forced tier never falls back silently. An unrecognized override raises ValueError.

• See which tier engaged: model[0].auto_model.attention_impl after a forward pass. 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 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 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).

• eager runs in bf16 too (because the stored weights are bf16), numerically equivalent to the varlen tiers, just without their memory and throughput wins. The model loads and encodes on any host regardless of which tier engages.

Usage

Identical to the original. The query prompt must include the task-instruction prefix "Represent this query for searching relevant code: "; documents need no prefix.

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("handwoven8588/CodeRankEmbed-flash-attn", trust_remote_code=True)
queries = ["Represent this query for searching relevant code: Calculate the n-th factorial"]
codes   = ["def fact(n):\n    if n < 0:\n        raise ValueError\n    return 1 if n == 0 else n * fact(n - 1)"]

q = model.encode(queries, normalize_embeddings=True)
d = model.encode(codes,   normalize_embeddings=True)

With flash_attn installed, opt in to flattened batches:

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):

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

Every number below compares this repo with the original model, 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'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 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

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

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).

• 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

1. Weights: fp32 → bf16. flash_attn and torch_varlen only accept half precision and the 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 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 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 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.
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

MIT — same license as nomic-ai/CodeRankEmbed (see NOTICE). The weights, tokenizer, and the bulk of the modeling file are a verbatim derivative of nomic-ai/CodeRankEmbed; the modeling file derives from Tri Dao's BERT implementation, and CodeRankEmbed was trained by the CoRNStack team (Suresh et al., 2025). Cite their work:

@misc{suresh2025cornstackhighqualitycontrastivedata,
      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},
}