From 3a90a98e7a60a77c61dd4e3d2675fd8b3024e08f Mon Sep 17 00:00:00 2001 From: xrip Date: Wed, 9 Sep 2026 11:36:37 +0300 Subject: [PATCH] llama: add opt-in SM75 recurrent state replay --- ggml/include/ggml-cuda-gdn-transaction.h | 14 +++ ggml/src/ggml-cpu/ggml-cpu.cpp | 3 + ggml/src/ggml-cuda/gated_delta_net.cu | 134 +++++++++++++++++++++-- ggml/src/ggml-cuda/gated_delta_net.cuh | 3 + ggml/src/ggml-cuda/ggml-cuda.cu | 13 +++ sm75/test-gdn-transaction.cpp | 89 +++++++++++++++ sm75/test-mtp-state.cpp | 85 ++++++++++++++ src/llama-graph.cpp | 7 ++ src/llama-graph.h | 1 + src/llama-memory-recurrent.cpp | 76 ++++++++++++- src/llama-memory-recurrent.h | 11 ++ src/models/delta-net-base.cpp | 7 +- src/models/qwen35.cpp | 9 +- 13 files changed, 440 insertions(+), 12 deletions(-) create mode 100644 ggml/include/ggml-cuda-gdn-transaction.h create mode 100644 sm75/test-gdn-transaction.cpp create mode 100644 sm75/test-mtp-state.cpp diff --git a/ggml/include/ggml-cuda-gdn-transaction.h b/ggml/include/ggml-cuda-gdn-transaction.h new file mode 100644 index 000000000..195b3bbc7 --- /dev/null +++ b/ggml/include/ggml-cuda-gdn-transaction.h @@ -0,0 +1,14 @@ +#pragma once + +#include + +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); diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 1df0f2bb9..2d6121601 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/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]; diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 1b431a724..6498c6a14 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/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(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 +__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 diff --git a/ggml/src/ggml-cuda/gated_delta_net.cuh b/ggml/src/ggml-cuda/gated_delta_net.cuh index f9bf43706..be801c6fb 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cuh +++ b/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 { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 8a8bf1af4..87d986938 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/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; } diff --git a/sm75/test-gdn-transaction.cpp b/sm75/test-gdn-transaction.cpp new file mode 100644 index 000000000..886a2dd2d --- /dev/null +++ b/sm75/test-gdn-transaction.cpp @@ -0,0 +1,89 @@ +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cuda-gdn-transaction.h" + +#include +#include +#include +#include +#include + +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 dist(-0.1f, 0.1f); + std::vector initial; + for (auto tensor : {q, k, v, g, b, s}) { + std::vector 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 expected(ggml_nelements(ref)); + std::vector 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 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); +} diff --git a/sm75/test-mtp-state.cpp b/sm75/test-mtp-state.cpp new file mode 100644 index 000000000..bb02238be --- /dev/null +++ b/sm75/test-mtp-state.cpp @@ -0,0 +1,85 @@ +#include "llama.h" +#include "ggml-backend.h" + +#include +#include +#include +#include +#include +#include + +static void require(bool ok, const char * what) { + if (!ok) { + std::fprintf(stderr, "FAIL: %s\n", what); + std::exit(1); + } +} + +static std::vector save(llama_context * ctx) { + std::vector 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 & 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 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"); +} diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index d86890745..f8082a670 100644 --- a/src/llama-graph.cpp +++ b/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 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]); diff --git a/src/llama-graph.h b/src/llama-graph.h index cc4110639..d9bf3d954 100644 --- a/src/llama-graph.h +++ b/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 diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 57919accf..428e22256 100644 --- a/src/llama-memory-recurrent.cpp +++ b/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 #include +#include #include #include #include @@ -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::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 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]; } diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 4abb3f5cf..f27dd84f5 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -4,6 +4,7 @@ #include "llama-graph.h" #include "llama-memory.h" +#include #include #include #include @@ -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 r_l; std::vector s_l; + std::vector txn_l; // a second conv history that must stay replicated across devices, so it cannot share the r row std::vector 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> 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; diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ad6612647..227cc30d6 100644 --- a/src/models/delta-net-base.cpp +++ b/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 { diff --git a/src/models/qwen35.cpp b/src/models/qwen35.cpp index a1e263500..72506d515 100644 --- a/src/models/qwen35.cpp +++ b/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);