Is AMX supported for LLM inference?

[Hyungyo1]

Describe the issue

Hi,
I have a quick question regarding the LLM inference on CPUs using this extension.
I've been digging into the LLM inference case, and it seems like the kernels written in C++ do not run on AMX (AVX512 is the only one I see). For example, _IPEXlinearReluCPU calls the torch.ops.torch_ipex.tpp_linear_relu C++ code which doesn't seem to be running on AMX.
Is there any LLM layer that runs on AMX, and if so, which C++ code implements it?

Thank you.

[jingxu10]

if you use data type bfloat16 or int8, pytorch and ipex will use AMX.

[Hyungyo1]

Thanks for your response. If it's possible, could you please point out which C++ kernel code implements GEMM on AMX?

[jgong5]

Please note that you will need the 4-th generation xeon or beyond to take advantage of AMX. The tpp kernel you refer to would invoke the micro-kernels in the libxsmm which would leverage AMX on the CPU platforms having the AMX HW support. See BrgemmTPP at

intel-extension-for-pytorch/csrc/cpu/tpp/kernels/TPPGEMMKrnl.h

Line 137 in 46c870e

(BrgemmTPP<T, T>(BSb, Hk, Hc, Hc, Hk * Hc, C, Hk, K, 1.0, 0, Ncb)));

and its implementation here:

intel-extension-for-pytorch/csrc/cpu/tpp/xsmm_functors.h

Line 1835 in 46c870e

class BrgemmTPP {

which calls into libxsmm.

[hezhiqian01]

@jgong5 IPEX calls the oneDNN kernels, doesn't it?

[jgong5]

@jgong5 IPEX calls the oneDNN kernels, doesn't it?

Not always. We have multiple choices of kernels for GEMMs, some implemented with oneDNN and others implemented with TPP/intrinsics kernels.

[NeoZhangJianyu]

@Hyungyo1
Could you feedback?
If no more questions, we will close this issue.

[Hyungyo1]

@NeoZhangJianyu
Yes, my question is answered. Thank you.