cuda: fold the recurrent-state gather into the GATED_DELTA_NET kernel, per-context registry (+3.8% TG on Ada)

[professorpalmer]

Summary

Split out of #215 at the reviewers' request so the activation-layout change there can be compared with #218 on its own. This PR is only the recurrent-state gather fusion for GATED_DELTA_NET, with the registry moved into the backend context (the P1 from the #215 review).

build_rs gathers each layer's live recurrent state via GET_ROWS(cache, s_copy) into a temp (3 MB per layer for Bonsai 2 27B) that only the GDN kernel reads. Traced in-graph with CUPTI on an RTX 4070, the gather kernel costs ~8 us/layer, and its dirty output sits in L2 until the FFN weight stream evicts it, which made the gate/up GEMVs in GDN layers ~15% slower than the identical GEMVs in attention layers.

ggml_cuda_try_gdn_gather_skip recognises GET_ROWS -> [RESHAPE] -> GATED_DELTA_NET src[5] with no other consumer of the gathered tensor, skips the GET_ROWS node, and records (cache base, ids, row stride) for that GDN node. The kernel then reads cache[s_ids[seq]] directly: no temp, no kernel, no writeback. Single-sequence only (with several sequences a gathered row could alias a row another sequence writes in the same op). GGML_CUDA_GDN_GATHER_FUSION=0 disables it.

Registry ownership

Registrations live in ggml_backend_cuda_context::gdn_gather_context (ggml_cuda_gdn_gather_context, next to the existing ggml_cuda_stream_context), keyed by node pointer, reset by the evaluating context at the start of each graph evaluation/capture and read by ggml_cuda_op_gated_delta_net through the ctx it already receives. No static state, no cross-context or cross-thread visibility: two contexts decoding at once each see only their own skips, so the interleaving described in the #215 review (A registers, B clears, A looks up) cannot occur.

Measurements

RTX 4070 12 GB, Ternary-Bonsai-2-27B-PTQ1_0.gguf, CUDA 13.3, llama-bench -ngl 99 -fa 1 -ctk q4_0 -ctv q4_0 -r 3, stock clocks, three alternating rounds, means:

	tg128	pp512
#215 without this (GGML_CUDA_GDN_GATHER_FUSION=0)	56.4	1256
#215 + this PR	58.5 (+3.8%)	1259

The gain is from removing 64 gather kernels and their L2 writebacks per token, so it is independent of which PTQ1_0 activation layout lands; it applies equally on top of #218.

Correctness

• test-backend-ops -b CUDA0 -o GATED_DELTA_NET: 39/39 supported cases vs CPU. GET_ROWS untouched.
• 400 greedy tokens byte-identical with GGML_CUDA_GDN_GATHER_FUSION=0 and =1.

Scope

CUDA graph capture and eager evaluation both go through the same skip; under graphs the kernel arguments (cache base, ids pointer) are baked into the captured graph exactly as the GET_ROWS node's would have been, and s_copy is a graph input whose contents are read at kernel time. HIP/MUSA: the fusion is pure host-side graph logic plus two kernel arguments; the kernel change is a pointer select on state_in_offset. Not tested on HIP hardware.

[professorpalmer]

@MrFadiAi thanks — same fusion, two backend shapes, useful to have both written down.

Why this PR is a graph skip and not ggml_gated_delta_net_rows. On CUDA the rows variant (src[6] = int32 index) is not in supports_op. qwen35.cpp already documents that and forces the gathered form on any non-Metal GPU (gdn_state_rows_dev_ok = false unless every device is MTL), so a CUDA tree that built the rows op would silently move the whole recurrent node to the CPU. The skip in #220 is the CUDA-side equivalent: leave build_rs / GET_ROWS in the graph, do not evaluate it, and have the GDN kernel index cache[s_ids[seq]] the same way your rows kernel does. Net effect matches #187: no 3 MB temp, no gather launch, no L2 writeback. Single-sequence only, same alias reason you would have on a multi-seq rows launch.

n_rs_seq / non-speculative rollback. Checked. This path does not touch the ring and does not force a rollback window. We only skip a GET_ROWS that build_rs already emitted; n_rs_seq stays whatever the server set (0 on a plain slot). The corruption you saw is the rows graph (build_rs_cache_view + src[6]) expecting the speculative scheduler’s accept/reject. qwen35.cpp already gates that on cparams.n_rs_seq > 0. If CUDA later grows a real gated_delta_net_rows we will keep that gate; it is not a condition of this skip. Happy to add a one-line comment in ggml_cuda_try_gdn_gather_skip pointing at that, so nobody “fixes” this into rows without it.

