0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

4bitのKVキャッシュをバイト列で運ぶ:AMDとLMCacheのGPU・CPU間転送

0
Posted at

長い会話を何度も処理する推論サーバーでは、モデルの重みがGPUに収まっても、過去の計算結果を保存するKVキャッシュが足りなくなる。
AMDは2026年8月28日、BF16と4bitが混在するKVキャッシュを、形式を保ったままGPUとCPUメモリの間で移す構成を公開した。[1]
仕組みの中心は、転送側が量子化データを数値へ戻さず、必要なバイト列を正しい位置へ運ぶことにある。

この記事では、キャッシュの保存形式、コピー範囲、再利用時の検証を順に追う。公式の性能値に加えて、公開ハーネスの設定を確認し、バイト列の保存・復元を説明する小さなPythonコードを実行した。MI355Xでの推論やGPU転送ベンチマークは実施していない。

再計算を減らすには、保存量と転送量の両方が効く

自己回帰型LLMでは、過去のトークンから計算したKeyとValueを保存し、次のトークンを生成するときに再利用する。会話が長くなり、同時に扱う要求も増えると、その保存量が増える。

さらに、次のターンでも同じ会話の先頭部分を処理するなら、保存済みのKVを使うことで入力部分の再計算を減らせる。LMCacheは、このKVをCPUメモリやほかの保存先にも置き、必要になったときに戻す機能を提供する。[2]

ここで、キャッシュミスを一種類にまとめると判断を誤る。

必要なKVの場所 次に必要な処理
GPUのHBMにある GPU上のブロックを再利用する
GPUにはないがCPUのDRAMにある GPUへ転送してから再利用する
再利用できる保存先にない 不足する範囲をprefillで再計算する

HBMはGPUに近い高速メモリ、DRAMはここではホストCPU側のメモリを指す。CPUに保存先を増やせば再計算を減らせる可能性があるが、代わりに転送が発生する。

量子化は、一つのトークンに必要な保存量を小さくする。これによって同じメモリに多くのKVを残せるだけでなく、退避したKVを戻すときの転送量も減らせる。効果が出るかどうかは、実際に再利用される文脈がどこまで収まるかに依存する。

4bitという名前だけでは、コピーするサイズは決まらない

今回の構成は、全レイヤーを同じ形式にするものではない。AMDが示すMiniMax-M2.5向けの配置は、先頭と末尾の計4層をBF16に保ち、中間58層を4bitへ量子化する。[1]

レイヤー群 格納用dtype 1トークン・1ヘッドあたりの表記
境界の4層 BF16 512バイト
中間の58層 packed uint8 134バイト

AMD記事のレイアウト表による。[1]

uint8 と書いてあっても、数値が8bit量子化されているという意味ではない。4bitの符号やインデックスを、一つの8bitコンテナーへ二つ詰めることができる。格納用dtypeと、値を表すビット数は分けて読む必要がある。

また、パックされた領域には、復元に必要なスケールなどの補助情報が含まれる場合がある。コピー側が「4bitだから要素数を2で割ればよい」と決め打ちすると、その補助情報やアラインメントを落としかねない。

上の表のサイズを使い、各層のトークン数とヘッド数が共通と仮定して合計すると、全層BF16に対する単純な容量比は次のようになる。

\begin{aligned}
S&=\frac{62\times512}
{4\times512+58\times134}\\
&\approx3.23.
\end{aligned}

これは表から計算したレイアウト上の比率で、実機で測定した速度向上率ではない。「16bitから4bitへ変えたので、サーバー全体のメモリが4分の1になる」という計算にはならないことが分かる。

モデルの重み、ワークスペース、管理情報、確保単位の切り上げは、この比率の外にある。実際のGPUメモリ消費や保持できるトークン数は、起動ログとアロケーターの値で確認する。

転送側は、量子化の解釈を引き受けない

AMDの説明では、コネクターはBF16とpacked uint8を別の転送グループとして扱い、それぞれのストライドに合わせてコピーする。パックされたスロットの中身は、数値として解釈せずに運ぶ。[1]

概念的な経路は次のようになる。

ただし、「中身を解釈しない」と「メタデータが不要」は違う。一般にこの方式を実装するには、少なくとも次の情報が要る。

  • どのレイヤー、どのKVブロックへ戻すか。
  • 有効なバイト数と、コピー元・コピー先の配置。
  • トークン方向やヘッド方向のストライド。
  • パック形式を読む側が期待する形式識別子と版。
  • データ生成側と利用側が共有するモデルや分割の条件。

例えば、同じ連続領域に見えても、一方が「Kの全要素、その後にV」、もう一方が「トークンごとにKとVを交互に配置」という形式なら、単なる要素数の一致では接続できない。

保存側は完全なスロットを保持し、利用側はその形式を理解して読む。この責任分担なら、転送のたびに量子化形式を解釈し直す処理を避けられる。ただし、読み手の互換性を確認する工程は残る。

