Stores inputs to the `memory_efficient_attention` operators

Holds a CUDAEvent. Why does it need to be a tensor? So that its `gradient` can also hold a CUDAEvent for overlaps in the BW pass

A lower-triangular (aka causal) mask with an additive bias

RMS Normalization layer along the last dimension. This is similar to torch.nn.functional.normalize but with eps being added instead of m

A causal masking. This mask is exactly the same as :attr:`LowerTriangularMask` when there is the same number of queries and keys. Wh

Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalMask`, except that each block is causal. This mask allows for a non-causal prefix

(Experimental feature) Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalCausalMask`. This makes the mask "local" and the attention

(Experimental feature) Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalCausalMask`. This makes the mask "local" and the attention

Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalMask`, except that each block is causal. Queries and Keys are each divided into the

Like :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalCausalLocalAttentionPaddedKeysMask`, except that this is non-causal. A query Q in b

(Internal) Flexible equivalent of _PaddedSeqLenInfo. There are two distinct semantics. (1) For non-paged masks: Represents the divis

(Internal) Represents the division of a dimension into blocks which are padded out to the same total length. For example, to represent

Like :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalCausalWithOffsetPaddedKeysMask`, except with a window size. A query Q in block i ca

Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalCausalMask`, except k/v is gappy. A query Q in block i cannot attend to a key wh

Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalCausalMask`, except an offset on causality is allowed for each block and we support p

Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalMask`, except k/v is gappy. A query Q in block i only attends to a key which is

A block-diagonal mask that can be passed as ``attn_bias`` argument to :attr:`xformers.ops.memory_efficient_attention`. Queries and Keys

Same as :attr:`xformers.ops.fmha.attn_bias.BlockDiagonalMask`, except we support padding for k/v The keys and values are divided into bl

Each of k/v_fp8_scales is an int32 tensor of shape (1, B * Mkv, Hq), or (1, page_size * max_pages_per_lane, Hq) in the paged case. Each i

A local attention mask The query at position :math:`q` can attend the key at position :math:`k` if :math:`q - window\\_left <= k + s <=

A mask that combines both :attr:`LowerTriangularFromBottomRightMask` and local attention. A query whose distance from the final query is

A lower-triangular (aka causal) mask A query Q cannot attend to a key which is farther from the initial key than Q is from the initial q

Same as BlockDiagonalCausalWithOffsetGappyKeysMask, but for paged attention. block_tables has shape [batch_size, max_num_pages] and K/V have

Same as BlockDiagonalCausalWithOffsetPaddedKeysMask, but for paged attention. block_tables has shape [batch_size, max_num_pages] and K/V have

Equivalent BlockDiagonalGappyKeysMask, but for paged attention. block_tables has shape [batch_size, max_num_pages] and K/V have shape [1,

Same as BlockDiagonalPaddedKeysMask, but for paged attention. block_tables has shape [batch_size, max_num_pages] and K/V have shape [1, m

This class is used to represent a partial attention output, which is returned by `memory_efficient_attention_partial_autograd`. Attribut

Generation stats, split by phases.

Dispatcher to automatically select the best operator in :attr:`xformers.ops.swiglu`

Encoding/decoding text using SentencePiece.

tree_choices: definition of the tree, tuples sorted by length, each corresponding to a node. See the docstring of TreeAttnMetadata.from_t

Set up a communication ring and perform fused ops on it Stores the persistent state needed to support a ring of connections between processes

(Internal) Represents the division of a dimension into blocks. For example, to represents a dimension of length 7 divided into three blo

Can be applied to rank based on absolute value

Base class for a custom bias that can be applied \ as the attn_bias argument in :attr:`xformers.ops.memory_efficient_attention`. That

Base class for any attention operator in xFormers See: - :attr:`xformers.ops.fmha.cutlass.FwOp` - :attr:`xformers.ops.fmha.cutlass.BwOp`