Python API Reference#

This section provides documentation for the TensorRT Edge-LLM Python package.

Python workflows use tensorrt_edgellm.quantization for checkpoint quantization, tensorrt_edgellm for ONNX export, and experimental.server for the experimental high-level API and OpenAI-compatible server.

Experimental Server#

Checkpoint-native inference server for TensorRT Edge-LLM.

Public API:

from experimental.server import LLM, SamplingParams

llm = LLM(model="Qwen/Qwen3.5-0.8B")
outputs = llm.generate(["Hello!"], SamplingParams(max_tokens=64))
print(outputs[0].text)

# Or start an OpenAI-compatible server:
llm.serve(port=8000)

class experimental.server.LLM( model: str, *, cache_dir: str = '', engine_cache_max_size_gb: float = 50.0, clear_engine_cache: bool = False, max_input_len: int = 4096, max_batch_size: int = 1, max_kv_cache_capacity: int = 8192, draft_top_k: int | None = None, draft_step: int | None = None, verify_tree_size: int | None = None, build_options: BuildOptions | None = None, speculative_config: Any | None = None, context_cache_config: ContextCacheConfig | Mapping[str, Any] | None = None, )[source]#

Bases: object

Checkpoint-direct entry point for offline and HTTP inference.

model accepts a local checkpoint or Hugging Face ID. The experimental builder compiles every model component into a profile-specific cache bundle and externalizes every supported weight kind.

runtime_kind = 'chat'#: Selects the HTTP contract without a hierarchy of capability flags.

__init__( model: str, *, cache_dir: str = '', engine_cache_max_size_gb: float = 50.0, clear_engine_cache: bool = False, max_input_len: int = 4096, max_batch_size: int = 1, max_kv_cache_capacity: int = 8192, draft_top_k: int | None = None, draft_step: int | None = None, verify_tree_size: int | None = None, build_options: BuildOptions | None = None, speculative_config: Any | None = None, context_cache_config: ContextCacheConfig | Mapping[str, Any] | None = None, )[source]#

close() → None[source]#: Drain active work and release native engines and device memory.

generate( prompts: str | List[str] | List[List[Dict[str, Any]]], sampling_params: SamplingParams | None = None, *, tools: Sequence[Dict[str, Any]] | None = None, tool_choice: str | Dict[str, Any] | None = None, tool_parser: str = 'auto', reasoning_parser: str = 'none', ) → List[CompletionOutput][source]#

Generate completions for the given prompts.

Parameters:

prompts – A single prompt string, a list of prompt strings, or a list of OpenAI-style message lists.
sampling_params – Sampling configuration. Defaults to SamplingParams().
tools – Optional OpenAI-compatible tool definitions.
tool_choice – Optional OpenAI-compatible tool choice.

Returns:

List of CompletionOutput objects, one per prompt.

chat( messages: List[Dict[str, Any]], sampling_params: SamplingParams | None = None, *, tools: Sequence[Dict[str, Any]] | None = None, tool_choice: str | Dict[str, Any] | None = None, tool_parser: str = 'auto', reasoning_parser: str = 'none', ) → CompletionOutput[source]#

Single-turn chat completion (convenience wrapper).

Parameters:

messages – OpenAI-style message list.
sampling_params – Sampling configuration.
tools – Optional OpenAI-compatible tool definitions.
tool_choice – Optional OpenAI-compatible tool choice.

Returns:

A single CompletionOutput.

generate_stream( messages: List[Dict[str, Any]], sampling_params: SamplingParams | None = None, *, tools: Sequence[Dict[str, Any]] | None = None, tool_choice: str | Dict[str, Any] | None = None, prebuilt_request: Any | None = None, ) → Iterator[StreamDelta][source]#

Stream generation deltas for a single message list.

Runs handleRequest in a background thread with a StreamChannel attached, yielding StreamDelta objects as tokens are produced.

generate_stream_with_audio( messages: List[Dict[str, Any]], sampling_params: SamplingParams | None = None, *, audio_params: AudioParams | None = None, prebuilt_request: Any | None = None, ) → Iterator[StreamDelta][source]#

Stream text and audio deltas for a single Omni request.

Runs the Thinker-Talker streaming pipeline in a background thread. Text deltas arrive through a StreamChannel and PCM chunks through an AudioStreamChannel; the two are interleaved into one generator. Admission follows generate_stream: the HTTP layer owns the gate when it passes prebuilt_request; otherwise it is acquired here.

generate_speech_stream( text: str, audio_params: AudioParams | None = None, ) → Iterator[StreamDelta][source]#

Standalone TTS on the Omni stack: synthesize text directly.

No Thinker generation pass — the input text goes straight to the Talker. Yields audio-only StreamDeltas.

