第8回「EVO-X2（Ryzen AI Max+ 395 / 128GB）+ Qwen3.8-Flash-Nextをいじっていたら、禁断のllama.cpp改造に手を出していた話」
第9回「EVO-X2（Ryzen AI Max+ 395 / 128GB）+ Qwen3.8-Flash-Next用のllama.cpp改造再開でベースのリポジトリを変更して、ROCm版を初ビルドした話」
の続き。

前回は環境を整備しただけで終わってしまったので、今回はガッツリ高速化をやりたいと思った。

第8回とその後のテストでは、LaurentZuijdwijkのStrix Halo専用forkベースの改造（r2）で、QSA grouped-unionによるPP高速化や、MTPによるTG高速化などを確認している。
他にもAgentionAI版モデルを動かすためのROCmFPx対応など、r2にはあって第9回からの本家(upsteram)llama.cppベースの改造版(r3)へまだ移植していない機能がいくつも残っている。

これらを全部移植してから次に進んだ方が良いのかもしれないけど、前回r2とr3を比較して、長コンテキストでのTGが妙に遅いのが気になったので、そこを調べて新しい高速化項目を探したくなった。今回はここを深堀りしてみる。

以降、Prompt Processing（入力を読み込む処理）をPP、Token Generation（出力生成）をTGと表記する。どちらも単位はtok/sで、高いほど速い。

結論

前回の本家(upsteram)llama.cpp b11247 + PLE16テーブル対応(COMMON-001)では、コンテキストが長くなるほどTGが大きく低下していた。

まずVulkan版でr2とr3のdecode処理をプロファイリングして遅くなっている原因を調べ、 COMMON-004: incremental pooled-key cache を実装した。

VulkanのTGは、

64k  17.20 -> 25.95 tok/s
128k 12.00 -> 24.00 tok/s
256k  7.30 -> 21.03 tok/s

まで改善した。

ROCmでもCOMMON-004の効果はあったが、コンテキスト長256kではまだ11.60 tok/sまで落ちる。

そこで今度はRadeon GPU ProfilerでROCmを調べ、COMMON-005: selected QSA K/V gather を実装すると、ROCm TGは、

64k  20.55 -> 22.60 tok/s  (+10.0%)
128k 16.43 -> 20.41 tok/s  (+24.3%)
256k 11.68 -> 17.37 tok/s  (+48.7%)

まで改善した。

最終的なTGの推移はこうなった。

Patch	Backend	64k TG	128k TG	256k TG
COMMON-001	Vulkan	17.20	12.00	7.30
COMMON-004	Vulkan	25.95	24.00	21.03
COMMON-005 ON	Vulkan	25.42	24.12	21.69
COMMON-001	ROCm	14.86	9.91	5.74
COMMON-004	ROCm	20.62	16.44	11.60
COMMON-005 ON	ROCm	22.60	20.41	17.37

特にROCm 256kは、

5.74 -> 11.60 -> 17.37 tok/s

と段階的に改善できた。

COMMON-005後にもコンテキスト長に比例する処理が1つ残ったので、次は COMMON-006: block-domain top-k を考えたが、最新upstreamを調べると、ほぼ同じ方向のQSA/k-pool再設計がすでに入っていたため、旧b11247へ独自実装するのは中止した。

第9回と今回のllama.cpp b11247ベース版(r3)はCOMMON-001/004/005を実装して効果を検証完了したところで凍結し、次回はまたupstreamを入れ替えることにした。

ソース、Windowsビルド、検証スクリプトなどはここで公開している。

main ブランチはStrix Halo専用forkベースの安定版（第8回の内容+α）でいったん凍結し、第9回からの作業は r3/upstream-first ブランチで行っている。今回の作業も r3/upstream-first ブランチで行った。

検証環境

速度に関係するところだけ先に。

