pytorch / pytorch/TensorRT

✨[Feature] Multi-Framework Runner Support in tools

Open
#4,220 3 comments 0 reactions 1 assignee View on GitHub

@narendasan is already working on this.

Since Apr 28, 2026.

feature request
Dominant language
Python
Stars
3k
Forks
410
Avg merge
3d 18h
Merged PRs (30d)
78

Description

Is your feature request related to a problem? Please describe.
The current tools/llm directory only supports text-generation models (LLMs and VLMs). The transformer architecture now underpins a much wider range of model families — diffusion models, speech models, encoder models, video generation — and users optimizing those models with Torch-TRT have no equivalent tooling. Each model type requires different I/O handling, evaluation metrics, and benchmarking methodology that can't be shoehorned into run_llm.py.

Describe the solution you'd like
Extend the tools/ directory with one entry point per task strategy, alongside the existing llm/ directory, each with its own run_.py entry point and benchmarking logic:

tools/
llm/ # existing: LLMs and VLMs
diffusion/ # new
audio/ # new
encoder/ # new
video/ # new
neural_operator/ #new

The model types to add:

Image

Each runner should expose the same --benchmark flag and report latency/throughput in
model-appropriate units (tokens/s for text, images/s for diffusion, real-time factor for
audio).

Describe alternatives you've considered

  • Extend run_llm.py with --model-type — a single script becomes unwieldy; diffusion models have fundamentally different I/O (image tensors, noise schedules, guidance scale) that doesn't map onto the LLM text-generation loop.
  • Leave each model type to users — this is the status quo and means there is no canonical Torch-TRT path for non-LLM workloads, even though TRT has published optimizations for all of these families (diffusion, BERT, Whisper).
  • Only add diffusion — diffusion models are the most visible gap, but ASR (Whisper) and encoder models (BERT) are equally well-established TensorRT use cases with existing NVIDIA blog coverage and deserve first-class support.

Additional context

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.