Exp 2026 07 14 Attention Perchan Scale and CNA Strided Read - oRKLLM/ork-driver GitHub Wiki

Exp 2026-07-14 — Per-channel-scaled fp16 matmul on NPU + CNA strided-read discovery

Campaign to build a fully on-NPU per-channel-scaled fp16 matmul (the attention A·V-normalize: out[m][n] = (Σ_k A[m][k]·B[k][n]) · scale[n]), and the NPU-hardware discoveries it produced. Board RK3588 10.3.0.236; ork-driver feat/attn-primitives.

Result — three working closes, all bit-exact on silicon

Validated MKN {8,32,64}/{16,64,128}/{32,128,256} = 512/2048/8192 exact:

  1. ork_npu_mm_perchan_f16 — fp16 matmul (contiguous fp16 out) → atom-8 per-channel EW-mul SDP (ork_npu_mul_perchan_f16, contiguous in/out, atom-8 internal repack). O(M·N); a tiny CPU repack.
  2. ork_npu_mm_perchan_f16_diag — per-channel scale as a diagonal 2nd matmul (A·B)·diag(scale), device-resident (G stays on NPU), 2 submits. O(M·N²); pure NPU, zero CPU touch of the intermediate.
  3. ork_npu_mm_perchan_f16_diag + ORK_DIAG_CHAIN — the two matmuls PC-chained into ONE submit (both uniform enable=0xd, chain like run_chain_i8). O(M·N²); pure-NPU single-submit (M4).

NPU hardware discoveries

1. The fp16 matmul CAN emit fp16 output (was thought impossible)

Earlier "fp16 matmul has no 2-byte writeout" was three fixable bugs, found by decoding a captured vendor fp16→fp16 matmul (~/rknn_sdk/cap_fp16f16.dec):

  • Use the MATMUL output stage, not the conv's: 0x4040 BS_CFG = 0x53 (BS FULLY bypassed), 0x40c0=0x20 — the conv's BS-active 0x20150 HANGS the matmul.
  • 0x4084 DPU_OUT_CVT_SCALE = 0x00010001 — FP32TOFP16_EN (bit16) enables the fp32→fp16 writeout (was 0x1).
  • sched=0 for small K; the fp16 matmul writes CONTIGUOUS [M][N], not the int-path atom-8. ork_npu_probe_f16_mm_f16out = 512/512 bit-exact.

2. fp16 matmul writes CONTIGUOUS only (atom-8 output hangs)

Forcing the fp16 matmul to the atom-8 output geometry (surf stride M·16, 0x4050=0x2) HANGS the CNA→DPU path (errno=110). There is no vendor fp16-matmul-atom-8 reference — the vendor always writes contiguous and reshapes. (int8/int16 matmuls DO write atom-8 via set_i16_out; fp16 does not.)

3. The DPU RDMA has NO independent SRC line/surf stride

Only 0x5040 EW_SURF_STRIDE (operand), 0x5048 SRC_DMA_CFG.LINE_NOTCH, 0x504c SURF_NOTCH. So the DPU/SDP cannot read a raw contiguous [M][N] as its surface layout — confirmed exhaustively (operand=1.0 isolation, rows-as-WIDTH and rows-as-HEIGHT, per-channel and per-element): the SDP reads row0/line0 bit-exact but never applies a per-row skip on raw contiguous. The NOTCH only adjusts within an already-reshaped buffer.

4. ★ The CNA (matmul activation reader) DOES read strided / non-contiguous input

synth sets 0x107c CNA_DMA_CON1.LINE_STRIDE = K/8 (per-row stride in 8-channel/16-byte surface units). Overriding it to pitch/8 and storing the activation at row pitch S>K (junk 0xdead padding between rows) → the matmul reads only the real [M][K] rows at the strided offsets, 512/512 bit-exact (probe ORK_F16_ROWPITCH). The matmul can consume arbitrary row-strided / non-contiguous activations directly — the DENSIFY lever: permuted-Q and KV-cache-view attention operands can be fed to the matmul WITHOUT a CPU gather/densify. (0x1080 SURF_STRIDE = channel-group stride, for the K-dim, untested here.)

5. The vendor per-channel scale is matmul → separate SDP (NOT fused); reshape = a CNA permutation-conv

Captured a vendor Gemm × [1,1,N] (build_gemm_mul.py, Colima VM). The matmul task is a PLAIN fp16 matmul (BS/BN/EW all bypassed, contiguous out); the scale is a SEPARATE fp16 EW-mul SDP (task13: EW_CFG=0x20800384, ERDMA=0x8000000a per-element, FEATURE_MODE=0x17849 FLYING_MODE). Between them, small enable=0xd CNA convolutions with a permutation weight (WEIGHT_KERNELS=64, WEIGHT_WIDTH=8) do the contiguous→atom-8 reshape, reading contiguous via the CNA LINE_STRIDE (0x107c) + SURF_STRIDE (0x1080). So the on-chip fp16 "fused matmul+per-channel-EW" is NOT a vendor datapath (our fused attempt hung).

Consequences / next

  • The diagonal-matmul close pays O(M·N²) for purity; the SDP close is O(M·N) with a tiny repack. Pick per use.
  • The CNA strided-read (#4) is the reusable lever for attention: feed permuted-Q / KV-cache views straight to the bmm via 0x107c, removing the CPU densify (the 2026-07-13 attention-on-NPU marshal).
  • Registers belong on the ISA Reference; feat/attn-primitives, probes mm_perchan_f16{,_diag}_probe, f16_mm_f16out_probe (ORK_F16_ROWPITCH), mul_perchan_f16_contig_probe.

Densify-drop primitive (applies the CNA strided-read)

ork_npu_probe_f16_stridedA — a fp16 matmul reading its activation DIRECTLY from a DMA buffer at row pitch apitch>K via CNA LINE_STRIDE (0x107c=apitch/8), NO host→Af gather. Validated BIT-EXACT on attention-shaped cases (M=queries, K=head_dim, N=out, apitch = strided into a wider tensor): {8,32,64,p64}/{16,64,128,p256}/{32,128,256,p512} = 512/2048/8192 exact. The permuted-Q / KV-cache-view sits strided in a DMA buffer and the matmul reads the view IN PLACE — dropping the CPU densify (bmm_gather_f16, ~16% of the pp512 attention profile) that ork_bmm_fp16_strided does today. Fork Stage-2 wiring: stage KV-cache/Q in an ork_dma_alloc buffer + pass the ggml nb[] row stride as apitch. tools/stridedA_bmm_probe.c.