llama: add opt-in SM75 recurrent state replay

ggml/include/ggml-cuda-gdn-transaction.h

@@ -0,0 +1,14 @@
+#pragma once
+
+#include <cstdint>
+
+struct ggml_cuda_gdn_replay_args {
+    const float * prior;
+    float * state;
+    const float * log;
+    int64_t head_count;
+    int32_t state_dim;
+    int32_t n_keep;
+};
+
+using ggml_cuda_gdn_replay_t = bool (*)(const ggml_cuda_gdn_replay_args *, int32_t);

ggml/src/ggml-cpu/ggml-cpu.cpp

@@ -422,6 +422,9 @@ static ggml_backend_buffer_t ggml_backend_cpu_device_buffer_from_host_ptr(ggml_b
 }
 
 static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
+    if (op->op == GGML_OP_GATED_DELTA_NET && op->src[6] != nullptr) {
+        return false;
+    }
     const struct ggml_tensor * src0 = op->src[0];
     const struct ggml_tensor * src1 = op->src[1];
 

ggml/src/ggml-cuda/gated_delta_net.cu

@@ -27,7 +27,10 @@ gated_delta_net_cuda(const float * q,
                                      const uint3   rq3_magic,
                                      float         scale,
                                      int64_t       state_slot_stride,
-                                     int           K) {
+                                     int           K,
+                                     float *       txn_log,
+                                     float *       txn_prior,
+                                     int           txn_slots) {
     const uint32_t h_idx    = blockIdx.x;
     const uint32_t sequence = blockIdx.y;
     // each warp owns one column, using warp-level primitives to reduce across rows
@@ -61,6 +64,13 @@ gated_delta_net_cuda(const float * q,
     }
 
     for (int t = 0; t < n_tokens; t++) {
+        const int log_start = n_tokens > txn_slots ? n_tokens - txn_slots : 0;
+        if (txn_prior != nullptr && t == log_start) {
+#pragma unroll
+            for (int r = 0; r < rows_per_lane; r++) {
+                txn_prior[state_in_offset + col * S_v + r * warp_size + lane] = s_shard[r];
+            }
+        }
         const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1;
         const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1;
         const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1;
@@ -95,6 +105,22 @@ gated_delta_net_cuda(const float * q,
             // delta[col] = (v[col] - g * kv[col]) * beta
             float delta_col = (v_t[col] - g_val * kv_col) * beta_val;
 
+            if (txn_log != nullptr && t >= log_start) {
+                float * record = txn_log + (t - log_start) * H * (1 + 2 * S_v);
+                if (col == 0) {
+                    if (lane == 0) {
+                        record[h_idx] = g_val;
+                    }
+#pragma unroll
+                    for (int r = 0; r < rows_per_lane; r++) {
+                        record[H + h_idx * S_v + r * warp_size + lane] = k_reg[r];
+                    }
+                }
+                if (lane == 0) {
+                    record[H + H * S_v + h_idx * S_v + col] = delta_col;
+                }
+            }
+
             // fused: S[i][col] = g * S[i][col] + k[i] * delta[col]
             // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i]
             float attn_partial = 0.0f;
@@ -176,7 +202,8 @@ static void launch_gated_delta_net(
         int64_t sv1,   int64_t sv2, int64_t sv3,
         int64_t sb1,   int64_t sb2, int64_t sb3,
         int64_t neqk1, int64_t rq3,
-        float scale, int64_t state_slot_stride, int K, cudaStream_t stream) {
+        float scale, int64_t state_slot_stride, int K, cudaStream_t stream,
+        float * txn_log = nullptr, float * txn_prior = nullptr, int txn_slots = 0) {
     //TODO: Add chunked kernel for even faster pre-fill
     const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
     const int num_warps = 4;
@@ -192,26 +219,26 @@ static void launch_gated_delta_net(
             ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t>, launch_params,
                 q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, txn_log, txn_prior, txn_slots);
             break;
         case 32:
             ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t>, launch_params,
                 q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, txn_log, txn_prior, txn_slots);
             break;
         case 64: {
             ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t>, launch_params,
                 q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, txn_log, txn_prior, txn_slots);
             break;
         }
         case 128: {
             ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t>, launch_params,
                 q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, txn_log, txn_prior, txn_slots);
             break;
         }
         default:
@@ -286,6 +313,21 @@ static void ggml_cuda_op_gated_delta_net_impl(
     const int K = ggml_get_op_params_i32(dst, 0);
     const bool keep_rs = K > 1;
 
+    float * txn_log = nullptr;
+    float * txn_prior = nullptr;
+    int txn_slots = 0;
+    if (dst->src[6] != nullptr) {
+        GGML_ASSERT(!kda && n_seqs == 1 && K == 1 && dst->src[7] != nullptr);
+        GGML_ASSERT(dst->src[6]->type == GGML_TYPE_F32 && dst->src[7]->type == GGML_TYPE_F32);
+        GGML_ASSERT(ggml_is_contiguous(dst->src[6]) && ggml_is_contiguous(dst->src[7]));
+        GGML_ASSERT(dst->src[6]->ne[0] == H * (1 + 2 * S_v));
+        GGML_ASSERT(ggml_nelements(dst->src[7]) == H * S_v * S_v);
+        txn_log = (float *) dst->src[6]->data;
+        txn_prior = (float *) dst->src[7]->data;
+        txn_slots = dst->src[6]->ne[1];
+        GGML_ASSERT(txn_slots > 0);
+    }
+
     // recurrent state -> gdn_out tail (after attention scores), or the cache when fusing
     float * state_d           = dst_d + S_v * H * n_tokens * n_seqs;
     int64_t state_slot_stride = S_v * S_v * H * n_seqs;
@@ -312,7 +354,7 @@ static void ggml_cuda_op_gated_delta_net_impl(
         } else {
             launch_gated_delta_net<false, false>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
                 S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream);
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, stream, txn_log, txn_prior, txn_slots);
         }
     }
 }
