Ensemble simulations of small ODEs on GPUs

Hello, I have a system of ODEs that are relatively small (a matrix ODE with matrix sizes of less than 20x20).

\hat{\rho}(t) = \hat{A}(t,\vec{p})\hat{\rho(t)}

I want to solve this ODE for various initial conditions $\hat{\rho}_0$ and parameter vectors $\vec{p}$ but I want to take advantage of the GPU because I would like to simultaneously solve $10^3$ to $10^4$ instances of this ODE at once.

I can define a function that solves the ODE:

def solve_ode(p):
    rhs = lambda t,y, args: right_hand_side_function(t,p)
    
    term = ODETerm(rhs)
    solver = Tsit5()
    saveat = . . .
    stepsize_controller = . . . 
    
    sol = diffeqsolve(term, solver, . . . )
    return sol.ys[-1]

And I can vmap this function for various input parameters $\vec{p}$, but some small benchmarks have shown that it's much faster on CPU than GPU and my guess is just because the ODE system size is so small that it doesn't really take advantage of the GPU.

Would it be better to generate a single ODE function that takes in all possible parameters and evolves the many ODEs in parallel in a block-diagonal way as if they were a single large set of coupled equations? In this case the ODEs will then need to have the same step size, max steps, etc. which would probably mean that some of the small ODEs are being overly-sampled on an extra fine step size that is determined by the stiffest small ODE.

Or is there some better way of dealing with this?

The issue does not name a file, test, or entry point to change. Start by reproducing the vmap benchmark around solve_ode, diffeqsolve, and Tsit5 on CPU and GPU, then compare it with a batched or block-diagonal formulation; done would require an agreed implementation and evidence that ensemble simulations perform as intended.