支持的运营和硬件

我们目前支持以下 GPU 内核：

• tokamax.dot_product_attention（FlashAttention）。
• tokamax.gated_linear_unit（门控线性单元 [SwiGLU 等]）。
• tokamax.layer_norm （层归一化和均方根归一化）。

对于 GPU 和 TPU，以下各项均适用：

• tokamax.ragged_dot （混合专家）。