llama: GDN rollback by replay, one-graph MTP draft chain, K/V-only MTP prompt pass, sync trim (port of Torchit1/arc-b580)

ggml/include/ggml-rpc.h

@@ -8,10 +8,10 @@ extern "C" {
 
 #define RPC_PROTO_MAJOR_VERSION    5
 #define RPC_PROTO_MINOR_VERSION    1
-#define RPC_PROTO_PATCH_VERSION    0
+#define RPC_PROTO_PATCH_VERSION    1
 
 #ifdef  __cplusplus
-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
+static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
 #endif
 
 #define GGML_RPC_MAX_SERVERS       16

ggml/include/ggml.h

@@ -596,6 +596,8 @@ extern "C" {
 
         GGML_OP_GLU,
 
+        GGML_OP_GATED_DELTA_NET_REPLAY, // ARC-LAB: appended last so the existing op numbers do not move
+
         GGML_OP_COUNT,
     };
 
@@ -2628,6 +2630,34 @@ extern "C" {
             struct ggml_tensor  * dt_bias,
             struct ggml_tensor  * a);
 
+    // ARC-LAB: gated delta net with rollback by replay (independent implementation of the technique described by
+    // llamADALL). Same recurrence as ggml_gated_delta_net (scalar gate only), but instead of K state snapshots the op
+    // keeps ONE base state plus, for each later token, its update record: the per-token update is
+    // S <- g * S + k (x) delta, so a record is (delta [S_v], k [S_v], g [1]) per head, R = H*(2*S_v + 1) floats:
+    //   record = [ delta: S_v*H | k: S_v*H (head h holds k of head h % H_k) | g: H ]   (g = the decay factor exp(gate))
+    //
+    //   q, k, v, g, beta : as ggml_gated_delta_net (g: [1, H, n_tokens, n_seqs])
+    //   state : [S_v, S_v, H, n_seqs] stored base state
+    //   rec   : F32 [R*(K-1), 2, n_seqs] stored records, two banks per seq (a pass reads one bank, the caller
+    //           stores the new records in the other one)
+    //   ctrl  : I32 [2, n_seqs] = {n_rec, bank}: records 0..n_rec-1 of `bank` are first re-applied to state
+    //
+    // output, packed: [ attn scores S_v*H*n_tokens*n_seqs | base D*n_seqs | records (K-1)*R*n_seqs ], D = S_v*S_v*H
+    //   with m = min(n_tokens - 1, K - 1): base = state after token n_tokens-1-m (the oldest state a rollback of up
+    //   to m tokens can need), record j = token n_tokens-m+j for j < m; records m..K-2 are left unwritten.
+    //   Rolling back r <= m tokens = re-applying m - r records to base on the next call.
+    GGML_API struct ggml_tensor * ggml_gated_delta_net_replay(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * q,
+            struct ggml_tensor  * k,
+            struct ggml_tensor  * v,
+            struct ggml_tensor  * g,
+            struct ggml_tensor  * beta,
+            struct ggml_tensor  * state,
+            struct ggml_tensor  * rec,
+            struct ggml_tensor  * ctrl,
+            int                   K);
+
     // DSA lightning indexer
     //
     // q:       [n_embd_idx, n_head_idx, n_batch, ne3 ]

ggml/src/ggml-backend-impl.h

@@ -270,6 +270,11 @@ extern "C" {
 #    define GGML_BACKEND_DL_SCORE_IMPL(score_fn)
 #endif
 
+    // ARC-LAB: sync/upload attribution (GGML_SYCL_SYNC_TRACE): the scheduler tags the blocking copies it makes with a
+    // cause (e.g. "split input <name>"); a backend that counts synchronizes / uploads reads it. Thread-local; null = none.
+    GGML_API void         ggml_sync_trace_set_cause(const char * cause);
+    GGML_API const char * ggml_sync_trace_get_cause(void);
+
 #ifdef  __cplusplus
 }
 #endif

ggml/src/ggml-backend.cpp

@@ -1591,6 +1591,26 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
     return true;
 }
 
