Expand description
Fast implementations of commonly used multi-op functions.
Enums§
- Scaled
DotProduct Attention Mask - Mask modes for scaled dot product attention.
Functions§
- layer_
norm - Layer normalization.
- layer_
norm_ device Deprecated - Compatibility shim for
layer_norm. - rms_
norm - Root Mean Square normalization (RMS norm).
- rms_
norm_ device Deprecated - Compatibility shim for
rms_norm. - rope
- Optimized implementation of
NN.RoPE. - rope_
device Deprecated - Compatibility shim for
rope. - rope_
dynamic - Optimized implementation of
NN.RoPEwith dynamic (array) offset. - rope_
dynamic_ device Deprecated - Compatibility shim for
rope_dynamic. - scaled_
dot_ product_ attention - A fast implementation of multi-head attention:
O = softmax(Q @ K.T, dim=-1) @ V - scaled_
dot_ product_ attention_ device Deprecated - Compatibility shim for
scaled_dot_product_attention.