diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d875a15ed066..7d518edf64c9 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2620,6 +2620,31 @@ extern "C" { struct ggml_tensor * rows, int n_snap_slots); + // cache-writing form (CPU backend only): like ggml_gated_delta_net_rows, but + // the op MUTATES `states` instead of returning state planes -- the result is + // the attention scores only, [S_v*H_v, n_tokens*n_seqs]. For sequence i the + // kernel reads cache row rows_in[i] (I32, n_seqs entries), updates the live + // state in place in row rows_out[0*n_seqs + i] (I64, min(n_tokens, K)*n_seqs + // entries; no copy at all when rows_in[i] == rows_out[i]) and writes snapshot + // slot s (state s tokens back) for s in 1..min(n_tokens, K)-1 straight into + // row rows_out[s*n_seqs + i]. The caller guarantees that rows_in[i] and every + // rows_out[s*n_seqs + i] are planes of ONE cell (row % mem_size equal, checked + // by the kernel) and that the cells of different sequences are pairwise + // distinct, so no work unit ever reads a row another unit writes. Removes the + // gather, the dst state planes and the write-back copy from recurrent decode. + GGML_API struct ggml_tensor * ggml_gated_delta_net_cache( + 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, // 2D cache view [S_v*S_v*H_v, n_rows], mutated + struct ggml_tensor * rows_in, // I32 [n_seqs] + struct ggml_tensor * rows_out, // I64 [min(n_tokens, K)*n_seqs] + int n_snap_slots, + int mem_size); + // fold the per-head gate activations into a gated_delta_net op (scalar gate only): // beta -> sigmoid(beta), g -> a[h] * softplus(g + dt_bias[h]) // dt_bias and a are F32 with H_v elements; removes four elementwise ops per layer diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 9cd2c2d39a5a..3a6e21591542 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -3011,7 +3011,9 @@ struct ggml_cplan ggml_graph_plan( { const int64_t S_v = node->src[2]->ne[0]; const int64_t K = ggml_get_op_params_i32(node, 0); - const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0); + // cache-write mode (op_params[2]) computes on the cache row: no state scratch + const bool cache_write = ggml_get_op_params_i32(node, 2) != 0; + const int64_t per_thread = S_v + ((K > 1 && !cache_write) ? S_v * S_v : 0); cur = per_thread * sizeof(float) * n_tasks; } break; case GGML_OP_COUNT: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 4fc7514bb3fb..83a616e101a6 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -10819,11 +10819,26 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( const int64_t state_seq_stride = src_rows ? 0 : (int64_t) (src_state->nb[3] / 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); + // cache-write mode (op_params[2], ggml_gated_delta_net_cache): the live state is + // updated in place in cache row rows_out[seq] and snapshot slots 1..K-1 are + // written straight into cache rows rows_out[slot*n_seqs + seq]; the result holds + // the attention scores only (state_out_base is unused in this mode). + const bool cache_write = ggml_get_op_params_i32(dst, 2) != 0; + GGML_ASSERT(!cache_write || state_rows_idx != nullptr); + const int64_t * rows_out = cache_write ? (const int64_t *) dst->src[9]->data : nullptr; + float * cache = cache_write ? (float *) src_state->data : nullptr; + const int64_t mem_size = cache_write ? ggml_get_op_params_i32(dst, 3) : 1; + + // the legacy K > 1 path stages the state through a per-thread scratch; the + // cache-write path never does (it computes on the cache row itself). + // ggml-cpu.c (graph plan work size) mirrors this expression. + const bool use_state_work = K > 1 && !cache_write; + + const int64_t per_thread = S_v + (use_state_work ? S_v * S_v : 0); const int ith = params->ith; float * delta = (float *)params->wdata + ith * per_thread + CACHE_LINE_SIZE_F32; - float * state_work = K > 1 ? (delta + S_v) : nullptr; + float * state_work = use_state_work ? (delta + S_v) : nullptr; // output layout: [attn_scores | new_states] // attn_scores: S_v * H * n_tokens * n_seqs floats @@ -10855,19 +10870,34 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( const int64_t iq3 = iv3 / rq3; const int64_t ik3 = iv3 / rk3; - // For K=1, write directly to the single output slot to avoid an extra memcpy at the end. - // For K>1, work in scratch and copy out per-token when the slot is in range. - float * s_out = (K > 1) - ? state_work - : state_out_base + (iv3 * H + iv1) * S_v * S_v; - - // copy input state into the working buffer and operate in-place. - // scratch mode: state layout [S_v, S_v, H, n_seqs], seq iv3 starts at - // iv3 * state_seq_stride. rows mode: cache row state_rows_idx[iv3]. + // input state: scratch mode: state layout [S_v, S_v, H, n_seqs], seq iv3 + // starts 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)); + + float * s_out; + if (cache_write) { + // operate directly on the live cache row rows_out[iv3]. rows_in[iv3] is + // either that same row (steady state: no copy at all) or another plane + // of the same cell (pending rollback / freshly zeroed cell): copy it in + // first -- distinct rows, so plain memcpy is safe. + const int64_t row_out = rows_out[iv3]; + GGML_ASSERT(row_out % mem_size == (int64_t) state_rows_idx[iv3] % mem_size); + s_out = cache + row_out * state_row_size + iv1 * S_v * S_v; + if (s_out != s_in) { + memcpy(s_out, s_in, S_v * S_v * sizeof(float)); + } + } else { + // For K=1, write directly to the single output slot to avoid an extra memcpy at the end. + // For K>1, work in scratch and copy out per-token when the slot is in range. + s_out = (K > 1) + ? state_work + : state_out_base + (iv3 * H + iv1) * S_v * S_v; + + // copy input state into the working buffer and operate in-place. + memcpy(s_out, s_in, S_v * S_v * sizeof(float)); + } // attn output pointer for first token of this (head, seq) float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v; @@ -10924,7 +10954,20 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( attn_data += S_v * H; // advance to next token - if (K > 1) { + if (cache_write) { + // snapshot slot s (state s tokens back) -> cache row rows_out[s*n_seqs + seq]. + // slot 0 is s_out itself (already in place). slot <= n_tokens-1 and + // slot < K keep the index below min(n_tokens, K)*n_seqs. Within this + // unit the rollback source plane was fully copied into s_out before + // the token loop, so overwriting it here is read-before-write safe. + if (K > 1) { + const int64_t slot = n_tokens - 1 - t; + if (slot >= 1 && slot < K) { + float * snap = cache + rows_out[slot * n_seqs + iv3] * state_row_size + iv1 * S_v * S_v; + memcpy(snap, s_out, S_v * S_v * sizeof(float)); + } + } + } else if (K > 1) { const int64_t target_slot = n_tokens - 1 - t; if (target_slot >= 0 && target_slot < K) { float * curr_state_o = state_out_base + target_slot * state_size_per_snap + diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index fa2a9c44f5cf..c233176541b3 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6423,6 +6423,75 @@ struct ggml_tensor * ggml_gated_delta_net_rows( return result; } +struct ggml_tensor * ggml_gated_delta_net_cache( + 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_in, + struct ggml_tensor * rows_out, + int n_snap_slots, + int mem_size) { + 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_in)); + GGML_ASSERT(ggml_is_contiguous(rows_out)); + + 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_in->type == GGML_TYPE_I32); + GGML_ASSERT(rows_out->type == GGML_TYPE_I64); + + 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 full state + GGML_ASSERT(states->ne[0] == S_v * S_v * H); + GGML_ASSERT(rows_in->ne[0] == n_seqs); + + const int64_t K = n_snap_slots; + GGML_ASSERT(K >= 1); + GGML_ASSERT(mem_size >= 1); + GGML_ASSERT(rows_out->ne[0] == MIN(n_tokens, K) * n_seqs); + + // the state stays in the cache: the result carries the attention scores only + const int64_t ne[4] = { S_v * H, n_tokens * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_i32(result, 0, (int32_t) K); + ggml_set_op_params_i32(result, 2, 1); // cache-write flag (op_params[1] is the raw-gates bit) + ggml_set_op_params_i32(result, 3, mem_size); // kernel same-cell check: row % mem_size + + 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_in; + // src[7] and src[8] are reserved for ggml_gated_delta_net_set_raw_gates + result->src[9] = rows_out; + + return result; +} + void ggml_gated_delta_net_set_raw_gates( struct ggml_tensor * gdn, struct ggml_tensor * dt_bias, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 99f263e32bd7..002c599eddb2 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -417,6 +417,11 @@ bool llm_graph_input_rs::can_reuse_rs(const llama_memory_recurrent_context * mct res &= head == mctx_cur->get_head(); res &= rs_z == mctx_cur->get_rs_z(); + // a graph built for the in-place GDN path must not be reused by a ubatch + // whose cells were reordered / freshly zeroed (and vice versa); the memory + // was already updated by apply() -> find_slot before can_reuse runs + res &= rs_inplace == mctx_cur->rs_inplace_ok(params.ubatch.n_seqs); + return res; } @@ -3573,6 +3578,8 @@ static std::unique_ptr build_rs_inp_impl( inp->head = mctx_cur->get_head(); inp->rs_z = mctx_cur->get_rs_z(); + inp->rs_inplace = mctx_cur->rs_inplace_ok(n_seqs); + return inp; } @@ -3601,28 +3608,38 @@ 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 { + int32_t n_seqs, + bool inplace) 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(); + // the in-place consumer mutates the cache rows of the ubatch: it is only + // entered when every main row is the identity and the extra range below is + // empty (rs_inplace_ok), which makes the ordering hazard described next + // unreachable for it + GGML_ASSERT(!inplace || (inp->rs_inplace && n_rs == (uint32_t) n_seqs)); + 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, inside the GDN op). // - // KNOWN LIMITATION (tracked follow-up): build_rs gathers the main rows - // BEFORE this extra relocation, so an overlapping main row is read before - // being overwritten. rows mode defers the main read into the consumer, and - // s_copy() maps a main row to an arbitrary cache slot (idx*size + src0), - // which can fall inside the extra destination [rs_head+n_seqs, rs_head+n_rs) - // during a cache reorder -- so this relocation could clobber a main row the - // consumer will later read. Not reachable on the current single-sequence - // decode path, but it is a real multi-sequence hazard; the correct fix is - // to order the relocation AFTER the GDN read (build_rs's read-before-write - // ordering), which is a graph-dependency refactor left as follow-up. + // KNOWN LIMITATION of the Metal rows mode (tracked follow-up): build_rs + // gathers the main rows BEFORE this extra relocation, so an overlapping main + // row is read before being overwritten. rows mode defers the main read into + // the consumer, and s_copy() maps a main row to an arbitrary cache slot + // (idx*size + src0), which can fall inside the extra destination + // [rs_head+n_seqs, rs_head+n_rs) during a cache reorder -- so this relocation + // could clobber a main row the consumer will later read. The CPU in-place + // mode never enters this function with a non-empty extra range (see the + // assert above; transition ubatches take the gathered build_rs path), so the + // hazard only remains for the Metal rows + SET_ROWS mode on a cache reorder; + // the correct fix there is to order the relocation AFTER the GDN read + // (build_rs's read-before-write ordering), a graph-dependency refactor left + // as follow-up. 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)); diff --git a/src/llama-graph.h b/src/llama-graph.h index e85c1aa20d1f..c3b7e13af6ba 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -346,6 +346,10 @@ class llm_graph_input_rs : public llm_graph_input_i { ggml_tensor * s_write_rows = nullptr; int64_t s_write_K = 0; + // every ubatch seq reads its own cell and n_rs == n_seqs: the GDN op may + // update the cache rows in place (rs_inplace_ok); part of the reuse key + bool rs_inplace = false; + const llama_memory_recurrent_context * mctx; // used in view offsets, need to match for valid graph reuse @@ -1437,11 +1441,14 @@ struct llm_graph_context { // (state_size, n_rows) 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 per layer per decode. + // inplace: the consumer is the cache-writing op (ggml_gated_delta_net_cache); + // requires inp->rs_inplace, i.e. an empty extra range and identity main rows. ggml_tensor * build_rs_cache_view( llm_graph_input_rs * inp, ggml_tensor * s, int32_t state_size, - int32_t n_seqs) const; + int32_t n_seqs, + bool inplace = false) const; // I64 cache-row indices for the rows-mode snapshot scatter: row of // (slot s, seq i) = s*mem_size + head + i, matching the strided-cpy diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7e..dc70452295e8 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1,5 +1,6 @@ #include "llama-memory-recurrent.h" +#include "ggml-alloc.h" #include "ggml-backend.h" #include "llama-impl.h" #include "llama-io.h" @@ -8,15 +9,66 @@ #include #include +#include +#include #include #include #include #include +#ifdef __linux__ +#include +#endif + // // llama_memory_recurrent // +#ifdef __linux__ +// One anonymous, private mapping for every tensor of ctx, placed exactly as +// ggml_backend_alloc_ctx_tensors_from_buft would place them (same padded sizes, +// same alignment, same order) but wrapped with ggml_backend_cpu_buffer_from_ptr, +// which does not own the memory: the caller munmaps it after the buffer is +// freed. Returns nullptr (and touches nothing) when the mapping cannot be made. +static ggml_backend_buffer_t rs_alloc_lazy_zero_host_buffer( + ggml_context * ctx, ggml_backend_buffer_type_t buft, void ** ptr_out, size_t * size_out) { + const size_t size = ggml_backend_alloc_ctx_tensors_from_buft_size(ctx, buft); + if (size == 0) { + return nullptr; + } + + void * ptr = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (ptr == MAP_FAILED) { + LLAMA_LOG_WARN("%s: anonymous mmap of %zu bytes failed (%s), falling back to the eager clear\n", + __func__, size, strerror(errno)); + return nullptr; + } + + // mmap returns page-aligned memory, which satisfies the CPU buffer alignment + ggml_backend_buffer_t buf = ggml_backend_cpu_buffer_from_ptr(ptr, size); + if (!buf) { + munmap(ptr, size); + return nullptr; + } + + ggml_tallocr talloc = ggml_tallocr_new(buf); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { + if (t->data == nullptr && t->view_src == nullptr) { + if (ggml_tallocr_alloc(&talloc, t) != GGML_STATUS_SUCCESS) { + ggml_backend_buffer_free(buf); + munmap(ptr, size); + return nullptr; + } + } + } + + *ptr_out = ptr; + *size_out = size; + + return buf; +} +#endif + llama_memory_recurrent::llama_memory_recurrent( const llama_model & model, ggml_type type_r, @@ -106,13 +158,62 @@ llama_memory_recurrent::llama_memory_recurrent( } // allocate tensors and initialize the buffers to avoid NaNs in the padding + // + // Host (CPU) buffers, Linux: by default the whole RS buffer is one anonymous + // private mmap (rs_alloc_lazy_zero_host_buffer) and is NOT memset. The + // kernel hands out zero pages on first touch, so the "no NaN padding" + // guarantee of the eager clear holds without making every cell of every + // snapshot plane resident at construction: the buffer is + // mem_size*(1+n_rs_seq) rows per layer, so --parallel 16 with n_rs_seq = 2 + // pinned ~4.7 GiB of the 15 GiB here (evicting the mmap'd weights that + // decode streams on every token) while one active slot needs ~0.3 GiB. + // Residency now follows the cells that were actually written. + // + // Why skipping the memset is exact -- every row the graph ever reads is + // zero by construction or was written before it is read: + // - a fresh sequence reads the rs_z row, which the graph zeroes in place + // (ggml_scale_inplace by 0 in build_rs / build_rs_cache_view) before any + // consumer reads it; 0 * (a kernel zero page) == 0 either way; + // - every other main row is src0 = src of a used cell (find_slot), i.e. a + // row written by an earlier step's write-back / in-place update, by the + // extra-row relocation, or by state_read_data; + // - a rollback read of snapshot plane r (1 <= r <= n_rs_seq) only happens + // after a ubatch of that sequence wrote slot r of the same cell: seq_rm + // only accepts rollback <= n_rs_seq, split_equal keeps the trailing + // n_rs_seq + 1 tokens of a sequence in one ubatch and the op writes slots + // 0..min(T, K)-1, so the plane was written before it is read; + // - the alignment padding between tensors is never read. + // clear(true) drops the pages with madvise(MADV_DONTNEED) instead of a + // memset (a private anonymous mapping then reads back as zero). + // LLAMA_RS_EAGER_ZERO=1 restores the eager alloc + memset path. +#ifdef __linux__ + static const bool rs_eager_zero = getenv("LLAMA_RS_EAGER_ZERO") != nullptr; +#endif + for (auto & [buft, ctx] : ctx_map) { - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + ggml_backend_buffer_t buf = nullptr; + bool lazy = false; + +#ifdef __linux__ + if (!rs_eager_zero && buft == ggml_backend_cpu_buffer_type()) { + void * ptr = nullptr; + size_t size = 0; + buf = rs_alloc_lazy_zero_host_buffer(ctx.get(), buft, &ptr, &size); + if (buf) { + mmaps.push_back({ buf, ptr, size }); + lazy = true; + } + } +#endif if (!buf) { - throw std::runtime_error("failed to allocate buffer for rs cache"); + buf = ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), buft); + if (!buf) { + throw std::runtime_error("failed to allocate buffer for rs cache"); + } + ggml_backend_buffer_clear(buf, 0); } - ggml_backend_buffer_clear(buf, 0); - LLAMA_LOG_INFO("%s: %10s RS buffer size = %8.2f MiB\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0); + LLAMA_LOG_INFO("%s: %10s RS buffer size = %8.2f MiB%s\n", __func__, ggml_backend_buffer_name(buf), ggml_backend_buffer_get_size(buf)/1024.0/1024.0, + lazy ? " (anonymous mmap, lazily zeroed)" : ""); ctxs_bufs.emplace_back(std::move(ctx), buf); } @@ -127,6 +228,17 @@ llama_memory_recurrent::llama_memory_recurrent( } } +llama_memory_recurrent::~llama_memory_recurrent() { + // the from_ptr buffers do not own their memory: free them first, then unmap + ctxs_bufs.clear(); +#ifdef __linux__ + for (const auto & m : mmaps) { + munmap(m.ptr, m.size); + } +#endif + mmaps.clear(); +} + void llama_memory_recurrent::clear(bool data) { for (int32_t i = 0; i < (int32_t) size; ++i) { cells[i].pos = -1; @@ -140,7 +252,20 @@ void llama_memory_recurrent::clear(bool data) { if (data) { for (auto & [_, buf] : ctxs_bufs) { - ggml_backend_buffer_clear(buf.get(), 0); + bool dropped = false; +#ifdef __linux__ + for (const auto & m : mmaps) { + if (m.buf == buf.get()) { + // drop the pages instead of memset: a private anonymous + // mapping reads back as zero and stays non-resident + dropped = madvise(m.ptr, m.size, MADV_DONTNEED) == 0; + break; + } + } +#endif + if (!dropped) { + ggml_backend_buffer_clear(buf.get(), 0); + } } } @@ -1270,3 +1395,30 @@ int32_t llama_memory_recurrent_context::s_copy(int i) const { } return (int32_t)(idx * mem->size) + src0; } + +bool llama_memory_recurrent_context::rs_inplace_ok(uint32_t n_seqs) const { + // the reserve context spans the whole cache (head = 0, n_rs = size) + if (is_full) { + return false; + } + + // the extra relocation range [head + n_seqs, head + n_rs) must be empty + if (mem->n != n_seqs) { + return false; + } + + // every ubatch cell must read its own state. After find_slot this holds for + // every cell that already owned its state (src0 = src = own index) and for + // the one-fresh-seq case where rs_z is the seq's own cell (the graph zeroes + // that row before the op reads it). It fails on a seq_cp fan-out, on a cell + // reorder and when two fresh seqs share rs_z -- those ubatches take the + // gathered path. + for (uint32_t i = 0; i < n_seqs; ++i) { + const uint32_t cell_idx = mem->head + i; + if (mem->cells[cell_idx].src0 != (int32_t) cell_idx) { + return false; + } + } + + return true; +} diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..78aa31dccc4c 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -26,7 +26,7 @@ class llama_memory_recurrent : public llama_memory_i { uint32_t n_rs_seq, const layer_filter_cb & filter); - ~llama_memory_recurrent() = default; + ~llama_memory_recurrent(); // // llama_memory_i @@ -121,6 +121,18 @@ class llama_memory_recurrent : public llama_memory_i { // ggml contexts for the KV cache along with the allocated backend buffers: std::vector> ctxs_bufs; + // host RS buffers backed by an anonymous private mmap (kernel zero pages, + // resident only where written) and wrapped with + // ggml_backend_cpu_buffer_from_ptr, which does not own the memory: unmapped + // in the destructor after the buffers are freed. Empty when + // LLAMA_RS_EAGER_ZERO=1 or on non-Linux hosts (eager alloc + memset). + struct rs_mmap_region { + ggml_backend_buffer_t buf; + void * ptr; + size_t size; + }; + std::vector mmaps; + size_t total_size() const; size_t size_r_bytes() const; @@ -173,6 +185,14 @@ class llama_memory_recurrent_context : public llama_memory_context_i { int32_t s_copy(int i) const; + // true when the GDN op may update this ubatch's cache rows in place: not the + // full (reserve) context, the cell range is exactly the ubatch (n == n_seqs) + // and every cell reads its own state (src0 == own index; a pending rollback + // still reads a plane of the same cell). Side-effect free on purpose: unlike + // s_copy() it never consumes the per-seq rollback index, so it can be + // evaluated at graph build and again in can_reuse. + bool rs_inplace_ok(uint32_t n_seqs) const; + private: const llama_memory_status status; diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ae67400c6131..26799679f6af 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -539,7 +539,8 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( ggml_tensor * b, ggml_tensor * s, int il, - ggml_tensor * state_rows) { + ggml_tensor * state_rows, + bool state_inplace) { const auto * mctx_cur = inp->mctx; const auto kv_head = mctx_cur->get_head(); const uint32_t mem_size = mctx_cur->get_size(); @@ -553,7 +554,49 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( const bool keep = cparams.n_rs_seq > 0; - GGML_ASSERT(state_rows == nullptr || keep); // rows mode is a ring-path optimization + GGML_ASSERT(state_rows == nullptr || keep || state_inplace); // Metal rows mode is a ring-path optimization + GGML_ASSERT(!state_inplace || state_rows != nullptr); + + if (state_inplace) { + // in-place mode (CPU): the cache-writing op reads each seq's state at + // cache row state_rows[seq] (rows_in = inp->s_copy_main: the live row, or + // the rollback plane of the same cell), updates the live row + // write_rows[seq] in place and writes snapshot slots 1..min(T, K)-1 + // straight into cache rows write_rows[slot*n_seqs + seq] -- the same + // rows the SET_ROWS / strided cpy below would have written. The result + // carries the attention scores only: no gather, no dst state planes, no + // write-back copy. Same float instruction sequence on the same inputs as + // the paths below, so the outputs and the resulting cache rows are + // bit-identical (tests/test-gdn-inplace.cpp). + const int64_t K = cparams.n_rs_seq + 1; + + const bool raw = gdn_raw_beta && gdn_raw_alpha && gdn_raw_dt_bias && gdn_raw_a; + ggml_tensor * gg = raw ? gdn_raw_alpha : g; + ggml_tensor * bb = raw ? gdn_raw_beta : b; + + ggml_tensor * write_rows = build_rs_write_rows(inp, K, n_seq_tokens, n_seqs); + + ggml_tensor * gdn_out = ggml_gated_delta_net_cache(ctx0, q, k, v, gg, bb, s, state_rows, write_rows, + (int) K, (int) mem_size); + if (raw) { + ggml_gated_delta_net_set_raw_gates(gdn_out, gdn_raw_dt_bias, gdn_raw_a); + } + if (n_seq_tokens > 1) { + res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il}); + } else { + res->add_fused_node({LLM_FUSED_OP_GDN_AR, gdn_out, il}); + } + + ggml_tensor * output = ggml_view_4d(ctx0, gdn_out, + S_v, H_v, n_seq_tokens, n_seqs, + ggml_row_size(gdn_out->type, S_v), + ggml_row_size(gdn_out->type, S_v * H_v), + ggml_row_size(gdn_out->type, S_v * H_v * n_seq_tokens), + 0); + cb(output, "attn_output", il); + + return output; + } if (!keep) { auto attn_out = build_delta_net(q, k, v, g, b, s, il); diff --git a/src/models/models.h b/src/models/models.h index a82133099d5c..1ae68a28edea 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -84,6 +84,10 @@ struct llm_build_delta_net_base : public llm_graph_context { // live state directly at cache row state_rows[seq] (inp->s_copy_main) -- // no gathered scratch; the snapshot write becomes a SET_ROWS the Metal // backend can fold into the fused op's epilogue. + // state_inplace (CPU only, requires state_rows): the cache-writing op + // ggml_gated_delta_net_cache updates the live rows in place and writes the + // snapshot slots straight into the cache -- no dst state planes, no SET_ROWS, + // no write-back cpy. Only valid on ubatches with inp->rs_inplace. // set per layer before build_recurrent_attn: the fused GDN op then receives the // pre-activation beta / alpha and folds sigmoid / softplus into its prologue // (ggml_gated_delta_net_set_raw_gates); the non-fused paths keep the activated g / b @@ -102,7 +106,8 @@ struct llm_build_delta_net_base : public llm_graph_context { ggml_tensor * b, ggml_tensor * s, int il, - ggml_tensor * state_rows = nullptr); + ggml_tensor * state_rows = nullptr, + bool state_inplace = false); }; struct llm_build_rwkv6_base : public llm_graph_context { @@ -2267,8 +2272,9 @@ struct llama_model_qwen35 : public llama_model_base { graph(const llama_model & model, const llm_graph_params & params); // device-dependent path choices, scanned once per graph build (not per layer) - bool gdn_state_rows_dev_ok = true; // every GPU device is Metal: fused GDN may read state rows in place - bool gdn_raw_gates_dev_ok = true; // every device is CPU/Metal/CUDA/ROCm/MUSA: fused GDN takes raw gates + bool gdn_state_rows_dev_ok = true; // every GPU device is Metal: fused GDN may read state rows in place + bool gdn_state_inplace_dev_ok = true; // no GPU/IGPU device at all: the in-place (cache-writing) GDN is CPU-only + bool gdn_raw_gates_dev_ok = true; // every device is CPU/Metal/CUDA/ROCm/MUSA: fused GDN takes raw gates private: ggml_tensor * build_layer_attn( llm_graph_input_attn_kv * inp_attn, diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index c5f816e96268..073a6e95b63e 100644 --- a/src/models/qwen35.cpp +++ b/src/models/qwen35.cpp @@ -149,13 +149,18 @@ llama_model_qwen35::graph::graph(const llama_model & model, const llm_graph_para 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) { - gdn_state_rows_dev_ok = false; - gdn_raw_gates_dev_ok = false; + gdn_state_rows_dev_ok = false; + gdn_state_inplace_dev_ok = false; + gdn_raw_gates_dev_ok = false; break; } // integrated GPUs (e.g. unified-memory CUDA devices) report IGPU, not GPU const bool is_gpu = ggml_backend_dev_type(ldev.dev) == GGML_BACKEND_DEVICE_TYPE_GPU || ggml_backend_dev_type(ldev.dev) == GGML_BACKEND_DEVICE_TYPE_IGPU; + if (is_gpu) { + // the cache-writing GDN op mutates one of its sources; only the CPU backend implements it + gdn_state_inplace_dev_ok = false; + } if (is_gpu && strcmp(reg_name, "MTL") != 0) { gdn_state_rows_dev_ok = false; } @@ -472,18 +477,40 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( // ring path: read per-seq live state directly from the cache inside the // fused GDN op (rows mode) instead of a gather 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. + // GGML_GDN_STATE_GATHER=1 restores the legacy gathered path (A/B) for both + // rows modes below. + // - in-place (CPU only, ggml_gated_delta_net_cache): the op updates the + // cache rows themselves -- no gather, no dst state planes, no write-back. + // Taken on every ubatch whose sequences all read their own cell + // (inp->rs_inplace); transition ubatches (a seq joins / leaves / is + // copied / reordered) fall back to the gathered path, which is exact and + // ordering-safe. Part of the graph reuse key (can_reuse_rs). + // - Metal rows + SET_ROWS (the src[6] variant, ring path only): kept for + // machines whose GPU devices are all Metal; other GPU backends reject it + // in supports_op, which would silently move the whole recurrent op to CPU. static const bool gdn_state_rows_env = getenv("GGML_GDN_STATE_GATHER") == nullptr; - const bool gdn_state_rows = gdn_state_rows_env && gdn_state_rows_dev_ok && cparams.n_rs_seq > 0; + // the in-place op always runs the fused kernel: only take it where the legacy + // path runs the fused kernel too (the ring path always does; the K == 1 path + // only with the fused_gdn_* probes on), so the A/B stays bit-identical + const bool fused_ok = cparams.n_rs_seq > 0 || (n_seq_tokens == 1 ? cparams.fused_gdn_ar : cparams.fused_gdn_ch); + + // GGML_GDN_INPLACE_K1_ONLY=1 restricts the in-place path to K == 1 (no + // speculation); by default the MTP speculative mode (K = n_rs_seq + 1) takes + // it too: slot 0 is updated in place, slots 1..K-1 are written straight into + // the snapshot planes of the same cell, and a pending rollback reads plane + // rs_idx of the same cell (rows_in != rows_out: copy-then-in-place). + static const bool gdn_inplace_k1_only_env = getenv("GGML_GDN_INPLACE_K1_ONLY") != nullptr; + + const bool gdn_state_inplace = gdn_state_rows_env && gdn_state_inplace_dev_ok && inp->rs_inplace && fused_ok && + (cparams.n_rs_seq == 0 || !gdn_inplace_k1_only_env); + const bool gdn_state_rows_metal = gdn_state_rows_env && gdn_state_rows_dev_ok && !gdn_state_inplace_dev_ok && + cparams.n_rs_seq > 0; + const bool gdn_state_rows = gdn_state_inplace || gdn_state_rows_metal; ggml_tensor * state; if (gdn_state_rows) { - state = build_rs_cache_view(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = build_rs_cache_view(inp, ssm_states_all, hparams.n_embd_s(), n_seqs, gdn_state_inplace); cb(state, "state_cache_view", il); } else { state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); @@ -560,7 +587,7 @@ ggml_tensor * llama_model_qwen35::graph::build_layer_attn_linear( 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, - gdn_state_rows ? inp->s_copy_main : nullptr); + gdn_state_rows ? inp->s_copy_main : nullptr, gdn_state_inplace); // 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/CMakeLists.txt b/tests/CMakeLists.txt index 67ff970be0b3..17aca7d3bdb4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -169,6 +169,11 @@ endif () llama_build(test-recurrent-state-rollback.cpp) +# byte-for-byte check of the cache-writing GDN op (ggml_gated_delta_net_cache) +# against the legacy gathered and rows+set_rows paths; CPU only, tiny synthetic +# shapes, no model needed +llama_build_and_test(test-gdn-inplace.cpp) + if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) llama_build_and_test(test-unicode.cpp) diff --git a/tests/test-gdn-inplace.cpp b/tests/test-gdn-inplace.cpp new file mode 100644 index 000000000000..c52eeb008af7 --- /dev/null +++ b/tests/test-gdn-inplace.cpp @@ -0,0 +1,375 @@ +// Byte-for-byte check of the cache-writing GDN op (ggml_gated_delta_net_cache) +// against the two legacy paths the graph used to emit, on identical inputs: +// +// A (gathered): ggml_get_rows(cache, rows_in) -> ggml_gated_delta_net(K) +// -> write-back cpy of the min(T, K) state planes into the +// cache rows slot*mem_size + head + seq +// (src/models/delta-net-base.cpp, !keep and keep/gathered) +// B (rows): ggml_gated_delta_net_rows(cache view, rows_in, K) +// -> ggml_set_rows(cache, snapshots, write_rows) +// (the Metal rows mode) +// C (in-place): ggml_gated_delta_net_cache(cache view, rows_in, write_rows, K) +// -- the state never leaves the cache +// +// For every case the attention scores of C must equal A and B byte for byte +// and the WHOLE cache tensor after C must equal the cache after A and after +// B byte for byte (touched rows identical, untouched rows untouched). +// +// Matrix: K in {1, 3}, T in {1, 3}, n_seqs in {1, 4, 8}, rows_in in plane 0 +// (identity: rows_in == rows_out) and in plane r in 1..K-1 (rollback read of +// the same cell), raw gates on/off, H_k == H_v and H_k == H_v/2, 1 and 4 +// threads. KDA (vector gate) off, as in the qwen35 graph. +// +// CPU only; runs on tiny synthetic shapes (S_v = 32, H_v = 4) and needs no +// model. + +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include +#include + +struct gdn_case { + int64_t S_v; + int64_t H_v; + int64_t H_k; + int64_t T; + int64_t n_seqs; + int64_t K; + int64_t mem_size; + int64_t head; + int64_t r_in; // plane read for every seq: rows_in[i] = r_in*mem_size + head + i + bool raw_gates; + int n_threads; +}; + +// deterministic input set shared by the three paths +struct gdn_inputs { + std::vector q, k, v, g, beta, dt_bias, a, cache; + std::vector rows_in; + std::vector rows_out; +}; + +static gdn_inputs make_inputs(const gdn_case & c, uint32_t seed) { + std::mt19937 rng(seed); + auto fill = [&](std::vector & dst, size_t n, float lo, float hi) { + std::uniform_real_distribution d(lo, hi); + dst.resize(n); + for (size_t i = 0; i < n; ++i) { + dst[i] = d(rng); + } + }; + + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t n_rows = c.mem_size * c.K; + + gdn_inputs in; + fill(in.q, c.S_v * c.H_k * c.T * c.n_seqs, -1.0f, 1.0f); + fill(in.k, c.S_v * c.H_k * c.T * c.n_seqs, -1.0f, 1.0f); + fill(in.v, c.S_v * c.H_v * c.T * c.n_seqs, -1.0f, 1.0f); + if (c.raw_gates) { + // pre-activation: g -> a[h]*softplus(g + dt_bias[h]), beta -> sigmoid(beta) + fill(in.g, c.H_v * c.T * c.n_seqs, -2.0f, 2.0f); + fill(in.beta, c.H_v * c.T * c.n_seqs, -3.0f, 3.0f); + fill(in.dt_bias, c.H_v, -1.0f, 1.0f); + fill(in.a, c.H_v, -1.0f, -0.1f); + } else { + // activated: g = log decay in (-2, 0), beta in (0, 1) + fill(in.g, c.H_v * c.T * c.n_seqs, -2.0f, 0.0f); + fill(in.beta, c.H_v * c.T * c.n_seqs, 0.0f, 1.0f); + } + // every row of every plane is random so a wrong row read/write shows up + fill(in.cache, D * n_rows, -1.0f, 1.0f); + + in.rows_in.resize(c.n_seqs); + for (int64_t i = 0; i < c.n_seqs; ++i) { + in.rows_in[i] = (int32_t) (c.r_in * c.mem_size + c.head + i); + } + + // set_input_rs mapping: row of (slot s, seq i) = s*mem_size + head + i + const int64_t n_write = std::min(c.T, c.K); + in.rows_out.resize(n_write * c.n_seqs); + for (int64_t r = 0; r < n_write * c.n_seqs; ++r) { + in.rows_out[r] = (r / c.n_seqs) * c.mem_size + c.head + (r % c.n_seqs); + } + + return in; +} + +struct gdn_tensors { + ggml_tensor * q; + ggml_tensor * k; + ggml_tensor * v; + ggml_tensor * g; + ggml_tensor * beta; + ggml_tensor * dt_bias; + ggml_tensor * a; + ggml_tensor * cache; // 2D [D, mem_size*K] + ggml_tensor * rows_in; // I32 [n_seqs] + ggml_tensor * rows_out; // I64 [min(T,K)*n_seqs] +}; + +static gdn_tensors make_tensors(ggml_context * ctx, const gdn_case & c, const gdn_inputs & in) { + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t n_rows = c.mem_size * c.K; + + gdn_tensors t; + t.q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c.S_v, c.H_k, c.T, c.n_seqs); + t.k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c.S_v, c.H_k, c.T, c.n_seqs); + t.v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, c.S_v, c.H_v, c.T, c.n_seqs); + t.g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, c.H_v, c.T, c.n_seqs); + t.beta = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, c.H_v, c.T, c.n_seqs); + t.dt_bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, c.H_v); + t.a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, c.H_v); + t.cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, n_rows); + t.rows_in = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, c.n_seqs); + t.rows_out = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, (int64_t) in.rows_out.size()); + + memcpy(t.q->data, in.q.data(), ggml_nbytes(t.q)); + memcpy(t.k->data, in.k.data(), ggml_nbytes(t.k)); + memcpy(t.v->data, in.v.data(), ggml_nbytes(t.v)); + memcpy(t.g->data, in.g.data(), ggml_nbytes(t.g)); + memcpy(t.beta->data, in.beta.data(), ggml_nbytes(t.beta)); + if (c.raw_gates) { + memcpy(t.dt_bias->data, in.dt_bias.data(), ggml_nbytes(t.dt_bias)); + memcpy(t.a->data, in.a.data(), ggml_nbytes(t.a)); + } else { + memset(t.dt_bias->data, 0, ggml_nbytes(t.dt_bias)); + memset(t.a->data, 0, ggml_nbytes(t.a)); + } + memcpy(t.cache->data, in.cache.data(), ggml_nbytes(t.cache)); + memcpy(t.rows_in->data, in.rows_in.data(), ggml_nbytes(t.rows_in)); + memcpy(t.rows_out->data, in.rows_out.data(), ggml_nbytes(t.rows_out)); + + return t; +} + +static ggml_context * new_ctx() { + ggml_init_params params = { + /*.mem_size =*/ 64u * 1024u * 1024u, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ false, + }; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + fprintf(stderr, "ggml_init failed\n"); + exit(1); + } + return ctx; +} + +struct path_result { + std::vector attn; // attention scores, S_v*H_v*T*n_seqs floats + std::vector cache; // whole cache tensor after the graph ran +}; + +static void snapshot(path_result & r, const ggml_tensor * attn_src, size_t attn_bytes, const ggml_tensor * cache) { + r.attn.resize(attn_bytes); + memcpy(r.attn.data(), attn_src->data, attn_bytes); + r.cache.resize(ggml_nbytes(cache)); + memcpy(r.cache.data(), cache->data, ggml_nbytes(cache)); +} + +// A: gathered legacy path, exactly as build_recurrent_attn's !keep (K == 1) +// and keep/gathered (K > 1) branches emit it +static path_result run_gathered(const gdn_case & c, const gdn_inputs & in) { + ggml_context * ctx = new_ctx(); + gdn_tensors t = make_tensors(ctx, c, in); + + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t attn_score_elems = c.S_v * c.H_v * c.T * c.n_seqs; + const int64_t state_size_per_snap = D * c.n_seqs; + const int64_t n_written = std::min(c.T, c.K); + const size_t row_size = D * sizeof(float); + + ggml_cgraph * gf = ggml_new_graph(ctx); + + // build_rs main gather (identity or rollback plane) -> 4D state + ggml_tensor * gathered = ggml_get_rows(ctx, t.cache, t.rows_in); + ggml_tensor * s4d = ggml_reshape_4d(ctx, gathered, c.S_v, c.S_v, c.H_v, c.n_seqs); + + ggml_tensor * gdn_out = ggml_gated_delta_net(ctx, t.q, t.k, t.v, t.g, t.beta, s4d, c.K); + if (c.raw_gates) { + ggml_gated_delta_net_set_raw_gates(gdn_out, t.dt_bias, t.a); + } + + if (c.K == 1) { + // build_delta_net_fused new_state view + ggml_cpy into rows [head, head + n_seqs) + ggml_tensor * new_state = ggml_view_4d(ctx, gdn_out, + c.S_v, c.S_v, c.H_v, c.n_seqs, + ggml_row_size(gdn_out->type, c.S_v), + ggml_row_size(gdn_out->type, c.S_v * c.S_v), + ggml_row_size(gdn_out->type, c.S_v * c.S_v * c.H_v), + ggml_row_size(gdn_out->type, c.S_v * c.H_v * c.T * c.n_seqs)); + ggml_build_forward_expand(gf, + ggml_cpy(ctx, new_state, + ggml_view_2d(ctx, t.cache, D, c.n_seqs, t.cache->nb[1], c.head * row_size))); + } else { + // strided 3D cpy: snapshot slot s -> plane s of the cache at head + ggml_tensor * src = ggml_view_3d(ctx, gdn_out, + D, c.n_seqs, n_written, + ggml_row_size(gdn_out->type, D), + ggml_row_size(gdn_out->type, state_size_per_snap), + ggml_row_size(gdn_out->type, attn_score_elems)); + ggml_tensor * dst = ggml_view_3d(ctx, t.cache, + D, c.n_seqs, n_written, + t.cache->nb[1], + (size_t) c.mem_size * row_size, + (size_t) c.head * row_size); + ggml_build_forward_expand(gf, ggml_cpy(ctx, src, dst)); + } + + if (ggml_graph_compute_with_ctx(ctx, gf, c.n_threads) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "gathered graph compute failed\n"); + exit(1); + } + + path_result r; + snapshot(r, gdn_out, attn_score_elems * sizeof(float), t.cache); + ggml_free(ctx); + return r; +} + +// B: rows mode + SET_ROWS, exactly as the keep/rows branch emits it +static path_result run_rows(const gdn_case & c, const gdn_inputs & in) { + ggml_context * ctx = new_ctx(); + gdn_tensors t = make_tensors(ctx, c, in); + + const int64_t D = c.S_v * c.S_v * c.H_v; + const int64_t attn_score_elems = c.S_v * c.H_v * c.T * c.n_seqs; + const int64_t n_written = std::min(c.T, c.K); + + ggml_cgraph * gf = ggml_new_graph(ctx); + + ggml_tensor * gdn_out = ggml_gated_delta_net_rows(ctx, t.q, t.k, t.v, t.g, t.beta, t.cache, t.rows_in, (int) c.K); + if (c.raw_gates) { + ggml_gated_delta_net_set_raw_gates(gdn_out, t.dt_bias, t.a); + } + + ggml_tensor * snaps = ggml_view_2d(ctx, gdn_out, + D, n_written * c.n_seqs, + ggml_row_size(gdn_out->type, D), + ggml_row_size(gdn_out->type, attn_score_elems)); + + ggml_build_forward_expand(gf, ggml_set_rows(ctx, t.cache, snaps, t.rows_out)); + + if (ggml_graph_compute_with_ctx(ctx, gf, c.n_threads) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "rows graph compute failed\n"); + exit(1); + } + + path_result r; + snapshot(r, gdn_out, attn_score_elems * sizeof(float), t.cache); + ggml_free(ctx); + return r; +} + +// C: the cache-writing op +static path_result run_inplace(const gdn_case & c, const gdn_inputs & in) { + ggml_context * ctx = new_ctx(); + gdn_tensors t = make_tensors(ctx, c, in); + + const int64_t attn_score_elems = c.S_v * c.H_v * c.T * c.n_seqs; + + ggml_cgraph * gf = ggml_new_graph(ctx); + + ggml_tensor * gdn_out = ggml_gated_delta_net_cache(ctx, t.q, t.k, t.v, t.g, t.beta, t.cache, + t.rows_in, t.rows_out, (int) c.K, (int) c.mem_size); + if (c.raw_gates) { + ggml_gated_delta_net_set_raw_gates(gdn_out, t.dt_bias, t.a); + } + + if (ggml_nelements(gdn_out) != attn_score_elems) { + fprintf(stderr, "in-place result carries %lld elements, expected scores only (%lld)\n", + (long long) ggml_nelements(gdn_out), (long long) attn_score_elems); + exit(1); + } + + ggml_build_forward_expand(gf, gdn_out); + + if (ggml_graph_compute_with_ctx(ctx, gf, c.n_threads) != GGML_STATUS_SUCCESS) { + fprintf(stderr, "in-place graph compute failed\n"); + exit(1); + } + + path_result r; + snapshot(r, gdn_out, attn_score_elems * sizeof(float), t.cache); + ggml_free(ctx); + return r; +} + +static size_t count_diff(const std::vector & a, const std::vector & b) { + if (a.size() != b.size()) { + return std::max(a.size(), b.size()); + } + size_t n = 0; + for (size_t i = 0; i < a.size(); ++i) { + n += a[i] != b[i]; + } + return n; +} + +int main() { + const int64_t S_v = 32; + const int64_t H_v = 4; + + int n_cases = 0; + int n_failed = 0; + uint32_t seed = 1234; + + for (int64_t K : { (int64_t) 1, (int64_t) 3 }) { + for (int64_t T : { (int64_t) 1, (int64_t) 3 }) { + for (int64_t n_seqs : { (int64_t) 1, (int64_t) 4, (int64_t) 8 }) { + for (int64_t r_in = 0; r_in < K; ++r_in) { + for (int raw = 0; raw < 2; ++raw) { + for (int64_t H_k : { H_v, H_v / 2 }) { + for (int n_threads : { 1, 4 }) { + gdn_case c; + c.S_v = S_v; + c.H_v = H_v; + c.H_k = H_k; + c.T = T; + c.n_seqs = n_seqs; + c.K = K; + c.mem_size = n_seqs + 3; // cells before and after the ubatch stay untouched + c.head = 2; + c.r_in = r_in; + c.raw_gates = raw != 0; + c.n_threads = n_threads; + + const gdn_inputs in = make_inputs(c, seed++); + + const path_result A = run_gathered(c, in); + const path_result B = run_rows(c, in); + const path_result C = run_inplace(c, in); + + const size_t d_attn_A = count_diff(C.attn, A.attn); + const size_t d_attn_B = count_diff(C.attn, B.attn); + const size_t d_cache_A = count_diff(C.cache, A.cache); + const size_t d_cache_B = count_diff(C.cache, B.cache); + + // the check is not vacuous: the in-place op must have changed the live rows + std::vector cache0((const uint8_t *) in.cache.data(), + (const uint8_t *) in.cache.data() + in.cache.size() * sizeof(float)); + const bool touched = count_diff(C.cache, cache0) != 0; + + const bool ok = d_attn_A == 0 && d_attn_B == 0 && d_cache_A == 0 && d_cache_B == 0 && touched; + + printf("K=%lld T=%lld n_seqs=%lld r_in=%lld raw=%d H_k=%lld nth=%d : attn diff A/B = %zu/%zu, " + "cache diff A/B = %zu/%zu, cache touched = %d -> %s\n", + (long long) K, (long long) T, (long long) n_seqs, (long long) r_in, raw, (long long) H_k, n_threads, + d_attn_A, d_attn_B, d_cache_A, d_cache_B, (int) touched, ok ? "OK" : "FAIL"); + + n_cases += 1; + n_failed += !ok; + }}}}}}} + + printf("%d cases, %d failed\n", n_cases, n_failed); + + return n_failed == 0 ? 0 : 1; +}