CodeRankEmbed-flash-attn

CodeRankEmbed-flash-attn

A bf16 quantization of nomic-ai/CodeRankEmbed with a three-tier attention dispatch built into a custom modeling_hf_nomic_bert.py shipped in this repo. It is not a finetune — the weights are the original CodeRankEmbed weights cast to bf16 (no further training). Two of the three tiers replace the original eager O(seq²) attention with an O(N) unpadded path; the third keeps the original eager algorithm as the correctness reference and universal fallback.

Why

nomic-ai/CodeRankEmbed loads through trust_remote_code, and its attention path is eager only — activation memory grows as batch × heads × seq², which OOMs at large batches even though the model is only 137M params. This repo adds two attention paths that compute the same attention in O(N) memory by packing unpadded sequences, so the large batches that OOM the eager path run comfortably — with parity embeddings (no quality change):

• 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 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-tier attention, 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 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, it logs one WARNING naming the tier and the error and steps down to the next 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)).

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

Parity & performance

The weights are the original CodeRankEmbed weights (bf16-cast), so embeddings match the fp32 original to within bf16 precision. The table below re-measures all three tiers after adding the dispatch: each tier's output on a 64-code-snippet corpus (fp32-renormalized), compared by cosine similarity against this same repo's pre-dispatch output on the same device (flash_attn on CUDA, eager on CPU — flash_attn can't run on CPU). eager and flash_attn each run the same computation as before the dispatch was added, so their cosines are parity checks (round to 1.000000 at 6 decimal places; true values are ≥ 0.9999997). torch_varlen is a different, torch-native kernel, and its cosine is the real signal.

Cosines rounded to 6 decimal places; peak VRAM rounded to the nearest MiB. eager always runs on CPU in this protocol (it is the universal fallback tier); the GPU named in the first column is the host each row's measurement ran on, not the device eager used on that row. torch was 2.12.1+cu130 for every row except the standalone CPU-only row (torch 2.14.0+cpu), which used a separate transformers==5.17.0 install to check the eager path against a newer transformers.

Separately, both varlen tiers (torch_varlen, flash_attn) are gated in this repo's downstream test suite at min cosine > 0.997 against the fp32 nomic-ai/CodeRankEmbed reference (not re-measured here — see the table above for this repo's own numbers), and stay under 20 GB peak VRAM at batch size 256.

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 upstream custom from_pretrained instantiated the model fp32 and load_state_dict-ed the checkpoint into fp32 params, ignoring torch_dtype. The copy here adds the standard transformers dtype resolution (explicit arg → config.torch_dtype → checkpoint dtype) so the model loads in its declared dtype.
3. Three-tier attention dispatch: NomicBertAttention.forward now selects one of three attention implementations at call time — torch_varlen (torch's own torch.nn.attention.varlen.varlen_attn, no third-party kernel, needs torch ≥ 2.10.0), flash_attn (the original flash-attn varlen-packed kernel, kept as a fallback for older torch builds that have it installed), and eager (the original padded attention algorithm, numerically unchanged, and the default off CUDA sm_80+; it now converts a raw 2-D [B, S] mask to the additive form before adding it, for the case where a device_map split hands it a mask a varlen tier upstream never passes). 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.

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 Text and Code Retrieval},
  author = {Suresh, K N Q and Wang, Xiang and Khan, Saqib and others},
  year   = {2025},
}