Repository navigation
AutoencoderKlMaisi: tensors with any dimension ≥ 500 are concatenated on the CPU, 40–165× slower even when they fit on the GPU #9250
Description
Describe the bug
In monai/apps/generation/maisi/networks/autoencoderkl_maisi.py, MaisiGroupNorm3D.forward (line 89) and MaisiConvolution._concatenate_tensors (line 214) switch to a CPU path when max(tensor.size()) >= 500. The tensor-parallel splits are moved to the CPU one by one and concatenated there, with empty_cache() and gc.collect() after each split, and the result is moved back to the GPU.
This is reached by any input with a spatial dimension of 500 or more. That includes 512×512 in-plane, the usual size of clinical CT. With such an input, encode + decode is 40–165× slower than at 496×496, even on an 80 GB H100 where the tensors fit in GPU memory.
The same code is on dev (lines 89 and 214 at 7f6bf2db). #8735 touched this path, but only for CPU inputs.
To Reproduce
Random weights are enough, so no checkpoint or data are needed. The network definition is the one in NV-Generate's configs/config_network_rflow.json.
import time, torch
from monai.apps.generation.maisi.networks.autoencoderkl_maisi import AutoencoderKlMaisi
ae = AutoencoderKlMaisi(
spatial_dims=3, in_channels=1, out_channels=1, latent_channels=4,
num_channels=[64, 128, 256], num_res_blocks=[2, 2, 2], norm_num_groups=32, norm_eps=1e-6,
attention_levels=[False, False, False], with_encoder_nonlocal_attn=False,
with_decoder_nonlocal_attn=False, use_checkpointing=False, use_convtranspose=False,
norm_float16=True, num_splits=4, dim_split=1,
).cuda().eval()
for n in (496, 512): # 496 < 500 <= 512
x = torch.rand(1, 1, n, n, 16, device="cuda")
with torch.inference_mode(), torch.autocast("cuda", dtype=torch.float16):
ae.decode(ae.encode(x)[0]) # warm-up
torch.cuda.synchronize(); t = time.perf_counter()
ae.decode(ae.encode(x)[0])
torch.cuda.synchronize()
print(f"{n}x{n}x16: {time.perf_counter() - t:.1f} s")H100 80 GB SXM (driver 580.173.02). The z = 64 lines come from the same code with 64 slices:
496x496x16: 0.29 s
512x512x16: 47.8 s
496x496x64: 0.96 s
512x512x64: 104.1 s
RTX 4070 Ti SUPER 16 GB:
496x496x16: 1.3 s
512x512x16: 53.4 s
On real data (whole-volume encode + decode on the H100):
- ten 512×512 CTs from the Medical Segmentation Decathlon (36–171 slices) took 86–227 s each, with a peak of 44 GB;
- volumes up to 462 voxels wide took 0.4–2.7 s;
- a 512×512×266 cone-beam CT took 326 s, with a peak of 67 GB.
Expected behavior
When the concatenated tensor fits in free GPU memory, the concatenation stays on the GPU. Alternatively, the threshold could be configurable, for example as a constructor argument, so that 512-wide volumes don't take the slow path by default.
If the 500 threshold is a deliberate memory safeguard, a note in the docstring would help users understand the slowdown.
Environment
Output of python -c "import monai; monai.config.print_debug_info()" on the RTX machine. The MONAI __file__ line is removed; the H100 run used the same MONAI, PyTorch and CUDA versions.
================================
Printing MONAI config...
================================
MONAI version: 1.6.1
Numpy version: 2.2.6
Pytorch version: 2.11.0+cu130
MONAI flags: HAS_EXT=False, USE_COMPILED=False, USE_META_DICT=False
MONAI rev id: 7af492176792f9069b8f1253d099abd05c5e769d
Optional dependencies:
Pytorch Ignite version: NOT INSTALLED or UNKNOWN VERSION.
ITK version: NOT INSTALLED or UNKNOWN VERSION.
Nibabel version: 5.4.2
scikit-image version: 0.26.0
scipy version: 1.17.1
Pillow version: 13.0.0
Tensorboard version: NOT INSTALLED or UNKNOWN VERSION.
gdown version: NOT INSTALLED or UNKNOWN VERSION.
TorchVision version: 0.26.0+cu130
tqdm version: 4.70.1
lmdb version: NOT INSTALLED or UNKNOWN VERSION.
psutil version: NOT INSTALLED or UNKNOWN VERSION.
pandas version: NOT INSTALLED or UNKNOWN VERSION.
einops version: 0.8.2
transformers version: NOT INSTALLED or UNKNOWN VERSION.
mlflow version: NOT INSTALLED or UNKNOWN VERSION.
pynrrd version: NOT INSTALLED or UNKNOWN VERSION.
clearml version: NOT INSTALLED or UNKNOWN VERSION.
================================
Printing GPU config...
================================
Num GPUs: 1
Has CUDA: True
CUDA version: 13.0
cuDNN enabled: True
NVIDIA_TF32_OVERRIDE: None
TORCH_ALLOW_TF32_CUBLAS_OVERRIDE: None
cuDNN version: 91900
Current device: 0
Library compiled for CUDA architectures: ['sm_75', 'sm_80', 'sm_86', 'sm_90', 'sm_100', 'sm_120']
GPU 0 Name: NVIDIA GeForce RTX 4070 Ti SUPER
GPU 0 Is integrated: False
GPU 0 Is multi GPU board: False
GPU 0 Multi processor count: 66
GPU 0 Total memory (GB): 15.6
GPU 0 CUDA capability (maj.min): 8.9
Additional context
We hit this while evaluating the MAISI / NV-Generate CT VAE on dental cone-beam CT. Happy to test a fix.
Activity
slavaheroes commentedon Nov 4, 2026
I'd like to work on this issue, so if nobody is already working on this, could it be assigned to me? I have an access to RTX 4090 and RTX 5090 GPUs.
I'll start by reproducing the slowdown, then work on a fix. Looks like the problem is located in:
My approach would be to try concatenation on the GPU first, then fall back to CPU only if it causes CUDA OOM. Also, I'll include correctness tests and report latency/memory benchmarks.
If you prefer another approach, your feedback would be welcome.
ikchain commentedon Nov 4, 2026
Thanks @slavaheroes, great to see this picked up so quickly! Assignment is up to the maintainers, but the approach sounds good to us.
Two small thoughts, since you asked:
Once you have a branch, we’re happy to run it on our 512-wide CBCT and MSD CT volumes and report timings and peak memory here.
slavaheroes commentedon Nov 5, 2026
Hi @ikchain , Thanks for your offer to test. I've opened a draft PR: #9255
The implementation is on https://github.com/slavaheroes/MONAI/tree/9250-maisi-concatenation
ikchain commentedon Nov 6, 2026
Thanks @slavaheroes! We tested the PR (
333f968) against MONAI 1.6.1 on an RTX 4070 Ti SUPER 16 GB that also runs the desktop. We used the network from the issue, seeded random weights shared by both versions, and timed encode + decode. Memory figures are peak allocated memory.Without memory pressure (median of 3 runs after a warm-up):
Outputs are bitwise identical.
Near the memory limit (about 12 GiB free, no allocator limit), real memory pressure triggered the CPU fallback and the runs completed. Each run used a fresh process. MONAI 1.6.1 was run once, since the slow path takes over a minute; the PR was timed after a warm-up:
At 512×512×48 both versions go OOM (MONAI 1.6.1 with about 0.25 GiB less free memory).
One thing we can't explain yet: in the two fallback runs, the outputs are not bitwise identical to MONAI 1.6.1 (max abs diff 0.0037 and 0.0039, under fp16 autocast), while at 16 slices, without fallback, they were. The concatenation itself looks exact, so this may come from elsewhere; we haven't isolated it. Have you seen anything similar in your fallback tests?
These are random inputs on a 16 GB card. The real 512-wide volumes from the issue need an 80 GB GPU and aren't covered here.