@@ -325,3 +367,81 @@ void ggml_cuda_op_gated_delta_net_fused_cache(
         ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_cuda_gated_delta_net_fused_cache cache) {
     ggml_cuda_op_gated_delta_net_impl(ctx, dst, &cache);
 }
+
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+template <int S>
+__global__ void gdn_replay_cuda(ggml_cuda_gdn_replay_args a) {
+    const int head = blockIdx.x;
+    const int col = blockIdx.y * blockDim.y + threadIdx.y;
+    const int lane = threadIdx.x;
+    constexpr int warp = S < 32 ? S : 32;
+    const int64_t offset = (int64_t) head * S * S + col * S;
+    float state[S / warp];
+#pragma unroll
+    for (int r = 0; r < S / warp; ++r) {
+        state[r] = a.prior[offset + r * warp + lane];
+    }
+    for (int t = 0; t < a.n_keep; ++t) {
+        const float * record = a.log + t * a.head_count * (1 + 2 * S);
+        const float decay = record[head];
+        const float delta = record[a.head_count * (1 + S) + head * S + col];
+#pragma unroll
+        for (int r = 0; r < S / warp; ++r) {
+            const float key = record[a.head_count + head * S + r * warp + lane];
+            if constexpr (S <= 32) {
+                state[r] = __fmaf_rn(key, delta, __fmul_rn(decay, state[r]));
+            } else {
+                state[r] = __fmaf_rn(decay, state[r], __fmul_rn(key, delta));
+            }
+        }
+    }
+#pragma unroll
+    for (int r = 0; r < S / warp; ++r) {
+        a.state[offset + r * warp + lane] = state[r];
+    }
+}
+
+bool ggml_cuda_gdn_replay(const ggml_cuda_gdn_replay_args * args, int32_t count) {
+    if (args == nullptr || count <= 0) {
+        return false;
+    }
+    int device = -1;
+    for (int i = 0; i < count; ++i) {
+        const auto & a = args[i];
+        if (a.head_count <= 0 || a.n_keep < 0 ||
+                (a.state_dim != 16 && a.state_dim != 32 && a.state_dim != 64 && a.state_dim != 128)) {
+            return false;
+        }
+        for (const float * ptr : {a.prior, (const float *) a.state, a.log}) {
+            cudaPointerAttributes attr = {};
+            if (ptr == nullptr || cudaPointerGetAttributes(&attr, ptr) != cudaSuccess) {
+                cudaGetLastError();
+                return false;
+            }
+            if (attr.type != cudaMemoryTypeDevice || (device >= 0 && device != attr.device)) {
+                return false;
+            }
+            device = attr.device;
+        }
+    }
+    const int previous_device = ggml_cuda_get_device();
+    ggml_cuda_set_device(device);
+    // Replay is outside the compute graph and must wait for all state writes.
+    CUDA_CHECK(cudaDeviceSynchronize());
+    for (int i = 0; i < count; ++i) {
+        const auto & a = args[i];
+        const dim3 grid(a.head_count, a.state_dim / 4, 1);
+        const dim3 block(a.state_dim < 32 ? a.state_dim : 32, 4, 1);
+        const ggml_cuda_kernel_launch_params launch(grid, block, 0, nullptr);
+        switch (a.state_dim) {
+            case 16:  ggml_cuda_kernel_launch(gdn_replay_cuda<16>,  launch, a); break;
+            case 32:  ggml_cuda_kernel_launch(gdn_replay_cuda<32>,  launch, a); break;
+            case 64:  ggml_cuda_kernel_launch(gdn_replay_cuda<64>,  launch, a); break;
+            case 128: ggml_cuda_kernel_launch(gdn_replay_cuda<128>, launch, a); break;
+        }
+    }
+    CUDA_CHECK(cudaDeviceSynchronize());
+    ggml_cuda_set_device(previous_device);
+    return true;
+}
+#endif

