X86 Quantization on test_matmul.py?

[Ryan Taylor (rtayl)]

Is anyone doing quantization on x86 for fp32 to int8, I've been looking at some examples and I'm trying to do this for the test_matmul.py in the triton_shared/python/examples dir but am not able to do it. Does triton even support int8 at all? For example, if I modify the MLIR file after --triton-to-linalg-experimental by changing f32 to i8 I run into translation errors when trying to lower to affine loops. If I change float32 to int8 at the source level for test_matmul.py I run into errors about int8 not being supported in triton_shared: RuntimeError: "normal_kernel_cpu" not implemented for 'Char'
. I have seen this: https://pytorch.org/blog/int8-quantization/ but could not get this to work for test_matmul.py.

[Renat Idrisov (parsifal-47)]

Is anyone doing quantization on x86 for fp32 to int8, I've been looking at some examples and I'm trying to do this for the test_matmul.py in the triton_shared/python/examples dir but am not able to do it. Does triton even support int8 at all? For example, if I modify the MLIR file after --triton-to-linalg-experimental by changing f32 to i8 I run into translation errors when trying to lower to affine loops. If I change float32 to int8 at the source level for test_matmul.py I run into errors about int8 not being supported in triton_shared: RuntimeError: "normal_kernel_cpu" not implemented for 'Char' . I have seen this: https://pytorch.org/blog/int8-quantization/ but could not get this to work for test_matmul.py.

Hi, yes, Triton supports int8. Could you please share the error that you see? Maybe I can help fixing it.

[Renat Idrisov (parsifal-47)]

sure:

    a = torch.randint(low=-128, high=128, size=(rows1, cols1), device='cpu', dtype=torch.int8)
    b = torch.randint(low=-128, high=128, size=(rows2, cols2), device='cpu', dtype=torch.int8)

Shouldn't this be -128, 127?

similarly to range() in python, the upper bound is exclusive

[Renat Idrisov (parsifal-47)]

Is this the internal error you are getting?:

AssertionError('Loop-carried variable accumulator has initial type <[16, 16], int8> but is re-assigned to <[16, 16], int32> in loop! Please make sure that the type stays consistent.')

It looks like tl.dot is perhaps converting the int8 to an int32 and returning that type back to accumulator in the loop. Is that what you've seen?

EDIT; Tried to add an ".to(int8)" to tl.dot call but that got errors. This seemed to get further still error so not sure this is the right thing here.

yes, I saw this error and I tried some simple workarounds like changing acc to int32, but it looks like something wrong is happening internally, need to look at internals

[Renat Idrisov (parsifal-47)]

I have found that tt.dot wasn't implemented properly, here is the fix:
#170

[Ryan Taylor (rtayl)]

I have found that tt.dot wasn't implemented properly, here is the fix: #170

This doesn't seem to have fixed the issue, I am still getting the same error.

[Renat Idrisov (parsifal-47)]

a = torch.randint(low=-128, high=128, size=(rows1, cols1), device='cpu', dtype=torch.int8)
b = torch.randint(low=-128, high=128, size=(rows2, cols2), device='cpu', dtype=torch.int8)

sorry, I forgot to follow up, the kernel now works for me, below are the changes, if these changes do not work, try to purge triton cache with "rm -rf ~/.triton/cache/*"

-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.int32)

-    c = accumulator.to(tl.float32)
+    c = accumulator.to(tl.int32)

-    a = torch.randn((rows1, cols1), device=device, dtype=torch.float32)
-    b = torch.randn((rows2, cols2), device=device, dtype=torch.float32)

+    a = torch.randint(low=-128, high=128, size=(rows1, cols1), device='cpu', dtype=torch.int8)
+    b = torch.randint(low=-128, high=128, size=(rows2, cols2), device='cpu', dtype=torch.int8)

please let me know if it works for you, thank you!