ggml-rocket: K-segmentation for K>8192 (int8-partial float-accumulate)

Enables MUL_MAT with K>ROCKET_K_MAX (e.g. gemma-4 ffn_down K=12288) on the
NPU via 6144-wide K-chunks, mirroring ggml-rknpu2 compute_k_segments.
Validated correct (coherent long-prompt greedy output; int8 2-way partials
lossless enough). pp512-NEUTRAL on gemma-4-E2B (46.61 vs 46.69) — kept on a
branch as a capability, not defaulted. May help models with heavy K>8192.

ggml/src/ggml-rocket/ggml-rocket.cpp

@@ -103,9 +103,16 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
     const int64_t r2 = ne12 / ne02;
     const int64_t r3 = ne13 / ne03;
 
-    ctx->wf.resize((size_t)N * K);
-    ctx->aq.resize((size_t)M * K);
-    ctx->wq.resize((size_t)N * K);
+    // K-segmentation: split K into chunks <= KLIM and float-accumulate int8
+    // partials (mirrors ggml-rknpu2 compute_k_segments). KLIM <= ROCKET_K_MAX,
+    // clear of the ~11008 single-op correctness cliff, and %256 so every
+    // segment boundary is block-aligned for any quant type.
+    const int64_t KLIM   = 6144;
+    const int64_t KS_MAX = (K < KLIM ? K : KLIM);
+
+    ctx->wf.resize((size_t)N * KS_MAX);
+    ctx->aq.resize((size_t)M * KS_MAX);
+    ctx->wq.resize((size_t)N * KS_MAX);
     ctx->yq.resize((size_t)M * N);
     ctx->ws.resize((size_t)N);
     ctx->os.resize((size_t)N);
@@ -119,13 +126,17 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
             const char * a_plane = (const char *) src1->data + i12*nb12 + i13*nb13;
                   char * d_plane = (      char *)  dst->data + i12*nb2  + i13*nb3;
 
