Moving a running model's layers between GPUs mid-reply (8 ms, no re-prefill)

I've been working on an inference engine called jitllm. The idea came from my background in operating systems and JIT compilers: generate kernels for whatever hardware is there, and treat layers a bit like processes the engine can place and move around.

In the video, a cloud VM with T4s: mid-reply, all layers of Llama-3.2-3B move from one GPU to the other and back. And my puny laptop where the model is split across CPU, CUDA and Vulkan (iGPU), rearranged while it's generating. The KV cache is paged and moves with the layers, so the prompt isn't re-processed and the output matches a run that didn't move.

What excites me is the room the design leaves. Generated kernels mean every improvement reaches every machine, and the engine can fuse and tune for hardware it has never seen. Paging and relocation turn placement into scheduling: models bigger than memory, mismatched GPUs, work that shifts as conversations grow.

Code -> https://github.com/jitllm/jitllm