Custom cuda matrix mult is very slow

Thomas_Ahle: <p>I’m trying to write a “fused” cuda kernel for low rank matrix multiplication.<br>
This is the benchmarking code:</p>
<pre><code class="lang-auto">import torch
import torch.utils.benchmark as benchmark

n, m, k = 5000, 4000, 1000
A = torch.randn(n, k, device='cuda')
B = torch.randn(k, m, device='cuda')
v = torch.randn(m, device='cuda')

t0 = benchmark.Timer(
    stmt='lowrank.multiply_low_rank_matrix(A, B, v)',
    setup='''
        from torch.utils.cpp_extension import load
        lowrank = load('lowrank', ['lowrank.cpp', 'lowrank.cu'])
        print('Done')
        ''',
    globals=dict(A=A, B=B, v=v))

t1 = benchmark.Timer(
    stmt='A @ (B @ v)',
    globals=dict(A=A, B=B, v=v))

print(t0.timeit(100))
print(t1.timeit(100))
</code></pre>
<p>When I run it, I get that the naive version <code>A @ (B @ v)</code> is about 16 times faster:</p>
<pre><code class="lang-auto">$ py lowrank_bench.py

&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f5aa667ab30&gt;
lowrank.multiply_low_rank_matrix(A, B, v)
setup:
  from torch.utils.cpp_extension import load
  lowrank = load('lowrank', ['lowrank.cpp', 'lowrank.cu'])

  4.76 ms
  1 measurement, 100 runs , 1 thread
&lt;torch.utils.benchmark.utils.common.Measurement object at 0x7f5aa6678910&gt;
A @ (B @ v)
  249.51 us
  1 measurement, 100 runs , 1 thread
</code></pre>
<p>I’m wondering if I’m doing something wrong with benchmarking, or if my code skills are just bad.<br>
My kernel looks roughly like the code below.<br>
It creates an intermediate vector for <code>Bv</code> (the same way I assume pytorch <code>A@(B@v)</code> has to), and uses a call to <code>__syncthreads()</code> to make sure <code>Bv</code> is computed before it starts multiplying with <code>A</code>.</p>
<pre><code class="lang-auto">template &lt;typename scalar_t&gt;
__global__ void multiply_low_rank_matrix_cuda_kernel(
    const torch::PackedTensorAccessor&lt;scalar_t,2,torch::RestrictPtrTraits&gt; A,
    const torch::PackedTensorAccessor&lt;scalar_t,2,torch::RestrictPtrTraits&gt; B,
    const torch::PackedTensorAccessor&lt;scalar_t,1,torch::RestrictPtrTraits&gt; vector,
    torch::PackedTensorAccessor&lt;scalar_t,1,torch::RestrictPtrTraits&gt; result) {

  const int thread_id = blockIdx.x * blockDim.x + threadIdx.x;
  __shared__ scalar_t intermediate_vector[1024]; // Shared memory for intermediate vector

  // Phase 1: Each thread computes one element of the intermediate vector (Bv)
  if(thread_id &lt; B.size(0)) {
    intermediate_vector[thread_id] = 0;
    for (int j = 0; j &lt; B.size(1); ++j) {
      intermediate_vector[thread_id] += B[thread_id][j] * vector[j];
    }
  }
  __syncthreads(); // Synchronize threads to ensure the intermediate vector is fully computed

  // Phase 2: Each thread computes one element of the final output vector (Av)
  if(thread_id &lt; A.size(0)) {
    result[thread_id] = 0;
    for (int j = 0; j &lt; A.size(1); ++j) {
      result[thread_id] += A[thread_id][j] * intermediate_vector[j];
    }
  }
}
</code></pre>
<p>My end goal is to write something more advanced with cuda, so I want to make sure I’m following best practices and not doing anything stupid at this stage.</p>

ptrblck: <p>I would assume your implementation isn’t optimal and you could check e.g. <a href="https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html">this blog post</a> which explains tiles and quantization effects and e.g. the <a href="https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#shared-memory">shared memory section of the CUDA programming guide</a> for more information.</p>

Armin: <p>Did you figure out the problem? I’m facing the same issue!</p>