Home 論文解説: Llumnix — ライブマイグレーションによるLLM推論の動的スケジューリング
投稿
キャンセル

📄 論文解説: Llumnix — ライブマイグレーションによるLLM推論の動的スケジューリング

本記事は Llumnix: Dynamic Scheduling for Large Language Model Serving(arXiv:2401.12843、2024年1月公開)の解説記事です。

論文概要(Abstract)

Llumnixは、LLM推論サービングにおける動的スケジューリングシステムである。著者らは、既存のLLMサービング基盤(vLLM等)が採用する静的なロードバランシング(ラウンドロビン等)では、リクエスト間の出力長の分散やバッチサイズの変動により、インスタンス間の負荷が不均衡になる問題を指摘している。Llumnixはリクエスト単位でのKVキャッシュ・ライブマイグレーションを実現し、推論中のリクエストを別インスタンスへ動的に移動することで負荷を平準化する。

この記事は Zenn記事: Azure OpenAI負荷分散設計:API ManagementとPTUスピルオーバーで可用性99.9%を実現する の深掘りです。

情報源

  • arXiv ID: 2401.12843
  • URL: https://arxiv.org/abs/2401.12843
  • 著者: Biao Sun, Ziming Huang, Hanyu Zhao, Wencong Xiao, Xinhe Wan, Minyi Guo, Wei Lin(Alibaba Group, Shanghai Jiao Tong University)
  • 発表年: 2024
  • 分野: cs.DC(分散コンピューティング), cs.LG(機械学習)

背景と動機(Background & Motivation)

LLM推論サービングにおいて、リクエストのスケジューリングは大きな課題である。LLM推論は2つのフェーズ(Prefill: プロンプト処理、Decode: トークン逐次生成)で構成され、各リクエストの処理時間は出力トークン数に依存する。出力トークン数は事前に予測困難であり、同時に処理されるリクエストのバッチサイズも動的に変化する。

既存のLLMサービングシステム(vLLM、Orca等)は、フロントのロードバランサーでリクエストを各インスタンスに振り分けた後は、各インスタンスが独立にスケジューリングを行う。ラウンドロビンやランダム分散では、リクエスト間の処理時間の不均一性によりホットインスタンス(過負荷)とアイドルインスタンス(遊休)が発生する。この問題は、Azure OpenAIのようなマルチインスタンス構成でも同様に発生し得る。

主要な貢献(Key Contributions)

著者らの主要な貢献は以下の3点である。

  • 貢献1: リクエスト単位のKVキャッシュ・ライブマイグレーション機構。推論中のリクエストを別GPUインスタンスに移動し、KVキャッシュをP2Pで転送する。再計算なしで推論を継続可能。
  • 貢献2: 動的スケジューリングアルゴリズム。全インスタンスのキュー長・GPU利用率・メモリ使用量をリアルタイム監視し、マイグレーションのコスト(転送時間)と利益(負荷平準化)を比較して移動判断を行う。
  • 貢献3: vLLMバックエンド上でのプロトタイプ実装と実トレースによる評価。ShareGPT/LMSYS-Chat-1Mトレースで、ラウンドロビン比TTFTを最大6.7倍改善。

技術的詳細(Technical Details)

アーキテクチャ

Llumnixは以下のコンポーネントで構成される。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
                    ┌──────────────────┐
                    │  Global Scheduler│
                    │  (状態監視+判断) │
                    └───────┬──────────┘
                            │
            ┌───────────────┼───────────────┐
            ▼               ▼               ▼
    ┌──────────────┐ ┌──────────────┐ ┌──────────────┐
    │  Instance 1  │ │  Instance 2  │ │  Instance 3  │
    │  (vLLM)      │ │  (vLLM)      │ │  (vLLM)      │
    │  ┌─────────┐ │ │  ┌─────────┐ │ │  ┌─────────┐ │
    │  │KV Cache │ │ │  │KV Cache │ │ │  │KV Cache │ │
    │  │Manager  │ │ │  │Manager  │ │ │  │Manager  │ │
    │  └─────────┘ │ │  └─────────┘ │ │  └─────────┘ │
    └──────────────┘ └──────────────┘ └──────────────┘

Global Scheduler: 全インスタンスの状態(キュー長、GPU利用率、KVキャッシュ使用量)を集約し、マイグレーション判断を行う。

KVキャッシュ・ライブマイグレーション

LLM推論では、各リクエストのKVキャッシュがGPUメモリに保持される。リクエストを別インスタンスに移動するには、このKVキャッシュを転送する必要がある。

KVキャッシュのサイズは以下で計算される:

\[\text{KV Cache Size} = 2 \times L \times H \times d_k \times s \times \text{dtype\_size}\]

