大規模モデル推論高速化技術の全景

推論高速化の重要性

大規模モデルの推論遅延はユーザー体験に直接影響します。研究によると、応答時間が100ms増えるごとに、ユーザー満足度は3〜5%低下します。リアルタイム対話アプリケーションでは、最初のトークンまでの遅延（TTFT）が200ms以内なら優秀、500ms以内なら許容可能、1秒を超えるとユーザーは明らかに待たされていると感じます。

KV Cache：最も基本的な高速化技術

KV CacheはTransformer推論における中核的な最適化です。KV Cacheがない場合、新しいトークンを生成するたびに、過去のすべてのトークンのKeyとValueを再計算する必要があります。KV Cacheは計算済みのK、Vをキャッシュし、各ステップでは新しいトークンのみを計算します：

# KV Cacheの動作原理
class KVCacheAttention:
    def __init__(self):
        self.k_cache = None
        self.v_cache = None

    def forward(self, q, k, v, use_cache=True):
        if use_cache and self.k_cache is not None:
            # 新しいK、Vをキャッシュに追加
            k = torch.cat([self.k_cache, k], dim=-2)
            v = torch.cat([self.v_cache, v], dim=-2)

        # キャッシュを更新
        self.k_cache = k
        self.v_cache = v

        # アテンションを計算
        attn = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)
        attn = F.softmax(attn, dim=-1)
        return torch.matmul(attn, v)

# 効果：O(n²)をO(n)に削減、推論速度10〜100倍向上

PagedAttention：効率的なKV Cache管理

vLLMが提案したPagedAttentionは、KV Cacheをページ単位で管理し、メモリ断片化の問題を解決します：

# PagedAttentionのメモリ管理
class PagedKVCache:
    def __init__(self, block_size=16, num_blocks=1024):
        self.block_size = block_size
        self.free_blocks = list(range(num_blocks))
        self.block_table = {}  # seq_id -> [block_ids]

    def allocate(self, seq_id, num_tokens):
        num_blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
        blocks = self.free_blocks[:num_blocks_needed]
        self.free_blocks = self.free_blocks[num_blocks_needed:]
        self.block_table[seq_id] = blocks
        return blocks

# 効果：メモリ利用率が20〜40%からほぼ100%に向上

投機的デコーディング

小規模モデルで複数のトークンを「推測」し、大規模モデルで一度に検証することで、2〜3倍のデコード高速化を実現します：

class SpeculativeDecoder:
    def __init__(self, target_model, draft_model):
        self.target_model = target_model  # 大規模モデル
        self.draft_model = draft_model    # 小規模モデル

    def generate(self, prompt, max_tokens=100):
        tokens = tokenize(prompt)

        while len(tokens) < max_tokens:
            # 1. 小規模モデルがK個の候補トークンを生成
            draft_tokens = self.draft_model.generate(tokens, k=5)

            # 2. 大規模モデルがすべての候補を一度に検証
            logits = self.target_model.forward(tokens + draft_tokens)

            # 3. 一致するトークンを受け入れ、不一致を拒否
            accepted = self._verify_and_accept(logits, draft_tokens)
            tokens.extend(accepted)

            if len(accepted) < len(draft_tokens):
                # 不一致の位置から再サンプリング
                correction = self.target_model.sample(logits[len(accepted)])
                tokens.append(correction)

        return detokenize(tokens)

オペレーター融合

複数の小さなオペレーターを1つの大きなオペレーターに融合し、メモリの読み書きを削減します：

# 融合前：5回のメモリ読み書き
# LayerNorm → Linear → Dropout → ReLU → Linear

# 融合後：1回のメモリ読み書き
# FusedMLP (LayerNorm + Linear + GELU + Linear)

# torch.compileを使用して自動融合
import torch

@torch.compile
def fused_mlp(x, w1, w2, b1, b2):
    return torch.nn.functional.linear(
        torch.nn.functional.gelu(
            torch.nn.functional.linear(x, w1, b1)
        ),
        w2, b2
    )

# 効果：遅延が20〜30%削減、メモリ帯域幅が40%節約

技術選定の比較

実践的なアドバイス

1. 測定してから最適化：プロファイラを使用して真のボトルネックを見つける
2. シンプルから始める：KV CacheとFlashAttentionはコストパフォーマンスが最も高い最適化です
3. 量子化を優先：INT8量子化は精度損失がほぼなく、高速化効果が顕著です
4. 組み合わせて使用：複数の技術を組み合わせると効果的です
5. 継続的な監視：最適化後も遅延とスループットを監視し続ける

まとめ

推論高速化はシステムエンジニアリングの取り組みです。KV Cacheと量子化から始め、徐々に高度な技術を導入することをお勧めします。核心的な原則は、計算の削減（量子化、投機的デコーディング）、メモリアクセスの削減（FlashAttention、オペレーター融合）、並列性の向上（連続バッチ処理）です。