NVIDIA / NVIDIA/TensorRT-LLM

Support NumPy token IDs in executor APIs

Open
#16,910 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

LLM API
Dominant language
Python
Stars
14.7k
Forks
2.8k
Avg merge
2d 23h
Merged PRs (30d)
489

Description

Baseten.co would like the executor APIs to accept NumPy token ID arrays directly.

Our tokenizer can return a contiguous token buffer with tokenizer.encode_to_numpy(), but TensorRT-LLM currently expects list[int]. At 200k tokens, crossing this boundary costs roughly 17 ms per materialization, and our request path currently pays that cost twice: about 34 ms total per request.

Accepting a contiguous integer NumPy array—or, more generally, a compatible Python buffer—would let the executor consume the token IDs without converting them through a Python list. That should make this boundary effectively zero-copy.

Would you be open to adding NumPy/buffer support alongside the existing list[int] API? We would be happy to contribute a small implementation and regression test.

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.

Research direction

Start at the executor APIs that currently accept list[int] token IDs and trace how tokenizer.encode_to_numpy() values reach them. Add a regression test covering a contiguous NumPy integer array or compatible buffer, and verify that executor requests accept it without list materialization while preserving existing list[int] behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
backend-api-design, machine-learning, performance
Issue type
Feature
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.