ここで、

  • $L$: Transformerの層数
  • $H$: アテンションヘッド数
  • $d_k$: ヘッドあたりの次元数
  • $s$: シーケンス長(これまでに生成されたトークン数)
  • $\text{dtype_size}$: データ型のバイト数(FP16の場合2バイト)
  • 係数2はKey, Valueの2つ分

例えばLLaMA-2-13Bの場合、$L=40, H=40, d_k=128$であり、シーケンス長1024トークンのリクエストのKVキャッシュは約640MBとなる。

マイグレーション判断アルゴリズム

著者らは、マイグレーションのコスト(KVキャッシュ転送時間によるレイテンシ増加)と利益(負荷平準化によるキュー遅延削減)を比較するアルゴリズムを提案している。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def should_migrate(source: Instance, target: Instance, request: Request) -> bool:
    """マイグレーション判断

    Args:
        source: 移動元インスタンス
        target: 移動先インスタンス
        request: 移動候補リクエスト

    Returns:
        マイグレーションすべきかどうか
    """
    # マイグレーションコスト:KVキャッシュ転送時間
    kv_size = compute_kv_cache_size(request)
    transfer_time = kv_size / network_bandwidth

    # マイグレーション利益:キュー遅延の改善
    source_queue_delay = estimate_queue_delay(source)
    target_queue_delay = estimate_queue_delay(target)
    benefit = source_queue_delay - target_queue_delay

    return benefit > transfer_time * alpha  # alphaは調整パラメータ

スケジューリングポリシー

著者らは3つのスケジューリングポリシーを比較している。

ポリシー説明適用場面
Load-Awareキュー長に基づく配分一般的なワークロード
Deficit-Based不足分を動的に補填バーストトラフィック
Hybrid両者の組み合わせ本番推奨

実装のポイント(Implementation)

P2P KVキャッシュ転送

KVキャッシュの転送は、GPU間の直接通信(NVLink/PCIe経由のP2P転送)で行われる。著者らの実装ではvLLM上でRayを使用し、インスタンス間のGPU通信を管理している。

転送中も推論は継続可能(ライブマイグレーション)であり、転送完了後にリクエストの処理が移動先インスタンスに引き継がれる。著者らの報告では、マイグレーションのオーバーヘッドは平均200ms程度とのことである(論文§5.4)。

vLLMとの統合

LlumnixはvLLMのスケジューラAPIにフックを挿入する形で実装されている。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
# Llumnixのスケジューラフック(概念的なコード)
class LlumnixScheduler:
    """Llumnixのグローバルスケジューラ

    vLLMインスタンス群を統括し、動的マイグレーションを管理する。
    """
    def __init__(self, instances: list[vLLMInstance]):
        self.instances = instances
        self.monitor = StateMonitor(instances)

    def schedule(self, new_request: Request) -> vLLMInstance:
        """新規リクエストの配置先を決定"""
        loads = self.monitor.get_instance_loads()
        return min(self.instances, key=lambda i: loads[i.id])

    def rebalance(self) -> list[Migration]:
        """既存リクエストの再配置を計画"""
        imbalance = self.monitor.compute_imbalance()
        if imbalance < threshold:
            return []
        return self._plan_migrations(imbalance)

実験結果(Results)

著者らの実験は、LLaMA-2-13B/70Bモデル、A100 GPUクラスターで実施されている。

主要な実験結果(論文Table 1, Figure 7より):

指標ラウンドロビンLlumnix改善率
TTFT P501.2秒0.18秒6.7x
TTFT P998.5秒1.1秒7.7x
JCT P5015.2秒8.9秒1.7x
JCT P9945.3秒3.8秒12x

※ TTFT = Time to First Token, JCT = Job Completion Time

著者らは、改善の主因としてホットインスタンスの解消を挙げている。ラウンドロビンでは一部のインスタンスにリクエストが集中し、キュー待ち時間が膨れ上がる。Llumnixのライブマイグレーションにより、過負荷インスタンスから遊休インスタンスへリクエストを移動し、全体のキュー遅延を平準化している。

実運用への応用(Practical Applications)

Azure OpenAIの負荷分散との関連

Azure OpenAI ServiceはマネージドAPIであり、ユーザーがKVキャッシュ層を直接制御することはできない。しかし、Llumnixの設計思想はAzure API Managementのバックエンドプール構成に応用可能である。

  • 「ホットインスタンスを避ける」思想: ラウンドロビンではなく、負荷を考慮したルーティング(Priority-based, Weighted)を採用することで、特定エンドポイントへの偏りを軽減する。
  • マイグレーションコスト vs 利益の比較: Circuit BreakerのtripDuration(フェイルオーバーのコスト)と、スロットリングによるレイテンシ増加(放置した場合の損失)を比較する設計判断に、Llumnixのコスト・利益フレームワークを適用できる。
  • バースト耐性: Llumnixの実験で示されたバースト係数(ピーク/平均比5〜15倍)は、Azure OpenAI PTUの過剰プロビジョニング率を決定する際の参考値として有用である。

