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, max_image_tokens: int | None = None, max_image_tokens_per_image: int | None = None, draft_top_k: int | None = None, draft_step: int | None = None, verify_tree_size: int | None = None, max_verify_tree_size: int | None = None, max_draft_tree_size: int | None = None, build_options: BuildOptions | None = None, speculative_config: Any | None = None, context_cache_config: ContextCacheConfig | Mapping[str, Any] | None = None, enable_in_flight_batching: bool = False, on_progress: Callable[[str], None] | 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, max_image_tokens: int | None = None, max_image_tokens_per_image: int | None = None, draft_top_k: int | None = None, draft_step: int | None = None, verify_tree_size: int | None = None, max_verify_tree_size: int | None = None, max_draft_tree_size: int | None = None, build_options: BuildOptions | None = None, speculative_config: Any | None = None, context_cache_config: ContextCacheConfig | Mapping[str, Any] | None = None, enable_in_flight_batching: bool = False, on_progress: Callable[[str], None] | None = None, )[source]#

property speech_available: bool#

the Omni stack is in the bundle and the runtime is the blocking one that drives it. Under in-flight batching the engine owns the runtime and serves text only.

Type:: Whether this process can produce speech

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, on_metrics: Callable[[GenerationMetrics], None] | None = None, ) → Iterator[StreamDelta][source]#

Stream generation deltas for a single message list.

Yields StreamDelta objects as tokens are produced. Under in-flight batching the request is submitted to the engine and this thread reads the attached StreamChannel directly; on the blocking path handleRequest runs in a background thread instead. on_metrics receives final native throughput and wall-clock timings after the producer stops, including when the stream is closed. Collection is disabled when on_metrics is omitted. Per-request metrics do not enable aggregate counters or CUDA event profiling. Cancellation or closure before iteration reports empty metrics once, without starting a request. cancel() requests cancellation; close() waits for the blocking worker, or the request engine outcome when metrics are requested.

validate_context( messages: List[Dict[str, Any]], sampling_params: SamplingParams | None = None, ) → int[source]#

Count a text conversation and reject turns outside the engine limits.

Reserve the requested output budget without truncating the prompt. This check does not generate tokens or change the context cache.

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 (or when in-flight batching left the speech stack unloaded).

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 max_kv_cache_capacity: int#: Total token capacity 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_scheduling_metrics()[source]#: Engine scheduling counters as a dict, or None on the blocking path.

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, seed: int | None = None, max_tokens: int = 2048, enable_thinking: bool = False, reasoning_effort: str = '', disable_spec_decode: bool = False, num_logprobs: int = 0, stop: List[str] = <factory>, logit_bias: Dict[int, float]=<factory>, guided_decoding: Tuple[str, str] | None=None, 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#

seed: int | None = None#

max_tokens: int = 2048#

enable_thinking: bool = False#

reasoning_effort: str = ''#

disable_spec_decode: bool = False#

num_logprobs: int = 0#

stop: List[str]#

logit_bias: Dict[int, float]#

guided_decoding: Tuple[str, str] | None = None#

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, seed: int | None = None, max_tokens: int = 2048, enable_thinking: bool = False, reasoning_effort: str = '', disable_spec_decode: bool = False, num_logprobs: int = 0, stop: List[str] = <factory>, logit_bias: Dict[int, float]=<factory>, guided_decoding: Tuple[str, str] | None=None, 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#

class experimental.server.CancellableIterator(iterator, cancel)[source]#

Bases: object

Wrap a generator with cancellation that can interrupt native waits.

Parameters:

iterator – Generator to consume and close.
cancel – Callable that requests producer cancellation without waiting for iteration to finish. It must tolerate repeated calls.

cancel() signals the producer without closing an executing generator. close() signals cancellation once and closes the generator when it is not executing in another thread. The caller must wait for active iteration to finish before releasing resources owned by that generator.

__init__(iterator, cancel) → None[source]#

cancel() → None[source]#: Request cancellation without closing an executing iterator.

close() → None[source]#

class experimental.server.GenerationMetrics( prompt_tokens: int | None = None, processed_prompt_tokens: int | None = None, prompt_time: float = 0.0, generated_tokens: int = 0, time_to_first_token: float | None = None, generation_time: float = 0.0, total_time: float = 0.0, decoded_tokens: int = 0, )[source]#

Bases: object

Native throughput timings and client-observed latency in seconds.

Native counts and phase times are independent of stream consumption. Incomplete native snapshots fall back to received token counts; unavailable phase times remain zero. TTFT and total_time include admission, rendering, and consumer overhead. generated_tokens includes prefill output; decoded_tokens and generation_time cover only work after prefill. Use decode_tokens_per_second for throughput.

