easydel.layers.caching.transformer.transformer_cache

easydel.layers.caching.transformer.transformer_cache#

class easydel.layers.caching.transformer.transformer_cache.TransformerCache(views: 'tp.List[tp.Optional[TransformerCacheView]]')[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(metadata: TransformerCacheMetaData, mesh: Mesh, quantizer: Optional[object] = None, dtype: Optional[dtype] = None, key_values_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)#

to_dict()#: Convert the instance to a dictionary for JSON serialization.

to_json(**kwargs)#: Convert the instance to a JSON string.

views: List[Optional[TransformerCacheView]]#

class easydel.layers.caching.transformer.transformer_cache.TransformerCacheMetaData(partition_axis: PartitionAxis, batch_size: int, sequence_length: int, num_hidden_layers: int, num_heads: Optional[int], head_dim: Optional[int], key_heads: Optional[int], value_heads: Optional[int], key_dim: Optional[int], value_dim: Optional[int], update_causal_mask: bool, create_attention_bias: bool)[source]#

Bases: BaseCacheMetadata

Metadata for transformer cache configuration.

batch_size: int#

classmethod create(partition_axis: PartitionAxis, batch_size: int, sequence_length: int, num_hidden_layers: int, num_heads: Optional[int] = None, head_dim: Optional[int] = None, key_heads: Optional[int] = None, value_heads: Optional[int] = None, key_dim: Optional[int] = None, value_dim: Optional[int] = None, update_causal_mask: bool = True, create_attention_bias: bool = True) → TransformerCacheMetaData[source]#

Create a TransformerCacheMetaData instance with validation.

Parameters

partition_axis – Partition axis.
batch_size – Size of the batch.
sequence_length – Length of the sequence.
num_hidden_layers – number of hidden layers.
num_heads – Number of attention heads.
head_dim – Dimension of each head.
key_heads – Number of key heads.
value_heads – Number of value heads.
key_dim – Dimension of keys.
value_dim – Dimension of values.
update_causal_mask – Whether to update causal mask.
create_attention_bias – Whether to create attention bias.

Returns

TransformerCacheMetaData instance

Raises

ValueError – If required parameters are missing or invalid.

create_attention_bias: bool#

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: Optional[int]#

key_dim: Optional[int]#

key_heads: Optional[int]#

num_heads: Optional[int]#

num_hidden_layers: int#

partition_axis: PartitionAxis#

replace(**kwargs)#

sequence_length: int#

to_dict()#: Convert the instance to a dictionary for JSON serialization.

to_json(**kwargs)#: Convert the instance to a JSON string.

update_causal_mask: bool#

value_dim: Optional[int]#

value_heads: Optional[int]#

class easydel.layers.caching.transformer.transformer_cache.TransformerCacheView(key: 'tp.Union[cx.Array, ImplicitArray]', value: 'tp.Union[cx.Array, ImplicitArray]', index: 'tp.Union[cx.Array, ImplicitArray]', metadata: 'TransformerCacheMetaData', layer_index: 'tp.Optional[int]' = None)[source]#

Bases: BaseCacheView

concatenate_to_cache(query: Union[Array, ndarray, bool, number], key: Union[Array, ndarray, bool, number], value: Union[Array, ndarray, bool, number], cache_metadata: Optional[TransformerMetadata], attention_mask: Union[Array, ndarray, bool, number], kv_sharding: NamedSharding, quantizer: object, causal_mask: Optional[Union[Array, ndarray, bool, number]] = None, token_type_ids: Optional[Union[Array, ndarray, bool, number]] = None) → Tuple[Union[Array, ndarray, bool, number], Union[Array, ndarray, bool, number], Union[Array, ndarray, bool, number]][source]#

Updates the KV cache functionally and returns the updated tensors along with the appropriate attention mask.

Parameters

query – Current query states.
key – Current key states to add to the cache.
value – Current value states to add to the cache.
cache_metadata – Optional metadata. If provided and contains slot/length info, enables pooled caching.
attention_mask – Base attention mask.
kv_sharding – NamedSharding spec for the cache tensors.
quantizer – Quantizer for the cache.
causal_mask – Optional causal mask.
token_type_ids – Optional token type IDs for segment masking.

Returns

Updated key cache tensor (functional update).
Updated value cache tensor (functional update).
Final attention mask to be used (either original or calculated).

Return type

Tuple[Array, Array, Array]

classmethod from_dict(data)#: Create an instance from a dictionary (deserialization).

classmethod from_json(json_str)#: Create an instance from a JSON string.

index: Union[Array, ndarray, bool, number, ImplicitArray]#

classmethod init(metadata: TransformerCacheMetaData, quantizer: object, key_values_partition_specs: PartitionSpec, dtype: dtype, mesh: Mesh, 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

key: Union[Array, ndarray, bool, number, ImplicitArray]#

layer_index: Optional[int] = None#

metadata: TransformerCacheMetaData#

replace(**kwargs)#

to_dict()#: Convert the instance to a dictionary for JSON serialization.

to_json(**kwargs)#: Convert the instance to a JSON string.

value: Union[Array, ndarray, bool, number, ImplicitArray]#

class easydel.layers.caching.transformer.transformer_cache.TransformerMetadata[source]#

Bases: BaseRunTimeMetadata

holds optional metadata for attention runtime

classmethod from_dict(data)#: Create an instance from a dictionary (deserialization).

classmethod from_json(json_str)#: Create an instance from a JSON string.

replace(**kwargs)#

to_dict()#: Convert the instance to a dictionary for JSON serialization.

to_json(**kwargs)#: Convert the instance to a JSON string.

easydel.layers.caching.transformer.transformer_cache

Contents

easydel.layers.caching.transformer.transformer_cache#