easydel.inference.vinference.utilities

class easydel.inference.vinference.utilities.SampleState(current_length: Union[Array, NamedSharding], sequences: Union[Array, NamedSharding], running_token: Union[Array, NamedSharding], is_sequence_finished: Union[Array, NamedSharding], prng_key: Union[PRNGKey, NamedSharding], model_kwargs: Union[Dict[str, Array], NamedSharding], tokens_per_second: Optional[float] = -inf, generated_tokens: Optional[int] = 0, padded_length: Optional[int] = 0, _time_spent_computing: Optional[float] = 0.0, _compile_config: Optional[vInferencePreCompileConfig] = None)

Bases: object

Represents the state of the sampling process during token generation within the vInference engine.

This class encapsulates all necessary information to pause and resume the generation loop. It tracks the progress of generation, including the tokens generated so far, the current position, completion status, random number generator state, and any model-specific state (like attention caches).

current_length

The current length of the generated sequences (number of tokens generated so far).

Type

Union[jax.Array, jaxlib.xla_extension.NamedSharding]

sequences

The tensor holding the generated token IDs for each sequence in the batch. Shape: (batch_size, max_sequence_length).

Type

Union[jax.Array, jaxlib.xla_extension.NamedSharding]

running_token

The most recently generated token for each sequence. Used as input for the next step. Shape: (batch_size, 1).

Type

Union[jax.Array, jaxlib.xla_extension.NamedSharding]

is_sequence_finished

A boolean tensor indicating whether each sequence in the batch has reached an end-of-sequence (EOS) token or the maximum generation length. Shape: (batch_size,).

Type

Union[jax.Array, jaxlib.xla_extension.NamedSharding]

prng_key

The JAX pseudo-random number generator key used for stochastic sampling.

Type

Union[jax._src.random.PRNGKey, jaxlib.xla_extension.NamedSharding]

model_kwargs

A dictionary containing any additional arguments required by the model for the next generation step (e.g., attention cache/past key-values). The structure depends on the specific model implementation.

Type

Union[Dict[str, jax.Array], jaxlib.xla_extension.NamedSharding]

tokens_per_second

Estimated generation speed in tokens per second. Defaults to -inf.

Type

Optional[float]

generated_tokens

The total count of tokens generated across all sequences in the current generation process up to this state. Defaults to 0.

Type

Optional[int]

padded_length

The target length to which sequences are padded. This might be different from max_sequence_length in some scenarios. Defaults to 0.

Type

Optional[int]

_time_spent_computing

Internal tracker for the cumulative computation time spent to reach this state. Defaults to 0.0.

Type

Optional[float]

_compile_config

The vInferencePreCompileConfig instance used for pre-compiling the functions associated with this generation state. Defaults to None.

Type

Optional[easydel.inference.vinference.utilities.vInferencePreCompileConfig]

current_length: Union[Array, NamedSharding]

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

generated_tokens: Optional[int] = 0

is_sequence_finished: Union[Array, NamedSharding]

model_kwargs: Union[Dict[str, Array], NamedSharding]

padded_length: Optional[int] = 0

prng_key: Union[PRNGKey, NamedSharding]

replace(**kwargs)

Creates a new instance with specified fields replaced.

running_token: Union[Array, NamedSharding]

sequences: Union[Array, NamedSharding]

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.

tokens_per_second: Optional[float] = -inf

easydel.inference.vinference.utilities.create_sampling_step(sampling_params: SamplingParams, eos_token_id: Array, pad_token_id: Array) → Callable

Creates a callable function that performs a single step of token generation (sampling).

This factory function returns a sampling_step function tailored with the provided logit processors/warpers and token IDs. The returned function is designed to be used within a generation loop (e.g., jax.lax.scan).

Parameters

• sampling_params – A SamplingParams.

• eos_token_id – A JAX array containing the token ID(s) representing the end-of-sequence. Generation stops for a sequence once an EOS token is sampled.

• pad_token_id – The JAX array representing the padding token ID. Once a sequence is finished (EOS sampled), subsequent steps will generate this token.

Returns

A callable function sampling_step(graphdef, graphstate, graphother, state) which takes the model’s NNX graph components (graphdef, graphstate, graphother) and the current SampleState, performs one generation step, and returns the updated SampleState.

