Speed up ModernBERT inference by 1.15x on CUDA

[hotchpotch]

What does this PR do?

This PR speeds up CUDA ModernBERT inference by splitting the GELU MLP input
projection so that the CUDA linear path can use cuBLASLt's fused GELU epilogue.

Method

ModernBERT's GELU MLP computes the following operation:

GELU(x @ Wi[:intermediate_size]) * (x @ Wi[intermediate_size:])

The checkpoint stores both projections in one Wi matrix. For CUDA devices
with GELU activation, this change loads the two halves as separate linear
layers:

1. The first half is a linear layer with the GELU activation, allowing the
   CUDA implementation to fuse GELU into the cuBLASLt epilogue.
2. The second half is loaded as the gate projection without an activation.
3. The two outputs are multiplied before the existing Wo projection.

The transformation preserves the model computation while improving the
matmul/activation execution for the CUDA workload measured here.

Compatibility and scope

• The split is enabled only for CUDA devices and GELU activation.
• CPU, Metal, and non-GELU ModernBERT paths keep the original full-width
  projection and activation flow.
• Model weights and the attention, normalization, pooling, and API behavior
  are unchanged.
• The implementation is limited to
  backends/candle/src/models/modernbert.rs; no new dependency is required.

Validation

• cargo test -p text-embeddings-backend-candle --lib modernbert
• Added coverage that CPU keeps the unsplit projection.
• Added a device-independent test comparing the split and unsplit MLP
  computations.
• On an NVIDIA GeForce RTX 5090, a controlled Alibaba-NLP/gte-modernbert-base
  benchmark using 20,000 Natural Questions examples improved throughput by
  1.15x over main (fc071b1cb6e1b091b67f20868de7c5982aa7d4d0).
• Under the same setup, lightonai/modernbert-embed-large with mean pooling
  improved throughput by 1.17x over the same main baseline.

Detailed benchmark methodology, environment, results, and numerical comparison
are provided in the first PR comment to keep this description focused.

Before submitting

• Did you read the contributor guideline?
• Was this discussed/approved via a GitHub issue or the forum?
• Did you make sure to update the documentation with your changes? No user-facing documentation changes are required for this internal CUDA optimization.
• Did you write any new necessary tests?

Who can review?

@Narsil @alvarobartt

[hotchpotch]

Benchmark details

This comment contains the benchmark methodology and full measurement details
for the result summarized in the PR description.

Headline result (Alibaba-NLP/gte-modernbert-base): 3,318.7 → 3,814.9
texts/s median (1.15x, +14.95%).

Environment and measurement procedure

Environment

Item	Value
GPU	NVIDIA GeForce RTX 5090, physical GPU 1
Compute capability	12.0 (sm_120)
NVIDIA driver	580.126.09
Driver-reported CUDA version	13.0 (nvidia-smi)
CUDA toolkit on the benchmark host	12.8.1
CUDA build target	CUDA_COMPUTE_CAP=120
Rust	1.92.0
Benchmark client	Python 3.12.12, datasets==5.0.0, numpy==2.5.0, requests==2.34.2
Model	Alibaba-NLP/gte-modernbert-base
Model storage	Same local snapshot for both revisions
Dtype / pooling	float16 / cls

GPU 1 was selected explicitly with CUDA_VISIBLE_DEVICES=1. The monitor
verified that it had no compute processes and at most 500 MiB allocated before
starting each revision; the inference server was the only compute process on
that GPU during the measurements.

Workload and procedure

• Dataset: sentence-transformers/natural-questions
• Configuration: pair
• Split: train[:20000]
• Text column: answer
• Client batch size: 1024
• Client workers: 16
• Warm-up: 4 batches, excluded from timing
• Endpoint: POST /embed
• Request options: normalize=true, truncate=true
• Server max_batch_tokens: 16,384
• Server max_client_batch_size: 1,024
• Server max_concurrent_requests: 16,384
• Three timed passes per revision; the median is reported below

The baseline is the exact repository main commit
fc071b1cb6e1b091b67f20868de7c5982aa7d4d0. The optimized binary was built
from the PR working tree. Both revisions used the same model snapshot, GPU,
dtype, server settings, request batches, and dataset examples.

Throughput

Revision	Run 1 (texts/s)	Run 2 (texts/s)	Run 3 (texts/s)	Median (texts/s)	Median elapsed
main (fc071b1)	3,354.5	3,318.7	3,292.3	3,318.7	6.026 s
Optimized	3,814.9	3,852.7	3,668.7	3,814.9	5.243 s

The optimized implementation is 1.15x faster than main on this workload
(+14.95% median throughput). The three runs show normal scheduling noise,
which is why the median rather than a single best run is used.

Output comparison

The optimized responses had shape (20000, 768). The first optimized pass was
compared with embeddings saved from the first baseline pass:

Check	Value
Minimum cosine similarity	0.999958217
Mean cosine similarity	0.999998808
Maximum absolute difference	0.001933098
Mean absolute difference	0.000042098
Maximum norm error from 1	1.19e-7

The outputs are not bit-identical because the cuBLASLt GELU epilogue changes
the floating-point operation order. The embeddings remain effectively
equivalent and normalized; attention, pooling, and model weights are not
changed.

Reproduction commands

Reproduction

Download the model once and build the CUDA binary with the target architecture:

hf download Alibaba-NLP/gte-modernbert-base --local-dir /tmp/gte-modernbert-base
CUDA_COMPUTE_CAP=120 cargo build --release --bin text-embeddings-router -F candle-cuda

