AI-SQLエンジン「Quail」公開 —⁠—LLMによるデータの絞り込み⁠⁠・結合を効率化

Full Stack Data LabのShreya Shankar氏らは9月24日、Modalとの共同開発によるオープンソースのAI-SQLエンジン「Quail」を公式ブログで発表した。SQLに自然言語で条件を書き、大規模言語モデル(LLM)の判定でデータを絞り込んだり、結び付けたりできる。データを処理する順序とLLMの動かし方を一体で最適化し、計算結果の再利用などによって大量の判定を効率よく実行する。ソースコードはMITライセンスでGitHubに公開されている。

自然言語の条件で行を絞り込み⁠、データを結合

AI-SQLでは、自然言語で書いた条件を使ってデータを絞り込める。Quailは、各行が条件を満たすかをLLMで判定するフィルター処理をサポートする。結合処理にも対応し、2つのデータセットの行同士が指定した関係を満たすかをLLMで判定して、該当する組み合わせを取り出す。

たとえば、データベースに保存された映画レビューを内容で絞り込む場合、接続ライブラリーなどでレビュー本文とIDを取得し、Apache Arrow形式のテーブルに整える。そのデータをPythonからQuailのセッションに名前付きで登録し、自然言語の条件を含むSQLを渡す。Quailが判定を行う順序やデータの処理方法を決め、GPU上でモデルを動かして、条件に合うレビューを絞り込む。

以下では、公式ブログの例をもとに、データと設定を簡略化したコードを示す。データベースから取得する代わりに、説明用のレビュー2件をPython内で用意する。Python 3.12とquail-engineを導入し、H100 SXMを利用できる環境を前提とする。

まず、レビューを格納するテーブルを作り、EngineConfigでモデルにQwen3 4B FP8、GPUにH100 SXMを1基指定する。

import pyarrow as pa
import quail

reviews = pa.table({
    "id": ["r1", "r2"],
    "review": [
        "The ending was moving. I recommend watching this movie.",
        "The soundtrack was pleasant, but I would not recommend it.",
    ],
})

config = quail.EngineConfig(
    gpus=1, model="qwen3-4b-fp8", backend="quail", device="h100-sxm"
)

続いて、このデータと設定を使ってセッションを開く。session.registerでテーブルをreviewsという名前で登録し、2つの自然言語条件を指定したSQLを実行する。

with quail.Session(config) as session:
    session.register(
        "reviews", quail.DocumentProvider.from_table(reviews, id_col="id")
    )
    query = session.sql(r"""
        SELECT r.id
        FROM reviews AS r
        WHERE AI.IF(PROMPT(
            'Does this review discuss the ending of the movie?\n\n{0}',
            r.review
        ))
        AND AI.IF(PROMPT(
            'Does the reviewer recommend watching the movie?\n\n{0}',
            r.review
        ))
    """, dialect="bq")
    print(query.run().collect())

PROMPTは質問文の{0}に各行のレビュー本文を埋め込む。各AI.IFで、レビューが「結末について言及している」か、「⁠映画を推薦している」かをLLMに判定させる。Quailは、ANDで結んだ両方の条件を満たす行を絞り込む。

query.run()でクエリを実行し、collect()で結果をPythonプログラムに取り込む。この例ではSELECT r.idと指定しているため、条件を満たしたレビューのIDを含むArrow形式のテーブルが返る。

GPUの待機と同じ入力の再計算を減らす仕組み

レビューの件数や判定条件が増えると、LLMによる判定回数も増える。大量の判定を短時間で終えるには、GPUの待機や同じ入力の再計算を減らすことが重要になる。AI-SQLでは、LLMへの判定要求を事前にほぼ把握できるため、Quailはクエリ全体を見通して、処理の順序とキャッシュの再利用を計画する。

開発側はQuailの比較対象として、AI-SQLで必要になるLLMの判定処理を、汎用のLLM推論エンジンvLLMで実行する構成を用意した。同じモデルに同じ判定をさせる際の、処理の進め方やキャッシュ管理の違いを比較している。

vLLMを使った構成では、設定を調整しても、CPUが判定要求の管理や実行順の調整に時間を費やし、GPUに待機時間が生じた。この待機時間を減らすため、Quailでは複数の行をバッチという処理単位にまとめ、GPUが現在のバッチを処理している間にCPUが次のバッチを準備する仕組みを採った。

また、vLLMでは後の判定で必要になる計算結果がキャッシュから破棄され、余分な再計算も発生した。この再計算を減らすため、Quailでは、LLMが各行のテキストを処理した際の計算結果を保持するKVキャッシュをGPUメモリー上に残し、次の判定でも再利用する仕組みを採った。たとえば、先ほどのレビューの例なら、1つ目の条件を満たしたレビューのキャッシュを2つ目の条件の判定で再利用する。後で使わないキャッシュは解放し、再利用する分の領域を確保する。

Quailはキャッシュを保持するだけでなく、結合処理中に読み出す回数も抑える。一方のデータセットの1行をもう一方の複数行と比較する場合、比較元の行は毎回同じでも、個別に処理するとそのKVキャッシュを何度も読み出すことになる。そこで、比較元の行が同じ処理をまとめ、キャッシュの繰り返し読み出しを減らしている。

29クエリの評価で1.84倍高速化⁠、vLLMが上回るケースも

開発側はQuailの性能を、映画レビュー、医療レポート、法的文書、AIエージェントの実行履歴などを扱う29クエリのベンチマーク「QUAIL-B」で評価した。モデルにはQwen3 4B FP8、GPUにはH100を1基使用し、QuailとvLLMを同じ物理GPU上で順番に実行した。モデルとプロンプトを統一し、フィルターや結合の組み合わせと順序を示す論理クエリ計画も、Quailが選んだものにそろえた。

全29クエリの比較には、主な入力データがベンチマークの基準規模の10%となる条件(スケール係数0.1)を用いた。その結果、Quailは27クエリで、フィルターや結合を1つ終えてから次へ進むvLLMの比較構成よりも短い時間で処理を終えた。各クエリで何倍速くなったかを幾何平均でまとめると、1.84倍だったという。

別途、基準規模の全量を使う条件(スケール係数1.0)で、医療レポートを扱う「BIO-4」を評価した。このクエリは、レポート5000件と有害反応を表す4144件の用語を使い、フィルター処理と2段階の結合を行う。実行時間はvLLMの6.84時間に対しQuailは29.26分で、14.04倍高速だった。Quailの入力処理速度は毎秒1903万トークンに達し、毎分10億トークンを超えたという。[1]

一方、AIエージェントの実行履歴を分析する「AGENT-1」では、vLLMが2.32倍高速だった。このデータでは、別々の行に格納された実行履歴でも、冒頭に同じ内容が含まれている。vLLMの自動プレフィックスキャッシュは、その一致を検出し、冒頭部分の計算結果を再利用できた。Quailは同じ行のデータを繰り返し使う場合のキャッシュ再利用には対応するが、別々の行のテキストを照合して共通する冒頭部分を見つける機能は未対応という。開発側は、この機能を今後の開発項目に挙げている。

このほか、AI-SQLで実行できる処理の種類や、対応するモデルとハードウェアを増やす方針も示した。公式ブログでは、Quailを試せるPlaygroundも案内している。

QuailのPlayground画面

おすすめ記事

記事・ニュース一覧