Поддерживаемые операции и оборудование

В настоящее время мы поддерживаем следующие ядра графических процессоров:

• tokamax.dot_product_attention ( FlashAttention ).
• tokamax.gated_linear_unit ( линейные блоки управления (SwiGLU и т. д.)).
• tokamax.layer_norm ( Нормализация слоев и среднеквадратическая нормализация ).

И следующее для GPU и TPU:

• tokamax.ragged_dot ( Смесь экспертов ).