Matrix multiplication implementation in PyTorch

Areg_Petrosyan: <p>I read about <code>torch.autocast</code> and how FP16 matrix multiplication is faster than FP32 on CUDA. However, on my Mac M1 (Intel chip), a 100x100 matrix multiplication takes 50 times longer in FP16 than FP32. Reducing the matrix size decreases the computation time, but FP32 still remains faster.</p>
<p>I’m trying to understand how <code>torch.mm</code> is implemented, but I couldn’t find the actual implementation code. Could you guide me on where to find it?<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/1/1164c0316f8a3e99f93bb18951f2174bd5a5b9cd.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/1164c0316f8a3e99f93bb18951f2174bd5a5b9cd" title="2024-12-21 21.59.23"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/1164c0316f8a3e99f93bb18951f2174bd5a5b9cd_2_556x500.jpeg" alt="2024-12-21 21.59.23" data-base62-sha1="2tRYpX1PHHMzeqaIF31yB7KNAqV" width="556" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/1164c0316f8a3e99f93bb18951f2174bd5a5b9cd_2_556x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/1164c0316f8a3e99f93bb18951f2174bd5a5b9cd_2_834x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/1/1164c0316f8a3e99f93bb18951f2174bd5a5b9cd_2_1112x1000.jpeg 2x" data-dominant-color="1F2020"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">2024-12-21 21.59.23</span><span class="informations">1280×1150 92.3 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>

blackberry: <p>Afaik, you need a GPU with tensor cores to be able to benefit from faster operations with float16.</p>

Areg_Petrosyan: <p>Yes, you are correct that in GPUs, F16 is faster than F32. However, I need to understand why the opposite is true for CPUs. F16 occupies half the memory of F32, and logically, when the same operation is applied to X and X/2, I would expect X/2 to operate faster. But why is this not the case for CPUs?</p>

albanD: <p>Hey!</p>
<p>Long story short, this is where most of the mm calls end up at on MPS device: <a href="https://github.com/pytorch/pytorch/blob/2ed4d65af0a1993c0df7b081f4088d0f3614283e/aten/src/ATen/native/mps/operations/LinearAlgebra.mm#L72" class="inline-onebox">pytorch/aten/src/ATen/native/mps/operations/LinearAlgebra.mm at 2ed4d65af0a1993c0df7b081f4088d0f3614283e · pytorch/pytorch · GitHub</a><br>
And for CPU it ends up in <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/CPUBlas.cpp" class="inline-onebox">pytorch/aten/src/ATen/native/CPUBlas.cpp at main · pytorch/pytorch · GitHub</a></p>
<p>As you can see in both cases, we just end up calling into Apple-provided libraries (either MPS kernel or Accelerate gemm kernel). So it is most likely that the concern is within these libraries.<br>
In some cases (like the Metal gemm you can find) we write our own kernels when they behave better than vendor-provided ones but we try to avoid unless really necessary.</p>
<p>In your particular case, it is most likely that the library ends up doing a lot of expensive upcasting for no good reason and thus is slower <img src="https://discuss.pytorch.org/images/emoji/apple/confused.png?v=12" title=":confused:" class="emoji" alt=":confused:" loading="lazy" width="20" height="20"><br>
That  being said, there can also sometimes be packaging issues with Accelerate/MKL/neither so if you see significant differences in an easy-to-reproduce environment, please submit an issue to we can track it and try and work around it!</p>

swolchok: <p>I’ve been working on aarch64 GEMM-related stuff lately, so I’d like to provide some additional color.</p>
<ol>
<li>
<p>The MPS device does not appear to be relevant to this example.</p>
</li>
<li>
<p>For CPU, <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/CPUBlas.cpp#L367" rel="noopener nofollow ugc">the half @ half → half GEMM implementation in CPUBlas</a> (which I believe is the relevant one) does not call into Accelerate, because (to the best of my knowledge) Accelerate does not support half or bfloat16 for GEMM. Instead, it calls into <code>gemm_stub</code>, which is implemented in ATen/native/cpu/BlasKernel.cpp. From there, depending on argument transposition, you will hit one of the naive GEMM kernels in that file. IIRC, the best one is <a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/cpu/BlasKernel.cpp#L375" rel="noopener nofollow ugc">gemm_transa_</a> (corresponding to the most common transposition case in LLMs, <code>activations @ weights.t()</code>), which still isn’t very good but at least uses our accelerated fp16 dot product kernel in ATen/native/cpu/ReducedPrecisionFloatGemvFastPathKernel.cpp. Note that matrix-vector multiplication of the form <code>vector @ matrix.t()</code> should hit the GEMV kernel in ReducedPrecisionFloatGemvFastPathKernel.cpp as a special case, because it is relevant to LLM decoding.</p>
</li>
</ol>
<p>I am interested in getting better GEMMs for half and bfloat16 into PyTorch core, but 1) I am struggling with prioritization given that, for LLM inference, it is unclear to me personally why 8-bit quantization wouldn’t be the obvious better performance/accuracy tradeoff instead, 2) our custom matrix multiplication infrastructure lives in torchao under <code>torchao/experimental/kernels/cpu/aarch64</code> (and only contains low-bit quantized kernels, but that’s the infrastructure we’d likely extend to half/bfloat16) and therefore currently isn’t usable from PyTorch core, and 3) XNNPACK has fp16 kernels but bf16 is planned-but-not-implemented, so if we wanted to simply use XNNPACK we would be incentivized to wait for their bf16 implementation. The issue seems to be that most BLAS implementations don’t have 16-bit floating-point support.</p>
<p>P.S. I meant to provide a lot more links, but as a new user on this forum I am prohibited from doing so.</p>