自前LLMクラスターでの活用

vLLMベースの自前クラスターを運用する場合、LlumnixのOSS実装(GitHub: AlibabaPAI/llumnix、Apache 2.0ライセンス)を直接導入可能である。

関連研究(Related Work)

  • vLLM (Kwonら, 2023, arXiv:2309.06180): PagedAttentionによるKVキャッシュ管理。Llumnixのバックエンドエンジンとして使用されている。
  • Orca (Yuら, 2022): Iteration-levelスケジューリングを提案。Llumnixはさらにインスタンス間の動的再配置を加えている。
  • DistServe (Zhongら, 2024, arXiv:2403.02310): Prefill/Decodeの分離によるgoodput最適化。Llumnixとは直交するアプローチであり、組み合わせが可能。

Production Deployment Guide

AWS実装パターン(コスト最適化重視)

Llumnixの動的スケジューリングの考え方をAWS上のLLM推論基盤で実現する構成を示す。

トラフィック量別の推奨構成:

規模月間リクエスト推奨構成月額コスト主要サービス
Small~3,000Serverless$50-150Lambda + Bedrock
Medium~30,000Hybrid$500-1,500ECS Fargate + vLLM + ALB
Large300,000+Container$3,000-8,000EKS + vLLM + Karpenter + Spot GPU

Large構成(Llumnix相当)の詳細 (月額$3,000-8,000):

  • EKS: コントロールプレーン ($72/月)
  • EC2 Spot GPU: g5.xlarge × 4台 (平均$1,200/月、Spot最大90%削減)
  • Karpenter: GPU需要に応じた自動スケーリング
  • vLLM + Llumnix: OSSスタック、追加ライセンスコスト不要
  • ALB + Target Group: ヘルスチェックベースの初期ルーティング

コスト試算の注意事項:

  • 上記は2026年2月時点のAWS ap-northeast-1料金に基づく概算値です
  • GPU Spot Instancesの価格は需給により変動します
  • 最新料金は AWS料金計算ツール で確認してください

Terraformインフラコード

Large構成: EKS + vLLM + GPU Spot

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
module "eks" {
  source  = "terraform-aws-modules/eks/aws"
  version = "~> 20.0"

  cluster_name    = "llm-inference"
  cluster_version = "1.31"

  vpc_id     = module.vpc.vpc_id
  subnet_ids = module.vpc.private_subnets

  cluster_endpoint_public_access = true
  enable_cluster_creator_admin_permissions = true
}

# Karpenter Provisioner(GPU Spot優先)
resource "kubectl_manifest" "gpu_provisioner" {
  yaml_body = <<-YAML
    apiVersion: karpenter.sh/v1alpha5
    kind: Provisioner
    metadata:
      name: gpu-spot
    spec:
      requirements:
        - key: karpenter.sh/capacity-type
          operator: In
          values: ["spot"]
        - key: node.kubernetes.io/instance-type
          operator: In
          values: ["g5.xlarge", "g5.2xlarge"]
      limits:
        resources:
          cpu: "64"
          nvidia.com/gpu: "8"
      providerRef:
        name: default
      ttlSecondsAfterEmpty: 60
  YAML
}

コスト最適化チェックリスト

  • GPU Spot Instances使用で最大90%削減
  • Karpenterでアイドル時スケールダウン(60秒タイムアウト)
  • vLLM PagedAttentionでGPUメモリ効率最大化
  • バッチサイズ最適化でスループット向上
  • CloudWatch GPU利用率メトリクスで監視
  • AWS Budgets月額予算設定

まとめと今後の展望

Llumnixは、LLM推論サービングにおける動的負荷分散の学術的基盤を提供する重要な研究である。著者らが示したKVキャッシュ・ライブマイグレーションの手法は、静的ロードバランシングの限界を克服し、TTFTを最大6.7倍改善するという結果を報告している(論文Table 1より)。

Azure OpenAI Serviceの文脈では、Llumnixの知見を以下の形で応用可能である:

  • バックエンドプールの重み設計における「負荷考慮型ルーティング」の理論的根拠
  • PTUの過剰プロビジョニング率の決定(バースト係数の参考値)
  • Circuit Breakerのパラメータチューニング(コスト・利益フレームワーク)

今後の研究方向として、著者らはマルチモデル環境でのスケジューリング最適化やエネルギー効率を考慮した配置戦略を挙げている。

参考文献


:::message この記事はAI(Claude Code)により自動生成されました。論文の内容は原著者の主張を要約・解説したものであり、本記事の著者が独自に実験を行ったものではありません。正確な数値・詳細は原論文をご確認ください。 :::

この投稿は CC BY 4.0 でライセンスされています。