easydel.layers.caching._abstracts

easydel.layers.caching._abstracts#

class easydel.layers.caching._abstracts.BaseCache[source]#

Bases: ABC

Abstract base class for the main cache container.

Manages a sequence of cache views (typically one per layer) and provides initialization methods.

abstract classmethod init_cache(metadata: BaseCacheMetadata, *args, **kwargs) → BaseCache[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

abstract classmethod init_empty(*args, **kwargs) → BaseCache[source]#

Initialize an empty cache container.

Parameters

  • *args – Additional positional arguments

  • **kwargs – Additional keyword arguments

Returns

Cache instance with uninitialized views

views: Sequence[Optional[BaseCacheView]]#
class easydel.layers.caching._abstracts.BaseCacheMetadata[source]#

Bases: ABC

Abstract base class defining the interface for cache metadata.

Concrete implementations should provide: - Required configuration parameters for cache initialization - Validation logic in the create() method - Any metadata needed during cache operations

abstract classmethod create(*args, **kwargs) → BaseCacheMetadata[source]#

Factory method to create validated metadata instance.

Parameters

  • *args – Positional arguments for metadata creation

  • **kwargs – Keyword arguments for metadata creation

Returns

Instance of concrete metadata implementation

Raises

ValueError – If any validation checks fail

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.

replace(**kwargs)#

Creates a new instance with specified fields replaced.

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.layers.caching._abstracts.BaseCacheView[source]#

Bases: ABC

Abstract base class for a single cache view (typically per layer).

Responsible for: - Storing cached key/value states - Tracking current cache position - Updating cache with new states

abstract concatenate_to_cache(*args, **kwargs) → Any[source]#

Update cache with new states.

Parameters

  • *args – Typically includes new tensors

  • **kwargs – Additional parameters for cache update

Returns

  • anything

Return type

Tuple containing

abstract classmethod init(metadata: BaseCacheMetadata, *args, **kwargs) → BaseCacheView[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]#
metadata: BaseCacheMetadata#
class easydel.layers.caching._abstracts.BaseRunTimeMetadata[source]#

Bases: object

Abstract base class for optional runtime metadata used during attention computation.

This can hold dynamic information needed during the forward pass that isn’t known at cache initialization time.

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.

replace(**kwargs)#

Creates a new instance with specified fields replaced.

to_dict() → Dict[str, Any]#

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str#

Serializes the PyTree object to a JSON string.

Contents