list_voices() → List[str][source]#: Speaker names accepted as voice; empty when not Omni-capable.

serve( host: str = '0.0.0.0', port: int = 8000, *, served_model_name: str = '', api_key: str = '', reasoning_parser: str = 'auto', tool_call_parser: str = 'auto', enable_auto_tool_choice: bool = False, max_queued_requests: int = 16, queue_timeout: float = 600.0, allowed_local_media_path: str | None = None, ) → None[source]#: Start the HTTP frontend for this runtime.

property model_dir: str#: Path to the resolved model checkpoint.

property model_id: str#: User-facing model identifier supplied at initialization.

property bundle_dir: str#: Profile-specific engine bundle selected from the cache.

property cache_dir: str#: Root containing downloaded checkpoints and built bundles.

property max_batch_size: int#: Maximum batch size supported by the loaded engine.

property video_capable: bool#: Whether this model bundle supports video input.

property has_draft_model: bool#: Whether speculative decoding is active.

property context_cache_enabled: bool#: Whether this runtime reuses matching text prefixes.

get_context_cache_metrics()[source]#: Return native reuse counters, or None when reuse is disabled.

property bundle_layout: BundleLayout#: Immutable component contract for the selected model bundle.

class experimental.server.TTS( model: str, *, cache_dir: str = '', engine_cache_max_size_gb: float = 50.0, clear_engine_cache: bool = False, max_input_len: int = 4096, max_batch_size: int = 1, max_kv_cache_capacity: int = 8192, build_options: BuildOptions | None = None, )[source]#

Bases: object

Checkpoint-direct serving for a model-owned TTS component stack.

runtime_kind = 'tts'#

has_draft_model = False#

__init__( model: str, *, cache_dir: str = '', engine_cache_max_size_gb: float = 50.0, clear_engine_cache: bool = False, max_input_len: int = 4096, max_batch_size: int = 1, max_kv_cache_capacity: int = 8192, build_options: BuildOptions | None = None, ) → None[source]#

generate_speech_stream( text: str, audio_params: AudioParams | None = None, ) → Iterator[StreamDelta][source]#: Synthesize text; yields audio-only StreamDeltas.

list_voices() → List[str][source]#: Speaker names accepted as voice.

close() → None[source]#: Drain active speech generation and release native resources.

property model_id: str#

property model_dir: str#

property bundle_dir: str#

property cache_dir: str#

property bundle_layout: BundleLayout#: Immutable component contract for the selected model bundle.

serve(host: str = '0.0.0.0', port: int = 8000) → None[source]#: Start the HTTP server (speech endpoint only).

class experimental.server.SamplingParams( temperature: float = 0.7, top_p: float = 0.9, top_k: int = 50, max_tokens: int = 2048, enable_thinking: bool = False, disable_spec_decode: bool = False, num_logprobs: int = 0, stop: List[str] = <factory>, logit_bias: Dict[int, float]=<factory>, skip_special_tokens: bool = True, reuse_context: bool = True, cache_generated_tokens: bool = True, )[source]#

Bases: object

Sampling parameters for one generation request.

temperature: float = 0.7#

top_p: float = 0.9#

top_k: int = 50#

max_tokens: int = 2048#

enable_thinking: bool = False#

disable_spec_decode: bool = False#

num_logprobs: int = 0#

stop: List[str]#

logit_bias: Dict[int, float]#

skip_special_tokens: bool = True#

reuse_context: bool = True#

cache_generated_tokens: bool = True#

__init__( temperature: float = 0.7, top_p: float = 0.9, top_k: int = 50, max_tokens: int = 2048, enable_thinking: bool = False, disable_spec_decode: bool = False, num_logprobs: int = 0, stop: List[str] = <factory>, logit_bias: Dict[int, float]=<factory>, skip_special_tokens: bool = True, reuse_context: bool = True, cache_generated_tokens: bool = True, ) → None#

class experimental.server.CompletionOutput( text: str = '', token_ids: List[int] = <factory>, prompt_tokens: int | None = None, finish_reason: str | None = None, logprobs: List[List[LogprobEntry]] = <factory>, tool_calls: Dict[str, ~typing.Any]]=<factory>, reasoning: str | None = None, )[source]#

Bases: object

Output of a single generation request.

text: str = ''#

token_ids: List[int]#

prompt_tokens: int | None = None#

finish_reason: str | None = None#

logprobs: List[List[LogprobEntry]]#

tool_calls: List[Dict[str, Any]]#

reasoning: str | None = None#

__init__( text: str = '', token_ids: List[int] = <factory>, prompt_tokens: int | None = None, finish_reason: str | None = None, logprobs: List[List[LogprobEntry]] = <factory>, tool_calls: Dict[str, ~typing.Any]]=<factory>, reasoning: str | None = None, ) → None#

