Skip to content

Support broadcast add & mul on CUDA (fixed) #2192

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Jul 14, 2023
Merged

Support broadcast add & mul on CUDA (fixed) #2192

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a53a59a
Support broadcast add & mul on CUDA (fixed)
li-plus Jul 12, 2023
4fc4014
Merge branch 'master' into bcast-cuda
ggerganov Jul 14, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 12 additions & 21 deletions ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -252,13 +252,13 @@ struct ggml_tensor_extra_gpu {
cudaEvent_t events[GGML_CUDA_MAX_DEVICES]; // events for synchronizing multiple GPUs
};

static __global__ void add_f32(const float * x, const float * y, float * dst, const int k) {
static __global__ void add_f32(const float * x, const float * y, float * dst, const int kx, const int ky) {
const int i = blockDim.x*blockIdx.x + threadIdx.x;

if (i >= k) {
if (i >= kx) {
return;
}
dst[i] = x[i] + y[i];
dst[i] = x[i] + y[i%ky];
}

static __global__ void add_f16_f32_f16(const half * x, const float * y, half * dst, const int k) {
Expand Down Expand Up @@ -1996,9 +1996,9 @@ static __global__ void scale_f32(const float * x, float * dst, const float scale
dst[i] = scale * x[i];
}

static void add_f32_cuda(const float * x, const float * y, float * dst, const int k, cudaStream_t stream) {
const int num_blocks = (k + CUDA_ADD_BLOCK_SIZE - 1) / CUDA_ADD_BLOCK_SIZE;
add_f32<<<num_blocks, CUDA_ADD_BLOCK_SIZE, 0, stream>>>(x, y, dst, k);
static void add_f32_cuda(const float * x, const float * y, float * dst, const int kx, const int ky, cudaStream_t stream) {
const int num_blocks = (kx + CUDA_ADD_BLOCK_SIZE - 1) / CUDA_ADD_BLOCK_SIZE;
add_f32<<<num_blocks, CUDA_ADD_BLOCK_SIZE, 0, stream>>>(x, y, dst, kx, ky);
}

static void add_f16_f32_f16_cuda(const half * x, const float * y, half * dst, const int k, cudaStream_t stream) {
Expand Down Expand Up @@ -2610,17 +2610,15 @@ inline void ggml_cuda_op_add(
GGML_ASSERT(src1_ddf_i != nullptr);
GGML_ASSERT(dst_ddf_i != nullptr);

// TODO: support broadcasting
GGML_ASSERT(ggml_nelements(src0) == ggml_nelements(src1));

const int64_t ne00 = src0->ne[0];
const int64_t i01_diff = i01_high - i01_low;

// const int64_t ne10 = src1->ne[0];
const int64_t ne10 = src1->ne[0];
const int64_t ne11 = src1->ne[1];

// compute
if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
add_f32_cuda(src0_ddf_i, src1_ddf_i, dst_ddf_i, ne00*i01_diff, cudaStream_main);
add_f32_cuda(src0_ddf_i, src1_ddf_i, dst_ddf_i, ne00*i01_diff, ne10*ne11, cudaStream_main);
} else if (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
add_f16_f32_f16_cuda((half *) src0_ddq_i, src1_ddf_i, (half *) dst_ddf_i, ne00*i01_diff, cudaStream_main);
} else {
Expand All @@ -2644,19 +2642,12 @@ inline void ggml_cuda_op_mul(
GGML_ASSERT(dst_ddf_i != nullptr);

const int64_t ne00 = src0->ne[0];
const int64_t i01_diff = i01_high - i01_low;

const int64_t ne10 = src1->ne[0];
const int64_t ne11 = src1->ne[1];

for (int64_t i01 = i01_low; i01 < i01_high; i01++) {
const int64_t i11 = i1*ne11 + i01%ne11; // broadcast src1 across src0

float * src0_ddf_i01 = src0_ddf_i + i01*ne00;
float * src1_ddf_i01 = src1_ddf_i + i11*ne10;
float * dst_ddf_i01 = dst_ddf_i + i01*ne00;

// compute
mul_f32_cuda(src0_ddf_i01, src1_ddf_i01, dst_ddf_i01, ne00, ne10, cudaStream_main);
}
mul_f32_cuda(src0_ddf_i, src1_ddf_i, dst_ddf_i, ne00*i01_diff, ne10*ne11, cudaStream_main);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The broadcasting logic here is still incorrect. The implementation on master broadcasts the values per row while this broadcasts the values after flattening both tensors. As long as ne11 == 1 this doesn't make a difference but I don't think that this is the implementation that we should be using.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I noticed that the column number of src0 and src1 is ensured to be the same by checking ggml_can_repeat_rows, so broadcasting after flattening 2d sub blocks should be the same as broadcasting every row. Did I miss something?

https://github.com/ggerganov/llama.cpp/blob/32c54116318929c90fd7ae814cf9b5232cd44c36/ggml.c#L5228-L5235

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, according to the commit history the following seems to have happened: I implemented broadcasting for multiplication both in dimension 0 and dimension 1. @ggerganov then added the additional requirement that dimension 0 must be equal via ggml_can_repeat_rows which effectively limits broadcasting to dimension 1. If that is indeed the specification to which broadcasting should be implemented then the broadcasting logic in this PR is correct.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For now let's broadcast only over dimension 1 - later we'll fix the TODOs and support dimension 0 broadcasts.
Can also obsolete ggml_scale() when we do that.


(void) dst;
(void) src0_ddq_i;
Expand Down