項目	内容
PC	GMKtec NucBox EVO-X2
APU	AMD Ryzen AI MAX+ 395 / Radeon 8060S
RAM	128GB UMA
UMA Frame Buffer	96GB
OS	Windows 11 Pro 25H2
モデル	Unsloth Qwen3.8-Flash-Next UD-IQ3_XXS
モデル配置	PLE16変換版
KV cache	f16
ubatch	1024
MTP	OFF
Vulkan compiler	LLVM / Clang 20.1.8
ROCm	10.0.0 / TheRock
ROCm compiler	AMD Clang 23.0.0
GPU target	gfx1151
主な入力	61,789 / 126,253 / 255,181 tokens

詳細は記事末尾にまとめる。

QSAとQSA indexerについて

今回の話では QSA と QSA indexer が何度も出てくるので、先にざっくり説明しておく。

QSA（Qwen Sparse Attention） は、長いcontext全部にそのままAttentionするのではなく、まず「今のtokenから見て重要そうな場所」を選び、その候補だけにAttentionする仕組み。

Qwen3.8-Flash-NextはGated DeltaNet（GDN）とQSAを組み合わせたhybrid構成で、このモデルでは48 layerのうち4 layerごとに1つ、合計12個のQSA layerが入っている。

普通のfull attentionなら、

query
  ↓
context内の全K/Vを見る
  ↓
Attention

となるところを、QSAでは大まかに、

query
  ↓
QSA indexerで重要そうなblockを探す
  ↓
選ばれたtoken/blockだけに絞る
  ↓
本体のAttention

という2段階に分ける。

QSA indexer

QSA indexer は、本体のAttentionより先に動く軽量な検索器のようなもの。

今回のモデルでは、indexer用のkeyを4 token単位のmicro-blockへまとめて、各blockの代表vectorを作る。

概念的には、

4 token分のindexer key
  ↓
average pooling(平均値プーリング、サイズ縮小)
  ↓
RMS norm(二乗平均平方根(RMS)をつかった正規化)
  ↓
RoPE(Rotary Position Embedding)で位置情報を付加
  ↓
queryとのscoreを計算
  ↓
高scoreのblockをtop-k選択

という流れ。

このblock単位の代表vectorが、この記事で何度も出てくる pooled summary。
全tokenを1個ずつ検索するより、先に4 token単位へ圧縮してから探した方が、長コンテキストでindexer自体の負荷を抑えられる。

今回の設定では indexer_top_k = 2048、compression ratioが4なので、最終的に本体Attentionへ渡す候補はおおむね約2,048 token分（実装上は現在の未完成block分を含めて最大約2,051 cell）になる。

つまり256k contextでも、QSAの狙いとしては、

256k全部へAttention

ではなく、

256kからindexerで重要箇所を検索
  ↓
約2k token分へ絞る
  ↓
そこだけAttention

にしたい。

今回のCOMMON-004とCOMMON-005は、この2段階の別々の場所を高速化している。

COMMON-004
QSA indexer側
pooled summaryを毎token全再計算しない

COMMON-005
QSA本体Attention側
選択済みK/Vだけをcompact化してAttentionする

Vulkan版のプロファイリングでTGが遅くなった原因を調べる

まず、前回気になった「r3の長コンテキストTGがr2より妙に遅い」の原因をVulkan版で調べた。

Vulkan版llama.cppには、環境変数

$env:GGML_VK_PERF_LOGGER = "1"

を設定すると、Vulkan backend内の処理時間を細かく記録できるperformance loggerがある。

通常のPP/TG計測と違ってlogger自体のオーバーヘッドがかなりあるので、logger ON時のtok/sを通常計測値と比較する用途では使わない。

今回は、同じ64k入力でr2とr3の単語を連続出力している安定状態のデータを取り、

1 tokenあたりのGPUグラフ実行時間
各operationがどの程度時間を使っているか
r2とr3でどのoperationが増えたか

を見ることにした。

比較した結果がこちら。

Build	QSA union	PP（logger ON）	TG（logger ON）	steady GPU graph
r2	OFF	221.44	7.78	44.18 ms/token
r2	ON	263.28	7.77	44.31 ms/token
r3 + COMMON-001	-	259.73	6.62	59.13 ms/token

