Metadata-Version: 2.5
Name: flash-aurora
Version: 0.1.5
Summary: Efficient inference and serving for Microsoft Aurora, an Earth system AI model.
License-Expression: MIT
License-File: LICENSE
License-File: flash_aurora/LICENSE
License-File: flash_aurora/models/aurora/LICENSE.txt
License-File: flash_aurora/models/aurora_v1p5/LICENSE.txt
Requires-Python: <3.14,>=3.12
Requires-Dist: cdsapi
Requires-Dist: cfgrib
Requires-Dist: eccodes
Requires-Dist: ecmwf-api-client>=1.7.0
Requires-Dist: ecmwflibs
Requires-Dist: einops>=0.8.2
Requires-Dist: fiona>=1.10.0
Requires-Dist: gcsfs>=2026.6.0
Requires-Dist: huggingface-hub>=1.8.0
Requires-Dist: ipykernel<7,>=6.30.1
Requires-Dist: matplotlib>=3.10.8
Requires-Dist: netcdf4
Requires-Dist: numpy
Requires-Dist: nvidia-cutlass-dsl[cu13]==4.7.1
Requires-Dist: pandas>=3.0.1
Requires-Dist: pydantic>=2.12.5
Requires-Dist: pyzmq>=26.0.0
Requires-Dist: quack-kernels==0.6.5
Requires-Dist: rasterio>=1.4.0
Requires-Dist: requests>=2.33.0
Requires-Dist: scipy>=1.17.1
Requires-Dist: timm>=1.0.26
Requires-Dist: torch==2.14.0
Requires-Dist: tqdm>=4.67.0
Requires-Dist: triton==3.8.0
Requires-Dist: xarray
Requires-Dist: zarr>=3.2.1
Provides-Extra: dev
Requires-Dist: flash-attn-4[cu13]>=4.0.0b31; extra == 'dev'
Requires-Dist: ipywidgets>=8.1.0; extra == 'dev'
Requires-Dist: pytest>=9.0.2; extra == 'dev'
Description-Content-Type: text/markdown

# Flash-Aurora