class experimental.server.StreamDelta( text: str = '', token_ids: List[int] = <factory>, prompt_tokens: int | None = None, finished: bool = False, finish_reason: str | None = None, logprobs: List[List[LogprobEntry]] = <factory>, audio_bytes: bytes | None = None, )[source]#

Bases: object

Single delta from a streaming generation.

Text deltas carry text/token_ids; audio deltas (Omni streaming) carry audio_bytes (int16 LE mono PCM) instead. finished marks the end of the text stream; generator exhaustion ends the audio stream.

text: str = ''#

token_ids: List[int]#

prompt_tokens: int | None = None#

finished: bool = False#

finish_reason: str | None = None#

logprobs: List[List[LogprobEntry]]#

audio_bytes: bytes | None = None#

__init__( text: str = '', token_ids: List[int] = <factory>, prompt_tokens: int | None = None, finished: bool = False, finish_reason: str | None = None, logprobs: List[List[LogprobEntry]] = <factory>, audio_bytes: bytes | None = None, ) → None#

class experimental.server.AudioParams( voice: str = '', talker_temperature: float = 0.9, talker_top_k: int = 50, talker_top_p: float = 1.0, repetition_penalty: float = 1.05, max_audio_length: int = 4096, codec_chunk_frames: int = 10, talker_prefill_threshold: int = 4, )[source]#

Bases: object

Talker / vocoder knobs for one Omni audio-output request.

voice: str = ''#

talker_temperature: float = 0.9#

talker_top_k: int = 50#

talker_top_p: float = 1.0#

repetition_penalty: float = 1.05#

max_audio_length: int = 4096#

codec_chunk_frames: int = 10#

talker_prefill_threshold: int = 4#

__init__( voice: str = '', talker_temperature: float = 0.9, talker_top_k: int = 50, talker_top_p: float = 1.0, repetition_penalty: float = 1.05, max_audio_length: int = 4096, codec_chunk_frames: int = 10, talker_prefill_threshold: int = 4, ) → None#

experimental.server.load_model(**kwargs)[source]#: Select the model-specific runtime from provider checkpoint metadata.

class experimental.server.ContextCacheConfig( enabled: bool = False, max_records: int = 1024, recurrent_snapshot_pool_bytes: int = 0, partial_kv_snapshot_pool_bytes: int = 0, )[source]#

Bases: object

Deployment-scoped KV-cache reuse configuration.

enabled: bool = False#

max_records: int = 1024#

recurrent_snapshot_pool_bytes: int = 0#

partial_kv_snapshot_pool_bytes: int = 0#

classmethod parse( value: Mapping[str, Any] | ContextCacheConfig | None, ) → ContextCacheConfig[source]#: Normalize an HLAPI mapping without accepting ignored keys.

__init__( enabled: bool = False, max_records: int = 1024, recurrent_snapshot_pool_bytes: int = 0, partial_kv_snapshot_pool_bytes: int = 0, ) → None#

experimental.server.clear_engine_cache(cache_dir: str = '') → int[source]#: Remove published engine bundles without deleting checkpoints.

experimental.server.prune_engine_cache( cache_dir: str = '', max_size_bytes: int = 53687091200, keep: str = '', ) → int[source]#: Evict least-recently-used bundles until the cache is within its cap.

Quantization#

Standalone quantization for TensorRT Edge-LLM.

Decoupled from the ONNX exporter — runs in a clean venv with only torch, transformers, and modelopt.

tensorrt-edgellm-quantize –help

Checkpoint Exporter#

TensorRT Edge-LLM Python package.

The checkpoint exporter is optional. Its PyTorch and ONNX modules are loaded only when an export API is requested, so the checkpoint-direct builder and server can import the package without installing the legacy export toolchain.

class tensorrt_edgellm.AutoModel[source]#

Bases: object

HuggingFace-style factory that dispatches on model_type.

classmethod from_pretrained( model_dir: str, device: str = 'cpu', key_remap=None, key_prefix: str | None = None, eagle_base: bool = False, eagle_draft_dir: str | None = None, reduced_vocab_dir: str | None = None, mtp_base: bool = False, mtp_tree_base: bool = False, mtp_draft: bool = False, tp_size: int = 1, tp_rank: int = 0, dflash_base: bool = False, dflash_tree_base: bool = False, dflash_draft: bool = False, dflash_draft_dir: str | None = None, jetspec_base: bool = False, jetspec_tree_base: bool = False, jetspec_draft: bool = False, jetspec_draft_dir: str | None = None, dspark_base: bool = False, dspark_draft: bool = False, dspark_draft_dir: str | None = None, gemma4_mtp_base: bool = False, gemma4_mtp_draft: bool = False, gemma4_kv_sharing_map: list[dict] | None = None, gemma4_target_kv_cache_quant: str | None = None, num_decoder_layers: int | None = None, extra_configs: dict | None = None, ) → torch.nn.Module[source]#

