Yeah, that's why I don't like vmap so much. It just never gets you an optimal algorithm. On CPUs, you don't need to do the same number of steps each solve: you

Yeah, that's why I don't like vmap so much. It just never gets you an optimal algorithm. On CPUs, you don't need to do the same number of steps each solve: you might as well fully decouple all of the solves and max out the threads. Broadcasting DiffEq.solve will do this fully decoupled version, but it's better to parallelize it either with multithreading and distributed (or polyester). That's of course the better thing to do on CPU, and then you don't need to worry about filling kernels or whatnot because you can just oversubscribe if you have to. On GPU vmap isn't as bad at face value, I thought it would be fairly close to optimal, but when we dug in we found it wasn't the right approach for the ODEs we were looking at either. Julia has KernelAbstractions.jl which can do similar code transformations as vmap, but when digging into the CUDA profiler we found it takes a pretty substantial number of trajectories to fill the kernels. There were a few more optimizations we could do to it, but it was reaching its limit rather quickly. Meanwhile, it turns out that if you can just generate .ptx kernels for the ODE, minimize the usage of global memory by keeping all ODE parts locally in registers, and then only couple within warps, you can beat that by about 100x. So some of our other GPU stuff has gone stale for a bit while we've been working to get this approach completed. It was somewhat of a disappointing conclusion because it means we have to really specialize some codes for GPU, but anyways, full details will come probably around the end of summer.