Start the server on the selected GPU:

CUDA_VISIBLE_DEVICES=1 \
  target/release/text-embeddings-router \
  --model-id /tmp/gte-modernbert-base \
  --dtype float16 \
  --pooling cls \
  --port 18080 \
  --prometheus-port 19090 \
  --max-batch-tokens 16384 \
  --max-client-batch-size 1024 \
  --max-concurrent-requests 16384 \
  --disable-spans

Run the client with the same options used for the measurement:

python benchmark_gte_modernbert_nq.py \
  --url http://127.0.0.1:18080/embed \
  --split 'train[:20000]' \
  --batch-size 1024 \
  --workers 16 \
  --warmup-batches 4 \
  --label optimized \
  --save-embeddings /tmp/optimized-embeddings.npy

The Python process is only the HTTP benchmark client; model inference runs in
the Rust/Candle server. The exact client used for this measurement is available
as the public gist
benchmark_gte_modernbert_nq.py.

For the baseline, run the same server and client commands from the main
commit, save its embeddings, and pass the saved file to the optimized run with
--compare-to. The benchmark client loads the dataset, sends all requests,
excludes the warm-up batches, reports throughput, and computes the cosine and
absolute-difference metrics shown above.

Additional model result (lightonai/modernbert-embed-large): 1,176.9 →
1,379.2 texts/s median (1.17x, +17.18%). Detailed results:

Additional model validation details

Additional model validation: lightonai/modernbert-embed-large

The same procedure was also run with lightonai/modernbert-embed-large. This
model uses mean pooling and produces 1,024-dimensional embeddings, so the
server was started with --pooling mean; all other benchmark settings were
unchanged.

The model can be prepared with:

hf download lightonai/modernbert-embed-large \
  --local-dir /tmp/lightonai-modernbert-embed-large \
  --exclude 'onnx/*'

For this validation, replace the GTE model path in the server command with
/tmp/lightonai-modernbert-embed-large and use --pooling mean.

Revision	Run 1 (texts/s)	Run 2 (texts/s)	Run 3 (texts/s)	Median (texts/s)	Median elapsed
main (fc071b1)	1,192.3	1,176.9	1,169.0	1,176.9	16.993 s
Optimized	1,382.8	1,379.2	1,372.4	1,379.2	14.501 s

The optimized implementation is 1.17x faster on this model (+17.18%
median throughput). The output shape was (20000, 1024).

Check	Value
Minimum cosine similarity	0.995162606
1st percentile cosine similarity	0.999869466
Median cosine similarity	0.999999225
Mean cosine similarity	0.999998271
Samples below 0.999 cosine similarity	2 / 20,000
Samples below 0.99 cosine similarity	0 / 20,000
Maximum absolute difference	0.016730726
Mean absolute difference	0.000033006
Maximum norm error from 1	1.19e-7

As with gte-modernbert-base, the outputs are not bit-identical because of
the cuBLASLt GELU epilogue's floating-point operation order, but the outputs
remain normalized and effectively equivalent for this validation.

[alvarobartt]

Awesome thanks a lot @hotchpotch! Before we merge, do you mind moving the tests into backends/candle/tests/test_modernbert.rs?

[Copilot]

Pull request overview

This PR optimizes ModernBERT’s CUDA GELU MLP path by splitting the combined Wi projection into two linear layers so the CUDA linear implementation can use cuBLASLt’s fused GELU epilogue, improving inference throughput while keeping non-CUDA and non-GELU behavior unchanged.

Changes:

• Split Wi into an activated GELU projection (wi) plus a separate gate projection (gate) on CUDA+GELU, then multiply and apply Wo.
• Preserve the original unsplit projection + separate activation flow for CPU/Metal and non-GELU activations.
• Add unit tests to ensure CPU keeps the unsplit path and to validate split vs. unsplit MLP equivalence.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[hotchpotch]

Good catch. Addressed in 5f76b6c by requiring
get_cublas_lt_wrapper().is_some() in should_split_wi, so the split is only
selected when the fused GELU path is available.

I also rebuilt with CUDA and reran the controlled GTE ModernBERT benchmark:
the reviewed build remained 1.1508x (+15.08%) faster than main by median
throughput, with the same embedding-comparison metrics as before.

[hotchpotch]

Thanks, @alvarobartt! I moved the integration-friendly MLP equivalence coverage to
backends/candle/tests/test_modernbert.rs.

One note: the source-local test module also contained two tests that directly
exercise private ModernBertMLP implementation details:

• cpu_modernbert_uses_unsplit_gelu_projection, which checks that CPU loading
  keeps gate unset and retains the separate GELU activation.
• modernbert_split_mlp_matches_unsplit_mlp, which constructs the split and
  unsplit private MLP representations and compares their forward outputs.

These tests could not be moved verbatim because integration tests are compiled
as a separate crate and cannot access those private fields and types. I
therefore removed the direct implementation-level tests for now, they are in
this removed test module.

If retaining that private implementation coverage is important, I can either
keep those tests next to the implementation, or introduce a narrowly scoped
public test-support API so they can live entirely in
backends/candle/tests/test_modernbert.rs. Please let me know which approach
you would prefer.

[alvarobartt]

Oh ACK @hotchpotch then feel free to add those back for ModernBERT and we'll then merge, apologies for the confusion, and thanks again for your contribution 🤗

[hotchpotch]

@alvarobartt
I've restored the tests in modernbert.rs. Thank you for the review!