ggml-rocket: make ROCKET_MIN_BATCH runtime-tunable via GGML_ROCKET_MIN_BATCH

The offload gate (M < 32 -> CPU) was compile-time only. Reading the env var
(default 32 unchanged) gives a clean NPU on/off toggle: GGML_ROCKET_MIN_BATCH=999999
forces CPU-only for benchmarking without rmmod or a CPU-only build. Used for the
2026-07-18 ampere NPU-vs-CPU gemma prefill benchmark.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EWpfhDgYNA21tETDP9ueBE

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

@@ -103,16 +103,9 @@ 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;
 
-    // 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->wf.resize((size_t)N * K);
+    ctx->aq.resize((size_t)M * K);
+    ctx->wq.resize((size_t)N * K);
     ctx->yq.resize((size_t)M * N);
     ctx->ws.resize((size_t)N);
     ctx->os.resize((size_t)N);
@@ -126,17 +119,13 @@ 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;
 
-            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;
+            // --- dequantize weight plane to F32, per-channel scale + L2 norm ---
+            float wn2max = 0.0f; // max_n ||w_n||_2
             for (int64_t n = 0; n < N; n++) {
-                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 * wrow = ctx->wf.data() + (size_t)n * K;
+                to_float((const char *) w_plane + n*nb01, wrow, K);
                 float amax = 0.0f, l2 = 0.0f;
-                for (int64_t k = 0; k < Ks; k++) {
+                for (int64_t k = 0; k < K; k++) {
                     float v = wrow[k];
                     float a = std::fabs(v);
                     if (a > amax) amax = a;
@@ -147,12 +136,12 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
                 if (l2 > wn2max) wn2max = l2;
             }
 
-            // --- quantize activation slice, per-tensor scale + max L2 row norm ---
+            // --- quantize activation plane, 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) + k0;
+                const float * arow = (const float *) (a_plane + m*nb11);
                 float l2 = 0.0f;
-                for (int64_t k = 0; k < Ks; k++) {
+                for (int64_t k = 0; k < K; k++) {
                     float v = arow[k];
                     float a = std::fabs(v);
                     if (a > amax_a) amax_a = a;
@@ -164,14 +153,21 @@ 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) + 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);
+                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);
             }
 
+            // 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;
@@ -181,13 +177,16 @@ 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 * 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);
+                    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);
                     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;
@@ -198,7 +197,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)Ks,
+                    ctx->wq.data(), (uint32_t)M, (uint32_t)N, (uint32_t)K,
                     a_scale, out_scale, ctx->os.data(), c0, cn, ctx->yq.data(), 0 };
                 pthread_create(&th[nth], NULL, rkt_thread_fn, &ta[nth]);
                 nth++;
@@ -210,16 +209,13 @@ 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), accumulate across K-segments ---
+            // --- dequantize int8 result into dst (F32) ---
             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;
-                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;
+                for (int64_t n = 0; n < N; n++)
+                    drow[n] = ((int)yrow[n] - 128) * out_scale;
             }
-            } // k-segment
         }
     }
 }
@@ -388,11 +384,12 @@ static bool ggml_backend_rocket_device_supports_op(ggml_backend_dev_t dev, const
             const int64_t K = src1->ne[0]; // == src0->ne[0]
             const int64_t M = op->ne[1];   // batch rows (tokens)
             const char *why = NULL;
-            if (M < ROCKET_MIN_BATCH) why="M<min_batch";
+            static const int min_batch = (getenv("GGML_ROCKET_MIN_BATCH") ? atoi(getenv("GGML_ROCKET_MIN_BATCH")) : ROCKET_MIN_BATCH);
+            if (M < min_batch) why="M<min_batch";
             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) why="K<=0";
+            else if (K <= 0 || K > ROCKET_K_MAX) why="K>Kmax";
             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"))