Exp 2026 07 10 Forward Softmax and RKNPU Chaining

2026-07-10 — Forward softmax on-NPU + the RKNPU multi-task hardware-chaining rules

Synthesis of the forward-softmax reverse-engineering (FWD_SOFTMAX_RE_WIP.md). Two deliverables:

1. Forward softmax runs on the NPU — the vendor's 9-task graph, replayed as ONE hardware-chained submit, bit-exact (9/9), ~471 µs (ork_npu_replay_softmax_f16, tools/re/softmax_replay.c, ork-driver main).
2. The reusable RKNPU multi-task chaining rules — including the ping-pong / LUT-SRAM-commit timing rule — that any future on-NPU op-graph (fused attention, norms, activations) needs. These generalise beyond softmax.

Standalone softmax is still submit-floor-bound (~471 µs vs ~µs CPU); the production win for these ops is fusing into an adjacent matmul's SDP output stage (Exp 2026-07-10 Fused Output Activation No-Crossing Chain, M4.8). This page is the chaining RE that a multi-op on-NPU graph depends on.

1. The forward-softmax op-graph (from the vendor capture)

nn.Softmax (rknn-toolkit2, .239) compiles to a 9-task graph per pass (captured op_softmax.rknn → rknpu_dump.so → decode_reg). Safe-softmax dataflow, buffers = one contiguous IOVA image (0xfffae000 out · 0xfffb6000 in · 0xfffbe000 scratch · 0xffff8000 regcmd+weights):

task	enable	kind	role
0,1	0x18	SDP 0x53+0x4088	row-max (input in 2 halves)
2	0xd	matmul	broadcast max
3	0x18	SDP	x − max
4	0x18	LUT-load (1097 reg)	stream exp LUT → DPU SRAM
5	0x18	SDP	exp (reads the LUT)
6	0xd	matmul (exp·ones)	Σ (the sum reduction)
7	0xd	matmul	1/Σ broadcast
8	0x18	SDP	exp·(1/Σ) → output

Softmax discriminator regs (vs a plain reduce): 1001:0x4040=0x53 (reduce/accumulate mode), 0x4070=0x383, 0x4088=0x40000000 (exp-normalize, softmax-unique), 2001:0x5044 reduction descriptor. Geometry maps exactly onto attention scores [M,N]: 0x4058 = 0x403c = 2001:0x5014 = N−1 (reduction/key length), 0x4030 = min(M,128)−1 (query rows, tiles at 128). On-NPU for N up to 1024 (refutes the CDP LRN n≤9 cap — softmax rides a matmul-class reduction, not the fixed-function LRN unit). reducemax fell to CPU; softmax/ reducemean/lrn stay on-NPU.

2. RKNPU multi-task hardware chaining — how it works (vendor rknpu_job_commit_pc)

The kernel programs the hardware from the FIRST task only, then the hardware walks the chain:

• PC_DATA_ADDR = first_task->regcmd_addr; PC_DATA_AMOUNT from first_task->regcfg_amount; INT_MASK from the LAST task; PC_TASK_CONTROL = task_number (bounds the walk).
• Per task, the HW reads the in-regcmd chain descriptor: 0101:0x0010 = next task's regcmd addr, 0101:0x0014 = next task's register-AMOUNT = (n+3)/2 (== kernel (n+EXTRA+scale-1)/scale-1, EXTRA=4, RK3588 pc_data_amount_scale=2; verified 69→0x24, 1097→0x226). The chain descriptor sits at word regcfg_amount*2 (right after the counted regs). run_chain_i8's hardcoded 0x0014=0x0037 is this rule for its uniform 108-reg tasks.

3. The chaining RULES (what it took to chain the softmax; reusable)

Chaining an arbitrary op-graph (not just uniform matmuls) requires, in order of how load-bearing they were:

1. 0101:0x0014 next-amount must be correct =(next_regcfg_amount+3)/2, recomputed for the ACTUAL next task (the captured value only matches capture order). Zeroing it → HW reads next-amount 0 → hangs at the first transition.
2. Task stride is CONTENT-driven, not fixed alignment. run_chain_i8's uniform 896-byte stride works only because an int8-matmul regcmd is always 108 regs (shape lives in reg values, not count). A mixed graph uses variable spacing sized to each task's regcmd (the LUT table size sets task4's slot). Requirement: slot ≥ regcmd size + 0x0010 aimed at where the next task actually sits. (Alignment per se was a red herring.)
3. Contiguous single-buffer image + single-delta rebase. The vendor's 5 buffers are ONE contiguous IOVA block, and task4's LUT-load writes INTO the regcmd region (0x4020=0xffffab00 = task0's freed regcmd space). Splitting into separate IN/OUT/SCR/LUT/regcmd buffers with per-buffer rebase BREAKS the LUT chain. Fix: allocate ONE buffer for the whole image, place every task/buffer at its captured relative offset, and blanket-rebase every in-image reference by a single delta (BIG.dma − 0xfffae000) — incl. 1001:0x4110 (the exp-LUT read ptr into the h2 staging area; whitelisting addr-regs and excluding it regressed the chain).
4. ★ PING-PONG OFF for chains containing a LUT-load (or any un-banked side effect) — the SRAM-timing rule. Submit flags bit RKNPU_JOB_PINGPONG (1<<2; the 0x4 in the common 0x5) double-buffers the register banks (NVDLA PRODUCER/CONSUMER groups) and swaps to the next task the instant the current task's REGISTER CONFIG is done. That pipelines register-config-only tasks (uniform matmuls) — but a LUT-load streams ~1000 entries into DPU LUT SRAM, a long-latency side effect NOT in the register bank, so ping-pong advances to the exp consumer before the LUT SRAM commits → the exp reads a half-written LUT → a non-deterministic stall (flaky task-counter 4↔5; fewer tasks stalling earlier than more = the race giveaway). Submit with flags = RKNPU_JOB_PC = 0x1 (not 0x5) → the HW fully completes each task, incl. the SRAM commit, before advancing. The vendor's captured softmax uses flags=0x1; run_chain_i8 safely keeps 0x5 because its tasks are register-config-only. This is the answer to the "SRAM commit timing" question: the mechanism WAS a commit race, but the fix was disabling the ping-pong that skips the commit wait — not adding a barrier.

4. Result + honest scope

ork_npu_replay_softmax_f16 default = the single-submit hardware chain (contiguous image, 0x0014 recompute, ping-pong off): 9/9 bit-exact, ~471 µs, 5× the per-task path (ORK_SM_PERTASK=1 for the per-task fallback; ORK_SM_1SUBMIT/ORK_SM_NTASK/ORK_SM_SPLIT preserve the RE bisect paths). It replays the captured geometry (N=64) — parameterising for arbitrary (M,N) uses the §1 register map; the uniform-input check (→ 1/N) is layout-agnostic, varied-input needs the NPU cube layout. As a standalone op it's still submit-floor-bound, so production still prefers the fused output-stage path (M4.8). The durable value is the chaining rules (§2–§3) — the recipe for driving ANY multi-task on-NPU op-graph, which the fused/pipelined future depends on. Sources: the mainline accel/rocket driver + vendor rknpu kernel (rknpu_job_commit_pc) — see AGENTS.md §4.