توسعه هسته FlashAttention با کدهای SQL؛ بررسی پروژه xarray-sql

معرفی پروژه xarray-sql و مسئله آن

زبان SQL زبان مشترک دنیای داده است، اما دیتاست‌های علمی و مقیاس‌پذیر در حوزه هوش مصنوعی و هواشناسی (مانند دیتاست‌های مبتنی بر Xarray) معمولاً در قالب‌های خاصی ذخیره می‌شوند و برای کاربران SQL به‌سادگی قابل دسترس نیستند. پروژه xarray-sql تلاشی است برای حل این مسئله؛ این کتابخانه یک رابط SQL برای دیتاست‌های آرایه‌ای فراهم می‌کند تا توسعه‌دهندگان بتوانند با تبدیل ساختارهای Xarray به جداول (Relations)، پرس‌وجوهای خود را روی موتورهای پایگاه داده دلخواه اجرا کنند.

نگاهی به پشت صحنه و ارکان پروژه

این ابزار بخشی از یک ایده بزرگ‌تر در سازمان xqlsystems است که تلاش می‌کند مدل رابطه‌ای (Relational Model) را به دنیای یادگیری ماشین و پردازش تنسورها بیاورد. در دموهای اخیر، این تیم توانسته است تنها با حدود ۲۵ خط SQL (که توجه تنسور متداول را بیان می‌کند)، یک هسته FlashAttention را به صورت خودکار استخراج کند. این معماری روی سه پایه اصلی بنا شده است:

• xarray-sql: کتابخانه‌ای برای نمایش کارآمد دیتاست‌های آرایه‌ای مقیاس‌پترابایت در فضای ابری به شکل جدول.
• ddx: یک افزونه پایگاه داده برای اضافه کردن قابلیت مشتق‌گیری خودکار (Automatic Differentiation) در سبک JAX به SQL استاندارد.
• einfold: یک برنامه‌ریز ماژولار پرس‌وجو برای انقباضات سریع تنسورها در SQL خالص.

نحوه کارکرد فنی و معماری

در نسخه‌های اولیه، این ابزار از Dask برای تبدیل تکه‌های (chunks) دیتاست Xarray به Dask DataFrame استفاده می‌کرد. اما در تکامل‌های بعدی، معماری به سمت استفاده از DataFusion و PyArrow حرکت کرد.

مکانیزم اصلی این کتابخانه مبتنی بر ترجمه تکه‌های دیتاست به Arrow RecordBatches است. با استفاده از یک TableProvider در DataFusion، موتور پایگاه داده می‌تواند آرایه‌های زیربنایی دیتاست را به پارتیشن‌های DataFusion تبدیل کند. این ترجمه همچنین به عنوان یک pyarrow.dataset.Dataset با قابلیت اعمال پیش‌فرض شرط‌ها (Predicate pushdown) عرضه می‌شود که به موتورهایی مانند DuckDB و Polars اجازه می‌دهد بدون نیاز به کدهای اختصاصی، از این دیتاست‌ها استفاده کنند.

نحوه استفاده و شروع سریع

برای نصب این کتابخانه کافی است دستور زیر را اجرا کنید:

pip install xarray-sql

در ادامه یک نمونه ساده از ثبت دیتاست و اجرای پرس‌وجو آورده شده است:

import xarray as xr
import xarray_sql as xql

# بارگذاری یک دیتاست نمونه
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
''')

clim_ds = clim.to_dataset(dims=["month"])

سیگنال‌های بلوغ و وضعیت مخزن

این پروژه تحت مجوز متن‌باز Apache-2.0 منتشر شده و در گیت‌هاب بیش از ۱۴۰ ستاره (Stars) دریافت کرده است. فعالیت توسعه روی آن جریان دارد و جامعه کاربری آن در حال آزمایش ایده‌های نوآورانه در تقاطع تئوری پایگاه داده و سیستم‌های یادگیری ماشین است.

چرا این پروژه در هکرنیوز مطرح شد؟

جامعه برنامه‌نویسان در هکرنیوز همواره به ایده‌های ساختارشکنانه که ابزارهای سنتی را در حوزه‌های جدید به کار می‌گیرند علاقه نشان می‌دهند. استفاده از SQL و جبر رابطه‌ای برای برنامه‌نویسی شتاب‌دهنده‌ها و استخراج هسته‌هایی مانند FlashAttention، دیدگاهی تازه و چالش‌برانگیز در برابر رویکردهای سنتی یادگیری ماشین ایجاد کرده است. در بحث‌های این پست نیز کاربران به قدرت و زیبایی زبان SQL در مدل‌سازی داده‌ها اشاره کرده‌اند.

در یک نگاه

• نام مخزن: xqlsystems/xarray-sql
• زبان اصلی: پایتون (Python)
• مجوز: Apache-2.0
• هدف اصلی: اتصال دیتاست‌های علمی Xarray به دنیای پایگاه‌های داده و SQL
• تکنولوژی‌های کلیدی: Apache Arrow ،DataFusion ،DuckDB ،Dask