Repository navigation
llama: in-place recurrent (GDN) state update for multi-sequence decode + lazily zeroed RS buffer - #207
llama: in-place recurrent (GDN) state update for multi-sequence decode + lazily zeroed RS buffer#207alenizy361 wants to merge 4 commits into
Conversation
…ate) New op form ggml_gated_delta_net_cache (op_params[2] = 1, src[6] = rows_in I32[n_seqs], src[9] = rows_out I64[min(T,K)*n_seqs]): the CPU kernel reads each sequence's state at cache row rows_in[seq], updates the live row rows_out[seq] in place (no copy when the rows coincide) and writes snapshot slots 1..min(T,K)-1 straight into the cache; the result carries the attention scores only. No get_rows gather, no dst state planes, no write-back cpy / set_rows per layer: 6 -> 2 state-sized moves per seq. Graph side: llama_memory_recurrent_context::rs_inplace_ok(n_seqs), a side-effect-free predicate (never calls s_copy, which consumes rs_idx): !is_full && n == n_seqs && cells[head+i].src0 == head+i. Stored in llm_graph_input_rs::rs_inplace and added to can_reuse_rs so a reused graph can never be in the wrong mode. build_rs_cache_view(inplace) asserts the extra relocation range is empty, which makes the documented rows-mode ordering hazard unreachable for the in-place consumer. Model side (qwen35): gdn_state_inplace_dev_ok (no GPU/IGPU device), the gate gdn_state_rows_env && inplace_dev_ok && inp->rs_inplace && fused_ok && cparams.n_rs_seq == 0 (Stage 7A: K == 1 only), the Metal rows+SET_ROWS mode kept for Metal-only machines; build_recurrent_attn gains the state_inplace branch, every legacy branch is untouched and serves the transition ubatches (join/leave/seq_cp/reorder) and the env switch. GGML_GDN_STATE_GATHER=1 still forces the legacy gathered path (A/B). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R
… + 1) Drop the cparams.n_rs_seq == 0 restriction of the qwen35 gate: with --spec-type draft-mtp the cache-writing op now updates slot 0 in place on the live row and writes slots 1..min(T, K)-1 straight into the snapshot planes of the same cell (rows_out = the existing s_write_rows mapping slot*mem_size + head + seq); a pending rollback reads plane rs_idx of the same cell (rows_in != rows_out -> copy-then-in-place in the kernel). No dst state planes, no SET_ROWS, no per-thread state scratch. Transition ubatches still take the gathered K-slot path. GGML_GDN_INPLACE_K1_ONLY=1 restores the previous commit's K == 1 gate so the two can be bisected on the same binary; GGML_GDN_STATE_GATHER=1 still forces the legacy gathered path for everything. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R
…paths
tests/test-gdn-inplace.cpp (llama_build_and_test) runs, on identical
synthetic inputs and in three separate ggml contexts, the gathered path
(get_rows -> ggml_gated_delta_net(K) -> write-back cpy, exactly as
build_recurrent_attn emits it for K == 1 and K > 1), the rows path
(ggml_gated_delta_net_rows -> ggml_set_rows) and the new
ggml_gated_delta_net_cache, then memcmp's the attention scores and the
WHOLE cache tensor (touched rows identical, untouched rows untouched) and
checks the in-place op actually modified the cache.
Matrix: K in {1, 3}, T in {1, 3}, n_seqs in {1, 4, 8}, rows_in in plane 0
(identity) and in every rollback plane 1..K-1 of the same cell, raw gates
on/off, H_k == H_v and H_k == H_v/2, 1 and 4 threads: 192 cases, 0 bytes
differ. CPU only, no model.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R
The recurrent cache is mem_size*(1+n_rs_seq) rows per layer and was memset
at construction (ggml_backend_buffer_clear), so every cell of every
snapshot plane was resident from the start: --parallel 16 with n_rs_seq = 2
pinned ~4.7 GiB of the 15 GiB here while one active slot needs ~0.3 GiB,
evicting the mmap'd weights that decode streams on every token.
For the CPU buffer type on Linux the buffer is now one anonymous private
mmap wrapped with ggml_backend_cpu_buffer_from_ptr (tensors placed with a
ggml_tallocr exactly as ggml_backend_alloc_ctx_tensors_from_buft would),
with no eager memset: the kernel hands out zero pages on first touch, so
residency follows the cells actually written. clear(true) drops the pages
with madvise(MADV_DONTNEED) instead of a memset; the destructor unmaps
after freeing the buffers (from_ptr buffers do not own their memory).
Non-host bufts, non-Linux hosts and LLAMA_RS_EAGER_ZERO=1 keep today's
alloc + clear.
Exactness argument (in the ctor comment): every row the graph reads is
either the rs_z row zeroed in-graph before the read, a row written by an
earlier step (src0 = src of a used cell), a snapshot plane written by the
same sequence before a rollback reads it (rollback <= n_rs_seq, tail kept
in one ubatch, slots 0..min(T,K)-1 written), or a row written by
state_read_data; padding is never read.
The startup line keeps its format ('RS buffer size = N MiB') and appends
'(anonymous mmap, lazily zeroed)' when the lazy path is active.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01QrdPnP5KL449GsoWBZx52R
|
Note on the numbers: the tok/s figures in the description were measured with the AVX2 PQ2_0 kernels from #206 applied on top (this branch alone still runs the generic scalar PQ2_0 path on AVX2-only CPUs, where everything is compute-bound and the state-copy saving is not visible). The bit-identity checks (PPL, greedy md5) hold on this branch by itself. Full setup and raw tables: https://github.com/alenizy361/bonsai-cpu-speedup/blob/master/results/measurements.md |
bri-prism
left a comment
bri-prism
left a comment
There was a problem hiding this comment.
Agent review: posted by the maintainer's coding agent at their request.
No findings in this source pass. The GDN implementation is also included in #208; testing that head on native ARM64 CPU passed all 360 cases in its expanded test-gdn-inplace matrix. This is validation of the shared kernel, not a claim that this exact head received the full #208 test suite.
The Linux anonymous-mmap/lazy-zero path was inspected but not executed on the macOS review host. Please retain Linux allocation/clear validation and model-backed multi-sequence transition/rollback coverage. Coordinate landing with #208, which contains this change.
Reviewed commit: c875ea463d607e4c2d2f5743cca2a8813610520f.
|
Linux coverage you asked for, run on x86-64 (Ubuntu, GNU ld, 6 threads) at the exact commit you reviewed,
Every case reports
Model-backed on Before reporting that as a regression I re-ran the same test, same model, same flags, on Byte-identical failure, identical logit values. So the dirty-context path is already broken on the base and this PR neither causes nor worsens it. I did not chase the root cause further, since it is out of scope here — happy to open a separate issue for it if you'd like, and I can bisect it against upstream if that's useful. Still not covered, honestly. The Linux anonymous-mmap/lazy-zero allocation path is exercised implicitly by these runs (they ran on Linux with the new RS buffer), but I have not written a test that asserts the pages are actually lazily faulted rather than eagerly zeroed — so "inspected but not executed" still applies to that specific claim. Same for the server-level mixed drafting/non-drafting slot behaviour, which belongs on #208 anyway. One process note: |
| 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; |
| static const bool rs_eager_zero = getenv("LLAMA_RS_EAGER_ZERO") != nullptr; | ||
| #endif | ||
|
|
||
| for (auto & [buft, ctx] : ctx_map) { |
| // 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. |
|
Tested on an Intel box (Core Ultra, Arc B390 Xe3 iGPU), Windows 11, driver 32.0.101.8724, GCC 16.2 ucrt64, Release, static. Head Correctness: Throughput,
So in-place is about 3.7% ahead of gather at 8 slots on CPU, and neutral against One pre-existing item, not caused by this PR: Not run: 16 slots, a larger GDN model on CPU, or MTP speculation, so the larger multi-sequence gains in the description are untested here. |
|
@alenizy361 @jasontitus, #207, #208 and #276 all change the same state-mode decision in Could you two agree on that shape here? Our suggested order is #207 first (it introduces the selector), then #208 rebased on it, then #276 rebased onto the selector. Open items before #207 can go in:
For #208, two things we saw in the code: |
|
This needs a rebase from you, and not a mechanical one. Since you opened it, GDN rollback on prism moved from snapshot planes to replay (#358), and #354 changed how a pending rollback survives seq_cp. The in-place op writes speculative slots into snapshot planes and reads a pending rollback from plane rs_idx. On prism, qwen35 now picks between the replay gather (base row in plane 0 plus re-applied records) and the rows path, so the rule for which state the in-place update may overwrite needs to be worked out again against replay. The conflicts are in llama-memory-recurrent, llama-graph and qwen35.cpp. Two suggestions: (1) split the lazily zeroed RS buffer into its own PR. It's independent, conflicts only on an include, and the RSS saving at --parallel 16 is worth taking now. (2) When you rebase the in-place op, keep test-gdn-inplace and add a replay-mode case, so the byte-for-byte check covers the rollback path prism actually uses. |
Summary
For hybrid/recurrent models (
qwen35: 48 GDN layers),build_rsgathers every sequence's recurrent state withggml_get_rowsand the model writes the new state back withggml_cpyon every step. For multi-sequence decode that is most of the non-GEMM time:perfon a 16-sequence step showed__memmove_avx_unaligned_erms18 % +ggml_compute_forward_get_rows8 % (≈ 2 × 16 × 48 × 2 MiB of state traffic per step). Upstream has an open request for the single-sequence subset of this (ggml-org#28841); this PR handles any number of sequences.ggml_gated_delta_net_cache(q, k, v, g, beta, states /*2D cache view*/, rows_in, rows_out, K, mem_size): the CPU kernel reads cache rowrows_in[seq], updates cache rowrows_out[seq]in place (a copy only when the two differ, i.e. a rollback plane or a fresh cell) and writes snapshot slots1..min(T,K)-1straight into the same cell's planes. The attention output is the only result tensor. Token loop untouched → same float sequence as the existing kernel.llama_memory_recurrent_context::rs_inplace_ok(n_seqs):!is_full && n == n_seqs && cells[head+i].src0 == head+i(side-effect free; never callss_copy). Fresh sequences,seq_cpfan-out, cell reorder, pending relocation and the reserve context fail the predicate and take the untouched gathered path. The decision is part of the graph reuse key (can_reuse_rs).GGML_GDN_STATE_GATHER=1forces the legacy path (same-binary A/B);GGML_GDN_INPLACE_K1_ONLY=1restricts it to non-speculative ubatches.K = n_rs_seq + 1): the snapshot planes are written by the same op (rows_out= the existings_write_rowsmapping), so--spec-type draft-mtpkeeps its per-token rollback.mmapinstead of an eagermemset, so untouched cells/planes never become resident.--parallel 16+ MTP on a 15.6 GB machine went from 12.8 GB boot RSS / heavy swapping / 0.89 tok/s to 8.4 GB / 7.0 tok/s.LLAMA_RS_EAGER_ZERO=1restores the memset. Exactness argument in the constructor comment (every read row is either zeroed in-graph before the read, written by an earlier step, or a plane just copied).tests/test-gdn-inplace.cpp(registered): the new op vs the gathered path vs the rows+set_rowspath on identical synthetic inputs, K ∈ {1,3}, T ∈ {1,3}, n_seqs ∈ {1,4,8}, rollback planes, raw gates on/off, 1 and 4 threads — attention output and the whole cache compared byte for byte (192 cases, 0 failed; 360 with the follow-up PR).Measurements (i7-13620H, single-channel DDR5, Bonsai-2 27B PQ2_0, CPU only, t=12)
GGML_GDN_STATE_GATHER=1)Bit-identity: perplexity identical to all printed digits with the env switch on/off; 120-token greedy server output md5-identical (
cacc21fd31be) between the two paths.Notes