ggml/src/ggml-cuda/gated_delta_net.cuh

@@ -1,5 +1,8 @@
 #include "common.cuh"
 #include "ggml.h"
+#include "ggml-cuda-gdn-transaction.h"
+
+bool ggml_cuda_gdn_replay(const ggml_cuda_gdn_replay_args * args, int32_t count);
 
 // fused-kernel recurrent-state output; strides in elements (per-seq stride is always D, set in-kernel)
 struct ggml_cuda_gated_delta_net_fused_cache {

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -5777,8 +5777,21 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t
     GGML_UNUSED(reg);
 }
 
+static bool ggml_backend_cuda_gdn_txn_supported(ggml_backend_dev_t dev) {
+    const auto * ctx = (const ggml_backend_cuda_device_context *) dev->context;
+    return ggml_cuda_info().devices[ctx->device].cc == GGML_CUDA_CC_TURING;
+}
+
 static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) {
     GGML_UNUSED(reg);
+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+    if (strcmp(name, "ggml_cuda_gdn_replay") == 0) {
+        return (void *) ggml_cuda_gdn_replay;
+    }
+    if (strcmp(name, "ggml_cuda_gdn_txn_supported") == 0) {
+        return (void *) ggml_backend_cuda_gdn_txn_supported;
+    }
+#endif
     if (strcmp(name, "ggml_backend_comm_init") == 0) {
         return (void *)ggml_backend_cuda_comm_init;
     }

sm75/test-gdn-transaction.cpp

@@ -0,0 +1,89 @@
+#include "ggml.h"
+#include "ggml-backend.h"
+#include "ggml-cuda-gdn-transaction.h"
+
+#include <algorithm>
+#include <cstdio>
+#include <cstring>
+#include <random>
+#include <vector>
+
+int main() {
+    ggml_backend_load_all();
+    auto dev = ggml_backend_dev_by_name("CUDA0");
+    if (!dev) return 2;
+    auto backend = ggml_backend_dev_init(dev, nullptr);
+    auto replay = (ggml_cuda_gdn_replay_t) ggml_backend_reg_get_proc_address(
+        ggml_backend_dev_backend_reg(dev), "ggml_cuda_gdn_replay");
+    if (!backend || !replay) return 2;
+    int checks = 0;
+    for (int dim : {16, 32, 64, 128}) {
+        for (int tokens : {1, 2, 4, 9, 32}) {
+            const int heads = 4;
+            const int slots = 4;
+            const int state_size = heads * dim * dim;
+            auto ctx = ggml_init({4 * 1024 * 1024, nullptr, true});
+            auto q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dim, 2, tokens, 1);
+            auto k = ggml_dup_tensor(ctx, q);
+            auto v = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dim, heads, tokens, 1);
+            auto g = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 1, heads, tokens, 1);
+            auto b = ggml_dup_tensor(ctx, g);
+            auto s = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dim, dim, heads, 1);
+            auto log = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, heads * (1 + 2 * dim), slots);
+            auto prior = ggml_dup_tensor(ctx, s);
+            auto state = ggml_dup_tensor(ctx, s);
+            auto ref = ggml_gated_delta_net(ctx, q, k, v, g, b, s, tokens + 1);
+            auto txn = ggml_gated_delta_net(ctx, q, k, v, g, b, s, 1);
+            txn->src[6] = log;
+            txn->src[7] = prior;
+            auto gf_ref = ggml_new_graph(ctx);
+            auto gf_txn = ggml_new_graph(ctx);
+            ggml_build_forward_expand(gf_ref, ref);
+            ggml_build_forward_expand(gf_txn, txn);
+            auto buffer = ggml_backend_alloc_ctx_tensors(ctx, backend);
+            if (!buffer) return 2;
+            ggml_backend_buffer_clear(buffer, 0);
+            std::mt19937 rng(1234 + tokens);
+            std::uniform_real_distribution<float> dist(-0.1f, 0.1f);
+            std::vector<float> initial;
+            for (auto tensor : {q, k, v, g, b, s}) {
+                std::vector<float> data(ggml_nelements(tensor));
+                for (auto & x : data) x = dist(rng);
+                if (tensor == g) for (auto & x : data) x = -0.1f + x;
+                if (tensor == b) for (auto & x : data) x += 0.5f;
+                if (tensor == s) initial = data;
+                ggml_backend_tensor_set(tensor, data.data(), 0, ggml_nbytes(tensor));
+            }
+            if (ggml_backend_graph_compute(backend, gf_ref) != GGML_STATUS_SUCCESS ||
+                ggml_backend_graph_compute(backend, gf_txn) != GGML_STATUS_SUCCESS) return 2;
+            const int attn = dim * heads * tokens;
+            std::vector<float> expected(ggml_nelements(ref));
+            std::vector<float> actual(ggml_nelements(txn));
+            ggml_backend_tensor_get(ref, expected.data(), 0, ggml_nbytes(ref));
+            ggml_backend_tensor_get(txn, actual.data(), 0, ggml_nbytes(txn));
+            if (memcmp(expected.data(), actual.data(), actual.size() * sizeof(float))) {
+                std::printf("FAIL forward dim=%d tokens=%d\n", dim, tokens);
+                return 1;
+            }
+            const int logged = std::min(tokens, slots);
+            for (int keep = logged; keep >= 0; --keep) {
+                ggml_cuda_gdn_replay_args args = {(const float *) prior->data, (float *) state->data,
+                    (const float *) log->data, heads, dim, keep};
+                if (!replay(&args, 1)) return 2;
+                std::vector<float> got(state_size);
+                ggml_backend_tensor_get(state, got.data(), 0, ggml_nbytes(state));
+                const float * want = keep == 0 && logged == tokens ? initial.data() :
+                    expected.data() + attn + (logged - keep) * state_size;
+                if (memcmp(want, got.data(), state_size * sizeof(float))) {
+                    std::printf("FAIL replay dim=%d tokens=%d keep=%d\n", dim, tokens, keep);
+                    return 1;
+                }
+                ++checks;
+            }
+            ggml_backend_buffer_free(buffer);
+            ggml_free(ctx);
+        }
+    }
+    ggml_backend_free(backend);
+    std::printf("PASS: %d bitwise replay checks and 20 forward checks\n", checks);
+}

