Home 論文解説: Intelligent Router for LLM Workloads — ワークロード特性を活用したLLM推論の負荷分散
投稿
キャンセル

📄 論文解説: Intelligent Router for LLM Workloads — ワークロード特性を活用したLLM推論の負荷分散

本記事は Intelligent Router for LLM Workloads: Improving Performance Through Workload-Aware Load Balancing の解説記事です。

論文概要(Abstract)

LLM推論ワークロードにはprefillフェーズとdecodeフェーズという異なるリソース要求を持つ2つの段階が存在する。著者らは、既存のスケジューリング手法がこれらのフェーズ特性を考慮していない問題を指摘し、ヒューリスティック誘導型の強化学習ベースインテリジェントルーターを提案している。応答長予測器とワークロード混合影響推定器という2つの新規コンポーネントにより、公開データセットで11%以上、実運用クラウドデータで7.8%のエンドツーエンドレイテンシ削減を達成したと報告されている。

この記事は Zenn記事: Azure OpenAI負荷分散2026年版 の深掘りです。

情報源

  • arXiv ID: 2408.13510
  • URL: https://arxiv.org/abs/2408.13510
  • 著者: Kunal Jain, Anjaly Parayil, Ankur Mallick, Esha Choukse, Xiaoting Qin, Jue Zhang, Íñigo Goiri, Rujia Wang, Chetan Bansal, Victor Rühle, Anoop Kulkarni, Steve Kofsky, Saravan Rajmohan
  • 発表年: 2024年8月(v2: 2025年1月改訂)
  • 分野: cs.DC(分散・並列・クラスタコンピューティング), eess.SY(システム・制御)

背景と動機(Background & Motivation)

LLMの推論処理は、プロンプト全体を並列処理するprefillフェーズと、トークンを逐次生成するdecodeフェーズに分かれる。prefillはGPU演算リソースを集中的に消費し、decodeはメモリ帯域幅がボトルネックとなる。この二相性は、Azure OpenAIのModel Routerのようなルーティングシステムの設計において重要な設計考慮点である。

著者らは、既存のロードバランサーがLLMワークロードをモノリシックな単一ジョブとして扱い、これらのフェーズ特性を無視している問題を指摘している。たとえば、ラウンドロビンや最短キュー方式では、あるインスタンスにprefill重いリクエストが集中し、別のインスタンスがアイドル状態になる非効率が発生しうる。

さらに著者らは、個々のインスタンスのスケジューラ最適化よりも、インスタンス間のクロスインスタンス負荷分散のほうがレイテンシ改善効果が大きいことを実験的に示しており、これはAzure OpenAIのAPIレベルでのルーティング設計が重要であることの学術的裏付けとなっている。

主要な貢献(Key Contributions)

  • 貢献1: LLM推論ワークロードのprefill/decodeフェーズ特性を定量的に分析し、クロスインスタンス負荷分散の優位性を実証
  • 貢献2: 応答長予測器(Response-Length Predictor)を組み込んだワークロード認識型ルーターの設計
  • 貢献3: ワークロード混合影響推定器(Workload-Mixing Impact Estimator)による、異種ワークロード混在時のパフォーマンス劣化の定量化手法
  • 貢献4: ヒューリスティック誘導型強化学習(Heuristic-Guided RL)フレームワークの提案と、公開ベンチマーク・本番データでの有効性検証

技術的詳細(Technical Details)

ワークロード特性の二相性

LLM推論は以下の2フェーズで構成される。

Prefillフェーズ: 入力プロンプト全体を並列に処理してKVキャッシュを構築する。計算量はプロンプト長$n$に対して$O(n^2 \cdot d)$($d$はモデル次元数)であり、GPU FLOPSがボトルネックとなる。

Decodeフェーズ: トークンを自己回帰的に1つずつ生成する。各ステップの計算量は$O(n \cdot d)$だが、KVキャッシュの読み出しが必要なためメモリ帯域幅がボトルネックとなる。

この二相性により、同一インスタンス上でprefill重いリクエストとdecode重いリクエストが混在すると、リソース競合によるパフォーマンス劣化が発生する。

強化学習ベースルーターの設計

著者らは、ルーティング問題を以下のマルコフ決定過程(MDP)として定式化している。

\[\text{MDP} = (\mathcal{S}, \mathcal{A}, P, R, \gamma)\]

ここで、

  • $\mathcal{S}$: 状態空間 — 各インスタンスのキュー長、現在のフェーズ構成(prefill/decode比率)、GPU利用率
  • $\mathcal{A}$: 行動空間 — ルーティング先インスタンスの選択($K$台のインスタンスから1台を選択)
  • $P$: 状態遷移確率
  • $R$: 報酬関数 — エンドツーエンドレイテンシの負値 $R = -\text{latency}(r)$
  • $\gamma$: 割引率

応答長予測器(Response-Length Predictor)

ルーティング決定時にリクエストの処理コストを事前推定するため、著者らは応答長予測器を導入している。この予測器は入力プロンプトの特徴量から予測応答トークン数$\hat{L}$を推定する。

\[\hat{L} = f_\phi(\text{features}(x))\]

ここで$x$は入力プロンプト、$f_\phi$は学習済み予測モデルである。予測された応答長はdecodeフェーズの所要時間推定に直結し、ルーティングの品質を左右する。

ワークロード混合影響推定器

異なるタイプのワークロードが同一インスタンス上で混在する際のパフォーマンス影響を定量化するコンポーネントである。たとえば、長いprefillリクエストが処理中のインスタンスに短いdecodeリクエストを送った場合の遅延増加を推定する。

これにより、ルーターは単純なキュー長だけでなく、現在処理中のワークロード構成との適合性を考慮したルーティングが可能になる。

ヒューリスティック誘導型強化学習

純粋なRL学習は探索空間が広く収束が遅いため、著者らはヒューリスティック(最短キュー方式など)をRLの初期方策として組み込む手法を採用している。これにより、学習の安定性と収束速度を改善しつつ、ヒューリスティック単体を超える性能を実現している。

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
class IntelligentRouter:
    """LLMワークロード認識型インテリジェントルーター"""

    def __init__(self, num_instances: int, predictor: ResponseLengthPredictor):
        self.num_instances = num_instances
        self.predictor = predictor
        self.policy = RLPolicy(state_dim=num_instances * 3, action_dim=num_instances)

    def route(self, request: LLMRequest) -> int:
        """リクエストを最適なインスタンスにルーティング"""
        predicted_length = self.predictor.predict(request.prompt)
        instance_states = self._get_instance_states()
        mixing_impacts = self._estimate_mixing_impacts(request, instance_states)

        state = self._build_state(instance_states, predicted_length, mixing_impacts)
        action = self.policy.select_action(state)
        return action

    def _estimate_mixing_impacts(
        self, request: LLMRequest, states: list[InstanceState]
    ) -> list[float]:
        """各インスタンスにリクエストを追加した場合の影響度を推定"""
        impacts = []
        for state in states:
            current_phase_mix = state.prefill_ratio
            new_phase = "prefill"  # 新規リクエストはprefillから開始
            impact = self._compute_contention(current_phase_mix, new_phase)
            impacts.append(impact)
        return impacts

実装のポイント(Implementation)

著者らの論文(16ページ、10図)から読み取れる実装上の注意点を以下に整理する。

応答長予測器の精度とルーティング品質のトレードオフ: 予測器が高精度であるほどルーティング品質は向上するが、予測器自体のレイテンシがルーティングのオーバーヘッドとなる。著者らはこのトレードオフを考慮し、軽量な予測モデルを採用している。

RL方策の学習安定性: LLM推論クラスタは動的な環境であり、ワークロードパターンは時間帯によって変化する。ヒューリスティック誘導により初期方策の品質を担保しつつ、オンラインで方策を更新する設計が重要である。

スケーラビリティ: インスタンス数$K$が増加すると行動空間も拡大する。著者らは状態表現の工夫により、$K$に対する線形スケーラビリティを実現していると考えられる。

フェーズ状態の観測: 各インスタンスのprefill/decode比率をリアルタイムで取得する必要があり、推論エンジン(vLLM等)との連携が前提となる。

Production Deployment Guide

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

本論文のインテリジェントルーターをAWS上で実装する場合のパターンを示す。

規模月間リクエスト推奨構成月額コスト概算主要サービス
Small~3,000 (100/日)Serverless$50-150Lambda + Bedrock + DynamoDB
Medium~30,000 (1,000/日)Hybrid$300-800Lambda + ECS Fargate + ElastiCache
Large300,000+ (10,000/日)Container$2,000-5,000EKS + Karpenter + EC2 Spot

Small構成の詳細 (月額$50-150):

  • Lambda: ルーティングロジック実行(1GB RAM, 30秒タイムアウト, $20/月)
  • Bedrock: Claude 3.5 Haiku推論(Prompt Caching有効, $80/月)
  • DynamoDB: ルーティング状態・応答長予測キャッシュ(On-Demand, $10/月)
  • CloudWatch: レイテンシ監視($5/月)

Large構成の詳細 (月額$2,000-5,000):

  • EKS: ルーター + 推論ワーカー管理($72/月 コントロールプレーン)
  • EC2 Spot Instances: g5.xlarge × 2-4台(平均$800/月, 最大90%削減)
  • Karpenter: GPU需要に応じた自動スケーリング
  • Bedrock Batch: 非リアルタイム処理で50%割引活用

コスト試算の注意事項: 上記は2026年7月時点のAWS ap-northeast-1(東京)リージョン料金に基づく概算値です。実際のコストはトラフィックパターン、リージョン、バースト使用量により変動します。最新料金は AWS料金計算ツール で確認してください。

Terraformインフラコード

Small構成 (Serverless): Lambda + Bedrock + DynamoDB

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
module "vpc" {
  source  = "terraform-aws-modules/vpc/aws"
  version = "~> 5.0"

  name = "llm-router-vpc"
  cidr = "10.0.0.0/16"
  azs  = ["ap-northeast-1a", "ap-northeast-1c"]
  private_subnets = ["10.0.1.0/24", "10.0.2.0/24"]

  enable_nat_gateway   = false
  enable_dns_hostnames = true
}

resource "aws_iam_role" "lambda_router" {
  name = "lambda-llm-router-role"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action    = "sts:AssumeRole"
      Effect    = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
    }]
  })
}

resource "aws_iam_role_policy" "bedrock_invoke" {
  role = aws_iam_role.lambda_router.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect   = "Allow"
      Action   = ["bedrock:InvokeModel", "bedrock:InvokeModelWithResponseStream"]
      Resource = "arn:aws:bedrock:ap-northeast-1::foundation-model/anthropic.claude-3-5-haiku*"
    }]
  })
}

resource "aws_lambda_function" "router" {
  filename      = "router.zip"
  function_name = "llm-intelligent-router"
  role          = aws_iam_role.lambda_router.arn
  handler       = "index.handler"
  runtime       = "python3.12"
  timeout       = 60
  memory_size   = 1024

  environment {
    variables = {
      DYNAMODB_TABLE      = aws_dynamodb_table.router_state.name
      ENABLE_RL_ROUTING   = "true"
      RESPONSE_PREDICTOR  = "lightweight"
    }
  }
}

resource "aws_dynamodb_table" "router_state" {
  name         = "llm-router-state"
  billing_mode = "PAY_PER_REQUEST"
  hash_key     = "instance_id"

  attribute {
    name = "instance_id"
    type = "S"
  }

  ttl {
    attribute_name = "expire_at"
    enabled        = true
  }
}

運用・監視設定

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
import boto3

cloudwatch = boto3.client('cloudwatch')

cloudwatch.put_metric_alarm(
    AlarmName='router-latency-p99',
    ComparisonOperator='GreaterThanThreshold',
    EvaluationPeriods=2,
    MetricName='Duration',
    Namespace='AWS/Lambda',
    Period=300,
    Statistic='p99',
    Threshold=30000,
    AlarmDescription='ルーター応答P99レイテンシ異常'
)

cloudwatch.put_metric_alarm(
    AlarmName='bedrock-token-spike',
    ComparisonOperator='GreaterThanThreshold',
    EvaluationPeriods=1,
    MetricName='TokenUsage',
    Namespace='AWS/Bedrock',
    Period=3600,
    Statistic='Sum',
    Threshold=500000,
    ActionsEnabled=True,
    AlarmActions=['arn:aws:sns:ap-northeast-1:123456789:cost-alerts'],
    AlarmDescription='Bedrockトークン使用量異常'
)

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

  • ~100 req/日 → Lambda + Bedrock (Serverless) - $50-150/月
  • ~1000 req/日 → ECS Fargate + Bedrock (Hybrid) - $300-800/月
  • 10000+ req/日 → EKS + Spot Instances (Container) - $2,000-5,000/月
  • EC2: Spot Instances優先(最大90%削減、Karpenter自動管理)
  • Reserved Instances: 1年コミットで最大72%削減
  • Bedrock Batch API: 50%割引(非リアルタイム処理)
  • Prompt Caching: 30-90%削減(システムプロンプト固定)
  • Lambda: メモリサイズ最適化(CloudWatch Insights分析)
  • ECS/EKS: アイドルタイムのスケールダウン(夜間0台)
  • AWS Budgets: 月額予算設定(80%で警告、100%でアラート)
  • CloudWatch アラーム: レイテンシ・トークン使用量スパイク検知
  • Cost Anomaly Detection: 自動異常検知
  • 日次コストレポート: SNS/Slackへ自動送信
  • 未使用リソース削除: Trusted Advisor活用
  • タグ戦略: 環境別(dev/staging/prod)でコスト可視化
  • ライフサイクルポリシー: S3古いキャッシュ自動削除(30日)
  • 開発環境: 夜間停止(Auto Start/Stop)
  • モデル選択: 開発時Haiku ($0.25/MTok)、本番Sonnet ($3/MTok)
  • トークン数制限: max_tokens設定で過剰生成防止
  • セキュリティ: IAM最小権限、KMS暗号化、Secrets Manager使用

実験結果(Results)

著者らは公開データセットおよび実運用クラウドプロバイダーのワークロードデータで評価を実施している。

評価環境提案手法のレイテンシ削減率
公開ベンチマーク(混合データセット)11%以上
実運用クラウドプロバイダーデータ7.8%

著者らは、クロスインスタンス負荷分散が個別インスタンスのスケジューラ最適化よりもレイテンシ改善効果が大きいことを実験的に確認している。これは、インスタンス間のワークロード偏りがレイテンシの主要因であり、各インスタンス内の処理順序最適化だけでは解消できないことを示唆している。

実運用データ(7.8%改善)が公開ベンチマーク(11%改善)よりも改善幅が小さい理由として、実運用環境ではワークロードの多様性がより高く、予測精度の影響が顕在化することが考えられる。

実運用への応用(Practical Applications)

本論文の知見は、Azure OpenAIのModel Routerの設計思想と直接的に関連する。Model Routerはプロンプトの複雑さやタスク種別に基づいてモデルを選択するが、本論文のアプローチはさらに低レベルのインスタンスレベルでのワークロード特性を考慮したルーティングを提案している。

Azure OpenAIとの接続点:

  • Model RouterのBalanced/Cost/Qualityモードは、本論文の「タスク特性に応じた最適化」と共通する設計思想を持つ
  • PTU(Provisioned Throughput Units)のSpillover設計は、本論文のワークロード混合影響推定と類似した問題を扱っている
  • APIM AI Gatewayのバックエンドロードバランサーに、本論文の知見を応用することで、より精密なルーティングが実現可能

プロダクション適用の考慮点:

  • 応答長予測器の学習には、対象ドメインのワークロードログが必要
  • RL方策のウォームスタートにはヒューリスティック(ラウンドロビン等)が有効
  • prefill/decodeフェーズ情報の取得には推論エンジンとの連携が必要

関連研究(Related Work)

  • RouteLLM (Ong et al., 2024): プロンプトの難易度に基づいてLLMを選択するルーティング手法。本論文が「どのインスタンスに送るか」を扱うのに対し、RouteLLMは「どのモデルに送るか」を扱う
  • Orca (Yu et al., 2022): LLM推論のイテレーションレベルスケジューリングを提案。本論文の「インスタンス内スケジューリングよりクロスインスタンスルーティングが重要」という知見の比較対象
  • vLLM (Kwon et al., 2023): PagedAttentionによるメモリ効率化を実現した推論エンジン。本論文のルーターはvLLMのようなエンジン上でのデプロイを想定

まとめと今後の展望

本論文は、LLM推論のprefill/decodeフェーズ特性を活用したインテリジェントルーティングの有効性を実証した研究である。ヒューリスティック誘導型RLと応答長予測器の組み合わせにより、既存手法を上回るレイテンシ削減を達成している。

Azure OpenAIのModel RouterやAPIM AI Gatewayのような商用サービスにおいても、本論文で提案されたワークロード認識型ルーティングの設計原則は適用可能であり、今後の負荷分散設計における重要な学術的基盤となっている。

参考文献

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

NeurIPS 2024論文解説: HippoRAG — 海馬記憶理論に基づくナレッジグラフ×PPRのRAGフレームワーク

論文解説: StepChain GraphRAG — 質問分解×BFS推論によるマルチホップQA精度向上