Cross-backend tensor/layer/batch split

[phhusson]

I've been working on cross-backend dynamic tensor/layer/batch split.

The basic idea is that some backends can share their memory and events with other backends, which allows dynamic splits, while keeping async graph.

• layer-split: NPU can compute just some ops, let the GPU do other ops, without breaking the async graph -- this requires cross-backend activation sharing and cross-backend event sharing
• tensor-split: NPU can compute half the row, GPU the other half, then go back to GPU graph, without breaking the async graph
• mode-split [1]: NPU handles big batches (for pp), but let GPU do the decode, and NPU handles speculation when speculator is confident for >= 4 tokens -- this requires cross-backend weight & activation sharing, but is doable without cross-backend event sharing

I've taken NPU/GPU as an example for clarity, but it can actually be done with any backends. It can be done with GPU+CPU (making ggml-cpu async), or ROCm+Vulkan on a single GPU (benefitting of ROCm's better optimized kernels, and of Vulkan faster dispatch)

On Linux, buffer-sharing is implemented using dma-buf, event-sharing is done using dma fences fd. On MacOS buffer-sharing is implemented using IOSurface, and event-sharing IOSurfaceSharedEvent.

I've done crude PoCs where I let GPT-6 Astra run loose on optimizing pp speed on that idea (those have a very horrible code/git history) [2]:

• On bonsai-llama.cpp on Apple Mac Mini M4 16GB, tensor+mode-split GPU + NPU made PQ2_0 go from 60tok/s pp to 115tok/s. (sorry that's not a mainline llama.cpp), without changing decode speed thanks to mode-split
• On Ryzen 8845HS, Qwen3.8-Flash-Next-UD-IQ1_S went from 150tok/s to 330tok/s using GPU/NPU tensor+modesplit, without changing decode speed thanks to mode-split
• On Ryzen 8845HS, ROCm/Vulkan tensor-split got 80% of the speed improvement of vulkan->rocm during pp (100% should be achievable with more work), while retaining vulkan decode speed

tensor+mode-split here means that during pp the we do tensor-split gpu+npu, then during tg it's just gpu.

I've made a tree that I hope is clean enough to understand the concepts, that implements cross-backend dynamic layer/batch split:
https://github.com/phhusson/llama.cpp/commits/33e610a24363acb3acbe4b6159a58f465afa69a1/

This tree doesn't bring any actual improvement, it is really architecture-only. (Well it does bring ROCm pp +Vulkan tg)

The mode-split is cheating the tensor-split scheduler by having the "batch-sized" backend refuse the op. We end-up with --device ROCm0,Vulkan0 -ts 1,0 -ot 'blk\..*=DMA_BUF', and ROCm0 is modified to return support_op false for small batches. This is okay-ish on NPUs, but it's semantically non-sense for ROCm.

There is no tensor-split in this tree. Ideally the scheduler would make an actually informed decision to decide tensor-split, layer-split, and mode-split, but it will require much more scaffolding worth discussing.

An additional remark that isn't visible in this tree: It might be reasonable to modify the GPU backend for better NPU co-work:

• PQ2_0 on ANE repacks the weight in memory as multi-planar (one plane for scales, one plane for qweights), so MTL needs to support that format. No infra change for this: ANE has the priority in --device, so when a weight is loaded it's ANE who decides how it gets loaded; it just needs MTL to know about this format
• AMDXDNA works in bf16, and wastes a lot of cycles when converting to/from fp32, so it actually makes sense to add Vulkan fp32//bf16 conversions. This requires much deeper changes
• AMD XDNA1 is limited to 6 simultaneous kernels to keep the whole graph async, so if a model use 10 different quants, there is a choice of "which quants do the most batched MACs" for layer-split.

I have a rather good understanding of all commits Linux/infra-side, and written their commit messages. I have a very poor understanding of the commits Apple-side, but I think they are useful to see what does cross-backend buffer sharing looks like on another OS. I've written this whole message without any LLM.

Apple-side uses private APIs. The nVidia dma-buf allocator uses an unstable API. Allocating dma-buf from amdgpu's gtt and using it in nvidia proprietary vulkan works, and is fun, but crashes the driver quite easily.

I think this work makes it much easier to integrate NPUs into llama.cpp, because a contributor has to contribute just one working matmul on the NPU to be able to provide speed-ups.

My questions:

• Is this of interest to ggml-org?
• Is this of interest to other people?
• Is there a priority between tensor/layer/batch split?
• Who should control modesplit? Should it be the backend (like my current tree, just fix the semantics of supports_op), should it be the scheduler (who will need information from the backend, maybe preferred_min_batch_size?), or should it be the user? (a --device ROCm0,Vulkan0 -sm batch_size -ts 8,7)
• How to expose to the user tensor+mode-split?
• What kind of split PR should this be? A nice first PR can be ROCm/Vulkan; it is pretty small, it doesn't require event-sharing, it is pretty useful. But it requires the ugly GGML_DMABUFT=1, the ugly GGML_ROCM_MIN_N=8, sync kinda works by accident (well that shouldn't be too hard to fix).

[1] Feel free to rename it in any other way; I tried "batch-split" but this can be understood as splitting on tokens
[2] You can find my various npu shitty codes and some explanations at https://github.com/phhusson/npu_experimentations