New mma_atoms and copy_atoms in bmg_grouped_gemm_fp8 #579
Conversation
nsingh-habana
changed the title
nsingh-habana
force-pushed
the
grouped_gemm_fp8_new_atoms
branch
from
0959d75 to
1d6d2f7
Compare
| Tensor tCrA_fp16 = make_fragment_like<half_t>(tCrA); | ||
| Tensor tCrB_fp16 = make_fragment_like<half_t>(tCrB); |
There was a problem hiding this comment.
@rolandschulz & @petercad, please advise whether such a redesign would make sense -
For FP8xFP8 GEMM on Xe2, FP8 is converted into FP16.
reorders in the new API allow multiple dtypes to share the same GEMM (MMA collectives) code. They're no-ops if dtype conversion is not needed. So, perhaps, they could share the same code? I could be wrong, but it seems this is (partly) what @petercad had in mind regarding the rearchitecture.
While it's indeed possible to use if constexpr with compile-time evaluated expressions to add multiple dtypes' GEMMs support in the same file, reorders seem to make things even simpler. It was earlier decided to choose readability & debuggability over reducing code duplication, which is why the legacy FP8xFP8 GEMM currently has a separate implementation with duplicated code.
@nsingh-habana, please also share your thoughts on it.
Thanks!
2 participants
No description provided.