L2へ保存するときに圧縮する経路とも区別する

2026年9月10日に確認したLMCache公式文書には、L2保存先ごとにFP8やTurboQuantの変換処理を有効にする serde がある。こちらは保存前に変換し、読み込み時に復元する経路だ。[3]

今回注目するのは、GPU上ですでに量子化された混在形式を、そのまま転送対象にする設計である。両者とも保存容量を減らすが、どこで変換し、どの形式をGPUへ返すかが異なる。

導入時は「LMCacheが量子化をサポートする」という説明だけで済ませず、使うコネクター、GPU側のKV形式、保存先の変換設定を一組として確認する。一般のインストール手順と、特定のフォークを使う実験構成を混ぜないことが大切だ。

バイトが同じなら、何が保証されるのか

量子化前のKVを $X$、量子化とパックをまとめた処理を $Q$、保存と復元の往復を $T$ とする。

転送がバイト単位で正しいという主張は、

T(Q(X))=Q(X)

である。量子化結果を失わず戻せた、という意味だ。

ここから、量子化前の数値まで完全に戻ったとは言えない。復号処理を $D$ とすれば、

D(T(Q(X)))=D(Q(X))

にはなるが、右辺が $X$ に等しいとは限らない。量子化による近似と、転送による破損は別々に検証する必要がある。

実装テストも分けると原因を探しやすい。

検証 比較対象 主に検出したい問題
バイト往復 保存直前と復元直後の全バイト コピー範囲や配置の誤り
隣接領域 別レイヤーや未使用領域の内容 上書き、サイズ超過
推論の経路比較 同じ量子化形式でHBM経由とDRAM経由 復元後の参照や同期の誤り
量子化の品質比較 BF16と量子化KVでの出力品質 近似による精度変化

同じ文字列の回答が出たという試験だけでは、使われなかったKV領域の破損を見逃す可能性がある。反対に、生成にランダム性があれば、出力文字列の違いだけで転送失敗とも断定できない。

そのため、バイト一致、経路を通った証拠、回答品質という順に検証対象を切り分ける。

小さな保存・復元コードで、責任範囲を確かめる

次のコードは、混在形式の領域をバイト数に従って退避する、本稿独自の説明用例である。TurboQuantの実際の134バイト形式や、LMCacheのGPUコネクターを再実装したものではない。

二つのBF16値に対応するバイト列と、架空の4bitコード8個にfloat32のスケールを付けた領域を用意する。保存・復元側は、その数値的な意味を使わずにコピーする。

from struct import pack

codes = [1, 14, 2, 13, 3, 12, 4, 11]
packed = bytes(
    codes[i] | (codes[i + 1] << 4)
    for i in range(0, len(codes), 2)
)

payloads = {
    "bf16": bytes.fromhex("80 3f 00 40"),
    "packed4": packed + pack("<f", 0.125),
}
slots = {"bf16": (0, 4), "packed4": (16, 8)}
arena = bytearray([0xA5] * 32)

for name, (offset, size) in slots.items():
    assert len(payloads[name]) == size
    arena[offset:offset + size] = payloads[name]

saved = {
    name: bytes(arena[offset:offset + size])
    for name, (offset, size) in slots.items()
}

# 利用領域を消して、退避したバイト列から戻す。
for offset, size in slots.values():
    arena[offset:offset + size] = bytes(size)
for name, (offset, size) in slots.items():
    arena[offset:offset + size] = saved[name]

for name, (offset, size) in slots.items():
    assert bytes(arena[offset:offset + size]) == payloads[name]
assert arena[4:16] == bytes([0xA5] * 12)

print("bf16:", saved["bf16"].hex())
print("packed4:", saved["packed4"].hex())
print("round-trip: exact")
print("payload bytes:", sum(map(len, saved.values())))
print("arena bytes:", len(arena))

Python標準ライブラリで実行した結果は次のとおりだった。

bf16: 803f0040
packed4: e1d2c3b40000003e
round-trip: exact
payload bytes: 12
arena bytes: 32

packed4 の後半に付けたスケールも、バイト列の一部として保存されている。隙間の監視用パターン 0xA5 も変わっていない。

また、有効データは12バイトでも、確保した領域は32バイトある。圧縮後のペイロード量と、アロケーターが確保する量は同じとは限らないことも、この小例で確認できる。

実際のGPU実装では、ここに非同期コピーの完了確認、バッファーの寿命管理、参照中のブロックの保護、複数要求の競合処理が加わる。このコードが確かめたのはCPUメモリ上のバイト往復だけである。

公表された性能値を、転送経路の実験として読む

AMDの主比較は、MiniMax-M2.5を2枚のMI355Xで動かし、同じGPUメモリ設定と128GBのCPU保存層を使ったものだ。約100kトークンの文脈を含む会話負荷を20分間与え、次の値を報告している。[1]

