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

Code before:
    const size_t nb02, const size_t nb03,
    const size_t nb12, const size_t nb13,
    const size_t nb2, const size_t nb3,
    const int n, const int k) {
    const int batch_idx = blockIdx.x;
    const int lane      = threadIdx.x;
    const int col_idx   = threadIdx.y;

    // A block processes one batch, k warps process k columns
    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*)((const char *)A + i02 * nb02 + i03 * nb03);
    const float* const B_batch = (const float*)((const char *)B + i02 * nb12 + i03 * nb13);
    float*             X_batch = (float*)      ((char *)X + i02 * nb2  + i03 * nb3);

    __shared__ float sA[MAX_N_FAST * MAX_N_FAST];
    __shared__ float sX[MAX_N_FAST * MAX_K_FAST];

    // Load A into shared memory (coalesced)
    #pragma unroll
    for (int i = threadIdx.x + threadIdx.y * blockDim.x; i < n * n; i += blockDim.x * blockDim.y) {
        sA[i] = A_batch[i];
    }

    // Load B into shared memory (coalesced)
    #pragma unroll
    for (int i = threadIdx.x + threadIdx.y * blockDim.x; i < n * k; i += blockDim.x * blockDim.y) {
        sX[i] = B_batch[i];
    }
    __syncthreads();

    // Each warp (32 threads with same col_idx) solves one column
    for (int row = 0; row < n; ++row) {
        float sum = 0.0f;

        // Parallel reduction for sum
        for (int j = lane; j < row; j += WARP_SIZE) {
            sum += sA[row * n + j] * sX[j * k + col_idx];
        }

        sum = warp_reduce_sum(sum);

        // Lane 0 computes and stores the final result for the current row
        if (lane == 0) {
            const float b_val = sX[row * k + col_idx]; // Value from B

Review comment (suggestion):
These cannot be unrolled so you can remove
```suggestion
```

Code after:
            const float b_val      = sXt[col_idx * n + row];
            const float a_diag     = sA[row * n + row];
            // no safeguards for division by zero because that indicates corrupt
            // data anyway
            sXt[col_idx * n + row] = (b_val - sum) / a_diag;
        }
    }

    __syncthreads();

#pragma unroll
    for (int i = 0; i < rows_per_warp; i++) {
        const int i0 = lane + i * WARP_SIZE;
        if (i0 < n) {
            X_batch[i0 * k + col_idx] = sXt[col_idx * n + i0];
        }
    }
}
#ifdef __clang__
#    pragma clang diagnostic pop
#endif  // __clang__

static void solve_tri_f32_cuda(const float * A,
                               const float * B,
                               float *       X,
                               int           n,
                               int           k,
                               int64_t       ne02,
                               int64_t       ne03,
                               size_t        nb02,
                               size_t        nb03,
                               size_t        nb12,
                               size_t        nb13,
                               size_t        nb2,
                               size_t        nb3,
                               cudaStream_t  stream) {
    const uint3 ne02_fd = init_fastdiv_values((uint32_t) ne02);
    dim3        threads(WARP_SIZE, k);
    dim3        grid(ne02 * ne03);
    if (n == 64) {
        switch (k) {
            case 32:
                solve_tri_f32_fast<64, 32>
                    <<<grid, threads, 0, stream>>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0);
                break;
            case 16:
                solve_tri_f32_fast<64, 16>
                    <<<grid, threads, 0, stream>>>(A, B, X, ne02_fd, nb02, nb03, nb12, nb13, nb2, nb3, 0, 0);
                break;
            case 14:
                solve_tri_f32_fast<64, 14>