Technical Summary: Performance Analysis and Optimization of 3D Generative Diffusion Models across GPU Architectures

Problem Statement

High-fidelity 3D medical image synthesis, particularly for MRI, relies heavily on Denoising Diffusion Probabilistic Models (DDPMs) like Med-DDPM due to their superior structural fidelity and training stability compared to GANs. However, the deployment of these models is severely constrained by immense computational costs. A single 3D diffusion sample requires hundreds of sequential U-Net evaluations, each involving complex 3D convolutions, normalization layers, and elementwise operations. This results in a highly heterogeneous GPU execution pattern characterized by frequent kernel launches, alternating compute-bound and memory-bound phases, and significant resource demands.

Existing medical imaging research typically evaluates diffusion models solely on image quality metrics, lacking a systems-level perspective. Consequently, the specific architectural bottlenecks—ranging from memory access patterns and tensor-layout conversions to underutilized Tensor Core pipelines—remain poorly understood. This paper addresses the gap by investigating why training 3D diffusion models demands hundreds of GPU-hours even for modest datasets and how current hardware architectures (Volta, Ampere, Hopper) handle these specific workloads.

Methodology

The authors conducted a comprehensive, full-stack performance analysis of the state-of-the-art Med-DDPM framework across three generations of NVIDIA GPUs:

• Hardware: NVIDIA V100 (Volta), A100 (Ampere), and H100 (Hopper).
• Software Stack: PyTorch 2.0.1, CUDA 12.9, cuDNN 8.9.2, and Nsight Compute (ncu) profiler.
• Experimental Setup: Training was executed with identical hyperparameters (25 epochs, batch size 1, 128³ input size) across all platforms. To ensure bit-identical workloads, the authors utilized containerized environments (Docker on Indiana Jetstream2/TACC Stampede3; Singularity/Apptainer on SDSC Expanse) to eliminate variability from system libraries.
• Profiling Strategy: The study utilized Nsight Compute to extract microarchitectural metrics, including:
  • Kernel-level runtime breakdowns and instruction mixes.
  • Dynamic instruction profiles (FP, INT, load/store, control).
  • Priority-score estimates (combining kernel duration and estimated speedup potential).
  • SM-level utilization, warp occupancy, and stall-type breakdowns (long-latency dependencies, math-pipeline throttling, divergent branches).
  • Memory hierarchy behavior (L1/L2 hit rates, DRAM bandwidth utilization).

Key Contributions

1. System-Level Characterization: The paper presents the first detailed microarchitectural characterization of Med-DDPM across Volta, Ampere, and Hopper architectures, revealing that training is overwhelmingly dominated by cuDNN 3D convolution and implicit-GEMM kernels.
2. Bottleneck Identification: Through kernel-mix breakdowns and priority-score analysis, the authors identify that inefficiencies arise not just from memory access irregularities but also from fragmented warp scheduling, cache residency effects, and limited Tensor Core utilization on Ampere/Hopper architectures.
3. Optimization Evaluation: The study evaluates two architecture-aware optimizations:
   • TF32 Tensor Core Activation: Enabling TF32 precision for matrix multiplication and 3D convolutions on Ampere and Hopper.
   • 3D Channels-Last Layout: Restructuring the memory layout from NCHW to NHWC to improve data locality for tensor cores.
4. Performance Insights: The paper demonstrates that while TF32 activation significantly reduces SM cycles and dynamic instructions, the channels-last layout fundamentally shifts the workload to a memory-bound regime dominated by small elementwise kernels, exposing new bottlenecks related to kernel fusion and launch overhead.

Results

The evaluation compared four configurations: Baseline, OPT1 (TF32 only), OPT2 (Channels-Last only), and OPT12 (Combined).

• TF32 Activation (OPT1):

  • Efficiency: On A100 and H100, enabling TF32 reduced SM cycles by up to 5× and dynamic instructions by up to 100× compared to the baseline.
  • Utilization: Tensor Core utilization increased significantly from ~1.45× to 9.98× on A100. Conversely, CUDA core utilization dropped as the workload shifted to dense MMA (Matrix Multiply-Accumulate) instructions.
  • Memory: While L1 cache hit rates decreased (due to larger, contiguous global memory tiles streaming through SMs), DRAM bandwidth utilization increased, indicating that the Tensor Core path aggressively pulls data to feed the pipelines.
  • Stalls: On A100, all stall categories decreased significantly. On H100, math-pipeline throttling dropped, but long-latency dependency stalls increased, reflecting the deeper dependency chains of Tensor Core operand availability.

• Channels-Last Layout (OPT2):

  • Kernel Shift: This optimization collapsed the dominance of large conv/GEMM kernels, shifting execution almost entirely to small, memory-bound elementwise and layout-transform kernels (e.g., padding, NHWC↔NCHW conversions).
  • Performance: While this reduced total cycles and instructions drastically (down to 0.01×–0.02× of baseline), it resulted in underutilized compute engines. Tensor Core utilization remained near baseline levels (~1.5%), and SM throughput dropped by roughly 50%.
  • Memory: L2 cache hit rates increased dramatically (up to 85%), and DRAM bandwidth utilization collapsed, confirming the workload became cache-resident but lacked arithmetic intensity to saturate the hardware.

• Combined (OPT12):

  • The combination partially restored Tensor Core usage on A100 but failed to achieve peak throughput on H100 or V100, where the workload remained dominated by elementwise and layout transforms. The authors note that further kernel fusion is required to realize the full potential of this layout.

Significance and Claims

The paper claims that the primary barrier to scaling 3D medical diffusion models is not merely algorithmic but systemic, rooted in the mismatch between diffusion workloads and GPU microarchitectures.

• Architectural Reality: The study establishes that simply moving to newer GPU generations (e.g., from V100 to H100) does not automatically resolve bottlenecks if the workload remains dominated by non-compute micro-operations, layout conversions, or control flow.
• Optimization Hierarchy: The authors conclude that TF32 Tensor Core activation is a robust, immediate optimization baseline that compresses arithmetic work without degrading synthesis quality. In contrast, while the channels-last layout improves cache locality, it currently exposes a new regime of memory-bound micro-kernels that require further fusion and architectural tuning to be effective.
• Practical Guidance: The findings provide a roadmap for optimizing general 3D diffusion models: align precision modes and memory formats with the underlying GPU's Tensor Core capabilities, and recognize that reducing cycle counts does not always equate to proportional speedups if the resulting workload fails to saturate compute or memory engines.

The paper modestly asserts that these insights are essential for making the practical deployment of high-fidelity 3D medical generators feasible, moving beyond image-quality metrics to address the computational realities of training.