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?

@plasmon

KVキャッシュを光で読んだらメモリ転送が16分の1になった

0
Posted at

KVキャッシュを光で読んだらメモリ転送が16分の1になった

LLMの長文推論で何がボトルネックか。計算ではない。メモリ帯域だ。

Transformerのデコードステップでは、1トークン生成するたびにKVキャッシュ全体をスキャンする。コンテキスト長nに対してO(n)のメモリ読み出しが毎ステップ発生する。GPUの演算器がどれだけ速くても、このメモリ帯域のO(n)は消えない。

2026年3月のArXiv論文（arXiv:2603.21576, Park & Park）が提案したPRISMは、このKVキャッシュのブロック選択をフォトニクス（光回路）で実行し、メモリアクセスをO(1)にした。

結果: 64Kトークンでメモリトラフィック16分の1。ブロック選択に限ればエネルギー効率は4桁改善。精度は100%維持。

なぜメモリ帯域がLLM推論のボトルネックなのか

デコードの構造的問題

# Transformerデコード1ステップの処理
def decode_one_token(query, kv_cache):
    # query: 現在のトークン (1個)
    # kv_cache: 過去のすべてのトークン (n個)

    # Step 1: queryとKVキャッシュ全体の類似度計算
    scores = query @ kv_cache.keys.T  # O(n) のメモリ読み出し

    # Step 2: Softmax
    weights = softmax(scores)

    # Step 3: 加重和
    output = weights @ kv_cache.values  # O(n) のメモリ読み出し

    return output

# コンテキスト長が伸びるとStep 1とStep 3が線形に重くなる
# 計算量もO(n)だが、実際のボトルネックはメモリ読み出し速度

数字で見る帯域地獄

# 64Kコンテキストでの帯域消費
# Qwen2.5-7B: hidden=3584, layers=28, GQA (KVヘッド=4, head_dim=128)
context_length = 64_000  # 64Kトークン
kv_dim = 4 * 128         # GQA: 4 KVヘッド × 128次元 = 512
num_layers = 28
bytes_per_element = 2    # FP16

# 1デコードステップあたりのKVキャッシュ読み出し量
kv_read_per_step = (
    context_length * kv_dim * 2  # K + V
    * num_layers * bytes_per_element
)
# = 64000 * 512 * 2 * 28 * 2 = 3.67 GB

# RTX 4060の帯域: 272 GB/s
# 1ステップの理論最小時間: 3.67 / 272 = 0.013秒
# → 64Kでも帯域だけなら~75 t/s 出せる計算になるが...

# GQAなしのMHAモデル (LLaMA-1 7B相当: hidden=4096, layers=32, 全ヘッドKV)
# 64000 * 4096 * 2 * 32 * 2 = 33.6 GB → 帯域上限8 t/s
# GQAが帯域問題をどれだけ緩和しているかが分かる

# ただしモデルが大きくなるとGQAでも帯域が壁になる:
# 70Bモデル (GQA, KV dim=1024, 80層): 64000*1024*2*80*2 = 16.8 GB
# RTX 4090でも: 16.8 / 1008 = 0.017秒 = ~60 t/s が上限

GPUの演算能力は年々倍増するが、メモリ帯域の伸びはそれよりはるかに遅い。この乖離がLLM長文推論の根本的な壁になっている。

これはメモリの中で計算すればGPUは要らないのか（別記事予定）で議論するPIM（Processing-in-Memory）が攻めている同じ問題だ。PRISMは別のアプローチ — メモリを読む量自体を減らす。

既存の解決策とその限界

Top-K / Sparse Attention

手法 方法 問題
Top-K Attention 類似度上位K個のKVブロックだけ読む 上位を知るために全体スキャンO(n)が必要
Sliding Window 直近Wトークンだけ参照 長距離情報を完全に捨てる。O(W)だが精度犠牲
H2O (Heavy-Hitter Oracle) 累積Attention Scoreが高いトークンを保持 スコア追跡自体がO(n)

共通の問題: どのKVブロックが重要かを判定する段階でO(n)のメモリスキャンが発生する。 ブロック選択後の計算をいくら効率化しても、選択自体のO(n)は消えない。

フォトニクスアクセラレータの既存研究

光回路でAttention計算を高速化する研究は存在する（Lightening TransformerやPhotonic Tensor Coreなど）。しかし、これらも密行列演算（dense attention）を光で加速するアプローチであり、O(n)のメモリスケーリングは変わらない。

