[JAX] Integrate BF16 Grouped GEMM with on-device group sizes#2680
Draft
jberchtold-nvidia wants to merge 3 commits intoNVIDIA:mainfrom
Draft
[JAX] Integrate BF16 Grouped GEMM with on-device group sizes#2680jberchtold-nvidia wants to merge 3 commits intoNVIDIA:mainfrom
jberchtold-nvidia wants to merge 3 commits intoNVIDIA:mainfrom
![pre-commit-ci[bot]](https://avatars.githubusercontent.com/in/68672?v=4&size=32){kind=avatar}
