# BLASST: Dynamic BLocked Attention Sparsity via Softmax Thresholding > [!info] Talk metadata > - **会議:** [[MLSys2026]] Day 2 (May 19 / Tue)、Grand Ballroom 1、Best Paper Session（08:45 PDT 開始、第3発表、Best Paper 受賞） > - **登壇者:** Cameron Shinn（UC Davis, John Owens 研究室 5 年目 PhD。共同筆頭著者・登壇者） > - **全著者:** Jiayi Yuan¹*, Cameron Shinn²*, Kai Xu³, Jingze Cui³, George Klimiashvili³, Guangxuan Xiao³, Perkz Zheng³, Bo Li³, Yuxin Zhou³, Zhouhai Ye³, Weijie You³, Tian Zheng³, Dominic Brown³, Pengbo Wang³, Markus Hoehnerbach⁴, Richard Cai³, Julien Demouth³, John D. Owens², Xia Hu¹, Song Han³, Timmy Liu³, Huizi Mao³（*: equal contribution） > - **所属:** ¹Rice University、²UC Davis、³NVIDIA、⁴Meta > - **URL:** https://mlsys.org/virtual/2026/session/3659 > - **関連研究・統合先:** TensorRT-LLM / FlashInfer（kernel 提供）、NVIDIA ModelOpt（閾値キャリブレーションツール・レシピ）、SGLang（対応済み）、vLLM（PR 進行中、coming soon）。FlashAttention-4 (Zadouri et al., 2026) の warp specialization 上に実装 > [!abstract] 概要（論文 PDF アブストラクト） > 大規模言語モデル(LLM)の長コンテキスト推論需要の高まりは、self-attention に内在する計算・メモリのボトルネックを深刻化させている。この課題に対し、固定スカラー閾値のみで attention ブロックをスキップして推論を高速化する、drop-in 型の動的スパース attention 機構 BLASST を提案する。本手法は既存研究の採用障壁を取り除くことで実用的な推論デプロイを対象とする。すなわち BLASST は学習不要、高コストな事前計算パスの回避、prefill と decode 双方の高速化、全主要 attention 変種(MHA, GQA, MQA, MLA)対応、modern hardware への最適化サポート、既存フレームワークへの容易な統合を実現する。これは online softmax 統計を再利用して無視可能な attention スコアを特定し、softmax・value ブロックロード・後続行列積をスキップすることで達成される。無視可能なレイテンシオーバーヘッドの最適化 kernel として BLASST アルゴリズムを実証する。自動閾値キャリブレーション手順は、最適閾値とコンテキスト長の間に単純な反比例関係を明らかにし、モデルごとに prefill 用・decode 用それぞれ単一の閾値のみを必要とする。ベンチマーク精度を維持しつつ、modern GPU 上で prefill は 71.9% スパース性で 1.52倍、decode は 73.2% スパース性で 1.48倍の高速化を実証する。 ## 背景: LLM のスケールと attention の支配 - LLM 推論のコンテキスト長は劇的に増大。system prompt トークン、reasoning / chain-of-thought トークン、content トークン(コード・画像・動画)のいずれも増加。VLM・diffusion model の普及も多トークン消費を後押し。 - コンテキスト長 $n$ に対し attention は $O(n^2)$ の計算・メモリアクセスを要する。スライドの latency 内訳(Llama-3-8B on A100、出典 LServe, Yang et al., 2025)では、入力長を 8K→128K に伸ばすと attention が prefill・decode いずれも支配的になる(decode では 128K で約 85% を占める)。長コンテキスト model のデプロイを困難にする根本要因。 - ℹ️ FlashAttention とその後継はメモリ帯域利用・kernel fusion を最適化したが、full attention matrix の計算は残り、二次計算量は未解決。 ## attention は本質的にスパース・既存手法の採用障壁 - attention は本質的にスパースで特定パターンが観測される: global / windowed / random sparsity (BigBird, Zaheer et al., 2020)、attention sink (StreamingLLM, Xiao et al., 2023)、Λ-shape / vertical-slash / block-sparse パターン (MInference, Jiang et al., 2024)、DeepSeek 系の dynamic sparse 変種 (HCA/CSA/DSA)。スパース性活用の研究蓄積は大きい。 - 既存手法は **5 つの採用障壁**を抱える: 1. スパースパターン決定のための事前計算が高速化分を相殺。 2. 追加層・新アーキテクチャの学習が必要。 3. prefill / decode のどちらか一方のみに限定。 4. modern GPU 向けに最適化・サポートされていない。 5. attention API への侵襲的改変でフレームワーク採用が困難。 > [!note] 関連手法との位置づけ(論文 Table 1・最終スライド) > Accelerates Prefill / Accelerates Decode / No Training / No Pre-Computation の 4 軸で、H2O・SnapKV・RocketKV・Quest・DuoAttention・DSA・MInference・SpargeAttention・XAttention と比較。**BLASST は 4 軸すべてを満たす唯一の手法**。最も設計が近いのは SpargeAttention だが、(1) BLASST は専用 kernel で prefill・decode 双方を最適化(SpargeAttention は prefill のみ)、(2) スキップ判定を計算済み統計で行いオーバーヘッドゼロ(SpargeAttention は別の prediction step)、(3) decode kernel が HBM からの Value ロードをスキップ、の 3 点で異なる。 ## 提案: BLASST アルゴリズム - **softmax thresholding による dynamic block sparse attention**。精度影響は最小、prefill・decode 双方を約ノーオーバーヘッドで高速化、学習不要、Blackwell / Hopper kernel サポート、単純な kernel backend swap で既存フレームワークにシームレス統合。 - **コア発想:** FlashAttention の block-wise online softmax 中、各ブロックの寄与は softmax 正規化後に無視可能になることが多い。block の local maximum が running maximum より閾値 $\lambda$ 分だけ著しく小さければ、正規化後の post-softmax 値はほぼゼロになり、当該 block の計算を完全にスキップできる。 - **スキップ判定式(論文 Eq.2、Algorithm 1 line 7):** online softmax がすでに計算する local row max $\tilde{m}^{(j)}_i$ と running row max $m^{(j)}_i$ を再利用し、行方向に block max を取る。 $\tilde{m}^{(j)}_i - m^{(j)}_i < \ln(\lambda)$ が成立すれば、exponentiation 後に $\exp(\tilde{m}^{(j)}_i - m^{(j)}_i) < \lambda \approx 0$ となり、この block で **softmax・$V$ ロード・後続の $\tilde{P}_{ij}V_j$ 行列積をスキップ**する。 - ℹ️ 三段階近似: 各スコア $S_{ij}$ の理想的重要度は(未知の)global max に対する相対値だが、true max のオンライン計算は高コストなため running max を扱いやすい proxy として用い、さらに kernel 内の効率的な block 単位判定のため token 単位を block-local max に置換 → 安価な `(block_max - running_max) < ln(λ)` 条件を得る。 - **FlashAttention への変更は極小**(スライド・論文 Algorithm 1、変更は緑色で 4 行程度): online softmax の rowmax 演算を再利用、既存 FA と同じ block size を用いタイル計算に直接マッピング。 ## 閾値キャリブレーション - 適切な閾値 $\lambda$ はモデル・コンテキスト長・目標精度に依存。スパース性が増えると精度は低下する。論文 Figure 2(左)は、全データセット・コンテキスト長で精度劣化パターンが一貫し、約 60–70% スパース性まで精度が安定、それ以降で急落することを示す。**精度劣化はスパース率そのもので主に決まり、データセット種別やシーケンス長では決まらない**。 - ⚠️ 一定の固定閾値はコンテキスト長間でスパース率が大きくばらつく(論文 Table 6: 目標 50% に対し固定 $\lambda=10^{-3}$ では 4K で 23.09%、64K で 74.63%)。固定スパース率を保つには適応的キャリブレーションが必須。 - **キャリブレーション手順(Algorithm 2、4 ステップ):** (1) 1 パスで複数候補閾値を一括スイープ(calibration set を 1 回実行)、(2) (閾値, シーケンス長) ごとに観測スパース率を記録、(3) 指数モデル $\lambda \cdot L = a \cdot \exp(b \cdot S)$ をフィット(出力はモデルごとに 2 スカラー $a, b$)、(4) 任意の目標スパース率 $S$ を選べば閾値が自動設定されコンテキスト長 $L$ に応じて再スケール。 - **コンテキスト長との反比例関係(論文 Eq.3):** 最適閾値は $\lambda = a/L$($a$ はモデル固有のスケール係数)。理論的根拠: attention スコアは行方向に和が 1 へ正規化されるため、長いシーケンスほどトークンあたり平均スコアが小さく、比例して小さい閾値を要する。 - ℹ️ prefill 用に 1 スカラー、decode 用にもう 1 スカラーを使うのが一般的(スライドのデプロイ例でも prefill / decode 別閾値を指定)。キャリブレーション済み $a$ はタスク種別をまたいで安定転移する(論文 Appendix A Table 12)。 ## kernel 設計: skip predicate reduction - **FlashAttention の前提:** $N^2$ attention matrix を HBM に保存・ロードせず、周辺演算 $\mathrm{softmax}(QK^\top)V$ を 1 kernel に fusion し online softmax を用いる。FlashAttention-4 (Zadouri et al., 2026) では Blackwell 向けに MMA と softmax を critical path 上で専用 warp に specialize。 - **skip predicate reduction(スライド該当図):** block max を critical path 外(correction warpgroup)で計算。各 thread が自分の block row のスキップ predicate を計算 → warp 内で `VOTE` 命令(`__all(skip)`)により単一値へ reduction → warp ごとに 1 thread が shared memory(SMEM)へ `atomicAnd(skip)` を発行し warp group 間で reduce。これで他 warp が読むタイル全体のスキップ条件が得られる。判定命令は既存演算の背後に隠れオーバーヘッド無視可能。 - **prefill(compute-bound): matmul と softmax をスキップ。** CUDA core(softmax)・tensor core(matmul)スループットが律速。スキップ対象 block の MMA(BMM2)と softmax を動的に飛ばし、後続演算を前倒しスケジューリングしてパイプラインを圧縮(論文 Figure 3 では 50% スパース性で 18 → 14 time unit)。なお prefill kernel では Value block は HBM からロードし続ける(メモリ帯域が律速でない・prefetch パイプラインが予測可能アクセスで恩恵・条件付き V ロードのレイテンシが節約を上回るため)。 - **decode(memory-bound): $V$ ロードをスキップ。** KV cache フェッチの HBM 帯域が律速。pruned block の memory-intensive な $V_j$ ロードをスキップし HBM ボトルネックに直接対処。**$K$/$V$ ロードの並べ替え**を行う: 既定の FlashAttention は $K$ と $V$ のロードを交互に行うが、$V$ は skip 条件計算(BMM1 完了)後でないと prefetch できない。そこで batched load scheduling に再設計し $K_1^\top Q, K_2^\top Q, \dots, K_B^\top Q$ を連続ロード、閾値判定を通った $V_j$ だけをまとめて発行することでパイプライン bubble を解消(論文 Figure 4 では 38 → 31 time unit、$B$ 個の $S_j$ 用 SMEM buffer が必要だがクエリ長 1 なので小さい)。 - ℹ️ MLA (Multi-head Latent Attention) は decode でも compute-bound になり得るため、MLA では計算側のスキップでも高速化を得られる。 ## 結果 > [!note] 精度維持(論文/スライド Table 2、Llama-3.1-8B / Qwen3-8B、Dense / 50% / 75%) > | Model | Sparsity | Prefill RULER-32K | Prefill LongBench | Decode MATH500 | Decode AIME2024 | Decode GPQA | P+D RULER-32K | P+D LongBench | > |---|---|---|---|---|---|---|---|---| > | Llama-3.1-8B | Dense | 92.33 | 31.40 | 73.40 | 46.66 | 46.71 | 92.33 | 31.40 | > | | 50% | 91.81 | 31.80 | 73.71 | 46.15 | 46.31 | 91.79 | 32.40 | > | | 75% | 91.67 | 31.80 | 73.89 | 46.01 | 45.95 | 91.67 | 31.80 | > | Qwen3-8B | Dense | 91.90 | 33.60 | 95.87 | 75.00 | 61.21 | 91.90 | 33.60 | > | | 50% | 92.08 | 35.10 | 96.23 | 76.50 | 61.56 | 92.07 | 33.30 | > | | 75% | 92.11 | 34.40 | 96.07 | 75.33 | 61.51 | 91.74 | 33.10 | > 約 75% スパース性でも精度劣化は最小。Qwen3-8B では MATH500(96.23 vs 95.87)・AIME2024(76.50 vs 75.00)で 50% スパース性が dense baseline を上回る。 - **prefill-only 精度(Table 3、Llama-3.1-8B、約 50% スパース性):** RULER 平均 92.87、LongBench Overall 31.8。全スパース手法中で最良の overall 精度、dense(RULER 93.21 / LongBench 31.4)に肉薄。MInference(RULER 84.15)・FlexPrefill(87.72)を大きく上回り、閾値ベース手法が proxy 重要度推定に勝ることを示す。 - **decode-only 精度(Table 4、Qwen3-8B、約 50% スパース性):** RULER-32K 91.55 / LongBench 33.90 / MATH500 96.23 / AIME2024 76.50 / LiveCodeBench 54.15 / GPQA 61.51、平均 68.97。全ベンチで dense(平均 68.57)を一致または上回る。参照として Quest(平均 60.75)・RocketKV(66.91)。 - **kernel 性能(Table 5・スライド、SoL FlashAttention-3 BF16 baseline 比):** 近損失精度(約 50% 目標)で Blackwell(B200)約 1.33倍(prefill)/ 1.25倍(decode)。高スパース性(約 70%)で prefill 1.52倍・decode 1.48倍へ増加。Hopper(H200)は prefill 71.0% スパース性で最大 1.52倍。**0% スパース性で性能劣化なし(0.96–1.00倍)**、スキップ判定計算が Tensor Core(prefill)/ HBM ロード(decode)の背後に隠れることを検証。 - ℹ️ end-to-end(Figure 5、Qwen3-30B-A3B-Instruct-2507、TensorRT-LLM in-flight batching、concurrency 64): TTFT・TPOT で 1.1倍高速化、LongBench V1 精度の低下は marginal。 ## 使い方・フレームワーク対応 - 数行で有効化。NVIDIA ModelOpt が閾値をキャリブレーション(`mtsa.sparsify(model, config=SKIP_SOFTMAX_CALIB)`)。 - TensorRT-LLM サーバ(`SkipSoftmaxAttentionConfig` に `threshold_scale_factor`、prefill / decode 別指定可)または SGLang サーバ(環境変数 `SGLANG_SKIP_SOFTMAX_PREFILL/DECODE_THRESHOLD_SCALE_FACTOR`)でデプロイ。閾値環境変数で kernel backend を swap。 - kernel は TensorRT-LLM と FlashInfer に提供。閾値キャリブレーションツール・レシピは ModelOpt に。SGLang は対応済み、vLLM は PR 進行中(coming soon)。 ## 結論 - BLASST は学習不要・事前計算なし・proxy スコアなしで prefill・decode 双方を高速化する drop-in 型 dynamic block sparse attention。online softmax 統計の再利用により最小精度損失とゼロに近いオーバーヘッドを両立。 - 自動キャリブレーションが閾値とコンテキスト長の反比例関係 $\lambda = a/L$ を明らかにし、モデルごと prefill / decode 各 1 スカラーで任意の目標スパース率を予測可能に制御。 - modern GPU(Blackwell / Hopper)向け最適化 kernel と主要フレームワーク統合により、実用デプロイの障壁を取り除く。 ## Q&A - **Q1（質問者名不鮮明 ※トランスクリプトは "Speaker 1"）— decode-only でスパース化が baseline を上回る理由:** なぜスパース化した方が良い性能が出るのか。 - **A:** model をスパース化したときによく見られる効果。ノイズを削ることで正則化効果が働き、無視できるデータを除いて実際のシグナルに集中することで baseline を上回ることがある。確たる説明は持たないが、これが観測された事実。(論文では(1)長コンテキストで低 attention block の pruning が確率質量を最重要トークンに集中させる implicit denoising、(2)長生成 reasoning で冗長・有害な中間 reasoning step のスキップ、と補足) - **Q2（Roshan, Microsoft Research）— 構造化スパース性(causal / sliding window)への適用:** causal attention や sliding window attention のような構造化スパース性にはどう適用するか。 - **A:** causal attention では動作する(LLM decoding 推論で常用)。windowed attention でも動作し得るが、現時点で実装は提供していない。比較的容易に適応可能だが追加作業が必要。 - **Q3（Roshan, Microsoft Research）— 事前構築スパースパターンとの比較:** FlashAttention のようにスパースパターンを事前構築する手法との比較は。 - **A:**（unstructured sparsity の文脈で）現時点で当該の数値は持っていない。優劣を示唆したくない。ポスターに来てもらえれば議論したい。