PagedAttention — vLLMの仮想メモリ着想によるKVキャッシュ管理

論理ブロック 0 → 物理ブロック 7  （トークン 0, 1, 2, 3）
論理ブロック 1 → 物理ブロック 3  （トークン 4, 5, _, _）

import numpy as np
import matplotlib.pyplot as plt

class BlockManager:
    """KVブロックの物理メモリ管理を行うクラス。
    OSのページフレームアロケータに相当する。"""

    def __init__(self, num_blocks, block_size, num_layers, num_heads, head_dim):
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim

        # フリーリスト
        self.free_blocks = list(range(num_blocks))

        # 参照カウント
        self.ref_count = [0] * num_blocks

        # 物理ブロックストレージ（簡略化: 1層1ヘッドのみ）
        # 実際は全層・全ヘッド分
        self.k_cache = np.zeros((num_blocks, block_size, head_dim))
        self.v_cache = np.zeros((num_blocks, block_size, head_dim))

    def allocate(self):
        """フリーリストから物理ブロックを1つ割り当てる。"""
        if not self.free_blocks:
            raise RuntimeError("No free blocks available")
        block_id = self.free_blocks.pop(0)
        self.ref_count[block_id] = 1
        return block_id

    def free(self, block_id):
        """物理ブロックを解放する（参照カウントを減らす）。"""
        self.ref_count[block_id] -= 1
        if self.ref_count[block_id] == 0:
            self.free_blocks.append(block_id)

    def add_ref(self, block_id):
        """参照カウントを増やす（COW用）。"""
        self.ref_count[block_id] += 1

    def copy_on_write(self, block_id):
        """参照カウント>1のブロックをコピーする。"""
        if self.ref_count[block_id] <= 1:
            return block_id  # コピー不要

        # 新しいブロックを割り当ててコピー
        new_block_id = self.allocate()
        self.k_cache[new_block_id] = self.k_cache[block_id].copy()
        self.v_cache[new_block_id] = self.v_cache[block_id].copy()

        # 元のブロックの参照カウントを減らす
        self.ref_count[block_id] -= 1

        return new_block_id

    @property
    def num_free(self):
        return len(self.free_blocks)

    @property
    def utilization(self):
        return 1.0 - len(self.free_blocks) / self.num_blocks

class Request:
    """1つの推論リクエストを管理するクラス。"""

    def __init__(self, request_id, block_manager, block_size):
        self.request_id = request_id
        self.block_manager = block_manager
        self.block_size = block_size
        self.block_table = []  # 論理ブロック → 物理ブロックのマッピング
        self.num_tokens = 0
        self.slot_in_last_block = 0

    def append_token(self, key, value):
        """新しいトークンのKVをキャッシュに追加する。"""
        bm = self.block_manager

        if self.slot_in_last_block == 0 or self.slot_in_last_block >= self.block_size:
            # 新しいブロックが必要
            if self.block_table and self.slot_in_last_block >= self.block_size:
                pass  # 既存ブロックは満杯
            new_block = bm.allocate()
            self.block_table.append(new_block)
            self.slot_in_last_block = 0

        # 最後のブロックにKVを書き込み
        physical_block = self.block_table[-1]

        # COWチェック
        if bm.ref_count[physical_block] > 1:
            physical_block = bm.copy_on_write(physical_block)
            self.block_table[-1] = physical_block

        bm.k_cache[physical_block, self.slot_in_last_block] = key
        bm.v_cache[physical_block, self.slot_in_last_block] = value

        self.slot_in_last_block += 1
        self.num_tokens += 1

    def get_all_kv(self):
        """全てのKVキャッシュを論理順序で取得する。"""
        bm = self.block_manager
        keys = []
        values = []

        for i, block_id in enumerate(self.block_table):
            if i < len(self.block_table) - 1:
                # 満杯のブロック
                keys.append(bm.k_cache[block_id])
                values.append(bm.v_cache[block_id])
            else:
                # 最後のブロック（部分的に埋まっている）
                keys.append(bm.k_cache[block_id, :self.slot_in_last_block])
                values.append(bm.v_cache[block_id, :self.slot_in_last_block])

        return np.concatenate(keys, axis=0), np.concatenate(values, axis=0)

    def release(self):
        """全ブロックを解放する。"""
        for block_id in self.block_table:
            self.block_manager.free(block_id)
        self.block_table = []
        self.num_tokens = 0