-            // --- dequantize weight plane to F32, per-channel scale + L2 norm ---
-            float wn2max = 0.0f; // max_n ||w_n||_2
+            for (int64_t k0 = 0; k0 < K; k0 += KLIM) {
+            const int64_t Ks     = (K - k0 < KLIM ? K - k0 : KLIM);
+            const bool    kfirst = (k0 == 0);
+
+            // --- dequantize weight slice [k0,k0+Ks) to F32, per-channel scale + L2 norm ---
+            float wn2max = 0.0f;
             for (int64_t n = 0; n < N; n++) {
-                float * wrow = ctx->wf.data() + (size_t)n * K;
-                to_float((const char *) w_plane + n*nb01, wrow, K);
+                float * wrow = ctx->wf.data() + (size_t)n * Ks;
+                to_float((const char *) w_plane + n*nb01 + ggml_row_size(type, k0), wrow, Ks);
                 float amax = 0.0f, l2 = 0.0f;
-                for (int64_t k = 0; k < K; k++) {
+                for (int64_t k = 0; k < Ks; k++) {
                     float v = wrow[k];
                     float a = std::fabs(v);
                     if (a > amax) amax = a;
@@ -136,12 +147,12 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
                 if (l2 > wn2max) wn2max = l2;
             }
 
-            // --- quantize activation plane, per-tensor scale + max L2 row norm ---
+            // --- quantize activation slice, per-tensor scale + max L2 row norm ---
             float amax_a = 0.0f, ym2max = 0.0f;
             for (int64_t m = 0; m < M; m++) {
-                const float * arow = (const float *) (a_plane + m*nb11);
+                const float * arow = (const float *) (a_plane + m*nb11) + k0;
                 float l2 = 0.0f;
-                for (int64_t k = 0; k < K; k++) {
+                for (int64_t k = 0; k < Ks; k++) {
                     float v = arow[k];
                     float a = std::fabs(v);
                     if (a > amax_a) amax_a = a;
@@ -153,21 +164,14 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
             const float a_scale = (amax_a > 0.0f ? amax_a : 1e-6f) / 127.0f;
             const float inv_a   = 1.0f / a_scale;
             for (int64_t m = 0; m < M; m++) {
-                const float * arow = (const float *) (a_plane + m*nb11);
-                uint8_t * qrow = ctx->aq.data() + (size_t)m * K;
-                for (int64_t k = 0; k < K; k++) qrow[k] = rocket_q8(arow[k], inv_a);
+                const float * arow = (const float *) (a_plane + m*nb11) + k0;
+                uint8_t * qrow = ctx->aq.data() + (size_t)m * Ks;
+                for (int64_t k = 0; k < Ks; k++) qrow[k] = rocket_q8(arow[k], inv_a);
             }
 
-            // Cauchy-Schwarz bound on |output| -> never clips int8 output.
             float out_scale = (wn2max * ym2max) / 127.0f;
             if (!(out_scale > 0.0f)) out_scale = 1e-6f;
 
-            // The NPU applies ONE OUT_CVT scale per op == per tile of
-            // ROCKET_TILE_N output columns. So every column in a tile must be
-            // quantized with the SAME weight scale, otherwise off-lead columns
-            // ride the tile-lead's scale and can clip. Use the tile's max
-            // per-channel scale (block boundaries must match rkt_gemm_plan's
-            // column tiling: [0,TILE_N),[TILE_N,2*TILE_N),...).
             for (int64_t c0 = 0; c0 < N; c0 += ROCKET_TILE_N) {
                 int64_t c1 = c0 + ROCKET_TILE_N;
                 if (c1 > N) c1 = N;
@@ -177,16 +181,13 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
                 if (!(wsb > 0.0f)) wsb = 1e-6f;
                 const float inv_w = 1.0f / wsb;
                 for (int64_t n = c0; n < c1; n++) {
-                    const float * wrow = ctx->wf.data() + (size_t)n * K;
-                    uint8_t * qrow = ctx->wq.data() + (size_t)n * K;
-                    for (int64_t k = 0; k < K; k++) qrow[k] = rocket_q8(wrow[k], inv_w);
+                    const float * wrow = ctx->wf.data() + (size_t)n * Ks;
+                    uint8_t * qrow = ctx->wq.data() + (size_t)n * Ks;
+                    for (int64_t k = 0; k < Ks; k++) qrow[k] = rocket_q8(wrow[k], inv_w);
                     ctx->os[n] = out_scale / wsb;
                 }
             }
 
-            // Split N columns across ROCKET_NFD fds == NPU cores (a single fd
-            // caps at 2 cores via the DRM sched tie-break). Column ranges are
-            // TILE_N-aligned so the per-tile weight-scale blocks don't split.
             pthread_t th[ROCKET_NFD];
             struct rkt_thread_arg ta[ROCKET_NFD];
             uint32_t per = ((uint32_t)N + ROCKET_NFD - 1) / ROCKET_NFD;
@@ -197,7 +198,7 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
                 if (c0 >= (uint32_t)N) break;
                 uint32_t cn = (c0 + per <= (uint32_t)N) ? per : ((uint32_t)N - c0);
                 ta[nth] = (struct rkt_thread_arg){ ctx->fd[fdi], ctx->aq.data(),
-                    ctx->wq.data(), (uint32_t)M, (uint32_t)N, (uint32_t)K,
+                    ctx->wq.data(), (uint32_t)M, (uint32_t)N, (uint32_t)Ks,
                     a_scale, out_scale, ctx->os.data(), c0, cn, ctx->yq.data(), 0 };
                 pthread_create(&th[nth], NULL, rkt_thread_fn, &ta[nth]);
                 nth++;
@@ -209,13 +210,16 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
             }
             GGML_ASSERT(rc == 0 && "rkt_npu_matmul failed");
 
-            // --- dequantize int8 result into dst (F32) ---
+            // --- dequantize int8 result into dst (F32), accumulate across K-segments ---
             for (int64_t m = 0; m < M; m++) {
                 float * drow = (float *) (d_plane + m*nb1);
                 const uint8_t * yrow = ctx->yq.data() + (size_t)m * N;
-                for (int64_t n = 0; n < N; n++)
-                    drow[n] = ((int)yrow[n] - 128) * out_scale;
+                if (kfirst)
+                    for (int64_t n = 0; n < N; n++) drow[n]  = ((int)yrow[n] - 128) * out_scale;
+                else
+                    for (int64_t n = 0; n < N; n++) drow[n] += ((int)yrow[n] - 128) * out_scale;
             }
+            } // k-segment
         }
     }
 }
@@ -388,7 +392,7 @@ static bool ggml_backend_rocket_device_supports_op(ggml_backend_dev_t dev, const
             else if (src1->type != GGML_TYPE_F32) why="src1!=F32";
             else if (!ggml_is_contiguous(src0)) why="src0 noncontig";
             else if (!ggml_is_contiguous(src1)) why="src1 noncontig";
-            else if (K <= 0 || K > ROCKET_K_MAX) why="K>Kmax";
+            else if (K <= 0) why="K<=0";
             else if (K % 16 != 0) why="K%16";
             else if (ggml_get_type_traits(src0->type)->to_float == NULL) why="src0 no to_float";
             if (getenv("GGML_ROCKET_DEBUG"))