⭐️ Overview

packed-encoders is a high-performance execution engine for ModernBERT backbones and, as its second architecture, the Qwen3.5 hybrid behind topk-embed-v1. One call — pe.pack(model) — monkeypatches a live Hugging Face ModernBertModel in place with fused CuteDSL LayerNorm, RoPE, and GeGLU kernels, cuBLAS GEMMs, packed (padding-free) attention with per-GPU calibrated dispatch, and optional CUDA graphs.

The kernels consume Hugging Face's exact weight layout, so nothing is re-packed and no wrapper class exists. state_dict, save_pretrained, from_pretrained, and gradient checkpointing all remain HF's own (or any other wrappers on top of HF). pe.unpack(model) restores the stock forward.

Inference: 1.9–2.1x document encoding and 2.5–2.6x query encoding vs eager Hugging Face on an RTX 5090 (1.3–1.5x vs torch.compile max-autotune), at cosine ≥ 0.999 output parity.

Indexing: 2.9–3.3x documents/second on a real long-document corpus, with ~26% less peak reserved VRAM.

Training: 7.0x median GradCache step time on a B200 running the Agent-ModernColBERT recipe — 4.6x wall-clock even after paying all compile warmup — with ~37% less per-step reserved memory and a bit-identical initial loss.

topk-embed-v1 (xsmall, small): 7.5–16x queries/second and 1.7–4.5x documents/second at batch 8 over the model as shipped, end to end through encode on L40S and H100, at mean token cosine ≥ 0.9998 to the model as shipped. The model's own loader options (bf16, FlashAttention 2, torch.compile) do not speed it up.

Check the benchmark section for the full tables.

 

Why packing?

Encoder runtimes were never meant to be rectangular, we're just inheriting decoder tooling and living with that. Queries are ~30 tokens, documents are hundreds to thousands, and contrastive batches mix both. Padding every sequence to the batch maximum spends real FLOPs, memory, and attention work on tokens that don't exist.

packed-encoders removes that waste in three layers, each usable independently:

1. Fused tail — CuteDSL LayerNorm / residual-LayerNorm / RoPE / GeGLU kernels plus cuBLAS GEMMs replace the stock layer internals. The kernels are shape-agnostic (the token count is a runtime grid dimension), so one compiled artifact serves every batch shape with no recompiles.
2. Packed execution — the padded [B, S] forward internally unpads to a flat [total_tokens] buffer and runs variable-length attention confined by cu_seqlens; pipelines that already collate without padding can call packed_forward() directly and skip the pad/unpad boundary entirely.
3. Dispatch and graphs — a per-GPU calibrated policy picks between a specialized packed Triton attention kernel and FlashAttention per batch, and optional bucketed CUDA graphs collapse the kernel-launch floor for short-sequence workloads.

Installation

Install the base runtime with CUDA 12.8:

uv pip install --torch-backend cu128 packed-encoders

From source, the lockfile selects Torch 2.11 / Triton 3.6 / CUDA 12.8:

git clone git@github.com:pau-mensa/packed-encoders.git
cd packed-encoders
uv sync --locked --no-dev                                # base runtime
uv sync --locked --no-dev --extra fa2                     # compiled FA2, including RTX 5090
uv sync --locked --no-dev --extra fa2 --extra pylate       # PyLate / ModernColBERT
uv sync --locked --no-dev --extra qwen3_5 --extra fa2      # topk-embed-v1, complete environment
uv sync --locked --no-dev --extra fa4                     # CuteDSL FA4 on sm_90 / sm_100

FA2 wheels come from Astral's GPU index, with explicit Torch ABI dependencies. The locked 2.11 wheels cover Linux x86-64/aarch64 and CPython 3.10–3.14, replacing the old Python-3.11-only GitHub wheel. Only FA2 is sourced from this index. For an existing environment or a published package, uv pip needs the index explicitly for published packages (project source settings are not wheel metadata):

uv pip install --no-sources --torch-backend cu128 \
  --index https://wheels.astral.sh/simple/cu128/ \
  "packed-encoders[fa2,qwen3_5]"
# Use -e ".[fa2,qwen3_5]" instead of the package name for an editable checkout.

