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

Code before:
                                 const float *               B,
                                 float *                     X,
                                 int                         n,
                                 int                         k,
                                 int64_t                     ne02,
                                 int64_t                     ne03,
                                 size_t                      s02,
                                 size_t                      s03,
                                 size_t                      s12,
                                 size_t                      s13,
                                 size_t                      s2,
                                 size_t                      s3,
                                 cudaStream_t                stream) {
    const float   alpha         = 1.0f;
    const int64_t total_batches = ne02 * ne03;
    if (total_batches == 0) {
        return;
    }

    // Bulk copy B -> X (contiguous tensors)
    if (X != B) {
        const int64_t total_elements_BX = n * k * total_batches;
        CUDA_CHECK(cudaMemcpyAsync(X, B, total_elements_BX * sizeof(float), cudaMemcpyDeviceToDevice, stream));
    }

    int id = ggml_cuda_get_device();

    ggml_cuda_pool_alloc<const float *> A_ptrs_alloc(ctx.pool(id), total_batches);
    ggml_cuda_pool_alloc<float *>       X_ptrs_alloc(ctx.pool(id), total_batches);

    const float ** A_ptrs_dev = A_ptrs_alloc.get();
    float **       X_ptrs_dev = X_ptrs_alloc.get();

    get_batch_pointers<<<(total_batches + 255) / 256, 256, 0, stream>>>(A, X, A_ptrs_dev, X_ptrs_dev, ne02,
                                                                        total_batches, s02, s03, s2, s3);

    CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream));

    // Yes, this is necessary, without this we get RMSE errors
    CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(id), CUBLAS_DEFAULT_MATH));
    CUBLAS_CHECK(cublasStrsmBatched(ctx.cublas_handle(id), CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N,
                                    CUBLAS_DIAG_NON_UNIT, k, n, &alpha, A_ptrs_dev, n, X_ptrs_dev, k, total_batches));

    // revert to standard mode from common.cuh
    CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(id), CUBLAS_TF32_TENSOR_OP_MATH));

    GGML_UNUSED_VARS(s12, s13);
}

// ======================
// Fast Kernel (n <= 64, k <= 32) - Warp-based parallel reduction

Review comment (suggestion):
```suggestion
    const int id = ggml_cuda_get_device();
```

Code after:
                                 const float *               B,
                                 float *                     X,
                                 int                         n,
                                 int                         k,
                                 int64_t                     ne02,
                                 int64_t                     ne03,
                                 size_t                      s02,
                                 size_t                      s03,
                                 size_t                      s12,
                                 size_t                      s13,
                                 size_t                      s2,
                                 size_t                      s3,
                                 cudaStream_t                stream) {
    const float   alpha         = 1.0f;
    const int64_t total_batches = ne02 * ne03;
    if (total_batches == 0) {
        return;
    }

    // Bulk copy B -> X (contiguous tensors)
    if (X != B) {
        const int64_t total_elements_BX = n * k * total_batches;
        CUDA_CHECK(cudaMemcpyAsync(X, B, total_elements_BX * sizeof(float), cudaMemcpyDeviceToDevice, stream));
    }

    const int id = ggml_cuda_get_device();

    ggml_cuda_pool_alloc<const float *> A_ptrs_alloc(ctx.pool(id), total_batches);
    ggml_cuda_pool_alloc<float *>       X_ptrs_alloc(ctx.pool(id), total_batches);

    const float ** A_ptrs_dev = A_ptrs_alloc.get();
    float **       X_ptrs_dev = X_ptrs_alloc.get();

    get_batch_pointers<<<(total_batches + 255) / 256, 256, 0, stream>>>(A, X, A_ptrs_dev, X_ptrs_dev, ne02,
                                                                        total_batches, s02, s03, s2, s3);

    CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream));

    // Yes, this is necessary, without this we get RMSE errors
    CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(id), CUBLAS_DEFAULT_MATH));
    CUBLAS_CHECK(cublasStrsmBatched(ctx.cublas_handle(id), CUBLAS_SIDE_RIGHT, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N,
                                    CUBLAS_DIAG_NON_UNIT, k, n, &alpha, A_ptrs_dev, n, X_ptrs_dev, k, total_batches));

    // revert to standard mode from common.cuh
    CUBLAS_CHECK(cublasSetMathMode(ctx.cublas_handle(id), CUBLAS_TF32_TENSOR_OP_MATH));

    GGML_UNUSED_VARS(s12, s13);
}

// ======================
// Fast Kernel (n <= 64, k <= 32) - Warp-based parallel reduction