[![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.21860697.svg)](https://doi.org/10.5281/zenodo.21860697)

Inference and serving for [Microsoft Aurora](https://github.com/microsoft/aurora) and other geospatial foundation models. A request takes an analysis cube in and writes NetCDF or GeoTIFF out. The tensors keep a lat/lon layout, so the engine uses shape-specialized kernels, named mixed precision (`bf16_mixed@fp32` by default), and a job-level GPU scheduler rather than an LLM serving loop.

[Walkthrough slides](https://catmanjr.github.io/Flash-Aurora-walkthrough/) · [Tutorial](docs/tutorial.md) · [Benchmark tables](docs/benchmarks.md)

## Install

```bash
pip install flash-aurora
```

From source:

```bash
git clone https://github.com/CatManJr/Flash-Aurora.git
cd Flash-Aurora
uv sync
```

Set `CUTE_DSL_ARCH` when the CuTe kernels need an explicit GPU architecture (`sm_89` on RTX 4090, `sm_120a` on Blackwell). The published measurements use PyTorch 2.14.0, Triton 3.8.0, and CuTe DSL 4.7.1, which `uv.lock` pins. If PyPI is slow, point `uv` at a mirror with `UV_DEFAULT_INDEX` or a local `uv.toml`; keep the committed `uv.lock` on the official index.

## Quick start

```python
from datetime import datetime
from pathlib import Path

from flash_aurora import AuroraEngine, DataDownloader

engine = AuroraEngine.from_preset(
    "era5_pretrained",
    asset_root=Path("/path/to/assets"),
    inference_precision="bf16_mixed@fp32",
)
downloader = DataDownloader.from_preset(
    "era5_pretrained",
    asset_root=engine.config.asset_root,
)
request = downloader.ingest_request(
    datetime(2023, 1, 1, 6),
    time_index=1,
    download=True,
)
batch = engine.prepare(request, rollout_steps=4)
forecasts = list(engine.rollout_stream(batch, steps=4))
engine.release_gpu(move_model_to_cpu=True)
engine.close()
```

Scheduler loopback, two-GPU placement, and the notebook index: [docs/tutorial.md](docs/tutorial.md).

## What it does

Production mixed precision is faster than unfused FP32 and stays closer to that twin than framework autocast. One forward step on RTX PRO 6000 Blackwell, measured with PyTorch 2.14.0 and CuTe DSL 4.7.1, is about $680$--$1000\,\mathrm{ms}$ with a TF32 Perceiver (`bf16_mixed@tf32`) and about $830$--$1150\,\mathrm{ms}$ with an FP32 Perceiver (`bf16_mixed@fp32`), versus about $1.9$--$2.5\,\mathrm{s}$ for unfused FP32. Aurora 1.5 ensemble is the slow end of both ranges, because its stochastic MLP stays on the tier's TF32 matmul. Each bar is a separate process (`--isolate-tiers`).

<p align="center">
  <img src="https://raw.githubusercontent.com/CatManJr/Flash-Aurora/master/docs/image/e2e_latency_by_tier_all_presets.svg" alt="One-step end-to-end forward latency by precision tier" width="95%"/>
</p>

Recommended tiers stay within per-variable tolerances versus the unfused FP32 reference (seed 42). The one exception is `bf16_mixed@tf32` on `aurora_v1p5_ensemble`, where `scaled_tp_1h` is about $1.1\times$ the tolerance. `bf16@*` is not a production path because the full BF16 GEMM path corruptted 
the prediction in experiments.
<p align="center">
  <img src="https://raw.githubusercontent.com/CatManJr/Flash-Aurora/master/docs/image/precision_mean_rel_stacked_by_model.svg" alt="Stacked mean relative error by precision tier and preset" width="95%"/>
</p>

Window-attention kernels are short-window CuTe DSL (`N = 144` on the default $0.25^{\circ}$ encoder), not a generic LLM attention stack.

<p align="center">
  <img src="https://raw.githubusercontent.com/CatManJr/Flash-Aurora/master/docs/image/window_attn_cute_vs_sdpa_blackwell.svg" alt="CuTe DSL window attention versus PyTorch SDPA on Blackwell" width="95%"/>
</p>

Serving is one GPU per job. A ZeroMQ coordinator fills idle workers; it does not batch tokens inside one forward. [Scheduler notebooks](docs/example_scheduler_distributed_workers.ipynb).

<table width="100%">
  <tr>
    <th width="50%" align="center">One job per worker</th>
    <th width="50%" align="center">Refill while <code>hres_0.1</code> is pending</th>
  </tr>
  <tr>
    <td width="50%" valign="top"><img src="https://raw.githubusercontent.com/CatManJr/Flash-Aurora/master/docs/image/4_workers.png" width="100%"/></td>
    <td width="50%" valign="top"><img src="https://raw.githubusercontent.com/CatManJr/Flash-Aurora/master/docs/image/4_workers_refill.png" width="100%"/></td>
  </tr>
</table>

A preset that does not fit one GPU can run encoder / backbone / decoder on two devices in the same process (`DistributedConfig`). [ROI export](docs/example_roi_export.ipynb) clips on the egress path so a region of interest does not require a global dump.

## Presets

| Preset | Grid | Source |
| ------ | ---- | ------ |
| `era5_pretrained` | $721 \times 1440$ | CDS ERA5 |
| `aurora_v1p5` / `aurora_v1p5_ensemble` | $721 \times 1440$ | CDS ERA5 (extended) |
| `small_pretrained` | $400 \times 800$ | CDS ERA5 |
| `hres_t0_finetuned` / `tc_tracking` | $721 \times 1440$ | WeatherBench2 HRES |
| `hres_0.1` | $1801 \times 3600$ | IFS analysis |
| `cams` | $451 \times 900$ | CAMS |
| `wave` | $721 \times 1440$ | WB2 + MARS |

`wave` usually needs a hand-placed MARS cache; see [example_wave.ipynb](docs/example_wave.ipynb). Notebooks for every preset are listed in the [tutorial](docs/tutorial.md#tutorial-notebooks).

## License

[MIT](LICENSE). Aurora code is derived from [Microsoft Aurora](https://github.com/microsoft/aurora) (MIT), frozen at **v1.8.0**, with Aurora 1.5 from tag `v2.0.1`. Kernel files under `flash_aurora.models.ops` follow their per-file headers (including NVIDIA BSD-3-Clause where noted).

Bodnar et al., *A Foundation Model for the Earth System*, Nature (2025). [doi:10.1038/s41586-025-09005-y](https://doi.org/10.1038/s41586-025-09005-y). Upstream docs: [microsoft.github.io/aurora](https://microsoft.github.io/aurora).