• The package requires Torch 2.11 and Triton 3.6 on Python 3.10–3.14. Users retaining Torch 2.8 should pin packed-encoders==0.1.0.
• The qwen3_5 extra selects Torch 2.11, Transformers 5.9, FLA, the SentenceTransformers multi-vector wrapper and matching torchvision. It needs no --no-deps workaround and can be combined with fa2 or fa4.
• PyLate and qwen3_5 need separate environments: PyLate currently pins SentenceTransformers 5.3, while topk uses 6.x. Request the extras individually; --all-extras is intentionally invalid. The sentence-transformers extra can be combined with either.
• FlashAttention is optional: the packed Triton kernel ships in-tree and SDPA needs no extra dependency. On large token budgets FA is heavily recommended. The FA4 GPU paths still require validation on their target hardware; the local 5090 correctness suite validates the compiled FA2 path.
• Add --group test for pytest or omit --no-dev for benchmark/development dependencies. GPU tests require model downloads on their first run. For CPU CI, use uv pip install --no-sources --torch-backend cpu -e . pytest.

For the tested environments and remaining performance checks, see the Torch installation review.

Requirements

• An NVIDIA GPU of compute capability sm_80, sm_89, sm_90, sm_100, or sm_120, with a working ptxas (CuteDSL JIT-compiles at runtime).
• bf16 weights for inference, or fp32 master weights under bf16 autocast for training.

⚡️ Quick Start

import packed_encoders as pe
from transformers import AutoModel

model = AutoModel.from_pretrained("answerdotai/ModernBERT-base", dtype="bfloat16").cuda()
pe.pack(model)                      # validated fused forward, installed in place
out = model(input_ids=ids, attention_mask=mask).last_hidden_state

pack() locates the encoder inside whatever you hand it — a raw HF model, a SentenceTransformer, or a PyLate ColBERT:

  from sentence_transformers import SentenceTransformer
  st = SentenceTransformer("answerdotai/ModernBERT-base").cuda().bfloat16()
+ pe.pack(st)
  emb = st.encode(texts)

  from pylate import models
  colbert = models.ColBERT(model_name_or_path="lightonai/GTE-ModernColBERT-v1")
+ pe.pack(colbert)
  train_loss = losses.CachedContrastive(model=colbert, ...)

topk-embed-v1 packs the same way, with CUDA graphs on by default:

  from sentence_transformers.multi_vector_encoder import MultiVectorEncoder
  model = MultiVectorEncoder("topk-io/topk-embed-v1-xsmall", trust_remote_code=True, device="cuda",
                             model_kwargs={"torch_dtype": torch.bfloat16})
+ pe.pack(model)
  docs = model.encode(texts, task="document", batch_size=8)

Usage

Attention backends

pe.pack(model, attention_backend="auto")   # default when an optimized kernel is importable

Backend	What it is
"auto"	Calibrated per-batch choice between the packed Triton kernel and FlashAttention. Exact-card policies exist for RTX 5090, A100-SXM4-40GB, L40S, H200, and B200; uncalibrated cards conservatively prefer Flash.
"triton"	The in-tree packed short-attention kernel — no-grad bf16 ModernBERT-base geometry (12 × 64 heads), global or local-64 attention, Smax ≤ 128. Invariant misses prefer Flash, then SDPA.
"flash"	FlashAttention with sliding-window pruning; varlen kernel on padded batches. FA2 wheel on sm_120, CuteDSL FA4 on sm_90/sm_100. Raises if no kernel is present.
"sdpa"	Dependency-free dense-mask attention. The explicit safe choice and the final fallback.

The dispatch policy uses only host-visible shapes (never device values, so it stays CUDA-graph-capturable) and was fit per card against measured crossovers: on the RTX 5090 the Triton kernel wins below 20,736 live tokens and FA2 above; at the kernel level the specialized Triton kernel is 1.4–2.1x faster than FA2 in its envelope. Policies are exact-card and backend-version specific by design — see the benchmarks/ guide for the formulas, sweeps, parity numbers, and reproduction commands.

CUDA graphs

The package also offers the possibility of using CUDA graphs for short-sequence inference, which is usually host-launch-bound; the bucketed graph runner collapses the whole forward to a single replay:

# Rectangular graphs over padded (B, S) buckets:
pe.pack(model, cuda_graph=True, attention_backend="sdpa")

# Packed graphs (auto/triton/flash) need explicit batch and sequence bounds:
pe.pack(model, cuda_graph=pe.GraphConfig(
    pad_to=64, max_batch=256, max_seq=128, max_graphs=32,
))

