diff --git a/ggml/src/ggml-rocket/ggml-rocket.cpp b/ggml/src/ggml-rocket/ggml-rocket.cpp index 20e57f8e3..468a08a7e 100644 --- a/ggml/src/ggml-rocket/ggml-rocket.cpp +++ b/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"))