class easydel.inference.vinference.utilities.vInferenceConfig(max_new_tokens: int = 64, streaming_chunks: int = 16, num_return_sequences: Optional[Union[int, Dict[int, int]]] = 1, pad_token_id: Optional[int] = None, bos_token_id: Optional[int] = None, eos_token_id: Optional[Union[int, List[int]]] = None, partition_rules: Optional[Tuple[Tuple[str, Any]]] = None, partition_axis: Optional[PartitionAxis] = None, _loop_rows: Optional[int] = None, sampling_params: Optional[SamplingParams] = None)

Bases: object

Configuration class for the vInference engine, controlling the overall generation process.

This class holds parameters that define how the generation loop behaves, including length constraints, token control, sharding strategies, and sampling settings.

max_new_tokens

The maximum number of new tokens to generate, excluding the initial prompt tokens. Defaults to 64.

Type

int

streaming_chunks

The number of generation steps to compile and execute together as a single unit. Larger chunks can improve performance on TPUs by reducing compilation overhead and kernel launch times, but may increase memory usage. Defaults to 16.

Type

int

num_return_sequences

The number of sequences to generate and return. Can be: - An integer: Generate this many sequences for all inputs. - A dictionary mapping precompile hash (from vInferencePreCompileConfig)

to an integer: Generate a specific number of sequences based on the compilation configuration. Defaults to 1.

Type

Optional[Union[int, Dict[int, int]]]

pad_token_id

The token ID used for padding sequences. If None, the model’s default pad token ID might be used, or padding might not be applied.

Type

Optional[int]

bos_token_id

The token ID representing the beginning-of-sequence. May be used implicitly by the model or generation logic.

Type

Optional[int]

eos_token_id

The token ID(s) representing the end-of-sequence. Generation stops for a sequence when one of these tokens is sampled. Can be a single integer or a list/tuple of integers.

Type

Optional[Union[int, List[int]]]

partition_rules

A tuple of custom sharding rules (regex pattern, PartitionSpec) to apply to the model’s parameters and intermediate states (like attention cache). If None, default rules based on partition_axis are generated. Example: ((“.*kernel.*”, PartitionSpec(“fsdp”, None)), …)

Type

Optional[Tuple[Tuple[str, Any]]]

partition_axis

A PartitionAxis object defining the logical names for sharding axes (e.g., ‘batch’, ‘sequence’, ‘head’). Required if partition_rules is None, used to generate default sharding rules.

Type

Optional[eformer.escale.partition.manager.PartitionAxis]

_loop_rows

(Internal) The calculated number of iterations needed in the generation loop based on max_new_tokens and streaming_chunks. Automatically computed in __post_init__.

Type

Optional[int]

sampling_params

A SamplingParams object containing parameters for the sampling process itself (e.g., temperature, top_k, top_p, repetition penalty). If None, a default SamplingParams instance with max_tokens set to max_new_tokens is created in __post_init__.

Type

Optional[easydel.inference.utilities.SamplingParams]

bos_token_id: Optional[int] = None

eos_token_id: Optional[Union[int, List[int]]] = None

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

get_partition_rules(runtime_config: Optional[vInferencePreCompileConfig] = None) → Tuple[Tuple[str, Any], ...]

Generates or retrieves the sharding partition rules for the vInference engine.

If self.partition_rules is already set (custom rules provided), it returns them directly.

Otherwise, it constructs a default set of partition rules based on the axis names defined in self.partition_axis. These default rules aim to provide sensible sharding for common model components: - Input sequences (sequences, running_token) are sharded along batch and sequence axes. - Attention masks and position IDs are sharded similarly. - Past key-value states (attention cache), including common quantized formats

(8-bit, NF4), are sharded across batch, key sequence, head, and attention dimension axes.

• Any parameters/states not matching the specific rules are replicated by default (.*).

Parameters

runtime_config – An optional vInferencePreCompileConfig. Currently unused in the default rule generation but available for potential customization in subclasses or future versions.

Returns

• A regex pattern (string) matching parameter or state names.

• A jax.sharding.PartitionSpec defining how the matched items should be sharded.

Return type

A tuple of partition rules. Each rule is a tuple containing

Raises

AssertionError – If self.partition_rules is None and self.partition_axis is also None, as axis names are required to generate default rules.

max_new_tokens: int = 64

num_return_sequences: Optional[Union[int, Dict[int, int]]] = 1

pad_token_id: Optional[int] = None

partition_axis: Optional[PartitionAxis] = None