PRISMの着眼点は違う。密行列演算ではなく、ブロック選択という粗い類似度検索に光を使う。

PRISMのアーキテクチャ

核心: ブロック選択を光でO(1)にする

従来のBlock Sparse Attention:
  Query → [電子回路: 全ブロックと類似度計算 O(n)] → Top-K選択 → 精密計算

PRISM:
  Query → [光回路: 全ブロックと同時に類似度計算 O(1)] → Top-K選択 → 精密計算
         ↑
         ここが違う

なぜ光だとO(1)になるのか。光の物理的性質を使う。ただし注意: O(1)は時間計算量であり、ハードウェアリソース（マイクロリングの数）はブロック数に比例してO(n)で必要。時間を買うために面積を払う、という古典的なトレードオフだ。

Broadcast-and-Weight (放送と重み付け)

# 電子回路でのブロック選択
def electronic_block_select(query, block_keys, k):
    """n個のブロックと順番に類似度計算 → O(n)"""
    scores = []
    for block_key in block_keys:  # n回のメモリ読み出し
        score = dot_product(query, block_key)
        scores.append(score)
    return top_k(scores, k)

# 光回路でのブロック選択 (PRISMの原理)
def photonic_block_select(query_light, block_modulators, k):
    """光を全ブロックに同時に送り、並列に類似度計算 → O(1)"""
    # Step 1: queryを光信号に変換
    # Step 2: 光を全マイクロリング共振器に同時にブロードキャスト
    # Step 3: 各共振器がブロックキーで光を変調 (重み付け)
    # Step 4: 全結果を同時に光検出器で読み取り
    # → 全ブロックの類似度が1クロックで確定
    return top_k(all_scores, k)  # O(1)

光は電気と違い、1本の導波路に複数の波長を同時に通せる（波長分割多重: WDM）。各波長が異なるブロックの類似度計算を担当すれば、全ブロックを物理的に並列処理できる。

TFLN マイクロリング共振器

PRISMは薄膜ニオブ酸リチウム（TFLN: Thin-Film Lithium Niobate）のマイクロリング共振器を使う。

PRISMはTFLN（薄膜ニオブ酸リチウム）のマイクロリング共振器を使う。シリコンより桁違いに大きい電気光学係数で、67-100 GHz帯域のns応答が可能。数十μm径で数千個をチップに集積できる。

ポイント: ブロック選択は粗い類似度検索であり、4-6ビット精度で十分（これは設計上の選択であり、TFLNの限界ではない）。フル精度のAttention計算は選択されたブロックだけ電子回路（GPU）で実行する。光は精度が低い仕事を超高速でこなし、電子回路は精度が必要な仕事を少量だけやる。

ベンチマーク: 100%精度で16倍削減

Needle-in-a-Haystack テスト

Qwen2.5-7Bを使い、PRISM（k=32ブロック）でNeedle-in-a-Haystackテストを実施。4K/8K/16K/32K/64Kの全コンテキスト長で精度100%。k=32は全ブロックのごく一部だが、ブロック選択の判断が十分正確であることを示している。

メモリトラフィック削減

コンテキスト長 メモリトラフィック削減
4K ~2倍
16K ~6倍
64K 16倍

コンテキストが長いほど削減率が大きい。PRISMの光回路コストは固定（O(1)）なので、nが増えるほど得をする。

具体例: 64Kトークンを128トークンごとのブロックに分割すると500ブロック。Full Attentionでは500ブロック全部読む。PRISMはこの中から上位32ブロック（6.4%）だけを光で選ぶ。残り468ブロックのメモリ転送がゼロになる。

エネルギー効率

ブロック選択操作のみのエネルギー比較（LLM推論全体ではない）:

コンテキスト長 GPU PRISM 差
4K ~1 mJ ~0.1 μJ ~10,000倍
64K ~16 mJ ~0.1 μJ ~160,000倍

GPUはO(n)で比例増加、PRISMはO(1)で固定。

4桁の効率差はブロック選択操作に限った数字だが、理由は単純: GPUはnが増えると比例してエネルギーを消費し、光回路はnに依存しない。コンテキストが長いほどこの差は広がる。ただし、LLM推論全体で見ると効果は薄まる — FFN計算やAttention本体は依然として電子回路が担当するため。メモリトラフィックの16倍削減の方が、実用上のインパクトは大きい。

