Devotd Github
新着推論・サービング

flashinfer-ai/flashinfer — Cake GQA 学習アテンションの境界検査修正

FlashInfer は、Cake の可変長 GQA 学習アテンションにおけるホスト側境界検証の欠落、int32 タイルカウンタの折り返し、apply の引数個数の問題を修正する正確性修正コミットを公開した。

リポジトリ情報flashinfer-ai/flashinfer ↗

FlashInfer は、Cake の可変長 GQA 学習アテンションにおけるホスト側境界検証の欠落、int32 タイルカウンタの折り返し、apply の引数個数の問題を修正する正確性修正コミットを公開した。

GitHub の説明: FlashInfer: Kernel Library for LLM Serving

スター
★ 6.6K
フォーク
1.6K
言語
Cuda
ライセンス
Apache-2.0
最新リリース
v0.7.2
原文を見る ↗GitHub で見る ↗

flashinfer-ai/flashinfer の最新コミット(2026-11-08)は、LLM サービング向けカーネルライブラリ FlashInfer の `cake_gqa_nondet_train` パスにある3つの境界チェック欠陥を修正する修正コミットである。このパスはヘッド次元256向けに構築された Cake の可変長因果 GQA 学習アテンションであり、長さが異なるシーケンスを1つのバッチとしてまとめて処理する。今回の変更はカーネルも実行計画も変えない正確性の修正である。

修正された3つは、いずれも入力境界の検証である。backward は、ホスト側のシーケンス境界リスト `cu_seqlens_q_host` と `cu_seqlens_k_host` が両方とも渡されるか両方とも無いこと、および被演算子のサイズと一致しているかを、メモリ確保や query・key が無い場合の早期リターンより前に確認する。デバイスのタイルカウンタは int32 であり、2^31 タイルを超えると値が折り返すため、範囲検査ではこの場合に int64 で境界表を作り直して不正な範囲を弾く。最後に `host_boundaries` の既定値を `None` とし、backward が forward の引数ごとに勾配を1つずつ返すようにすることで、引数が9個である `apply` 呼び出しが引き続き動作するようになった。テスト側には130行が追加され、折り返した int32 タイルテーブルに対する検査を子プロセスで実行している。

性能の変化はない。9種のシェイプ × 3種の演算(前方・後方・ステップ)を GB200、GB300、R200において FlashInfer main と同じモードで比較したところ、すべての行が A/A ノイズバンドの範囲内に収まった(GB200 ±0.69%–±3.34%、GB300 ±0.12%–±1.19%、R200 ±0.27%–±9.94%)。これは正確性修正として意図された結果である。

ただし、証拠がカバーできる範囲は狭い。int64 パスは計測したシェイプのいずれでも一度も発動せず、2^31 タイルまで大きくなり得るシェイプの挙動も計測によって確認されていない。小さい呼び出し8個の layout においても、GB300 の前方2行(`tails_small` 0.950x、`tile_edges` 0.980x)が ±3.21%、±0.91% のバンドを外れたが、前方コードに触れていない変更であることから、14~31 us の呼び出しにおける計測ノイズとして報告される。参考として、決定的な Cake パッケージ(`cake_gqa_det_train`、PR #9121)に対するステップは 0.992x–1.039x、行ごとに最速の同精度非 Cake 実装に対するステップは 1.050x–1.262x の範囲であった。つまり、このコミットは速度競争ではなく、不正な範囲を弾く安全装置に近い。