Input:
Q (2x4): [[1.0, 0.0, 0.0, 0.0],
          [0.0, 1.0, 0.0, 0.0]]
K (3x4): [[1.0, 0.0, 0.0, 0.0],
          [0.0, 1.0, 0.0, 0.0],]
          [0.0, 0.0, 1.0, 0.0]]
V (3x4): [[1.0, 2.0, 3.0, 4.0],
          [5.0, 6.0, 7.0, 8.0],]
          [9.0, 10.0, 11.0, 12.0]]

Output:
output (2x4): [[4.29, 5.29, 6.29, 7.29],
               [5.0, 6.0, 7.0, 8.0]]

#include <cuda_runtime.h>

__device__ __forceinline__ constexpr int compute_offset(int row, int col,
                                                        int cols) {
  return row * cols + col;
}

// A: M x K
// B: K x N
// => output: AB, M x N
// A: M x K
// B: K x N
// => output: AB^T, M x N
template <const int block_size, const bool transpose, const bool norm>
__global__ void multiple_kernel(const float* A, const float* B, float* output,
                                int M, int N, int K) {
  __shared__ float tile_A[block_size][block_size];
  __shared__ float tile_B[block_size][block_size];

  const auto col = threadIdx.x + blockIdx.x * blockDim.x;
  const auto row = threadIdx.y + blockIdx.y * blockDim.y;

  const auto local_row = threadIdx.y;
  const auto local_col = threadIdx.x;

  float sum = 0.0f;

#pragma unroll
  for (int i = 0; i < K; i += block_size) {
    if (row < M && local_col + i < K) {
      tile_A[local_row][local_col] = A[compute_offset(row, local_col + i, K)];
    } else {
      tile_A[local_row][local_col] = 0.0f;
    }

    if constexpr (transpose) {
      if (col < N && local_row + i < K) {
        tile_B[local_row][local_col] = B[compute_offset(col, local_row + i, K)];
      } else {
        tile_B[local_row][local_col] = 0.0f;
      }
    } else {
      if (local_row + i < K && col < N) {
        tile_B[local_row][local_col] = B[compute_offset(local_row + i, col, N)];
      } else {
        tile_B[local_row][local_col] = 0.0f;
      }
    }

    __syncthreads();

#pragma unroll
    for (int j = 0; j < block_size; j++) {
      sum = fmaf(tile_A[local_row][j], tile_B[j][local_col], sum);
    }

    __syncthreads();
  }

  if (row < M && col < N) {
    if constexpr (norm) {
      output[compute_offset(row, col, N)] = sum * rsqrtf(static_cast<float>(K));
    } else {
      output[compute_offset(row, col, N)] = sum;
    }
  }
}

struct SoftmaxState {
  float max;
  float sum;
};

__device__ __forceinline__ SoftmaxState init_state() {
  return {-INFINITY, 0.0f};
}

__device__ __forceinline__ SoftmaxState mono_state(float value) {
  return {value, 1.0f};
}

__device__ __forceinline__ SoftmaxState combine(SoftmaxState a,
                                                SoftmaxState b) {
  if (a.max == -INFINITY) {
    return b;
  }
  if (b.max == -INFINITY) {
    return a;
  }

  SoftmaxState out;
  out.max = fmaxf(a.max, b.max);
  out.sum = a.sum * __expf(a.max - out.max) + b.sum * __expf(b.max - out.max);
  return out;
}

__device__ __forceinline__ SoftmaxState warp_reduce(SoftmaxState state) {
#pragma unroll
  for (int offset = 16; offset > 0; offset >>= 1) {
    SoftmaxState other;
    other.max = __shfl_down_sync(0xffffffff, state.max, offset);
    other.sum = __shfl_down_sync(0xffffffff, state.sum, offset);
    state = combine(state, other);
  }
  return state;
}

template <const int block_size>
__device__ __forceinline__ SoftmaxState block_reduce(SoftmaxState state) {
  __shared__ SoftmaxState warp_results[block_size];

  const auto lane = threadIdx.x & 31;
  const auto warp = threadIdx.x >> 5;
  const auto num_warps = (blockDim.x + 31) >> 5;

  state = warp_reduce(state);

  if (lane == 0) {
    warp_results[warp] = state;
  }

  __syncthreads();

  SoftmaxState block_state = init_state();
  if (warp == 0) {
    if (lane < num_warps) {
      block_state = warp_results[lane];
    }

    block_state = warp_reduce(block_state);
  }

  return block_state;
}

// softmax(input)
template <const int block_size>
__global__ void softmax_kernel(float* input, int M, int N) {
  const auto row = blockIdx.x;

  if (row >= M) {
    return;
  }

  auto state = init_state();

#pragma unroll
  for (int col = threadIdx.x; col < N; col += blockDim.x) {
    const int idx = compute_offset(row, col, N);
    state = combine(state, mono_state(input[idx]));
  }

  state = block_reduce<block_size>(state);

  __shared__ SoftmaxState final_state;
  if (threadIdx.x == 0) {
    final_state = state;
  }
  __syncthreads();

#pragma unroll
  for (int col = threadIdx.x; col < N; col += blockDim.x) {
    const int idx = compute_offset(row, col, N);
    input[idx] = __expf(input[idx] - final_state.max) / final_state.sum;
  }
}

// Q, K, V, output are device pointers
// Q: M x d
// K: N x d
// QK^T: M x N
// V: N x d
// output: M x d
extern "C" void solve(const float* Q, const float* K, const float* V,
                      float* output, int M, int N, int d) {
  auto total_items = M * N;
  float* attention;
  cudaMalloc(&attention, total_items * sizeof(float));
  cudaMemset(attention, 0, total_items * sizeof(float));

  // QK^T / sqrt(d)
  {
    constexpr auto block_size = 32;
    const auto grid_dim = dim3((N + block_size - 1) / block_size,
                               (M + block_size - 1) / block_size);
    const auto block_dim = dim3(block_size, block_size);

    multiple_kernel<block_size, true, true>
        <<<grid_dim, block_dim>>>(Q, K, attention, M, N, d);
  }

  // row-independent softmax(attention)
  {
    constexpr auto block_size = 256;
    const auto grid_dim = dim3(M);
    const auto block_dim = dim3(block_size);

    softmax_kernel<block_size><<<grid_dim, block_dim>>>(attention, M, N);
  }

  // softmax(attention) * V
  {
    constexpr auto block_size = 32;
    const auto grid_dim = dim3((d + 31) / 32, (M + 31) / 32);
    const auto block_dim = dim3(32, 32);

    multiple_kernel<block_size, false, false>
        <<<grid_dim, block_dim>>>(attention, V, output, M, d, N);
  }

  cudaDeviceSynchronize();
}