Construct and load a model from model_dir.

Reads config.json via ModelConfig, looks up the model class in the registry (falling back to the built-in CausalLM), instantiates it, moves it to device, and loads safetensors weights.

Parameters:

model_dir – Local HF checkpoint directory.
device – Target device (e.g. "cpu", "cuda:0").
key_remap – Optional callable (key: str) -> Optional[str]. Passed through to load_weights() for checkpoint key remapping (e.g. TTS talker codec_embedding → embed_tokens).
key_prefix – Explicit checkpoint key prefix to strip (e.g. "talker."). Passed through to load_weights().
eagle_base – When True, export as EAGLE3 base model with extra tree-attention inputs and hidden_states output.
eagle_draft_dir – Optional EAGLE3 draft checkpoint directory. Gemma4 EAGLE3 uses it to select the draft-trained target hidden layers for base hidden_states output.
reduced_vocab_dir – Optional directory containing vocab_map.safetensors.
mtp_base – When True, export the standard Qwen3.5 text model as the MTP base variant.
mtp_draft – When True, build the dedicated Qwen3.5 dense MTP draft model from the base checkpoint config.
tp_size – Tensor-parallel world size. When >1 the config is reduced to per-rank shapes via ModelConfig.for_rank(), and weights are sharded on assignment. Default 1 = no TP.
tp_rank – This rank’s index in [0, tp_size).
mtp_tree_base – When True, MTP base export adds DDTree parent/depth metadata inputs for Qwen3.5 hybrid state execution (MTP tree drafting).
dflash_base – When True, export as DFlash base model.
dflash_tree_base – When True, add DDTree parent/depth metadata inputs for Qwen3.5 hybrid state execution.
dflash_draft – When True, build the DFlash draft model.
dflash_draft_dir – Path to the DFlash draft checkpoint directory.
jetspec_base – When True, export as JetSpec base model.
jetspec_tree_base – When True, add DDTree parent/depth metadata inputs for JetSpec branching-tree verification.
jetspec_draft – When True, build the JetSpec draft model.
jetspec_draft_dir – Path to the JetSpec draft checkpoint directory.
dspark_base – When True, export as DSpark base model.
dspark_draft – When True, build the DSpark draft backbone model.
dspark_draft_dir – Path to the DSpark draft checkpoint directory.
gemma4_mtp_base – Export a Gemma4 target checkpoint as the base engine for paired Gemma4 MTP.
gemma4_mtp_draft – Export/load a paired Gemma4 assistant checkpoint.
gemma4_kv_sharing_map – Validated assistant-layer to target-layer map for Gemma4 MTP draft runtime config.
gemma4_target_kv_cache_quant – Target/base KV-cache quantization mode inherited by Gemma4 MTP draft inputs that alias target KV cache.
num_decoder_layers – When set, truncate the model to only the first N decoder layers (few-layer numeric validation). Only supported for the plain default CausalLM path (e.g. Qwen3); rejected for eagle/mtp/dflash/jetspec/dspark and registered non-default variants. The checkpoint’s extra-layer weights are simply skipped by the loader.

Returns:

Loaded nn.Module in eval mode.

