rosenblatt/docs/architecture.md

Architecture — CPU + NPU mix for llama.cpp on RK3588

The split

llama.cpp's compute graph is built around ggml ops. We don't replace llama.cpp's whole engine — we register a new device backend (in ggml's ggml-backend abstraction) named rknpu and selectively offload the ops that are worth the round-trip cost.

Goes to NPU (heavy, dense, INT8-friendly)

• MUL_MAT (matrix-matrix multiply) — the workhorse, dominates wall time. Both attention Q @ K^T, attn @ V, and FFN up_proj, down_proj, gate_proj are this shape.
• MUL_MAT_ID (MoE-style mixture matmul) — when we eventually try a mixture-of-experts model. Phase-1+ scope.

Stays on CPU (small, op-specific, or per-token)

• Embedding lookup (GET_ROWS) — random-access gather, NPU has no fast path
• RMS_NORM / LAYER_NORM — per-token reduction + element-wise
• ROPE — small, per-head, lots of trig
• SOFT_MAX — small, per-head
• Activations (SILU, GELU) — element-wise, cheap on NEON
• SCALE, ADD, MUL (element-wise) — cheap on NEON
• Sampling, KV cache update, tokenization — entirely host

KV cache: open question

Two options:

1. CPU-resident: lives in normal Linux memory; NPU pulls activations from CPU and pushes results back per layer.
2. NPU-resident: allocated in dmabuf, NPU reads K/V across layers without round-trips. Cheaper per-layer but constrains model size to NPU-accessible memory.

Phase 5 (Plan) picks based on Phase-3 (Analyze) findings on the DMA cost.

---

Memory mapping

ggml's ggml_backend_buffer_t abstracts the buffer pool. We implement:

• alloc_buffer(size) → allocate a dmabuf of size bytes
• free_buffer(buffer) → release dmabuf
• set_tensor / get_tensor → CPU → NPU memcpy
• cpy_tensor → device-internal copy

The dmabuf approach lets us share buffers between CPU producer (e.g. embedding lookup output) and NPU consumer (matmul) without an extra copy — mmap on CPU side, DMA-import on NPU side.

If Tomeu's accel uAPI uses dmabuf natively, we follow that. If it doesn't, we go through /dev/dri/renderD* with a thin shim.

---

Quantization strategy

llama.cpp ships Q4_K_M as the default for ~2B models. Q4_K_M is a 4-bit weight quantization with per-group scale + min, no per-channel scale. The NPU expects INT8 (or INT16) tensors with per-channel scale factors.

Two paths:

1. Dequantize on CPU per-layer: unpack Q4_K_M → INT8 right before the matmul; ship INT8 to NPU. Adds a per-layer CPU pre-pass.
2. Dequantize once at load time: unpack the entire weight tensor to INT8 + per-channel scales at model-load. Permanent ~2x memory cost (Q4_K_M is ~5 bits/weight effective; INT8 is 8 bits/weight), but no per-layer CPU work.

Phase-1 choice: (2) — straightforward, makes the NPU path the only thing happening at inference time, easier to profile. The memory cost on 1.5B is ~1.5 GB INT8 weights vs ~900 MB Q4_K_M — boltzmann has 32 GB, this isn't the constraint.

Phase-2+: revisit (1) if we go for larger models or if INT8 turns out to be quality-loss-meaningful on the small ones.

---

Backend interface — concrete

Mirroring ggml's existing ggml-cuda.cu / ggml-metal.m shape, we add:

ggml-rknpu.h       — public API: ggml_backend_rknpu_init() etc.
ggml-rknpu.c       — backend implementation: device registration, op
                     dispatch table, memory management
ggml-rknpu-ops.c   — per-op kernels: matmul tiled to NPU's preferred
                     shape, INT8 quant pre-pass

In llama.cpp/ggml/src/ggml-rknpu/. Out-of-tree initially; if upstream-acceptable, send for review after Phase 8.

---

Failure handling

llama.cpp's backend abstraction already supports falling back to CPU on per-op basis — ggml_backend_dev_supports_op(). We declare MUL_MAT supported for INT8 / FP16 inputs with the right shape constraints, and let the framework route everything else to CPU.

If the NPU driver returns an error mid-inference (timeout, DMA fence wait fail, etc.), the strategy is abort the inference, log, return error to caller. We don't try to silently fall back to CPU mid-stream because the state would be corrupted (NPU may have partially written to the dmabuf).

---

Phase-1 milestone

A npu-probe userspace binary that:

1. Opens the NPU device (whatever the mainline path is — likely /dev/accel/accelN or /dev/dri/renderD*)
2. Allocates two small INT8 input tensors + one output (e.g. 64x64)
3. Submits a matmul via the uAPI
4. Waits, reads back, compares to a CPU reference

This proves the substrate is alive before we touch llama.cpp. If it doesn't work, we're back in kernel-driver land, not llama.cpp land.