How to achieve BLAS-level matmul performance on 4080/4090/A100 (pre Hopper Arch) ThunderKittens?

[fire-emblem]

Hi,

I’m trying to implement a high-performance matmul kernel for large matrices (4096x4096x4096) on an RTX 4080 using ThunderKittens. I’ve studied kernels/matmul/educational/level_04.cu, but my performance is only about 1/10th of peak, far behind cuBLAS or torch.matmul.

I noticed that other matmul examples require features based on Hopper, such as TMA or WGMMA.

What main optimizations or directions would you recommend to approach BLAS-level performance on Ada GPUs like the 4080?

Also, could you provide more demo and kernel examples specifically for pre-Hopper architectures?

Thanks!

[tomflinda]

Similar senario to me, I implemented a similar matmul kernel for large matrices (8192x8192x8192) on A100-PCIE-40GB based on kernels/matmul/educational/level_04.cu, shown as follows:

struct matmul_globals {
  using sub_tile = st_bf<BLOCK_SIZE, BLOCK_SIZE>;
  using sub_tile_float = kittens::st<float, BLOCK_SIZE, BLOCK_SIZE>;
  using tile_gl = gl<bf16, 1, 1, -1, -1, sub_tile>;
  using tile_gl_c = gl<float, 1, 1, -1, -1, sub_tile_float>;
  tile_gl A;
  tile_gl B;
  tile_gl_c C;
  int g_N;
};

__global__ void micro_kernel(const __grid_constant__ matmul_globals g,
                             float alpha, float beta) {
  extern __shared__ alignment_dummy __shm[];
  shared_allocator al((int *)&__shm[0]);
  st_bf<BLOCK_SIZE, BLOCK_SIZE> &As =
      al.allocate<st_bf<BLOCK_SIZE, BLOCK_SIZE>>();
  st_bf<BLOCK_SIZE, BLOCK_SIZE> &Bs =
      al.allocate<st_bf<BLOCK_SIZE, BLOCK_SIZE>>();

  rt_bf<BLOCK_SIZE, BLOCK_SIZE> A_reg;
  rt_bf<BLOCK_SIZE, BLOCK_SIZE> B_reg;
  rt_bf<BLOCK_SIZE, BLOCK_SIZE, ducks::rt_layout::col> B_reg_col;
  int row = blockIdx.y;

  zero(C_accum);
  int num_tiles = (g.g_N + BLOCK_SIZE - 1) / BLOCK_SIZE;
  for (int tile = 0; tile < num_tiles; ++tile) {
    load(As, g.A, {0, 0, row, tile});
    load(Bs, g.B, {0, 0, tile, col});
    __syncthreads();
    load(A_reg, As);
    load(B_reg, Bs);
    swap_layout(B_reg_col, B_reg);
    __syncthreads();
    mma_AB(C_accum, A_reg, B_reg_col, C_accum);
    __syncthreads();
  }
  store(g.C, C_accum, {0, 0, row, col});
}

// Launch Kernel
void dispatch_micro(kittens::bf16 *d_A, kittens::bf16 *d_B, float *d_C,
                    float alpha, float beta) {
  // Create global memory objects directly in dispatch
  using a_gl = matmul_globals::tile_gl;
  using b_gl = matmul_globals::tile_gl;
  using c_gl = matmul_globals::tile_gl_c;
  a_gl a_arg{d_A, nullptr, nullptr, N, N};
  b_gl b_arg{d_B, nullptr, nullptr, N, N};
  c_gl c_arg{d_C, nullptr, nullptr, N, N};
  matmul_globals g{a_arg, b_arg, c_arg, N};

  // Calculate grid and block dimensions
  dim3 grid(N / BN, M / BM, 1);
	@@ -83,8 +77,11 @@ void dispatch_micro(kittens::bf16 *d_A, kittens::bf16 *d_B, float *d_C,
  unsigned long mem_size =
      32768; // No dynamic shared memory used in this example

  cudaFuncSetAttribute(micro_kernel,
                       cudaFuncAttributeMaxDynamicSharedMemorySize, mem_size);

  // Launch kernel
  micro_kernel<<<grid, block, mem_size>>>(g, alpha, beta);

  cudaDeviceSynchronize();
}

And its performance only get to 30311.6 GFLOPS, far from the peek formance 300TF.

[fattorib]

A bit late, but on Ada (RTX 4070) I wrote this kernel on an older version of Thunderkittens which achieves parity with cuBLAS and Triton. Just a note that it probably won't compile on the latest version of TK, but you can probably use it as a starting point to get a faster implementation