bf16 state pool. Orthogonal and not in this PR. The CUDA kernel still takes float * for curr_state and the skip matcher requires cache->type == F32. Halving the state RMW is the larger iGPU lever you measured; on a 4070 the gather was the thing CUPTI showed (8 µs/layer + dirty L2 in front of the FFN GEMVs, +3.8% TG). A LLAMA_SSM_BF16_STATE CUDA path is a follow-up, not a reason to hold #220.

Vulkan kernel layout translating 1:1 is noted. Once CUDA has a rows supports_op, the skip becomes redundant and can come out.

[Copilot]

Copilot review overview

🔵 Needs a closer look

Graph elision and CUDA capture behavior require targeted regression coverage and hardware validation.

Review effort: Balanced
Findings: 2

Open (2)

• Add coverage for fused GET_ROWS reshape gated-delta graph · New
• Correct comment about temporary allocation remaining reserved · New

What changed in this PR

Fuses recurrent-state gathering into CUDA GATED_DELTA_NET execution to reduce kernel launches and memory traffic.

Changes:

• Detects eligible single-sequence gather graphs.
• Stores gather metadata per backend context.
• Reads selected cache rows directly in the CUDA kernel.

File	Description
ggml/​src/​ggml-cuda/​ggml-cuda.cu	Detects and skips eligible gathers.
ggml/​src/​ggml-cuda/​gated_delta_net.cuh	Documents the fused-gather interface.
ggml/​src/​ggml-cuda/​gated_delta_net.cu	Adds direct cache-row reads.
ggml/​src/​ggml-cuda/​common.cuh	Adds the per-context gather registry.

---

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

[professorpalmer]

Fair ask. test-backend-ops already covers GDN numerics on the kernel, but it does not emit this GET_ROWS+RESHAPE+GDN graph, so it never hits the matcher or CUDA-graph replay of the skipped launch. I am not stuffing a graph-capture case into this PR; I will add one as a follow-up that runs with the fusion on and off. Disable with GGML_CUDA_GDN_GATHER_FUSION=0 in the meantime.

[professorpalmer]

Correct, and the comment was wrong. 80132df now says the allocator still reserved the unused temp; the fusion only skips the kernel and its traffic.

[bri-prism]

Reviewed 80132df against the previous global-registry concern. The gather map is now backend-context-owned and reset for graph evaluation/capture. The matcher checks both GET_ROWS and intermediate RESHAPE consumer counts. I also checked an explicitly exported reshape with a real GGML graph: ggml_set_output propagates the output flag to the underlying gather, and the matcher declines the skip. No new finding in this focused review. CUDA execution, graph replay and concurrent-context stress have not been rerun here; the performance figures from the older combined branch should not be treated as measurements of this exact head.

[professorpalmer]

Thank you for walking the context-owned map and the ggml_set_output reshape case. That matches the intent: skip only a private GET_ROWS whose only consumer is the GDN node, and decline when the gather is a real graph output.

Agreed that the older combined-branch timings should not be read as this head. The +3.8% TG figure was measured with the fusion on the 4070 before the split; this PR is the same matcher and kernel path, now with the registry on ggml_backend_cuda_context. I still owe a capture/replay test (fusion on and off) as a follow-up. Disable with GGML_CUDA_GDN_GATHER_FUSION=0 until that lands.

[professorpalmer]

Landed as part of #221 (5244cea on prism), so closing this one. The registry is on ggml_backend_cuda_context (ggml_cuda_gdn_gather_context), as reviewed here, and GGML_CUDA_GDN_GATHER_FUSION=0 turns it off.

Still open, and I have not forgotten it: the GET_ROWS -> RESHAPE -> GATED_DELTA_NET test with a nonzero cache row, run eagerly and through CUDA graph capture/replay, with the fusion on and off. It is not in prism yet. I will send it as its own small PR against prism instead of reopening this one.

[professorpalmer]

The test I owed here is #296: GET_ROWS -> RESHAPE -> GATED_DELTA_NET with a nonzero cache row, in test-backend-ops, evaluated four times so CUDA runs it eagerly, captures a graph and replays it, with a different row on each run. Passes with the fusion on, with GGML_CUDA_GDN_GATHER_FUSION=0 and with GGML_CUDA_DISABLE_GRAPHS=1; a kernel hard-wired to one row fails it.