+// ARC-LAB: see ggml-backend-impl.h
+static thread_local char g_sync_trace_cause[128];
+static thread_local bool g_sync_trace_has_cause = false;
+
+void ggml_sync_trace_set_cause(const char * cause) {
+    g_sync_trace_has_cause = cause != nullptr;
+    if (cause) {
+        snprintf(g_sync_trace_cause, sizeof(g_sync_trace_cause), "%s", cause);
+    }
+}
+
+const char * ggml_sync_trace_get_cause(void) {
+    return g_sync_trace_has_cause ? g_sync_trace_cause : nullptr;
+}
+
+static bool ggml_sync_trace_on() {
+    static const bool on = getenv("GGML_SYCL_SYNC_TRACE") && atoi(getenv("GGML_SYCL_SYNC_TRACE")) != 0;
+    return on;
+}
+
 static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
     GGML_ASSERT(sched);
     struct ggml_backend_sched_split * splits = sched->splits;
@@ -1609,6 +1629,9 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
         // ensure the previous split's async work has completed before we start
         // this split, the allocator may have reused buffer regions across splits
         if (split->n_inputs == 0 && prev_backend_id >= 0 && prev_backend_id != split_backend_id) {
+            if (ggml_sync_trace_on()) {
+                ggml_sync_trace_set_cause("sched: backend switch");
+            }
             if (sched->events[prev_backend_id][sched->cur_copy] != NULL) {
                 ggml_backend_event_synchronize(sched->events[prev_backend_id][sched->cur_copy]);
             } else {
@@ -1621,6 +1644,12 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
             ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]);
             struct ggml_tensor * input = split->inputs[input_id];
             struct ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);
+            if (ggml_sync_trace_on()) {
+                char cause[128];
+                snprintf(cause, sizeof(cause), "sched: split input %s%s", input->flags & GGML_TENSOR_FLAG_INPUT ? "" : "(computed) ",
+                         input->name);
+                ggml_sync_trace_set_cause(cause);
+            }
 
             if (input->flags & GGML_TENSOR_FLAG_INPUT) {
                 // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
@@ -1739,6 +1768,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
             }
         }
 
