Dealing with type conversion and relayout for NPU backends (how to know the data type during buffer creation/rw)?

[marty1885]

Hi all,

I'm trying to build a NPU backend for the Rochchip RK3588 and Tenstorrent cards. I run into a problem with the buffer allocation interface. These accelerator hardware usually demands activation and weights are laid out in the hardware's native format and have limited data type support. My idea was to perform the processing during tensor read and write (deuanltize the buffer, convert to a format the NPU supports then relayout).

However, ggml_backend_buffer_type_i::alloc_buffer nor ggml_backend_buffer_i::set_tensor tells me the data type of the tensor and/or of the source buffer. Thus it is impossible to perform the processing here.

The best I can think of is to re-layout during graph execution. However this feels wrong to me.

[slaren]

set_tensor does tell you the data type in tensor->type. Similarly, you can use get_alloc_size to determine how much memory needs to be allocated for a tensor.

[slaren]

The backend does not set the size of the buffer, but the size of the buffer is calculated from the calls to get_alloc_size with the tensors that will be allocated on the buffer. The backend only needs to implement get_alloc_size and return the amount of memory it needs for a tensor. In set_tensor you can apply any transformations to the data.

[marty1885]

Thanks. A follow up question but related. If I'm understanding correctly. alloc_buffer allocates a large buffer then ggml_tallocr chunks the allocation. Is it possible to work around this?

The NPUs I'm working with either only exposes a high level allocation API (allocates memory for a tensor and does not allow views into the buffer) or have basically hardware tagged data types (bank size must align with X elements of the underlying data type).

For example, the NPU would generate the wrong output if alloc_buffer I set the bank size to 1KB but the data requires bank multiple of 768 bytes.

[slaren]

Can you link me a reference to this API so that I can understand the problem better?

[marty1885]

For the Rockchip NPU the situation easier. The allocation API returns rknn_tensor_mem which contains the virtual address of the allocated buffer. There is no guarantee offsetting into the virtual address will also write to the corresponding physical memory (undocumented and because of IOMMU), And it has a 4GB limit due to data type. But in reality it can support much larger sizes using the IOMMU.

https://github.com/airockchip/rknn-toolkit2/blob/77b71094e08391c543d9c65fea5f7cf98cc16eee/rknpu2/runtime/Linux/librknn_api/include/rknn_api.h#L391-L399

Tenstorrent side is a bit convoluted. Tenstorrent devices don't do the usual byte addressing. Instead if reads in units of DRAM banks and tiles (32x32 grid of values). First, Tenstorrent's buffer allocation gives you a buffer but does not allow sub-buffering.

https://github.com/tenstorrent/tt-metal/blob/ae7c15b141fd78c5e53af900cbe49dd594fd28df/tt_metal/host_api.hpp#L222

It wouldn't be an issue on GPUs. But I can't read from arbitrary offsets within the device kernel. Their hardware reads and write memory in a unit of tiles. In BufferConfig, size % page_size == 0 and page_size % packed_buffer_size_bytes<dtype>() == 0 in the kernel. Else buffer read becomes unreliable. However, due to support for types like bfloat8 and bfloat4 (similar to Q8 and Q4 in GGML). This leads to the tile size in bytes not nicely matching up with the common 1K or 4K border.

https://github.com/tenstorrent/tt-metal/blob/ae7c15b141fd78c5e53af900cbe49dd594fd28df/tt_metal/impl/buffers/buffer.hpp#L115-L121
https://github.com/tenstorrent/tt-metal/blob/ae7c15b141fd78c5e53af900cbe49dd594fd28df/tt_eager/tensor/tensor_utils.hpp#L64-L80

Kernel side code (read will fail if page_size mismatches with tile_size_bytes and data_format).
https://github.com/marty1885/ttVecAdd/blob/eb1d94f2103759f58988441070a36f3465ef9901/vecadd_kernels/interleaved_tile_read.cpp#L24-L33

[slaren]

For the Rockship NPU you can implement the function get_max_size. This establishes the maximum buffer size, and applications that use ggml_backend_alloc_ctx_tensors_from_buft will respect this limit and allocate multiple buffers if necessary.

For the tenstorrent, using offsets is a fundamental part of the memory allocation strategy of ggml-alloc. What you could possibly do, assuming that it is possible to write custom kernels at all, would be to set get_alignment to as big as it needs to be for every data type, and then pass the buffer offsets of the tensors as a kernel parameter, and write the kernel to respect this offset.

[Cikaros]

Hello @marty1885 , did the branch you pulled support RK3588 at that time?
I want to try to build your project, how do I do it?
This helps me see if it's worth buying OrangePi 5 Max. Thanks!

[marty1885]

My branch is quite old, before GGML officially supports backend. Building should be easy

1. Install RKNPU2 1.6.0 to /usr/{include,lib}
2. Build with cmake .. -DLLAMA_USE_RKNPU2=ON (or something like that, I forgot)
3. Pass -ngl <something larger then 0> to bin/main.

I would wait if the only reason to buy an Orange Pi Max is to run LLMs on the NPU. I am waiting for Tomeu to release the reversed NPU documents and rebuild it against the fully open source driver. Until then I'm focusing on the Tenstorrent port.

[Cikaros]

What you say makes a lot of sense.
Wait for...

[System64fumo]

@marty1885 The open source NPU driver has been out for a bit, Would you be willing to take a look at it?

https://blog.tomeuvizoso.net/2025/07/rockchip-npu-update-6-we-are-in-mainline.html

Here's the output from my orange pi 5 plus running 6.18.0-rc4:

dmesg | grep -i rocket
[    8.688474] [drm] Initialized rocket 0.0.0 for rknn on minor 0
[    8.689048] rocket fdab0000.npu: Rockchip NPU core 0 version: 1179210309
[    8.689671] rocket fdac0000.npu: Rockchip NPU core 1 version: 1179210309
[    8.690189] rocket fdad0000.npu: Rockchip NPU core 2 version: 1179210309

[marty1885]

Not really having the time. The Tenstorrent backend looks very promising and should serve as a good baseline for future NPU backends (which AFAIK there are a few private forks being based on my work). I am planning on upstreaming it end of Q1 2026. And hopefully another talk at FOSDEM 2027.

Tenstorrent is actively supporting OSS developers and users. As well as a very public model of operation. I work for Tenstorrent now. And I would rather spending my energy helping a company that's doing OSS then a company that does everything in private and doesn't even bother supporting 3rd party distro maintainers (wink wink Ubuntu image wink).

[marty1885]

Also from my understanding from RKLLM. It's a massive hack on top of (my) llama.cpp with changes to the execution model in order to utilize the NPU 3 cores. I would have to do the same in order to utilize the NPU to the same efficiency nomatter who's driver I am using. And that patch is unlikely be be upstreamable.