3090x3で 314 万コンテキスト; KVキャッシュをホストRAMに置けるアーキテクチャの偉大さ

長い文脈が載らない理由は、大体の場合モデルではなく KV キャッシュのほうにある。 VRAM ギリギリのサイズの gguf をダウンロードして、痛い目にあった人は多分いる。

KV キャッシュは厄介なものである。多くの人は wikitext-2 @512 とかいう当てにならない指標と、 real world use からかけ離れたクソみたいなベンチマークを引っ張り出して、8bit がロスレスだと勘違いして、痛い目を見る。 KV キャッシュのキモい所はそこだけじゃない。モデルによってコンテキスト毎に消費する KV キャッシュサイズは大きく違うし、 なんなら KV キャッシュの量子化の敏感さも違う。

これは VRAMのない人にとっては問題だ。KV キャッシュ事情には嫌でも触れなければいけないのだから。だが、 Qwen Sparse Attention（以降 QSA ）はそれをある程度は解決した。

私の環境は RTX 3090 が 3 枚、VRAM 合計 70.68 GiB。Qwen3.8-Flash-Next の EXL3 4.05bpw 量子化を 載せると、重み・リカレント状態・ワークスペースまで込みで 64.35 GiB を占める。キャッシュに残るのは 6.33 GiB しかない。fp16 の KV は 1 トークンあたり 27.75 KiB なので、上限は 約 234,000 トークン。

私はキャッシュの量子化はしない宗教に入っている。量子化すれば載るコンテキストは増えるが、品質を捨てることになる。 重みを CPU にオフロードしてもコンテキストは増えるが、速度を捨てることになる。これらどちらかを取らねばならないのが これまでの Qwen のアーキテクチャだったが、この qwen4exp アーキテクチャは違う。

三つめの道があるのだ。Qwen Sparse Attention なら、品質も速度もほとんど落とさずにコンテキストを伸ばせる。 キャッシュのほとんどは今ホスト RAM にあり、上限は 3,145,728 トークン、100 万トークンの文脈で デコードは 50.7 tok/s 出る。

環境

RTX 3090を 3 枚、RAM は DDR 4 3200 128GB、OS は Arch Linux。

普段は vLLM を使用していたが、今回は EXL3 を使用した。理由は低 bit での品質が高いからである。

なぜ

QSA には indexer という、履歴のどこを読むかを選ぶ仕組みがあり、そこには固定の予算がある。 このモデルの text_config では indexer_budget: 2048、indexer_compress_ratio: 4。 文脈が 4 * block_topk + 3 = 2051 トークンを超えると、QSA は履歴を全部読むのをやめて、 indexer が選んだところだけを読むようになる。つまり 1 トークン生成するのに読む K/V はきっかり 2048 トークン分で、それ以上にはならない。 後ろに 3,000 トークンあろうが 300 万トークンあろうが同じである。

話はこれでほぼ終わっている。デコード 1 ステップで RAM から読むバイト数を出すと:

K_pad(≈2048) × head_dim(256) × 2 B × 2 (K と V) × num_kv_heads(2)  = 4 MiB / 層
× フルアテンション 12 層                                            = 48 MiB / step

1 トークンあたり 48 MiB で、そこから先は増えない。この機械で実測した RAM から GPU への転送速度は PCIe 4.0 x16 で 26.50 GB/s なので、転送は 1.9 ms/token。デコードは元々 15〜20 ms 台に いる。入る余地がある。

他の多くのモデルではこうはいかない。読み出しが文脈長に比例して伸びるので、キャッシュのオフロードは 考えるだけ損する。Sparse アテンションは演算を減らすために設計されたものだが、ついでに キャッシュがRAMであることを許してくれる。

カーネルを書き換えずに済む

二つめの幸運は、gather カーネルが読む場所をどう決めているかにある。_qsa_sparse_split_kernel から:

k_ptrs = k_cache + ((tok[None, :] * n_kv_heads + kv_head) * head_dim + offs_d[:, None])
v_ptrs = v_cache + ((tok[:, None] * n_kv_heads + kv_head) * head_dim + offs_d[None, :])

先頭のアドレスに、そこからのズレを足しているだけである。だから k_cache が VRAM ではなく、 GPU から直接読めるようにマップしたホスト RAM を指していても、カーネルは無改造で動く。実際に動いた。 MoE の CPU オフロード用に pinned_cuda_view() が既にあったからである。

実装はキャッシュのクラス 1 つで済む。alloc() の時点で:

1. K/V を置くメモリを匿名 mmap で確保する（ページ境界が保証され、ゼロ埋め済みなので 24 GiB の memset が要らない）
2. 対象デバイスをカレントにして cudaHostRegister(ptr, nbytes, PORTABLE | MAPPED)
3. 返ってきたデバイスポインタを CUDA テンソルとして包む

下流の sparse_attend も get_kv も ext.paged_kv_cache_update も、気づかないまま通る。 CUDA グラフのキャプチャも生き残った。EXL3_BC_ATTN_TRACE=1 で見ると 12 層すべてが built で、 グラフ化に失敗して通常実行に落ちた層はない。

