[None][feat] add native Qwen3-Reranker support to the PyTorch backend

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

@@ -24,6 +24,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 cross-encoder reranking models (currently the Qwen3-Reranker family), the ``trtllm-serve rerank`` subcommand starts a server that exposes ``/rerank``, ``/v1/rerank``, and the Cohere-compatible ``/v2/rerank`` endpoints with native dynamic batching. See :doc:`Reranking <../../features/reranking>` for details.
+
 Starting a Server
 -----------------
 

docs/source/features/reranking.md

@@ -0,0 +1,115 @@
+# Reranking (Cross-Encoder Models)
+
+`trtllm-serve` can serve **cross-encoder reranking models** through OpenAI/Cohere-style
+**`POST /rerank`**, **`POST /v1/rerank`**, and **`POST /v2/rerank`** endpoints with
+**native dynamic batching** — coalescing many independent concurrent (query, document)
+pairs into a single forward pass for high throughput, the same way
+[embeddings serving](embeddings.md) does.
+
+Currently supported: the **Qwen3-Reranker family** (`Qwen3-Reranker-0.6B`, `-4B`,
+`-8B`). These ship as a `Qwen3ForCausalLM` decoder scored as a cross-encoder: the
+model judges whether a document is relevant to a query and answers `"yes"` or `"no"`;
+the reranking score is `yes_logit - no_logit` at the last token, read directly out of
+the model's own `lm_head` (no decode loop, sampling, or detokenization needed).
+
+## Quick start
+
+Launch a rerank server with the `rerank` subcommand:
+
+```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
+```
+
+Send a request with `curl` (any of `/rerank`, `/v1/rerank`, `/v2/rerank` works):
+
+```bash
+curl http://localhost:8000/v1/rerank \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "<model>",
+    "query": "What is the capital of France?",
+    "documents": [
+      "Paris is the capital and most populous city of France.",
+      "Berlin is the capital of Germany."
+    ]
+  }'
+```
+
+The response is Cohere's `/v2/rerank` shape:
+
+```json
+{
+  "id": "rerank-...",
+  "model": "<model>",
+  "results": [
+    {"index": 0, "relevance_score": 8.4},
+    {"index": 1, "relevance_score": -6.1}
+  ],
+  "usage": {"total_tokens": 96}
+}
+```
+
+`results` is sorted by descending `relevance_score`.
+
+## Request fields
+
+| Field | Type | Notes |
+|---|---|---|
+| `model` | str | Model name. |
+| `query` | str | The search query. |
+| `documents` | list[str \| {"text": str}] | Candidate documents, as plain strings or Cohere-style `{"text": ...}` objects. |
+| `top_n` | int (optional) | Return only the top-`n` scoring results. Defaults to all documents. |
+| `return_documents` | bool (default `false`) | Include each document's text in the response results. |
+| `instruction` | str (optional) | TRT-LLM extension. Overrides the Qwen3-Reranker task instruction embedded in the scoring prompt (default: a generic web-search-retrieval instruction). |
+
+## Dynamic batching
+
+Rerank serving reuses the same in-server dynamic batcher as [embeddings
+serving](embeddings.md#dynamic-batching), with identical `--max_batch_size` /
+`--max_queue_delay` / `--max_queue_size` flags and Triton `dynamic_batching`
+migration mapping. Each `(query, document)` pair is tokenized into one prompt and
+coalesced into the batch independently — a single request with many `documents`
+still benefits from batching alongside other concurrent requests.
+
+## Prompt construction and truncation
+
+Each `(query, document)` pair is scored using the official Qwen3-Reranker prompt: a
+fixed system instruction, an `<Instruct>/<Query>/<Document>` user turn, and an
+assistant preamble ending right where the model emits its `yes`/`no` answer. Only the
+**document** is truncated to fit the model's `max_seq_len`; the query and instruction
+are always kept intact, so a truncated request still asks the same question, just
+against less of the candidate document. If the instruction and query alone exceed
+`max_seq_len`, the request fails with `400` rather than silently dropping part of the
+query.
+
+## Error handling
+
+| Condition | HTTP status |
+|---|---|
+| `documents` is empty | 400 |
+| Instruction + query alone exceed `max_seq_len` | 400 |
+| Request queue full (`--max_queue_size` reached) | 429 |
+| Invalid request body | 400 |
+
+Rerank responses are unary (non-streaming).
+
+Notes:
+
+- The rerank path uses the synchronous `llm.encode()` fast path (`EncoderExecutor`): a
+  single forward pass per batch, **no KV cache, sampler, or decode loop**.
+- **Single-GPU per server.** Like the `embeddings` command, the encode path runs
+  in-process and does not use the multi-GPU worker proxy, so `rerank` does not expose
+  tensor/pipeline parallelism flags. To scale out, follow the same [data-parallel
+  replica pattern](embeddings.md#scaling-out-across-gpus) as embeddings serving — run
+  one single-GPU `trtllm-serve rerank` instance per GPU behind a load balancer.
+
+## Relationship to `llm.encode()`
+
+The server reuses the existing Python `llm.encode()` API (`LLM(..., encode_only=True)`)
+under the hood, with the model's `architectures` and the tokenizer-resolved
+`reranking_token_true_id`/`reranking_token_false_id` overridden via `model_kwargs`; the
+only addition is the async coalescing layer, prompt construction, and the HTTP surface.

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-cache-compression.md
    features/long-sequence.md

docs/source/models/supported-models.md

@@ -55,6 +55,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

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

tensorrt_llm/_torch/models/modeling_qwen3.py

@@ -363,3 +363,68 @@ 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 (0.6B / 4B / 8B).
+
+    The Qwen3 decoder backbone used as a cross-encoder reranker: `forward`
+    projects only the last token of each (query, document) sequence through
+    the ordinary lm_head, then returns the scalar `yes_logit - no_logit`, per
+    the official Qwen3-Reranker scoring convention. This keeps the single
+    causal-LM forward pass and last-token lm_head projection that a normal
+    context phase already does, but skips the decode loop, sampling, and
+    detokenization a generation-path score would otherwise need.
+
+    `reranking_token_true_id`/`reranking_token_false_id` (the tokenizer's
+    "yes"/"no" token ids) must be present on `pretrained_config`; the
+    `trtllm-serve rerank` command resolves them from the model's tokenizer
+    and injects them via `model_kwargs` before this class is constructed
+    (see `commands/serve.py`).
+    """
+
+    def __init__(self, model_config: ModelConfig[Qwen3Config]):
+        super().__init__(
+            Qwen3Model(model_config),
+            config=model_config,
+            hidden_size=model_config.pretrained_config.hidden_size,
+            vocab_size=model_config.pretrained_config.vocab_size,
+        )
+
+        pretrained_config = model_config.pretrained_config
+        token_true_id = getattr(pretrained_config, "reranking_token_true_id",
+                                None)
+        token_false_id = getattr(pretrained_config,
+                                 "reranking_token_false_id", None)
+        if token_true_id is None or token_false_id is None:
+            raise ValueError(
+                "Qwen3ForTextReranking requires reranking_token_true_id and "
+                "reranking_token_false_id on pretrained_config (the "
+                "tokenizer's \"yes\"/\"no\" token ids). Construct this model "
+                "through `trtllm-serve rerank`, which resolves and injects "
+                "them via model_kwargs.")
+        self.token_true_id = token_true_id
+        self.token_false_id = token_false_id
+
+    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)
+
+        # Last-token-only vocab projection (the same context-logits gather the
+        # generation path uses for its final token), then read out just the
+        # two rows this score needs.
+        logits = self.logits_processor(hidden_states,
+                                       self.lm_head,
+                                       attn_metadata,
+                                       return_context_logits=False)
+        return logits[:, self.token_true_id] - logits[:, self.token_false_id]