onnxruntime
e266b17e - Add BF16 routing support to CUDA MoE (#32659)

Commit
57 days ago
Add BF16 routing support to CUDA MoE (#32659) ### Description This pull request adds support for the `bfloat16` (BF16) data type in the Mixture of Experts (MoE) CUDA kernel implementation. The main changes ensure that both the sparse mixer and softmax top-k operations can now handle BF16 input tensors, improving compatibility and performance on hardware supporting BF16. Support for BF16 data type: * Added a check for BF16 input tensors using `IsDataType<BFloat16>()` in the MoE kernel (`moe.cc`). * Updated the sparse mixer path to call `LaunchSparseMixerTop2` with BF16 data when appropriate. * Updated the softmax top-k path to call `LaunchSoftmaxTopK` with BF16 data when appropriate. ### Motivation and Context BF16 routing support was missing in the CUDA MoE kernel.
Parents
Loading