Skip to content

Commit 182655e

Browse files
hipuddinghipudding
committed
cann: add SWIGLU_OAI and ADD_ID op support for gpt-oss MoE
- Implement ggml_cann_swiglu_oai using aclnnClippedSwiglu (concat gate/value halves via aclnnCat, then split inside the kernel) - Implement ggml_cann_add_id via aclnnIndexSelect + aclnn_add to scatter expert outputs back onto the hidden state - Register both ops in compute_forward and supports_op - Add debug log in ggml_backend_dev_supports_op for unsupported ops - Detect ARM fp16 vector arithmetic feature in CPU CMake
1 parent 886c3a2 commit 182655e

5 files changed

Lines changed: 165 additions & 1 deletion

File tree

‎ggml/src/ggml-backend.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -622,7 +622,11 @@ ggml_backend_buffer_t ggml_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t d
622622

623623
bool ggml_backend_dev_supports_op(ggml_backend_dev_t device, const struct ggml_tensor * op) {
624624
GGML_ASSERT(device);
625-
return device->iface.supports_op(device, op);
625+
bool result = device->iface.supports_op(device, op);
626+
if (!result) {
627+
GGML_LOG_DEBUG("%s: backend %s does not support op %s\n", __func__, ggml_backend_dev_name(device), ggml_op_name(op->op));
628+
}
629+
return result;
626630
}
627631