r2では第8回で入れてPP高速化に効果があったQSA grouped-unionがONでもOFFでもTG側のGPUグラフ実行時間はほぼ同じだった。つまり、r3のTGが遅い原因は「r2で入れたgrouped-unionをまだ移植していないから」ではない。

安定生成中の1トークンあたりのグラフ実行時間は、r2の44.18 ms/tokenに対して、r3は59.13 ms/token。約14.95 ms/token増えている。
増えた処理を分解すると、

処理	r3で増えた時間
CONT	+11.38 ms/token
full-block indexer RMS norm	+2.18 ms/token
QSA関連RoPE	+0.59 ms/token
合計	+14.15 ms/token

ここで出てくる処理を簡単に説明すると、

CONT
ggml(テンソル計算ライブラリの名前)のtensorを連続したメモリ配置（contiguous）へ並べ直す処理。今回のr3では、QSA indexer用のsummaryを毎decodeで作り直す途中で大きなtensorの再配置が繰り返され、ここが一番大きな差になっていた。
full-block indexer RMS norm
QSA indexerが使う、完成済みblockのsummary vectorへRMS normalizationをかける処理。block数はcontext長と一緒に増えるため、全blockへ毎token適用すると長コンテキストほど重くなる。
QSA関連RoPE
indexer用のkey / summaryへ位置情報を反映するRotary Position Embedding（RoPE）の処理。これも全block分を毎decodeやり直していたため、context長に応じたコストになっていた。
QSA indexerのpooled summary
QSAが「長いKV cacheのどこを見るか」を決めるために、複数tokenをblock単位で圧縮して作る代表vector。全tokenを直接比較する代わりに、このsummaryを使って候補blockを絞り込む。COMMON-004では、完成済みblockのsummaryをcacheして再利用する。

約14.95 ms/tokenの差のうち14.15 ms、**約94.6%**をこの3つで説明できた。

ソースとgraphを追ってみると、r2とr3ではQSA indexerのpooled summaryの扱いが違っていた。
r2では、最初のdecodeで完成済みblockのsummaryを作った後、それを再利用している。
r3では、decodeするたびにraw indexer cacheからcontext全体の完成済みblockをもう一度pooling → RMS norm → RoPEしていた。
長いcontextほど毎tokenの再計算量が増えるので、64k→128k→256kとTGが急落していたのも説明できる。
最初に、ここを直すことにした。

incremental pooled-key cache（COMMON-004）

対策は単純で、

一度計算した完成済みblockのpooled keyをcacheして、変化したblockだけ更新する

というもの。これを管理ID COMMON-004 とした。

各QSA layerごとに、完成済みposition blockのsummary rowをpersistent cacheへ保存する。

decode時には、

(従来)
raw indexer K cache
  ↓
全blockをpooling
  ↓
全blockをRMS norm
  ↓
全blockをRoPE
  ↓
QSA top-k

(COMMON-004)
persistent pooled-key cache
  ↑
新しく完成 / 無効化されたblockだけ更新
  ↓
QSA top-k

という形にした。

state loadやsequence editでcache内容を信用できなくなった場合はinvalidateし、必要になったところで再構築する。

また、QSA layerだけを対象にしてrecurrent layerは除外。speculative decodeのdirty writeでも範囲を壊さないようにした。

同一binaryでON/OFFテストできるように、

LLAMA_QSA_NO_POOLED_CACHE=1

でCOMMON-004を無効化できるようにもした。

decode向けには、

LLAMA_QSA_POOLED_MAX_TOKENS=32

をdefaultにして、大きなprefill ubatchでは従来pathへ戻るようにしている。

実装commitは、

6559dd272fd0d5f553823e8851c78b9edc2a5016
qwen4exp: cache pooled QSA keys across decode steps

速度計測は今までと同じ日本語テキスト要約タスクを用いて計測した。

COMMON-004実装結果

まず、同一binaryでCOMMON-004だけをON/OFFした速度計測結果。

Vulkan

