delta_net: fill the conv slot index on the host

The rollback slot index of the conv windows is a constant per graph
shape. It is now an I32 input of the recurrent memory filled on the
host, instead of an arange and an argsort computed on the device, and
the meta backend no longer needs a mirrored split for ARANGE.

ggml/src/ggml-backend-meta.cpp

@@ -700,6 +700,26 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
             }
             GGML_ABORT("fatal error");
         }
+        // permuted view: split the view dim that spans the src split axis, other dims must be inside one element of it or outside the span
+        if (ggml_is_permuted(tensor) && !ggml_is_permuted(tensor->src[0]) && axis >= 0 && axis < GGML_MAX_DIMS-1) {
+            const size_t nb_axis = tensor->src[0]->nb[axis];
+            const size_t nb_next = tensor->src[0]->nb[axis+1];
+            for (int dim = 0; dim < GGML_MAX_DIMS; dim++) {
+                if (tensor->ne[dim] == 1 || tensor->nb[dim] < nb_axis || tensor->nb[dim]*tensor->ne[dim] != nb_next) {
+                    continue;
+                }
+                bool others_ok = true;
+                for (int i = 0; i < GGML_MAX_DIMS; i++) {
+                    if (i == dim || tensor->ne[i] == 1) {
+                        continue;
+                    }
+                    others_ok = others_ok && (tensor->nb[i]*tensor->ne[i] <= tensor->nb[dim] || tensor->nb[i] % nb_next == 0);
+                }
+                if (others_ok) {
+                    return {ggml_backend_meta_split_axis(dim), {0}, {1}, 1};
+                }
+            }
+        }
         if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED || src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_PARTIAL) {
             return src_ss[0];
         }

src/llama-graph.cpp

@@ -396,6 +396,17 @@ void llm_graph_input_cls::set_input(const llama_ubatch * ubatch) {
     }
 }
 
+static void set_input_s_slot_rev(ggml_tensor * s_slot_rev) {
+    if (s_slot_rev) {
+        GGML_ASSERT(ggml_backend_buffer_is_host(s_slot_rev->buffer));
+        int32_t * data = (int32_t *) s_slot_rev->data;
+
+        for (int64_t i = 0; i < s_slot_rev->ne[0]; ++i) {
+            data[i] = (int32_t) (s_slot_rev->ne[0] - 1 - i);
+        }
+    }
+}
+
 void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {
     GGML_UNUSED(ubatch);
 
@@ -410,6 +421,8 @@ void llm_graph_input_rs::set_input(const llama_ubatch * ubatch) {
             data[i] = mctx->s_copy(i);
         }
     }
+
+    set_input_s_slot_rev(s_slot_rev);
 }
 
 bool llm_graph_input_rs::can_reuse(const llm_graph_params & params) {
@@ -1188,6 +1201,8 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
             data[i] = mctx->get_recr()->s_copy(i);
         }
     }
+
+    set_input_s_slot_rev(inp_rs->s_slot_rev);
 }
 
 bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) {
@@ -1231,6 +1246,8 @@ void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) {
             data[i] = mctx->get_recr()->s_copy(i);
         }
     }
+
+    set_input_s_slot_rev(inp_rs->s_slot_rev);
 }
 
 bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) {
@@ -1304,6 +1321,8 @@ void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) {
             data[i] = mctx->get_recr()->s_copy(i);
         }
     }
+
+    set_input_s_slot_rev(inp_rs->s_slot_rev);
 }
 
 bool llm_graph_input_mem_hybrid_iswa::can_reuse(const llm_graph_params & params) {

src/llama-graph.h

@@ -294,6 +294,9 @@ public:
     ggml_tensor * s_copy_main; // I32 [n_seqs]
     ggml_tensor * s_copy_tail; // I32 [n_rs - 1]
 
+    // I32 [n_slot] = {n_slot - 1, ..., 0}, rollback slot of each conv window, created by the graphs that write them
+    ggml_tensor * s_slot_rev = nullptr;
+
     const llama_memory_recurrent_context * mctx;
 
     // used in view offsets, need to match for valid graph reuse

src/models/delta-net-base.cpp

@@ -496,26 +496,35 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state(
         // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside
         //   the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument
 
-        const int64_t K = (int64_t) cparams.n_rs_seq + 1;
+        const int64_t K  = (int64_t) cparams.n_rs_seq + 1;
+        const int64_t n  = conv_input->ne[0] - conv_states->ne[0]; // n_seq_tokens
+        const int64_t Kw = std::min(n + 1, K);                     // windows that exist in this ubatch
 
-        for (int64_t t = 1; t <= K; ++t) {
-            const int64_t s_idx  = std::max<int64_t>(0, conv_input->ne[0] - conv_states->ne[0] - K + t);
-            const int64_t s_slot = K - t;
-
-            ggml_tensor * conv_state_last =
-                ggml_view_3d(ctx0, conv_input,
-                        conv_kernel_size - 1, conv_channels, n_seqs,
-                        conv_input->nb[1], conv_input->nb[2],
-                        ggml_row_size(conv_input->type, s_idx));
-
-            ggml_tensor * conv_state_update =
-                ggml_view_2d(ctx0,
-                        conv_states_all, row_count, n_seqs,
-                        conv_states_all->nb[1],
-                        (s_slot * mem_size + kv_head) * row_size);
-
-            ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update));
+        // {Kw-1, ..., 0}, filled on the host and shared by every layer of the graph
+        if (inp->s_slot_rev == nullptr) {
+            inp->s_slot_rev = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, Kw);
+            ggml_set_input(inp->s_slot_rev);
+            ggml_set_name(inp->s_slot_rev, "rs_s_slot_rev");
         }
+
+        // window w starts at column n - (Kw - 1) + w (ascending time) and is rollback slot Kw - 1 - w
+        // slots reaching back past this ubatch are left unwritten, like the ssm snapshots in build_recurrent_attn
+        // the windows overlap, so the view is created with one window and widened after the size check
+        ggml_tensor * windows = ggml_view_4d(ctx0, conv_input,
+                conv_kernel_size - 1, 1, conv_channels, n_seqs,
+                ggml_element_size(conv_input), conv_input->nb[1], conv_input->nb[2],
+                ggml_row_size(conv_input->type, n - (Kw - 1)));
+        windows->ne[1] = Kw;
+
+        // [d-1, slot, channel, seq] view of the cells of this ubatch
+        ggml_tensor * slots = ggml_view_4d(ctx0, conv_states_all,
+                conv_kernel_size - 1, K, conv_channels, n_seqs,
+                (size_t) mem_size * row_size,
+                ggml_row_size(conv_states_all->type, conv_kernel_size - 1),
+                conv_states_all->nb[1],
+                kv_head * row_size);
+
+        ggml_build_forward_expand(gf, ggml_set_rows(ctx0, slots, windows, inp->s_slot_rev));
     }
 
     return conv_input;