easydel.layers.caching.mamba2.__init__

easydel.layers.caching.mamba2.__init__#

class easydel.layers.caching.mamba2.__init__.Mamba2Cache(views: List[Optional[easydel.layers.caching.mamba2.mamba2_cache.Mamba2CacheView]])[source]#

Bases: BaseCache

classmethod from_dict(data)#

Create an instance from a dictionary (deserialization).

classmethod from_json(json_str)#

Create an instance from a JSON string.

classmethod init_cache(num_hidden_layers: int, metadata: Mamba2CacheMetaData, dtype: Optional[dtype] = None, partition_specs: Optional[PartitionSpec] = None)[source]#

Initialize a complete cache with views for all layers.

Parameters

  • metadata – Configuration metadata

  • *args – Additional positional arguments

  • **kwargs – Additional keyword arguments

Returns

Fully initialized cache instance

classmethod init_empty(num_hidden_layers)[source]#

Initialize an empty cache container.

Parameters

  • *args – Additional positional arguments

  • **kwargs – Additional keyword arguments

Returns

Cache instance with uninitialized views

replace(**kwargs)#
reset() → Mamba2Cache[source]#

Reset all cache views to their initial state.

Returns

Reset MambaCache

to_dict()#

Convert the instance to a dictionary for JSON serialization.

to_json(**kwargs)#

Convert the instance to a JSON string.

update_conv_state(layer_idx: int, new_conv_state: Union[Array, ndarray, bool, number], cache_position: Union[Array, ndarray, bool, number]) → Mamba2Cache[source]#

Update the convolutional state for a specific layer.

Parameters

  • layer_idx – Index of the layer to update

  • new_conv_state – New state to be inserted

  • cache_position – Position in the cache to update

Returns

Updated MambaCache

update_seq(num)[source]#
update_ssm_state(layer_idx: int, new_ssm_state: Union[Array, ndarray, bool, number]) → Mamba2Cache[source]#

Update the SSM state for a specific layer.

Parameters

  • layer_idx – Index of the layer to update

  • new_ssm_state – New SSM state to replace the current one

Returns

Updated MambaCache

views: List[Optional[Mamba2CacheView]]#
class easydel.layers.caching.mamba2.__init__.Mamba2CacheMetaData(partition_axis: PartitionAxis, num_hidden_layers: int, batch_size: int, intermediate_size: int, num_heads: int, head_dim: int, state_size: int, conv_kernel_size: int, n_groups: int)[source]#

Bases: BaseCacheMetadata

Metadata for Mamba2 cache configuration.

batch_size: int#
conv_kernel_size: int#
classmethod create(parition_axis: PartitionAxis, num_hidden_layers: int, batch_size: int, intermediate_size: int, num_heads: int, head_dim: int, state_size: int, conv_kernel_size: int, n_groups: int) → Mamba2CacheMetaData[source]#

Create a Mamba2CacheMetaData instance with validation.

classmethod from_dict(data)#

Create an instance from a dictionary (deserialization).

classmethod from_json(json_str)#

Create an instance from a JSON string.

head_dim: int#
intermediate_size: int#
n_groups: int#
num_heads: int#
num_hidden_layers: int#
partition_axis: PartitionAxis#
replace(**kwargs)#
state_size: int#
to_dict()#

Convert the instance to a dictionary for JSON serialization.

to_json(**kwargs)#

Convert the instance to a JSON string.

class easydel.layers.caching.mamba2.__init__.Mamba2CacheView(conv_states: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, eformer.jaximus._imus.ImplicitArray], ssm_states: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, eformer.jaximus._imus.ImplicitArray], positions: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number], seqlen_offset: int, metadata: easydel.layers.caching.mamba2.mamba2_cache.Mamba2CacheMetaData, layer_index: Optional[int] = None)[source]#

Bases: BaseCacheView

concatenate_to_cache(*args, **kwargs)[source]#

Update cache with new states.

Parameters

  • *args – Typically includes new tensors

  • **kwargs – Additional parameters for cache update

Returns

  • anything

Return type

Tuple containing

conv_states: Union[Array, ndarray, bool, number, ImplicitArray]#
classmethod from_dict(data)#

Create an instance from a dictionary (deserialization).

classmethod from_json(json_str)#

Create an instance from a JSON string.

classmethod init(metadata: Mamba2CacheMetaData, partition_specs: PartitionSpec, dtype: dtype, layer_index: Optional[int] = None)[source]#

Initialize a new cache view instance.

Parameters

  • metadata – Configuration metadata for the cache

  • *args – Additional positional arguments

  • **kwargs – Additional keyword arguments

Returns

Initialized cache view instance

layer_index: Optional[int] = None#
metadata: Mamba2CacheMetaData#
positions: Union[Array, ndarray, bool, number]#
replace(**kwargs)#
reset() → Mamba2CacheView[source]#

Reset both conv and ssm states to zeros.

seqlen_offset: int#
ssm_states: Union[Array, ndarray, bool, number, ImplicitArray]#
to_dict()#

Convert the instance to a dictionary for JSON serialization.

to_json(**kwargs)#

Convert the instance to a JSON string.

update_conv_state(new_conv_state: Union[Array, ndarray, bool, number], cache_position: Union[Array, ndarray, bool, number]) → Mamba2CacheView[source]#

Update the convolutional state of the cache.

update_ssm_state(new_ssm_state: Union[Array, ndarray, bool, number]) → Mamba2CacheView[source]#

Update the SSM state of the cache.

class easydel.layers.caching.mamba2.__init__.Mamba2Metadata[source]#

Bases: BaseRunTimeMetadata

Contents