Context	COMMON-004 OFF PP	ON PP	OFF TG	ON TG	TG改善
64k	269.18	262.89	17.13	25.95	+51.5%
128k	159.39	156.47	11.91	24.00	+101.6%
256k	102.77	98.63	7.27	21.03	+189.3%

ROCm

Context	COMMON-004 OFF PP	ON PP	OFF TG	ON TG	TG改善
64k	348.90	356.55	14.82	20.62	+39.1%
128k	262.67	261.67	9.89	16.44	+66.2%
256k	167.07	166.18	5.52	11.60	+110.1%

どちらのbackendでもTG改善はかなり大きい。

特にVulkanでは、64k→256kで急落していたTGが、

25.95 -> 24.00 -> 21.03 tok/s

となり、長コンテキストでの低下はかなり小さくなった。

一方、Vulkan PPはCOMMON-004をONにすると約2〜4%低下する傾向がある。
今回の目的はTGのボトルネック解消なので、PP側はここでは深追いしないことにする。

参考として、過去のr2 stable QSA unionと比べるとこうなる。

Build	64k PP	64k TG	128k PP	128k TG	256k PP	256k TG
r2 stable QSA union	274.67	23.55	222.77	21.65※	177.01	18.84
r3 COMMON-004 Vulkan	262.89	25.95	156.47	24.00	98.63	21.03

※r2 128k TGは第8回の記事本文には掲載していない計測値。

TGだけを見ると、COMMON-004後のr3 Vulkanはr2を約10〜12%上回った。

PPは逆で、r3にはr2のQSA grouped-unionをまだ移植していないため、特に128k/256kではr2の方がかなり速い。

ROCmの方もCOMMON-004は効果があったが、

64k  20.62
128k 16.44
256k 11.60

と、Vulkanよりコンテキスト長による低下がまだかなり大きい。
ROCmには別のコンテキスト依存コストが残っているらしい。

なので今度はROCm側をprofileしてみることにした。

ROCm版でもプロファイリングしてTGが遅い原因を調べる

ROCmの方はVulkanと違って、環境変数設定だけでプロファイルログを取る仕組みがないので、AMDの Radeon Developer Panel / Radeon GPU Profiler を導入した。

導入はここからWindows版zipをダウンロードして適当な場所に展開するだけ。
展開したフォルダから RadeonDeveloperPanel.exe を起動する。

Radeon Developer PanelはAMD driverと接続して、VulkanだけでなくHIP compute workloadのRGP profileも取得できる。

Radeon Developer Panelの設定

今回の使い方はざっくり、

RadeonDeveloperPanel.exe を起動
CONNECTIONタブでLocalへ接続
CAPTUREタブでProfilingを選択
ApplicationsのAPIでHIPを選択して、llama-cli.exe を対象として認識させる
Auto captureのdelay / dispatch範囲を設定
いつも通り別PowerShellから llama-cli.exe を実行

という流れ。

Auto captureのdelayはプロファイル開始時間、dispatch範囲は何回分データを取るかの設定で、安定生成中の１トークン分のデータを取りたいので、delayはprefillが終わってgeneration開始から終了までの間に設定したい。
ここは入力長によってprefill時間が全然違うので、128k用と256k用でタイマーを変えている。
delayは1トークン分の範囲を取りたいので、何度かデータを取ってみて、最終的には4096に設定した。

captureに成功すると .rgp profileが一覧に出るので、ダブルクリックするとRadeon GPU Profilerが開く。(スクショの青枠部分にrgpの一覧が出る。スクショでは実行前なのでまだ何も表示されていない)

今回主に見たのは、Radeon GPU Profilerで表示される項目のうち、

OVERVIEW / Most expensive events
どのkernelが重いかを見る
EVENTS / Wavefront occupancy
kernel dispatchの繰り返し方やGPU上での実行状況を見る
EVENTS / Event timing
個々のkernelの実行時間を比較する

の3つ。

今回のボトルネック特定には、特にEvent timingの128k / 256k比較が効いた。

128kと256kを比較してみる

COMMON-004 ONの状態でsteady single-token decodeをprofileした結果、代表的なkernelはこうなった。

