easydel.layers.attention_operator.modules.ring

Contents

easydel.layers.attention_operator.modules.ring#

class easydel.layers.attention_operator.modules.ring.Carry(numerator, denominator, max_so_far)[source]#

Bases: NamedTuple

denominator: Array#: Alias for field number 1

max_so_far: Array#: Alias for field number 2

numerator: Array#: Alias for field number 0

class easydel.layers.attention_operator.modules.ring.RingAttn(metadata: AttentionMetadata)[source]#

Bases: AttentionImpl

forward_cpu(*args, **kwargs) → AttentionOutput[source]#

forward_cuda(*args, **kwargs) → AttentionOutput[source]#

forward_gpu(*args, **kwargs) → AttentionOutput[source]#

forward_native(q: Array, k: Array, v: Array, mask: Optional[Array] = None, bias: Optional[Array] = None, init_bias: Optional[Callable[[], Array]] = None, deterministic: bool = False, dropout_rng: Optional[PRNGKey] = None, causal: bool = True, **ignore) → AttentionOutput[source]#

forward_rocm(*args, **kwargs) → AttentionOutput[source]#

forward_tpu(q: Array, k: Array, v: Array, mask: Optional[Array] = None, bias: Optional[Array] = None, init_bias: Optional[Callable[[], Array]] = None, deterministic: bool = False, dropout_rng: Optional[PRNGKey] = None, causal: bool = True, **ignore) → AttentionOutput[source]#

get_impl_metadata() → AttentionMetadata[source]#

classmethod get_impl_name() → Union[str, Tuple[str]][source]#

easydel.layers.attention_operator.modules.ring.blockwise_attn(query, key, value, bias=None, deterministic=True, dropout_rng=None, attn_pdrop=0.0, causal=True, query_chunk_size=2048, key_chunk_size=2048, dtype=<class 'jax.numpy.float32'>, policy=False, precision=None, float32_logits=True, prevent_cse=True)[source]#