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 .rknn with size=63 falls back to CPU while size=5 runs 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 .239 Colima VM (build recipe = torch nn.LocalResponseNorm via torch.onnx.export, do_quantization=False; a hand-built ONNX LRN node 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.