Test Medusa-style tree attention.

Count flops for self-attention. NB: We can assume that value_shape == key_shape

Heuristic to decide whether to use Triton Split-k or default (Flash Attention) for prefix attention.

Wrap a module with selective activation checkpointing. It behaves similarly to PyTorch's checkpoint_wrapper, but gives the possibility t

Computes a SwiGLU block given the weights/bias of the 3 linear layers. - It is recommended to keep ``op=None`` so the best implementatio

Computes a SwiGLU block given the weights/bias of the 3 linear layers. :Equivalent pytorch code: .. code-block:: python x1

A wrapper around tree_attention which constructs the biases. Arguments are the same as in xformers.ops.tree_attention.tree_attention, but ins

Does exactly the same as :attr:`torch.unbind` for the forward. In backward, avoids a :attr:`torch.cat` if the gradients are already multi

For a given factor, operate on BMHK attention output as follows: If sequence length is M and there are H * factor heads, redi

Unpacks the ScaledTensor by returning its data and scale as a tuple. Returns: - A tuple of (data, scale), both of which are t

Each thread has data for an 8x8 area of the input tensor Due to the very specific format of the metadata, 32 consecutive bits of the metadata tensor w