PFlash：RTX 3090で128Kにおいてllama.cpp比10倍のプリフィル高速化

やあ、仲間のLlamaたち。前回投稿した内容に対する、たくさんの温かい言葉と素晴らしいフィードバックをありがとう。共有して役立つかもしれない、また新しいものを用意しました。いつもながら時間は貴重なので、手短にいきます。

量子化した27Bターゲット向けのロングコンテキストデコードに対して、C++/CUDAのみで推測プリフィル（speculative prefill）を構築しました。小さなドラフタがプロセス内で、全文プロンプトに対してトークンの重要度をスコアリングします。重いターゲットは、本当に関係のあるスパンだけをプリフィルします。

リポジトリ: github.com/Luce-Org/lucebox-hub（オープンソース、MIT）。

Qwen3.6-27B Q4_K_M、RTX 3090、シングルショットでの対決: 24.8 s（TTFT） vs バニラllama.cppで約257 s ＝128Kで約10.4×（64Kでは13.5 s vs 134.95 s ＝10.0×）。NIAHのretrievalはエンドツーエンドで保持されます。推論ループ内にPythonはありません。Tritonもありません。PyTorchもありません。

問題

24 GBの3090上でのQ4_K_M Qwen3.6-27Bは高速にデコードできます（DFlashのspec decodeで約74 tok/s） が、プリフィルはO(S²)にスケールします。131Kトークンのプロンプトでは、バニラのllama.cppはコールドで248.4 sかかります（llama-bench pp131072 --no-warmup -r 1、527.6 tok/s）。最初のトークンまで何も起きない時間を4.1分見つめることになります。デコード自体は速いのに、その待ちがUXを壊します。ウォーム状態の定常値はまだマシ（128Kで169.3 s）ですが、それでもつらく、コンテキストを伸ばすほど二次的に増えていきます。

肩の上に立って

この取り組みは、2つの最近の論文に基づいています。どちらもとても読みごたえがあります:

• Speculative Prefill（Liu et al、arXiv 2502.02789）およびCross-Family Speculative Prefill（SambaNova、ICLR 2026）。洞察: 長いプロンプトに対する小さなドラフトモデルの注意（attention）パターンが、答えにとって重要なトークンを忠実に予測できる。ドラフトを走らせてトークンごとの重要度をスコアし、上位のスパンだけを残して残りを捨てる。
• FlashPrefill（Fan et al、2026）。ブロックスパース注意により、ドラフタ自身が128KでO(S²)を払わないようにする。
• FA-2由来のsm_80+ sparse forward向けに、mit-han-lab/Block-Sparse-Attention（BSA）。
• 実行時はggml / llama.cpp。libggml*.aはリンクしますが、libllamaはリンクしません。

私たちの貢献は、この2つのアルゴリズムをC++/CUDAでインプロセス合成し、24GBのコンシューマカード上で動かしたことです。私たちの認識では、この2つの論文は先にオープン実装で組み合わせられていませんでした。

構築したもの

• インプロセス合成。ドラフタのforward（カスタムQwen3-0.6B BF16 ggmlグラフ）、FlashPrefillのスコアリング、スパース注意、ターゲットのプリフィル、DFlashのspec decodeはすべて、同一のggmlアロケータを共有する1つのC++/CUDAプロセス内で動きます。サブプロセスなし、IPCなし、Pythonなし、Tritonなし、推論ループ内にPyTorchなし。
• FlashPrefillのCUDA移植。リファレンス（qhfan/FlashPrefill）はTritonです。私たちはスクラッチから4つのCUDAカーネル（mean_K、score、select、sparse_fwd）を書き、スパースforwardはmit-han-lab/Block-Sparse-Attentionを通してディスパッチしました。BSAはlibtorchのC++拡張として提供されますが、24GBの推論ループに2GBのlibtorchを持ち込むのは現実的でないため、dflash/deps/bsa_stubs/配下の3ヘッダのATen/c10スタブセット経由で配線しました。
• 24GBメモリのオーケストレーション。ドラフタ（1.3GBの重み＋KV＋128Kで約600MBのBSAスクラッチ）とDFlashデーモン（15GBのターゲット＋3GBのドラフト＋3GBのKV）は、3090上では同居できません。デーモンは、stdinプロトコルでステージ間に重みを退避（park）・復帰（unpark）し、解放します。1リクエストあたり約3秒で、パイプライン全体を1台のコンシューマカードに収めています。

セットアップ

bash

