llm-inference-batching-scheduler

LLM 推論バッチングスケジューラー

このスキルでは、コスト、パディングの無駄、レイテンシを最小限に抑えながらリクエストをバッチにまとめる必要がある、LLM 推論のバッチングとスケジューリングの最適化問題を解くための指針を示します。

問題の理解

実装する前に、問題の構造を十分に分析してください。

制約の分析

1. すべてのハード制約を特定する - 次の正確な上限を確認します。

   • 許容される一意なシェイプの最大数（例：すべてのバケットでシェイプ数が ≤ 8）
   • レイテンシのしきい値（P95、P99）
   • コスト予算のしきい値
   • パディング比率の上限

2. ハードな上限を早い段階で計算する - コーディングする前に、次を計算します。

   • アライメント要件による最小パディング量
   • カバレッジに必要な最小バッチ数
   • 制約下で達成可能な最大効率

3. コスト関数を分解する - 各要素を理解します。

   • バッチごとのオーバーヘッド（バッチごとの固定コスト）
   • シェイプのコンパイルコスト（シーケンス長に対して二次関数となることが多い）
   • プリフィル／デコードのコスト（リクエストごとに変動）
   • 次のように記述します：Cost ≈ overhead × num_batches + shape_compile_cost + prefill_cost + decode_cost

データ分析

1. リクエストの分布をプロファイリングする - 次を調べます。

   • プロンプト長（prompt_len）の分布
   • 生成長（gen_len）の分布
   • 指標に不均衡な影響を与える可能性のある外れ値を特定します

2. カバレッジ要件を確認する - 次を満たしていることを確認します。

   • 各バケットで最大の prompt_len が、選択したシェイプでカバーされている
   • gen_len が極端な値となるエッジケースが処理されている

実装方針

再利用可能な評価基盤を構築する

パラメーターの調整を繰り返す前に、体系的な評価ハーネスを作成します。

1. Write a function that takes parameters (shape_list, gen_bucket_sizes) and returns all metrics
2. Include automatic constraint verification with assertions
3. Enable rapid parameter comparison without manual re-runs

パラメーターの探索戦略

手当たり次第の試行錯誤は避け、代わりに次の方法を使います。

1. パラメーター空間が小さい場合はグリッドサーチを行う - パラメーターの範囲が限られている場合（例：gen_bucket_size が [15, 50]）、組み合わせを体系的に評価します

2. 単一パラメーターには二分探索を使う - 他のパラメーターを固定して一つのパラメーターを最適化する場合、二分探索で最適値を見つけます

3. 最適化の探索空間を記録する - どのパラメーターの組み合わせがどの指標値を生み出すかを記録し、トレードオフを把握します

シェイプ選択の指針

シーケンス長のバケット分けに使用するシェイプを選ぶ際は、次の点に注意してください。

1. 長さの分布を分析する - 最も頻度の高い長さでパディングを最小限に抑えるシェイプを選びます
2. 2 の累乗または幾何級数を検討する - カバレッジとシェイプ数のバランスを取れることがよくあります
3. 両方のバケットを合わせて考慮する - シェイプをバケット間で共有する場合は、個別ではなく全体を最適化します

生成長のバケット分け

gen_len のバケット分けパラメーターは、パディング比率に大きく影響します。

1. バケットが小さい = パディング比率は低くなりますが、バッチ数が増えます（コスト増）
2. バケットが大きい = バッチ数は減りますが、gen_len のばらつきによるパディングが増えます
3. パディング予算を計算して逆算し、最適なバランスを見つけます

検証戦略

制約の早期チェック

出力を生成したら、すぐに次を検証します。

1. All request IDs appear exactly once
2. Number of unique shapes ≤ limit
3. Each request's prompt_len ≤ assigned shape
4. No missing shapes that would cause alignment failures

指標の検証

解決策を完成とする前に、次を確認します。

1. 提供されている場合は、公式の評価スクリプトを実行します
2. すべての指標を、すべてのしきい値と照合します
3. 各バケットを個別に確認します。一つのバケットを通過しても、他のバケットも通過するとは限りません

よくある検証の失敗

次の問題に注意してください。

• カバレッジに穴ができる原因となるシェイプの不足（例：必要なときにシェイプ 2048 がない）
• バッチごとのオーバーヘッドを無駄にする、リクエストが 1 件だけのバッチ
• バケットを個別に最適化したことで発生する、シェイプ制約への違反

よくある落とし穴

性急な最適化

• 誤り：数学的な制約を理解する前に実装に取りかかる
• 対処法：事前に正確な予算を計算する時間を取ります（例：「bucket_1 で許容されるパディングトークンは最大 25,735 個」）

コスト分析が不十分

• 間違い：どのコスト要素が支配的なのかを理解していません
• 対処法：最適化の前に、コストの内訳全体を計算して記録します

バケットを個別に最適化

• 間違い：制約が両方にまたがっているのに、各バケットを個別に最適化しています
• 対処法：特に形状の共通制約については、共同最適化を検討します

手動でのパラメーター調整の繰り返し

• 間違い：パラメーターを手動で何度も変更し、そのたびに再実行しています
• 対処法：複数の組み合わせを自動でテストするパラメーター掃引スクリプトを作成します

制約違反の発見が遅い

• 間違い：広範な反復を行った後に初めて制約を確認しています
• 対処法：出力生成の直後にアサーションチェックを追加します

ベースライン実装を無視

• 間違い：非効率な点を理解するために、提供されたベースラインコードを分析していません
• 対処法：baseline_packer.py（または同等のもの）を調べ、具体的に何が最適でないのかを特定します

最適化のトレードオフ

最適化の前に、以下のトレードオフを明示的に記録してください。

要素	減少するもの	増加するもの
形状を増やす	パディング率	形状コンパイルコスト
生成バケットを小さくする	パディング率	バッチ数（コスト）
バッチサイズを大きくする	バッチごとのオーバーヘッド	パディングの可能性
形状の間隔を狭める	パディング	必要な形状数

推奨ワークフロー

1. 分析 - 問題を解析し、制約を特定して、厳密な上限を計算します
2. 計測 - 制約チェックを自動で行う評価ハーネスを構築します
3. ベースライン - ベースラインソリューションを実行し、その欠点を把握します
4. 分解 - コストを要素に分解し、支配的な項を特定します
5. 探索 - ランダムな探索ではなく、体系的なパラメーター探索（グリッド／二分探索）を使用します
6. 検証 - すべてのバケットについて、すべての制約と指標を確認します
7. 反復 - 失敗した場合は、しきい値から最も遠い指標を特定し、その改善に注力します