Accurate weather and climate prediction relies on data assimilation (DA), which estimates the Earth system state by integrating observations with models. While exascale computing has significantly advanced earth simulation, scalable and accurate inference of the Earth system state remains a fundamental bottleneck, limiting uncertainty quantification and prediction of extreme events. We introduce a unified one-stage generative DA framework that reformulates assimilation as Bayesian posterior sampling, replacing the conventional forecast–update cycle with compute-dense, GPU-efficient inference. At the core is STORM, a novel spatiotemporal transformer with a global attention linear-complexity scaling algorithm that breaks the quadratic attention barrier. On 32,768 GPUs of the Frontier supercomputer, our method achieves 63% strong scaling efficiency and 1.6 ExaFLOP sustained performance. We further scale to 20 billion spatiotemporal tokens, enabling km-scale global modeling over 177k temporal frames, regimes previously unreachable, establishing a new paradigm for Earth system prediction.

(a) Tiling achieves linear complexity but limits interactions to local regions, while halo overlap improves continuity but remains local. STORM averages denoised outputs (gradients) in overlapping regions and propagates them iteratively across tiles, enabling global context with linear complexity. Hanning weighting stabilizes boundary interactions. (b) Hierarchical mapping of parallelism strategies to supercomputer hardware, with communication frequencies shown on the right.

Our diffusion-based DA framework provides a stable-in-time estimation of hurricane track and intensity, without the need for a computationally expensive physical model. The forecast–analysis cycles are carried out seamlessly within a unified, data-driven framework that avoids expensive I/O operations for updating the model state.

Hurricane track skill: Ensemble trajectories of the four simulated hurricanes, computed from the maximum surface wind speed in each ensemble member. The black tracks correspond to STORM forecast-only , whereas the blue tracks are associated with a data assimilation experiment in which 20% of the ERA5 variables are observed. The reference hurricane tracks (red curves) are obtained from the Best Track dataset provided by National Hurricane Center.

Regional high-resolution temperature simulations of 2-m temperature every 6 hours. Similar to the global case, regional systems such as HRRR are limited by small ensemble sizes due to the high cost of physics-based simulations, restricting their ability to capture uncertainty and nonlinear error growth. STORM relaxes this constraint by enabling large-ensemble, long-context inference within a unified framework. Prediction with DA effectively correct errors by incorporating observational constraints. The assimilated fields closely match the ground truth, and regional biases are largely eliminated.