marfrit/rk-llama.cpp
1
0
Fork 0
Code Issues Pull Requests Actions 33 Packages Projects Releases Wiki Activity

ggml-rocket: multi-fd (3-core) column-split + task_num param

Browse Source 
Integrate the second proven lever: split each MUL_MAT's N columns across
ROCKET_NFD=3 independent rocket_open() fds (TILE_N-aligned ranges), run them in
3 pthreads. A single fd/entity caps at 2 NPU cores (DRM sched tie-break, Fable);
3 fds get all 3 cores (33/33/33, ~6.4x aggregate proven standalone). task_num
is now a builder parameter (was a global) — no TLS (broke ggml LTO) and no
cross-thread race. Column ranges align to TILE_N so the per-tile weight-scale
blocks aren't split. Correctness re-verified (kscan all shapes; gemma4 pp512
clean).

Measured: NEUTRAL on the current hybrid (repack on -> only 2 bf16 ops offload,
multi-fd overhead ~cancels the gain on so few ops: pp512 ~48 vs 49). Full offload
(GGML_NO_REPACK) with 3fd+CBUF is now HOST-BOUND at 20.6 pp512 (186% CPU) — the
backend re-dequantizes every Q8_0 weight to F32 per forward on one thread. This
is the definitive evidence that the last required piece is the **prepack weight
buffer-type** (dequant+int8-pack ONCE at load into device BOs, size-partitioned
across the 3 fds under the 4GiB/fd IOVA budget). All NPU-side levers (BO pool,
weight cache, CBUF reuse, 3-core) are now in place and proven; the host-side
per-forward dequant is the remaining wall, and only prepack removes it.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EWpfhDgYNA21tETDP9ueBE
This commit is contained in:
mfritsche
2026-07-14 13:25:51 +02:00
parent 412cba2601
commit 288d70413a
8 changed files with 75 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-rocket/CMakeLists.txt
+3
View File
@@ -19,3 +19,6 @@ target_include_directories(ggml-rocket
PUBLIC
${CMAKE_CURRENT_SOURCE_DIR}
)
find_package(Threads REQUIRED)
target_link_libraries(ggml-rocket PRIVATE Threads::Threads)
ggml/src/ggml-rocket/ggml-rocket.cpp
+49 -18
View File
@@ -27,6 +27,7 @@
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <pthread.h>
extern "C" {
#include "librocket.h"
@@ -41,9 +42,10 @@ extern "C" {
#define ROCKET_TILE_N 128
#define ROCKET_K_MAX 8192
#define ROCKET_MIN_BATCH 32 // below this, CPU wins (decode is bandwidth-bound)
#define ROCKET_NFD 3 // one fd/entity caps at 2 NPU cores (DRM sched tie-break); use 3 for all cores
struct ggml_backend_rocket_context {
int fd = -1;
int fd[ROCKET_NFD] = { -1, -1, -1 };
// reusable scratch (grown as needed, never shrunk)
std::vector<float> wf; // dequantized weight plane [N*K]
std::vector<uint8_t> aq; // quantized activations [M*K]
@@ -60,6 +62,20 @@ static inline uint8_t rocket_q8(float v, float inv_s) {
return (uint8_t)(q + 128); // symmetric int8 stored as uint8, zp=128
}
struct rkt_thread_arg {
int fd; const uint8_t *X, *Wc; uint32_t M, N, K;
float in_scale, out_scale; const float *out_scales;
uint32_t col_start, col_num; uint8_t *Y; int rc;
};
static void * rkt_thread_fn(void *a) {
struct rkt_thread_arg *t = (struct rkt_thread_arg *)a;
t->rc = rkt_npu_matmul(t->fd, t->X, t->Wc, NULL, t->M, t->N, t->K,
128, 128, 128, t->in_scale, 1.0f, t->out_scale,
t->out_scales, ROCKET_TILE_M, ROCKET_TILE_N,
t->col_start, t->col_num, t->Y);
return NULL;
}
static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struct ggml_tensor * dst) {
const struct ggml_tensor * src0 = dst->src[0]; // weight [K, N]
const struct ggml_tensor * src1 = dst->src[1]; // activ [K, M]
@@ -168,16 +184,29 @@ static void ggml_backend_rocket_mul_mat(ggml_backend_rocket_context * ctx, struc
}
}
static const bool dbg = getenv("GGML_ROCKET_DEBUG") != NULL;
if (dbg) fprintf(stderr, "[rocket] mul_mat M=%lld N=%lld K=%lld plane(%lld,%lld) type=%s ...",
(long long)M, (long long)N, (long long)K,
(long long)i12, (long long)i13, ggml_type_name(type)), fflush(stderr);
int rc = rkt_npu_matmul(ctx->fd, ctx->aq.data(), ctx->wq.data(), NULL,
(uint32_t)M, (uint32_t)N, (uint32_t)K,
128, 128, 128,
a_scale, 1.0f, out_scale, ctx->os.data(),
ROCKET_TILE_M, ROCKET_TILE_N, ctx->yq.data());
if (dbg) fprintf(stderr, " rc=%d\n", rc);
// 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;
per = ((per + ROCKET_TILE_N - 1) / ROCKET_TILE_N) * ROCKET_TILE_N;
int nth = 0;
for (int fdi = 0; fdi < ROCKET_NFD; fdi++) {
uint32_t c0 = (uint32_t)fdi * per;
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,
a_scale, out_scale, ctx->os.data(), c0, cn, ctx->yq.data(), 0 };
pthread_create(&th[nth], NULL, rkt_thread_fn, &ta[nth]);
nth++;
}
int rc = 0;
for (int i = 0; i < nth; i++) {
pthread_join(th[i], NULL);
if (ta[i].rc) rc = ta[i].rc;
}
GGML_ASSERT(rc == 0 && "rkt_npu_matmul failed");
// --- dequantize int8 result into dst (F32) ---
@@ -260,14 +289,16 @@ static ggml_guid_t ggml_backend_rocket_guid(void) {
}
ggml_backend_t ggml_backend_rocket_init(void) {
int fd = rocket_open("/dev/accel/accel0");
if (fd < 0) {
GGML_LOG_ERROR("%s: failed to open /dev/accel/accel0 (%d)\n", __func__, fd);
return NULL;
}
ggml_backend_rocket_context * ctx = new ggml_backend_rocket_context;
ctx->fd = fd;
for (int i = 0; i < ROCKET_NFD; i++) {
int fd = rocket_open("/dev/accel/accel0");
if (fd < 0) {
GGML_LOG_ERROR("%s: failed to open /dev/accel/accel0 (%d)\n", __func__, fd);
delete ctx;
return NULL;
}
ctx->fd[i] = fd;
}
ggml_backend_t backend = new ggml_backend {
/* .guid = */ ggml_backend_rocket_guid(),
ggml/src/ggml-rocket/rkt_gemm.c
+7 -4
View File
@@ -6,22 +6,25 @@
#include "rkt_matmul.h"
int rkt_gemm_plan(uint32_t M, uint32_t N, uint32_t tile_m, uint32_t tile_n,
uint32_t col_start, uint32_t col_num,
struct rkt_gemm_tile *tiles, int max_tiles)
{
if (M == 0 || N == 0 || tile_m == 0 || tile_n == 0)
return -1;
uint32_t col_end = col_start + col_num;
if (col_end > N) col_end = N;
int t = 0;
/* column-major: consecutive tiles share the column tile, so the weight
* pack+write can be cached across all row tiles of a column. */
for (uint32_t c = 0; c < N; c += tile_n) {
/* column-major over [col_start,col_end): consecutive tiles share the column
* tile (weight pack+write cached across a column's row tiles). */
for (uint32_t c = col_start; c < col_end; c += tile_n) {
for (uint32_t r = 0; r < M; r += tile_m) {
if (t >= max_tiles)
return -1;
tiles[t].row = r;
tiles[t].col = c;
tiles[t].m = (r + tile_m <= M) ? tile_m : (M - r);
tiles[t].n = (c + tile_n <= N) ? tile_n : (N - c);
tiles[t].n = (c + tile_n <= col_end) ? tile_n : (col_end - c);
t++;
}
}
ggml/src/ggml-rocket/rkt_gemm.h
+1
View File
@@ -28,6 +28,7 @@ struct rkt_gemm_tile {
* rkt_gemm_op_fits() to confirm a tile actually fits one rocket op.
*/
int rkt_gemm_plan(uint32_t M, uint32_t N, uint32_t tile_m, uint32_t tile_n,
uint32_t col_start, uint32_t col_num,
struct rkt_gemm_tile *tiles, int max_tiles);
/* True (1) if an m x n, full-k matmul fits a single rocket op. */
ggml/src/ggml-rocket/rkt_matmul.c
+3 -5
View File
@@ -83,8 +83,6 @@ static unsigned calc_weights_banks(const struct op *o)
return banks;
}
int rkt_g_task_num = 0; /* experiment: CBUF weight-reuse multi-task chaining */
int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
uint32_t M, uint32_t N, uint32_t K,
uint64_t input_dma, uint64_t weights_dma,
@@ -93,7 +91,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
int32_t weight_zero_point,
int32_t output_zero_point,
float input_scale, float weights_scale,
float output_scale, uint64_t bias_dma)
float output_scale, uint64_t bias_dma, int task_num)
{
if (M == 0 || N == 0 || K == 0)
return -1;
@@ -210,7 +208,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
c.reuse_weights_cbuf = reuse_weights_cbuf;
c.addition_input = o.addition_input;
c.add_tensor = o.add_tensor;
c.task_num = rkt_g_task_num;
c.task_num = task_num;
/* ---- CORE/DPU/PC params ---- */
struct coredpu_params d;
@@ -260,5 +258,5 @@ int rkt_build_matmul_regcmd(uint64_t *out, int out_capacity,
input_dma, weights_dma, output_dma,
input_zero_point, weight_zero_point,
output_zero_point,
1.0f, 1.0f, 1.0f, 0);
1.0f, 1.0f, 1.0f, 0, 0);
}
ggml/src/ggml-rocket/rkt_matmul.h
+1 -1
View File
@@ -51,7 +51,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
int32_t input_zero_point, int32_t weight_zero_point,
int32_t output_zero_point,
float input_scale, float weights_scale,
float output_scale, uint64_t bias_dma);
float output_scale, uint64_t bias_dma, int task_num);
int rkt_build_matmul_regcmd(uint64_t *out, int out_capacity,
uint32_t M, uint32_t N, uint32_t K,
ggml/src/ggml-rocket/rkt_npu_matmul.c
+9 -11
View File
@@ -38,8 +38,6 @@ __attribute__((destructor)) static void rkt_prof_dump(void){
g_setup_ns/1e9,g_submit_ns/1e9,g_prep_ns/1e9,g_free_ns/1e9);
}
extern int rkt_g_task_num;
int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
uint32_t M, uint32_t N, uint32_t K,
uint64_t input_dma, uint64_t weights_dma,
@@ -47,7 +45,7 @@ int rkt_build_matmul_regcmd_scaled(uint64_t *out, int out_capacity,
int32_t input_zero_point, int32_t weight_zero_point,
int32_t output_zero_point,
float input_scale, float weights_scale,
float output_scale, uint64_t bias_dma);
float output_scale, uint64_t bias_dma, int task_num);
struct rkt_bo { uint32_t h; uint64_t dma, off; void *map; uint32_t sz; };
@@ -87,7 +85,8 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
const int32_t *bias, uint32_t M, uint32_t N, uint32_t K,
uint8_t izp, uint8_t wzp, uint8_t ozp,
float in_scale, float wt_scale, float out_scale, const float *out_scales,
uint32_t tile_m, uint32_t tile_n, uint8_t *Y)
uint32_t tile_m, uint32_t tile_n,
uint32_t col_start, uint32_t col_num, uint8_t *Y)
{
int ret = -1;
int32_t *bias0 = NULL;
@@ -110,15 +109,15 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
uint8_t *wpk = calloc(1, max_wsz);
int32_t *bpk = malloc((size_t)tile_n * sizeof(int32_t));
struct rkt_gemm_tile *tiles = malloc(
((M + tile_m - 1) / tile_m) * ((N + tile_n - 1) / tile_n) *
((M + tile_m - 1) / tile_m) * ((col_num + tile_n - 1) / tile_n) *
sizeof(struct rkt_gemm_tile));
uint64_t *rc = malloc(4096 * sizeof(uint64_t));
if (!Xt || !Wt || !ipk || !wpk || !bpk || !tiles || !rc)
goto out;
int nt = rkt_gemm_plan(M, N, tile_m, tile_n, tiles,
int nt = rkt_gemm_plan(M, N, tile_m, tile_n, col_start, col_num, tiles,
(int)(((M + tile_m - 1) / tile_m) *
((N + tile_n - 1) / tile_n)));
((col_num + tile_n - 1) / tile_n)));
if (nt <= 0)
goto out;
@@ -168,12 +167,12 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
bo_write(fd, &rin[cnt], ipk, rkt_raw_input_size(1, m, K));
bo_write(fd, &ro[cnt], NULL, rkt_raw_output_size(1, m, n));
rkt_g_task_num = cnt; /* task 0 loads weights, 1+ reuse CBUF */
/* task 0 loads weights DDR->CBUF, tasks 1+ reuse SRAM */
int nw = rkt_build_matmul_regcmd_scaled(
rc, 4096, m, n, K, rin[cnt].dma, w.dma, ro[cnt].dma,
izp, wzp, ozp, in_scale, wt_scale,
out_scales ? out_scales[cc] : out_scale, b.dma);
if (nw < 0) { rkt_g_task_num = 0; goto out; }
out_scales ? out_scales[cc] : out_scale, b.dma, cnt);
if (nw < 0) goto out;
bo_write(fd, &rreg[cnt], rc, (unsigned)nw * sizeof(uint64_t));
tasks[cnt].regcmd = (uint32_t)rreg[cnt].dma;
@@ -182,7 +181,6 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
meta[cnt].r = r; meta[cnt].m = m; meta[cnt].n = n;
cnt++; t++;
}
rkt_g_task_num = 0;
struct drm_rocket_job job;
memset(&job, 0, sizeof job);
ggml/src/ggml-rocket/rkt_npu_matmul.h
+2 -1
View File
@@ -30,6 +30,7 @@ int rkt_npu_matmul(int fd, const uint8_t *X, const uint8_t *Wc,
const int32_t *bias, uint32_t M, uint32_t N, uint32_t K,
uint8_t izp, uint8_t wzp, uint8_t ozp,
float in_scale, float wt_scale, float out_scale, const float *out_scales,
uint32_t tile_m, uint32_t tile_n, uint8_t *Y);
uint32_t tile_m, uint32_t tile_n,
uint32_t col_start, uint32_t col_num, uint8_t *Y);
#endif /* RKT_NPU_MATMUL_H */