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#

vLLM-style inference server for TensorRT Edge-LLM.

Public API:

from experimental.server import LLM, SamplingParams

llm = LLM(model="Qwen/Qwen3-1.7B")
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 = '', *, onnx_dir: str = '', visual_onnx_dir: str = '', audio_onnx_dir: str = '', engine_dir: str = '', multimodal_engine_dir: str = '', visual_engine_dir: str = '', max_input_len: int = 4096, max_batch_size: int = 1, max_kv_cache_capacity: int = 8192, eagle_engine_dir: str = '', draft_top_k: int = 10, draft_step: int = 6, verify_tree_size: int = 60, talker_engine_dir: str = '', code_predictor_engine_dir: str = '', code2wav_engine_dir: str = '', )[source]#

Bases: object

vLLM-style entry point for TensorRT Edge-LLM inference.

Three initialization modes (exactly one of model, onnx_dir, or engine_dir must be provided):

HuggingFace checkpoint — exports ONNX, builds engine, loads:
```
llm = LLM(model="Qwen/Qwen3-1.7B")
```
ONNX directory — builds engine from ONNX, loads:
```
llm = LLM(onnx_dir="/path/to/onnx")
```

Pre-built engine — loads directly:

llm = LLM(engine_dir="/path/to/engine")
llm = LLM(engine_dir="...", multimodal_engine_dir="...")

See experimental.server.engine_layout for the expected directory layouts.

text_capable = True#: Distinguishes full LLM servers from TTS-only ones in the API layer.

__init__( model: str = '', *, onnx_dir: str = '', visual_onnx_dir: str = '', audio_onnx_dir: str = '', engine_dir: str = '', multimodal_engine_dir: str = '', visual_engine_dir: str = '', max_input_len: int = 4096, max_batch_size: int = 1, max_kv_cache_capacity: int = 8192, eagle_engine_dir: str = '', draft_top_k: int = 10, draft_step: int = 6, verify_tree_size: int = 60, talker_engine_dir: str = '', code_predictor_engine_dir: str = '', code2wav_engine_dir: str = '', )[source]#

count_prompt_tokens( messages: List[Dict[str, Any]], *, tools: Sequence[Dict[str, Any]] | None = None, tool_choice: str | Dict[str, Any] | None = None, tool_config: ToolConfig | None = None, enable_thinking: bool = False, ) → int | None[source]#: Best-effort prompt token count via the HF tokenizer: exact for tool-templated requests, within a few tokens for plain ones (HF vs C++ template). Multimodal placeholders are counted once, not expanded, so multimodal prompts are undercounted. None when counting is unavailable.

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, ) → 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, admission_handoff: Any | None = None, ) → Generator[StreamDelta, None, None][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, admission_handoff: Any | None = None, ) → Generator[StreamDelta, None, None][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, *, admission_handoff: Any | None = None, ) → Generator[StreamDelta, None, None][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, *, enable_batching: bool = False, batch_timeout_ms: float = 10.0, max_queue_batch_size: int | None = None, request_queue_size: int | None = None, allowed_local_media_path: str | None = None, ) → None[source]#

Start an OpenAI-compatible HTTP server.

Parameters:

host – Bind address.
port – Bind port.
enable_batching – Batch compatible non-streaming HTTP requests.
batch_timeout_ms – Maximum time to wait for compatible requests.
max_queue_batch_size – Optional cap for queued HTTP micro-batches.
request_queue_size – Max concurrently admitted requests (queued + running) before the server returns backpressure. None uses the server default.
allowed_local_media_path – Directory HTTP clients may reference local media from. Unset rejects bare paths and file:// URLs.

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

property engine_dir: str#: Path to the TensorRT engine directory.

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

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

property omni_capable: bool#: Whether the Omni audio-output stack is loaded.

class experimental.server.TTS( talker_engine_dir: str, code_predictor_engine_dir: str | None = None, code2wav_engine_dir: str | None = None, tokenizer_dir: str = '', model: str | None = None, )[source]#

Bases: object

TTS-only serving for Qwen3-TTS-style engine sets.

Loads Talker + CodePredictor + Code2Wav without a Thinker/text engine. serve() exposes /v1/audio/speech; chat endpoints return 400.

Example:

from experimental.server import TTS

tts = TTS(talker_engine_dir="/engines/qwen3-tts/talker")
tts.serve(port=8000)

code_predictor_engine_dir / code2wav_engine_dir default to the talker directory’s siblings; tokenizer_dir defaults to the talker directory itself (the standard export layout ships tokenizer files there).

text_capable = False#

omni_capable = True#

has_draft_model = False#

__init__( talker_engine_dir: str, code_predictor_engine_dir: str | None = None, code2wav_engine_dir: str | None = None, tokenizer_dir: str = '', model: str | None = None, ) → None[source]#

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

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

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>, )[source]#

Bases: object

Sampling parameters (mirrors vLLM’s 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]#

logit_bias: Dict[int, float]#

__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>, ) → None#

class experimental.server.CompletionOutput( text: str = '', token_ids: List[int] = <factory>, 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]#

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>, 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>, 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]#

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>, 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#

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#

Checkpoint loader and ONNX exporter for causal LMs.

Supports common HF architectures and FP8, NVFP4, INT4 (AWQ / GPTQ), INT8 SmoothQuant, and mixed-precision checkpoints when described by config.json / hf_quant_config.json.

Quick start:

from tensorrt_edgellm import AutoModel, export_onnx

model = AutoModel.from_pretrained("/path/to/checkpoint")
export_onnx(model, "output/model.onnx", model_dir="/path/to/checkpoint")

Config: checkpoint.checkpoint_utils.load_checkpoint_config_dicts() / checkpoint.checkpoint_utils.load_config_dict(). Weights: checkpoint.loader.load_weights(). Export sidecars: checkpoint.checkpoint_utils.write_runtime_artifacts().

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, 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.
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/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.

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', ) → 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_tp{N}_rank{R}.json" for per-rank TP exports so each rank is self-describing.

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.

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, 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#

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_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, 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, 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#

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, excluded: List[str] = <factory>, layer_overrides: dict = <factory>, is_mixed_precision: bool = False, ) → None#

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.