SaveHiddenStatesDecodingConfig

class tensorrt_llm.llmapi.SaveHiddenStatesDecodingConfig(

*,

max_draft_len: Annotated[int, Ge(ge=0)] | None = None,

max_total_draft_tokens: int | None = 1,

speculative_model: str | Path | None = None,

moe_backend: Literal['AUTO', 'CUTLASS', 'CUTEDSL', 'CUTEDSL_FC12', 'TRTLLM', 'DEEPGEMM', 'DENSEGEMM', 'VANILLA', 'TRITON', 'MARLIN', 'MEGAMOE_DEEPGEMM', 'MEGAMOE_CUTEDSL'] | None = None,

max_concurrency: Annotated[int, Gt(gt=0)] | None = None,

draft_len_schedule: dict[int, int] | None = None,

load_format: str | None = None,

acceptance_rate_window_size: Annotated[int, Ge(ge=0)] | None = None,

acceptance_rate_threshold: Annotated[float | None, Ge(ge=0.0), Le(le=1.0)] = None,

use_rejection_sampling: bool = False,

allow_advanced_sampling: bool = False,

advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL,

enable_penalty: bool = False,

decoding_type: Literal['SaveState'] = 'SaveState',

output_directory: str,

write_interval: int = 20,

file_prefix: str = 'data',

eagle3_layers_to_capture: Set[int] | None = None,

)

Bases: DecodingBaseConfig

field acceptance_rate_threshold: float | None = None

The threshold for average true acceptance rate (accepted_draft_tokens / drafted_tokens); speculation will be disabled permanently once the rolling average over the last N speculation-enabled decoding iterations (N = acceptance_rate_window_size) drops below this value.

Constraints:

• ge = 0.0

• le = 1.0

field acceptance_rate_window_size: NonNegativeInt | None = None

The rolling average window size (N) for acceptance rate across speculation-enabled decoding iterations. If not set or set to 0, the feature is disabled. PyTorch backend only.

field advanced_sampling_mode: AdvancedSamplingMode = AdvancedSamplingMode.FULL

Deploy-time specialization of the one-model advanced sampler that skips disabled filter kernels. FULL (default): per-row top_k/top_p/min_p in one fused kernel, the only mode accepting min_p. NO_TOPK: skip top_k. NO_TOPP: skip top_p. NO_TOPK_NO_TOPP: skip both.

field allow_advanced_sampling: bool = False

DEPRECATED: no-op kept for backward compatibility. Will be removed in a future release. Non-greedy sampling is now auto-detected per request; this flag no longer has any effect.

field decoding_type: Literal['SaveState'] = 'SaveState'

field draft_len_schedule: dict[int, int] | None = None

Developer interface: dynamically adjust draft length based on active batch size in runtime.Maps batch size to draft lengths.For example: draft_len_schedule = {4:4, 8:2, 32:1} - Batch sizes 1-4: use draft_len=4 - Batch sizes 5-8: use draft_len=2 - Batch sizes 9-32: use draft_len=1 - Batch sizes 33+: use draft_len=0 (implicit, speculation disabled). Mutually exclusive with max_concurrency since draft_len_schedule implicitly support max concurrency control.

field eagle3_layers_to_capture: Set[int] | None = None

Set of target model layer indices to capture hidden states from for EAGLE3 draft model training. Use -1 to indicate the final post-norm hidden state. If not provided, defaults to capturing 3 intermediate layers plus the post-norm hidden state. When provided, -1 is automatically added if not present.

field enable_penalty: bool = False

If true, enables the occurrence penalties (repetition / presence / frequency) for one-model speculative decoding. Off by default because the penalties need a [num_seq_slots, vocab_size] occurrence-count workspace that is allocated up front (CUDA graphs capture fixed buffer addresses). While off, a request that asks for any of these penalties is rejected at admission rather than silently decoded without them.

field file_prefix: str = 'data'

Prefix for output filenames. Files are saved as ‘<file_prefix>_<iteration>.pt’ containing input_ids and hidden_state tensors.