with pe.no_cuda_graph(model):          # bypass for a one-off odd shape
    out = model(input_ids=huge_ids, attention_mask=huge_mask)
pe.set_cuda_graph(model, False)        # or flip after the fact

• Graphs engage only when gradients and CUDA autocast are both disabled; out-of-bucket shapes fall back to the eager fused forward automatically.
• cuda_graph=True only captures calls with S ≤ 64 (cuda_graph_seq_cutoff), which keeps graphs only on short S regimes.
• The attention backend is resolved before capture and baked into each graph, so replay does zero dispatch work.
• Kill switch: PACKED_ENCODERS_GRAPH=0 disables all graph runners globally.

Training has its own captured forward+backward runner:

pe.pack(model, train_cuda_graph=pe.TrainGraphConfig(max_seq=64))

with two caller invariants: optimizer.zero_grad(set_to_none=False) (the backward graph writes into persistent .grad) and fixed (B, S) shapes per captured step.

Check if your encoder is supported with validate()

pack() runs validate() by default (validate=False to skip). It checks:

1. the model architecture (config.model_type, and with it the split-half RoPE convention the kernels assume);
2. the compute capability against the validated matrix: sm_80, sm_89, sm_90, sm_100, sm_120;
3. that the CuteDSL toolchain can JIT a kernel on this machine;
4. that the fused forward tracks stock HF within the bf16 band (cosine ≥ 0.997 at S = 128/512/2048).

The packed forward requires input_ids; inputs_embeds, attention/hidden-state outputs, and non-zero token_type_ids are not supported.

Fully packed execution (advanced)

pack() keeps the public padded interface. If your pipeline already collates without padding, packed_forward() is the direct entry:

import torch
from packed_encoders.config import ModernBertParams
from packed_encoders.forward import packed_forward

params = ModernBertParams.from_hf_config(model.config)
# three sequences of lengths 5, 3, 8 — concatenated, zero pad tokens
packed_ids   = torch.cat([ids_a, ids_b, ids_c])                             # [16]
cu_seqlens   = torch.tensor([0, 5, 8, 16], dtype=torch.int32, device="cuda")
position_ids = torch.cat([torch.arange(n) for n in (5, 3, 8)]).cuda()       # [16]

hidden = packed_forward(model, params, packed_ids, cu_seqlens, 8, position_ids)  # [16, H]

 

Benchmarks

All numbers come from real Tevatron/AgentIR-data queries and long documents, are PyTorch-vs-PyTorch comparisons, and report peak reserved VRAM. Every accelerated path is parity-checked against eager stock before a result is published (cosine ≥ 0.999 on these runs). Configs, runners, protocols, and remote launchers live in benchmarks/.

Inference (RTX 5090, batch 8)

Queries are capped at 128 tokens (dispatch-bound; CUDA graphs enabled on the packed paths), documents at 2048 (compute-bound; graphs deliberately off). Sequences/second, higher is better:

Endpoint	Model family	Eager HF	torch.compile	pe.pack	Fully packed
Queries (≤128)	ColBERT (PyLate)	1,370	2,614	3,149	3,572 (+160%)
Queries (≤128)	single-vector (ST)	1,441	2,833	3,507	3,602 (+150%)
Documents (≤2048)	ColBERT (PyLate)	207	281	379	388 (+87%)
Documents (≤2048)	single-vector (ST)	232	329	473	480 (+106%)

On documents the packed path is faster and leaner than eager (0.80-0.85 vs 0.82-1.04 GB reserved). On queries the graphed packed path deliberately trades memory for throughput (~1.2-1.4 GB reserved vs ~0.4-0.5 eager).

Indexing (RTX 5090, batch 128, documents ≤2048)

Model family	Eager HF	torch.compile	pe.pack	Fully packed	Peak reserved
ColBERT (PyLate)	143 docs/s	199	387	418 (2.9x)	7.85 → 5.80 GB (−26%)
single-vector (ST)	160 docs/s	243	522	527 (3.3x)	7.87 → 5.75 GB (−27%)

The fully packed late-interaction index is also 2.15x smaller on disk (60 MB vs 129 MB for the same corpus) because no rectangular token tensor or mask is stored.

Training (NVIDIA B200)

Multi-vector (late interaction)