class tensorrt_edgellm.ModelConfig( model_type: str, hidden_size: int, num_hidden_layers: int, num_attention_heads: int, num_key_value_heads: int, intermediate_size: int, head_dim: int, rms_norm_eps: float, vocab_size: int, rope_theta: float, max_position_embeddings: int, default_attention_scale: float, rope_scaling: dict | None = None, original_max_position_embeddings: int | None = None, partial_rotary_factor: float = 1.0, global_head_dim: int | None = None, num_global_key_value_heads: int | None = None, hidden_activation: str = 'silu', num_code_groups: int = 0, sliding_rope_config: dict | None = None, full_rope_config: dict | None = None, has_qk_norm: bool = False, has_value_norm: bool = False, attention_bias: bool = False, attention_scaling: float | None = None, attention_k_eq_v: bool = False, encoder_layer_scalars: List[float] = <factory>, decoder_layer_scalars: List[float] = <factory>, self_conditioning_size: int = 0, diffusion: DiffusionConfig | None = None, embedding_scale: float = 1.0, final_logit_softcapping: float | None = None, torch_dtype: str = 'bfloat16', tie_word_embeddings: bool = False, sliding_window_size: int = -1, skip_softmax_scale_factor: float = 0.0, use_vision_bidirectional_attention: bool = False, layer_types: List[str] = <factory>, attention_layer_types: List[str] = <factory>, num_deepstack_features: int = 0, accept_hidden_layer: int = -1, quant: QuantConfig = <factory>, mamba_cfg: MambaConfig | None = None, gdn_cfg: GdnConfig | None = None, attn_output_gate: bool = False, mtp_num_hidden_layers: int | None = None, mtp_use_dedicated_embeddings: bool = False, mtp_hybrid_override_pattern: str | None = None, mtp_layer_types: List[str] = <factory>, mtp_base: bool = False, root_model_type: str = '', raw_layer_types: List[str] = <factory>, rope_parameters: dict | None = None, backbone_hidden_size: int = 0, gemma4_mtp_base: bool = False, gemma4_mtp_draft: bool = False, assistant_hidden_size: int = 0, shares_target_kv: bool = False, has_own_kv_cache: bool = True, constant_draft_positions: bool = False, returns_feedback_hidden: bool = False, use_ordered_embeddings: bool = False, num_centroids: int = 0, centroid_intermediate_top_k: int = 0, sparse_logits_enabled: bool = False, kv_sharing_map: List[dict] = <factory>, mtp_tree_base: bool = False, draft_vocab_size: int | None = None, target_hidden_size: int | None = None, is_eagle3_draft_flag: bool = False, eagle3_target_layer_ids: List[int] = <factory>, eagle_base: bool = False, dflash_base: bool = False, dflash_tree_base: bool = False, is_dflash_draft_flag: bool = False, dflash_target_layer_ids: List[int] = <factory>, dflash_block_size: int = 16, dflash_mask_token_id: int = 248070, dflash_fc_native_precision: bool = False, jetspec_base: bool = False, jetspec_tree_base: bool = False, is_jetspec_draft_flag: bool = False, jetspec_target_layer_ids: List[int] = <factory>, jetspec_block_size: int = 16, jetspec_mask_token_id: int = 151669, jetspec_causal_head: bool = False, dspark_base: bool = False, is_dspark_draft_flag: bool = False, dspark_target_layer_ids: List[int] = <factory>, dspark_block_size: int = 7, dspark_mask_token_id: int = 151669, dspark_enable_confidence_head: bool = False, dspark_confidence_head_with_markov: bool = False, dspark_markov_head_type: str = '', dspark_markov_rank: int = 0, num_experts: int = 0, n_routed_experts: int = 0, num_experts_per_tok: int = 0, moe_intermediate_size: int = 0, moe_shared_expert_intermediate_size: int = 0, moe_latent_size: int | None = None, routed_scaling_factor: float = 1.0, n_group: int = 1, topk_group: int = 1, decoder_sparse_step: int = 1, mlp_only_layers: List[int] = <factory>, norm_topk_prob: bool = True, reduced_vocab_size: int | None = None, hidden_size_per_layer_input: int = 0, vocab_size_per_layer_input: int = 0, num_kv_shared_layers: int = 0, use_double_wide_mlp: bool = False, enable_moe_block: bool = False, mapping: Mapping = <factory>, )[source]#

Bases: object

Flat model hyper-parameter config consumed by module builders.

model_type: str#

hidden_size: int#

num_hidden_layers: int#

num_attention_heads: int#

num_key_value_heads: int#

intermediate_size: int#

head_dim: int#

rms_norm_eps: float#

vocab_size: int#

rope_theta: float#

max_position_embeddings: int#

default_attention_scale: float#

rope_scaling: dict | None = None#

original_max_position_embeddings: int | None = None#

partial_rotary_factor: float = 1.0#

hidden_activation: str = 'silu'#

num_code_groups: int = 0#

sliding_rope_config: dict | None = None#

full_rope_config: dict | None = None#

has_qk_norm: bool = False#

has_value_norm: bool = False#

attention_bias: bool = False#

attention_scaling: float | None = None#

global_head_dim: int | None = None#

num_global_key_value_heads: int | None = None#

attention_k_eq_v: bool = False#

encoder_layer_scalars: List[float]#

decoder_layer_scalars: List[float]#

self_conditioning_size: int = 0#

diffusion: DiffusionConfig | None = None#

embedding_scale: float = 1.0#

final_logit_softcapping: float | None = None#

torch_dtype: str = 'bfloat16'#

tie_word_embeddings: bool = False#

sliding_window_size: int = -1#

skip_softmax_scale_factor: float = 0.0#

use_vision_bidirectional_attention: bool = False#

layer_types: List[str]#

attention_layer_types: List[str]#

num_deepstack_features: int = 0#

accept_hidden_layer: int = -1#

quant: QuantConfig#

mamba_cfg: MambaConfig | None = None#

gdn_cfg: GdnConfig | None = None#

attn_output_gate: bool = False#

mtp_num_hidden_layers: int | None = None#

mtp_use_dedicated_embeddings: bool = False#

mtp_hybrid_override_pattern: str | None = None#

