Exp 2026 06 13 QKV GateUp Fusion - oRKLLM/ork-driver GitHub Wiki

Experiment โ€” QKV / gate-up matmul fusion

Date: 2026-06-13 ยท Status: โš–๏ธ NEUTRAL on 1.7B โ€” not worth the complexity (opt-in, off by default) ยท Path: ggml-ork

Hypothesis

If the per-matmul NPU cost is a fixed submit-latency floor, then fusing independent same-input matmuls (Q/K/V; gate/up) into one wider matmul โ€” concatenating weights along N, quantizing the shared activation once, one submit, then scattering the wide result โ€” should cut submits/token and speed up decode.

Implementation

  • Commit: 1a0a39e8 (oRKLLM/llama.cpp-rockchip, ork-backend)
  • Symbol: ggml-ork.cpp : ggml_backend_ork_mul_mat_group_i8 + detection in graph_compute (adjacent same-src1 2D MUL_MATs โ†’ grouped call).
  • Knob: ORK_FUSE=1 (off by default).

Detection in graph_compute โ€” group adjacent MUL_MATs sharing src1:

const int fuse = (ctx->qbits == 8) && (getenv("ORK_FUSE") != nullptr);
// ...
struct ggml_tensor * grp[16]; int ng = 1; grp[0] = node;
if (fuse && node->ne[2] == 1 && node->ne[3] == 1) {
    while (i + ng < cgraph->n_nodes && ng < 16) {
        struct ggml_tensor * nj = cgraph->nodes[i + ng];
        if (nj->op == GGML_OP_MUL_MAT && nj->src[1] == node->src[1] &&
            nj->src[0]->ne[0] == node->src[0]->ne[0] && nj->ne[2] == 1 && nj->ne[3] == 1)
            grp[ng++] = nj;
        else break;
    }
}
if (ng >= 2) { ggml_backend_ork_mul_mat_group_i8(ctx, grp, ng); i += ng - 1; }

The grouped path โ€” concat weights along N, quantize the shared activation once, ONE submit, scatter:

int Ntot = 0, off[16];
for (int i = 0; i < ng; i++) { off[i] = Ntot; Ntot += (int) g[i]->src[0]->ne[1]; }
// ...build/pack the fused weight [K, Ntot] once, keyed by g[0]->src[0]->data...
ctx->last_src1 = nullptr;                      // group overwrote ctx->ai โ€” kill reuse cache
// quantize shared activation once into ai/as ...
if (ork_mm_run_i8(ctx->npu, ow.w, M, ai, ci)) return false;   // ONE submit for all ng matmuls
for (int i = 0; i < ng; i++) {                 // scatter + dequant into each dst
    const int Ni = (int) g[i]->src[0]->ne[1]; const int o = off[i];
    float * dbase = (float *) g[i]->data;
    for (int m = 0; m < M; m++) { const float rs = as[m]; const int32_t * cr = ci + (size_t) m*Ntot + o;
        float * dr = dbase + (size_t) m*Ni; for (int n = 0; n < Ni; n++) dr[n] = rs * bs[o+n] * (float) cr[n]; }
}

Benchmark (standardized โ€” 2026-06-13)

Qwen3-1.7B-w8a8, llama-bench (decode-only -p 0 -n 128 to isolate the decode arm, prefill from the full matrix), -r 2, DDR 2112 / CPU 2.4 GHz. Both arms all-NPU (ORK_MINM=1).

Config pp512 tg128
all-NPU (ORK_MINM=1) 131.42 ยฑ 2.29 5.03 ยฑ 0.15
+ fusion (ORK_FUSE=1) 136.03 ยฑ 1.53 5.10 ยฑ 0.02
librkllmrt (closed ref, same model) 183.9 11.71

(librkllmrt has no fusion toggle โ€” closed reference. Note: the closed runtime likely does fuse QKV/gate-up into single pre-packed weights at model-build, avoiding the runtime scatter that makes ORK_FUSE a wash here.)

Historical result (non-standard)

An earlier run on Qwen2.5-0.5B via llama-cli showed decode 9.4 โ†’ 6.4 โ€” slower, and matmul count dropping 7951 โ†’ 6821.

Verdict

โš–๏ธ Revised under the standard. On the reference 1.7B model with llama-bench, fusion is neutral โ€” decode 5.10 vs 5.03 and prefill 136 vs 131 are both within run-to-run noise, not the clear loss seen earlier on 0.5B/llama-cli. The result is model- and method-dependent; the earlier "ruled out, slower" verdict was a non-standard measurement. Net: fusion buys nothing measurable here and adds complexity (wide matmul + strided scatter), so it stays off by default behind ORK_FUSE=1 โ€” kept runnable for larger models / a tuned scatter. Either way it is not a decode lever on this model.

โš ๏ธ **GitHub.com Fallback** โš ๏ธ