Exp 2026 07 10 On NPU Attention Batched GEMM and Norm - oRKLLM/ork-driver GitHub Wiki
2026-07-10 — ⚖️ On-NPU batched-GEMM (attention/GDN) + normalization primitives: built + validated, but perf-gated
Commit: ork-driver main 32db14c (+ f4b0154 ork_bench fix) · src/npu.c, src/neon_activations.{c,h},
include/ork_npu.h, examples/test_bmm.c (in make test), tools/re/test_rsqrt_lut.c (make test_rsqrt_lut).
Env gate: ORK_NORM_NPU (norm), ORK_ATTN (ggml-ork batched-MUL_MAT dispatch, fork side).
Goal: get the ops the static-weight matmul path can't serve — attention (QKᵀ/AV) and the Gated-Delta-Net chunked matmuls, and RMSNorm/L2Norm — onto the NPU, to cut the CPU wall on the qwen35 SSM-hybrid target (Ornith-9B). Outcome: the primitives are built + bit-exact-validated and landed as gated infrastructure, but each standalone op is submit-floor-bound and loses to the fused CPU/NEON pass — the durable wins are the primitives themselves and the register-level RE (below), not an end-to-end speedup.
What's new (validated, in make test)
| primitive | what | accuracy vs CPU |
|---|---|---|
ork_bmm_i8/i4/fp16 |
batched dynamic GEMM — both operands dynamic (attention QKᵀ/AV, GDN chunked matmuls). Per-batch ephemeral pack + submit. | i8/i4 exact, fp16 4e-4 |
ork_npu_rmsnorm_f16 / ork_npu_l2norm_f16 (gated) |
full-channel Σx² reduction on NPU as a matmul contraction (any n via K-split) + rsqrt on NPU |
1.3e-3 (n=128), 1.5e-3 (n=4096) |
ork_mm_build_f16_rsqrt_lut |
bakes 1/√(ss/n+ε) into the fp16 fused-output PWL LUT — the fp16 fused-into-the-reduce twin of the existing standalone ork_npu_rsqrt_i8/i16 (Exp-2026-07-04). FUSED (K=n≤2048, one submit) 6e-4; DECOUPLED (K=512, any n) 5e-4 |
|
ork_l2norm_f32 |
NEON CPU L2 normalize (the GDN q/k norm; ork_rmsnorm_f32 already existed) |
exact |
K=32 was degenerate. The decoupled rsqrt-on-ss first used a K=32 matmul → the fp16 fused-silu tiling
returns acc≈0 → constant output (~10% err). Fix: K=512 (the LUT builder's geometry) with the scalar ss
fed dense + normalized (A[m,k]=ss/G ∀k, weight =−S·G/512 → acc=−S·ss, A~O(1), small weight).
Why it's perf-gated (the measurements)
ORK_ATTN routes batched MUL_MAT to ork_bmm, measured with bench_monitored.sh:
| model | batched matmuls route? | NPU avg/peak | prefill | verdict |
|---|---|---|---|---|
| Ornith-9B (qwen35 GDN) | yes (GDN chunks are ggml_cont'd) |
1.6% / 19% (vs 1.5% off) | 27.6 (vs 27.3) | neutral — GDN matmuls are a tiny slice; the CPU wall is the sequential recurrence (gated_delta_net/solve_tri/cumsum/exp) |
| Qwen3-1.7B (pure attn) | no — real QKᵀ/AV read non-contiguous KV-cache/permuted-Q views → ggml_is_contiguous declines |
unchanged | unchanged | attention core not offloaded (needs strided-operand support) |
Norm standalone: ~4 sub-ops (square/reduce/rsqrt/scale, chainable to ~1 submit) at the ~365 µs submit floor vs a ~µs fused NEON pass → loses. The win only materializes by fusing into an adjacent matmul's output stage.
RE landed (open-source cross-reference — the high-value part)
- The RK3588 NPU is NVDLA-derived; cross-referenced the mainline
drivers/accel/rocket/rocket_registers.h(Tomeu Vizoso) + NVDLA docs. Our numeric regs are now NAMED: our "0x40xx" = the DPU output stage —0x4020 DST_BASE,0x4024 DST_SURF_STRIDE,0x4030/34/3c DATA_CUBE_WIDTH/HEIGHT/CHANNEL(geometry),0x4040 DPU_BS_CFG,0x4060 DPU_BN_CFG,0x4070 DPU_EW_CFG (+LUT). - Why full-channel norm can't use the fixed-function LRN op: NVDLA CDP (the LRN unit) caps its reduction
window at n∈{3,5,7,9} — it can't do a
d_model=2048–4096 reduction. This is why a hand-built LRN.rknnwithsize=63falls back to CPU whilesize=5runs on NPU, and why the norm reduction must instead ride a matmul contraction (K unbounded via K-split) — which is exactly what shipped. - RE box: rknn-toolkit2 2.3.2 in the
.239Colima VM (build recipe = torchnn.LocalResponseNormviatorch.onnx.export,do_quantization=False; a hand-built ONNXLRNnode routes to CPU).
Verdict
Landed as gated, validated infrastructure. Attention/GDN/norm on the NPU is not an end-to-end win on
RK3588 — small-M/small-op work is submit-floor-bound (consistent with the decode M=1 result), and the qwen35
CPU wall is the sequential SSM recurrence, not matmuls. ork_bmm is the reusable batched-GEMM building block
(useful for a large-contiguous-batched workload / a future fused-into-matmul norm); the register-name RE
accelerates any future on-NPU output-stage work. Overlap note: the standalone SDP activation/rsqrt/exp family
already shipped in Exp-2026-07-04; this session adds the batched
GEMM + the matmul-reduction norm + the fp16 fused-into-reduce rsqrt twin.