diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d29efecffc80..915ce0fe6e1f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2557,6 +2557,22 @@ extern "C" { struct ggml_tensor * beta, struct ggml_tensor * state); + // rows-indexed state read: instead of a gathered/contiguous (D, K, n_seqs) + // scratch, the op reads each sequence's live state directly from `states` + // (2D cache view, D-wide rows) at row `rows[seq]` (I32, n_seqs entries). + // Removes the per-layer get_rows + slot-0 cpy from recurrent decode graphs. + // Output layout is identical to ggml_gated_delta_net with K = n_snap_slots. + GGML_API struct ggml_tensor * ggml_gated_delta_net_rows( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * states, + struct ggml_tensor * rows, + int n_snap_slots); + // custom operators typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index c26c3f1470d8..8c8cb827fc71 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -222,6 +222,15 @@ if (GGML_SCHED_NO_REALLOC) target_compile_definitions(ggml-base PUBLIC GGML_SCHED_NO_REALLOC) endif() +if (GGML_OPENMP AND EMSCRIPTEN) + # Emscripten/WASM cannot emit the common symbols that libomp's reduction + # helpers generate (e.g. .gomp_critical_user_.reduction.var), so an OpenMP + # build of ggml-quants.c fails to link. WASM has no host threads to gain + # from OpenMP anyway -- disable it rather than fail the build. + message(STATUS "ggml: disabling OpenMP for Emscripten/WASM target") + set(GGML_OPENMP OFF) +endif() + if (GGML_OPENMP) find_package(OpenMP) if (OpenMP_FOUND) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 05ed61a097b0..9b60d2379815 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -10559,11 +10559,16 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( const bool kda = (neg0 == S_v); - // state is 3D (S_v*S_v*H, K, n_seqs); K is the snapshot slot count. - const int64_t K = src_state->ne[1]; + // K is the snapshot slot count (op_params, shared by both op variants). + const int64_t K = ggml_get_op_params_i32(dst, 0); GGML_ASSERT(K >= 1); - // per-seq stride in floats (slot 0 of seq s lives at state + s * seq_stride) - const int64_t state_seq_stride = src_state->nb[2] / sizeof(float); + // rows mode (src[6] set): state is a 2D cache view (D, n_rows) and each + // sequence's live state is read at row rows[seq] -- no gathered scratch. + const ggml_tensor * src_rows = dst->src[6]; + const int32_t * state_rows_idx = src_rows ? (const int32_t *) src_rows->data : nullptr; + // scratch mode: per-seq stride in floats (slot 0 of seq s at s * seq_stride) + const int64_t state_seq_stride = src_rows ? 0 : (int64_t) (src_state->nb[2] / sizeof(float)); + const int64_t state_row_size = src_rows ? (int64_t) (src_state->nb[1] / sizeof(float)) : 0; const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0); const int ith = params->ith; @@ -10608,9 +10613,12 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( ? state_work : state_out_base + (iv3 * H + iv1) * S_v * S_v; - // copy input state into the working buffer and operate in-place - // state layout (D, K, n_seqs): slot 0 of seq iv3 starts at iv3 * state_seq_stride. - const float * s_in = state_in_base + iv3 * state_seq_stride + iv1 * S_v * S_v; + // copy input state into the working buffer and operate in-place. + // scratch mode: state layout (D, K, n_seqs), slot 0 of seq iv3 at + // iv3 * state_seq_stride. rows mode: cache row state_rows_idx[iv3]. + const float * s_in = state_rows_idx + ? state_in_base + (int64_t) state_rows_idx[iv3] * state_row_size + iv1 * S_v * S_v + : state_in_base + iv3 * state_seq_stride + iv1 * S_v * S_v; memcpy(s_out, s_in, S_v * S_v * sizeof(float)); // attn output pointer for first token of this (head, seq) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 66b0e0114796..13e1b8a2e737 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5417,6 +5417,11 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_RWKV_WKV7: return true; case GGML_OP_GATED_DELTA_NET: + // rows-indexed state read (src[6]) not implemented on CUDA yet; + // reject so it falls back instead of silently reading src[5] as a scratch + if (op->src[6] != NULL) { + return false; + } //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327 #ifdef GGML_USE_MUSA return false; diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index d550841a2a59..1232ec8522cf 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -3710,7 +3710,8 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons break; case GGML_OP_GATED_DELTA_NET: - supp = ggml_hexagon_supported_gated_delta_net(sess, op); + // rows-indexed state read (src[6]) not implemented here + supp = op->src[6] == NULL && ggml_hexagon_supported_gated_delta_net(sess, op); break; case GGML_OP_CUMSUM: diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 86a60c060565..f7e6dd8dbd08 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -583,7 +583,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_ return res; } -ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net(ggml_metal_library_t lib, const ggml_tensor * op) { +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net(ggml_metal_library_t lib, const ggml_tensor * op, bool write_rows) { char base[256]; char name[256]; @@ -591,8 +591,11 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net( const int ne20 = op->src[2]->ne[0]; // S_v const int ne21 = op->src[2]->ne[1]; // H const int ne30 = op->src[3]->ne[0]; // G - // state is src[5], 3D (S_v*S_v*H, K, n_seqs); K is the snapshot slot count. - const int K = op->src[5]->ne[1]; + // K (snapshot slot count) comes from op_params: in rows mode src[5] is the + // 2D cache view, so its ne[1] is the cache row count, not K. + const int K = ggml_get_op_params_i32(op, 0); + // rows mode: src[6] holds per-seq cache row indices for the state read + const bool has_rows = op->src[6] != NULL; const int nsg = op->src[2]->ne[0]/32; @@ -601,7 +604,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net( GGML_ASSERT(ne20 % 32 == 0); snprintf(base, 256, "kernel_gated_delta_net_%s_%d", ggml_type_name(op->src[0]->type), nsg); - snprintf(name, 256, "%s_ne20=%d_ne30=%d_K=%d", base, ne20, ne30, K); + snprintf(name, 256, "%s_ne20=%d_ne30=%d_K=%d_rows=%d_write_rows=%d", base, ne20, ne30, K, has_rows ? 1 : 0, write_rows ? 1 : 0); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); if (!res.pipeline) { @@ -610,6 +613,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net( ggml_metal_cv_set_int16(cv, ne20, FC_GATED_DELTA_NET + 0); ggml_metal_cv_set_int16(cv, ne30, FC_GATED_DELTA_NET + 1); ggml_metal_cv_set_int16(cv, K, FC_GATED_DELTA_NET + 2); + ggml_metal_cv_set_bool (cv, has_rows, FC_GATED_DELTA_NET + 3); + ggml_metal_cv_set_bool (cv, write_rows, FC_GATED_DELTA_NET_WRITE_ROWS); res = ggml_metal_library_compile_pipeline(lib, base, name, cv); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 04a9229b513f..6552fc0c789b 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -127,7 +127,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv 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); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv (ggml_metal_library_t lib, const struct ggml_tensor * op); -struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net (ggml_metal_library_t lib, const struct ggml_tensor * op, bool write_rows); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 89188fef29b1..b7c6c2fa8c3c 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -104,6 +104,7 @@ #define FC_SUM_ROWS 1400 #define FC_UPSCALE 1500 #define FC_GATED_DELTA_NET 1600 +#define FC_GATED_DELTA_NET_WRITE_ROWS (FC_GATED_DELTA_NET + 4) // op-specific constants #define OP_FLASH_ATTN_EXT_NQPSG 8 diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 879e826373a3..3c1eb188db00 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -13,6 +13,7 @@ #include #include #include +#include static ggml_metal_buffer_id ggml_metal_get_buffer_id(const ggml_tensor * t) { if (!t) { @@ -73,6 +74,14 @@ struct ggml_metal_op { return idxs.size(); } + bool is_fused_set_rows(const ggml_tensor * node) const { + return fused_set_rows.find(node) != fused_set_rows.end(); + } + + void mark_fused_set_rows(const ggml_tensor * node) { + fused_set_rows.insert(node); + } + ggml_tensor * node(int i) const { assert(i >= 0 && i < (int) idxs.size()); return ggml_graph_node(gf, idxs[i]); @@ -109,6 +118,7 @@ struct ggml_metal_op { // non-empty node indices std::vector idxs; + std::unordered_set fused_set_rows; }; ggml_metal_op_t ggml_metal_op_init( @@ -182,6 +192,13 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { return 1; } + // A rows scatter may be consumed by the preceding fused GDN epilogue. + // Keep the graph node for dependency construction, but do not encode a + // second copy/scatter kernel. + if (node->op == GGML_OP_SET_ROWS && ctx->is_fused_set_rows(node)) { + return 1; + } + switch (node->op) { case GGML_OP_NONE: case GGML_OP_RESHAPE: @@ -1591,6 +1608,55 @@ int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) { return 1; } +// The rows-mode GDN op produces attention output plus a trailing snapshot +// region. In the recurrent ring graph that region is viewed and later +// scattered back into the state cache by SET_ROWS. Keep the graph nodes (and +// therefore the dependency) but let the GDN epilogue perform that scatter so +// the 786K-element SET_ROWS dispatch disappears from the Metal command stream. +static int ggml_metal_gdn_write_rows( + ggml_metal_op_t ctx, + int idx, + ggml_tensor ** write_rows, + ggml_tensor ** state_dst, + ggml_tensor ** fused_set_rows) { + *write_rows = nullptr; + *state_dst = nullptr; + *fused_set_rows = nullptr; + + const ggml_tensor * gdn = ctx->node(idx); + if (gdn->op != GGML_OP_GATED_DELTA_NET || gdn->src[6] == nullptr || + getenv("GGML_GDN_WRITE_FOLD_DISABLE") != nullptr) { + return 1; + } + + for (int j = idx + 1; j < ctx->n_nodes(); ++j) { + ggml_tensor * set_rows = ctx->node(j); + if (set_rows->op != GGML_OP_SET_ROWS || set_rows->src[0] == nullptr) { + continue; + } + + // SET_ROWS receives a view into the GDN result. Follow the view chain + // because attention normalization and cache maintenance nodes may be + // ordered between the producer and this scatter in the graph. + const ggml_tensor * src = set_rows->src[0]; + while (src != nullptr && (src->op == GGML_OP_VIEW || src->op == GGML_OP_RESHAPE)) { + src = src->src[0]; + } + if (src != gdn || set_rows->src[1] == nullptr || set_rows->src[2] == nullptr || + set_rows->src[1]->type != GGML_TYPE_I64 || set_rows->src[2]->type != GGML_TYPE_F32 || + set_rows->src[2]->buffer == nullptr || set_rows->src[2]->data == nullptr) { + continue; + } + + *write_rows = set_rows->src[1]; + *state_dst = set_rows->src[2]; + *fused_set_rows = set_rows; + return 1; + } + + return 1; +} + int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -1607,7 +1673,22 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { GGML_TENSOR_LOCALS( int32_t, ne, op, ne); GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); - auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op); + ggml_tensor * write_rows = nullptr; + ggml_tensor * state_dst = nullptr; + ggml_tensor * fused_set_rows = nullptr; + const int n_fuse = ggml_metal_gdn_write_rows(ctx, idx, &write_rows, &state_dst, &fused_set_rows); + const bool has_write_rows = write_rows != nullptr; + + if (has_write_rows) { + ctx->mark_fused_set_rows(fused_set_rows); + // The future SET_ROWS is an explicit write dependency. Register its + // destination now and force a barrier before the in-kernel write so + // earlier cache maintenance cannot overlap it. + ggml_metal_op_concurrency_reset(ctx); + ggml_metal_op_concurrency_add(ctx, fused_set_rows); + } + + auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op, has_write_rows); int ida = 0; @@ -1657,13 +1738,20 @@ int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // gate ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // beta ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); // state + // rows (rows mode; bind state as a never-read placeholder otherwise -- + // the function constant compiles the rows path out entirely) + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6] ? op->src[6] : op->src[5]), ida++); + // write rows and destination are only consumed by the fused ring path; + // bind valid placeholders for the ordinary/scratch variants. + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(has_write_rows ? write_rows : (op->src[6] ? op->src[6] : op->src[5])), ida++); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(has_write_rows ? state_dst : op->src[5]), ida++); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst const int nsg = pipeline.nsg; ggml_metal_encoder_dispatch_threadgroups(enc, op->src[2]->ne[0]/nsg, op->src[2]->ne[1], op->src[2]->ne[3], 32, nsg, 1); - return 1; + return n_fuse; } int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) { diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index f5836fca94c1..61959d2a5fb6 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2622,6 +2622,8 @@ kernel void kernel_rwkv_wkv7_f32( constant short FC_gated_delta_net_ne20 [[function_constant(FC_GATED_DELTA_NET + 0)]]; constant short FC_gated_delta_net_ne30 [[function_constant(FC_GATED_DELTA_NET + 1)]]; constant short FC_gated_delta_net_K [[function_constant(FC_GATED_DELTA_NET + 2)]]; +constant bool FC_gated_delta_net_rows [[function_constant(FC_GATED_DELTA_NET + 3)]]; +constant bool FC_gated_delta_net_write_rows [[function_constant(FC_GATED_DELTA_NET + 4)]]; #if 1 template @@ -2633,13 +2635,18 @@ kernel void kernel_gated_delta_net_impl( device const char * g, device const char * b, device const char * s, + device const char * rows, + device const char * write_rows, + device char * state_dst, device char * dst, uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], uint3 ntg[[threads_per_threadgroup]]) { -#define S_v FC_gated_delta_net_ne20 -#define G FC_gated_delta_net_ne30 -#define K FC_gated_delta_net_K +#define S_v FC_gated_delta_net_ne20 +#define G FC_gated_delta_net_ne30 +#define K FC_gated_delta_net_K +#define HAS_ROWS FC_gated_delta_net_rows +#define WRITE_ROWS FC_gated_delta_net_write_rows const uint tx = tpitg.x; const uint ty = tpitg.y; @@ -2653,9 +2660,14 @@ kernel void kernel_gated_delta_net_impl( const float scale = 1.0f / sqrt((float)S_v); - // input state layout (D, K, n_seqs): per-seq stride is K*H*D; we read slot 0. + // input state read base. scratch mode: layout (D, K, n_seqs), per-seq + // stride K*H*D, slot 0. rows mode: s is a 2D cache view with D-wide + // contiguous rows; seq i23's live state is at cache row rows[i23]. // state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous - const uint state_in_base = (i23*K*args.ne21 + i21)*S_v*S_v + i20*S_v; + const uint state_seq_base = HAS_ROWS + ? ((uint)((device const int *) rows)[i23])*(uint)(args.ne21*S_v*S_v) + : (i23*K*args.ne21)*S_v*S_v; + const uint state_in_base = state_seq_base + i21*S_v*S_v + i20*S_v; device const float * s_ptr = (device const float *) (s) + state_in_base; float ls[NSG]; @@ -2736,11 +2748,30 @@ kernel void kernel_gated_delta_net_impl( if (K > 1) { const int target_slot = (int)t - shift; if (target_slot >= 0 && target_slot < (int)K) { + // always populate the op's own snapshot tail: the fold must + // not leave the documented output region uninitialized for + // other consumers (or output/eval callbacks) device float * dst_state = (device float *) (dst) + attn_size + (uint)target_slot * state_size_per_snap + state_out_base; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; dst_state[is] = ls[j]; } + + if (WRITE_ROWS) { + // additionally scatter into the state cache in place of + // the folded SET_ROWS. SET_ROWS receives only the trailing + // n_write snapshots when T < K; convert the absolute + // output slot back to the compact row-index input's + // slot-major coordinate. + const int write_slot = target_slot - max(0, (int)K - (int)args.ne22); + const uint64_t row = ((device const int64_t *) write_rows)[(uint)write_slot * args.ne23 + i23]; + device float * dst_rows = (device float *) state_dst + row * (uint64_t)(S_v * S_v * args.ne21) + + (uint) i21 * S_v * S_v + i20 * S_v; + FOR_UNROLL (short j = 0; j < NSG; j++) { + const short is = tx*NSG + j; + dst_rows[is] = ls[j]; + } + } } } } @@ -2751,11 +2782,24 @@ kernel void kernel_gated_delta_net_impl( const short is = tx*NSG + j; dst_state[is] = ls[j]; } + + if (WRITE_ROWS) { + // single snapshot slot: scatter it to the cache row in place of + // the folded SET_ROWS, same as the K > 1 branch above + const uint64_t row = ((device const int64_t *) write_rows)[i23]; + device float * dst_rows = (device float *) state_dst + row * (uint64_t)(S_v * S_v * args.ne21) + + (uint) i21 * S_v * S_v + i20 * S_v; + FOR_UNROLL (short j = 0; j < NSG; j++) { + const short is = tx*NSG + j; + dst_rows[is] = ls[j]; + } + } } #undef S_v #undef G #undef K +#undef WRITE_ROWS } typedef decltype(kernel_gated_delta_net_impl<4>) kernel_gated_delta_net_t; diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 2a41215fd13d..d21e84fe3d96 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -5135,6 +5135,10 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32); case GGML_OP_GATED_DELTA_NET: { + // rows-indexed state read (src[6]) not implemented here + if (op->src[6] != NULL) { + return false; + } // Match the Vulkan backend: only F32 -> F32, S_v in {16, 32, 64, 128}. if (op->src[0]->type != GGML_TYPE_F32 || op->type != GGML_TYPE_F32) { return false; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 3f246e8672d5..9e9242ba3b1c 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -5515,8 +5515,10 @@ static bool ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_RWKV_WKV6: case GGML_OP_RWKV_WKV7: case GGML_OP_GATED_LINEAR_ATTN: - case GGML_OP_GATED_DELTA_NET: return true; + case GGML_OP_GATED_DELTA_NET: + // rows-indexed state read (src[6]) not implemented here + return op->src[6] == NULL; case GGML_OP_SSM_CONV: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c1fc03b57cde..fd22aba117eb 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -16890,6 +16890,10 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm return true; // all inputs are contiguous, see ggml.c case GGML_OP_GATED_DELTA_NET: { + // rows-indexed state read (src[6]) not implemented on Vulkan yet + if (op->src[6] != nullptr) { + return false; + } const uint32_t S_v = op->src[2]->ne[0]; if (S_v != 32 && S_v != 64 && S_v != 128) { return false; diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp index c6cfb0bbbadc..f77e5a933f4a 100644 --- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -4329,6 +4329,10 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const break; case GGML_OP_GATED_DELTA_NET: { + if (op->src[6] != nullptr) { + supports_op = false; // rows-indexed state read not implemented here + break; + } const uint32_t s_v = (uint32_t) src2->ne[0]; supports_op = op->type == GGML_TYPE_F32 && src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 && src2->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 && diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 1de9882792b2..de0615fb116f 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6236,6 +6236,69 @@ struct ggml_tensor * ggml_gated_delta_net( result->src[4] = beta; result->src[5] = state; + // K for the output snapshot slots; kept in op_params so both op variants + // (scratch-state and rows-indexed) share one code path in the backends + ggml_set_op_params_i32(result, 0, (int32_t) K); + + return result; +} + +struct ggml_tensor * ggml_gated_delta_net_rows( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * states, + struct ggml_tensor * rows, + int n_snap_slots) { + GGML_ASSERT(ggml_is_contiguous_rows(q)); + GGML_ASSERT(ggml_is_contiguous_rows(k)); + GGML_ASSERT(ggml_is_contiguous_rows(v)); + GGML_ASSERT(ggml_is_contiguous(g)); + GGML_ASSERT(ggml_is_contiguous(beta)); + GGML_ASSERT(ggml_is_contiguous(states)); + GGML_ASSERT(ggml_is_contiguous(rows)); + + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(k->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32); + GGML_ASSERT(g->type == GGML_TYPE_F32); + GGML_ASSERT(beta->type == GGML_TYPE_F32); + GGML_ASSERT(states->type == GGML_TYPE_F32); + GGML_ASSERT(rows->type == GGML_TYPE_I32); + + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + + GGML_ASSERT(g->ne[0] == 1 || g->ne[0] == S_v); + GGML_ASSERT(beta->ne[0] == 1); + + // states is a 2D cache view (D, n_rows); each row is one sequence's state + GGML_ASSERT(states->ne[0] == S_v * S_v * H); + GGML_ASSERT(rows->ne[0] == n_seqs); + + const int64_t K = n_snap_slots; + GGML_ASSERT(K >= 1); + + const int64_t state_rows = K * S_v * n_seqs; + const int64_t ne[4] = { S_v * H, n_tokens * n_seqs + state_rows, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_GATED_DELTA_NET; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = g; + result->src[4] = beta; + result->src[5] = states; + result->src[6] = rows; + + ggml_set_op_params_i32(result, 0, (int32_t) K); + return result; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index b374ace4fc2d..a8eb28d9eecc 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2885,6 +2885,33 @@ ggml_tensor * llm_graph_context::build_rs( get_state_rows); } +ggml_tensor * llm_graph_context::build_rs_cache_view( + llm_graph_input_rs * inp, + ggml_tensor * s, + int32_t state_size, + int32_t n_seqs) const { + const auto * kv_state = inp->mctx; + + const uint32_t n_rs = kv_state->get_n_rs(); + const uint32_t rs_head = kv_state->get_head(); + const int32_t rs_zero = kv_state->get_rs_z(); + + ggml_tensor * states = ggml_reshape_2d(ctx0, s, state_size, s->ne[1]); + + // same cache hygiene as build_rs, minus the main gather (the consumer reads + // per-seq rows via inp->s_copy_main directly) + ggml_tensor * state_zero = ggml_view_1d(ctx0, states, state_size*(rs_zero >= 0), rs_zero*states->nb[1]*(rs_zero >= 0)); + ggml_build_forward_expand(gf, ggml_scale_inplace(ctx0, state_zero, 0)); + + ggml_tensor * states_extra = ggml_get_rows(ctx0, states, inp->s_copy_extra); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, + states_extra, + ggml_view_2d(ctx0, s, state_size, (n_rs - n_seqs), s->nb[1], (rs_head + n_seqs)*s->nb[1]))); + + return states; +} + ggml_tensor * llm_graph_context::build_rwkv_token_shift_load( llm_graph_input_rs * inp, const llama_ubatch & ubatch, diff --git a/src/llama-graph.h b/src/llama-graph.h index fe16e34f008f..ac8b9e2edee5 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1174,6 +1174,17 @@ struct llm_graph_context { int32_t n_seqs, const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const; + // like build_rs, but WITHOUT the main per-seq state gather: performs the + // rs_zero clear and the extra-states relocation, then returns the 2D + // (state_size, n_rs_total) cache view. For consumers that read per-seq + // state rows directly via inp->s_copy_main (e.g. ggml_gated_delta_net_rows), + // saving a get_rows + a downstream slot-0 cpy per layer per decode. + ggml_tensor * build_rs_cache_view( + llm_graph_input_rs * inp, + ggml_tensor * s, + int32_t state_size, + int32_t n_seqs) const; + ggml_tensor * build_rwkv_token_shift_load( llm_graph_input_rs * inp, const llama_ubatch & ubatch, diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index 01876e3010db..e65f55b212d1 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -545,17 +545,22 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( ggml_tensor * g, ggml_tensor * b, ggml_tensor * s, - int il) { + int il, + ggml_tensor * state_rows) { const auto * mctx_cur = inp->mctx; const auto kv_head = mctx_cur->get_head(); - const int64_t S_v = s->ne[0]; - const int64_t H_v = s->ne[2]; - const int64_t n_seqs = s->ne[3]; + // dims from v (always (S_v, H_v, T, B)): in rows mode `s` is the 2D cache + // view, so its shape no longer carries them + const int64_t S_v = v->ne[0]; + const int64_t H_v = v->ne[1]; + const int64_t n_seqs = v->ne[3]; const int64_t n_seq_tokens = q->ne[2]; const bool keep = cparams.n_rs_seq > 0; + GGML_ASSERT(state_rows == nullptr || keep); // rows mode is a ring-path optimization + if (!keep) { auto attn_out = build_delta_net(q, k, v, g, b, s, il); ggml_tensor * output = attn_out.first; @@ -578,12 +583,23 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( // snapshot slot 0 of each sequence (all backends -- see ggml_gated_delta_net); // slots 1..K-1 exist only to size the K-slot output. copy the current state // into slot 0 and leave the rest uninitialized instead of zero-padding, - // which would write D*K elements per layer on every decode - ggml_tensor * s_in = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, K, n_seqs); - ggml_tensor * s_in0 = ggml_view_3d(ctx0, s_in, D, 1, n_seqs, s_in->nb[1], s_in->nb[2], 0); - ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_reshape_3d(ctx0, s, D, 1, n_seqs), s_in0)); + // which would write D*K elements per layer on every decode. + // Keep a private scratch tensor per recurrent layer. Reusing one scratch + // across layers creates overlapping live ranges in the Metal scheduler; + // that splits the ring-enabled graph at every recurrent boundary. The + // extra memory is preferable to serializing all 48 GDN layers. + ggml_tensor * gdn_out; + if (state_rows) { + // rows mode: the fused op reads each seq's live state directly from the + // cache view at row state_rows[seq] -- no gather, no slot-0 cpy + gdn_out = ggml_gated_delta_net_rows(ctx0, q, k, v, g, b, s, state_rows, (int) K); + } else { + ggml_tensor * s_in = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, D, K, n_seqs); + ggml_tensor * s_in0 = ggml_view_3d(ctx0, s_in, D, 1, n_seqs, s_in->nb[1], s_in->nb[2], 0); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_reshape_3d(ctx0, s, D, 1, n_seqs), s_in0)); - ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s_in); + gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s_in); + } if (n_seq_tokens > 1) { cb(gdn_out, LLAMA_TENSOR_NAME_FGDN_CH, il); } else { diff --git a/src/models/models.h b/src/models/models.h index 1823c1d52b6f..b26a39c2f7bb 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -78,6 +78,11 @@ struct llm_build_delta_net_base : public llm_graph_context { // run delta-net attention and write the new recurrent state(s) back to ssm_states_all // s: (head_v_dim, head_v_dim, num_v_heads, n_seqs); returns output: (head_v_dim, num_v_heads, n_seq_tokens, n_seqs) + // + // state_rows (optional, ring path only): when set, `s` is instead the 2D + // cache view from build_rs_cache_view and the fused op reads each seq's + // live state directly at cache row state_rows[seq] (inp->s_copy_main) -- + // no gathered scratch, no slot-0 cpy. ggml_tensor * build_recurrent_attn( llm_graph_input_rs * inp, ggml_tensor * ssm_states_all, @@ -87,7 +92,8 @@ struct llm_build_delta_net_base : public llm_graph_context { ggml_tensor * g, ggml_tensor * b, ggml_tensor * s, - int il); + int il, + ggml_tensor * state_rows = nullptr); }; struct llm_build_rwkv6_base : public llm_graph_context { diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index bf34c8a8cfb5..37dfc51bc8d5 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -425,9 +425,39 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, conv_kernel_size, conv_channels, il); - ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); - state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); - cb(state, "state_predelta", il); + // ring path: read per-seq live state directly from the cache inside the + // fused GDN op (rows mode) instead of gather + slot-0 cpy per layer. + // GGML_GDN_STATE_GATHER=1 restores the legacy gathered path (A/B). + // rows mode (the src[6] variant) is implemented on CPU and Metal only; + // other GPU backends reject it in supports_op, which would silently move + // the whole recurrent op to CPU -- keep the gathered form unless every + // GPU device in the model is Metal. + static const bool gdn_state_rows_env = getenv("GGML_GDN_STATE_GATHER") == nullptr; + + bool gdn_state_rows_dev_ok = true; + for (const auto & ldev : model.devices) { + if (ldev.dev == nullptr || ggml_backend_dev_type(ldev.dev) != GGML_BACKEND_DEVICE_TYPE_GPU) { + continue; + } + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ldev.dev); + const char * reg_name = reg ? ggml_backend_reg_name(reg) : nullptr; + if (reg_name == nullptr || strcmp(reg_name, "Metal") != 0) { + gdn_state_rows_dev_ok = false; + break; + } + } + + const bool gdn_state_rows = gdn_state_rows_env && gdn_state_rows_dev_ok && cparams.n_rs_seq > 0; + + ggml_tensor * state; + if (gdn_state_rows) { + state = build_rs_cache_view(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + cb(state, "state_cache_view", il); + } else { + state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); + cb(state, "state_predelta", il); + } ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); cb(conv_output_proper, "conv_output_raw", il); @@ -485,7 +515,8 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( cb(k_conv, "k_conv_predelta", il); cb(v_conv, "v_conv_predelta", il); - ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); + ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il, + gdn_state_rows ? inp->s_copy_main : nullptr); // z: [head_dim, n_heads, n_tokens, n_seqs] -> [n_heads * n_tokens * n_seqs, head_dim] ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index be1978de81a1..18b87d63c26f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3865,16 +3865,17 @@ struct test_gated_delta_net : public test_case { const bool permuted; const bool kda; const int64_t K; // snapshot slot count: 1 = final-only, >1 = last K states + const bool rows_mode; // rows-indexed state read from a 2D cache view (src[6]) std::string vars() override { - return VARS_TO_STR9(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K); + return VARS_TO_STR10(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K, rows_mode); } test_gated_delta_net(ggml_type type = GGML_TYPE_F32, int64_t head_count = 4, int64_t head_size = 16, int64_t n_seq_tokens = 1, int64_t n_seqs = 1, - int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1) + int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1, bool rows_mode = false) : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs), - v_repeat(v_repeat), permuted(permuted), kda(kda), K(K) {} + v_repeat(v_repeat), permuted(permuted), kda(kda), K(K), rows_mode(rows_mode) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q; @@ -3896,14 +3897,26 @@ struct test_gated_delta_net : public test_case { const int64_t g_ne0 = kda ? head_size : 1; ggml_tensor * g = ggml_new_tensor_4d(ctx, type, g_ne0, head_count * v_repeat, n_seq_tokens, n_seqs); ggml_tensor * beta = ggml_new_tensor_4d(ctx, type, 1, head_count * v_repeat, n_seq_tokens, n_seqs); - ggml_tensor * state = ggml_new_tensor_3d(ctx, type, head_size * v_repeat * head_size * head_count, K, n_seqs); ggml_set_name(g, "g"); ggml_set_name(beta, "beta"); - ggml_set_name(state, "state"); // q/k are L2-normalised in qwen35/kimi-linear before delta_net q = ggml_l2_norm(ctx, q, 1e-6f); k = ggml_l2_norm(ctx, k, 1e-6f); - ggml_tensor * out = ggml_gated_delta_net(ctx, q, k, v, g, beta, state); + ggml_tensor * out; + if (rows_mode) { + // 2D cache view with more rows than sequences; per-seq state rows + // are picked via the I32 rows tensor (see initialize_tensors) + const int64_t D = head_size * v_repeat * head_size * head_count; + ggml_tensor * states = ggml_new_tensor_2d(ctx, type, D, n_seqs + 3); + ggml_tensor * rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs); + ggml_set_name(states, "state"); + ggml_set_name(rows, "rows"); + out = ggml_gated_delta_net_rows(ctx, q, k, v, g, beta, states, rows, K); + } else { + ggml_tensor * state = ggml_new_tensor_3d(ctx, type, head_size * v_repeat * head_size * head_count, K, n_seqs); + ggml_set_name(state, "state"); + out = ggml_gated_delta_net(ctx, q, k, v, g, beta, state); + } return out; } @@ -3916,6 +3929,13 @@ struct test_gated_delta_net : public test_case { init_tensor_uniform(t, 0.0f, 1.0f); } else if (strcmp(t->name, "v") == 0) { init_tensor_uniform(t, -0.3f, 5.0f); + } else if (strcmp(t->name, "rows") == 0) { + // deterministic, distinct, in-range cache rows (stride 2 over n_seqs+3) + std::vector idx(t->ne[0]); + for (int64_t i = 0; i < t->ne[0]; i++) { + idx[i] = (int32_t) ((i*2 + 1) % (t->ne[0] + 3)); + } + ggml_backend_tensor_set(t, idx.data(), 0, idx.size()*sizeof(int32_t)); } else { init_tensor_uniform(t); } @@ -9151,6 +9171,13 @@ static std::vector> make_test_cases_eval() { // overflow: n_tokens > K — only the last K snapshots kept. test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, /*K=*/3)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 16, 2, 1, false, false, /*K=*/4)); + // rows mode: state read directly from a 2D cache view at rows[seq] (src[6]) + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 1, 1, 1, false, false, /*K=*/2, /*rows=*/true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, false, /*K=*/4, /*rows=*/true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 4, 1, 1, false, false, /*K=*/4, /*rows=*/true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, /*K=*/3, /*rows=*/true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 16, 2, 1, false, false, /*K=*/4, /*rows=*/true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, true, /*K=*/4, /*rows=*/true)); #if 0 // these tests are disabled to save execution time, sbut they can be handy for debugging