easydel.layers.attention_operator.modules.splash

easydel.layers.attention_operator.modules.splash#

class easydel.layers.attention_operator.modules.splash.SplashAttn(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(*args, **kwargs) → AttentionOutput[source]#
forward_rocm(*args, **kwargs) → AttentionOutput[source]#
forward_tpu(q: Array, k: Array, v: Array, mask: Optional[Array] = None, causal: bool = True, **ignore) → AttentionOutput[source]#
get_impl_metadata() → AttentionMetadata[source]#
classmethod get_impl_name() → Union[str, Tuple[str]][source]#
Contents