mtp_layer_types: List[str]#

mtp_base: bool = False#

root_model_type: str = ''#

raw_layer_types: List[str]#

rope_parameters: dict | None = None#

backbone_hidden_size: int = 0#

gemma4_mtp_base: bool = False#

gemma4_mtp_draft: bool = False#

assistant_hidden_size: int = 0#

shares_target_kv: bool = False#

has_own_kv_cache: bool = True#

constant_draft_positions: bool = False#

returns_feedback_hidden: bool = False#

use_ordered_embeddings: bool = False#

num_centroids: int = 0#

centroid_intermediate_top_k: int = 0#

sparse_logits_enabled: bool = False#

kv_sharing_map: List[dict]#

mtp_tree_base: bool = False#

draft_vocab_size: int | None = None#

target_hidden_size: int | None = None#

is_eagle3_draft_flag: bool = False#

eagle3_target_layer_ids: List[int]#

eagle_base: bool = False#

dflash_base: bool = False#

dflash_tree_base: bool = False#

is_dflash_draft_flag: bool = False#

dflash_target_layer_ids: List[int]#

dflash_block_size: int = 16#

dflash_mask_token_id: int = 248070#

dflash_fc_native_precision: bool = False#

jetspec_base: bool = False#

jetspec_tree_base: bool = False#

is_jetspec_draft_flag: bool = False#

jetspec_target_layer_ids: List[int]#

jetspec_block_size: int = 16#

jetspec_mask_token_id: int = 151669#

jetspec_causal_head: bool = False#

dspark_base: bool = False#

is_dspark_draft_flag: bool = False#

dspark_target_layer_ids: List[int]#

dspark_block_size: int = 7#

dspark_mask_token_id: int = 151669#

dspark_enable_confidence_head: bool = False#

dspark_confidence_head_with_markov: bool = False#

dspark_markov_head_type: str = ''#

dspark_markov_rank: int = 0#

num_experts: int = 0#

n_routed_experts: int = 0#

num_experts_per_tok: int = 0#

moe_intermediate_size: int = 0#

moe_shared_expert_intermediate_size: int = 0#

moe_latent_size: int | None = None#

routed_scaling_factor: float = 1.0#

n_group: int = 1#

topk_group: int = 1#

decoder_sparse_step: int = 1#

mlp_only_layers: List[int]#

norm_topk_prob: bool = True#

reduced_vocab_size: int | None = None#

hidden_size_per_layer_input: int = 0#

vocab_size_per_layer_input: int = 0#

num_kv_shared_layers: int = 0#

use_double_wide_mlp: bool = False#

enable_moe_block: bool = False#

mapping: Mapping#

property tp_size: int#

property tp_rank: int#

property is_eagle3_draft: bool#

property is_mtp_draft: bool#: True for a derived MTP draft config built from a base checkpoint.

property is_gemma4_mtp_draft: bool#: True for a paired Gemma4 assistant draft checkpoint.

property is_diffusion_gemma: bool#

property is_dflash_draft: bool#

property is_jetspec_draft: bool#

property is_dspark_draft: bool#

property ple_enabled: bool#: True when Gemma4 per-layer embeddings are enabled.

property eagle3_target_hidden_size: int#

property eagle3_num_target_layers: int#

property is_hybrid: bool#

property is_nemotron_h: bool#

property num_attn_layers: int#

Total attention layers (includes Gemma4 sliding/full variants).

Note: layers may have heterogeneous head dims — use per-layer configs (kv_layer_configs) for allocation, not this count alone.

property num_mamba_layers: int#

property num_gdn_layers: int#

property num_mlp_layers: int#

property num_moe_layers: int#

property use_dual_rope: bool#

property compute_dtype: torch.dtype#

for_rank( rank: int, world: int, ) → ModelConfig[source]#

Return a per-rank copy of this config for TP.

Divides head and intermediate sizes by world so each rank’s model carries per-rank shapes.

Usage:

cfg = ModelConfig.from_pretrained(: path, lambda head_dim: 1.0 / (float(head_dim)**0.5)

).for_rank(rank, world) model = CausalLM(cfg) load_weights(model, path, mapping=cfg.mapping)

classmethod from_pretrained( model_dir: str, default_attention_scale: Callable[[int], float], ) → ModelConfig[source]#

Load a ModelConfig from a checkpoint directory.

Loads architecture hyper-parameters via AutoConfig (see checkpoint_utils.load_checkpoint_config_dicts()) and then either hf_quant_config.json or the embedded quantization_config block to determine the quantisation scheme.

default_attention_scale is a required model-family callable accepting head_dim. has_qk_norm is auto-detected by scanning the safetensors key index for .q_norm.weight entries; no model-type assumptions are made here.