Kernel	128k	256k	変化
`flash_attn_tile<256,256,1,4,false>`	約1.452 ms	約2.939 ms	約2.02倍
`k_get_rows_float`	約0.451 ms	約0.907 ms	約2.01倍
`mul_mat_vec_q<type14>`	約2.835 ms	約2.836 ms	ほぼ一定

kernel名も分かりにくいので簡単に説明すると、

flash_attn_tile<256,256,1,4,false>
ROCm/HIP側でFlash Attentionを実行しているGPU kernel。<...> はtileサイズなどの実装パラメータで、今回重要なのは「QSAのfull-attention部分を実行しているkernel」という点。128k→256kで約2倍になっていたため、full context長に比例する処理が残っていると分かった。
k_get_rows_float
指定したindexのrowをfloat tensorから取り出す get_rows 系kernel。今回のQSAでは、block単位で計算したscoreを cell_blk を使って全KV cellへ展開する処理に対応しており、展開先がcontext全体なので128k→256kでほぼ2倍になった。
mul_mat_vec_q
量子化weightのmatrixとvectorを掛けるquantized matrix-vector multiplication kernel。1 tokenずつ生成するdecodeでは頻繁に使われる主要な演算だが、今回の128k/256k比較では時間がほぼ一定だったため、「contextが長くなるほどTGが遅くなる原因」ではなかった。
mul_mat_vec_q<type14>
上の mul_mat_vec_q のうち、ggml内部の型ID 14 用に特殊化されたkernel。llama.cppでは型ID 14は Q6_K を表す。これはprofileに出たそのkernelがQ6_K tensorを処理していたという意味で、使用モデル全体がQ6_K量子化という意味ではない。

128kから256kへcontextを2倍にすると、Flash Attentionと k_get_rows_float がほぼ2倍。
mul_mat_vec_q はほぼ変わっていない。

profile全体では、

128k: 約62.11 ms/token
256k: 約85.79 ms/token
差:   約23.68 ms/token

だった。

Qwen3.8-Flash-Nextは48 layersで full_attention_interval = 4 なので、1 tokenあたり12個のfull-attention groupがある。
代表時間の増加を12倍すると、

Flash Attention:
(2.939 - 1.452) * 12 ≒ +17.84 ms/token

get_rows:
(0.907 - 0.451) * 12 ≒ +5.47 ms/token

合計 ≒ +23.31 ms/token

実測の+23.68 ms/tokenをほぼ全部説明できる。

なぜFlash Attentionがcontext長に比例していたのか

QSAではindexerが全KVから候補を選ぶが、このモデルでは、

indexer_top_k = 2048
compression ratio = 4

なので、最終的にattentionで有効になるトークン数は、

2048 + 4 - 1 = 2051

つまり128kでも256kでも、「本当に見たいtoken」は約2,051件程度しかない。
qwen4exp graph側も n_kv_max でこのsparse budgetをFlash Attentionへhintとして渡している。

ところがHIP backendを追ってみると、sparse Flash Attentionのmask compaction pathはHIPでは使われない。

そのためROCmでは、

論理上:
128k / 256kから約2051件を選択

実際のFlash Attention:
full K/Vをそのまま処理

となっていた。
contextを128k→256kへ2倍にしたらFlash Attentionがほぼ2倍になったのは、これなら納得できる。

Vulkanではbackend側にsparse Flash Attentionのmask compaction pathがあるため、選択された約2051件のみをAttentionに使っていたので、COMMON-004のみでTGの低速化はほぼ解消していた。

もう1つの k_get_rows_float は別の場所だった。
現在の build_qsa_top_k() は、

block score
  ↓
全KV cellへscoreを展開
  ↓
cell-level top-k

という流れで、top-kする前に一度context全体のcell数まで展開していた。
これが128k→256kで約0.451→0.907 msへ増えていた k_get_rows_float と対応する。

こちらは後の候補(COMMON-006)に回すことにして、まず時間の大部分を占めているFlash Attention側から直す。

