mirror of
https://github.com/LostRuins/koboldcpp.git
synced 2026-08-07 23:45:30 +00:00
metal: implement DeepSeek V4 hyper-connections (#26459)
- Implement GGML_OP_DSV4_HC_COMB, GGML_OP_DSV4_HC_PRE, and GGML_OP_DSV4_HC_POST with SIMDgroup register and shuffle optimized kernels. - Add Metal dispatch and support plumbing and test the production Sinkhorn iteration count and embedding width. Assisted-by: Codex Co-authored-by: Thiago Padilha <thiago@padilha.cc>
This commit is contained in:
Georgi Gerganov
GitHub
parent
bb4e0e1b3f
commit
fffbcbdb9d
8 changed files with 397 additions and 0 deletions
|
|
@ -477,6 +477,24 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_me
|
|||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) {
|
||||
const char * name = nullptr;
|
||||
|
||||
switch (op) {
|
||||
case GGML_OP_DSV4_HC_COMB: name = "kernel_dsv4_hc_comb_f32"; break;
|
||||
case GGML_OP_DSV4_HC_PRE: name = "kernel_dsv4_hc_pre_f32"; break;
|
||||
case GGML_OP_DSV4_HC_POST: name = "kernel_dsv4_hc_post_f32"; break;
|
||||
default: GGML_ABORT("fatal error");
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
|
||||
}
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) {
|
||||
GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
|
||||
|
|
|
|||
|
|
@ -124,6 +124,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_bl
|
|||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs);
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op);
|
||||
|
|
|
|||
|
|
@ -1299,6 +1299,42 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
|
|||
return false;
|
||||
}
|
||||
return has_simdgroup_mm; // TODO: over-restricted for vec-kernels
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
return has_simdgroup_reduction &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->src[2]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32 &&
|
||||
op->src[0]->ne[0] == 24 &&
|
||||
op->src[1]->ne[0] >= 3 &&
|
||||
op->src[2]->ne[0] == 24 &&
|
||||
ggml_is_contiguous_rows(op->src[0]) &&
|
||||
ggml_is_contiguous_rows(op->src[1]) &&
|
||||
ggml_is_contiguous_rows(op->src[2]);
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
return has_simdgroup_reduction &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32 &&
|
||||
op->src[0]->ne[1] == 4 &&
|
||||
op->src[1]->ne[0] == 4 &&
|
||||
ggml_is_contiguous_rows(op->src[0]) &&
|
||||
ggml_is_contiguous_rows(op->src[1]);
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
return has_simdgroup_reduction &&
|
||||
op->src[0]->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->src[2]->type == GGML_TYPE_F32 &&
|
||||
op->src[3]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32 &&
|
||||
op->src[1]->ne[1] == 4 &&
|
||||
op->src[2]->ne[0] == 4 &&
|
||||
op->src[3]->ne[0] == 4 &&
|
||||
op->src[3]->ne[1] == 4 &&
|
||||
ggml_is_contiguous_rows(op->src[0]) &&
|
||||
ggml_is_contiguous_rows(op->src[1]) &&
|
||||
ggml_is_contiguous_rows(op->src[2]) &&
|
||||
ggml_is_contiguous_rows(op->src[3]);
|
||||
case GGML_OP_SSM_CONV:
|
||||
case GGML_OP_SSM_SCAN:
|
||||
return has_simdgroup_reduction;
|
||||
|
|
|
|||
|
|
@ -1171,6 +1171,49 @@ typedef struct {
|
|||
int64_t val;
|
||||
} ggml_metal_kargs_memset;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_tokens;
|
||||
int32_t n_iter;
|
||||
uint64_t nb_m0;
|
||||
uint64_t nb_m1;
|
||||
uint64_t nb_s0;
|
||||
uint64_t nb_b0;
|
||||
uint64_t nb_d0;
|
||||
uint64_t nb_d1;
|
||||
uint64_t nb_d2;
|
||||
float eps;
|
||||
} ggml_metal_kargs_dsv4_hc_comb;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_embd;
|
||||
int32_t n_tokens;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_x2;
|
||||
uint64_t nb_w0;
|
||||
uint64_t nb_w1;
|
||||
uint64_t nb_d0;
|
||||
uint64_t nb_d1;
|
||||
} ggml_metal_kargs_dsv4_hc_pre;
|
||||
|
||||
typedef struct {
|
||||
int32_t n_embd;
|
||||
int32_t n_tokens;
|
||||
uint64_t nb_x0;
|
||||
uint64_t nb_x1;
|
||||
uint64_t nb_r0;
|
||||
uint64_t nb_r1;
|
||||
uint64_t nb_r2;
|
||||
uint64_t nb_p0;
|
||||
uint64_t nb_p1;
|
||||
uint64_t nb_c0;
|
||||
uint64_t nb_c1;
|
||||
uint64_t nb_c2;
|
||||
uint64_t nb_d0;
|
||||
uint64_t nb_d1;
|
||||
uint64_t nb_d2;
|
||||
} ggml_metal_kargs_dsv4_hc_post;
|
||||
|
||||
typedef struct {
|
||||
int32_t ne00;
|
||||
int32_t ne01;
|
||||
|
|
|
|||
|
|
@ -316,6 +316,12 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
|
|||
{
|
||||
n_fuse = ggml_metal_op_cumsum(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
{
|
||||
n_fuse = ggml_metal_op_dsv4_hc(ctx, idx);
|
||||
} break;
|
||||
case GGML_OP_SOFT_MAX:
|
||||
{
|
||||
n_fuse = ggml_metal_op_soft_max(ctx, idx);
|
||||
|
|
@ -1297,6 +1303,137 @@ int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) {
|
|||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
ggml_metal_encoder_t enc = ctx->enc;
|
||||
auto pipeline = ggml_metal_library_get_pipeline_dsv4_hc(ctx->lib, op->op);
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline);
|
||||
|
||||
switch (op->op) {
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
{
|
||||
const ggml_tensor * mixes = op->src[0];
|
||||
const ggml_tensor * scale = op->src[1];
|
||||
const ggml_tensor * base = op->src[2];
|
||||
|
||||
GGML_ASSERT(mixes->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(scale->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(base->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(mixes->ne[0] == 24);
|
||||
GGML_ASSERT(op->ne[0] == 4 && op->ne[1] == 4);
|
||||
|
||||
ggml_metal_kargs_dsv4_hc_comb args = {
|
||||
/*.n_tokens =*/ (int32_t) mixes->ne[1],
|
||||
/*.n_iter =*/ ggml_get_op_params_i32(op, 1),
|
||||
/*.nb_m0 =*/ mixes->nb[0],
|
||||
/*.nb_m1 =*/ mixes->nb[1],
|
||||
/*.nb_s0 =*/ scale->nb[0],
|
||||
/*.nb_b0 =*/ base->nb[0],
|
||||
/*.nb_d0 =*/ op->nb[0],
|
||||
/*.nb_d1 =*/ op->nb[1],
|
||||
/*.nb_d2 =*/ op->nb[2],
|
||||
/*.eps =*/ ggml_get_op_params_f32(op, 0),
|
||||
};
|
||||
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(mixes), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(scale), 2);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(base), 3);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4);
|
||||
|
||||
// One SIMDgroup owns one 4x4 Sinkhorn matrix. Packing up to four
|
||||
// independent tokens per threadgroup keeps both decode and prompt
|
||||
// dispatches compact without any threadgroup-memory synchronization.
|
||||
const int nsg = std::min(4, args.n_tokens);
|
||||
ggml_metal_encoder_dispatch_threadgroups(
|
||||
enc, (args.n_tokens + nsg - 1)/nsg, 1, 1, 32, nsg, 1);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
{
|
||||
const ggml_tensor * x = op->src[0];
|
||||
const ggml_tensor * weights = op->src[1];
|
||||
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(weights->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(x->ne[1] == 4);
|
||||
|
||||
ggml_metal_kargs_dsv4_hc_pre args = {
|
||||
/*.n_embd =*/ (int32_t) x->ne[0],
|
||||
/*.n_tokens =*/ (int32_t) x->ne[2],
|
||||
/*.nb_x0 =*/ x->nb[0],
|
||||
/*.nb_x1 =*/ x->nb[1],
|
||||
/*.nb_x2 =*/ x->nb[2],
|
||||
/*.nb_w0 =*/ weights->nb[0],
|
||||
/*.nb_w1 =*/ weights->nb[1],
|
||||
/*.nb_d0 =*/ op->nb[0],
|
||||
/*.nb_d1 =*/ op->nb[1],
|
||||
};
|
||||
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weights), 2);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3);
|
||||
|
||||
const int n_tiles = (args.n_embd + 31)/32;
|
||||
const int nsg = std::min(4, n_tiles);
|
||||
ggml_metal_encoder_dispatch_threadgroups(
|
||||
enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
{
|
||||
const ggml_tensor * x = op->src[0];
|
||||
const ggml_tensor * residual = op->src[1];
|
||||
const ggml_tensor * post = op->src[2];
|
||||
const ggml_tensor * comb = op->src[3];
|
||||
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(residual->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(post->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(comb->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(op->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(residual->ne[1] == 4);
|
||||
|
||||
ggml_metal_kargs_dsv4_hc_post args = {
|
||||
/*.n_embd =*/ (int32_t) x->ne[0],
|
||||
/*.n_tokens =*/ (int32_t) x->ne[1],
|
||||
/*.nb_x0 =*/ x->nb[0],
|
||||
/*.nb_x1 =*/ x->nb[1],
|
||||
/*.nb_r0 =*/ residual->nb[0],
|
||||
/*.nb_r1 =*/ residual->nb[1],
|
||||
/*.nb_r2 =*/ residual->nb[2],
|
||||
/*.nb_p0 =*/ post->nb[0],
|
||||
/*.nb_p1 =*/ post->nb[1],
|
||||
/*.nb_c0 =*/ comb->nb[0],
|
||||
/*.nb_c1 =*/ comb->nb[1],
|
||||
/*.nb_c2 =*/ comb->nb[2],
|
||||
/*.nb_d0 =*/ op->nb[0],
|
||||
/*.nb_d1 =*/ op->nb[1],
|
||||
/*.nb_d2 =*/ op->nb[2],
|
||||
};
|
||||
|
||||
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(x), 1);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(residual), 2);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(post), 3);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(comb), 4);
|
||||
ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5);
|
||||
|
||||
const int n_tiles = (args.n_embd + 31)/32;
|
||||
const int nsg = std::min(4, n_tiles);
|
||||
ggml_metal_encoder_dispatch_threadgroups(
|
||||
enc, (n_tiles + nsg - 1)/nsg, args.n_tokens, 1, 32, nsg, 1);
|
||||
} break;
|
||||
default:
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
|
||||
return 1;
|
||||
}
|
||||
|
||||
int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
|
|
|
|||
|
|
@ -54,6 +54,7 @@ int ggml_metal_op_cumsum (ggml_metal_op_t ctx, int idx);
|
|||
int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_ssm_scan (ggml_metal_op_t ctx, int idx);
|
||||
|
|
|
|||
|
|
@ -11278,3 +11278,162 @@ kernel void kernel_count_equal(
|
|||
typedef decltype(kernel_count_equal<int32_t>) kernel_count_equal_t;
|
||||
|
||||
template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
|
||||
|
||||
kernel void kernel_dsv4_hc_comb_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_comb & args,
|
||||
device const char * mixes,
|
||||
device const char * scale,
|
||||
device const char * base,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
constexpr ushort comb_offset = 2*hc;
|
||||
|
||||
const int it = tgpig.x*ntg.y + sgitg;
|
||||
if (it >= args.n_tokens) {
|
||||
return;
|
||||
}
|
||||
|
||||
float scale_lane = 0.0f;
|
||||
if (tiisg == 0) {
|
||||
scale_lane = *(device const float *) (scale + 2*args.nb_s0);
|
||||
}
|
||||
const float scale_comb = simd_shuffle(scale_lane, 0);
|
||||
|
||||
float v = 0.0f;
|
||||
if (tiisg < hc*hc) {
|
||||
v = *(device const float *) (mixes + (comb_offset + tiisg)*args.nb_m0 + it*args.nb_m1)*scale_comb
|
||||
+ *(device const float *) (base + (comb_offset + tiisg)*args.nb_b0);
|
||||
}
|
||||
|
||||
// Softmax across destinations (the four contiguous lanes for each source).
|
||||
float vmax = max(v, simd_shuffle_xor(v, 1));
|
||||
vmax = max(vmax, simd_shuffle_xor(vmax, 2));
|
||||
v = exp(v - vmax);
|
||||
|
||||
float sum = v + simd_shuffle_xor(v, 1);
|
||||
sum += simd_shuffle_xor(sum, 2);
|
||||
v = v/sum + args.eps;
|
||||
|
||||
// Normalize columns: equal destination indices are four lanes apart.
|
||||
sum = v + simd_shuffle_xor(v, 4);
|
||||
sum += simd_shuffle_xor(sum, 8);
|
||||
v /= sum + args.eps;
|
||||
|
||||
for (int i = 1; i < args.n_iter; ++i) {
|
||||
sum = v + simd_shuffle_xor(v, 1);
|
||||
sum += simd_shuffle_xor(sum, 2);
|
||||
v /= sum + args.eps;
|
||||
|
||||
sum = v + simd_shuffle_xor(v, 4);
|
||||
sum += simd_shuffle_xor(sum, 8);
|
||||
v /= sum + args.eps;
|
||||
}
|
||||
|
||||
if (tiisg < hc*hc) {
|
||||
const ushort idst = tiisg & 3;
|
||||
const ushort isrc = tiisg >> 2;
|
||||
*(device float *) (dst + idst*args.nb_d0 + isrc*args.nb_d1 + it*args.nb_d2) = v;
|
||||
}
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_pre_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_pre & args,
|
||||
device const char * x,
|
||||
device const char * weights,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
|
||||
const int it = tgpig.y;
|
||||
const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
|
||||
|
||||
float weight_lane = 0.0f;
|
||||
if (tiisg < hc) {
|
||||
weight_lane = *(device const float *) (weights + tiisg*args.nb_w0 + it*args.nb_w1);
|
||||
}
|
||||
|
||||
float w[hc];
|
||||
FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
|
||||
w[ih] = simd_shuffle(weight_lane, ih);
|
||||
}
|
||||
|
||||
if (i0 >= args.n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
device const char * xb = x + i0*args.nb_x0 + it*args.nb_x2;
|
||||
float result = 0.0f;
|
||||
FOR_UNROLL (ushort ih = 0; ih < hc; ++ih) {
|
||||
result = fma(*(device const float *) (xb + ih*args.nb_x1), w[ih], result);
|
||||
}
|
||||
|
||||
*(device float *) (dst + i0*args.nb_d0 + it*args.nb_d1) = result;
|
||||
}
|
||||
|
||||
kernel void kernel_dsv4_hc_post_f32(
|
||||
constant ggml_metal_kargs_dsv4_hc_post & args,
|
||||
device const char * x,
|
||||
device const char * residual,
|
||||
device const char * post,
|
||||
device const char * comb,
|
||||
device char * dst,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
ushort sgitg[[simdgroup_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr ushort hc = 4;
|
||||
|
||||
const int it = tgpig.y;
|
||||
const int i0 = ((int) tgpig.x*ntg.y + sgitg)*32 + tiisg;
|
||||
|
||||
float coeff_lane = 0.0f;
|
||||
if (tiisg < hc) {
|
||||
coeff_lane = *(device const float *) (post + tiisg*args.nb_p0 + it*args.nb_p1);
|
||||
} else if (tiisg < hc + hc*hc) {
|
||||
const ushort idx = tiisg - hc;
|
||||
const ushort idst = idx & 3;
|
||||
const ushort isrc = idx >> 2;
|
||||
coeff_lane = *(device const float *) (comb + idst*args.nb_c0 + isrc*args.nb_c1 + it*args.nb_c2);
|
||||
}
|
||||
|
||||
float post_reg[hc];
|
||||
float comb_reg[hc][hc];
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
post_reg[idst] = simd_shuffle(coeff_lane, idst);
|
||||
}
|
||||
FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
comb_reg[isrc][idst] = simd_shuffle(coeff_lane, hc + idst + hc*isrc);
|
||||
}
|
||||
}
|
||||
|
||||
if (i0 >= args.n_embd) {
|
||||
return;
|
||||
}
|
||||
|
||||
const float xv = *(device const float *) (x + i0*args.nb_x0 + it*args.nb_x1);
|
||||
float result[hc];
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
result[idst] = xv*post_reg[idst];
|
||||
}
|
||||
|
||||
device const char * rb = residual + i0*args.nb_r0 + it*args.nb_r2;
|
||||
FOR_UNROLL (ushort isrc = 0; isrc < hc; ++isrc) {
|
||||
const float rv = *(device const float *) (rb + isrc*args.nb_r1);
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
result[idst] = fma(rv, comb_reg[isrc][idst], result[idst]);
|
||||
}
|
||||
}
|
||||
|
||||
FOR_UNROLL (ushort idst = 0; idst < hc; ++idst) {
|
||||
*(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -8069,6 +8069,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
|||
test_cases.emplace_back(new test_dsv4_hc_comb(1, 1));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(17, 4));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(257, 8));
|
||||
test_cases.emplace_back(new test_dsv4_hc_comb(17, 20));
|
||||
|
||||
test_cases.emplace_back(new test_dsv4_hc_pre(1, 1));
|
||||
test_cases.emplace_back(new test_dsv4_hc_pre(31, 17));
|
||||
|
|
@ -8078,6 +8079,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
|||
test_cases.emplace_back(new test_dsv4_hc_post(1, 1));
|
||||
test_cases.emplace_back(new test_dsv4_hc_post(31, 17));
|
||||
test_cases.emplace_back(new test_dsv4_hc_post(128, 257));
|
||||
test_cases.emplace_back(new test_dsv4_hc_post(4096, 21));
|
||||
|
||||
// glu ops
|
||||
for (ggml_type type : {GGML_TYPE_F16, GGML_TYPE_F32}) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue