Home 論文解説: LLUMNIX — リクエストライブマイグレーションによるLLM推論の動的負荷分散
投稿
キャンセル

📄 論文解説: LLUMNIX — リクエストライブマイグレーションによるLLM推論の動的負荷分散

本記事は LLUMNIX: Dynamic Scheduling for Large Language Model Serving (arXiv:2401.11351) の解説記事です。

論文概要(Abstract)

LLUMNIXは、LLM推論サービングクラスタにおいて実行中のリクエストをKVキャッシュごとGPUインスタンス間でライブマイグレーションする機構を提案した論文である。従来のLLMサービングシステム(vLLM、Orca等)ではリクエストが割り当てられたGPUインスタンス上で完了まで処理されるため、バーストトラフィック時に特定インスタンスに負荷が集中する問題があった。著者らは、リクエスト単位のマイグレーションにより動的な負荷再分配を実現し、P99 TTFT(Time to First Token)を最大3.3倍削減、スループットを最大36%向上させたと報告している。

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

情報源

背景と動機(Background & Motivation)

LLMの推論サービングは、Prefillフェーズ(入力トークンの処理)とDecodeフェーズ(出力トークンの逐次生成)から構成される。vLLMに代表される既存サービングシステムでは、リクエストは受信時にスケジューラが選択した1つのGPUインスタンスに割り当てられ、そのインスタンス上で処理が完了するまで固定される。

この静的な割り当て方式には以下の課題がある。

  1. 負荷不均衡: バーストトラフィック時に一部のインスタンスにリクエストが集中し、他のインスタンスがアイドル状態になる。初回割り当て時にはリクエストの実行時間(出力長)が不明であるため、完全な負荷均等化は困難である。

  2. KVキャッシュの断片化: 短いリクエストが完了した後にメモリが断片化し、長いコンテキストの新規リクエストを受け入れられない状態が発生する。

  3. プリエンプションの非効率性: vLLMでは、メモリ不足時にリクエストをキューに戻すSwap/Recompute方式のプリエンプションが行われるが、これらはリクエストの再処理コストが高い。

Zenn記事で解説されているAzure API ManagementのPriority-based負荷分散(PTU優先→PAYG)はAPIゲートウェイ層での負荷分散であり、リクエストの初回ルーティング時に適用される。一方、LLUMNIXは推論エンジン層で、リクエストの実行中に動的な再配置を行うアプローチであり、両者は相補的な関係にある。

主要な貢献(Key Contributions)

  • 貢献1: LLM推論リクエストをKVキャッシュごとライブマイグレーションする機構の設計と実装。マイグレーション中もリクエストの処理を継続できるため、サービス中断が発生しない。

  • 貢献2: グローバルスケジューラによる動的負荷再分配アルゴリズム。キュー長、メモリ使用率、断片化率を入力として、マイグレーション対象のリクエストとターゲットインスタンスを選定する。

  • 貢献3: vLLMをバックエンドとした実装(Apache 2.0ライセンスでOSS公開)と、最大32インスタンス規模での評価。LLaMA-2-7B/13B/70Bを用いた実験で、P99 TTFTの3.3倍削減とスループット36%向上を確認。

技術的詳細(Technical Details)

アーキテクチャ

LLUMNIXは3層のアーキテクチャで構成される。

グローバルスケジューラ(Global Scheduler): クラスタ全体のリクエスト分配とマイグレーション判断を担当する。各インスタンスからキュー長 $Q_i$ とメモリ使用率 $M_i$ を定期的に収集し、負荷不均衡を検知する。

ローカルスケジューラ(Local Scheduler): 各GPUインスタンス上で動作し、割り当てられたリクエストのPrefill/Decodeスケジューリングを管理する。vLLMのスケジューラをベースとしている。

マイグレーションエンジン: KVキャッシュのGPU間転送を担当する。転送はページ単位で行われ、vLLMのPagedAttentionのメモリレイアウトと互換性がある。

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

著者らは、マイグレーションの判断に以下のコスト関数を用いている。

\[C_{\text{migrate}}(r, i \to j) = T_{\text{transfer}}(r) + T_{\text{interrupt}}(r) - G_{\text{balance}}(i, j)\]

ここで、

  • $r$: マイグレーション対象のリクエスト
  • $i, j$: ソース/ターゲットインスタンス
  • $T_{\text{transfer}}(r)$: KVキャッシュ転送時間(KVキャッシュサイズに比例)
  • $T_{\text{interrupt}}(r)$: マイグレーション中のリクエスト処理中断時間
  • $G_{\text{balance}}(i, j)$: 負荷均等化により得られるレイテンシ改善の推定値

$C_{\text{migrate}} < 0$ の場合にマイグレーションが実行される。つまり、マイグレーションによる負荷均等化の利得が、転送コストを上回る場合にのみ実行される。

KVキャッシュ転送の詳細

KVキャッシュの転送時間は以下の式で推定される。

\[T_{\text{transfer}}(r) = \frac{S_{\text{KV}}(r)}{B_{\text{net}}}\]

ここで、

  • $S_{\text{KV}}(r) = n_{\text{layers}} \times 2 \times d_{\text{head}} \times n_{\text{tokens}}(r) \times \text{sizeof}(\text{dtype})$: リクエスト $r$ のKVキャッシュサイズ
  • $n_{\text{layers}}$: Transformerのレイヤー数
  • $d_{\text{head}}$: ヘッド次元
  • $n_{\text{tokens}}(r)$: 現時点までに処理されたトークン数
  • $B_{\text{net}}$: GPU間ネットワーク帯域幅(NVLink/InfiniBand)

論文Table 2より、LLaMA-2-13B(40層、$d_{\text{head}} = 128$)で1024トークン処理済みのリクエストの場合、KVキャッシュサイズは約80MBとなり、NVLink(600GB/s)環境では転送時間は0.13ms、InfiniBand 200Gbps環境では約3.2msと報告されている。

負荷均衡指標

グローバルスケジューラは、以下の不均衡指標 $I$ を定期的に計算する。

\[I = \frac{\max_i Q_i - \min_i Q_i}{\bar{Q}}\]

ここで $Q_i$ はインスタンス $i$ の待機中リクエスト数、$\bar{Q}$ はクラスタ全体の平均キュー長である。$I$ が閾値 $\theta$ を超えた場合にマイグレーションがトリガーされる。論文の実験では $\theta = 2.0$ が推奨されている。

擬似コード

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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
from dataclasses import dataclass
from typing import Optional

@dataclass
class MigrationDecision:
    """マイグレーション判断の結果"""
    request_id: str
    source_instance: int
    target_instance: int
    estimated_cost_ms: float

def should_migrate(
    queue_lengths: list[int],
    memory_usage: list[float],
    threshold: float = 2.0
) -> Optional[MigrationDecision]:
    """グローバルスケジューラのマイグレーション判断ロジック

    Args:
        queue_lengths: 各インスタンスのキュー長
        memory_usage: 各インスタンスのメモリ使用率 (0.0-1.0)
        threshold: 不均衡閾値

    Returns:
        マイグレーション判断(不要ならNone)
    """
    avg_q = sum(queue_lengths) / len(queue_lengths)
    if avg_q == 0:
        return None

    max_q = max(queue_lengths)
    min_q = min(queue_lengths)
    imbalance = (max_q - min_q) / avg_q

    if imbalance < threshold:
        return None  # 負荷均衡 → マイグレーション不要

    # ソース: 最も負荷の高いインスタンス
    source = queue_lengths.index(max_q)
    # ターゲット: 最も負荷の低いインスタンス(メモリに余裕あり)
    candidates = [
        i for i, (q, m) in enumerate(zip(queue_lengths, memory_usage))
        if q == min_q and m < 0.8  # メモリ使用率80%未満
    ]
    if not candidates:
        return None

    target = candidates[0]

    return MigrationDecision(
        request_id="oldest_in_queue",  # 最も古いリクエストを選択
        source_instance=source,
        target_instance=target,
        estimated_cost_ms=3.2  # InfiniBand環境での推定転送時間
    )

実装のポイント(Implementation)

LLUMNIXの実装において、以下の点が実務上重要である。

vLLMとの統合: LLUMNIXはvLLM 0.3.xをベースに実装されている。PagedAttentionのブロック管理を拡張し、マイグレーション用のメモリ転送APIを追加している。ただし、vLLMのバージョンアップに追従する必要があるため、最新バージョン(v0.6+)との互換性は別途確認が必要である。

Ray Clusterの構成: マルチインスタンス管理にはRayを使用している。Ray Clusterのセットアップ、ヘッドノードの冗長化、ワーカーノードのオートスケール設定が前提となる。

ネットワーク要件: KVキャッシュの転送は高帯域幅ネットワークが望ましい。論文ではNVLink/InfiniBand環境で評価されているが、Ethernet(25-100GbE)環境でも動作する。ただし、転送時間が増加するため、マイグレーション閾値 $\theta$ の調整が必要になる。

マイグレーション中のリクエスト処理: マイグレーション中はソースインスタンスでDecodeが一時停止する(実測100ms未満)。ストリーミング出力のユーザーにはこの中断が短い遅延として見える可能性がある。

実験結果(Results)

著者らはA100 40GB GPUクラスタ(最大32インスタンス)で、LLaMA-2-7B/13B/70Bを用いた評価を行っている。

メインの実験結果(論文Table 3/Figure 6より):

構成指標vLLMLLUMNIX改善率
8インスタンス・バースト負荷P99 TTFT12.4s3.8s3.3x削減
8インスタンス・バースト負荷平均TBT48ms42ms12.5%改善
16インスタンス・定常負荷スループット1.0x1.36x36%向上
32インスタンス・混合負荷P99 TTFT8.7s3.1s2.8x削減

マイグレーション頻度: バースト負荷時に1分あたり平均5-15回のマイグレーションが発生し、転送コストは全体の推論時間の0.3%未満であったと報告されている。

スケーラビリティ: グローバルスケジューラの判断レイテンシは32インスタンスで0.5ms未満であり、インスタンス数に対してほぼ線形にスケールすることが確認されている。

実運用への応用(Practical Applications)

Zenn記事で解説されているAzure API Management + PTU/PAYGスピルオーバー構成は、APIゲートウェイ層での負荷分散である。一方、LLUMNIXは推論エンジン層で動的な負荷再配置を行う。実運用では、この2つのレイヤーを組み合わせることで、より精緻な負荷分散が可能になる。

組み合わせパターン: API Management(Priority-basedルーティング)→ PTUクラスタ内でLLUMNIX(リクエストマイグレーション)→ PAYG(スピルオーバー先)。PTUクラスタ内の複数GPUインスタンス間でLLUMNIXが動的負荷分散を行い、クラスタ全体の容量を超えた場合にAPI ManagementがPAYGへフェイルオーバーする構成である。

Azure環境での適用可能性: LLUMNIXはvLLM + Rayベースであるため、Azure Kubernetes Service(AKS)上にデプロイ可能である。ただし、Azure OpenAI Service(マネージドAPI)では内部の推論エンジンを制御できないため、自前でLLMを運用する場合に限定される。

Production Deployment Guide

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

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

規模月間リクエスト推奨構成月額コスト主要サービス
Small~3,000 (100/日)Serverless + Bedrock$50-200Lambda + Bedrock API
Medium~30,000 (1,000/日)ECS + vLLM$3,000-5,000ECS Fargate + g5.xlarge
Large300,000+ (10,000/日)EKS + LLUMNIX$10,000-25,000EKS + p4d/g5 Spot + Ray

Small構成: Bedrock APIを直接使用する場合、LLUMNIXは不要。マネージドサービス側でスケーリングが行われる。

Medium構成 (月額$3,000-5,000):

  • ECS Fargate: g5.xlarge(1 GPU)×2-4タスク ($2,400/月)
  • vLLM: コンテナイメージでデプロイ ($0/月)
  • Application Load Balancer: ラウンドロビン ($20/月)
  • CloudWatch: メトリクス収集 ($10/月)
  • S3: モデルウェイト保存 ($50/月)

Large構成 (月額$10,000-25,000):

  • EKS: コントロールプレーン ($72/月)
  • EC2 Spot: g5.xlarge × 8-16台 (平均$4,000/月、Spot利用で最大70%削減)
  • Ray Cluster: Kuberay Operator上でデプロイ ($0/月)
  • LLUMNIX: vLLM上でリクエストマイグレーション ($0/月)
  • EFS: 共有モデルストレージ ($200/月)
  • CloudWatch + X-Ray: 詳細監視 ($100/月)

コスト試算の注意事項: 上記は2026年2月時点のAWS ap-northeast-1(東京)リージョン料金に基づく概算値です。GPU インスタンスの料金はリージョンと可用性により大きく変動します。最新料金は AWS料金計算ツール で確認してください。

Terraformインフラコード

Large構成 (EKS + Ray + LLUMNIX):

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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
# --- EKSクラスタ ---
module "eks" {
  source  = "terraform-aws-modules/eks/aws"
  version = "~> 20.0"

  cluster_name    = "llm-serving-cluster"
  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

  eks_managed_node_groups = {
    gpu-workers = {
      instance_types = ["g5.xlarge"]
      capacity_type  = "SPOT"  # Spot Instancesで最大70%削減
      min_size       = 4
      max_size       = 16
      desired_size   = 8

      labels = {
        "nvidia.com/gpu" = "true"
        "workload"       = "llm-inference"
      }

      taints = [{
        key    = "nvidia.com/gpu"
        value  = "true"
        effect = "NO_SCHEDULE"
      }]
    }

    ray-head = {
      instance_types = ["m5.2xlarge"]
      capacity_type  = "ON_DEMAND"  # Ray Headはオンデマンド
      min_size       = 1
      max_size       = 1
      desired_size   = 1

      labels = {
        "ray.io/node-type" = "head"
      }
    }
  }
}

# --- Kuberay Operator ---
resource "helm_release" "kuberay" {
  name       = "kuberay-operator"
  repository = "https://ray-project.github.io/kuberay-helm/"
  chart      = "kuberay-operator"
  version    = "1.2.0"
  namespace  = "ray-system"

  create_namespace = true
}

# --- CloudWatchアラーム(GPU使用率監視) ---
resource "aws_cloudwatch_metric_alarm" "gpu_utilization" {
  alarm_name          = "llm-gpu-utilization-low"
  comparison_operator = "LessThanThreshold"
  evaluation_periods  = 3
  metric_name         = "GPUUtilization"
  namespace           = "Custom/LLMServing"
  period              = 300
  statistic           = "Average"
  threshold           = 30  # 30%未満でスケールダウン検討
  alarm_description   = "GPU使用率が低い(オーバープロビジョニング)"
}

# --- AWS Budgets ---
resource "aws_budgets_budget" "llm_serving" {
  name         = "llm-serving-monthly"
  budget_type  = "COST"
  limit_amount = "25000"
  limit_unit   = "USD"
  time_unit    = "MONTHLY"

  notification {
    comparison_operator       = "GREATER_THAN"
    threshold                 = 80
    threshold_type            = "PERCENTAGE"
    notification_type         = "ACTUAL"
    subscriber_email_addresses = ["ops@example.com"]
  }
}

運用・監視設定

CloudWatch Logs Insights クエリ:

1
2
3
4
5
6
7
8
9
-- LLUMNIX マイグレーション頻度の監視
fields @timestamp, migration_source, migration_target, kv_cache_size_mb, transfer_time_ms
| stats count() as migrations, avg(transfer_time_ms) as avg_transfer by bin(1h)
| filter migrations > 20  -- 1時間あたり20回超でアラート

-- GPU使用率の不均衡検知
fields @timestamp, instance_id, gpu_utilization
| stats max(gpu_utilization) as max_gpu, min(gpu_utilization) as min_gpu by bin(5m)
| filter (max_gpu - min_gpu) > 50  -- 50%以上の差異で検知

メトリクス収集コード(Python):

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
39
40
41
import boto3
from datetime import datetime

cloudwatch = boto3.client('cloudwatch')

def publish_llumnix_metrics(
    migrations_per_minute: int,
    avg_transfer_ms: float,
    imbalance_ratio: float,
    queue_lengths: list[int]
) -> None:
    """LLUMNIXの運用メトリクスをCloudWatchに送信"""
    cloudwatch.put_metric_data(
        Namespace='Custom/LLMServing',
        MetricData=[
            {
                'MetricName': 'MigrationRate',
                'Value': migrations_per_minute,
                'Unit': 'Count/Second',
                'Timestamp': datetime.utcnow()
            },
            {
                'MetricName': 'AvgTransferLatency',
                'Value': avg_transfer_ms,
                'Unit': 'Milliseconds',
                'Timestamp': datetime.utcnow()
            },
            {
                'MetricName': 'LoadImbalanceRatio',
                'Value': imbalance_ratio,
                'Unit': 'None',
                'Timestamp': datetime.utcnow()
            },
            {
                'MetricName': 'MaxQueueLength',
                'Value': max(queue_lengths),
                'Unit': 'Count',
                'Timestamp': datetime.utcnow()
            }
        ]
    )

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

アーキテクチャ選択:

  • ~100 req/日 → Bedrock API直接利用(LLUMNIX不要)
  • ~1,000 req/日 → ECS + vLLM + ALBラウンドロビン
  • 10,000+ req/日 → EKS + Ray + LLUMNIX(本論文の構成)

リソース最適化:

  • GPU: Spot Instances優先(g5.xlarge、最大70%削減)
  • Ray Head: On-Demandで安定性確保
  • Karpenterによる自動スケーリング
  • 夜間のスケールダウン(最小4台→最小2台)
  • EFS → S3 FSx for Lustreへの変更検討(モデルロード高速化)

LLMサービングコスト削減:

  • KVキャッシュサイズの監視と最大コンテキスト長の制限
  • マイグレーション閾値の調整(過度なマイグレーション防止)
  • バッチサイズの最適化(GPU使用率とレイテンシのトレードオフ)
  • 量子化(INT8/FP8)によるメモリ効率改善

監視・アラート:

  • CloudWatch: GPU使用率、キュー長、マイグレーション頻度
  • AWS Budgets: 月額予算設定(80%で警告)
  • Cost Anomaly Detection有効化
  • Prometheus + Grafana: vLLMメトリクスダッシュボード

セキュリティ:

  • EKSクラスタ: プライベートエンドポイント
  • IAM: 最小権限ロール
  • VPC: プライベートサブネット内にGPUノード配置
  • モデルウェイト: S3暗号化(SSE-KMS)

関連研究(Related Work)

  • vLLM (Kwon et al., 2023): PagedAttentionによるKVキャッシュの効率的なメモリ管理を提案。LLUMNIXはvLLMのスケジューラを拡張しており、PagedAttentionのブロック管理と互換性がある。

  • Orca (Yu et al., 2022): Iteration-level schedulingを提案し、Continuous Batchingの基盤を確立。LLUMNIXはOrcaのスケジューリングを前提とした上で、インスタンス間の動的再配置を追加している。

  • DistServe (Zhong et al., 2024): PrefillとDecodeを物理的に分離するアーキテクチャ。LLUMNIXはPrefill/Decodeを分離せず、リクエスト単位のマイグレーションで負荷分散するアプローチであり、設計思想が異なる。

まとめと今後の展望

LLUMNIXは、LLM推論クラスタにおける動的負荷分散の課題に対して、リクエストのライブマイグレーションという手法で取り組んでいる。著者らの実験では、バーストトラフィック時のP99 TTFTを3.3倍削減し、定常負荷時のスループットを36%向上させたと報告されている。

Zenn記事のAzure API Management負荷分散がAPIゲートウェイ層での制御であるのに対し、LLUMNIXは推論エンジン層での最適化であり、両者を組み合わせることでより精緻な負荷管理が可能になる。ただし、LLUMNIXの適用にはRay Clusterの運用コストとKVキャッシュ転送のネットワーク要件が伴うため、マネージドAPI(Azure OpenAI、Amazon Bedrock等)で十分な場合はそちらを優先すべきである。

参考文献

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