RTX 4060ユーザーにとっての意味

PRISMは研究段階のフォトニクスチップであり、今日買えるものではない。しかし、この研究が示す方向性はローカルLLMユーザーに直接関係する。

現在の長文推論の壁

RTX 4060 8GB（帯域272 GB/s）でQwen2.5-7B Q4_K_Mを動かした場合:

コンテキスト長 状態
4K 正常動作
8K 体感で遅くなる
16K KVキャッシュがVRAMを圧迫し著しく低速化
32K OOM or スワップ発生

PRISMが示唆する短期的な実践

PRISMそのものは使えないが、同じ原理（ブロック選択の効率化）をソフトウェアで近似するアプローチは存在する。

手法 方法 効果 精度
Quest (2024) ブロック統計量(min/max)でスキップ判定 KVアクセス大幅削減 Full Attentionに近い
RetrievalAttention (2024) KVキャッシュをベクトルDB化、近似最近傍探索 長文で大幅高速化 タスク依存
vLLM PagedAttention ページ単位のKVキャッシュ管理 VRAM有効利用率向上 損失なし

ソフトウェア手法はO(n)のまま定数を小さくする。PRISMはO(1)に根本的に変える。方向性は同じ — 全部読むな、必要なブロックだけ読め。

光コンピューティングとLLM推論の未来

CPOとPRISMの関係

CPOが消費電力を救えない理由（別記事予定）で議論する光接続（CPO: Co-Packaged Optics）は、チップ間のデータ転送を光で高速化する技術だ。PRISMはチップ内の演算自体を光で行う。

CPO:  チップ ←光→ チップ (転送を光で高速化)
PIM:  メモリ内で計算 (転送自体を減らす)
PRISM: 演算を光で実行 (選択をO(1)にして転送を減らす)

アプローチは違うが、すべて同じ敵 — メモリ帯域の壁 — と戦っている

実用化までの距離

PRISMは現在、シミュレーション＋個別コンポーネント実証の段階。フルシステム統合テスト、GPUとの物理インターフェース、量産プロセス、ソフトウェアスタック（ドライバ・コンパイラ）が未解決。楽観的に3-5年で研究プロトタイプ、現実的には5-10年で限定的な商用展開。

関連技術との比較:

  • CPO（光接続）: Intel/TSMCが2026-2027に限定量産予定
  • PIM（メモリ内計算）: Samsung HBM-PIMが2021年に発表、AMDと共同テスト実施済み
  • PRISM型（光演算）: アカデミック段階

実用化は遠い。しかし、LLMのコンテキスト長は年々倍増しており（GPT-4: 8K → GPT-4 Turbo: 128K → Gemini 1.5: 1M → Claude 4.x: 1M）、メモリ帯域の壁はますます深刻になる。電子回路の帯域改善では追いつかない成長速度に対して、光のO(1)スケーリングは原理的な解答を持っている。

PRISMの先に見えるもの

PRISMはシミュレーションと個別コンポーネントの実証段階であり、フルシステムの統合テストはまだない。GPUとの物理インターフェース、量産プロセス、ソフトウェアスタック — 実用化までに超えるべき壁は多い。正直、5年以内に手元に届くとは思えない。

しかし、この論文が重要なのは結果の数字ではなく、問いの立て方だ。既存のフォトニクス研究が「Attention計算を光で速くする」と問うていたのに対し、PRISMは「そもそも何を光でやるべきか」を問い直した。全体を加速するのではなく、ボトルネック（ブロック選択）だけを光に任せる。精度が要らない粗い仕事を光が引き受け、精度が要る仕事は電子回路に残す。

CPO（転送の光化）、PIM（メモリ内計算）、PRISM（選択の光化）。メモリの壁に対して3つの異なるレイヤーから攻撃が始まっている。次にRTX 4060で長文が速くなるとき、その裏では光か、メモリ内演算か、あるいはその両方が動いているかもしれない。

参考文献

  1. "PRISM: Breaking the O(n) Memory Wall in Long-Context LLM Inference via O(1) Photonic Block Selection" (2026) arXiv:2603.21576
  2. "Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference" (2024) arXiv:2406.10774
  3. "RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval" (2024) arXiv:2409.10516
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
Sign upLogin
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?