EncodeCudaGraphConfig#

class tensorrt_llm.llmapi.EncodeCudaGraphConfig(
*,
batch_sizes: ~typing.List[int] | None = None,
max_batch_size: ~typing.Annotated[int,
~annotated_types.Ge(ge=0)] = 0,
enable_padding: bool = False,
mode: ~typing.Literal['encode'] = 'encode',
num_tokens: ~typing.Annotated[~typing.List[~typing.Annotated[int,
~annotated_types.Gt(gt=0)]] | None,
~annotated_types.MinLen(min_length=1)] = None,
max_num_token: ~typing.Annotated[int,
~annotated_types.Ge(ge=0)] = 0,
seq_lens: ~typing.Annotated[~typing.List[~typing.Annotated[int,
~annotated_types.Gt(gt=0)]] | None,
~annotated_types.MinLen(min_length=1)] = None,
max_seq_len: ~typing.Annotated[int,
~annotated_types.Ge(ge=0)] = 0,
extra_model_inputs: ~typing.List[~tensorrt_llm.llmapi.llm_args.EncodeExtraInputSpec] = <factory>,
)[source]#

Bases: BaseCudaGraphConfig

CUDA graph configuration for encode-only requests (LLM.encode()).

A graph is captured at startup for every feasible combination of batch_sizes, num_tokens and seq_lens. A batch runs the graph for its (batch size, total tokens, longest request) shape, rounded up when enable_padding is set, and runs eagerly if no graph matches.

Example:

cuda_graph_config = EncodeCudaGraphConfig(
    batch_sizes=[1, 2, 4, 8],
    num_tokens=[128, 256, 512],
    seq_lens=[64, 128],
    enable_padding=True,
    extra_model_inputs=[
        EncodeExtraInputSpec(name="token_type_ids",
                             shape=("num_tokens",),
                             dtype="int32"),
    ],
)
llm = LLM(model, encode_only=True, cuda_graph_config=cuda_graph_config)
# token_type_ids: one value per token, prompts packed in order.
outputs = llm.encode(prompts, token_type_ids=token_type_ids)
field batch_sizes: List[int] | None = None#

List of batch sizes to create CUDA graphs for.

field enable_padding: bool = False#

If true, batches are rounded up to the nearest cuda_graph_batch_size. This is usually a net win for performance.

field extra_model_inputs: List[EncodeExtraInputSpec] [Optional]#

Tensor kwargs (beyond input_ids / position_ids) that LLM.encode() passes to the encoder forward() under CUDA graphs. Each is backed by a static buffer sized along its symbolic dim at the bucket maximum (num_tokens → max(num_tokens), batch_size → max(batch_sizes)). With encoder CUDA graphs enabled, every tensor kwarg passed to encode() must be declared here, and every declared input must be passed on every call. Tensors may be on host or device.

field max_batch_size: NonNegativeInt = 0#

Maximum batch size for CUDA graphs.

Constraints:

  • ge = 0

field max_num_token: Annotated[int, Ge(ge=0)] = 0#

Maximum total number of tokens for encoder CUDA graphs. If num_tokens is provided, must equal max(num_tokens); otherwise num_tokens is generated from this value. Ignored by a fixed-shape feature encoder.

Constraints:

  • ge = 0

field max_seq_len: Annotated[int, Ge(ge=0)] = 0#

Maximum per-request sequence length for encoder CUDA graphs. If seq_lens is provided, must equal max(seq_lens); otherwise seq_lens is generated from this value. Ignored by a fixed-shape feature encoder.

Constraints:

  • ge = 0

field mode: Literal['encode'] = 'encode'#

CUDA graph configuration mode.

field num_tokens: List[Annotated[int, Gt(gt=0)]] | None = None#

List of total token counts (sum of all per-request sequence lengths in a batch) to create encoder CUDA graphs for. Required for an encoder that packs a variable number of tokens per request; ignored by an encoder whose input is a fixed-shape per-request feature tensor, which derives this from the model.

Constraints:

  • min_length = 1

field seq_lens: List[Annotated[int, Gt(gt=0)]] | None = None#

List of max per-request sequence lengths to create encoder CUDA graphs for. Required for an encoder that packs a variable number of tokens per request; ignored by an encoder whose input is a fixed-shape per-request feature tensor, which derives this from the model.

Constraints:

  • min_length = 1

__init__(**data: Any) → None#

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

validator validate_base_cuda_graph_config  »  all fields#

Validate CUDA graph configuration.

Ensures that: 1. If batch_sizes is provided, max_batch_size is derived as max(batch_sizes).

If max_batch_size was already set it must be compatible (equal to max(batch_sizes)); otherwise an error is raised.

  1. If only max_batch_size is provided, batch_sizes is generated from it.

  2. If neither is provided, a default max_batch_size of 128 is used.

validator validate_encoder_cuda_graph_config  »  all fields[source]#