+        if (ggml_sync_trace_on()) {
+            ggml_sync_trace_set_cause(nullptr);
+        }
+
         if (!sched->callback_eval) {
             enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph);
             if (ec != GGML_STATUS_SUCCESS) {

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2112,6 +2112,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
             {
                 ggml_compute_forward_gated_delta_net(params, tensor);
             } break;
+        case GGML_OP_GATED_DELTA_NET_REPLAY:
+            {
+                ggml_compute_forward_gated_delta_net_replay(params, tensor);
+            } break;
         case GGML_OP_LIGHTNING_INDEXER:
             {
                 ggml_compute_forward_lightning_indexer(params, tensor);
@@ -2308,6 +2312,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_COUNT_EQUAL:
         case GGML_OP_SOLVE_TRI:
         case GGML_OP_GATED_DELTA_NET:
+        case GGML_OP_GATED_DELTA_NET_REPLAY:
         case GGML_OP_DSV4_HC_COMB:
         case GGML_OP_DSV4_HC_PRE:
         case GGML_OP_DSV4_HC_POST:
@@ -3045,6 +3050,12 @@ struct ggml_cplan ggml_graph_plan(
                         const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0);
                         cur = per_thread * sizeof(float) * n_tasks;
                     } break;
+                case GGML_OP_GATED_DELTA_NET_REPLAY:
+                    {
+                        // per thread: working state [S_v, S_v] + delta [S_v]
+                        const int64_t S_v = node->src[2]->ne[0];
+                        cur = (S_v * S_v + S_v) * sizeof(float) * n_tasks;
+                    } break;
                 case GGML_OP_COUNT:
                     {
                         GGML_ABORT("fatal error");

ggml/src/ggml-cpu/ops.cpp

@@ -11171,6 +11171,151 @@ void ggml_compute_forward_gated_delta_net(
     }
 }
 
+// ARC-LAB: ggml_compute_forward_gated_delta_net_replay (rollback by replay, see ggml_gated_delta_net_replay).
+// The per-token step is the scalar-gate path of ggml_compute_forward_gated_delta_net_one_chunk, operation for
+// operation, and a replayed record repeats exactly its state update (scale by g, then mad k * delta), so the replayed
+// state is bit-identical to the state that path computes for the same token.
+static void ggml_compute_forward_gated_delta_net_replay_f32(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    const ggml_tensor * src_q     = dst->src[0];
+    const ggml_tensor * src_k     = dst->src[1];
+    const ggml_tensor * src_v     = dst->src[2];
+    const ggml_tensor * src_g     = dst->src[3];
+    const ggml_tensor * src_beta  = dst->src[4];
+    const ggml_tensor * src_state = dst->src[5];
+    const ggml_tensor * src_rec   = dst->src[6];
+    const ggml_tensor * src_ctrl  = dst->src[7];
+
+    GGML_TENSOR_LOCALS(int64_t, neq, src_q, ne);
+    GGML_TENSOR_LOCALS(size_t,  nbq, src_q, nb);
+    GGML_TENSOR_LOCALS(int64_t, nek, src_k, ne);
+    GGML_TENSOR_LOCALS(size_t,  nbk, src_k, nb);
+    GGML_TENSOR_LOCALS(int64_t, nev, src_v, ne);
+    GGML_TENSOR_LOCALS(size_t,  nbv, src_v, nb);
+    GGML_TENSOR_LOCALS(size_t,  nbg, src_g, nb);
+    GGML_TENSOR_LOCALS(size_t,  nbb, src_beta, nb);
+
+    const int64_t S_v      = nev0;
+    const int64_t H        = nev1;
+    const int64_t n_tokens = nev2;
+    const int64_t n_seqs   = nev3;
+    const int64_t K        = ggml_get_op_params_i32(dst, 0);
+    const int64_t R        = H * (2 * S_v + 1);
+    const int64_t D        = S_v * S_v * H;
+
+    // base = state after token t_base; tokens after it leave a record
+    const int64_t m      = std::min<int64_t>(n_tokens - 1, K - 1);
+    const int64_t t_base = n_tokens - 1 - m;
+
+    float * attn_out_base = (float *) dst->data;
+    float * base_out      = attn_out_base + S_v * H * n_tokens * n_seqs;
+    float * rec_out       = base_out + D * n_seqs;
+
+    const int32_t * ctrl = (const int32_t *) src_ctrl->data;
+
+    const int64_t rq3 = nev3 / neq3;
+    const int64_t rk3 = nev3 / nek3;
+
+    const float scale = 1.0f / sqrtf((float) S_v);
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+
+    float * s     = (float *) params->wdata + ith * (S_v * S_v + S_v + CACHE_LINE_SIZE_F32);
+    float * delta = s + S_v * S_v;
+
+    const int64_t nr  = H * n_seqs;
+    const int64_t dr  = (nr + nth - 1) / nth;
+    const int64_t ir0 = dr * ith;
+    const int64_t ir1 = std::min(ir0 + dr, nr);
+
+    for (int64_t ir = ir0; ir < ir1; ++ir) {
+        const int64_t iv1 = ir % H; // head
+        const int64_t iv3 = ir / H; // sequence
+
+        const int64_t iq1 = iv1 % neq1;
+        const int64_t ik1 = iv1 % nek1;
+        const int64_t iq3 = iv3 / rq3;
+        const int64_t ik3 = iv3 / rk3;
+
+        // state layout: s[j*S_v + i] = S[i][j] (row j = value column j)
+        memcpy(s, (const char *) src_state->data + iv3 * src_state->nb[3] + iv1 * S_v * S_v * sizeof(float),
+               S_v * S_v * sizeof(float));
+
+        // replay: re-apply the stored updates of the tokens kept since the base state
+        const int32_t n_rec = ctrl[2 * iv3 + 0];
+        const int32_t bank  = ctrl[2 * iv3 + 1];
+        GGML_ASSERT(n_rec >= 0 && n_rec <= K - 1 && (bank == 0 || bank == 1));
+        const float * rec_in = (const float *) ((const char *) src_rec->data + iv3 * src_rec->nb[2] + bank * src_rec->nb[1]);
+        for (int32_t j = 0; j < n_rec; ++j) {
+            const float * r  = rec_in + j * R;
+            const float * rd = r + iv1 * S_v;
+            const float * rk = r + H * S_v + iv1 * S_v;
+            ggml_vec_scale_f32(S_v * S_v, s, r[2 * H * S_v + iv1]);
+            for (int64_t jj = 0; jj < S_v; ++jj) {
+                ggml_vec_mad_f32(S_v, &s[jj * S_v], rk, rd[jj]);
+            }
+        }
+
+        float * attn_data = attn_out_base + (iv3 * n_tokens * H + iv1) * S_v;
+
+        for (int64_t t = 0; t < n_tokens; t++) {
+            const float * q_d = (const float *)((const char *)src_q->data + iq3 * nbq3 + t * nbq2 + iq1 * nbq1);
+            const float * k_d = (const float *)((const char *)src_k->data + ik3 * nbk3 + t * nbk2 + ik1 * nbk1);
+            const float * v_d = (const float *)((const char *)src_v->data + iv3 * nbv3 + t * nbv2 + iv1 * nbv1);
+
+            const float beta_val = *(const float *)((const char *)src_beta->data + iv3 * nbb3 + t * nbb2 + iv1 * nbb1);
+            const float g0       = *(const float *)((const char *)src_g->data    + iv3 * nbg3 + t * nbg2 + iv1 * nbg1);
+            const float g_val    = expf(g0);
+
+            ggml_vec_scale_f32(S_v * S_v, s, g_val);
+
+            for (int64_t j = 0; j < S_v; ++j) {
+                float sum = 0.0f;
+                ggml_vec_dot_f32(S_v, &sum, 0, &s[j * S_v], 0, k_d, 0, 1);
+                delta[j] = (v_d[j] - sum) * beta_val;
+            }
+
+            for (int64_t j = 0; j < S_v; ++j) {
+                ggml_vec_mad_f32(S_v, &s[j * S_v], k_d, delta[j]);
+            }
+
+            for (int64_t j = 0; j < S_v; ++j) {
+                float sum = 0.0f;
+                ggml_vec_dot_f32(S_v, &sum, 0, &s[j * S_v], 0, q_d, 0, 1);
+                attn_data[j] = sum * scale;
+            }
+
+            attn_data += S_v * H;
+
+            if (t == t_base) {
+                memcpy(base_out + (iv3 * H + iv1) * S_v * S_v, s, S_v * S_v * sizeof(float));
+            } else if (t > t_base) {
+                float * r = rec_out + iv3 * (K - 1) * R + (t - t_base - 1) * R;
+                memcpy(r + iv1 * S_v,           delta, S_v * sizeof(float));
+                memcpy(r + H * S_v + iv1 * S_v, k_d,   S_v * sizeof(float));
+                r[2 * H * S_v + iv1] = g_val;
+            }
+        }
+    }
+}
+
+void ggml_compute_forward_gated_delta_net_replay(
+        const ggml_compute_params * params,
+        ggml_tensor * dst) {
+    switch (dst->src[0]->type) {
+        case GGML_TYPE_F32:
+            {
+                ggml_compute_forward_gated_delta_net_replay_f32(params, dst);
+            } break;
+        default:
+            {
+                GGML_ABORT("fatal error");
+            }
+    }
+}
+
 
 // ggml_compute_forward_dsv4_hc_comb
 

ggml/src/ggml-cpu/ops.h

@@ -105,6 +105,7 @@ void ggml_compute_forward_rwkv_wkv7(const struct ggml_compute_params * params, s
 void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_gated_delta_net_replay(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);

ggml/src/ggml-sycl/dequantize.hpp

@@ -67,6 +67,19 @@ static __dpct_inline__ int ptq1_0_trit(const uint8_t * qs, const uint8_t * qh, c
     return (int) ((v * 3) >> 8);
 }
 
+// ARC-LAB: f32 output whatever GGML_SYCL_F16 says, for GET_ROWS (MTP chain draft-token embeddings). (trit - 1) * d
+// with d widened from fp16 is exact in f32 (the product of 0 / +-1 and d), so the rows match the CPU's
+// dequantize_row_ptq1_0 bit for bit.
+static __dpct_inline__ void dequantize_ptq1_0_f32(const void *vx, const int64_t ib,
+                                                  const int iqs, float &v0, float &v1) {
+    const block_ptq1_0 * x = (const block_ptq1_0 *) vx;
+
+    const float d = (float) x[ib].d;
+
+    v0 = (float) (ptq1_0_trit(x[ib].qs, x[ib].qh, iqs + 0) - 1) * d;
+    v1 = (float) (ptq1_0_trit(x[ib].qs, x[ib].qh, iqs + 1) - 1) * d;
+}
+
 static __dpct_inline__ void dequantize_q4_0(const void *vx, const int64_t ib,
                                             const int iqs, dfloat2 &v) {
     const block_q4_0 * x = (const block_q4_0 *) vx;