sm75/test-mtp-state.cpp

@@ -0,0 +1,85 @@
+#include "llama.h"
+#include "ggml-backend.h"
+
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <string>
+#include <vector>
+
+static void require(bool ok, const char * what) {
+    if (!ok) {
+        std::fprintf(stderr, "FAIL: %s\n", what);
+        std::exit(1);
+    }
+}
+
+static std::vector<uint8_t> save(llama_context * ctx) {
+    std::vector<uint8_t> data(llama_state_seq_get_size(ctx, 0));
+    require(llama_state_seq_get_data(ctx, data.data(), data.size(), 0) == data.size(), "save");
+    return data;
+}
+
+int main(int argc, char ** argv) {
+    if (argc != 3) return 2;
+    ggml_backend_load_all();
+    llama_backend_init();
+    auto mp = llama_model_default_params();
+    mp.n_gpu_layers = 999;
+    auto model = llama_model_load_from_file(argv[1], mp);
+    require(model != nullptr, "model");
+    auto cp = llama_context_default_params();
+    cp.n_ctx = 1024;
+    cp.n_batch = 512;
+    cp.n_ubatch = 512;
+    cp.n_seq_max = 1;
+    cp.n_rs_seq = 3;
+    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
+    auto ctx = llama_init_from_model(model, cp);
+    require(ctx != nullptr, "context");
+    auto mem = llama_get_memory(ctx);
+    auto batch = llama_batch_init(512, 0, 1);
+    const int vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
+    auto decode = [&](int start, int count) {
+        batch.n_tokens = count;
+        for (int i = 0; i < count; ++i) {
+            batch.token[i] = 100 + (start + i) % 1000;
+            batch.pos[i] = start + i;
+            batch.n_seq_id[i] = 1;
+            batch.seq_id[i][0] = 0;
+            batch.logits[i] = i == count - 1;
+        }
+        require(llama_decode(ctx, batch) == 0, "decode");
+    };
+    auto record = [&](const std::string & name, const std::vector<uint8_t> & data) {
+        std::ofstream file(std::string(argv[2]) + "/" + name, std::ios::binary);
+        file.write((const char *) data.data(), data.size());
+        require(file.good(), "write");
+    };
+    decode(0, 32);
+    auto base = save(ctx);
+    record("base.bin", base);
+    for (int keep = 1; keep <= 4; ++keep) {
+        require(llama_state_seq_set_data(ctx, base.data(), base.size(), 0) == base.size(), "restore base");
+        decode(32, 4);
+        if (keep < 4) require(llama_memory_seq_rm(mem, 0, 32 + keep, -1), "rollback");
+        auto cut = save(ctx);
+        record("cut-" + std::to_string(keep) + ".bin", cut);
+        if (keep < 4) require(!llama_memory_seq_rm(mem, 0, 32 + keep - 1, -1), "reject second pending rollback");
+        decode(32 + keep, 1);
+        std::vector<float> logits(llama_get_logits_ith(ctx, -1), llama_get_logits_ith(ctx, -1) + vocab);
+        require(llama_state_seq_set_data(ctx, cut.data(), cut.size(), 0) == cut.size(), "restore cut");
+        decode(32 + keep, 1);
+        require(memcmp(logits.data(), llama_get_logits_ith(ctx, -1), vocab * sizeof(float)) == 0, "save/load logits parity");
+        record("after-" + std::to_string(keep) + ".bin", save(ctx));
+    }
+    require(llama_memory_seq_rm(mem, 0, 0, -1), "clear sequence");
+    decode(0, 32);
+    require(save(ctx) == base, "clear and restart");
+    llama_batch_free(batch);
+    llama_free(ctx);
+    llama_model_free(model);
+    llama_backend_free();
+    std::puts("PASS: accepted-prefix rollback, pending rollback rejection, state save/load, clear and restart");
+}

