AsyncLLM#

class tensorrt_llm.llmapi.AsyncLLM(
placement_groups: List[Any] | None = None,
placement_bundle_indices: List[List[int]] | None = None,
per_worker_gpu_share: float | None = None,
*args,
**kwargs,
)[source]#

Bases: LLM

AsyncLLM is a subclass of LLM that supports asynchronous setup, release and resume operations that are necessary for RL or agentic scenarios.

Currently, RL APIs are only supported with Ray orchestrator.

__init__(
placement_groups: List[Any] | None = None,
placement_bundle_indices: List[List[int]] | None = None,
per_worker_gpu_share: float | None = None,
*args,
**kwargs,
)[source]#
async collective_rpc(
method: str,
args: tuple[Any, ...] = (),
kwargs: dict | None = None,
unique_reply_rank: int | None = None,
target_ranks: int | list[int] | None = None,
) → list[Any][source]#

Execute an asynchronous RPC call on all GPU workers. Currently, this is only supported for RayExecutor.

Parameters:

  • method (str) – The name of the worker method to execute.

  • args (tuple[Any, ...]) – Positional arguments to pass to the worker method. Defaults to ().

  • kwargs (dict, optional) – Keyword arguments to pass to the worker method. Defaults to None.

  • unique_reply_rank (int, optional) – The rank of the worker that will be used to send the reply.

  • target_ranks (int | list[int] | None) – The ranks of the workers that will be used to send the reply.

Returns:

A list of results from each worker.

Return type:

list[Any]

encode(
inputs: str | List[int] | TextPrompt | TokensPrompt | Sequence[str | List[int] | TextPrompt | TokensPrompt],
add_special_tokens: bool = True,
batch_indexed_model_output: bool = True,
copy_logits_to_host: bool = True,
return_raw_logits: bool = False,
**model_kwargs: Any,
) → EncoderOutput | List[EncoderOutput] | Tensor#

prototype Encode inputs using an encoder-only model (PyTorch backend only).

Only available when encode_only=True is set in the LLM constructor.

Parameters:

  • inputs (tensorrt_llm.inputs.data.PromptInputs, Sequence[tensorrt_llm.inputs.data.PromptInputs]) – The prompt text or token ids. It can be a single prompt or batched prompts.

  • add_special_tokens (bool) – Whether to add special tokens (e.g., [CLS]/[SEP]) during tokenization. Defaults to True.

  • batch_indexed_model_output (bool) – If specified, assume batched model output indexed by request index, as opposed to token index. Defaults to True.

  • copy_logits_to_host (bool) – If set, copy logits from device to host. Otherwise, return a view into the on-device logits tensor. Defaults to True.

  • return_raw_logits (bool) – Whether to return the raw CPU logits tensor for the whole input batch. Defaults to False.

  • model_kwargs (Any) – Model-specific inputs passed through to the model’s forward(). Examples: token_type_ids (BERT), inputs_embeds (reward models).

Returns:

If return_raw_logits=True, returns the raw CPU logits tensor for the whole input batch. Otherwise, returns one EncoderOutput for a single input, or a list of EncoderOutput objects for batched inputs.

Return type:

Union[tensorrt_llm.llmapi.llm.EncoderOutput, List[tensorrt_llm.llmapi.llm.EncoderOutput], torch.Tensor]

Raises:

RuntimeError – If encode_only mode is not enabled.

generate(
inputs: str | List[int] | TextPrompt | TokensPrompt | Sequence[str | List[int] | TextPrompt | TokensPrompt],
sampling_params: SamplingParams | List[SamplingParams] | None = None,
use_tqdm: bool = True,
lora_request: LoRARequest | Sequence[LoRARequest] | None = None,
prompt_adapter_request: PromptAdapterRequest | Sequence[PromptAdapterRequest] | None = None,
kv_cache_retention_config: KvCacheRetentionConfig | Sequence[KvCacheRetentionConfig] | None = None,
disaggregated_params: DisaggregatedParams | Sequence[DisaggregatedParams] | None = None,
scheduling_params: SchedulingParams | List[SchedulingParams] | None = None,
conversation_params: ConversationParams | List[ConversationParams] | None = None,
cache_salt: str | Sequence[str] | None = None,
priority: float | List[float] = 0.5,
) → RequestOutput | List[RequestOutput]#