field load_format: str | None = None

The load format of the speculative model.

field max_concurrency: PositiveInt | None = None

When specified (>0), speculation will be disabled at batch sizes above this value. Otherwise, speculation will always be on. PyTorch backend only. Mutually exclusive with max_concurrency since draft_len_schedule implicitly supports max concurrency control.

field max_draft_len: NonNegativeInt | None = None

The maximum number of draft tokens.

field max_total_draft_tokens: int | None = 1

Internal field, not user-configurable. Fixed to 1 since this mode captures hidden states without draft token generation.

field moe_backend: _MoeBackend | None = None

MoE backend override for a neural draft model or embedded MTP layers on the PyTorch backend. None inherits the target model’s backend. AUTO resolves from the draft checkpoint or embedded MTP layer quantization, and a concrete backend applies only to the draft model or layers. Resolution may fall back based on model, quantization, and hardware support. Replacement-head MTP checkpoints are unsupported because their independent quantization metadata is not loaded. Nemotron-H embedded MTP layers must inherit the target backend because their checkpoint mapper uses a shared backend-dependent layout. Decoding methods without a neural draft model ignore this option.

field output_directory: str [Required]

Directory path where hidden states data files will be saved. The directory is created if it does not exist.

field speculative_model: str | Path | None = None

The speculative (draft) model. Accepts either (1) a HuggingFace Hub model ID (e.g. ‘yuhuili/EAGLE3-LLaMA3.1-Instruct-8B’), which will be automatically downloaded, or (2) a local filesystem path to a downloaded model directory. For one-model MTP, a non-target checkpoint provides either replacement MTP heads or a complete external draft model, depending on the target model implementation. Pointing it at the target checkpoint uses the target’s embedded mtp.* weights.

field use_rejection_sampling: bool = False

If true, enables rejection sampling for one-model speculative decoding paths when the batch contains any non-greedy request. All-greedy batches always take the argmax fast path regardless of this flag. Set to false (default) to use exact-match verification on non-greedy batches. The non-dynamic-tree one-model path requires FlashInfer.

field write_interval: int = 20

Number of requests to process before writing accumulated hidden states to disk. Lower values write more frequently but may impact performance.

__init__(**data: Any) → None

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

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

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

get_runtime_tokens_per_gen_step(

runtime_draft_len: int,

) → int

Total tokens per gen request for the current runtime draft length.

supports_backend(backend: str) → bool

Override if the speculation algorithm does not support a subset of the possible backends.

validator validate_draft_len_schedule_and_sort  »  draft_len_schedule

Validate and sort draft_len_schedule by batch size thresholds.

validator validate_max_concurrency_and_draft_len_schedule_mutually_exclusive  »  all fields

validator validate_rejection_sampling_config  »  all fields

Disable rejection sampling when SA-enhanced configurations are active.

Only silently disable a default-inherited value; an explicit use_rejection_sampling=True is preserved so TorchLlmArgs.validate_speculative_config can raise for the unsupported SA combination.

property is_linear_tree: bool

property needs_separate_draft_weights: bool

Whether draft weights must be loaded from speculative_model.

This includes external draft models and MTP head replacement checkpoints.

property num_capture_layers

Returns the number of layers to save. The following hidden states are saved: - If eagle3_layers_to_capture is None, save the eagle3 base set plus the post norm last hidden state. - Otherwise, save the specified layers plus the post norm last hidden state.

The saved data will contain two tensors, hidden_states and aux_hidden_states. * hidden_states will contain the last post norm state. * aux_hidden_states will contain all other captured layers. The last hidden state will also be included in this tensor if you explicitly captured layer -1.

Note that if you set layers to capture to {-1}, aux_hidden_states won’t exist.

property spec_dec_mode

property tokens_per_gen_step: int

Total tokens per gen request in one spec dec iteration (including golden token).

property uses_external_draft_model: bool

Whether speculative_model contains an external draft model.

property uses_replacement_heads: bool

Whether speculative_model contains replacement MTP heads.