スケーリング則への信念は、モデルパラメータとデータ規模での絶え間ない突破を駆り立てるだけでなく、同時にインフラストラクチャーエンジニアリングの限界をも絶えず押し広げています。このプロセスには不可避の痛みが伴い、それを私たちは「Scaling Pain(スケーリングの痛み)」と呼んでいます。
大規模モデルの応用が単純な対話から、より複雑で長期的なコーディングAIエージェントタスクへと全面的に移行するにつれ、私たちの推論インフラストラクチャーはかつてないプレッシャーに直面しており、毎日数億回ものコーディングAIエージェントの呼び出しを処理しています。ここ数週間、一部のユーザーがGLM-5シリーズモデルを使用して複雑なコーディングAIエージェントタスクを実行する際に、文字化け、同じ内容の繰り返し(復唱)、そして時折発生する珍しい漢字(生僻字)など、複数種類の異常に遭遇していました。これらの問題は標準的な推論環境では存在せず、高同時実行、長文脈のコーディングAIエージェントシナリオ下でのみトリガーされ、安定して再現することが困難でした。
私たちは数週間にわたる推論、原因特定、そしてストレステストを経て、最終的に相互に独立した複数の低レベルの競合状態バグ(レースコンディションバグ)を特定・修正し、それによって明らかになったシステムのボトルネックに対して的を絞った最適化を実施し、推論システムの安定性と効率を大幅に向上させました。
私たちはこの探求の中で得た経験と教訓を皆さまと共有し、コーディングAIエージェント推論のScaling Painを共に克服したいと思います。
オフライン再現から異常検知へ
3月以降、私たちはGLM-5のオンラインモニタリングとユーザーフィードバックから、以下の三種類の異常現象を確認しました:文字化け(ガーブルドアウトプット)、復唱(リピティション)、そして珍しい漢字(レアキャラクター)です。これらの現象は表面的には長文脈シナリオでよく見られる「知能低下」と似ていますが、モデル精度を低下させるいかなる最適化も導入していなかったため、より重要な問題は次の点でした:異常の原因はモデル自体にあるのか、それとも推論パイプラインにあるのか?モデルに原因がある場合、異常は特定の入力に対する安定的で再現可能な動作として現れるはずです。逆に、異常がシステム負荷や実行時の状態と相関している場合、推論インフラストラクチャー内のパイプラインや状態管理の問題を示している可能性が高くなります。
原因特定の初期段階では、まずユーザーから報告された不良事例をローカルでリプレイし、同じ一連のリクエストを数百回繰り返し推論させましたが、異常はまったく再現されず、問題はおそらくモデル自体ではないことを示していました。オンライン環境の負荷をさらに模擬するために、オンラインログを匿名化処理し、元の同時実行分布とリクエスト順序を可能な限り保持した上で、ローカルで全量リプレイを実施しました。最初は異常が再現されませんでしたが、さらにPrefill-Deocde分離(PD分離)比率を調整し、システム負荷を継続的に引き上げ、ピーク時のPrefill滞留とDecode側のKVキャッシュへの負荷を模擬したところ、ようやく約1万リクエストごとに3~5回の異常が安定して再現されました。この「リクエスト内容とは無関係で、システム負荷と相関する」という特徴は、問題が高負荷下の推論状態管理に起因する可能性を示唆しています。それと同時に、オフラインで再現された異常の頻度は、オンラインで報告された頻度を依然として下回っており、既存の検出方法には見落としがあるか、あるいはまだカバーしきれていないトリガーシナリオが存在する可能性を示していました。
異常な出力をいかに確実に識別するかが新たな課題となりました。三種類の異常のうち、復唱は比較的検出が容易でしたが、文字化けと珍しい漢字は厄介でした。正規表現や文字セットマッチングなどのヒューリスティックな手法を試したり、モデルによる判別方式も試しましたが、前者には明らかな見逃しや誤検出が発生し、後者は大規模なアブレーション実験の効率要件を満たすことが困難でした。これらの制約により、異常検出自体が原因特定プロセスにおけるボトルネックとなっていました。
図1:投機的サンプリング指標は異常検知の重要な参考情報となり得る
推論ログを繰り返し分析した結果、私たちは予想外の切り口を発見しました:投機的サンプリング(スペキュラティブデコーディング)の指標が、異常検知の重要な参考情報として利用できるということです。投機的サンプリングは本来、パフォーマンス最適化技術であり、最初にドラフトモデルで候補トークンを生成し、その後ターゲットモデルで検証・採否を決定することで、最終的な出力分布を変更することなくデコード効率を向上させます。図1に示すように、私たちは二つの指標(spec_accept_length:ターゲットモデルが連続して受け入れたドラフトトークンのプレフィックス長、spec_accept_rate:ドラフトトークンが受け入れられた割合)が、異常発生時に安定したパターンを示すことを観察しました:
- 文字化けと珍しい漢字:通常、極めて低いspec_accept_lengthを伴います。つまり、ドラフトモデルが生成した候補トークンのほぼ全てがターゲットモデルによって拒否され、ターゲットモデルが見ているKVキャッシュ状態とドラフトモデルの予期する状態との間に顕著な差異があることを示します。
- 復唱:通常、高いspec_accept_rateを伴います。これは、破損したKVキャッシュが、アテンションパターンを劣化させ、生成プロセスを高信頼度の繰り返しループに追い込んでいる可能性を示します。
上記の観察に基づき、私たちはさらにオンラインの異常監視戦略を実装しました:spec_accept_lengthが1.4未満、かつ生成された長さが128トークンを超える場合、またはspec_accept_rateが0.96を超える場合、システムは現在の生成を能動的に中止し、リクエストをロードバランサーに渡して再試行します。この戦略により、投機的サンプリングは単なるパフォーマンス最適化技術から、出力品質のリアルタイム監視シグナルへと拡張され、その後のアブレーション実験における重要なツールとなりました。
BugFix#1:PD分離アーキテクチャにおけるKVキャッシュの競合状態
異常出力と同時実行負荷に明らかな相関性があることを確認した後、私たちはその原因をさらに分析しました。リクエストのライフサイクルと、推論エンジンにおけるPD分離実行のタイミングを分析した結果、この問題はリクエストライフサイクルと、KVキャッシュの回収・再利用タイミングの不一致に起因し、それによって引き起こされるKVキャッシュ再利用時の競合であることが判明しました。
1. 原因分析:非同期Abort(中断)によって引き起こされるKVキャッシュ再利用の競合状態
テールレイテンシを制限するために、私たちは推論エンジンにタイムアウトベースのリクエスト終了メカニズムを導入しています。Prefillフェーズが規定時間内に完了しない場合、Decode側はリクエストに対してAbort(中断)を実行し、そのリクエストが占有していたKVキャッシュリソースを回収します。しかし、このAbortシグナルがPrefill側に正しく伝播されず、同時にDecode側もKVキャッシュを安全に回収・再利用できるかどうかを判断するための十分な情報を持っていませんでした。その結果、DecodeがAbortを実行し、対応するKVキャッシュ領域を新しいリクエストに割り当てた後も、以前に開始されたRDMA書き込みや実行中のPrefill計算が継続され、同期的にキャンセルされていませんでした。
図2:PD分離シナリオにおけるKVキャッシュ競合状態の概念図
図2は、PD分離アーキテクチャにおいて、二つのリクエストがPrefillとDecode間で相互作用するタイミング関係と、それによって引き起こされるKVキャッシュの競合状態を示しています。
初期段階で、Req1がPrefill-1 (P1) とDecode (D) に送信されます。スケジューリングやキューイングなどの理由により、Req1はP1側で一定時間待機した後にPrefill Forwardの実行を開始します。一方、Decode側は対応するKVキャッシュデータを一定期間受信できず、タイムアウトメカニズムがトリガーされ、Req1に対してAbortを実行します。
その後、Decode側はReq1が占有していたKVキャッシュスロットを回収しますが、P1には正しく通知しません。続いて、新しいリクエストReq2が到着し、Prefill-2 (P2) と Decodeに割り当てられます。メモリ再利用戦略により、Req2はReq1と同じKVキャッシュアドレスに割り当てられます。P2はPrefill ForwardとKV転送の実行を開始し、より短時間でこれを完了するため、Decode側は生成フェーズに入ります。
それと同時に、P1側でReq1に対して開始されたKVキャッシュへの書き込みは継続されており、そのデータはReq2によって再利用されたGPUメモリ領域に書き込まれ、Req2のKVキャッシュの一部を上書きします。最終的に、Req2はDecodeフェーズで上書きされたデータを読み取り、生成結果の異常を引き起こします。
2. 修正案:KVキャッシュ解放におけるタイミング整合性の保証
上記の競合状態を解消するため、私たちは推論エンジンにより厳格なタイミング制約を導入し、リクエスト終了とKVキャッシュ書き込み完了の間に明示的な同期関係を確立しました。
具体的には、DecodeがAbortをトリガーした後、Prefill側に通知を送信します。Prefillは、以下の条件が満たされた場合にのみ「解放可能」シグナルを返します:関連するRDMA書き込みがまだ開始されていない、または全てのコミット済み書き込みが完了している。Decodeはこの確認応答を受信した後にのみ、対応するKVキャッシュスロットの回収と再利用を許可します。このメカニズムにより、KV書き込みがGPUメモリ再利用の境界を越えることがなくなり、リクエスト間でのKVキャッシュの上書きを防止します。
修正効果:この修正の適用後、異常出力の発生率は約1万回に十数回程度から1万回に3回未満へと低下しました。この結果は、PD分離アーキテクチャにおいて、ノード間のデータ転送とGPUメモリ再利用に対して明確な整合性制約を確立し、類似の問題を回避する必要があることを示しています。
BugFix#2: HiCacheローディングタイミングの欠落
コーディングAIエージェントのシナリオは、入力長を顕著に増大させ(平均70Kトークン超)、同時に高いプレフィックス再利用率を伴います。この種のワークロードにより、HiCache(多階層KVキャッシュ)はオンラインサービスにおける重要な最適化手段となります。しかし、KVキャッシュのスワップインと計算が重複実行されるケースにおいて、既存の実装はデータが使用前にロードを完了していることを保証しておらず、準備が整っていないKVキャッシュがアクセスされる可能性がありました。
1. 原因分析:パイプライン同期の欠落によるread-before-ready(準備完了前の読み取り)
HiCacheの実行タイミング分析を通じて、私たちは問題をDSA HiCacheのキャッシュ読み取りパスに特定しました。システムはCPUメモリから履歴プレフィックスキャッシュを非同期でスワップイン(swap-in)し、Load StreamとForward Streamの重複実行によってスループットを向上させます。
図3(a)に示すように、Load StreamはKVキャッシュとIndexerキャッシュのロードを担当し、Forward StreamはIndex計算と後続のSparse Attentionを順次実行します。理論上、Forward Stream内のIndexer計算は、対応するIndexerキャッシュのロードが完了した後にのみ開始されるべきです。しかし、元の実装では、この依存関係が明示的に表現されていませんでした。
具体的には、Indexer演算子の起動時に、Load Indexer Cacheの完了に対する同期制約が確立されていませんでした(図3の赤い点線領域)。そのため、Forward StreamがLoad Streamによるデータロード完了よりも先に実行を開始し、結果として「Read-before-Ready」アクセスパターン(データが完全にロードされる前に読み取られる状態)が発生する可能性がありました。
この問題により、Index計算が不完全または未初期化のデータに基づいて実行され、後続のSparse Attentionの計算結果に影響を及ぼし、最終的に出力異常として顕在化します。
図3:HiCache読み取りパイプラインのタイミング異常と修正の概念図
2. 修正案:演算子パイプラインの原子性の再構築
上記の問題を解決するために、私たちはHiCacheの読み取りパイプラインを修正しました(図3(b) 参照)。データロードと計算の間に明示的な同期制約を導入しました:
- 明示的な同期制約:Indexer演算子の起動前に、Load Streamとの同期ポイントを導入し、該当する階層のIndexerキャッシュのロードが完了していることを確認します。Forward Streamはデータ準備が整った後にのみ計算を開始し、これによりread-before-readyのアクセスを回避します。
この修正の適用後、同一負荷条件下において、実行タイミングの不整合に起因する異常は完全に消失し、システムの挙動は安定化しました。この修正はPull Request #22811を通じてSGLangコミュニティに提出されました。
最適化:KVキャッシュ階層型ストレージ LayerSplit
上記の二つの競合状態問題は、共通するシステムボトルネックを明らかにしました:長文脈のコーディングAIエージェントサービングシナリオにおいて、Prefillフェーズがシステムパフォーマンスを支配しているということです。
Prefillの待ち行列によるTTFT(Time To First Token)を抑制するためにタイムアウトAbortを導入し、Prefill側のKVキャッシュ容量逼迫を緩和するためにHiCacheを導入しました。これらの状態整合性問題を修正した後、私たちはさらにボトルネックそのものに立ち返りました:どのようにしてPrefillのスループットを向上させ、Prefill側のKVキャッシュ用GPUメモリ負荷を低減するか。そのために、私たちはKVキャッシュの階層型ストレージソリューションLayerSplitを設計・実装しました。
コーディングAIエージェントのワークロードは通常、長いコンテキスト長と高いプレフィックスキャッシュヒット率を特徴とします。このシナリオでは、Prefillフェーズがシステムの主要なパフォーマンスボトルネックとなることが多いため、Context Parallel(CP)がオンラインPrefillノードの主要な並列化戦略となります。しかし、既存のSGLangオープンソース実装にはKVキャッシュの冗長保存問題があり、限られたKVキャッシュ容量がGPU計算リソース使用率の制限要因となっていました。
図4:LayerSplit、KVキャッシュ階層型ストレージソリューション
この問題に対処するため、私たちはKVキャッシュの階層型ストレージソリューション(LayerSplit)を設計・実装しました。このスキームでは、各GPUは全層のKVキャッシュを保存するのではなく、一部の層のKVキャッシュのみを保持します(図4(a)参照)。これにより、単一GPUカードのGPUメモリ占有量を大幅に削減します。
計算プロセスにおいて、異なるCPランクは図4(b)に示す方法で協調してPrefillを完了します:具体的には、ある層のKVキャッシュを保持しているランクが、Attention計算を実行する前に、その層のキャッシュを他の関連ランクにブロードキャストします。通信オーバーヘッドを削減するために、私たちはさらにKVキャッシュのブロードキャストとindexer計算を重複させるメカニズムを設計し、両者を時間的に相互に隠蔽します。最終的に、プロセス全体ではIndexer Cacheのブロードキャストによる追加オーバーヘッドのみが発生し、その規模はKVキャッシュの約1/8であるため、全体の通信コストは低く、パフォーマンスへの影響は無視できます。
図5:GLM-5.1 + LayerSplit 異なる長さにおけるスループット向上
図5は、キャッシュヒット率が90%に達する条件下で、この最適化がリクエスト長40kから120kの範囲でもたらすパフォーマンス向上を示しています。実験結果によると、システムスループットの向上幅は10%から132%の範囲であり、コンテキスト長が長くなるほど効果が顕著になります。全体として、この最適化はコーディングAIエージェントシナリオにおけるシステムの処理能力を大幅に向上させました。
まとめ
知能が真に高同時実行、長文脈のコーディングAIエージェントシナリオに突入した後、推論インフラストラクチャの課題はもはやスループット、レイテンシ、可用性だけではありません。その出力品質を維持することが極めて重要になります。スケーリング則へのあらゆる追求は、同等の強度を持つシステムエンジニアリングによって支えられなければなりません。これらの経験を共有することで、コミュニティが遠回りをするのを避け、AGIの未来を担うことができる推論インフラストラクチャを共に磨き上げる一助となれば幸いです。
謝辞
このブログでは、コーディングAIエージェントサービスにおける一連のシステムレベルの問題に関する研究を紹介しました。それには、これらの問題の再現と分析、ならびに対応する最適化案が含まれます。中科加禾(北京)科技有限公司、ならびに中国科学院計算技術研究所プロセッサチップ国家重点実験室チームの協力とサポートに感謝いたします。