train: the exact walk at her real window, exact on a GPU hybrid; the walk is its own module

ggml/include/ggml-opt.h

@@ -230,6 +230,14 @@ extern "C" {
     // for a GRAD leaf (ggml_set_grad) of the evaluated graph, or NULL when the graph has none.
     GGML_API struct ggml_tensor * ggml_opt_leaf_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * leaf);
 
+    // For the graph last given to ggml_opt_prepare_alloc: its backward ADDS the gradient on `tensor`
+    // (a node of that graph, flagged GRAD by this call) into `acc`, an F32 tensor of the same shape
+    // in a buffer the caller owns and zeroes, instead of keeping it in the compute buffer for
+    // ggml_opt_leaf_grad. Nothing gradient-sized outlives its consumer on the device: the
+    // accumulator may live in host memory while the graph runs on a GPU. The optimizer never
+    // zeroes it, except at its own first allocation, which precedes any accumulation.
+    GGML_API void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, struct ggml_tensor * acc);
+
     // ############################################################################
     // ## The high-level functions start here. They do not depend on any private ##
     // ## functions or structs and can be copied to and adapted for user code.   ##

ggml/include/ggml.h

@@ -654,7 +654,7 @@ extern "C" {
         GGML_TENSOR_FLAG_PARAM   =  4, // ...contains trainable parameters
         GGML_TENSOR_FLAG_LOSS    =  8, // ...defines loss for numerical optimization (multiple loss tensors add up)
         GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed
-        GGML_TENSOR_FLAG_GRAD    = 32, // ...needs its gradient computed, but no optimizer updates it (a leaf whose gradient is read back)
+        GGML_TENSOR_FLAG_GRAD    = 32, // ...needs its gradient computed, but no optimizer updates it (read back, or accumulated into a caller's tensor)
     };
 
     enum ggml_tri_type {
@@ -2464,6 +2464,18 @@ extern "C" {
     GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec(
             const struct ggml_tensor * a);
 
+    // the KV positions per tile of FLASH_ATTN_EXT's backward pass (0 = the default): the backward
+    // recomputes each tile's scores, so its memory is queries x tile, never queries x window
+    GGML_API void ggml_flash_attn_ext_set_back_tile(
+            struct ggml_tensor * a,
+            int32_t              tile);
+
+    // an upper bound on the graph nodes FLASH_ATTN_EXT's backward adds for n_kv KV positions at
+    // `tile` (0 = the default): a caller sizing a training graph counts them with its own
+    GGML_API size_t ggml_flash_attn_ext_back_n_nodes(
+            int64_t              n_kv,
+            int32_t              tile);
+
     GGML_API void ggml_flash_attn_ext_add_sinks(
             struct ggml_tensor * a,
             struct ggml_tensor * sinks);

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2622,7 +2622,12 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) {
         //    default: break;
         //}
 
-        ggml_metal_encode_mul_mm(ctx, op);
+        // GGML_PREC_F32 on an F32 x F32 product is honored with the float-staged kernel (CUDA honors
+        // the same request): the default stages through half, ~1e-3 relative per product, which the
+        // training backward cannot afford (flash attention's tiled backward, continuum 2026-10-11)
+        const bool fp32 = !props_dev->has_tensor && ggml_get_op_params_i32(op, 0) == GGML_PREC_F32 &&
+                          op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
+        ggml_metal_encode_mul_mm(ctx, op, fp32);
     } else {
         auto pipeline = ggml_metal_library_get_pipeline_mul_mv(lib, op);
 
@@ -3412,7 +3417,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         // the shared memory needed for the simdgroups to load the KV cache
         // each thread loads (dequantizes) 16 head elements, there are 32 threads in th SG
         //
-#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16))
+        // the F32 kernels hold the query as float (kernels/fa.metal, QH): twice the halves per element
+        const int qh = !use_kv_f16 && op->src[1]->type == GGML_TYPE_F32 ? 2 : 1;
+#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(qh*ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16))
 
         //int64_t nsgmax = 4;
         //

ggml/src/ggml-metal/kernels/fa.metal

@@ -273,12 +273,16 @@ void kernel_flash_attn_ext_impl(
     constexpr short SH  = 2*C; // shared memory per simdgroup (s_t == float)
 
     constexpr short TS = 2*SH;
-    constexpr short T  = DK + 2*PV; // shared memory size per query in (half)
+    // halves per query element: 1 for a half query, 2 for the F32 kernels' float query (Continuum:
+    // a training graph reads attention through the F32 kernels, and a query rounded to half moved
+    // its step, test-walk-exact on the M5 2026-10-11; the host sizes the memory the same way)
+    constexpr short QH = sizeof(q_t)/sizeof(half);
+    constexpr short T  = QH*DK + 2*PV; // shared memory size per query in (half)
 
     threadgroup q_t  * sq  = (threadgroup q_t  *) (shmem_f16 + 0*T); // holds the query data
     threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*T); // same as above but in q4_t
-    threadgroup o_t  * so  = (threadgroup o_t  *) (shmem_f16 + 0*T + Q*DK); // the result for all queries in 8x8 matrices (the O matrix from the paper)
-    threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 0*T + Q*DK);
+    threadgroup o_t  * so  = (threadgroup o_t  *) (shmem_f16 + 0*T + Q*QH*DK); // the result for all queries in 8x8 matrices (the O matrix from the paper)
+    threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 0*T + Q*QH*DK);
     threadgroup s_t  * ss  = (threadgroup s_t  *) (shmem_f16 + Q*T); // scratch buffer for attention, mask and diagonal matrix
     threadgroup s2_t * ss2 = (threadgroup s2_t *) (shmem_f16 + Q*T); // same as above but in s2_t
 
