From a5ee202a76277cb2c945323afdf495cebb87718b Mon Sep 17 00:00:00 2001 From: zc62 <17048823+zc62@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:44:19 -0500 Subject: [PATCH 1/2] df: bound CUDA tasks for integral builds, gradients and Hessians Before this change, the three-center builders and derivative contractions allow up to 4B orbital pairs per task, where B is the kernel capacity with s-type auxiliary shells. Each task also spans an auxiliary-shell group. Large groups and expensive angular channels produce long indivisible tasks, leaving workers idle as a launch finishes. This commit changes pair limits to min(4B, max(B//4, 16)) in CDERI builds, the shared first-derivative evaluator, restricted/unrestricted ground-state and TD J/K gradients, and restricted/unrestricted DF Hessian energy contractions. The smaller pair ranges expose more independently schedulable work. Each path uses its own kernel capacity; pair ordering and packed auxiliary layouts remain unchanged. Whole-auxiliary CDERI builds previously assign an entire homogeneous auxiliary group to each task, potentially spanning hundreds of shells. This commit also subdivides those groups into at most 16 shells when aux_batch_size is None, following df.j_engine_3c2e.Int3c2eOpt.contract_auxvec. The subdivisions are carried through auxiliary sorting and multi-device clone metadata. Explicit auxiliary batches and derivative auxiliary grouping retain their existing layouts. B/4 is a heuristic balancing parallel work against task setup. For B=128, 32 pairs times 16 auxiliary shells give 512 shell triples. Kernel capacities of 8 to 128 process that task in 64 to 4 passes. Pair subdivision also shortens derivative tasks that retain whole auxiliary groups. Further subdivision increases queue claims and may increase repeated setup and partially filled passes. The 16-pair floor is a separate heuristic limiting small-task setup cost, with the original 4B cap preserved when it is below 16. For the g-g pair in aug-cc-pVQZ, B=32: the floor raises B/4 from 8 to 16 pairs. Sweeps of floors 4, 8, 16 and 32 show low sensitivity in complete builds and gradients, supporting 16 as an intermediate choice. The 16-shell auxiliary limit follows the existing direct-J implementation. --- gpu4pyscf/df/grad/rhf.py | 4 +++- gpu4pyscf/df/grad/tdrhf.py | 8 ++++++-- gpu4pyscf/df/grad/uhf.py | 4 +++- gpu4pyscf/df/hessian/rhf.py | 10 ++++++++-- gpu4pyscf/df/hessian/uhf.py | 5 ++++- gpu4pyscf/df/int3c2e_bdiv.py | 16 +++++++++++++--- 6 files changed, 37 insertions(+), 10 deletions(-) diff --git a/gpu4pyscf/df/grad/rhf.py b/gpu4pyscf/df/grad/rhf.py index 6210d4d7e..9603261c0 100644 --- a/gpu4pyscf/df/grad/rhf.py +++ b/gpu4pyscf/df/grad/rhf.py @@ -147,8 +147,10 @@ def _jk_energy_per_atom(int3c2e_opt, dm, j_factor=1, k_factor=1, hermi=0, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc diff --git a/gpu4pyscf/df/grad/tdrhf.py b/gpu4pyscf/df/grad/tdrhf.py index 1bbf11fec..81f2158dd 100644 --- a/gpu4pyscf/df/grad/tdrhf.py +++ b/gpu4pyscf/df/grad/tdrhf.py @@ -132,8 +132,10 @@ def _jk_energy_per_atom(int3c2e_opt, dms, j_factor=None, k_factor=None, hermi=0, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc @@ -542,8 +544,10 @@ def _jk_energies_by_dm_factors(int3c2e_opt, dm_factors, j_factor, k_factor, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc diff --git a/gpu4pyscf/df/grad/uhf.py b/gpu4pyscf/df/grad/uhf.py index 67da9e48f..af0ad29a4 100644 --- a/gpu4pyscf/df/grad/uhf.py +++ b/gpu4pyscf/df/grad/uhf.py @@ -124,8 +124,10 @@ def _jk_energy_per_atom(int3c2e_opt, dm, j_factor=1, k_factor=1, hermi=0, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) aux_loc = auxmol.ao_loc diff --git a/gpu4pyscf/df/hessian/rhf.py b/gpu4pyscf/df/hessian/rhf.py index 091f550e8..fbbc0f312 100644 --- a/gpu4pyscf/df/hessian/rhf.py +++ b/gpu4pyscf/df/hessian/rhf.py @@ -188,8 +188,11 @@ def proc(aux_batch_iter): laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) l_ctr_aux_offsets = np.append(0, np.cumsum(auxmol.l_ctr_counts)) @@ -586,8 +589,11 @@ def proc(): laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) diff --git a/gpu4pyscf/df/hessian/uhf.py b/gpu4pyscf/df/hessian/uhf.py index 593407991..f31960888 100644 --- a/gpu4pyscf/df/hessian/uhf.py +++ b/gpu4pyscf/df/hessian/uhf.py @@ -163,8 +163,11 @@ def proc(): laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) l_ctr_aux_offsets = np.append(0, np.cumsum(auxmol.l_ctr_counts)) diff --git a/gpu4pyscf/df/int3c2e_bdiv.py b/gpu4pyscf/df/int3c2e_bdiv.py index 7832375ef..070c2c9d6 100644 --- a/gpu4pyscf/df/int3c2e_bdiv.py +++ b/gpu4pyscf/df/int3c2e_bdiv.py @@ -34,7 +34,7 @@ from gpu4pyscf.lib import multi_gpu from gpu4pyscf.gto.mole import ( PTR_BAS_COORD, SortedMole, RysIntEnvVars, PBCIntEnvVars, - extract_pgto_params, groupby) + extract_pgto_params, groupby, _split_l_ctr_groups) from gpu4pyscf.scf.jk import ( _nearest_power2, _scale_sp_ctr_coeff, _cache_q_cond_and_non0pairs, _check_rsh_factors, SHM_SIZE, libvhf_rys) @@ -175,8 +175,11 @@ def int3c2e_evaluator(self, ao_pair_batch_size=None, aux_batch_size=None, bas_ij_cache = {k: v for k, v in mol.iter_pair_by_l(self.bas_ij_cache)} else: bas_ij_cache = self.bas_ij_cache + # Bound pair work alongside the auxiliary-shell task subdivision. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - bas_ij_cache, nsp_per_block[0]*4) + bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], bas_ij_cache, cart) nao_pair = ao_pair_loc[-1].get() @@ -210,6 +213,10 @@ def int3c2e_evaluator(self, ao_pair_batch_size=None, aux_batch_size=None, uniq_l_ctr_aux = auxmol.uniq_l_ctr aux_loc = auxmol.ao_loc if aux_batch_size is None: + # Follow df.j_engine_3c2e.Int3c2eOpt.contract_auxvec's 16-shell tasks. + uniq_l_ctr_aux, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + l_ctr_aux_offsets = np.append(0, np.cumsum(aux_counts)) ksh_offsets_cpu = l_ctr_aux_offsets aux_splits = [0, len(ksh_offsets_cpu)-1] else: @@ -785,8 +792,11 @@ def _int3c2e_ip1_evaluator(int3c2e_opt, scheme, batch_size, laux = auxmol.uniq_l_ctr[:,0].max() shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() + # Bound serial work in each derivative-integral task. + pair_batch_size = np.minimum(nsp_per_block[0]*4, + np.maximum(nsp_per_block[0]//4, 16)) bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( - int3c2e_opt.bas_ij_cache, nsp_per_block[0]*4) + int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) nao_pair = int(ao_pair_loc[-1].get()) From 7967987ff56eea597254fb33c91f70ba55112ec4 Mon Sep 17 00:00:00 2001 From: zc62 <17048823+zc62@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:23:29 -0500 Subject: [PATCH 2/2] df: subdivide whole auxiliary groups in Coulomb-only contractions Before this change, the Coulomb-only derivative and Rys potential contractions assign each CUDA block an orbital-pair range and a whole homogeneous auxiliary-shell group. Large fitting bases lengthen the serial auxiliary loop inside each block. A long auxiliary group stays assigned to one block even when other SMs become idle. Subdividing that group allows its work to run across multiple SMs. This commit splits each auxiliary group into ranges of at most 16 shells with _split_l_ctr_groups, following Int3c2eOpt.contract_auxvec in df/j_engine_3c2e.py. It applies this subdivision to ground-state and TD Coulomb gradients, pure-J Hessians and multi-density direct SCF J builds. --- gpu4pyscf/df/grad/rhf.py | 7 +++++-- gpu4pyscf/df/grad/tdrhf.py | 12 +++++++++--- gpu4pyscf/df/hessian/rhf.py | 6 +++++- gpu4pyscf/df/int3c2e_bdiv.py | 5 ++++- 4 files changed, 23 insertions(+), 7 deletions(-) diff --git a/gpu4pyscf/df/grad/rhf.py b/gpu4pyscf/df/grad/rhf.py index 9603261c0..890e9b64b 100644 --- a/gpu4pyscf/df/grad/rhf.py +++ b/gpu4pyscf/df/grad/rhf.py @@ -27,7 +27,7 @@ SHM_SIZE, LMAX, L_AUX_MAX, THREADS, libvhf_rys, Int3c2eOpt, int2c2e) from gpu4pyscf.df import df from gpu4pyscf.df.df_jk import factorize_dm, _DFHF -from gpu4pyscf.gto.mole import SortedMole +from gpu4pyscf.gto.mole import SortedMole, _split_l_ctr_groups from gpu4pyscf.scf.jk import _check_rsh_factors __all__ = ['Gradients'] @@ -284,7 +284,10 @@ def _j_energy_per_atom(int3c2e_opt, dm, hermi=0, auxbasis_response=True, verbose shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( int3c2e_opt.bas_ij_cache, nsp_per_block[0]*16) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs diff --git a/gpu4pyscf/df/grad/tdrhf.py b/gpu4pyscf/df/grad/tdrhf.py index 81f2158dd..f77d8e6e7 100644 --- a/gpu4pyscf/df/grad/tdrhf.py +++ b/gpu4pyscf/df/grad/tdrhf.py @@ -27,7 +27,7 @@ from gpu4pyscf.df import df from gpu4pyscf.tdscf import rhf as tdrhf from gpu4pyscf.grad import tdrhf as tdrhf_grad -from gpu4pyscf.gto.mole import SortedMole +from gpu4pyscf.gto.mole import SortedMole, _split_l_ctr_groups from gpu4pyscf.scf.jk import _check_rsh_factors __all__ = ['Gradients'] @@ -270,7 +270,10 @@ def _j_energy_per_atom(int3c2e_opt, dms, j_factor, hermi=0, verbose=None): shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( int3c2e_opt.bas_ij_cache, nsp_per_block[0]*16) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs @@ -695,7 +698,10 @@ def _j_energies_per_atom(int3c2e_opt, dm_pairs, j_factor, shm_size_max = shm_size[:laux+1,:lmax+1,:lmax+1].max() bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( int3c2e_opt.bas_ij_cache, nsp_per_block[0]*16) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs diff --git a/gpu4pyscf/df/hessian/rhf.py b/gpu4pyscf/df/hessian/rhf.py index fbbc0f312..0dcb64069 100644 --- a/gpu4pyscf/df/hessian/rhf.py +++ b/gpu4pyscf/df/hessian/rhf.py @@ -40,6 +40,7 @@ from gpu4pyscf.hessian import rhf as rhf_hess from gpu4pyscf.hessian.rhf import _hcore_energy, _aggregate_to_atoms from gpu4pyscf.lib import multi_gpu +from gpu4pyscf.gto.mole import _split_l_ctr_groups num_devices = multi_gpu.num_devices @@ -595,7 +596,10 @@ def proc(): bas_ij_idx, shl_pair_offsets = mol.aggregate_shl_pairs( int3c2e_opt.bas_ij_cache, pair_batch_size) ao_pair_loc = get_ao_pair_loc(mol.uniq_l_ctr[:,0], int3c2e_opt.bas_ij_cache) - ksh_offsets_cpu = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + ksh_offsets_cpu = np.append(0, np.cumsum(aux_counts)) ksh_offsets_gpu = cp.asarray(ksh_offsets_cpu+mol.nbas, dtype=np.int32) int3c2e_envs = int3c2e_opt.int3c2e_envs diff --git a/gpu4pyscf/df/int3c2e_bdiv.py b/gpu4pyscf/df/int3c2e_bdiv.py index 070c2c9d6..4e0cbfb10 100644 --- a/gpu4pyscf/df/int3c2e_bdiv.py +++ b/gpu4pyscf/df/int3c2e_bdiv.py @@ -417,7 +417,10 @@ def contract_auxvec(self, auxvec): self.bas_ij_cache, nsp_per_block[0]*4) gout_stride = cp.asarray(gout_stride, dtype=np.int32) - l_ctr_aux_offsets = np.append(0, np.cumsum(auxmol.l_ctr_counts)) + # Split auxbasis into small batches for load balance + _, aux_counts = _split_l_ctr_groups( + auxmol.uniq_l_ctr, auxmol.l_ctr_counts, [16]*(laux+1)) + l_ctr_aux_offsets = np.append(0, np.cumsum(aux_counts)) ksh_offsets = cp.asarray(l_ctr_aux_offsets + mol.nbas, dtype=np.int32) log.debug1('sp_blocks = %d, ksh_blocks = %d, shm_size = %d B', len(shl_pair_offsets)-1, len(ksh_offsets)-1, shm_size_max)