metal: fix the recurrent-state snapshot ring's CPU fallback and wide-row set_rows

[bri-prism]

What

Three commits that make the recurrent-state snapshot ring (n_rs_seq > 0) usable on the
Metal backend, plus two correctness harnesses for the rollback path.

Why

With the snapshot ring enabled, decode on a 64-layer hybrid gated-delta-net model collapsed
by about 44 percent (22.4 tok/s against a 40.5 tok/s ring-disabled baseline). Profiling with
GGML_SCHED_DEBUG=2 showed the run was CPU bound, not GPU bound: 87s of user CPU time for a
17.7s wall-clock run, with about 5 cores busy. The graph dump explained it. With the ring off,
one node runs on the CPU backend. With the ring on, 49 do, and 48 of those are IM2COL, one
per recurrent layer. Graph split markers went from 16 to 784.

There turned out to be two independent causes, and fixing either one alone leaves most of the
regression in place.

How

1. metal: fix 1D im2col kernel/dispatch mismatch + wide-row f32 set_rows

Two kernel-level bugs. The im2col pipeline was selected by ne00*ne01 <= 1024 while the
dispatch geometry was chosen by KH*KW. For 1D im2col ne01 is IC rather than KH, so any 1D
im2col with KW*IC > 1024 paired kernel_im2col_ext with the wrong dispatch and produced
garbage. Separately, the generic set_rows kernel assigns one threadgroup per destination row,
which cannot saturate bandwidth for the very wide rows the ring writes.

2. metal: keep the recurrent conv-state write on the GPU

build_conv_state() materializes the n_write overlapping conv windows with a single
ggml_im2col over tail, a strided ggml_view_3d of conv_input. The Metal backend accepts
im2col only when its input is contiguous:

case GGML_OP_IM2COL:
    return ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32 && ...

so the op fell to the CPU backend, one node per recurrent layer, splitting the graph at every
one of them. Making the input contiguous with ggml_cont first keeps it on the GPU. At decode
n_write is min(n_seq_tokens, n_rs_seq + 1) == 1, where im2col would only repackage that
one window, so the reshape is taken directly.

CPU-resident nodes with the ring enabled go from 49 to 1, and graph split markers from 784
back to 16.

3. examples: add recurrent-state rollback correctness harnesses

A forward-only generation never reads the snapshot slots back, so it cannot catch a bad
snapshot write. llama-rs-rollback generates a reference stream, then regenerates while
periodically over-generating tokens, calling llama_memory_seq_rm() to discard them, and
continuing. llama-rs-rollback-multi does the same across concurrent sequences where only
sequence 0 rolls back, which covers the failure mode a single-sequence test cannot see, since
the ring's snapshot rows are seq-major.

Numbers

Metal, M5 Pro, 64-layer hybrid gated-delta-net model, with the ring enabled (n_rs_seq = 7)
against the ring-disabled baseline.

Decode, isolating the two kernel fixes (baseline 40.5 tok/s):

build	tg128	vs baseline
before	22.44	-44.6%
commit 1 only	27.80	-31.4%
commit 2 only	28.86	-28.8%
both	38.51	-4.9%

The two are strongly super additive. Removing the CPU fallback is what lets the faster
set_rows kernel actually pay off.

Full picture with both applied:

test	ring off	ring on	before
pp512	443.2	438.1 (-1.2%)	-26.5%
pp8	88.5	72.2 (-18.4%)	-26.5%
tg128	40.0	38.1 (-4.7%)	-44.6%

Prefill is back to parity. The residual pp8 gap is inherent snapshot work rather than a
fallback: an 8 token batch writes 8 snapshots, one per token, where pp512 amortizes the same
8 across 512 tokens.

Testing

• test-backend-ops test -b MTL0 passes for IM2COL, SET_ROWS, GATED_DELTA_NET, CONT, CPY and
  SSM_CONV. Commit 1 adds a 1D im2col case (KW=3, IC=1024) that fails before it.
• Greedy generation is byte identical to the ring-disabled baseline over 512 tokens, on two
  quantizations, at n_rs_seq of 1 and 7.
• Rollback harness passes at rewind depths 1, 4 and 6, and multi-sequence at S=3 and S=4.
  As a control, at n_rs_seq = 0 the partial rollback is refused, so the tests are not
  passing vacuously.
• Measured logit-KLD against the ring-disabled baseline is 5.0e-5 maximum with a median of
  zero and 100 percent same-top-1, which is about ten times below the evaluation noise floor.
  Note this needs -ub 1, since a default perplexity run has n_write > 1 and would exercise
  the im2col path instead.

Notes for review

The remaining pp8 overhead is the real cost of keeping a rollback ring, not a bug: a batch of
8 writes 8 snapshots. For block-verify speculative decoding at k=8 that puts break-even at
about 55 percent acceptance, down from about 62 percent before these commits.