Generate output for the given prompts in the synchronous mode. Synchronous generation accepts either single prompt or batched prompts.

Parameters:
Returns:

The output data of the completion request to the LLM.

Return type:

Union[tensorrt_llm.llmapi.llm.RequestOutput, List[tensorrt_llm.llmapi.llm.RequestOutput]]

generate_async(*args, **kwargs)[source]#

Generate output for the given prompt in the asynchronous mode. Asynchronous generation accepts single prompt only.

Parameters:

  • inputs (Union[tensorrt_llm.inputs.data.PromptInputs, tensorrt_llm.llmapi.llm.PreprocessedInputs]) – The prompt text or token ids, or a PreprocessedInputs returned by preprocess. If the latter, preprocessing will be skipped by this method.

  • sampling_params (tensorrt_llm.sampling_params.SamplingParams, optional) – The sampling params for the generation. Defaults to None. A default one will be used if not provided.

  • lora_request (tensorrt_llm.executor.request.LoRARequest, optional) – LoRA request to use for generation, if any. Defaults to None.

  • prompt_adapter_request (tensorrt_llm.executor.request.PromptAdapterRequest, optional) – Prompt Adapter request to use for generation, if any. Defaults to None.

  • streaming (bool) – Whether to use the streaming mode for the generation. Defaults to False.

  • kv_cache_retention_config (tensorrt_llm.bindings.executor.KvCacheRetentionConfig, optional) – Configuration for the request’s retention in the KV Cache. Defaults to None.

  • disaggregated_params (tensorrt_llm.disaggregated_params.DisaggregatedParams, optional) – Disaggregated parameters. Defaults to None.

  • trace_headers (Mapping[str, str], optional) – Trace headers. Defaults to None.

  • scheduling_params (tensorrt_llm.scheduling_params.SchedulingParams, optional) – Scheduling parameters. Defaults to None.

  • conversation_params (tensorrt_llm.conversation_params.ConversationParams, optional) – Conversation parameters. Defaults to None.

  • cache_salt (str, optional) – If specified, KV cache will be salted with the provided string to limit the kv cache reuse to the requests with the same string. Defaults to None.

  • priority (float) – The scheduling priority for the request, in the range [0, 1]. Higher values indicate higher priority. Defaults to 0.5.

Returns:

The output data of the completion request to the LLM.

Return type:

tensorrt_llm.llmapi.llm.RequestOutput

get_data_transceiver_state() → bytes#

prototype Get the serialized DataTransceiverState for arbitrary KV cache transfer.

Returns:

Serialized DataTransceiverState, or empty bytes if no transceiver is configured.

Return type:

bytes

get_kv_cache_capacity() → dict#

beta Get the runtime’s static primary/GPU KV cache capacity.

Raises:

RuntimeError – If called when encode_only=True.

Returns:

KV cache capacity. The returned capacity covers the primary

GPU KV cache pool only; CPU/host offload capacity is not included. e.g., {“maxNumBlocks”: …, “tokensPerBlock”: …, “maxNumTokens”: …}

Return type:

dict

get_kv_cache_events(
timeout: float | None = 2,
) → List[dict]#

beta Get iteration KV events from the runtime.

KV events are used to track changes and operations within the KV Cache. Types of events:

  • KVCacheCreatedData: Indicates the creation of cache blocks.

  • KVCacheStoredData: Represents a sequence of stored blocks.

  • KVCacheRemovedData: Contains the hashes of blocks that are being removed from the cache.

  • KVCacheUpdatedData: Captures updates to existing cache blocks.