src/llama-graph.cpp

@@ -340,6 +340,9 @@ void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {
         // assuming copy destinations ALWAYS happen ONLY on the cells between head and head+n
         for (uint32_t i = 0; i < n_rs; ++i) {
             data[i] = mctx->s_copy(i);
+            if (txn_copy) {
+                ((int32_t *) txn_copy->data)[i] = data[i] % mctx->get_size();
+            }
         }
     }
 }
@@ -3589,6 +3592,10 @@ static std::unique_ptr<llm_graph_input_rs> build_rs_inp_impl(
     inp->s_copy = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs);
     ggml_set_input(inp->s_copy);
     ggml_set_name(inp->s_copy, "rs_s_copy");
+    if (mctx_cur->txn_enabled()) {
+        inp->txn_copy = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs);
+        ggml_set_input(inp->txn_copy);
+    }
 
     inp->s_copy_main  = ggml_view_1d(ctx0, inp->s_copy, n_seqs, 0);
     inp->s_copy_extra = ggml_view_1d(ctx0, inp->s_copy, n_rs - n_seqs, n_seqs * inp->s_copy->nb[0]);

src/llama-graph.h

@@ -269,6 +269,7 @@ public:
     bool can_reuse(const llm_graph_params & params) override;
 
     ggml_tensor * s_copy;  // I32 [n_rs]
