EncodeExtraInputSpec

class tensorrt_llm.llmapi.EncodeExtraInputSpec(

*,

name: str,

shape: Tuple[Literal['num_tokens', 'batch_size'] | Annotated[int, Gt(gt=0)], ...],

dtype: str,

)

Bases: StrictBaseModel

Declares an extra encoder-forward tensor kwarg for CUDA graph capture.

One spec is required for every tensor kwarg passed to LLM.encode(…, **model_kwargs) (e.g. token_type_ids for BERT). The runner backs each with a static buffer sized at the bucket maximum, and warmup captures every bucket with a zero-filled stand-in so the graph sees the full forward signature before the first real call.

Only tensors can be declared: a non-tensor kwarg cannot be captured and instead forces that encode() call onto the eager path. Device is not part of the spec: pass tensors where they already live. Each replay copies them into the graph’s buffer asynchronously, H2D for host tensors and D2D for device tensors.

field dtype: str [Required]

Tensor dtype string accepted by tensorrt_llm (e.g. “int32”, “float32”, “bfloat16”). See tensorrt_llm._utils._str_to_torch_dtype_dict for the full list.

field name: str [Required]

Kwarg name as it appears in the encoder forward() signature (e.g. “token_type_ids”).

field shape: Tuple[Literal['num_tokens', 'batch_size'] | Annotated[int, Gt(gt=0)], ...] [Required]

Tensor shape. Use exactly one symbolic dim from {“num_tokens”, “batch_size”}: “num_tokens” scales with the packed token bucket (e.g. token_type_ids), “batch_size” scales with the request bucket (e.g. per-request features that are independent of sequence length). Remaining dims must be positive integer literals (e.g. hidden_size).

__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.

resolve_shape(

num_tokens: int,

batch_size: int,

) → Tuple[int, ...]

Substitute the symbolic dim with the matching size, leaving literals alone.

Each spec uses exactly one symbolic dim (validated). The caller passes both candidate sizes; only the one this spec actually uses is consumed.

symbolic_dim() → Tuple[str, int]

Return (name, axis) of the spec’s single symbolic dim.

The shape validator guarantees exactly one symbolic dim is present.

torch_dtype() → dtype