xarray-sql Derives FlashAttention Kernels from ~25 Lines of SQL

pip install xarray-sql

import xarray as xr
import xarray_sql as xql

# 2013-2014 年经纬度网格上的每日 4 次地表气温数据。
ds = xr.tutorial.open_dataset('air_temperature')

ctx = xql.Context()
ctx.from_dataset('air', ds, chunks=dict(time=100))

# 在 SQL 中计算气候态——即年均周期——：对全年每个月的所有网格单元和年份的平均气温进行计算。
clim = ctx.sql('''
  SELECT
    CAST(date_part('month', "time") AS INTEGER) AS month,
    AVG("air") AS air
  FROM "air"
  GROUP BY CAST(date_part('month', "time") AS INTEGER)
  ORDER BY month
''')

# 将结果往返转换回 Xarray。`month` 是一个派生列，因此将其命名为维度。
clim_ds = clim.to_dataset(dims=["month"])

# 将年均周期绘制为时间序列。
clim_ds["air"].plot()  # 在脚本中，调用 matplotlib.pyplot.show() 来显示图像

import duckdb

con = duckdb.connect()
xql.register(con, 'air', ds, chunks=dict(time=100))
rel = con.sql('SELECT time, AVG("air") AS air FROM air GROUP BY time ORDER BY time')
xql.to_dataset(rel, template=ds)   # 任何引擎的 Arrow 结果都可以往返转换

import xarray as xr
import xarray_sql as xql

# 打开 ARCO-ERA5 —— 一个自 1940 年以来拥有 273 个变量的气象数据集。
# 关闭 dask 意味着我们无需等待构建任务图。
ds = xr.open_zarr(
  'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3',
  chunks=None,  # 关闭 dask
  storage_options={'token': 'anon'}  # 匿名读取公共 GCS 存储桶 — 无需身份验证。
)

ctx = xql.Context()
# `chunks` 设置分区大小；如果不设置，分区将遵循存储自身的 chunk 结构（此处为每小时一个）。
ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={
    ('time', 'latitude', 'longitude'): 'surface',
    ('time', 'level', 'latitude', 'longitude'): 'atmosphere',
})
# 在我的机器上，注册过程大约耗时 10 秒。

# 注意：ARCO-ERA5 包含 262 个地表变量 + 11 个大气变量。该库会将列投影下推至 Zarr，
# 因此 SELECT 仅获取你请求的数据 — 但 `SELECT * FROM era5.surface` 会尝试拉取全年所有变量
# （从 GCS 下载数 TB 数据）。
#  ---> 始终 SELECT 特定列。 <---

# 计算 2020-01-01 上午纽约市的平均 2 米温度。该库会将针对维度列的 WHERE 子句下推以进行分区剪枝。
ctx.sql('''
  SELECT AVG("2m_temperature") - 273.15 AS avg_c
  FROM era5.surface
  WHERE time BETWEEN TIMESTAMP '2020-01-01'
                 AND TIMESTAMP '2020-01-01 05:00:00'
    AND latitude  BETWEEN 39 AND 40
    AND longitude BETWEEN 286 AND 287  -- ERA5 使用 0-360 经度
''').to_pandas()
#       avg_c
# 0  8.640069

> （注：更多技术实现细节与完整 API 文档请查阅原项目 README）

---
**作者发帖说明（HN）**：

在这个演示中，我们仅用约 25 行 SQL 自动推导出了一个 FlashAttention kernel，而这些 SQL 仅仅表达了标准的注意力机制（vanilla attention）。之所以能够实现这一点，得益于我们正在构建的三个可移植的 SQL 扩展：
- **xarray-sql**：一个库，能够高效地将云端的 PB 级数组数据集表示为表（关系）。
- **ddx**：一个数据库扩展，在上述数组数据模型的基础上，为标准 SQL 添加了类似 JAX 风格的自动微分功能。
- **einfold**：一个模块化查询规划器，同样基于上述数组数据模型，用于在标准 SQL 中执行快速的张量收缩运算。

我希望关系模型和代数能为我们如何开展加速器编程提供全新的视角。在此，我们的核心论点是：将张量程序中的逻辑层与物理层分离，既更正确，也更具可扩展性。

如果“苦涩的教训”是机器学习领域的中心教条，那么为何不用关系代数将其形式化呢？