[Bug] nb2d=0 can silently fall back to block size 1 and hang cuBLASMp TRSM in multi-GPU RT-TDDFT

[Stardust0831]

Describe the bug

In a multi-node LCAO RT-TDDFT calculation, the documented default nb2d=0 first selects a block size from nlocal, but may then silently fall back directly to the minimum block size (1 for nspin != 4) when the wavefunction distribution is incompatible with the process grid.

For a Si240 case on 96 MPI ranks / 96 V100 GPUs, this produces a valid-looking 1x1 distributed matrix descriptor. The calculation subsequently stops inside the collective cublasMpTrsm() call. A standalone cuBLASMp reproducer with the same matrix size, process grid, data type, and block size reproduces the stop, whereas explicitly setting nb2d=32 lets the ABACUS calculation complete the TRSM and the remaining propagation.

This report is primarily about the ABACUS block-size selection and fallback. The non-returning cublasMpTrsm() behavior for a 1x1 block appears to be an upstream cuBLASMp 0.9.1 issue, but ABACUS's current fallback selects that problematic and extremely inefficient block size even though a valid larger block (32) exists.

Documentation mismatch

The current manual describes nb2d as follows:

Type: Integer
Description: In LCAO calculations, the Hamiltonian and overlap matrices are distributed across 2D processor grid. This parameter controls the 2D block size for distribution.
Default: 0

Source:

abacus-develop/docs/advanced/input_files/input-main.md

Lines 853 to 857 in ff23741

	### nb2d
	- **Type**: Integer
	- **Description**: In LCAO calculations, the Hamiltonian and overlap matrices are distributed across 2D processor grid. This parameter controls the 2D block size for distribution.
	- **Default**: 0

. The same text and default are registered in

abacus-develop/source/source_io/module_parameter/read_input_item_elec_stru.cpp

