torad-labs · marcospaulo · Sep 30, 2026 · Sep 30, 2026 · Sep 30, 2026 · Sep 30, 2026
diff --git a/TORAD.md b/TORAD.md
diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
@@ -2467,6 +2467,15 @@ extern "C" {
     GGML_API bool ggml_flash_attn_ext_get_mask_prefix(
             const struct ggml_tensor * a);
 
+    // Use finite mask entries as a sparse K/V set. Set 0 to disable.
+    // n_kv_max must bound the number of finite entries in every mask row.
+    GGML_API void ggml_flash_attn_ext_set_n_kv_max(
+            struct ggml_tensor * a,
+            int32_t              n_kv_max);
+
+    GGML_API int32_t ggml_flash_attn_ext_get_n_kv_max(
+            const struct ggml_tensor * a);
+
     // TODO: needs to be adapted to ggml_flash_attn_ext
     GGML_API struct ggml_tensor * ggml_flash_attn_back(
            struct ggml_context * ctx,
@@ -2666,6 +2675,19 @@ extern "C" {
         struct ggml_tensor  * weights,
         struct ggml_tensor  * mask);
 
+    // the same with key i of stream s at row k_rows[i, s] of k (a view into a cache, no gathered copy), and the
+    // scores in f32 whatever k's type (as over the f32 rows ggml_get_rows would have copied):
+    // k:       [n_embd_idx, 1, n_rows, ne3]
+    // k_rows:  [n_kv, ne3] I32, each in [0, n_rows)
+    // mask, res: n_kv as above
+    GGML_API struct ggml_tensor * ggml_lightning_indexer_rows(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * q,
+        struct ggml_tensor  * k,
+        struct ggml_tensor  * k_rows,
+        struct ggml_tensor  * weights,
+        struct ggml_tensor  * mask);
+
     // DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880)
     // In short these operations are replacements for the original residual connection (x = transformer(x) + x)
     // using a richer representation through streams.

diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp
@@ -579,7 +579,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co
             }
         }
         if (ret.axis == GGML_BACKEND_SPLIT_AXIS_NONE) {
-            ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
+            // no sources (ARANGE): every device computes the same values
+            ret = {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
         }
         if (scalar_only && ret.axis >= 0 && ret.axis < GGML_MAX_DIMS) {
             ret = {GGML_BACKEND_SPLIT_AXIS_UNKNOWN, {0}, {1}, 1};
@@ -851,6 +852,8 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co
         for (size_t i = 0; i < 4; i++) {
             GGML_ASSERT(src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
         }
+        // ggml_lightning_indexer_rows: every device reads the same keys through the same rows
+        GGML_ASSERT(tensor->src[4] == nullptr || src_ss[4].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
         return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
     };
 

diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp
@@ -2386,9 +2386,10 @@ bool ggml_backend_compare_graph_backend(ggml_backend_t backend1, ggml_backend_t
 
     if (num_test_nodes != 0) {
         GGML_ASSERT(test_nodes);
-        // Compute the whole graph and only test the output for specific tensors
-        ggml_backend_graph_compute(backend1, g1);
+        // Compute the whole graph and only test the output for specific tensors: backend1's last, so its outputs are
+        // read as soon as its evaluation returns, and work it leaves running past the return reads as a wrong result
         ggml_backend_graph_compute(backend2, g2);
+        ggml_backend_graph_compute(backend1, g1);
 
         bool verified = false;
         for (int i = 0; i < g1->n_nodes; i++) {

diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp
@@ -12163,11 +12163,13 @@ void ggml_compute_forward_lightning_indexer(
     const ggml_tensor * k = dst->src[1];
     const ggml_tensor * w = dst->src[2]; // weights
     const ggml_tensor * m = dst->src[3]; // mask
+    const ggml_tensor * r = dst->src[4]; // ggml_lightning_indexer_rows: key i of stream s at k's row r[i, s]
 
     GGML_ASSERT(dst->type  == GGML_TYPE_F32);
     GGML_ASSERT(   q->type == GGML_TYPE_F32);
     GGML_ASSERT(   w->type == GGML_TYPE_F32);
     GGML_ASSERT(   m->type == GGML_TYPE_F16);
+    GGML_ASSERT(r == nullptr || r->type == GGML_TYPE_I32);
 
     GGML_TENSOR_LOCALS(int64_t, neq,  q, ne)
     GGML_TENSOR_LOCALS(size_t,  nbq,  q, nb)
@@ -12190,7 +12192,7 @@ void ggml_compute_forward_lightning_indexer(
     const int n_head    = q->ne[1];
     const int n_tokens  = q->ne[2];
     const int n_stream  = q->ne[3];
-    const int n_kv      = k->ne[2];
+    const int n_kv      = dst->ne[0];
 
     ggml_to_float_t const k_to_float = ggml_get_type_traits(k->type)->to_float;
     GGML_ASSERT((k->type == GGML_TYPE_F32 || k_to_float) && "lightning indexer: unsupported K-type");
@@ -12215,7 +12217,9 @@ void ggml_compute_forward_lightning_indexer(
             const ggml_fp16_t *   m_row = (ggml_fp16_t *) ((char *)   m->data + t*nbm1 + (s%nem3)*nbm3);
             float             * dst_row =       (float *) ((char *) dst->data + t*nb1  +        s*nb3 );
             for (int ik = ir0; ik < ir1; ++ik) {
-                char * k_row = (char *) k->data + ik*nbk2 + s*nbk3;
+                const int64_t i_row = r ? ((const int32_t *) ((const char *) r->data + s*r->nb[1]))[ik] : ik;
+                GGML_ASSERT(i_row >= 0 && i_row < nek2);
+                char * k_row = (char *) k->data + i_row*nbk2 + s*nbk3;
                 if (k_to_float) {
                     k_to_float(k_row, k_row_f32, n_embd);
                 } else {

diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu
@@ -542,6 +542,45 @@ void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst,
     }
 }
 
+// a thread an element of a token's sum, the slots in their order: each product and each sum rounded on its own, as the
+// MUL and the ADDs round them (never contracted into an FMA)
+static __global__ void k_moe_weighted_sum(const float * experts, const float * weights, float * dst, const int n_embd,
+        const int n_used, const int64_t se1, const int64_t se2, const int64_t sw1, const int64_t sw2, const int64_t sd1) {
+    ggml_cuda_pdl_lc();
+    const int     i = blockIdx.x*blockDim.x + threadIdx.x;
+    const int64_t t = blockIdx.y;
+    if (i >= n_embd) {
+        return;
+    }
+    ggml_cuda_pdl_sync();
+    const float * e = experts + t*se2 + i;
+    const float * w = weights + t*sw2;
+    float acc = __fmul_rn(e[0], w[0]);
+    for (int s = 1; s < n_used; ++s) {
+        acc = __fadd_rn(acc, __fmul_rn(e[s*se1], w[s*sw1]));
+    }
+    dst[t*sd1 + i] = acc;
+}
+
+void ggml_cuda_op_moe_weighted_sum(ggml_backend_cuda_context & ctx, const ggml_tensor * experts,
+        const ggml_tensor * weights, ggml_tensor * dst) {
+    GGML_ASSERT(experts->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32);
+    GGML_ASSERT(experts->nb[0] == sizeof(float) && dst->nb[0] == sizeof(float));
+    const int     n_embd   = (int) experts->ne[0];
+    const int     n_used   = (int) experts->ne[1];
+    const int64_t n_tokens = experts->ne[2];
+    GGML_ASSERT(weights->ne[0] == 1 && weights->ne[1] == n_used && weights->ne[2] == n_tokens);
+    GGML_ASSERT(dst->ne[0] == n_embd && dst->ne[1] == n_tokens);
+
+    constexpr int block_size = 128; // k_bin_bcast's
+    const ggml_cuda_kernel_launch_params params = ggml_cuda_kernel_launch_params(
+        dim3((n_embd + block_size - 1) / block_size, (unsigned) n_tokens, 1), dim3(block_size, 1, 1), 0, ctx.stream());
+    ggml_cuda_kernel_launch(k_moe_weighted_sum, params, (const float *) experts->data, (const float *) weights->data,
+        (float *) dst->data, n_embd, n_used, experts->nb[1] / (int64_t) sizeof(float),
+        experts->nb[2] / (int64_t) sizeof(float), weights->nb[1] / (int64_t) sizeof(float),
+        weights->nb[2] / (int64_t) sizeof(float), dst->nb[1] / (int64_t) sizeof(float));
+}
+
 void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
 

diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh
@@ -10,3 +10,10 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst
 
 void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse);
 void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse);
+
+// Routed experts' weighted sum as build_moe_ffn writes it, MUL(experts, weights) then the slots' views added in order
+// (ADD(ADD(v0, v1), v2), ...), in one launch: dst[i, t] = ((e[i,0,t]*w[0,t] + e[i,1,t]*w[1,t]) + ...), every product and
+// sum rounded as those nodes round them, so bit for bit their result. experts [n_embd, n_used, n_tokens], weights
+// [1, n_used, n_tokens], dst [n_embd, n_tokens], all F32 with contiguous rows; dst lies over neither input.
+void ggml_cuda_op_moe_weighted_sum(ggml_backend_cuda_context & ctx, const ggml_tensor * experts,
+        const ggml_tensor * weights, ggml_tensor * dst);