Nothing in CI exercises this. The ubuntu and windows jobs have no Metal runner, so the
kernels here are built and tested nowhere in the matrix, and the two new harnesses are
examples/ that only build on macOS. All the evidence above is from a local M5 Pro run.

[Copilot]

Pull request overview

Improves Metal recurrent-state snapshot performance and adds rollback validation harnesses.

Changes:

• Fixes Metal IM2COL selection and adds wide-row SET_ROWS.
• Keeps recurrent convolution snapshots GPU-resident.
• Adds single- and multi-sequence rollback harnesses.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/test-backend-ops.cpp	Adds wide-row and 1D IM2COL cases.
src/models/delta-net-base.cpp	Makes snapshot inputs contiguous.
ggml/src/ggml-metal/ggml-metal.metal	Adds wide-row Metal kernels.
ggml/src/ggml-metal/ggml-metal-ops.cpp	Dispatches wide-row kernels.
ggml/src/ggml-metal/ggml-metal-impl.h	Defines wide-row kernel arguments.
ggml/src/ggml-metal/ggml-metal-device.h	Declares pipeline lookup.
ggml/src/ggml-metal/ggml-metal-device.cpp	Registers pipelines and adjusts IM2COL selection.
examples/rs-rollback/rs-rollback.cpp	Adds single-sequence rollback harness.
examples/rs-rollback/rs-rollback-multi.cpp	Adds multi-sequence rollback harness.
examples/rs-rollback/CMakeLists.txt	Builds rollback executables.
examples/CMakeLists.txt	Registers the new examples.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[bri-prism]

This was already addressed in d451930, which landed after the commit this review ran against (a376425). The fixed 1024 threshold is gone. ggml_metal_op_im2col now builds the normal pipeline first, asks it for its own limit via ggml_metal_pipeline_max_theads_per_threadgroup, and derives use_ext from that, then passes the flag into ggml_metal_library_get_pipeline_im2col. The kernel choice and the dispatch geometry come from one device-aware decision, so they cannot disagree on a device whose limit is below 1024.

[bri-prism]

Also fixed in d451930. rs-rollback.cpp now counts rollbacks and fails when none happened, rather than reporting PASS on a run that never entered the branch. The failure message names the condition, that n_predict has to exceed stride plus rewind.

[bri-prism]

Same fix in d451930. The multi harness tracks n_rollbacks, and the result is now (n_rollbacks > 0) and (total_mismatch == 0), so a run too short to roll back reports FAIL instead of a vacuous PASS.

[bri-prism]

Fixed in d451930. The argument check now rejects n_seqs <= 0, rewind <= 0, stride <= 0 and n_predict <= 0 up front, before the prompt and output vectors are built, so run() can no longer index out[0] on an empty vector.

[khosravipasha]

Could you rebase this onto prism-v7 and retarget it there? The earlier base switch without a rebase pulled the entire old prism history into the diff, which is why it was flipped back for now.

From a quick check the fixes are still needed on prism-v7: supports_op still requires a contiguous im2col src1 (ggml-metal-device.m, GGML_OP_IM2COL case), and the delta-net conv-state write path is unchanged, so the ring-on CPU fallback should reproduce there. Heads up that the Metal shaders moved to per-op files under ggml/src/ggml-metal/kernels/ on that branch, so the ggml-metal.metal hunks need the same relanding treatment #121 went through.

Once it is on prism-v7 we will test on M5 Max and the CUDA box and merge.

[bri-prism]

Rebased onto prism-v7 and retargeted. The diff is now +607/-13 across 10 files instead of the old prism history.

Relanding notes, as you predicted: kernel_set_rows_f32_wide went into kernels/quantize.metal rather than being cherry-picked, get_pipeline_set_rows takes the op on this branch so the wide variant now matches it, and test_set_rows gained a src_type parameter so the new wide cases use it.

One thing did not come across, and it is worth flagging because you expected it to. The delta-net conv-state change is not included. It writes only the snapshots this ubatch produced, and that is only correct because the ring head rotated backwards, which needs rs_ring in llama_memory_recurrent. That does not exist on prism-v7. Porting it as is would leave the older snapshots stale and quietly give wrong recurrent state, so I left it out rather than have it look applied.

Related: the im2col stranding you expected to reproduce there does not arise the same way. The conv-state write path on prism-v7 uses ggml_cpy with views and has no im2col in it at all. So that fix needs either the ring to land first, or a different approach aimed at the cpy path.

Testing on M5: SET_ROWS 175/175, IM2COL 93/93, 3/3 backends. I also checked that kernel_set_rows_f32_wide actually compiles and runs on the 4096-wide cases rather than being quietly skipped.

The old branch tip is kept at backup/118-pre-v7-rebase if you want the delta-net work for reference.

[khosravipasha]

LGTM to merge now. We can test/optimize further after we get all change in.