@@ -913,7 +917,7 @@ kernel void kernel_flash_attn_ext(
     //float,  float4,    simdgroup_float8x8
 
 #define FA_TYPES_F32 \
-    half,   half4,     simdgroup_half8x8,  \
+    float,  float4,    simdgroup_float8x8, \
     float,  float4x4,  simdgroup_float8x8, \
     float,  float4x4,  simdgroup_float8x8, \
     float,             simdgroup_float8x8, \

ggml/src/ggml-opt.cpp

@@ -86,6 +86,9 @@ struct ggml_opt_context {
     bool                 period_fresh    = true; // the next backward starts a period (zeroed gradients)
     // the gradients of the evaluated graph's GRAD leaves, readable until the next alloc
     std::unordered_map<const struct ggml_tensor *, struct ggml_tensor *> leaf_grads;
+    // ggml_opt_accumulate_grad: GRAD nodes whose gradient adds into a caller's tensor, for the
+    // graph prepared last (ggml_opt_prepare_alloc clears them)
+    std::unordered_map<const struct ggml_tensor *, struct ggml_tensor *> next_grad_accs;
     std::vector<struct ggml_tensor *> grad_accs;
 
     int64_t iter               = 1;
@@ -305,6 +308,12 @@ static ggml_tensor * ggml_opt_view_root(ggml_tensor * t) {
     return t;
 }
 
+// A CPY whose destination is its own result (ggml_cast) writes no other tensor: it is a pure
+// conversion, recomputed like any op. A CPY into another tensor (a cache write) is a side effect.
+static bool ggml_opt_is_cast(const ggml_tensor * node) {
+    return node->op == GGML_OP_CPY && node->src[1] == node;
+}
+
 static ggml_tensor * ggml_opt_recompute_node(
         ggml_context * ctx, const std::unordered_set<ggml_tensor *> & forward,
         const std::unordered_set<ggml_tensor *> & keep, std::unordered_map<ggml_tensor *, ggml_tensor *> & clones,
@@ -320,9 +329,13 @@ static ggml_tensor * ggml_opt_recompute_node(
     }
     switch (node->op) {
         case GGML_OP_NONE:
-        case GGML_OP_CPY:      // a write into another tensor (a cache): side effect, never repeated
         case GGML_OP_SET_ROWS:
             return node;
+        case GGML_OP_CPY:      // a write into another tensor (a cache): side effect, never repeated
+            if (!ggml_opt_is_cast(node)) {
+                return node;
+            }
+            break;
         default:
             break;
     }
@@ -364,7 +377,7 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd,
         // 6.8e-4 with recompute on until in-place ops counted here).
         const bool view_op = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
                              node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE;
-        if (node->op == GGML_OP_CPY || node->op == GGML_OP_SET_ROWS || (node->view_src != nullptr && !view_op)) {
+        if ((node->op == GGML_OP_CPY && !ggml_opt_is_cast(node)) || node->op == GGML_OP_SET_ROWS || (node->view_src != nullptr && !view_op)) {
             mutated.insert(ggml_opt_view_root(node));
         }
     }
@@ -408,7 +421,9 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd,
             return;
         }
         for (int k = 0; k < GGML_MAX_SRC; ++k) {
-            place(t->src[k]);
+            if (t->src[k] != t) { // a cast is its own destination
+                place(t->src[k]);
+            }
         }
         place(t->view_src);
         placed.insert(t);
@@ -677,6 +692,8 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) {
             if (accumulate) {
                 opt_ctx->grad_accs[i] = opt_ctx->param_grad_acc.at(node);
             }
+        } else if (const auto it = opt_ctx->next_grad_accs.find(node); it != opt_ctx->next_grad_accs.end()) {
+            opt_ctx->grad_accs[i] = it->second;
         }
     }
 
@@ -686,7 +703,7 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) {
     opt_ctx->leaf_grads.clear();
     for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) {
         ggml_tensor * node = opt_ctx->gf->nodes[i];
-        if (node->flags & GGML_TENSOR_FLAG_GRAD) {
+        if ((node->flags & GGML_TENSOR_FLAG_GRAD) && !opt_ctx->next_grad_accs.count(node)) {
             ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_grad, node);
             if (grad) {
                 ggml_set_output(grad); // read back after the step: the allocator must not reuse it
@@ -811,6 +828,14 @@ void ggml_opt_set_next_step(ggml_opt_context_t opt_ctx, bool period_end, float l
     opt_ctx->next_extra_loss = extra_loss;
 }
 
+void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, struct ggml_tensor * acc) {
+    GGML_ASSERT(!opt_ctx->eval_ready && "register an accumulator before ggml_opt_alloc");
+    GGML_ASSERT(acc->type == GGML_TYPE_F32 && ggml_are_same_shape(tensor, acc));
+    GGML_ASSERT((acc->view_src ? acc->view_src : acc)->buffer != nullptr && "the accumulator lives in the caller's buffer");
+    ggml_set_grad(tensor);
+    opt_ctx->next_grad_accs[tensor] = acc;
+}
+
 struct ggml_tensor * ggml_opt_leaf_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * leaf) {
     const auto it = opt_ctx->leaf_grads.find(leaf);
     return it == opt_ctx->leaf_grads.end() ? nullptr : it->second;
@@ -950,6 +975,7 @@ void ggml_opt_prepare_alloc(
     opt_ctx->gf          = gf;
     opt_ctx->inputs      = inputs;
     opt_ctx->outputs     = outputs;
+    opt_ctx->next_grad_accs.clear(); // registrations belong to the graph prepared last
 }
 
 void ggml_opt_set_alloc_budget(ggml_opt_context_t opt_ctx, size_t bytes) {