selected QSA K/V gather（COMMON-005）

考え方は、

backendにsparse Flash Attentionを実装する代わりに、model graph側で先にK/Vを小さくしてしまう

というもの。

これには当時のupstream PR #28213
qwen4exp : gather-based sparse attention for QSA decode

を主な参考にした。

このPRも、QSA indexerでtop-kを選んでいるのに、その結果をfull KV上のmaskへ戻してしまうためattention costがcontext長と一緒に増える、という同じ問題を扱っていた。

r3では既存のtop-k結果を使って、

top-k（約2051 cells）
  ↓
get_rows(K)
  ↓
get_rows(V)
  ↓
get_rows(mask)
  ↓
compact K/V/mask
  ↓
約2051件に対して通常のFlash Attention

とした。

対象はまず、

qwen4exp
single-token decode
single-stream

に限定。

prefill / batched pathは従来のままfallbackさせ、MTPも今回のCOMMON-005には含めない。

同一binaryでON/OFFテストできるように、

QWEN4EXP_QSA_GATHER=0
QWEN4EXP_QSA_GATHER=1

で切り替えられるようにした。

これを管理ID COMMON-005 とした。

実装commitは、

d03c91b6342b099457de3508c5d67533a9a5f0ee
qwen4exp: gather selected QSA KV rows for decode

ROCmでの結果

COMMON-004は常にONにして、COMMON-005のGatherだけを同じbinaryでON/OFFした。
64kはOFF/ON 1組、128kと256kはOFF=A、ON=BとしてABBA計測した。
r2 Vulkanは第8回で計測した参考値。

Context	r2 Vulkan TG	ROCm Gather OFF TG	ROCm Gather ON TG	COMMON-005効果
64k	23.55	20.55	22.60	+10.0%
128k	21.65※	16.43	20.41	+24.3%
256k	18.84	11.68	17.37	+48.7%

※r2 128k TGは第8回本文には未掲載。

256kでは11.68→17.37 tok/sで約+49%。
長いほど効果が大きくなっている。

PP側は、

Context	Gather OFF PP	Gather ON PP
64k	359.37	357.27
128k	262.44	262.50
256k	167.58	167.99

で、ほぼ変わっていない。

decode latencyで見ると、

Context	Gather OFF	Gather ON
64k	48.66 ms/token	44.25 ms/token
128k	60.88 ms/token	49.00 ms/token
256k	85.62 ms/token	57.57 ms/token

128k→256kで増えるlatencyは、

OFF: +24.73 ms/token
ON :  +8.57 ms/token

となり、長文化による追加latencyを約65%削減できた。

Radeon GPU Profilerでも確認

COMMON-005前は、

flash_attn_tile<256,256,1,4,false>

128k: 約1.452 ms
256k: 約2.939 ms

だった。

COMMON-005後はcompact K/Vを使う別のtileになり、

flash_attn_tile<256,256,2,1,false>

128k: 約41.93 us
256k: 約42.15 us

となった。

128k代表値で約34倍級の短縮。

しかも128k→256kでほぼ増えていない。
ただし、

k_get_rows_float
128k: 約447.92 us
256k: 約907.13 us

はまだ約2倍のまま残っている。
これはCOMMON-005では触っていない、block scoreを全cellへ展開する処理なので予定通り。これは次のCOMMON-006候補になる。

Vulkanでも試す

念のためVulkan版もCOMMON-005を実装した状態でビルドして計測してみた。
Vulkanはもともとbackend側にsparse Flash Attentionがあるので、ROCmほど大きな改善は期待していなかった。
結果はこうなった。

Context	r2 Vulkan TG	Gather OFF TG	Gather ON TG	変化
64k	23.55	25.95	25.42	-2.0%
128k	21.65※	24.61	24.12	-2.0%
256k	18.84	21.135	21.69	+2.6%

64k/128kでは既存のVulkan sparse FAの方が約2%速い。
256kでは逆にmodel-side gatherが約2.6%速い。
どちらにしても差は数%で、ROCmのような「長くなるほど明確にCOMMON-005が効く」という結果ではなかった。

