Why not use an explicit `max_batch_size_x` and `max_batch_size_y`  and use a default of `1000` or `5000` for each for instance? I think that would have made it more intuitive for the user to control the memory usage: `1000x1000` is 8MB in double precision. `5000x5000` is 200MB. AFAIK, most GEMM from optimized blas implementations reach peak performance for sizes around `5000x5000`.