To enable KV events:

  • set event_buffer_max_size to a positive integer in the KvCacheConfig.

  • set enable_block_reuse to True in the KvCacheConfig.

Parameters:

timeout (float, optional) – Max wait time in seconds when retrieving events from queue. Defaults to 2.

Returns:

A list of runtime events as dict.

Return type:

List[dict]

get_kv_cache_events_async(
timeout: float | None = 2,
) → IterationResult#

beta Get iteration KV events from the runtime.

KV events are used to track changes and operations within the KV Cache. Types of events:

  • KVCacheCreatedData: Indicates the creation of cache blocks.

  • KVCacheStoredData: Represents a sequence of stored blocks.

  • KVCacheRemovedData: Contains the hashes of blocks that are being removed from the cache.

  • KVCacheUpdatedData: Captures updates to existing cache blocks.

To enable KV events:

  • set event_buffer_max_size to a positive integer in the KvCacheConfig.

  • set enable_block_reuse to True in the KvCacheConfig.

Parameters:

timeout (float, optional) – Max wait time in seconds when retrieving events from queue. Defaults to 2.

Returns:

An async iterable object containing runtime events.

Return type:

tensorrt_llm.executor.result.IterationResult

get_stats(timeout: float | None = 2) → List[dict]#

beta Get iteration statistics from the runtime. To collect statistics, call this function after prompts have been submitted with LLM().generate().

Parameters:

timeout (float, optional) – Max wait time in seconds when retrieving stats from queue. Defaults to 2.

Returns:

A list of runtime stats as dicts.

e.g., [{“cpuMemUsage”: …, “iter”: 0, …}, {“cpuMemUsage”: …, “iter”: 1, …}]

Return type:

List[dict]

get_stats_async(
timeout: float | None = 2,
) → IterationResult#

beta Get iteration statistics from the runtime. To collect statistics, you can call this function in an async coroutine or the /metrics endpoint (if you’re using trtllm-serve) after prompts have been submitted.

Parameters:

timeout (float, optional) – Max wait time in seconds when retrieving stats from queue. Defaults to 2.

Returns:

An async iterable object containing runtime stats.

Return type:

tensorrt_llm.executor.result.IterationResult

async pause_generation() → None[source]#

Abort all in-flight requests and block new ones until resume_generation() is called.

Sends abort signals then drains the executor. It only returns once the engine has no active or queued requests.

preprocess(
inputs: str | List[int] | TextPrompt | TokensPrompt,
sampling_params: SamplingParams | None = None,
disaggregated_params: DisaggregatedParams | None = None,
) → PreprocessedInputs#

prototype Preprocess raw prompts into token IDs and multimodal params.

Parameters:
Returns:

A preprocessed-inputs object that can be

passed directly to generate_async() as inputs.

Return type:

tensorrt_llm.llmapi.llm.PreprocessedInputs

async release(tags: list[str])[source]#

Release the GPU memory used by the LLM asynchronously.

Parameters:

tags – List of memory tag strings to release (e.g., [“model”, “kv_cache”]).

async resume(tags: list[str])[source]#

Resume the GPU memory used by the LLM asynchronously.

Parameters:

tags – List of memory tag strings to resume (e.g., [“model”, “kv_cache”]).

async resume_generation() → None[source]#

Allow new generation requests after a pause_generation() call.

async setup_async()[source]#

Setup the LLM asynchronously.

shutdown() → None#

beta None

async update_weights(weights: dict[str, str])[source]#

Update the weights of the LLM asynchronously.

Parameters:

weights – Dictionary mapping device UUIDs to IPC handles for weight tensors.

property disaggregated_params: dict#

beta None

property llm_id: str#

beta None

property startup_metrics: dict#

beta Cache and return rank-0 startup metrics.

Returns:

The cached metrics, or an empty dict when metrics retrieval fails.

Return type:

dict

property tokenizer: TokenizerBase | None#