easydel.layers.attention_operator.modules.flash#
- class easydel.layers.attention_operator.modules.flash.FlashAttn(metadata: AttentionMetadata)[source]#
Bases:
AttentionImpl- forward_cpu(*args, **kwargs) AttentionOutput[source]#
- 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[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, causal: bool = False, **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, causal: bool = False, **ignore) AttentionOutput[source]#
- get_impl_metadata() AttentionMetadata[source]#