# サブモジュール付きでクローン（llama.cpp/ggml + Block-Sparse-Attentionを含む） git clone --recurse-submodules https://github.com/Luce-Org/lucebox-hub cd lucebox-hub/dflash # dflash + BSAカーネル（sm_80+、~10分のコールドコンパイル。cutlassを取り込む） cmake -B build -S . -DCMAKE_BUILD_TYPE=Release 
 -DCMAKE_CUDA_ARCHITECTURES=86 
 -DDFLASH27B_ENABLE_BSA=ON cmake --build build --target test_dflash test_flashprefill_kernels -j # 重みを取得（ターゲット + ドラフタ + spec-decodeドラフト） huggingface-cli download unsloth/Qwen3.6-27B-GGUF Qwen3.6-27B-Q4_K_M.gguf --local-dir models/ huggingface-cli download Qwen/Qwen3-0.6B model.safetensors tokenizer.json --local-dir models/drafter/ huggingface-cli download z-lab/Qwen3.6-27B-DFlash --local-dir models/draft/ # ベンチ cd ../pflash && pip install -e . python tests/niah_gen.py --n 1 --ctx 131072 --out /tmp/niah_128k.jsonl python tests/bench_niah_cpp.py 
 --bin ../dflash/build/test_dflash 
 --target ../dflash/models/Qwen3.6-27B-Q4_K_M.gguf 
 --draft ../dflash/models/draft/model.safetensors 
 --drafter-gguf ../dflash/models/drafter/qwen3-0.6b.gguf 
 --cases /tmp/niah_128k.jsonl --keep-ratio 0.05 

数値

RTX 3090、Qwen3.6-27B Q4_K_Mターゲット、q4_0 KV、DFLASH_FP_USE_BSA=1 DFLASH_FP_ALPHA=0.85 keep_ratio=0.05でのシングルショット。NIAHのsingle-needleを、エンドツーエンドのretrievalチェックとして使用しています。ベースラインは、デフォルトのf16 KVを使うバニラllama.cppです（KVは単純比較ではないです。q4_0 KVは短いコンテキストで約3% ALを要し、8.56から8.33へ。ベンチマーク済み）。

コンテキスト	PFlash TTFT	llama.cpp コールド	スピードアップ（コールド）	llama.cpp ウォーム
64K	13.5 s	134.95 s	10.0x	（より小さい）
128K	24.8 s	248.4 s	10.0x	169.3 s

これらはコールドキャッシュの数値です（プロセス起動直後の最初のリクエスト）。ウォーム対ウォームの倍率が小さいのは、llama.cppが128Kでキャッシュが温まると約169 sに落ち着くためです。どちらも実測値で、あなたのワークロードによって正しい判断は変わります。エンジンを常駐させるならウォームを使ってください。

プリフィル後のデコードは、DDTreeを使う標準的なDFlashのspec-decode経路で、Qwen3.6-27B Q4_K_Mで約74 tok/sの持続を達成しています。

品質

NIAH single-needle（マジックキー＋7桁の答えがフィラー内にランダム配置）を、32Kから128Kまでのテストしたすべてのコンテキストでretrievalできています。keep_ratio=0.05、DFLASH_FP_ALPHA=0.85。

正直な注意点: NIAH single-needleは、私たちのようなattentionベースの選択手法に対して構造的に簡単なプローブです。アルゴリズムが「注意が高い単一のスパン」を見つけるのに適しているからです。RULERとNIAHのmulti-needleは次のリストです。公正な監査では、それらの数値を待つべきです。

なぜこのスタックが機能するのか

推測プリフィルは、品質面の問題を解決します。答えに関連する内容を失わずに、どう圧縮するのか？ FlashPrefillは、ドラフタのステップ内部における速度の問題を解決します。128Kでドラフタを、ボトルネックにならないほど速くするにはどうすればいいのか？ ターゲット側（DFlash spec decode）は変更されないため、両者はうまく合成できます。ターゲットには、フルのattentionが有効になった、はるかに短いプロンプトが渡されるだけです。

128Kでは、ドラフタのスコアリングが支配的なコストになっています（24.8 sのTTFTのうち約12 s）。圧縮された約6.5Kのサバイバーに対するターゲットのプリフィルは約10 s。残りの約3 sは、park/unpark/freeのダンスです。次に明らかなレバーは、より小さい、あるいは蒸留したドラフタですが、これはまだ行っていません。

チューニング

bash

DFLASH_FP_USE_BSA=1 # BSAを通してスパースなFAをフォワードディスパッチ（sm_80+、10xに必要） DFLASH_FP_ALPHA=0.85 # ブロック選択のしきい値。高いほど厳密＝Q行あたりのKブロック数が減る DFLASH_FP_PROFILE=1 # ステージごとの計測をログ出力（mean_K / score / select / forward） 

keep_ratio=0.05 がデフォルトです。0.02 だと目標のプリフィル時間が約10秒から約3秒に短縮されますが、針（ニュアンス）を失い始めます。 DFLASH_FP_ALPHA=0.99 は、128Kで約1秒短縮しますが、小さなNIAHマージンの損失が伴います。キャリブレーション領域です。

フィードバックは大歓迎です！