[BACKEND] Add support to convert INT8 MMAV2 accumulator layout to dot_operand layout

[tongyuantongyu]

Partial fix of #3580. Resolved the INT8 layoutC -> INT8 layoutA case.

• Port MMAV3's reg shuffling to support MMAV2 layout.
• Simplify the shuffling logic.

[ThomasRaoux]

I haven't reviewed in details as it is failing the tests. I'm not sure I understand why the code sequence for mma to dot_operand(fp8) has changed.

The FP16 case is a bit tough. convert_layout only knows that the tensor is in #mma layout, but has no idea which MMA exactly. #mma layouts of different MMAs are different. Guessing from element type is not reliable, as user may (and for INT8/FP8 MMA, have to) cast between types.

I don't understand, could you give an example of what MMA format is different based on the type?

[ThomasRaoux]

Why are those removed?

[tongyuantongyu]

These two corresponding to the prmt-s using selectorEx0 and selectorEx1. These two patterns are simply selecting one Value of the two, so a select is sufficient and no prmt is required here.

[ThomasRaoux]

let' not add it if it is not used and tested

[tongyuantongyu]

Done

[tongyuantongyu]

I don't understand, could you give an example of what MMA format is different based on the type?

You're indeed right. I was confused as there turns out to be another issue here (Detail reported in #3580 (comment)). Loading INT8 input and do both MMA in FP16 also get wrong result. convert8BitsMMAV2To16BitsDotOperand fixed the wrong order, and made me think the issue is MMA having different layouts.

[tongyuantongyu]

FAILED hopper/test_gemm.py::test_gemm[128-128-64-4-1-4096-1-1024-False-False-True-none-float32-False-3] - AssertionError: Tensor-likes are not close!

Mismatched elements: 1 / 4096 (0.0%)
Greatest absolute difference: 2.0 at index (289, 0) (up to 0.001 allowed)
Greatest relative difference: 2.0 at index (289, 0) (up to 0.01 allowed)

This failure seems flaky. I don't see a torch.manual_seed call in test_gemm.py, so maybe this is numerical instability with specific input value.

[ThomasRaoux]

I see a big performance regression in the fp8 flash attention tutorial with this patch (run python tutorials/06-fused-attention.py on h100 to see it). I'm not sure where this is coming from. Please fix it and I can take a deeper look at the changes.

[ThomasRaoux]

that doesn't make sense. That will change behavior of existing tests. If we want to tests i32 dtype it can be set in the config?

[tongyuantongyu]

Output type for INT8 MMA is always INT32 (https://github.com/openai/triton/blob/main/python/triton/language/semantic.py#L1358), and out_dtype is simply ignored, as opposed to FP MMA (https://github.com/openai/triton/blob/main/python/triton/language/semantic.py#L1367)

I changed it to tl.int32 here just for correctness. If you feel necessary, I can add a check that out_dtype must be tl.int32 for tl.int8 inputs, in this or a separate PR.

[tongyuantongyu]

I see a big performance regression in the fp8 flash attention tutorial with this patch

Sorry I don't have access to H100. I made an attempt to fix it, could you test if it fixes the regression? If it's still there, I'll try to revert all changes to MMAV3 part.

[ThomasRaoux]

I see a big performance regression in the fp8 flash attention tutorial with this patch

Sorry I don't have access to H100. I made an attempt to fix it, could you test if it fixes the regression? If it's still there, I'll try to revert all changes to MMAV3 part.

Thanks, that seems to fix it. I'll look more at the code sequence in a little bit to understand it.