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...