__init__( model_type: str, hidden_size: int, num_hidden_layers: int, num_attention_heads: int, num_key_value_heads: int, intermediate_size: int, head_dim: int, rms_norm_eps: float, vocab_size: int, rope_theta: float, max_position_embeddings: int, default_attention_scale: float, rope_scaling: dict | None = None, original_max_position_embeddings: int | None = None, partial_rotary_factor: float = 1.0, global_head_dim: int | None = None, num_global_key_value_heads: int | None = None, hidden_activation: str = 'silu', num_code_groups: int = 0, sliding_rope_config: dict | None = None, full_rope_config: dict | None = None, has_qk_norm: bool = False, has_value_norm: bool = False, attention_bias: bool = False, attention_scaling: float | None = None, attention_k_eq_v: bool = False, encoder_layer_scalars: List[float] = <factory>, decoder_layer_scalars: List[float] = <factory>, self_conditioning_size: int = 0, diffusion: DiffusionConfig | None = None, embedding_scale: float = 1.0, final_logit_softcapping: float | None = None, torch_dtype: str = 'bfloat16', tie_word_embeddings: bool = False, sliding_window_size: int = -1, skip_softmax_scale_factor: float = 0.0, use_vision_bidirectional_attention: bool = False, layer_types: List[str] = <factory>, attention_layer_types: List[str] = <factory>, num_deepstack_features: int = 0, accept_hidden_layer: int = -1, quant: QuantConfig = <factory>, mamba_cfg: MambaConfig | None = None, gdn_cfg: GdnConfig | None = None, attn_output_gate: bool = False, mtp_num_hidden_layers: int | None = None, mtp_use_dedicated_embeddings: bool = False, mtp_hybrid_override_pattern: str | None = None, mtp_layer_types: List[str] = <factory>, mtp_base: bool = False, root_model_type: str = '', raw_layer_types: List[str] = <factory>, rope_parameters: dict | None = None, backbone_hidden_size: int = 0, gemma4_mtp_base: bool = False, gemma4_mtp_draft: bool = False, assistant_hidden_size: int = 0, shares_target_kv: bool = False, has_own_kv_cache: bool = True, constant_draft_positions: bool = False, returns_feedback_hidden: bool = False, use_ordered_embeddings: bool = False, num_centroids: int = 0, centroid_intermediate_top_k: int = 0, sparse_logits_enabled: bool = False, kv_sharing_map: List[dict] = <factory>, mtp_tree_base: bool = False, draft_vocab_size: int | None = None, target_hidden_size: int | None = None, is_eagle3_draft_flag: bool = False, eagle3_target_layer_ids: List[int] = <factory>, eagle_base: bool = False, dflash_base: bool = False, dflash_tree_base: bool = False, is_dflash_draft_flag: bool = False, dflash_target_layer_ids: List[int] = <factory>, dflash_block_size: int = 16, dflash_mask_token_id: int = 248070, dflash_fc_native_precision: bool = False, jetspec_base: bool = False, jetspec_tree_base: bool = False, is_jetspec_draft_flag: bool = False, jetspec_target_layer_ids: List[int] = <factory>, jetspec_block_size: int = 16, jetspec_mask_token_id: int = 151669, jetspec_causal_head: bool = False, dspark_base: bool = False, is_dspark_draft_flag: bool = False, dspark_target_layer_ids: List[int] = <factory>, dspark_block_size: int = 7, dspark_mask_token_id: int = 151669, dspark_enable_confidence_head: bool = False, dspark_confidence_head_with_markov: bool = False, dspark_markov_head_type: str = '', dspark_markov_rank: int = 0, num_experts: int = 0, n_routed_experts: int = 0, num_experts_per_tok: int = 0, moe_intermediate_size: int = 0, moe_shared_expert_intermediate_size: int = 0, moe_latent_size: int | None = None, routed_scaling_factor: float = 1.0, n_group: int = 1, topk_group: int = 1, decoder_sparse_step: int = 1, mlp_only_layers: List[int] = <factory>, norm_topk_prob: bool = True, reduced_vocab_size: int | None = None, hidden_size_per_layer_input: int = 0, vocab_size_per_layer_input: int = 0, num_kv_shared_layers: int = 0, use_double_wide_mlp: bool = False, enable_moe_block: bool = False, mapping: Mapping = <factory>, ) → None#

class tensorrt_edgellm.QuantConfig( quant_type: str = 'fp16', group_size: int = 1, gptq_zero_point_offset: int = 1, kv_cache_quant: str | None = None, visual_mha_quant: str | None = None, excluded: List[str] = <factory>, layer_overrides: dict = <factory>, is_mixed_precision: bool = False, )[source]#

Bases: object

Quantization parameters extracted from the checkpoint config.

quant_type: str = 'fp16'#