The Agent-ModernColBERT recipe — lightonai/GTE-ModernColBERT-v1, GradCache (CachedContrastive), B=32 with 7 negatives, query/document caps 8192/4096, 326 optimizer steps over two epochs:

	Stock	pe.pack	Fully packed
Median step time	10,518 ms	1,812 ms (5.8x)	1,516 ms (7.0x)
Wall clock, including all compile warmup	65m 06s	16m 39s (3.91x)	14m 18s (4.55x)
Per-step peak reserved	26.7 GB	17.2 GB (−36%)	16.9 GB (−37%)

Single vector

gte-modernbert-base, full in-batch negatives, B=16, query/document caps 128/2048, 650 optimizer steps over two epochs:

	Stock	pe.pack	Fully packed
Median step time	1,216 ms	271 ms (4.5x)	270 ms (4.5x)
Per-step peak reserved	176.5 GB	91.8 GB (−48%)	91.2 GB (−48%)

All variants start from a bit-identical initial probe loss and follow comparable optimization trajectories under the same data and training configuration.

topk-embed-v1 (L40S and H100)

End to end model.encode (tokenisation included), items/second, bf16 pe.pack(model) vs the model as shipped in the same process. Data: 256 queries (mean 21 tokens) and 512 documents (mean 332 tokens, untruncated) sampled from BEIR SciFact, sorted by length as encode does. topk's own stack: torch 2.11, transformers 5.9, fla 0.5.1.

Endpoint	xsmall L40S	small L40S	xsmall H100	small H100
Queries, batch 8	116 → 1,386 (12.0x)	112 → 842 (7.5x)	105 → 1,679 (16.0x)	95 → 1,404 (14.8x)
Documents, batch 8	98 → 260 (2.7x)	77 → 134 (1.7x)	94 → 424 (4.5x)	84 → 298 (3.5x)
Queries, batch 128	1,370 → 2,620 (1.9x)	1,048 → 1,756 (1.7x)	1,337 → 3,056 (2.3x)	1,248 → 2,746 (2.2x)
Documents, batch 128	161 → 238 (1.5x)	102 → 142 (1.4x)	379 → 536 (1.4x)	267 → 370 (1.4x)

Mean token cosine to the model as shipped (bf16): 0.99984–0.99987. Against an fp32 forward, the per-token probe (benchmarks/qwen35_parity_probe.py) finds the packed model as close as the shipped bf16 one or closer. The shipped model is host-bound at small batches (on L40S its GPU is busy ~8% of a query batch), so most of the gain is CUDA graphs; the rewritten forward (merged GEMMs, in-kernel gates, fused norms) is what makes them cheap to capture and trims document GPU time by ~17%. The model's loader options change nothing here: it already loads bf16, and its hard-wired flex attention ignores attn_implementation; torch.compile only helped documents at batch 128 (155 → 199 on L40S xsmall, below pe.pack's 226 in the same run). benchmarks/practical_ladder.py runs this ladder; benchmarks/qwen35_topk_bench.py times the forward alone and A/Bs the engine's options.

Reproducibility

The runners and canonical configs are checked in; the figures above are snapshots from the published reference runs. See the protocols for inference, indexing, and training for the exact environment, methodology, and parity gates. Attention crossover calibration is documented in the same benchmarks/ guide.

 

Tests

uv run pytest tests/

The suite covers kernel numerics (LayerNorm, matmul, short attention — tile boundaries, raggedness, cross-sequence isolation, determinism), backend selection and dispatch-score logic, CUDA-graph bucketing for inference and training, and padded/packed forward-backward parity against stock Hugging Face.

Contributing

Contributions are welcome, especially for:

• Supporting more encoder families
• Reach a higher level of abstraction that can result in supporting a generic packed encoder runtime

Citation

@software{packedencoders,
  author = {Montserrat, Pau},
  title = {packed-encoders: a padding-free runtime for encoders},
  year = {2026},
  url = {https://github.com/pau-mensa/packed-encoders}
}

Acknowledgments

• Answer.AI / LightOn for ModernBERT, and LightOn for PyLate and the GTE-ModernColBERT models used throughout the showcase.
• IBM Research Israel for the Flash-MaxSim kernel the showcase's MaxSim scoring builds on.
• Tri Dao et al. for FlashAttention, and NVIDIA for CUTLASS / CuteDSL.

License

MIT. benchmarks/helpers/flash_maxsim_packed.py builds on flash-maxsim, which is licensed under Apache-2.0.