628632
bool ggml_backend_dev_supports_buft(ggml_backend_dev_t device, ggml_backend_buffer_type_t buft) {

‎ggml/src/ggml-cann/aclnn_ops.cpp‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,9 @@
3333
#include <aclnnop/aclnn_avgpool2d.h>
3434
#include <aclnnop/aclnn_batch_matmul.h>
3535
#include <aclnnop/aclnn_cast.h>
36+
#include <aclnnop/aclnn_cat.h>
3637
#include <aclnnop/aclnn_clamp.h>
38+
#include <aclnnop/aclnn_clipped_swiglu.h>
3739
#include <aclnnop/aclnn_constant_pad_nd.h>
3840
#include <aclnnop/aclnn_convolution.h>
3941
#include <aclnnop/aclnn_copy.h>
@@ -211,6 +213,68 @@ void ggml_cann_swiglu(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
211213
GGML_CANN_CALL_ACLNN_OP(ctx, SwiGlu, acl_src.get(), (int64_t)2, acl_dst.get());
212214
}
213215

216+
// SWIGLU_OAI (gpt-oss MoE FFN). Per element, with A the gate half and B the
217+
// value half:
218+
// A = min(A, limit); B = clamp(B, -limit, limit);
219+
// y = A * sigmoid(alpha * A) * (B + 1)
220+
// This is exactly aclnnClippedSwiglu (bias = 1, interleaved = false, A = first
221+
// half, B = second half). aclnnClippedSwiglu takes a single tensor and splits
222+
// it, so the two halves are concatenated as [gate; value] along ne[0] first.
223+
// `swapped` reverses which half is the gate.
224+
void ggml_cann_swiglu_oai(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
225+
ggml_tensor * src0 = dst->src[0];
226+
ggml_tensor * src1 = dst->src[1];
227+
228+
GGML_ASSERT(ggml_is_contiguous_1(src0));
229+
GGML_ASSERT(ggml_is_contiguous_1(dst));
230+
GGML_ASSERT(src0->type == GGML_TYPE_F32);
231+
GGML_ASSERT(dst->type == GGML_TYPE_F32);
232+
233+
const int32_t swapped = ggml_get_op_params_i32(dst, 1);
234+
const float alpha = ggml_get_op_params_f32(dst, 2);
235+
const float limit = ggml_get_op_params_f32(dst, 3);
236+
237+
// acl_a/acl_b are the gate/value candidates; swapped flips which is the gate.
238+
acl_tensor_ptr acl_a;
239+
acl_tensor_ptr acl_b;
240+
if (src1 == nullptr) {
241+
GGML_ASSERT(src0->ne[0] % 2 == 0);
242+
int64_t ne_h[GGML_MAX_DIMS] = { src0->ne[0] / 2, src0->ne[1], src0->ne[2], src0->ne[3] };
243+
size_t nb_h[GGML_MAX_DIMS] = { (size_t)src0->nb[0], (size_t)src0->nb[1],
244+
(size_t)src0->nb[2], (size_t)src0->nb[3] };
245+
const size_t off = ne_h[0] * ggml_element_size(src0);
246+
acl_a = ggml_cann_create_tensor(src0, ne_h, nb_h, GGML_MAX_DIMS, ACL_FORMAT_ND, 0); // first half
247+
acl_b = ggml_cann_create_tensor(src0, ne_h, nb_h, GGML_MAX_DIMS, ACL_FORMAT_ND, off); // second half
248+
} else {
249+
GGML_ASSERT(ggml_is_contiguous_1(src1));
250+
GGML_ASSERT(src0->type == src1->type);
251+
acl_a = ggml_cann_create_tensor(src0);
252+
acl_b = ggml_cann_create_tensor(src1);
253+
}
254+
if (swapped) {
255+
std::swap(acl_a, acl_b); // acl_a = gate (first half of the input to ClippedSwiglu)
256+
}
257+
258+
// concat [acl_a (gate); acl_b (value)] along ne[0] -> ClippedSwiglu splits it
259+
const size_t es = ggml_element_size(dst);
260+
int64_t cat_ne[GGML_MAX_DIMS] = { dst->ne[0] * 2, dst->ne[1], dst->ne[2], dst->ne[3] };
261+
size_t cat_nb[GGML_MAX_DIMS];
262+
cat_nb[0] = es;
263+
for (int i = 1; i < GGML_MAX_DIMS; i++) {
264+
cat_nb[i] = cat_nb[i - 1] * cat_ne[i - 1];
265+
}
266+
ggml_cann_pool_alloc cat_alloc(ctx.pool(), cat_nb[GGML_MAX_DIMS - 1] * cat_ne[GGML_MAX_DIMS - 1]);
267+
acl_tensor_ptr acl_x = ggml_cann_create_tensor(cat_alloc.get(), ggml_cann_type_mapping(dst->type),
268+
es, cat_ne, cat_nb, GGML_MAX_DIMS);
269+
acl_tensor_list_ptr list = ggml_cann_create_tensor_list(acl_a, acl_b);
270+
GGML_CANN_CALL_ACLNN_OP(ctx, Cat, list.get(), (int64_t)(GGML_MAX_DIMS - 1), acl_x.get());
271+
272+
acl_tensor_ptr acl_out = ggml_cann_create_tensor(dst);
273+
GGML_CANN_CALL_ACLNN_OP(ctx, ClippedSwiglu, acl_x.get(), (const aclTensor *) nullptr,
274+
(int64_t)(GGML_MAX_DIMS - 1), (double) alpha, (double) limit, (double) 1.0,
275+
/*interleaved*/ false, acl_out.get());
276+
}
277+
214278
// Fused GeGLU using aclnnGeGluV3: splits input along ne[0] (CANN last dim),
215279
// activates the LEFT half with GELU, multiplies by right half.
216280
// approximate: 0=tanh, 1=none(erf). activateLeft=true matches GGML convention.
@@ -3550,6 +3614,71 @@ void ggml_cann_geglu_quick(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
35503614
ggml_cann_op_unary_gated(gelu_quick_fn, ctx, dst);
35513615
}
35523616

3617+
/**
3618+
* @brief Adds expert-specific rows to a tensor (MoE) using the CANN backend.
3619+
*
3620+
* Implements GGML_OP_ADD_ID. For each row (i2, i1) of src0, the row of src1
3621+
* selected by ids[i1, i2] is added element-wise. Used to scatter expert
3622+
* outputs back onto the hidden state in Mixture-of-Experts models.
3623+
*
3624+
* Dimensions:
3625+
* - src0 (a) : [ne0, ne1, ne2, ne3], F32
3626+
* - src1 (b) : [ne0, ne11], F32 (ne11 rows to select from)
3627+
* - ids : [ne1, ne2], I32
3628+
* - dst : same shape as src0
3629+
*
3630+
* For each i2, IndexSelect gathers ne1 rows from src1 according to ids[:, i2]
3631+
* into a [ne0, ne1] scratch buffer, which is then added to every matching
3632+
* (i2, i3) slice of src0.
3633+
*/
3634+
void ggml_cann_add_id(ggml_backend_cann_context & ctx, ggml_tensor * dst) {
3635+
ggml_tensor * src0 = dst->src[0]; // a: [ne0, ne1, ne2, 1]
3636+
ggml_tensor * src1 = dst->src[1]; // b: [ne0, ne11]
3637+
ggml_tensor * ids = dst->src[2]; // ids: [ne1, ne2]
3638+
3639+
GGML_ASSERT(dst->type == GGML_TYPE_F32);
3640+
GGML_ASSERT(src0->type == GGML_TYPE_F32);
3641+
GGML_ASSERT(src1->type == GGML_TYPE_F32);
3642+
GGML_ASSERT(ids->type == GGML_TYPE_I32);
3643+
GGML_ASSERT(src0->ne[3] == 1);
3644+
GGML_ASSERT(dst->ne[3] == 1);
3645+
3646+
const int64_t ne0 = src0->ne[0];
3647+
const int64_t ne1 = src0->ne[1];
3648+
const int64_t ne2 = src0->ne[2];
3649+
3650+
// scratch buffer holding the selected src1 rows: [ne0, ne1]
3651+
ggml_cann_pool_alloc select_allocator(ctx.pool(), ne0 * ne1 * ggml_element_size(src1));
3652+
void * select_ptr = select_allocator.get();
3653+
3654+
int64_t select_ne[2] = { ne0, ne1 };
3655+
size_t select_nb[2];
3656+
select_nb[0] = src1->nb[0];
3657+
select_nb[1] = select_nb[0] * ne0;
3658+
3659+
int64_t slice_ne[2] = { ne0, ne1 };
3660+
size_t a_slice_nb[2] = { src0->nb[0], src0->nb[1] };
3661+
size_t d_slice_nb[2] = { dst->nb[0], dst->nb[1] };
3662+
3663+
acl_tensor_ptr acl_b = ggml_cann_create_tensor(src1, src1->ne, src1->nb, 2);
3664+
3665+
for (int64_t i2 = 0; i2 < ne2; i2++) {
3666+
// ids[:, i2]: ne1 indices into b's rows
3667+
acl_tensor_ptr select_index =
3668+
ggml_cann_create_tensor(ids, ids->ne, ids->nb, 1, ACL_FORMAT_ND, i2 * ids->nb[1]);
3669+
acl_tensor_ptr select_out = ggml_cann_create_tensor(select_ptr, ggml_cann_type_mapping(src1->type),
3670+
ggml_element_size(src1), select_ne, select_nb, 2);
3671+
// gather ne1 rows from b along its row axis (ACL dim 0 == ggml ne[1])
3672+
GGML_CANN_CALL_ACLNN_OP(ctx, IndexSelect, acl_b.get(), 0, select_index.get(), select_out.get());
3673+
3674+
acl_tensor_ptr acl_a_slice =
3675+
ggml_cann_create_tensor(src0, slice_ne, a_slice_nb, 2, ACL_FORMAT_ND, i2 * src0->nb[2]);
3676+
acl_tensor_ptr acl_d_slice =
3677+
ggml_cann_create_tensor(dst, slice_ne, d_slice_nb, 2, ACL_FORMAT_ND, i2 * dst->nb[2]);
3678+
aclnn_add(ctx, acl_a_slice.get(), select_out.get(), acl_d_slice.get());
3679+
}
3680+
}
3681+
35533682
/**
35543683
* @brief Performs expert-specific matrix multiplication (MoE) with
35553684
* floating-point precision using the CANN backend.

‎ggml/src/ggml-cann/aclnn_ops.h‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
void ggml_cann_repeat(ggml_backend_cann_context & ctx, ggml_tensor * dst);
7777

7878
void ggml_cann_swiglu(ggml_backend_cann_context & ctx, ggml_tensor * dst);
79+
void ggml_cann_swiglu_oai(ggml_backend_cann_context & ctx, ggml_tensor * dst);
7980
void ggml_cann_geglu(ggml_backend_cann_context & ctx, ggml_tensor * dst, int64_t approximate);
8081

8182
/**
@@ -959,6 +960,24 @@ void ggml_cann_gated_linear_attn(ggml_backend_cann_context & ctx, ggml_tensor *
959960
*/
960961
void ggml_cann_mul_mat_id(ggml_backend_cann_context & ctx, ggml_tensor * dst);
961962

963+
/**
964+
* @brief Adds expert-specific rows to a tensor (MoE) using the CANN backend.
965+
*
966+
* Implements GGML_OP_ADD_ID. For each row (i2, i1) of src0, the row of src1
967+
* selected by ids[i1, i2] is added element-wise and stored in dst.
968+
*
969+
* Dimensions:
970+
* - src0 (a) : [ne0, ne1, ne2, ne3], F32
971+
* - src1 (b) : [ne0, ne11], F32 (ne11 rows to select from)
972+
* - ids : [ne1, ne2], I32
973+
* - dst : same shape as src0
974+
*
975+
* @param ctx The CANN context used for operations.
976+
* @param dst The destination tensor where the result is stored. dst->op is
977+
* `GGML_OP_ADD_ID`.
978+
*/
979+
void ggml_cann_add_id(ggml_backend_cann_context & ctx, ggml_tensor * dst);
980+
962981
/**
963982
* @brief Performs fused ADD + RMS_NORM operation using the CANN backend.
964983
*

‎ggml/src/ggml-cann/ggml-cann.cpp‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1788,6 +1788,9 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context & ctx, struct gg
17881788
case GGML_OP_ADD1:
17891789
ggml_cann_binary_op<aclnn_add>(ctx, dst);
17901790
break;
1791+
case GGML_OP_ADD_ID:
1792+
ggml_cann_add_id(ctx, dst);
1793+
break;
17911794
case GGML_OP_SUB:
17921795
ggml_cann_binary_op<aclnn_sub>(ctx, dst);
17931796
break;
@@ -1872,6 +1875,9 @@ static bool ggml_cann_compute_forward(ggml_backend_cann_context & ctx, struct gg
18721875
case GGML_GLU_OP_SWIGLU:
18731876
ggml_cann_swiglu(ctx, dst);
18741877
break;
1878+
case GGML_GLU_OP_SWIGLU_OAI:
1879+
ggml_cann_swiglu_oai(ctx, dst);
1880+
break;
18751881
case GGML_GLU_OP_GEGLU_QUICK:
18761882
ggml_cann_geglu_quick(ctx, dst);
18771883
break;
@@ -2426,6 +2432,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
24262432
case GGML_GLU_OP_REGLU:
24272433
case GGML_GLU_OP_GEGLU:
24282434
case GGML_GLU_OP_SWIGLU:
2435+
case GGML_GLU_OP_SWIGLU_OAI:
24292436
case GGML_GLU_OP_GEGLU_ERF:
24302437
case GGML_GLU_OP_GEGLU_QUICK:
24312438
return true;
@@ -2470,6 +2477,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
24702477
default:
24712478
return false;
24722479
}
2480+
case GGML_OP_ADD_ID:
2481+
return op->src[0]->type == GGML_TYPE_F32 &&
2482+
op->src[1]->type == GGML_TYPE_F32 &&
2483+
op->src[2]->type == GGML_TYPE_I32;
24732484
// embedding
24742485
case GGML_OP_GET_ROWS:
24752486
{

‎ggml/src/ggml-cpu/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,6 +163,7 @@ function(ggml_add_cpu_backend_variant_impl tag_name)
163163
check_arm_feature(i8mm MATMUL_INT8 "#include <arm_neon.h>\nint main() { int8x16_t _a, _b; volatile int32x4_t _s = vmmlaq_s32(_s, _a, _b); return 0; }")
164164
check_arm_feature(sve SVE "#include <arm_sve.h>\nint main() { svfloat32_t _a, _b; volatile svfloat32_t _c = svadd_f32_z(svptrue_b8(), _a, _b); return 0; }")
165165
check_arm_feature(sme SME "#include <arm_sme.h>\n__arm_locally_streaming int main() { __asm__ volatile(\"smstart; smstop;\"); return 0; }")
166+
check_arm_feature(fp16 FP16_VECTOR_ARITHMETIC "#include <arm_neon.h>\nint main() { float16x8_t _a, _b; volatile float16x8_t _c = vaddq_f16(_a, _b); return 0; }")
166167

167168
list(APPEND ARCH_FLAGS "${ARM_NATIVE_FLAG}${ARM_NATIVE_FLAG_FIX}")
168169
else()

0 commit comments

Comments
 (0)