ここで zero-copy（GPU が RAM を直接読む方式）を選んだのは、デコードでは先読みができないからである。 層 i+1 がどのトークンを読むかは層 i の出力が出るまで決まらないので、層をまたいで先に転送を 始めておくことができない。かといって明示的に転送を指示する方式にすると、どこを読むかの情報を CPU 側に読み戻すための同期が層ごとに、1 トークンあたり 12 回入る。しかも CUDA グラフのキャプチャが 壊れる。zero-copy なら、転送の待ち時間はカーネルが並列に走らせている他の処理の裏に隠れる。

VRAM に残す必要があるもの

キャッシュのうち一箇所だけは読む量が一定ではない。ここまでホストに出したら終わる。

indexer は毎ステップ全ブロックをスコアリングする。そのスコアリングに使う pooled key （ブロックごとにキーをまとめたもの）だけは、読む量が文脈長に比例して増える。 100 万トークンでは 250,000 ブロック × 128 × 2 B = 64 MiB/層、12 層で 732 MiB を毎トークン全走査する。 VRAM の 835 GB/s なら 0.92 ms。PCIe の 26 GB/s では 28 ms になり、デコード 1 ステップ全体より重い。

というわけで、K と V と indexer の raw key はホスト RAM へ、pooled key は何があっても VRAM に残す。 最終的な配分は VRAM 0.75 KiB/token・RAM 27.00 KiB/token。全部 VRAM に置く 27.75 KiB/token から こう分かれた。

ちなみにこの欠点

「読み出しが一定」というのはトークン 1 個についての話である。プリフィルでは数千トークンを同時に 処理していて、そのトークンがそれぞれ勝手に自分の 2048 トークンを選ぶ。全部を足すと、結局ほぼ 履歴全体を読むことになる。素直にやると 4096 トークンのチャンク 1 つで 204 GiB を PCIe 越しに 読むことになる。100 万トークンのプリフィルは転送だけで 30 分かかる。

対処はループの順番を入れ替えることである。トークンごとに必要なところを引っぱってくるのをやめ、 履歴を一定サイズずつに区切って、その区間を VRAM 上の小さな作業用バッファへ 1 回だけコピーする。 全トークンをその区間のぶんだけまとめて計算し、softmax の途中結果をマージしながら次の区間へ進む。 外側のループが RAM から VRAM へのコピー、内側のループがトークンである。逆にすると、トークンの まとまりごとに同じ区間を何度も転送し直すことになる。

どれくらい効いたかは、EXL3_QSA_KVO_STATS=1 を付けると出る集計でわかる。60,000 トークンのプロンプト、-cs 262144:

103 分の 1 である。そして 1 トークンあたりの選択数が前提どおり、文脈長によって変わっていないことも 確認できる。

ここで一つ想定外があった。計画では作業用バッファを既定 2 GiB とするつもりだったが、 この機械はモデルが VRAM の 91% を占めているため、128 MiB でも読み込みが落ち、64 MiB でも プリフィル中に OOM した。採用値は 32 MiB。 これはほぼ無害だとわかった。バッファが小さいと 履歴を区切る回数が増えるだけで、PCIe を通るバイト数は変わらない。増えるのは、どのトークンを 読むかのリストをカーネルが読み直す回数だけである。

結果

デコード（tok/s）:

ベースラインの列が途中で切れているのは、そこから先へ行けないからである。23.4 万トークン付近で VRAM が尽きる。オフロード側の列がこの作業の全部で、2,048 から 1,048,320 トークンへ、 文脈が 512 倍になってデコードの低下は 6.2% に収まっている。残っている低下は pooled key を 毎回すべて読んでいるぶんで、文脈長に比例して増える。事前の見積りどおりである。

固定費は短文脈でのベースライン比 約 13% の低下。これが PCIe 転送のぶんで、最初に払って 以降は増えない。

メモリ（GiB）:

314 万は、n-gram テーブルを RAM に載せるのをやめてディスクから読むようにし、RAM を空けて到達した。 ここから先の律速は VRAM ではなく RAM である。GPU 側にはまだ 3.4 GiB 余っていて、 0.75 KiB/token なら 450 万トークン分ある。

規模感として、プリフィルは 32k で 1,640 tok/s、1M で 592 tok/s。つまり 100 万トークンの文脈を ゼロから作るのに 約 29.5 分かかる。作ってしまえば 50 tok/s でデコードできる。

ちなみに

• 262,144 を超えるなら YaRN が要る（config.json の rope_scaling）。ここの 1M / 2M / 3M の 数字はメモリと速度の測定であって、その長さでの出力品質については何も言っていない。
• -kvo は CPU ページキャッシュ（-ccs）と併用できない。 計画では「K/V が既にホストにあるなら、 退避するのは VRAM に残っている部分だけでよい」と考えていたが、これは成立しない。K/V を含まない 退避データから復元すると、前に使っていた別のシーケンスの K/V がそのまま残ってしまう。 正しく共存させるには K/V も含めて丸ごと退避して、RAM から RAM へコピーする必要がある。 めんどいのでやっていない。

使い方

qsa-kv-offload.patch を入手して、

cd exllamav3
git apply qsa-kv-offload.patch

を行う。

8 ファイルを変更し 2 ファイルを追加する。

使い方は -kvo を付けるだけ。