効率的な KV キャッシュ管理と Attention 計算

大規模言語モデル（LLM）の推論プロセスにおいて、Attention 層の最適化はレイテンシ削減の鍵となります。特に、Qwen3 などの最新アーキテクチャでは、Triton によるカスタムカーネルを用いた KV キャッシュの書き込みと、FlashAttention ライブラリを組み合わせたハイブリッドなアプローチが採用されています。この構成により、Prefill（初期入力処理）段階と Decode（トークン生成）段階の双方で高い効率性を実現します。

Triton カーネルによる KV キャッシュ書き込み

新生成された Key と Value ベクトルをグローバルな KV キャッシュ領域へ転送する処理は、Triton カーネルによって高速化されます。このカーネルは、各トークンに対応するキャッシュスロットへのマッピングを行い、メモリアクセスを最適化します。

@triton.jit
def write_kv_cache_triton(
    k_src_ptr, k_step,
    v_src_ptr, v_step,
    k_buf_ptr, v_buf_ptr,
    slot_map_ptr,
    hidden_size: tl.constexpr
):
    token_idx = tl.program_id(0)
    # Compute source offsets
    k_offs = token_idx * k_step + tl.arange(0, hidden_size)
    v_offs = token_idx * v_step + tl.arange(0, hidden_size)
    
    # Load vectors
    k_vec = tl.load(k_src_ptr + k_offs)
    v_vec = tl.load(v_src_ptr + v_offs)
    
    # Determine cache slot
    slot_id = tl.load(slot_map_ptr + token_idx)
    cache_offs = slot_id * hidden_size + tl.arange(0, hidden_size)
    
    # Persist to cache
    tl.store(k_buf_ptr + cache_offs, k_vec)
    tl.store(v_buf_ptr + cache_offs, v_vec)

ここで、hidden_size は num_heads * head_dim に相当し、トークンごとの展平された次元数を表します。slot_map_ptr により、バッチ内のトークンがキャッシュ上のどの物理位置に格納されるかが決定されます。これにより、動的なシーケンス長やバッチ拼接の際にも、重複計算なしでキャッシュ更新が可能になります。

Python 層からのカーネル呼び出し

Triton カーネルを PyTorch テンソルから起動するためのラッパー関数です。メモリレイアウトの連続性を確認し、適切なグリッドサイズでカーネルを実行します。

def invoke_kv_cache_write(k_tensor, v_tensor, k_buffer, v_buffer, slot_map):
    batch_sz, n_heads, head_dim = k_tensor.shape
    total_dim = n_heads * head_dim
    
    # Validate memory continuity
    assert k_tensor.stride(-1) == 1 and v_tensor.stride(-1) == 1
    assert k_tensor.stride(1) == head_dim and v_tensor.stride(1) == head_dim
    assert k_buffer.stride(1) == total_dim and v_buffer.stride(1) == total_dim
    assert slot_map.numel() == batch_sz
    
    # Execute kernel
    write_kv_cache_triton[(batch_sz,)](
        k_tensor, k_tensor.stride(0),
        v_tensor, v_tensor.stride(0),
        k_buffer, v_buffer,
        slot_map,
        total_dim
    )

アサートにより、テンソルの stride が期待通りであることを保証し、Triton 側でのメモリアクセスエラーを防ぎます。グリッドサイズ (batch_sz,) は、各ブロックが 1 つのトークンを処理することを意味します。

Attention モジュールの実装

実際の Attention 計算を行うモジュールでは、状況に応じて FlashAttention の異なる関数を使い分けます。GQA（Grouped Query Attention）や Prefix Cache にも対応可能な構造になっています。

class EfficientAttentionLayer(nn.Module):

    def __init__(
        self,
        n_heads,
        head_dim,
        scale_factor,
        n_kv_heads,
    ):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.scale = scale_factor
        self.n_kv_heads = n_kv_heads
        self.k_cache = self.v_cache = torch.tensor([])

    def forward(self, query, key, value):
        output: torch.Tensor
        # Reshape inputs to [tokens, heads, dim]
        query = query.reshape(-1, self.n_heads, self.head_dim)
        key = key.reshape(-1, self.n_kv_heads, self.head_dim)
        value = value.reshape(-1, self.n_kv_heads, self.head_dim)
        
        runtime_ctx = get_execution_context()
        k_buf = self.k_cache
        v_buf = self.v_cache
        
        # Update global KV cache
        invoke_kv_cache_write(key, value, k_buf, v_buf, runtime_ctx.slot_map)
        
        if runtime_ctx.phase == 'prefill':
            if runtime_ctx.block_tables is not None:
                # Use cached data for prefix sharing
                key, value = k_buf, v_buf
            
            output = flash_attn_varlen_func(
                query, key, value,
                max_seqlen_q=runtime_ctx.max_q_len,
                cu_seqlens_q=runtime_ctx.cum_q_lens,
                max_seqlen_k=runtime_ctx.max_k_len,
                cu_seqlens_k=runtime_ctx.cum_k_lens,
                softmax_scale=self.scale,
                causal=True,
                block_table=runtime_ctx.block_tables
            )
        else:
            # Decode phase: attend only to cached history
            output = flash_attn_with_kvcache(
                query.unsqueeze(1),
                k_buf, v_buf,
                cache_seqlens=runtime_ctx.current_lens,
                block_table=runtime_ctx.block_tables,
                softmax_scale=self.scale,
                causal=True
            )
        
        return output.reshape(-1, self.n_heads * self.head_dim)

コンテキストオブジェクト runtime_ctx には、シーケンス長情報やブロックテーブルなどが含まれています。Prefill 段階では flash_attn_varlen_func を使用し、変長バッチを効率的に処理します。一方、Decode 段階では flash_attn_with_kvcache を利用し、過去の KV キャッシュを参照しながら新しいトークンを生成します。

この実装により、num_kv_heads が num_heads より少ない GQA 構成や、複数のリクエストで共通の Prefix を共有するケースにも柔軟に対応できます。出力は再度展平され、隠れ層サイズ num_heads * head_dim のテンソルとして後続の層へ渡されます。