Qwen3 reranker computes full-sequence vocabulary logits, causing multi-GB VRAM spikes and OOMs

System Info / 系統信息

• OS: Linux
• GPU: NVIDIA RTX 5090 32 GB
• NVIDIA driver: 590.48.01
• CUDA: 13.0.2
• Model: Qwen3-Reranker-0.6B
• Engine: sentence_transformers
• Model format: pytorch
• dtype: BF16
• transformers: 5.13.1
• Registered max_tokens: 32768

Running Xinference with Docker? / 是否使用 Docker 运行 Xinfernece？

• docker / docker
• pip install / 通过 pip install 安装
• installation from source / 从源码安装

Version info / 版本信息

3.4.1.dev0+g99868ea70

Docker image: xprobe/xinference:latest

The command used to start Xinference / 用以启动 xinference 的命令

docker run -d --name xinference --restart unless-stopped
--gpus all
-e XINFERENCE_HOME=/data
-v xinference_data:/data
-p 127.0.0.1:9997:9997
xinference-local:latest
xinference-local -H 0.0.0.0

Reproduction / 复现过程

The Qwen3 reranker computes vocabulary logits for every sequence position although only the final position is used.

In xinference/model/rerank/sentence_transformers/core.py, the current implementation is essentially:

@torch.inference_mode()
def compute_logits(inputs, **kwargs):
    batch_scores = model(**inputs).logits[:, -1, :]

Qwen3ForCausalLM.forward() in Transformers 5.13.1 supports logits_to_keep, which slices the hidden states before lm_head:

slice_indices = slice(-logits_to_keep, None)
logits = self.lm_head(hidden_states[:, slice_indices, :])

With the current Xinference code, [:, -1, :] is applied only after the full [batch, seq_len, vocab_size] logits tensor has already been materialized.

For Qwen3-Reranker-0.6B, vocab_size = 151669.

I reproduced the issue with a micro-batch of 4 and a sequence length of 8265 tokens.

The returned reranking scores were bit-identical.

The full logits tensor alone accounts for approximately:

4 × 8265 × 151669 × 2 bytes ≈ 9.34 GiB

This matches the failed CUDA allocation observed during the incident:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 9.34 GiB.
GPU 0 has a total capacity of 31.35 GiB of which 1.36 GiB is free.

ERROR xinference.core.model
Exiting model subprocess via os._exit(1) to trigger pool recovery

The model subprocess exits and clients temporarily receive HTTP 500 responses while the model is being recovered.

A minimal change that avoids the full-sequence vocabulary projection is:

batch_scores = model(
    **inputs,
    logits_to_keep=1,
    use_cache=False,
).logits[:, -1, :]

logits_to_keep=1 is the important part of the change: the reranker only needs the final token logits, so projecting all sequence positions through the 151669-entry vocabulary is unnecessary.

use_cache=False is a secondary optimization because reranking performs one forward pass and does not reuse the KV cache.

Related issue / PR

This appears related to:

• #3735 — Qwen3-Reranker consumes too much VRAM https://github.com/xorbitsai/inference/issues/3735

• #3666 — ENH: Add 4-sample micro-batching to Qwen-3 reranker to reduce GPU memory https://github.com/xorbitsai/inference/pull/3666

PR #3666 introduced a micro-batch size of 4 to bound memory usage for long contexts. That optimization is useful but addresses a different dimension: it bounds the number of pairs processed concurrently, while the logits tensor still scales with seq_len × vocab_size inside each micro-batch.

Therefore, keeping the micro-batching and additionally using logits_to_keep=1 should be complementary.

The upstream Qwen reranker example also uses the model(**inputs).logits[:, -1, :] pattern, but its example uses max_length = 8192, whereas Xinference registers this model with max_tokens = 32768.

Expected behavior / 期待表现

The Qwen3 reranker should not materialize vocabulary logits for sequence positions whose values are never used.

Since reranking only consumes the logits from the final token, Xinference should request only that position from Qwen3ForCausalLM, for example with logits_to_keep=1.

For the reproduced input, this reduces the logits shape from:

(4, 8265, 151669)

to:

(4, 1, 151669)

while producing identical reranking scores and avoiding the multi-GB temporary allocation.

The existing micro-batching introduced in #3666 should remain, since it bounds transformer activation memory independently of this optimization.