group_size: int = 1#

gptq_zero_point_offset: int = 1#

kv_cache_quant: str | None = None#

visual_mha_quant: str | None = None#

excluded: List[str]#

layer_overrides: dict#

is_mixed_precision: bool = False#

property is_quantized: bool#

property uses_nvfp4_weights: bool#: True if any linear uses NVFP4 weights (dominant quant or layer override).

property uses_mxfp8_weights: bool#: True if any linear uses MXFP8 weights (dominant quant or layer override).

__init__( quant_type: str = 'fp16', group_size: int = 1, gptq_zero_point_offset: int = 1, kv_cache_quant: str | None = None, visual_mha_quant: str | None = None, excluded: List[str] = <factory>, layer_overrides: dict = <factory>, is_mixed_precision: bool = False, ) → None#

tensorrt_edgellm.export_onnx( model: CausalLM, output_path: str, model_dir: str = '', fp8_embedding: bool = False, reduced_vocab_dir: str = '', externalize_weights=None, config_filename: str = 'config.json', write_shared_artifacts: bool = True, ) → None[source]#

Export model to ONNX using the dynamo exporter.

Writes model.onnx, model.onnx.data, the runtime config (named config_filename), embedding.safetensors, and any tokenizer files present in model_dir to the same output directory.

Parameters:

model – A CausalLM with weights loaded.
output_path – Destination .onnx file path.
model_dir – Checkpoint directory (for tokenizer file copying). If empty, tokenizer files are skipped.
fp8_embedding – Quantize embedding.safetensors to FP8 E4M3 with per-row block scales.
reduced_vocab_dir – Directory containing vocab_map.safetensors when reduced vocabulary is enabled.
externalize_weights – Iterable of weight kinds to expose as fixed-shape ONNX inputs and save to safetensors external weight files. Supported kinds: int4_ffn, int4_moe, nvfp4_moe, lm_head, and all.
config_filename – Filename for the runtime config beside the ONNX. Use "config.json" for single-device exports or "config_world{N}.json" for multi-rank exports.
write_shared_artifacts – Emit shared embedding/tokenizer files. Set to False on non-rank-0 per-rank exports to avoid redundant rewrites of identical sidecar files.

tensorrt_edgellm.load_checkpoint_config_dicts( model_dir: str, ) → Tuple[Dict[str, Any], Dict[str, Any]][source]#

Return (root_dict, llm_dict) from the checkpoint config.

Tries AutoConfig.from_pretrained first (handles registered HF model types). Falls back to reading config.json directly for custom / not-yet-registered model types (e.g. qwen3_asr, qwen3_tts).

For multimodal models (e.g. Qwen2.5-VL, Qwen3-ASR), the LLM text config is promoted out of the nested sub-object by _promote_llm_subconfig(). Any fields lost during promotion are patched back from the raw JSON.

tensorrt_edgellm.load_config_dict(model_dir: str) → Dict[str, Any][source]#: Return only the promoted LLM config dict.

tensorrt_edgellm.load_weights( model: torch.nn.Module, model_dir: str, device: str = 'cpu', key_remap: Callable[[str], str | None] | None = None, key_prefix: str | None = None, pre_repack_hook: Callable[[torch.nn.Module], None] | None = None, mapping: Mapping | None = None, ) → None[source]#

Load all safetensors weights from model_dir into model in-place.

Parameters:

model – Module built by from_pretrained().
model_dir – Checkpoint directory that contains safetensors files.
device – Target device (e.g. "cpu", "cuda:0"). Tensors are moved here after loading.
key_remap – Optional callable (key: str) -> Optional[str]. Called on each checkpoint key after stripping the prefix. Return a new key to remap, the original key unchanged, or None to skip the tensor entirely.
key_prefix – Explicit checkpoint key prefix to strip (e.g. "talker." or "talker.code_predictor."). When provided, only keys starting with this prefix are loaded and auto-detection via _detect_key_prefix() is skipped.
pre_repack_hook – Optional callback invoked after raw checkpoint tensors are loaded and before quantized weights are repacked.
mapping – Parallel-placement config (default = no TP). Drives _shard_for_module() to slice each NVFP4 weight/scale to its per-rank shard before assignment. Must match ModelConfig.mapping used to build model.

tensorrt_edgellm.register_model( model_type: str, model_class: Type[torch.nn.Module], default_attention_scale: Callable[[int], float], ) → None[source]#

When AutoModel.from_pretrained() encounters a checkpoint whose model_type field equals model_type, it instantiates model_class instead of the built-in CausalLM.

Parameters:

model_type – Value of model_type in the checkpoint config.json.
model_class – nn.Module subclass; must accept a single ModelConfig as its constructor argument.
default_attention_scale – Function returning this family’s default for a given attention head dimension.