公表値 BF16 KV TurboQuant KV
20分間の完了要求数 189 385
p95 TTFT 140秒 53秒

TTFTは最初の出力トークンまでの時間である。完了数の比は約2.04倍、p95 TTFTの比は約2.64倍になる。

ただし、完了要求数と、遅延目標を満たした要求数は同じ集計ではない。この表だけから「385件すべてが60秒以内に応答した」とは読めない。p95の53秒も、すべての要求の待ち時間を表す値ではない。

そして、量子化が常に最速とも限らない。AMDは、比較対象のFP8も十分に保存できるメモリ条件では、FP8側が勝つ結果を示している。[1]

これは自然なトレードオフである。小さく保存する効果が薄くなると、パックされた値を読み出す処理やカーネルの効率が、性能を左右しやすくなる。

転送した時間と、利用者が余計に待った時間

一般的な見積もりとして、読み戻すデータ量を $B$、実効転送帯域を $\beta$、固定の準備時間を $t_0$ と置けば、

t_{\mathrm{copy}}\approx t_0+\frac{B}{\beta}

となる。小さなデータでは $t_0$ が、大きなデータでは転送量が支配しやすい。

ただし、このコピーを計算と重ねられるなら、計測したコピー時間のすべてが応答遅延へ足されるわけではない。一方、Attentionが必要とするブロックの到着を待っていれば、その待ちが性能を決める。

評価では、総転送バイト数、転送イベントの時間、待ちが実際に発生した区間をそれぞれ取る。コピー時間だけからTTFTの改善幅を計算しない。

公開ハーネスには、再現前に読むべき設定がある

本稿では公開された run_equal_memory.sh も確認した。参照したrawファイルの版は 5e5b69bf58a41a327d2db388e4b9dc86be7e76e0 である。[4]

両構成で gpu-memory-utilization=0.475 とCPU保存層の128GBを揃える一方、TurboQuant側にはdtypeに加え、Attentionバックエンド、ブロックサイズ、CUDA Graph関連設定、専用の環境変数を渡している。

従って、この比較は実用的な構成同士の比較として読むのが適切だ。KVのビット数だけを変えた効果を切り出すには、追加設定を揃えた対照実験も必要になる。

また、vLLMの gpu-memory-utilization はmodel executor全体のGPUメモリ割合であり、KV専用領域の割合ではない。同じ設定値でも、ほかの確保量が変わればKVに残る容量も変わり得る。公式文書には、KV容量を直接指定する別の kv-cache-memory-bytes もある。[5]

ハーネスは負荷クライアントをコミット指定なしで取得し、start-users=4max-users=32 を渡す。現在取得されるクライアントの動作まで固定されているとは限らないため、再実行時はクライアントの版と実際の同時実行数の推移も記録したい。[4]

これらは、公表値が間違っているという判定ではない。再現する条件と、数値から切り出せる因果関係の範囲を明確にするための確認である。

手元の環境で試す順序

最初に確かめたいのは、GPUから追い出されたKVがどれだけ再利用され、再prefillにどれだけ時間を使っているかだ。短い会話や少数の同時要求で、必要なKVが十分にGPUへ残るなら、退避や量子化を足す効果は小さい可能性がある。

次に、同じプロンプトを使って、HBM上での再利用と、CPUへ退避してからの再利用を分けて試す。二回目が速かっただけではCPU経由を検証したことにならない。HBMから実際に外れ、CPU側に残り、読み戻しが発生したことをログで確認する。

その後、量子化形式を変えて精度と性能を測る。同時実行数、文脈長、メモリ予算を変えながら、完了数、TTFT分布、再計算量、転送量を並べる。メモリに収まる条件と収まらない条件の両方を含めると、どこで方式の優劣が変わるかを判断しやすい。

本稿は、AMDのGPUベンチマークやフォークの実機動作を再現したものではない。特定の構成で得た速度や回答品質を、別のモデル・GPU・量子化形式へそのまま一般化することもできない。

それでも設計上の要点は明確だ。混在するKVを運ぶには、バイト列そのものだけでなく、その長さと配置、読み手との互換性を保つ必要がある。容量削減を推論の改善につなげるには、その転送によってどの再計算と待ち時間が減ったかまで測ることが欠かせない。

参考リンク

[1]AMD — 4-bit KV Caching in LMCache: Offloading Quantized KV Beyond HBM for Context-Heavy Agents on AMD MI355X(2026年8月28日)

[2]LMCache — Overview / KV cache management(2026年9月10日参照)

[3]LMCache — KV Cache Compression:L2 adapterのserde

[4]AMD公開ベンチマークハーネス — run_equal_memory.sh(確認したrawファイルの版に固定)

[5]vLLM — Engine Arguments:gpu-memory-utilization / kv-cache-memory-bytes

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?