def paged_attention(query, request, scale=None):
    """PagedAttentionによるAttention計算。"""
    keys, values = request.get_all_kv()
    d = query.shape[0]
    if scale is None:
        scale = 1.0 / np.sqrt(d)

    # Attention scores
    scores = keys @ query * scale  # (num_tokens,)

    # Softmax
    scores_max = np.max(scores)
    exp_scores = np.exp(scores - scores_max)
    attn_weights = exp_scores / np.sum(exp_scores)

    # 重み付き和
    output = values.T @ attn_weights  # (d,)
    return output

def simulate_memory_usage(num_requests, max_seq_len, actual_lengths,
                          block_size=16, head_dim=128):
    """従来方式とPagedAttentionのメモリ使用量を比較する。"""
    # 従来方式: 最大系列長分を事前確保
    traditional_per_request = max_seq_len * head_dim * 2 * 2  # KV, FP16
    traditional_total = num_requests * traditional_per_request
    traditional_used = sum(l * head_dim * 2 * 2 for l in actual_lengths)
    traditional_waste = traditional_total - traditional_used

    # PagedAttention: ブロック単位でオンデマンド割当
    paged_blocks_used = 0
    paged_internal_waste = 0
    for length in actual_lengths:
        blocks_needed = (length + block_size - 1) // block_size
        paged_blocks_used += blocks_needed
        waste_in_last = (blocks_needed * block_size - length)
        paged_internal_waste += waste_in_last

    block_bytes = block_size * head_dim * 2 * 2
    paged_total_used = paged_blocks_used * block_bytes
    paged_waste = paged_internal_waste * head_dim * 2 * 2

    return {
        'traditional_total': traditional_total,
        'traditional_used': traditional_used,
        'traditional_waste': traditional_waste,
        'paged_total': paged_total_used,
        'paged_used': traditional_used,
        'paged_waste': paged_waste,
    }


# シミュレーション
np.random.seed(42)
num_requests = 32
max_seq_len = 2048

# 実際の系列長は対数正規分布（短いリクエストが多い）
actual_lengths = np.random.lognormal(mean=5.0, sigma=1.0, size=num_requests).astype(int)
actual_lengths = np.clip(actual_lengths, 32, max_seq_len)

results = simulate_memory_usage(num_requests, max_seq_len, actual_lengths)

# 可視化
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# 左: 系列長の分布
axes[0].hist(actual_lengths, bins=20, color='#00d4ff', alpha=0.7, edgecolor='white')
axes[0].axvline(x=max_seq_len, color='#ef5350', linestyle='--',
                label=f'Max reserved ({max_seq_len})')
axes[0].axvline(x=np.mean(actual_lengths), color='#ffa726', linestyle='--',
                label=f'Mean actual ({np.mean(actual_lengths):.0f})')
axes[0].set_xlabel('Actual Sequence Length')
axes[0].set_ylabel('Number of Requests')
axes[0].set_title('Distribution of Actual Sequence Lengths')
axes[0].legend()
axes[0].grid(True, alpha=0.3)

# 中央: メモリ使用量の比較
to_mb = lambda x: x / (1024**2)
categories = ['Traditional', 'PagedAttention']
used = [to_mb(results['traditional_used']), to_mb(results['paged_used'])]
waste = [to_mb(results['traditional_waste']), to_mb(results['paged_waste'])]

bars_used = axes[1].bar(categories, used, label='Used', color='#66bb6a', alpha=0.8)
bars_waste = axes[1].bar(categories, waste, bottom=used, label='Wasted',
                         color='#ef5350', alpha=0.8)
