mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-07-21 10:15:53 +00:00
cuda : add support for f16->f16 GGML_OP_SET_ROWS (#25367)
This commit is contained in:
@@ -4709,10 +4709,16 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
|
||||
} break;
|
||||
case GGML_OP_SET_ROWS:
|
||||
{
|
||||
return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 ||
|
||||
op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 ||
|
||||
op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
return (
|
||||
(
|
||||
(op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 ||
|
||||
op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 ||
|
||||
op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL) &&
|
||||
op->src[0]->type == GGML_TYPE_F32
|
||||
) || (
|
||||
op->type == GGML_TYPE_F16 && op->src[0]->type == GGML_TYPE_F16
|
||||
)
|
||||
) &&
|
||||
(op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32);
|
||||
} break;
|
||||
case GGML_OP_SET:
|
||||
|
||||
@@ -322,17 +322,77 @@ static void set_rows_cuda(ggml_backend_cuda_context & ctx, const ggml_tensor * s
|
||||
}
|
||||
}
|
||||
|
||||
template<>
|
||||
void set_rows_cuda<half, int32_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
|
||||
const half * src0_d = (const half *)src0->data;
|
||||
const int32_t * src1_d = (const int32_t *)src1->data;
|
||||
|
||||
GGML_TENSOR_BINARY_OP_LOCALS
|
||||
|
||||
cudaStream_t stream = ctx.stream();
|
||||
|
||||
|
||||
if (dst->type == GGML_TYPE_F16) {
|
||||
set_rows_cuda(
|
||||
src0_d, src1_d, (half*)dst->data,
|
||||
ne00, ne01, ne02, ne03,
|
||||
ne10, ne11, ne12, ne13,
|
||||
nb01, nb02, nb03,
|
||||
nb10, nb11, nb12,
|
||||
nb1, nb2, nb3,
|
||||
stream
|
||||
);
|
||||
} else {
|
||||
GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
|
||||
}
|
||||
}
|
||||
|
||||
template<>
|
||||
void set_rows_cuda<half, int64_t>(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
|
||||
const half * src0_d = (const half *)src0->data;
|
||||
const int64_t * src1_d = (const int64_t *)src1->data;
|
||||
|
||||
GGML_TENSOR_BINARY_OP_LOCALS
|
||||
|
||||
cudaStream_t stream = ctx.stream();
|
||||
|
||||
|
||||
if (dst->type == GGML_TYPE_F16) {
|
||||
set_rows_cuda(
|
||||
src0_d, src1_d, (half*)dst->data,
|
||||
ne00, ne01, ne02, ne03,
|
||||
ne10, ne11, ne12, ne13,
|
||||
nb01, nb02, nb03,
|
||||
nb10, nb11, nb12,
|
||||
nb1, nb2, nb3,
|
||||
stream
|
||||
);
|
||||
} else {
|
||||
GGML_ABORT("unsupported type %s", ggml_type_name(dst->type));
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
void ggml_cuda_op_set_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
|
||||
const ggml_tensor * src0 = dst->src[0];
|
||||
const ggml_tensor * src1 = dst->src[1];
|
||||
|
||||
GGML_ASSERT(src0->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(src0->type == GGML_TYPE_F32 || (src0->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16));
|
||||
GGML_ASSERT(src1->type == GGML_TYPE_I64 || src1->type == GGML_TYPE_I32);
|
||||
|
||||
if (src1->type == GGML_TYPE_I64) {
|
||||
set_rows_cuda<float, int64_t>(ctx, src0, src1, dst);
|
||||
if (src0->type == GGML_TYPE_F32) {
|
||||
if (src1->type == GGML_TYPE_I64) {
|
||||
set_rows_cuda<float, int64_t>(ctx, src0, src1, dst);
|
||||
} else {
|
||||
set_rows_cuda<float, int32_t>(ctx, src0, src1, dst);
|
||||
}
|
||||
} else if (src0->type == GGML_TYPE_F16) {
|
||||
if (src1->type == GGML_TYPE_I64) {
|
||||
set_rows_cuda<half, int64_t>(ctx, src0, src1, dst);
|
||||
} else {
|
||||
set_rows_cuda<half, int32_t>(ctx, src0, src1, dst);
|
||||
}
|
||||
} else {
|
||||
set_rows_cuda<float, int32_t>(ctx, src0, src1, dst);
|
||||
GGML_ABORT("unsupported type %s", ggml_type_name(src0->type));
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user