prompt_tokens: int | None = None#

processed_prompt_tokens: int | None = None#

prompt_time: float = 0.0#

generated_tokens: int = 0#

time_to_first_token: float | None = None#

generation_time: float = 0.0#

total_time: float = 0.0#

decoded_tokens: int = 0#

property decode_tokens_per_second: float#: Native decode throughput, or zero when decode timing is unavailable.

__init__( prompt_tokens: int | None = None, processed_prompt_tokens: int | None = None, prompt_time: float = 0.0, generated_tokens: int = 0, time_to_first_token: float | None = None, generation_time: float = 0.0, total_time: float = 0.0, decoded_tokens: int = 0, ) → 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_tree_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, low_cpu_mem_usage: bool = False, ) → 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_tree_base – When True, add DDTree parent/depth metadata inputs for hybrid DSpark base verification.
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.
low_cpu_mem_usage – Initialize model tensors on the meta device and materialize them from the checkpoint during load.

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, qk_norm_post_rope: bool = False, has_value_norm: bool = False, attention_bias: bool = False, attention_scaling: float | None = None, attention_k_eq_v: bool = False, attention_sink_bias: 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, output_multiplier: float = 1.0, post_norm_eps: 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_version: DFlashVersion = DFlashVersion.V1, dflash_shift_label: bool = False, dflash_overflow_fp32: bool = False, dflash_fc_native_precision: bool = False, dflash2_target_layer_ids: List[int] = <factory>, dflash2_block_size: int = 8, dflash2_mask_token_id: int = 248070, dflash2_is_causal: bool = False, dflash2_conv_kernel_size: int = 2, dflash2_conv_group_size: int = 16, dflash2_selector_rank: int = 256, dflash2_selector_top_k: int = 16, 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, dspark_tree_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_sample_from_anchor: bool = True, dspark_confidence_head_with_markov: bool = False, dspark_markov_head_type: str = '', dspark_markov_rank: int = 0, dspark_fc_native_precision: bool = False, dspark_causal_proposal: bool = False, 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#

qk_norm_post_rope: 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#

attention_sink_bias: 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#

output_multiplier: float = 1.0#

post_norm_eps: 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_version: DFlashVersion = 1#

dflash_shift_label: bool = False#

dflash_overflow_fp32: bool = False#

dflash_fc_native_precision: bool = False#

dflash2_target_layer_ids: List[int]#

dflash2_block_size: int = 8#

dflash2_mask_token_id: int = 248070#

dflash2_is_causal: bool = False#

dflash2_conv_kernel_size: int = 2#

dflash2_conv_group_size: int = 16#

dflash2_selector_rank: int = 256#

dflash2_selector_top_k: int = 16#

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#

dspark_tree_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_sample_from_anchor: bool = True#

dspark_confidence_head_with_markov: bool = False#

dspark_markov_head_type: str = ''#

dspark_markov_rank: int = 0#

dspark_fc_native_precision: bool = False#

dspark_causal_proposal: bool = False#

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; the HunYuan .query_layernorm.weight spelling is honored only for model types whose loader remaps it (see _detect_has_qk_norm()).

__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, qk_norm_post_rope: bool = False, has_value_norm: bool = False, attention_bias: bool = False, attention_scaling: float | None = None, attention_k_eq_v: bool = False, attention_sink_bias: 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, output_multiplier: float = 1.0, post_norm_eps: 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_version: DFlashVersion = DFlashVersion.V1, dflash_shift_label: bool = False, dflash_overflow_fp32: bool = False, dflash_fc_native_precision: bool = False, dflash2_target_layer_ids: List[int] = <factory>, dflash2_block_size: int = 8, dflash2_mask_token_id: int = 248070, dflash2_is_causal: bool = False, dflash2_conv_kernel_size: int = 2, dflash2_conv_group_size: int = 16, dflash2_selector_rank: int = 256, dflash2_selector_top_k: int = 16, 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, dspark_tree_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_sample_from_anchor: bool = True, dspark_confidence_head_with_markov: bool = False, dspark_markov_head_type: str = '', dspark_markov_rank: int = 0, dspark_fc_native_precision: bool = False, dspark_causal_proposal: bool = False, 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>, layer_group_sizes: dict = <factory>, is_mixed_precision: bool = False, quantize_activations: bool = <factory>, )[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#

layer_group_sizes: dict#

is_mixed_precision: bool = False#

quantize_activations: bool#

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>, layer_group_sizes: dict = <factory>, is_mixed_precision: bool = False, quantize_activations: bool = <factory>, ) → 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.