axes[1].set_ylabel('Memory (MB)')
axes[1].set_title('Memory Usage Comparison')
axes[1].legend()
axes[1].grid(True, alpha=0.3, axis='y')

# 使用率を表示
trad_util = results['traditional_used'] / results['traditional_total'] * 100
paged_util = results['paged_used'] / results['paged_total'] * 100
axes[1].text(0, to_mb(results['traditional_total']) * 0.5,
             f'{trad_util:.1f}%', ha='center', va='center', fontsize=14,
             fontweight='bold', color='white')
axes[1].text(1, to_mb(results['paged_total']) * 0.5,
             f'{paged_util:.1f}%', ha='center', va='center', fontsize=14,
             fontweight='bold', color='white')

# 右: ブロックサイズによる内部断片化
block_sizes = [4, 8, 16, 32, 64, 128]
internal_wastes = []
for bs in block_sizes:
    total_waste = 0
    for length in actual_lengths:
        blocks = (length + bs - 1) // bs
        total_waste += blocks * bs - length
    waste_ratio = total_waste / sum(actual_lengths) * 100
    internal_wastes.append(waste_ratio)

axes[2].plot(block_sizes, internal_wastes, 'o-', color='#ffa726', markersize=8)
axes[2].set_xlabel('Block Size (tokens)')
axes[2].set_ylabel('Internal Fragmentation (%)')
axes[2].set_title('Internal Fragmentation vs Block Size')
axes[2].grid(True, alpha=0.3)
axes[2].set_xscale('log', base=2)

plt.tight_layout()
plt.savefig('paged_attention_memory.png', dpi=150, bbox_inches='tight')
plt.show()

print(f"従来方式 メモリ利用率: {trad_util:.1f}%")
print(f"PagedAttention メモリ利用率: {paged_util:.1f}%")
print(f"メモリ削減率: {(1 - results['paged_total']/results['traditional_total'])*100:.1f}%")

# COWシミュレーション
block_size = 4
head_dim = 8  # 簡略化
num_total_blocks = 100

bm = BlockManager(num_total_blocks, block_size, 1, 1, head_dim)

# プロンプト部分の処理（16トークン = 4ブロック）
prompt_request = Request(0, bm, block_size)
for t in range(16):
    key = np.random.randn(head_dim)
    value = np.random.randn(head_dim)
    prompt_request.append_token(key, value)

prompt_blocks = prompt_request.block_table.copy()
print(f"プロンプト後のブロック: {prompt_blocks}")
print(f"使用ブロック数: {num_total_blocks - bm.num_free}")

# 並列サンプリング: 4つの応答候補を生成
# 各候補はプロンプトのブロックを共有
parallel_requests = []
for i in range(4):
    req = Request(i + 1, bm, block_size)
    req.block_table = prompt_blocks.copy()
    req.num_tokens = 16
    req.slot_in_last_block = block_size  # 最後のブロックは満杯

    # 共有ブロックの参照カウントを増やす
    for block_id in prompt_blocks:
        bm.add_ref(block_id)

    parallel_requests.append(req)

blocks_after_share = num_total_blocks - bm.num_free
print(f"\n共有後の使用ブロック数: {blocks_after_share}（COWなし: {4 * 4 + 4} = 20）")

# 各候補に独自のトークンを追加
for i, req in enumerate(parallel_requests):
    for t in range(8):  # 8トークンずつ追加
        key = np.random.randn(head_dim) * (i + 1)
        value = np.random.randn(head_dim) * (i + 1)
        req.append_token(key, value)

blocks_after_generation = num_total_blocks - bm.num_free
naive_blocks = 4 * (4 + 2)  # 4候補 × (プロンプト4 + 応答2ブロック)
print(f"生成後の使用ブロック数: {blocks_after_generation}（COWなし: {naive_blocks}）")
print(f"メモリ節約率: {(1 - blocks_after_generation/naive_blocks)*100:.1f}%")