easydel.layers.attention_operator.modules.__init__

class easydel.layers.attention_operator.modules.__init__.FlashAttn(metadata: AttentionMetadata)

Bases: AttentionImpl

forward_cpu(*args, **kwargs) → AttentionOutput

forward_cuda(q: Array, k: Array, v: Array, mask: Optional[Array] = None, bias: Optional[Array] = None, init_bias: Optional[Callable[[], Array]] = None, causal: bool = False, **ignore) → AttentionOutput

forward_gpu(*args, **kwargs) → AttentionOutput

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

forward_rocm(*args, **kwargs) → AttentionOutput

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

get_impl_metadata() → AttentionMetadata

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

class easydel.layers.attention_operator.modules.__init__.RingAttn(metadata: AttentionMetadata)

Bases: AttentionImpl

forward_cpu(*args, **kwargs) → AttentionOutput

forward_cuda(*args, **kwargs) → AttentionOutput

forward_gpu(*args, **kwargs) → AttentionOutput

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

forward_rocm(*args, **kwargs) → AttentionOutput

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

get_impl_metadata() → AttentionMetadata

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

class easydel.layers.attention_operator.modules.__init__.ScaledDotProductAttn(metadata: AttentionMetadata)

Bases: AttentionImpl

forward_cpu(*args, **kwargs) → AttentionOutput

forward_cuda(q: Array, k: Array, v: Array, mask: Optional[Array] = None, bias: Optional[Array] = None, init_bias: Optional[Callable[[], Array]] = None, causal: bool = False, **ignore) → AttentionOutput

forward_gpu(*args, **kwargs) → AttentionOutput

forward_native(q: Array, k: Array, v: Array, mask: Optional[Array] = None, bias: Optional[Array] = None, init_bias: Optional[Callable[[], Array]] = None, causal: bool = False) → AttentionOutput

forward_rocm(*args, **kwargs) → AttentionOutput

forward_tpu(*args, **kwargs) → AttentionOutput

get_impl_metadata() → AttentionMetadata

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

class easydel.layers.attention_operator.modules.__init__.SplashAttn(metadata: AttentionMetadata)

Bases: AttentionImpl

forward_cpu(*args, **kwargs) → AttentionOutput

forward_cuda(*args, **kwargs) → AttentionOutput

forward_gpu(*args, **kwargs) → AttentionOutput

forward_native(*args, **kwargs) → AttentionOutput

forward_rocm(*args, **kwargs) → AttentionOutput

forward_tpu(q: Array, k: Array, v: Array, mask: Optional[Array] = None, causal: bool = True, **ignore) → AttentionOutput

get_impl_metadata() → AttentionMetadata

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

class easydel.layers.attention_operator.modules.__init__.VanillaAttn(metadata: AttentionMetadata)

Bases: AttentionImpl

forward_cpu(*args, **kwargs) → AttentionOutput

forward_cuda(*args, **kwargs) → AttentionOutput

forward_gpu(*args, **kwargs) → AttentionOutput

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, **ignore) → AttentionOutput

forward_rocm(*args, **kwargs) → AttentionOutput

forward_tpu(*args, **kwargs) → AttentionOutput

get_impl_metadata() → AttentionMetadata

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