Lines 1162 to 1179 in ff23741

	// LCAO
	{
	Input_Item item("nb2d");
	item.annotation = "matrix 2d division";
	item.category = "System variables";
	item.type = "Integer";
	item.description = "In LCAO calculations, the Hamiltonian and overlap matrices are distributed across 2D processor grid. This parameter controls the 2D block size for distribution.";
	item.default_value = "0";
	item.unit = "";
	item.availability = "";
	read_sync_int(input.nb2d);
	item.check_value = [](const Input_Item& item, const Parameter& para) {
	if (para.input.nb2d < 0)
	{
	ModuleBase::WARNING_QUIT("ReadInput", "nb2d should be greater than 0");
	}
	};
	this->add_item(item);

.

The documentation does not explain that 0 is an automatic-selection sentinel, which values are selected, or that ABACUS may replace the selected value with 1. Therefore a user cannot predict that the default will produce a 1x1 descriptor at a particular MPI rank count.

Relevant code path

At commit ff237418db665a9fe3befa7029e3364b085b7987, the automatic value is chosen only from nlocal. In this case nlocal=5280, so nb2d=0 becomes nb2d=64:

abacus-develop/source/source_lcao/LCAO_init_basis.cpp

Lines 23 to 40 in ff23741

	const int nlocal = PARAM.globalv.nlocal;
	int nb2d = PARAM.inp.nb2d;
	// autoset NB2D first
	if (nb2d == 0)
	{
	if (nlocal > 0)
	{
	nb2d = (PARAM.inp.nspin == 4) ? 2 : 1;
	}
	if (nlocal > 500)
	{
	nb2d = 32;
	}
	if (nlocal > 1000)
	{
	nb2d = 64;
	}
	}

The wavefunction distribution then checks whether every process-grid column can receive a band block. With nbands=480, nb=64, and dim1=12, there are only ceil(480 / 64) = 8 band blocks, so this returns 1:

abacus-develop/source/source_basis/module_ao/parallel_orbitals.cpp

Lines 222 to 250 in ff23741

	int Parallel_Orbitals::set_nloc_wfc_Eij(
	const int& N_A,
	std::ofstream& ofs_running,
	std::ofstream& ofs_warning)
	{
	ModuleBase::TITLE("Parallel_Orbitals", "set_nloc_wfc_Eij");
	// for wavefuncton , calculate nbands_loc
	this->nbands = N_A;
	int end_id = 0;
	int block = N_A / nb;
	if (block * nb < N_A)
	{
	block++;
	}
	if (dim1 > block)
	{
	ofs_warning << " cpu 2D distribution : " << dim0 << "*" << dim1 << std::endl;
	ofs_warning << " but, the number of bands-row-block is " << block << std::endl;
	if (nb > 1)
	{
	return 1;
	}
	else
	{
	ModuleBase::WARNING_QUIT("Parallel_Orbitals::set_nloc_wfc_Eij",
	"The number of columns of the 2D process grid exceeds the number of bands. "
	"Try launching the calculation with fewer MPI processes."
	);
	}

The caller does not try an intermediate block size. It immediately replaces 64 with 1:

abacus-develop/source/source_lcao/LCAO_init_basis.cpp

Lines 77 to 85 in ff23741

	int try_nb = pv.init(nlocal, nlocal, nb2d, DIAG_WORLD);
	try_nb += pv.set_nloc_wfc_Eij(PARAM.inp.nbands, GlobalV::ofs_running, GlobalV::ofs_warning);
	if (try_nb != 0)
	{
	// fall back to the minimum size, 1 or 2 (nspin=4)
	const int min_size = (PARAM.inp.nspin == 4) ? 2 : 1;
	pv.set(nlocal, nlocal, min_size, pv.blacs_ctxt);
	try_nb = pv.set_nloc_wfc_Eij(PARAM.inp.nbands, GlobalV::ofs_running, GlobalV::ofs_warning);
	}

The resulting matrix descriptor is later passed to the RT-TDDFT triangular solve here:

abacus-develop/source/source_lcao/module_rt/propagator_cn2.cpp

Lines 492 to 540 in ff23741

	// 7. Solve Triangular System (TRSM)
	size_t ws_trsm_dev = 0, ws_trsm_host = 0;
	std::complex<double> alpha_trsm = {1.0, 0.0};
	cublasMpTrsm_bufferSize(cublas_res.cublasmp_handle,
	CUBLAS_SIDE_LEFT,
	CUBLAS_FILL_MODE_UPPER,
	CUBLAS_OP_N,
	CUBLAS_DIAG_NON_UNIT,
	m_global,
	n_global,
	&alpha_trsm,
	d_Den,
	1,
	1,
	desc_blas,
	d_Num,
	1,
	1,
	desc_blas,
	CUBLAS_COMPUTE_64F,
	&ws_trsm_dev,
	&ws_trsm_host);
	void *d_work_trsm = nullptr, *h_work_trsm = nullptr;
	cudaMallocAsync(&d_work_trsm, ws_trsm_dev, cublas_res.stream);
	h_work_trsm = malloc(ws_trsm_host);
	cublasMpTrsm(cublas_res.cublasmp_handle,
	CUBLAS_SIDE_LEFT,
	CUBLAS_FILL_MODE_UPPER,
	CUBLAS_OP_N,
	CUBLAS_DIAG_NON_UNIT,
	m_global,
	n_global,
	&alpha_trsm,
	d_Den,
	1,
	1,
	desc_blas,
	d_Num,
	1,
	1,
	desc_blas,
	CUBLAS_COMPUTE_64F,
	d_work_trsm,
	ws_trsm_dev,
	h_work_trsm,
	ws_trsm_host);

The same fallback also applies when a user explicitly supplies a nonzero nb2d: an incompatible explicit value is silently replaced by the minimum instead of being rejected or reported clearly.

Environment

ABACUS commit: ff237418db665a9fe3befa7029e3364b085b7987
CUDA:         12.9.1
NCCL:         2.29.3
cuSOLVERMp:   0.9.0
cuBLASMp:     0.9.1
GPU:          NVIDIA V100
nodes:        6
MPI ranks:    96 (16 ranks per node, one GPU per rank)
BLACS grid:   8 x 12

The main input parameters are:

calculation    md
esolver_type   tddft
basis_type     lcao
device         gpu
ks_solver      cusolvermp
scf_nmax       1
md_nstep       1
# nb2d omitted, documented default is 0

The Si240 system has nlocal=5280 and nbands=480.

Steps to reproduce

1. Run the attached Si240 input on 96 MPI ranks / 96 GPUs, with 16 ranks per node and one GPU per rank.
2. Leave nb2d unset (default 0).
3. Observe the actual matrix distribution and entry/return markers around cublasMpTrsm().

The measured descriptor on all ranks is consistent with cublasMpNumroc:

global=5280x5280
block=1x1
source=0,0
grid=8x12
local=660x440
nloc=290400
lld=660
descriptor status=0,0
matrix copy status=0,0
den_nonfinite=0
num_nonfinite=0

The buffer-size query succeeds, but the collective call does not return:

Trsm_buffer_size status=0
Trsm_workspace_alloc
Trsm_call_enter

No Trsm_call_return marker is produced by any observed rank.

Minimal library-level reproduction

A standalone program using only MPI, NCCL, CUDA, and cuBLASMp reproduces the same behavior, without ABACUS or cuSOLVERMp:

cuBLASMp 0.9.1
CUDA_C_64F / CUBLAS_COMPUTE_64F
m=n=5280
grid=8x12
mb=nb=1
side=LEFT, uplo=UPPER, trans=N, diag=NON_UNIT

TRSM_ENTER ranks=96 grid=8x12 m=5280 n=5280 block=1
# no TRSM_RETURN

The same block=1 test also failed to return on 64 ranks / an 8x8 grid. In contrast, block=64 standalone tests completed on both 64 ranks (8x8) and 96 ranks (8x12). This indicates that the critical difference is the fallback block size, not the non-square 8x12 process grid by itself.

Verified workaround

Adding the following input parameter avoids the fallback:

nb2d 32

With the same Si240 case, ABACUS build, cuBLASMp 0.9.1, and 96 GPUs, the actual descriptor uses block=32x32 and the calculation completes the propagation:

Trsm_call_return local_s=2.034 s
Trsm_stream_sync=0.063 s
propagator=6.827 s
upsi=0.465 s
norm=0.173 s
ekb=0.158 s

Expected behavior / suggested changes

1. For automatic nb2d=0, select a practical block size intelligently based on the matrix dimensions, wavefunction/band distribution, and process grid, rather than falling back directly to the minimum block size when the initial choice is incompatible.
2. Print the final effective block size and the reason for any fallback in the normal running log.
3. If an explicitly supplied nonzero nb2d is incompatible, emit a clear warning or error rather than silently replacing it with 1.
4. If block size 1 is the only valid choice for a GPU/cuBLASMp calculation, warn that it can be prohibitively slow or unsupported and suggest reducing MPI ranks.
5. Update the nb2d manual entry to document the meaning and policy of the default value 0.

Longer term, it may be useful to decouple the block size used for square H/S matrices from the block size required by the rectangular wavefunction/band distribution, since a single value currently has to satisfy both layouts.