COMMON-001 → 004 → 005 の推移

今回のr3をまとめるとこうなる。

Patch	Backend	64k PP	64k TG	128k PP	128k TG	256k PP	256k TG
COMMON-001	Vulkan	267.90	17.20	159.56	12.00	103.75	7.30
COMMON-004	Vulkan	262.89	25.95	156.47	24.00	98.63	21.03
COMMON-005 ON	Vulkan	262.44	25.42	156.16	24.12	97.98	21.69
COMMON-001	ROCm	357.93	14.86	259.41	9.91	167.35	5.74
COMMON-004	ROCm	356.55	20.62	261.67	16.44	166.18	11.60
COMMON-005 ON	ROCm	357.27	22.60	262.50	20.41	167.99	17.37

PPは今回ほぼ改善対象にしていないので、大きくは変わっていない。
TGは段階的に改善できた。

特にROCm 256kは、

COMMON-001:  5.74 tok/s
COMMON-004: 11.60 tok/s
COMMON-005: 17.37 tok/s

と、最初の約3倍になった。

VulkanはCOMMON-004で大部分の問題が解消し、COMMON-005は既存sparse FAとほぼ互角。
同じmodel graphの改造でもbackendによって効果が全然違う、というのも面白い結果だった。

block-domain top-k（COMMON-006）実装は見送り

ROCmのプロファイリング結果から、もう1つ高速化案があった。

COMMON-005後も、

k_get_rows_float
128k: 約447.92 us
256k: 約907.13 us

とcontext長にほぼ比例して増えている。

原因は先ほどの、

現在
block score
  ↓
全KV cellへ展開
  ↓
cell-level top-k

という処理。

QSA indexerは最初からblock単位のscoreを持っているので、

候補
block score
  ↓
block-domainでtop-k
  ↓
選ばれたblockだけcell indexへ展開

にすれば、全context分のcellへ一度展開する処理を避けられそう。

これを COMMON-006: block-domain top-k 候補とした。

ただし、COMMON-004/005を実装した後で2026-10-02時点の最新upstreamを調査すると、upstream PR #29751 llama: fix qwen4exp でQSA/k-pool部分がかなり作り直されていた。

新しいupstreamでは、

pool / block score
  ↓
pool-domain top-k
  ↓
選択したpoolだけcell indexへ展開
  ↓
tail indexを追加

という、COMMON-006で考えていたのと実質同じ方向になっている。

旧b11247上にもう一度独自COMMON-006を作って、correctness確認と128k/256k A/Bを全部やった直後にupstreamを入れ替えるのはさすがに無駄が多い。

そのためCOMMON-006は実装せず、

新upstreamで本当に同じボトルネックが消えているかを先に測る

方針に変更した。

前回のupstreamベースへの切り替えからまだ1週間も経っていないのに、早速ベースを入れ替えることになった。

ベース入れ替えを前提にした設計にしておいて良かったとは言えるのだけど、r2の機能をr3に全部移植し終わる前に入れ替えになったのはちょっと残念な気もする。

r3は、

COMMON-001 PLE16 loader
COMMON-004 incremental pooled-key cache
COMMON-005 selected QSA K/V gather

までを実装、検証完了として残して、次は別ブランチでclean upstream baselineからやり直すことにした。

まとめ、今後の予定

今回は、前回見つけた「長コンテキストになるほどTGが妙に遅い」をprofilerで順番に追った。

最初のVulkan profileでは、r3が毎decodeでpooled summaryをfull-contextから作り直していることが分かり、COMMON-004でincremental cache化。
これでVulkan 256k TGは、

7.30 -> 21.03 tok/s

まで改善した。

ROCmはまだコンテキスト依存の低下が残ったので、今度はRadeon GPU Profilerで128kと256kを比較。
Flash Attentionと k_get_rows_float の2つがcontext長にほぼ比例して増えていることが分かった。
そのうちFlash Attention側をCOMMON-005でcompact K/V化した結果、ROCm 256k TGは、

