Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cuda/solve_tri.cu (CUDA)

Code before:
                                          const size_t nb13,
                                          const size_t nb2,
                                          const size_t nb3,
                                          const int    n_arg,
                                          const int    k_arg) {
    const int n = n_template == 0 ? n_arg : n_template;
    const int k = k_template == 0 ? k_arg : k_template;

    const int batch_idx = blockIdx.x;
    const int lane      = threadIdx.x;
    const int col_idx   = threadIdx.y;

    if (col_idx >= k) {
        return;
    }

    const uint2   i02_i03 = fast_div_modulo(batch_idx, ne02);
    const int64_t i02     = i02_i03.y;
    const int64_t i03     = i02_i03.x;

    const float * const A_batch = (const float *) (A + i02 * nb02 / sizeof(float) + i03 * nb03 / sizeof(float));
    const float * const B_batch = (const float *) (B + i02 * nb12 / sizeof(float) + i03 * nb13 / sizeof(float));
    float *             X_batch = (float *) (X + i02 * nb2 / sizeof(float) + i03 * nb3 / sizeof(float));

    __shared__ float sA[MAX_N_FAST * MAX_N_FAST];
    __shared__ float sXt[MAX_N_FAST * (MAX_K_FAST + 1)];

    const int offset = threadIdx.x + threadIdx.y * blockDim.x;

#pragma unroll
    for (int i = 0; i < n * n; i += k * WARP_SIZE) {
        int i0 = i + offset;
        if (i0 < n * n) {
            sA[i0] = A_batch[i0];
        }
    }

    const int rows_per_warp = (n + WARP_SIZE - 1) / WARP_SIZE;

#pragma unroll
    for (int i = 0; i < rows_per_warp; i++) {
        const int i0 = lane + i * WARP_SIZE;
        if (i0 < n) {
            sXt[col_idx * n + i0] = B_batch[i0 * k + col_idx];
        }
    }

    __syncthreads();

#pragma unroll
    for (int row = 0; row < n; ++row) {

Review comment (suggestion):
Simply changing the size of a 1D allocation like this does nothing to fix shared memory bank conflicts. You have to actually access elements with the padded stride. One way to do this automatically is to change the array shape to be 2D and to pad the last dimension.

Code after:
                                          const size_t nb13,
                                          const size_t nb2,
                                          const size_t nb3,
                                          const int    n_arg,
                                          const int    k_arg) {
    const int n = n_template == 0 ? n_arg : n_template;
    const int k = k_template == 0 ? k_arg : k_template;

    const int batch_idx = blockIdx.x;
    const int lane      = threadIdx.x;
    const int col_idx   = threadIdx.y;

    if (col_idx >= k) {
        return;
    }

    const uint2   i02_i03 = fast_div_modulo(batch_idx, ne02);
    const int64_t i02     = i02_i03.y;
    const int64_t i03     = i02_i03.x;

    const float * const A_batch = (const float *) (A + i02 * nb02 + i03 * nb03);
    const float * const B_batch = (const float *) (B + i02 * nb12 + i03 * nb13);
    float *             X_batch = (float *) (X + i02 * nb2 + i03 * nb3);

    __shared__ float sA[MAX_N_FAST * MAX_N_FAST];
    __shared__ float sXt[MAX_N_FAST * (MAX_K_FAST + 1)];

    const int offset = threadIdx.x + threadIdx.y * blockDim.x;

#pragma unroll
    for (int i = 0; i < n * n; i += k * WARP_SIZE) {
        int i0 = i + offset;
        if (i0 < n * n) {
            sA[i0] = A_batch[i0];
        }
    }

    const int rows_per_warp = (n + WARP_SIZE - 1) / WARP_SIZE;

#pragma unroll
    for (int i = 0; i < rows_per_warp; i++) {
        const int i0 = lane + i * WARP_SIZE;
        if (i0 < n) {
            sXt[col_idx * n + i0] = B_batch[i0 * k + col_idx];
        }
    }

    __syncthreads();

#pragma unroll
    for (int row = 0; row < n; ++row) {