partition_rules: Optional[Tuple[Tuple[str, Any]]] = None

replace(**kwargs)

Creates a new instance with specified fields replaced.

sampling_params: Optional[SamplingParams] = None

streaming_chunks: int = 16

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.

class easydel.inference.vinference.utilities.vInferencePreCompileConfig(batch_size: Union[int, List[int]] = 1, prefill_length: Optional[Union[int, List[int]]] = None, vision_included: Union[bool, List[bool]] = False, vision_batch_size: Optional[Union[int, List[int]]] = None, vision_channels: Optional[Union[int, List[int]]] = None, vision_height: Optional[Union[int, List[int]]] = None, vision_width: Optional[Union[int, List[int]]] = None, required_props: Optional[Union[Mapping[str, Dict[str, Any]], List[Mapping[str, Dict[str, Any]]]]] = None)

Bases: object

Configuration class for pre-compiling vInference functions.

This class holds parameters that define the shape and properties of inputs expected by the vInference engine during pre-compilation. It allows specifying different configurations, potentially in lists, to compile for multiple scenarios.

batch_size

Batch size or list of batch sizes for text generation.

Type

Union[int, List[int]]

prefill_length

Prefill sequence length or list of lengths. If None, it might be inferred or not used depending on the context.

Type

Optional[Union[int, List[int]]]

vision_included

Whether vision inputs are included in the model.

Type

Union[bool, List[bool]]

vision_batch_size

Batch size for vision inputs. Only relevant if vision_included is True.

Type

Optional[Union[int, List[int]]]

vision_channels

Number of channels for vision inputs. Only relevant if vision_included is True.

Type

Optional[Union[int, List[int]]]

vision_height

Height of vision inputs. Only relevant if vision_included is True.

Type

Optional[Union[int, List[int]]]

vision_width

Width of vision inputs. Only relevant if vision_included is True.

Type

Optional[Union[int, List[int]]]

required_props

Optional dictionary or list of dictionaries specifying required properties for advanced configuration (e.g., specific model arguments).

Type

Optional[Union[Mapping[str, Dict[str, Any]], List[Mapping[str, Dict[str, Any]]]]]

batch_size: Union[int, List[int]] = 1

classmethod create_optimized_compo(batch_size: Union[int, List[int]] = 1, max_prefill_length: int = 2048, min_prefill_length: int = 64, vision_included: Union[bool, List[bool]] = False, vision_batch_size: Optional[Union[int, List[int]]] = None, vision_channels: Optional[Union[int, List[int]]] = None, vision_height: Optional[Union[int, List[int]]] = None, vision_width: Optional[Union[int, List[int]]] = None, required_props: Optional[Union[Mapping[str, Dict[str, Any]], List[Mapping[str, Dict[str, Any]]]]] = None)

extract() → dict

Converts the configuration instance into a dictionary.

This method is useful for serialization or easily accessing all configuration values.

Returns

A dictionary representation of the vInferencePreCompileConfig instance.

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

get_default_hash() → int

Generates a unique integer hash representing the configuration.

This hash is calculated based on the string representation of all configuration attributes, ensuring that identical configurations produce the same hash. This is crucial for caching compiled functions based on their configuration.

Returns

An integer hash value representing the configuration.

get_standalones() → List[vInferencePreCompileConfig]

Generates a list of standalone configurations from a potentially multi-value config.

If any attribute in the current configuration is a list (indicating multiple scenarios), this method expands the configuration into multiple individual vInferencePreCompileConfig instances. Each resulting instance represents a single, specific compilation scenario.

If an attribute’s list is shorter than the longest list among all attributes, its last element is repeated to ensure all generated configurations have values for all attributes.

If the original configuration is already standalone (no list attributes), this method returns a list containing only the original instance.

Returns

A list of vInferencePreCompileConfig instances, each representing a single, standalone compilation scenario.

prefill_length: Optional[Union[int, List[int]]] = None

replace(**kwargs)

Creates a new instance with specified fields replaced.

required_props: Optional[Union[Mapping[str, Dict[str, Any]], List[Mapping[str, Dict[str, Any]]]]] = None

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.

vision_batch_size: Optional[Union[int, List[int]]] = None

vision_channels: Optional[Union[int, List[int]]] = None

vision_height: Optional[Union[int, List[int]]] = None

vision_included: Union[bool, List[bool]] = False

vision_width: Optional[Union[int, List[int]]] = None