cupy bug can cause compressed_cderi_j_only to crash

I'm running a semilocal PBC KRKS on an H200 like so:

from pyscf.pbc import gto                                                                    
from gpu4pyscf.pbc.dft import BeckeGrids
cell = gto.Cell()
a, atom = gto.fromfile('si.xyz')
cell.a = a
cell.atom = atom
cell.basis = 'cc-pvdz'
cell.verbose = 10
cell.max_memory = 30000
cell.build()


kpts = cell.make_kpts([6,6,6])
mf = cell.KRKS(xc='pbe', kpts=kpts).to_gpu().density_fit()
mf.grids = BeckeGrids(cell)
mf.run()

where si.xyz contains

2                                                                                            
Lattice="0.0 2.715 2.715 2.715 0.0 2.715 2.715 2.715 0.0" Properties=species:S:1:pos:R:3 pbc="T T T"
Si       0.00000000       0.00000000       0.00000000
Si       1.35750000       1.35750000       1.35750000

Unfortunately, I got a CUDA_ERROR_ILLEGAL_ADDRESS. (cupy 13.4.1 and cupy 13.6.0, cuda 12.6)

Stack trace with CUDA_LAUNCH_BLOCKING=1:

  File "/home/cgh42/src/gpu4pyscf/gpu4pyscf/pbc/df/df_jk.py", line 68, in get_j_kpts                              
    mydf.build(j_only=True, kpts_band=kpts_band)                                   
  File "/home/cgh42/src/gpu4pyscf/gpu4pyscf/pbc/df/df.py", line 142, in build                                                    
    self._cderi, self._cderip, self._cderi_idx = rsdf_builder.build_cderi(                           
                                                 ^^^^^^^^^^^^^^^^^^^^^^^^^                                         
  File "/home/cgh42/src/gpu4pyscf/gpu4pyscf/pbc/df/rsdf_builder.py", line 89, in build_cderi                                     
    cderi, cderip, cderi_idx = compressed_cderi_j_only(                                            
                               ^^^^^^^^^^^^^^^^^^^^^^^^                                           
  File "/home/cgh42/src/gpu4pyscf/gpu4pyscf/pbc/df/rsdf_builder.py", line 341, in compressed_cderi_j_only                        
    multi_gpu.run(proc, non_blocking=True)                                                                        
  File "/home/cgh42/src/gpu4pyscf/gpu4pyscf/lib/multi_gpu.py", line 36, in run
    return [func(*args, *kwargs)]                                                                  
            ^^^^^^^^^^^^^^^^^^^^                                                                     
  File "/home/cgh42/src/gpu4pyscf/gpu4pyscf/pbc/df/rsdf_builder.py", line 319, in proc                             
    j3c = aux_coeff.T.dot(j3c.sum(axis=2).T, out=j3c_buf)  
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^                                          
  File "cupy/_core/core.pyx", line 1810, in cupy._core.core._ndarray_base.dot                     
  File "cupy/_core/_routines_linalg.pyx", line 536, in cupy._core._routines_linalg.dot                             
  File "cupy/_core/_routines_linalg.pyx", line 626, in cupy._core._routines_linalg.tensordot_core                 
  File "cupy/_core/_routines_linalg.pyx", line 763, in cupy._core._routines_linalg.tensordot_core_v11
  File "cupy_backends/cuda/libs/cublas.pyx", line 1426, in cupy_backends.cuda.libs.cublas.gemmEx
  File "cupy_backends/cuda/libs/cublas.pyx", line 1431, in cupy_backends.cuda.libs.cublas.gemmEx
  File "cupy_backends/cuda/libs/cublas.pyx", line 511, in cupy_backends.cuda.libs.cublas._setStream
  File "cupy_backends/cuda/libs/cublas.pyx", line 491, in cupy_backends.cuda.libs.cublas.setStream
  File "cupy_backends/cuda/api/runtime.pyx", line 940, in cupy_backends.cuda.api.runtime.streamIsCapturing
  File "cupy_backends/cuda/api/runtime.pyx", line 146, in cupy_backends.cuda.api.runtime.check_status
cupy_backends.cuda.api.runtime.CUDARuntimeError: cudaErrorIllegalAddress: an illegal memory access was encountered

The bug is in j3c.sum(axis=2). When the batch size is large, cupy.sum crashes. (https://github.com/cupy/cupy/issues/9780)
As a workaround, I am limiting the batch size manually.
Do more experienced users have suggestions for dealing with this?

Start in gpu4pyscf/pbc/df/rsdf_builder.py at compressed_cderi_j_only and inspect the j3c.sum(axis=2) call around line 319. Reproduce the semilocal PBC KRKS example with CUDA_LAUNCH_BLOCKING=1, then verify the large-batch path completes without the CuPy illegal-address error. The issue identifies CuPy issue 9780 and manual batch limiting as relevant context.