5.74
  ↓ COMMON-004
11.60
  ↓ COMMON-005
17.37 tok/s

まで改善。
RGP上でもFlash Attentionは128k/256kとも約42 usまで縮み、コンテキスト長による増加をほぼ解消できた。

残ったblock-score展開はCOMMON-006として実装する予定だったけど、最新upstreamに同じ方向の実装がすでに入っていたので中止。

次回はベースの本家(upstream) llama.cppを最新版へ入れ替えて、まずclean baselineを取り直してから高速化の続きをやる予定。

実は実装、計測してから記事を書くまでに数日あいてしまったのだけど、その間に今回実装したCOMMON-004,005についても近い考え方がupstreamに既にマージされているかもしれない。
色々と独自実装するよりupstreamの更新を待った方が良いような気もしないでもないのだけど、遅い原因を調べてソースを追って修正すること自体がだんだん楽しくなってきている。

そろそろ本来の目的である日本語テキスト処理の検証も進めたいのだけど、高速化のほうが結果がハッキリ数字に出るから、ついハマってしまう。

付録

長文要約テストのコマンド例

.\llama-cli.exe `
    -m "(モデルのパス)" `
    -c (コンテキスト長) `
    -ngl 999 `
    -ncmoe 0 `
    -t 4 `
    -tb 4 `
    -b 2048 `
    -ub 1024 `
    -fa 1 `
    -ctk f16 `
    -ctv f16 `
    --temp 0.2 `
    --top-k 20 `
    --top-p 0.8 `
    --min-p 0.05 `
    --jinja `
    --single-turn `
    --reasoning off `
    -f 入力ファイルのパス `
    -n 1024 `
    -lv 4

テストデータについて

今回も前回までと同じ長文入力を使用した。

実入力token数は、

context	実入力
64k	61,789 tokens
128k	126,253 tokens
256k	255,181 tokens

256kでは要求context 262,144の約97.3%まで入力している。

前回までの記事と同じものなので詳細は省略する。

詳細な検証環境

Hardware

Device: NucBox_EVO-X2
CPU / APU: AMD Ryzen AI MAX+ 395 with Radeon 8060S
CPU clock displayed by Windows: 3.00 GHz
Installed memory: 128 GB
UMA Frame Buffer: 96 GB
Windows CPU側に見えるRAM: 約32 GB
GPU: AMD Radeon 8060S Graphics
GPU arch: gfx1151

OS

OS: Windows 11 Pro
Version: 25H2
OS Build: 26200.9168
Installed: 2025-10-24
Windows Feature Experience Pack: 1000.26100.344.0

ベースのllama.cpp

Upstream repository: ggml-org/llama.cpp
Build: b11247
Commit: 0bc845d356f437d5ce4fe975c36428f7522829cb
Commit title: vulkan : reuse descriptor sets when bindings are constant (#29280)
r3 branch: r3/upstream-first

Vulkan:

Compiler: LLVM / Clang 20.1.8
Vulkan SDK: 1.4.357.0

ROCm:

ROCm: 10.0.0 / TheRock
Compiler: AMD Clang 23.0.0
Host toolchain: Visual Studio 2022 Build Tools / MSVC 14.44
GPU target: gfx1151

COMMON-004

6559dd272fd0d5f553823e8851c78b9edc2a5016
qwen4exp: cache pooled QSA keys across decode steps

COMMON-005

d03c91b6342b099457de3508c5d67533a9a5f0ee
qwen4exp: gather selected QSA KV rows for decode

使用したモデル

Repository: unsloth/Qwen3.8-Flash-Next-GGUF
URL: https://huggingface.co/unsloth/Qwen3.8-Flash-Next-GGUF
Quantization: UD-IQ3_XXS
File size reported by llama.cpp: 76.32 GiB
BPW: 3.71
Model params: 176.94B

上記モデルを、LaurentZuijdwijkのStrix Halo専用fork内の

gguf-py\gguf\scripts\gguf_split_ple_heads.py

で変換したPLE16モデルを使用した。