Fix CUDA batched GEMV grid overflow

[jasp-nerd]

Proposed changes

Fixes #3858.

CUDA batched GEMV used grid.y as a linear batch index, so valid batches above
65,535 exceeded the CUDA grid-dimension limit before kernel execution.
gemv_gather used the same launch geometry.

This change splits the batch dimension over grid.y and grid.z, reconstructs
the logical batch index in both kernels, and ignores the one padded slot when
the rectangular grid is not exact. It also keeps empty gathered outputs from
reaching a zero-sized launch.

The regression tests cover ordinary and gathered GEMV at batch sizes 65,535,
65,536, and 65,537 using small analytical inputs, plus an empty gather.

Validation:

pre-commit hooks on changed files: passed
Python test syntax: passed
independent grid-mapping and value-oracle checks: passed
CPU C++ suite: 229 test cases, 3,070 assertions passed

This machine does not have a CUDA toolkit or NVIDIA GPU, so CUDA compilation
and execution are left to the repository CI matrix.

I used an AI coding assistant to help research and prepare this change; I
remain responsible for the submitted patch.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (not needed for this backend-only fix)

[zcbenz]

Sorry but we do not review code that never got tested.

[sashko-zakharchuk]

@jasp-nerd I have an RTX 5050 (sm_120) and verified this patch on it (CUDA 12.8), since it was closed for lack of a hardware run.

It applies cleanly on current main. Before the patch, mx.matmul(mx.zeros((65536, 225, 4)), mx.zeros((65536, 4, 1))) aborts with cudaGraphAddKernelNode ... invalid argument (with graphs off, cudaLaunchKernelExC ... invalid configuration argument), while 65535 and a square (65536, 8, 8) @ (65536, 8, 8) are fine. With the patch, batch 65535 / 65536 / 65537 / 200000 all run and match the CPU result for both matmul and gather_mm, and the empty gather_mm case returns an empty array instead of crashing. Both added tests pass.

Would it be possible to reopen this? Happy to help if anything else is needed.

[jasp-nerd]

Thanks so much for testing this, @sashko-zakharchuk! @zcbenz, would this hardware verification be enough to reopen the PR?

[zcbenz]

LGTM