bf16_mixed@fp32 / bf16_mixed@tf32
Hybrid backbone: BF16 attention and MLP paths, TF32 elsewhere; FP32 or TF32 encoder/decoder GEMMs.
1 · Background and motivation
What Aurora tries to solve
Official Microsoft Aurora demos across weather, air quality, waves, and tropical-cyclone tracking - examples of outputs that Flash-Aurora can export as GIS-ready layers.
Source: Microsoft Aurora documentation GIFs; Bodnar et al., Nature 2025 [5]; github.com/microsoft/aurora [6].
1 · Background and motivation
Research questions
LLM stacks already ship FlashAttention [7] and vLLM [8]; weather AI has Earth2Studio [9]. GFMs like Aurora still lack a dedicated inference and serving path for GIS practice. This project asks:
- RQ1 · Easy to use An end-to-end Engine from high-resolution initial conditions to GIS-ready ROI export
- RQ2 · Efficiency Custom GPU kernels and tiered mixed precision for Aurora throughput, memory, and numerical trade-offs
- RQ3 · Future trend Multi-GPU serving, desktop placement, and a reusable model factory toward unified GFM infrastructure
2 · Architecture
The Architecture
Three nested packages separate deployment, forecast execution, and model computation. The same Engine can run directly in a process or inside a Scheduler worker.
- Models: use published Aurora variants with custom GPU kernels and precision routing.
- Engine: prepares initial conditions, rolls out forecasts, and exports products.
- Scheduler: queues complete forecast jobs across GPU workers.
Flash-Aurora is not a new GFM. It extends the published Aurora family with inference and serving infrastructure. The packages nest like an egg. Read from inside out: Models are the yolk that run the forward path; Engine is the white that owns the forecast workflow; Scheduler is the shell that manages deployment.
3 · Efficient inference
Custom GPU kernels improve speed and enable finer-grained data-type control
The Swin3D backbone is a major forward-pass cost. Custom GPU kernels built with Triton [10] fuse window layout and AdaLN, while CuTe DSL [11] replaces short-window attention on production tiers. They avoid repeatedly reading and writing temporary data, while letting us choose data types at a finer grain for each compute step.
- Triton [10] fuses window layout and AdaLN, so intermediate tensors need not be materialized
- CuTe DSL [11] runs short-window attention with 144 tokens per window and tiled Q, K, and V data in shared memory
- Online softmax accumulates in FP32 without writing the full pairwise attention matrix to global memory
3 · Efficient inference
End-to-end mixed-precision: faster matrix math where safe
Tiers use backbone@encoder_decoder. The left token
selects backbone matrix-multiply precision; the right selects
encoder/decoder matrix-multiply precision. These are compute
policies, not a conversion of stored model weights. Lower-precision
Tensor Core matrix operations provide higher throughput than strict
FP32 in the heavy backbone, while the encoder/decoder can retain
FP32 where a stricter numerical path is needed.
tf32@fp32 / tf32@tf32
TF32 backbone GEMMs; FP32 or TF32 Tensor Core GEMMs in the encoder and decoder.
fp32@fp32 / fp32@tf32
Strict FP32 backbone GEMMs; FP32 or TF32 encoder/decoder GEMMs.
3 · Efficient inference
How much faster is one forecast step now?
Measured latency for one model.forward call, or one
rollout step, across presets and precision tiers. The tier-isolated
benchmark warms the model first, then compares compute only: it
excludes data download and GIS export.
model.forward step after warmup (NVIDIA RTX PRO 6000 Blackwell)
Preset tokens (era5_pretrained, hres_0.1, ...)
are the same ones mapped to published models and tasks on page 2.
3 · Efficient inference
How much numerical precision do we lose
Stacked per-variable mean relative error versus the PyTorch FP32 baseline (seed 42). Each segment is one output field; shorter stacks mean less aggregate drift. PyTorch autocast is a comparison baseline, not a quality threshold; it exceeds tolerance for some variables.
Preset tokens (era5_pretrained, hres_0.1, ...)
are the same ones mapped to published models and tasks on page 2.
3 · Efficient inference
From analysis data to GIS layers
Three linked stages: solve the inputs, load and run a model, then export predictions. Dashed parts are optional; for simple one-model tasks you do not need to launch the Scheduler.
4 · Efficient serving
Scheduler architecture
Clients submit whole forecast jobs to a coordinator. Long-lived workers keep models resident on a primary GPU and execute jobs sequentially. This is job-level scheduling, not tensor parallelism; optional pipeline parallelism is a separate in-process Engine mode.
A worker is bound to one preset and normally one GPU. A pipeline worker is the exception: one process can reserve multiple GPUs for one model.
4 · Efficient serving
Heterogeneous jobs and asynchronous refill
Worker-utilization traces: first assignment, then refill while a
slower hres_0.1 job is still pending. A completed worker
can accept its next queued job while another GPU remains busy.
Traces: docs/example_scheduler_distributed_workers.ipynb
on 4x NVIDIA RTX PRO 6000 Blackwell.
5 · Affordable AI
Fit large grids on smaller GPUs
For supported Aurora models, pipeline parallelism keeps one model in one process but places the encoder, backbone, and decoder on separate GPUs. It can avoid requiring one high-VRAM GPU for the entire model.
The trade-off is explicit: more GPU-to-GPU transfers and coordination, but a lower per-GPU VRAM requirement for the same forecast.
6 · Towards a unified inference infrastructure
Aurora 1.5 as an extensibility test
Day-one support shows that a closely related Aurora variant can
reuse the Engine, source profiles, egress, precision routing, and
much of the acceleration path. It does not demonstrate a general
runtime for unrelated GFMs. Broader unification remains a design
direction that requires standardized checkpoint, source, and
output contracts. Run output from
docs/example_aurora_v1p5.ipynb
(ERA5 IC, 2023-01-01).
