[#19510][feat] Add native Qwen3 reranking support

docs/source/commands/trtllm-serve/trtllm-serve.rst

@@ -28,6 +28,8 @@ The ``metrics`` endpoint provides runtime-iteration statistics such as GPU memor
 
 For encoder-only models (BERT-style classifiers, reward models, text-embedding models), the ``trtllm-serve embeddings`` subcommand starts a server that exposes an OpenAI-compatible ``/v1/embeddings`` endpoint with native dynamic batching. See :doc:`Embeddings <../../features/embeddings>` for details.
 
+For reranker models, the ``trtllm-serve rerank`` subcommand exposes ``/rerank``, ``/v1/rerank``, and Cohere-compatible ``/v2/rerank`` endpoints. See :doc:`Reranking <../../features/reranking>` for details.
+
 Starting a Server
 -----------------
 

docs/source/features/reranking.md

@@ -0,0 +1,186 @@
+# Reranking
+
+`trtllm-serve rerank` serves text reranking models through HTTP endpoints with
+native dynamic batching. For each request, the server pairs the query with every
+candidate document, evaluates the pairs in the encode-only path, and returns
+the documents ordered by relevance.
+
+The encode-only path performs one forward pass per batch, with no KV cache,
+sampler, or decode loop. Concurrent query-document pairs are coalesced into
+larger GPU batches.
+
+## Quick start
+
+Launch a server with a [supported reranker checkpoint](#supported-model-architectures):
+
+```bash
+trtllm-serve rerank <hf_model_or_path> \
+    --max_batch_size 32 \
+    --max_queue_delay 0.005 \
+    --max_queue_size 2048 \
+    --host 0.0.0.0 --port 8000
+```
+
+For example, the Qwen3-Reranker family is supported:
+
+```bash
+trtllm-serve rerank Qwen/Qwen3-Reranker-0.6B \
+    --max_batch_size 32 \
+    --max_queue_delay 0.005 \
+    --host 0.0.0.0 --port 8000
+```
+
+The server exposes the following routes:
+
+| Route | Compatibility |
+|---|---|
+| `POST /rerank` | Jina/vLLM-style request and response |
+| `POST /v1/rerank` | Alias of `/rerank` |
+| `POST /v2/rerank` | Cohere v2-compatible request and response |
+
+Send a request with `curl`:
+
+```bash
+curl http://localhost:8000/v1/rerank \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-Reranker-0.6B",
+    "query": "What is the capital of China?",
+    "documents": [
+      "The capital of China is Beijing.",
+      "Bananas grow in tropical climates."
+    ],
+    "top_n": 1,
+    "return_documents": true
+  }'
+```
+
+Results are ordered by descending relevance score. `index` always refers to
+the document's position in the original request.
+
+```json
+{
+  "id": "rerank-...",
+  "model": "Qwen/Qwen3-Reranker-0.6B",
+  "results": [
+    {
+      "index": 0,
+      "relevance_score": 0.98,
+      "document": {"text": "The capital of China is Beijing."}
+    }
+  ],
+  "usage": {"prompt_tokens": 91, "total_tokens": 91}
+}
+```
+
+## Request fields
+
+The two API styles share the query, documents, and result-ordering semantics,
+but differ in a few compatibility fields:
+
+| Field | `/rerank` and `/v1/rerank` | `/v2/rerank` |
+|---|---|---|
+| `model` | Optional string | Required string |
+| `query` | Required non-empty string | Required non-empty string |
+| `documents` | Non-empty list of strings or `{"text": "..."}` objects | Non-empty list of strings |
+| `top_n` | Optional non-negative integer; `0` returns all results | Optional positive integer |
+| `return_documents` | Optional boolean; defaults to `false` | Not part of the v2 response |
+| `max_tokens_per_doc` | Optional positive integer | Positive integer; defaults to `4096` |
+| `priority` | Not supported | Only `0` is currently supported |
+| `instruction` | Optional TensorRT-LLM extension | Optional TensorRT-LLM extension |
+
+`max_tokens_per_doc` limits document content, not the complete model input.
+The server preserves the query and any model-specific prompt suffix while
+truncating a document.
+
+## Dynamic batching
+
+The in-server batcher coalesces query-document pairs from concurrent HTTP
+requests before running the encode-only forward pass.
+
+| `trtllm-serve rerank` flag | Behavior |
+|---|---|
+| `--max_batch_size` | Maximum number of query-document pairs in one forward pass. A full batch is dispatched immediately. |
+| `--max_num_tokens` | Maximum total number of input tokens in one batch. |
+| `--max_queue_delay` | Maximum time, in seconds, that a pair waits for other work to join its batch. |
+| `--max_queue_size` | Maximum number of queued pairs. Additional requests are rejected with HTTP 429. |
+
+A batch is dispatched when it reaches `--max_batch_size`, adding another pair
+would exceed `--max_num_tokens`, or the queue-delay window expires. Tune batch
+size and delay together for the desired latency-throughput tradeoff.
+
+## Scoring contract
+
+The reranking service expects the model implementation to emit one scalar,
+pre-sigmoid relevance logit for each query-document pair. The server applies a
+numerically stable sigmoid, sorts the resulting probabilities, and selects
+`top_n` when requested.
+
+The HTTP and batching layers are shared infrastructure. Supporting another
+reranker family also requires model-specific input formatting and a TensorRT-LLM
+architecture that satisfies this scalar-output contract.
+
+## Supported model architectures
+
+The HTTP surface is designed for reranking generally. The model-specific
+preprocessing and scoring implementation currently supports the
+**Qwen3-Reranker family**:
+
+- `Qwen3-Reranker-0.6B`
+- `Qwen3-Reranker-4B`
+- `Qwen3-Reranker-8B`
+
+Only unquantized checkpoints are supported. Quantized checkpoints, including
+FP8 and NVFP4, are not supported: the scalar scoring head is derived directly
+from unquantized language-model or tied embedding weights.
+
+Qwen3-Reranker checkpoints are published as causal language models that score
+the tokens `yes` and `no`. TensorRT-LLM selects `Qwen3ForTextReranking` for
+these checkpoints, reads the token IDs from `1_LogitScore/config.json` when
+present, and derives a one-row scoring head at load time:
+
+```text
+W_score = W_yes - W_no
+relevance_logit = hidden_state_last @ W_score
+relevance_score = sigmoid(relevance_logit)
+```
+
+This is mathematically equivalent to the checkpoint's two-token softmax while
+avoiding a projection over the full vocabulary. The Qwen3 input formatter uses
+the model's official query-document prompt. The optional `instruction` field
+replaces its default web-search instruction.
+
+## Error handling
+
+| Condition | HTTP status |
+|---|---|
+| Invalid request body or unsupported `priority` | 400 |
+| Instruction and query alone exceed the model or batch token limit | 400 |
+| Document count exceeds a positive `--max_queue_size` | 400 |
+| Request queue full | 429 |
+| Model does not return one scalar per pair | 500 |
+
+Requests exceeding a positive `--max_queue_size` document count are rejected
+before tokenization. Split these into smaller requests; retrying the same
+oversized request will not help. This limit applies to all submitted documents,
+even when `top_n` requests fewer results.
+
+Over-long documents are truncated to fit the smaller of the model's sequence
+limit and the batch token limit, while preserving the instruction, query, and
+assistant suffix. Reranking responses are unary and non-streaming.
+
+## Deployment scope
+
+- One reranker model is served by each server instance. Generation, embeddings,
+  and reranking modes are not mixed in one instance.
+- The reranking command is single-GPU. Tensor, pipeline, and context parallel
+  values greater than one in `--config` are rejected at startup.
+- A single worker drives the GPU because the underlying encode executor is not
+  safe for concurrent calls. Increase throughput with dynamic batching rather
+  than additional in-process workers.
+- To scale across GPUs, run one server replica per GPU and distribute requests
+  across the replicas with a load balancer.
+
+The server uses the same synchronous `llm.encode()` path as other encode-only
+models, with an asynchronous batching layer and reranking HTTP APIs added on
+top.

docs/source/index.rst

@@ -70,6 +70,7 @@ Welcome to TensorRT LLM's Documentation!
    features/disagg-serving.md
    features/subagent-routing.md
    features/embeddings.md
+   features/reranking.md
    features/kvcache.md
    features/kv-hints.md
    features/kv-cache-compression.md

docs/source/models/supported-models.md

@@ -50,6 +50,7 @@ The following is a table of supported models for the PyTorch backend:
 | `Qwen2ForRewardModel`                | Qwen2-based                        | `Qwen/Qwen2.5-Math-RM-72B`                   |
 | `Qwen3ForCausalLM`                   | Qwen3                              | `Qwen/Qwen3-8B`                              |
 | `Qwen3ForTextEmbedding`              | Qwen3-Embedding                    | `Qwen/Qwen3-Embedding-8B`                    |
+| `Qwen3ForTextReranking`              | Qwen3-Reranker                     | `Qwen/Qwen3-Reranker-8B`                     |
 | `Qwen3MoeForCausalLM`                | Qwen3MoE                           | `Qwen/Qwen3-30B-A3B`                         |
 | `Qwen3NextForCausalLM`               | Qwen3Next                          | `Qwen/Qwen3-Next-80B-A3B-Thinking`           |
 | `Qwen3_5MoeForCausalLM`              | Qwen3.8-MoE, Qwen3.5-MoE           | `Qwen/Qwen3.8-2.4T-A95B`, `Qwen/Qwen3.5-397B-A17B` |

tensorrt_llm/_torch/model_config.py

@@ -530,7 +530,8 @@ def is_generation_model(model_architectures: Optional[List[str]],
         return model_architectures[0] not in [
             "BertForSequenceClassification", "Qwen2ForProcessRewardModel",
             "Qwen2ForRewardModel", "LlamaForTextEmbedding",
-            "PplxDeciderForDecision", "Qwen3ForTextEmbedding"
+            "PplxDeciderForDecision", "Qwen3ForTextEmbedding",
+            "Qwen3ForTextReranking"
         ]
         # TODO: should be 'not model_type == ModelType.ENCODER_ONLY'
         # once ModelType is used in pytorch flow.

tensorrt_llm/_torch/models/_arch_index.py

@@ -113,6 +113,7 @@ def is_builtin_zoo_module(module_name: str) -> bool:
     "Qwen2_5_VLForConditionalGeneration": "modeling_qwen2vl",
     "Qwen3ForCausalLM": "modeling_qwen3",
     "Qwen3ForTextEmbedding": "modeling_qwen3",
+    "Qwen3ForTextReranking": "modeling_qwen3",
     "Qwen3MoeForCausalLM": "modeling_qwen3_moe",
     "Qwen3NextForCausalLM": "modeling_qwen3_next",
     "Qwen3VLForConditionalGeneration": "modeling_qwen3vl",

tensorrt_llm/_torch/models/modeling_qwen3.py

@@ -1,4 +1,9 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import copy
+import json
+from pathlib import Path
 from typing import Optional
 
 import torch
@@ -18,7 +23,7 @@
 from ..modules.decoder_layer import DecoderLayer
 from ..modules.embedding import Embedding
 from ..modules.gated_mlp import GatedMLP
-from ..modules.linear import TensorParallelMode
+from ..modules.linear import Linear, TensorParallelMode
 from ..modules.rms_norm import RMSNorm
 from ..speculative import SpecMetadata
 from .modeling_speculative import SpecDecOneEngineForCausalLM
@@ -378,3 +383,118 @@ def forward(self,
         pooled = hidden_states[end_indices]
         normalized = torch.nn.functional.normalize(pooled.float(), p=2, dim=-1)
         return normalized.to(pooled.dtype)
+
+
+@register_auto_model("Qwen3ForTextReranking")
+class Qwen3ForTextReranking(DecoderModelForCausalLM[Qwen3Model, Qwen3Config]):
+    """Qwen3-Reranker family exposed as a scalar text-reranking model.
+
+    The published checkpoints are causal language models whose relevance
+    probability is the two-token softmax over ``yes`` and ``no``. A one-row
+    head with ``W_yes - W_no`` produces the equivalent pre-sigmoid logit while
+    avoiding a full vocabulary projection during inference.
+    """
+
+    # Published checkpoints record these values in 1_LogitScore/config.json.
+    _YES_TOKEN_ID = 9693
+    _NO_TOKEN_ID = 2152
+
+    def __init__(self, model_config: ModelConfig[Qwen3Config]) -> None:
+        nn.Module.__init__(self)
+        self.model_config = model_config
+
+        config = model_config.pretrained_config
+        self._yes_token_id, self._no_token_id = self._load_score_token_ids(
+            config)
+        self.model = Qwen3Model(model_config)
+        self.score = Linear(config.hidden_size,
+                            1,
+                            bias=False,
+                            dtype=config.torch_dtype)
+        self.prologue = []
+        self.epilogue = [self.score]
+
+    @classmethod
+    def _load_score_token_ids(cls, config: Qwen3Config) -> tuple[int, int]:
+        checkpoint_dir = getattr(config, "_name_or_path", None)
+        if not checkpoint_dir:
+            return cls._YES_TOKEN_ID, cls._NO_TOKEN_ID
+
+        score_config_path = (Path(checkpoint_dir) / "1_LogitScore" /
+                             "config.json")
+        try:
+            with score_config_path.open(encoding="utf-8") as config_file:
+                score_config = json.load(config_file)
+        except FileNotFoundError:
+            return cls._YES_TOKEN_ID, cls._NO_TOKEN_ID
+        except (OSError, json.JSONDecodeError) as error:
+            raise ValueError(f"Could not read Qwen3 reranker score config at "
+                             f"{score_config_path}: {error}") from error
+
+        try:
+            yes_token_id = score_config["true_token_id"]
+            no_token_id = score_config["false_token_id"]
+        except (KeyError, TypeError) as error:
+            raise ValueError(
+                f"Qwen3 reranker score config at {score_config_path} must "
+                "define true_token_id and false_token_id.") from error
+
+        for field_name, token_id in (("true_token_id", yes_token_id),
+                                     ("false_token_id", no_token_id)):
+            if (not isinstance(token_id, int) or isinstance(token_id, bool)
+                    or token_id < 0):
+                raise ValueError(
+                    f"{field_name} in {score_config_path} must be a "
+                    f"non-negative integer; got {token_id!r}.")
+
+        return yes_token_id, no_token_id
+
+    @staticmethod
+    def _derive_score_weight(
+            weights: dict,
+            tie_word_embeddings: bool,
+            yes_token_id: int = _YES_TOKEN_ID,
+            no_token_id: int = _NO_TOKEN_ID) -> tuple[str, torch.Tensor]:
+        """Build the scalar head from the checkpoint's yes/no token rows."""
+        source_names = ["lm_head.weight"]
+        if tie_word_embeddings:
+            source_names.append("model.embed_tokens.weight")
+        for source_name in source_names:
+            if source_name in weights:
+                source = weights[source_name]
+                yes = source[yes_token_id:yes_token_id + 1]
+                no = source[no_token_id:no_token_id + 1]
+                return source_name, yes - no
+        raise ValueError(
+            "Qwen3 reranker checkpoint must contain lm_head.weight unless "
+            "tie_word_embeddings is enabled and model.embed_tokens.weight "
+            "is present.")
+
+    def load_weights(self, weights: dict, **kwargs) -> None:
+        source_name, score_weight = self._derive_score_weight(
+            weights, self.model_config.pretrained_config.tie_word_embeddings,
+            self._yes_token_id, self._no_token_id)
+        weights["score.weight"] = score_weight
+        # An untied LM head is not otherwise consumed by this model. Releasing
+        # it here keeps peak host memory bounded for the 8B checkpoint.
+        if (source_name == "lm_head.weight"
+                and hasattr(weights, "mark_consumed")):
+            weights.mark_consumed("lm_head")
+        elif source_name == "lm_head.weight":
+            del weights[source_name]
+        super().load_weights(weights, **kwargs)
+
+    def forward(self,
+                attn_metadata: AttentionMetadata,
+                input_ids: torch.IntTensor,
+                position_ids: Optional[torch.IntTensor] = None,
+                inputs_embeds: Optional[torch.FloatTensor] = None,
+                **kwargs) -> torch.Tensor:
+        assert attn_metadata.seq_lens is not None
+
+        hidden_states = self.model(attn_metadata,
+                                   input_ids,
+                                   position_ids=position_ids,
+                                   inputs_embeds=inputs_embeds)
+        end_indices = torch.cumsum(attn_metadata.seq_lens, dim=0) - 1
+        return self.score(hidden_states[end_indices])