[None][feat] Add llm.encode() fast path for encoder-only models

[tingyangk]

@coderabbitai summary

Summary

Adds a dedicated llm.encode() API for encoder-only models (BERT, RoBERTa, reward models) that bypasses the decoder-oriented PyExecutor loop entirely.

Problem

The current LLM API routes encoder models through the same PyExecutor designed for autoregressive decoders, introducing siginificant CPU overhead per batch from scheduler, KV cache management, sampling, and request state machine — none of which apply to encoders. Encoder models need a simple, direct path to the model’s forward call with batch inference executed in a single pass.

Solution

A new execution path (encoder_only=True) that creates a lightweight EncoderExecutor instead of the full PyExecutor. The encode() method tokenizes, packs, and runs a single forward pass directly through ModelEngine.encoder_forward(), returning EncoderOutput with logits. This new API demonstrates a 3.92× speedup for the BERT 110M model (textattack/bert-base-uncased-yelp-polarity) in eager mode with batch size 10.

encode()   mean: 5.17ms  (p50=5.16ms)
generate() mean: 20.25ms (p50=20.02ms)
Speedup: 3.92x

Usage

# New dedicated path
llm = LLM(model="bert-base-uncased-yelp-polarity", encoder_only=True)
outputs = llm.encode(["Hello world", "Test sentence"])                                                                                                                 
print(outputs[0].logits)  # [num_classes] tensor
                                                                                                                                                                       
# Old path still works unchanged (no encoder_only flag)
llm = LLM(model="bert-base-uncased-yelp-polarity", disable_overlap_scheduler=True)
outputs = llm.generate(prompts, SamplingParams(return_context_logits=True))    

• encoder_only=True must be explicitly set. Default (None) uses the old generate() path.
• encoder_only=True creates only EncoderExecutor; False/None creates only PyExecutor. Mutually exclusive.
• generate()/generate_async() raise RuntimeError when encoder_only=True. encode() is the only API.
• Since llm.encode() reuses PyTorchModelEngine and its _forward_step() path, features like TorchCompileConfig are compatible.

Future Works

• Encoder CUDA graph integration — capture the encoder model to one single CUDA graph
• Triton backend update — add an encoder model example
• Parallelism supports (e.g. TP > 1) — expand the EncoderExecutor
• Other minor optimizations (e.g. batch tokenization, cache AttentionMetadata, etc)

Test Coverage

• tests/unittest/llmapi/test_llm_encode.py — 11 new tests:
  • Basic: single string, batch, token IDs, mixed input types
  • Correctness: logits compared against HuggingFace BertForSequenceClassification
  • Health check: _check_health() returns True for encoder-only LLM
• Existing tests unaffected (no encoder_only=True → old path):
  • tests/integration/defs/test_e2e.py::test_ptp_quickstart_bert
  • tests/unittest/llmapi/test_llm_pytorch.py::test_llm_reward_model

CC: @symphonylyh @amukkara @nvrohanv @schetlur-nv @juney-nvidia

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[nvrohanv]

Some comments on tokenization piece and handling of empty batch but overall looks good!

[nvrohanv]

Nit: How do we handle case where an empty batch is passed in, unless theres handling elsewhere I think this would cause Index Error. I'm guessing this logic exists elsewhere as well so this might be a general question of how we want to handle it.

[nvrohanv]

Is it faster to do the tokenization in a batch - I see that later we do the processing to turn it into a flat "packed" tensor? Especially for cases with larger batch sizes I'm curious if we could cut down on tokenization overhead this way. Not sure about the structure of input_processor and if it can do this well.

[nvrohanv]

Can we make this a bit clearer that it would work for decoder only models being run in "encoder-only" style

[amukkara]

Similar to mm_encoder_only above, should this field be bool with default=False?

[amukkara]

checkpoint loader logic is common across create_py_executor, create_encoder_executor, can be moved to helper function.

[amukkara]

Suggested change

	return self._forward_step(model_inputs, None, False)
	return self._forward_step(model_inputs, gather_ids=None, gather_context_logits=False)

[Funatiq]

I would suggest to enforce kwargs by adding *, after inputs here.

[pcastonguay]

@Superjomn could you review since it adds a new method to LLM API? Thx.

[Funatiq]

I would suggest to enforce kwargs by adding *, after inputs here.

[Funatiq]

How does this function relate to _forward_step_mm_encoder_only? The purpose seems similar. Does it make sense to unify these?