Does torchao support FP8 Grouped GEMM?

[zigzagcai]

Grouped GEMM kernels (https://github.com/fanshiqing/grouped_gemm) are used in many MoE models.

I just wander does torchao support FP8 kernels for Grouped GEMM, such like the three commonly used ops:

grouped_gemm.backend.gmm
grouped_gemm.ops.unpermute
grouped_gemm.ops.permute

[vkuzo]

hi @zigzagcai , we recently landed a grouped gemm API into core which includes fp8: pytorch/pytorch#148531 . We plan to provide wrappers in torchao, although we do not have them just yet. cc @drisspg

[zigzagcai]

hi @zigzagcai , we recently landed a grouped gemm API into core which includes fp8: pytorch/pytorch#148531 . We plan to provide wrappers in torchao, although we do not have them just yet. cc @drisspg

Thank you @vkuzo !
I just wander how can I use this aten newly needed grouped gemm ops?

[supriyar]

cc @HDCharles who has been looking into MoE quantization and grouped gemm recently

[HDCharles]

Hey,

I'm working to enable our existing quantization kernels to compose with group gemm its still in progress at the moment. As far as the core kernel, you can look at: https://github.com/pytorch/pytorch/pull/148531/files#diff-3f31c52b48cfddf8f4617d809f7695b2e4a1c78656f8c4b5143a4b45d01fcf23R1178

...for an example

[jeromeku]

@HDCharles @vkuzo

Interested in this as well and potentially helping tune the kernel.

There is a link mentioned in the grouped gemm PR describing the design of the grouped GEMM. How can I view the doc (access seems to be gated)?