+    ggml_tensor * txn_copy = nullptr; // S always uses the active row after replay
 
     // views of s_copy, computed once per graph
     // and shared across layers which use build_rs

src/llama-memory-recurrent.cpp

@@ -1,6 +1,7 @@
 #include "llama-memory-recurrent.h"
 
 #include "ggml-backend.h"
+#include "ggml-cuda-gdn-transaction.h"
 #include "llama-impl.h"
 #include "llama-io.h"
 #include "llama-batch.h"
@@ -8,6 +9,7 @@
 
 #include <algorithm>
 #include <cassert>
+#include <cstdlib>
 #include <cstring>
 #include <limits>
 #include <map>
@@ -35,6 +37,26 @@ llama_memory_recurrent::llama_memory_recurrent(
     this->n_rs_seq = n_rs_seq;
     rs_idx.assign(n_seq_max, 0);
 
+    const char * txn_env = std::getenv("LLAMA_SM75_GDN_TXN");
+    use_txn = txn_env && strcmp(txn_env, "1") == 0 && model.arch == LLM_ARCH_QWEN35 &&
+        offload && mem_size == 1 && n_seq_max == 1 && n_rs_seq > 1 && type_s == GGML_TYPE_F32 &&
+        hparams.ssm_d_state == 128;
+    ggml_backend_dev_t txn_device = nullptr;
+    for (int i = 0; use_txn && i < n_layer; ++i) {
+        if (filter && !filter(i)) {
+            continue;
+        }
+        auto * dev = model.dev_layer(i);
+        auto reg = ggml_backend_dev_backend_reg(dev);
+        auto supported = (bool (*)(ggml_backend_dev_t)) ggml_backend_reg_get_proc_address(reg, "ggml_cuda_gdn_txn_supported");
+        use_txn = supported && supported(dev) && ggml_backend_reg_get_proc_address(reg, "ggml_cuda_gdn_replay") != nullptr &&
+            (!txn_device || txn_device == dev);
+        txn_device = dev;
+    }
+    if (txn_env && strcmp(txn_env, "1") == 0) {
+        LLAMA_LOG_INFO("%s: GDN transaction %s\n", __func__, use_txn ? "enabled" : "unsupported, using snapshots");
+    }
+
     cells.clear();
     cells.resize(mem_size);
 
@@ -52,7 +74,7 @@ llama_memory_recurrent::llama_memory_recurrent(
         if (it == ctx_map.end()) {
             ggml_init_params params = {
                 // r and s per layer, plus the separate PLE conv row where the model has one
-                /*.mem_size   =*/ size_t((hparams.ple_conv_state() > 0 ? 3u : 2u)*n_layer*ggml_tensor_overhead()),
+                /*.mem_size   =*/ size_t((hparams.ple_conv_state() > 0 || use_txn ? 3u : 2u)*n_layer*ggml_tensor_overhead()),
                 /*.mem_buffer =*/ NULL,
                 /*.no_alloc   =*/ true,
             };
@@ -72,6 +94,7 @@ llama_memory_recurrent::llama_memory_recurrent(
 
     r_l.resize(n_layer);
     s_l.resize(n_layer);
+    txn_l.resize(n_layer);
     p_l.resize(n_layer);
 
     for (int i = 0; i < n_layer; i++) {
@@ -100,12 +123,18 @@ llama_memory_recurrent::llama_memory_recurrent(
 
         const uint32_t n_rows = mem_size * (1 + n_rs_seq);
         ggml_tensor * r = ggml_new_tensor_2d(ctx, type_r, hparams.n_embd_r(), n_rows);
-        ggml_tensor * s = ggml_new_tensor_2d(ctx, type_s, hparams.n_embd_s(), n_rows);
+        ggml_tensor * s = ggml_new_tensor_2d(ctx, type_s, hparams.n_embd_s(), use_txn ? 2 : n_rows);
         ggml_format_name(r, "cache_r_l%d", i);
         ggml_format_name(s, "cache_s_l%d", i);
         r_l[i] = r;
         s_l[i] = s;
 
+        if (use_txn) {
+            const int64_t heads = hparams.ssm_d_inner / hparams.ssm_d_state;
+            txn_l[i] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, heads * (1 + 2 * hparams.ssm_d_state), n_rs_seq + 1);
+            ggml_format_name(txn_l[i], "cache_txn_l%d", i);
+        }
+
         // the PLE history needs its own row: Meta must mirror it while the delta-net conv state next door stays split
         if (hparams.ple_conv_state() > 0 && hparams.is_ple(i)) {
             ggml_tensor * p = ggml_new_tensor_2d(ctx, type_r, hparams.ple_conv_state(), n_rows);
@@ -130,6 +159,14 @@ llama_memory_recurrent::llama_memory_recurrent(
         const size_t memory_size_s = size_s_bytes();
         const size_t memory_size_p = size_p_bytes();
 
+        if (use_txn) {
+            size_t log_bytes = 0;
+            for (const auto * log : txn_l) {
+                if (log) log_bytes += ggml_nbytes(log);
+            }
+            LLAMA_LOG_INFO("%s: GDN transaction log = %.2f MiB\n", __func__, log_bytes / 1024.0 / 1024.0);
+        }
+
         LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB, P (%s): %7.2f MiB\n", __func__,
                 (float)(memory_size_r + memory_size_s + memory_size_p) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq,
                 ggml_type_name(type_r), (float)memory_size_r / (1024.0f * 1024.0f),
@@ -139,6 +176,7 @@ llama_memory_recurrent::llama_memory_recurrent(
 }
 
 void llama_memory_recurrent::clear(bool data) {
+    txn_tokens = 0;
     for (int32_t i = 0; i < (int32_t) size; ++i) {
         cells[i].pos = -1;
         cells[i].seq_id.clear();
@@ -177,6 +215,7 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
     const bool rm_all = p0 == 0 && p1 == std::numeric_limits<llama_pos>::max();
     if (rm_all) {
         set_rs_idx(seq_id, 0);
+        txn_tokens = 0;
     }
 
     // models like Mamba or RWKV can't have a state partially erased at the end
@@ -196,6 +235,9 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos
                 // pending rollback is single-use
                 const bool pending = rs_idx[seq_id] != 0;
                 if (!pending && rollback >= 1 && rollback <= (llama_pos) n_rs_seq) {
+                    if (use_txn && !txn_rollback(rollback)) {
+                        return false;
+                    }
                     set_rs_idx(seq_id, (uint32_t) rollback);
                     cell.pos = p0 - 1;
                     return true;
@@ -846,6 +888,7 @@ void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq
 
 void llama_memory_recurrent::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
     GGML_UNUSED(flags);
+    txn_tokens = 0;
 
     uint32_t cell_count;
     io.read(&cell_count, sizeof(cell_count));
@@ -953,7 +996,7 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std::
             for (const auto & range : cell_ranges) {
                 const size_t range_size = range.second - range.first;
                 const size_t buf_size = range_size * s_size_row;
-                io.write_tensor(s_l[il], range.first * s_size_row, buf_size);
+                io.write_tensor(s_l[il], (use_txn ? 0 : range.first) * s_size_row, buf_size);
             }
         }
     } else {
@@ -1261,6 +1304,9 @@ bool llama_memory_recurrent_context::apply() {
     }
 
     mem->find_slot(ubatches[i_next]);
+    if (mem->txn_enabled()) {
+        mem->txn_batch(ubatches[i_next].n_tokens);
+    }
 
     return true;
 }
@@ -1299,6 +1345,30 @@ ggml_tensor * llama_memory_recurrent_context::get_s_l(int32_t il) const {
     return mem->s_l[il];
 }
 
+ggml_tensor * llama_memory_recurrent_context::get_txn_l(int32_t il) const {
+    return mem->txn_l[il];
+}
+
+bool llama_memory_recurrent::txn_rollback(uint32_t rollback) {
+    if (rollback > txn_tokens) {
+        return false;
+    }
+    std::vector<ggml_cuda_gdn_replay_args> args;
+    ggml_backend_reg_t reg = nullptr;
+    for (size_t il = 0; il < s_l.size(); ++il) {
+        if (!s_l[il]) {
+            continue;
+        }
+        const auto buft = ggml_backend_buffer_get_type(s_l[il]->buffer);
+        reg = ggml_backend_dev_backend_reg(ggml_backend_buft_get_device(buft));
+        float * state = (float *) s_l[il]->data;
+        args.push_back({state + hparams.n_embd_s(), state, (const float *) txn_l[il]->data,
+            hparams.ssm_d_inner / hparams.ssm_d_state, (int32_t) hparams.ssm_d_state, (int32_t) (txn_tokens - rollback)});
+    }
+    auto replay = reg ? (ggml_cuda_gdn_replay_t) ggml_backend_reg_get_proc_address(reg, "ggml_cuda_gdn_replay") : nullptr;
+    return replay && replay(args.data(), args.size());
+}
+
 ggml_tensor * llama_memory_recurrent_context::get_p_l(int32_t il) const {
     return mem->p_l[il];
 }

src/llama-memory-recurrent.h

@@ -4,6 +4,7 @@
 #include "llama-graph.h"
 #include "llama-memory.h"
 
+#include <algorithm>
 #include <map>
 #include <set>
 #include <vector>
@@ -78,6 +79,9 @@ public:
 
     void set_rs_idx(llama_seq_id seq_id, uint32_t idx);
 
+    bool txn_enabled() const { return use_txn; }
+    void txn_batch(uint32_t n_tokens) { txn_tokens = std::min(n_tokens, n_rs_seq + 1); }
+
     // computed before each graph build
     uint32_t n = 0;
 
@@ -111,6 +115,7 @@ public:
     // per layer
     std::vector<ggml_tensor *> r_l;
     std::vector<ggml_tensor *> s_l;
+    std::vector<ggml_tensor *> txn_l;
     // a second conv history that must stay replicated across devices, so it cannot share the r row
     std::vector<ggml_tensor *> p_l;
 
@@ -120,6 +125,10 @@ private:
 
     const uint32_t n_seq_max = 1;
 
+    bool use_txn = false;
+    uint32_t txn_tokens = 0;
+    bool txn_rollback(uint32_t rollback);
+
     // ggml contexts for the KV cache along with the allocated backend buffers:
     std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;
 
@@ -173,6 +182,8 @@ public:
 
     ggml_tensor * get_r_l(int32_t il) const;
     ggml_tensor * get_s_l(int32_t il) const;
+    ggml_tensor * get_txn_l(int32_t il) const;
+    bool txn_enabled() const { return mem->txn_enabled(); }
     ggml_tensor * get_p_l(int32_t il) const;
 
     int32_t s_copy(int i) const;

src/models/delta-net-base.cpp

@@ -561,10 +561,15 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn(
     }
 
     const int64_t D = S_v * S_v * H_v;
-    const int64_t K = cparams.n_rs_seq + 1;
+    const bool txn = mctx_cur->txn_enabled();
+    const int64_t K = txn ? 1 : cparams.n_rs_seq + 1;
 
     // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output.
     ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K);
+    if (txn) {
+        gdn_out->src[6] = mctx_cur->get_txn_l(il);
+        gdn_out->src[7] = ggml_view_1d(ctx0, ssm_states_all, D, D * sizeof(float));
+    }
     if (n_seq_tokens > 1) {
         res->add_fused_node({LLM_FUSED_OP_GDN_CH, gdn_out, il});
     } else {

src/models/qwen35.cpp

@@ -384,7 +384,14 @@ 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);
+    ggml_tensor * state;
+    if (mctx_cur->txn_enabled()) {
+        GGML_ASSERT(n_seqs == 1 && mctx_cur->get_size() == 1);
+        state = build_rs(ssm_states_all, inp->txn_copy, inp->s_copy_extra, hparams.n_embd_s(), n_seqs,
+            mctx_cur->get_n_rs(), mctx_cur->get_head(), mctx_cur->get_size(), mctx_cur->get_rs_z(), ggml_get_rows);
+    } 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);