Baseten、生産環境向けエージェントカーネル開発の課題と解決策を報告

近年、エージェントはアイデアの創出からゼロからのカーネル生成に至るまで、驚くほど高度なカーネル開発能力を発揮するようになりました。KernelBench といった既存のベンチマークにより、エージェントが独立した汎用問題においていかに効果的にカーネルを最適化できるかを評価することが容易になっています。

しかし、ベンチマークで勝利することと、実際のプロダクション環境へ最適化を適用してリリースすることは、まだ大きな隔たりがあります。

その理由はいくつかあります。

最適なカーネル設定は、本番環境でのワークロードに依存します。一般的なベンチマークで勝つカーネルが、特定のデプロイメントでは敗北することもあります。タイル形状やワープ特化戦略、CTA 構成といった最適化手法は、テンソル形状やバッチサイズ、シーケンス長の変化に対して異なる反応を示します。特に MoE や Attention のようなカーネルにおいて、この傾向は顕著です。

マイクロベンチマークで速いからといって、必ずしもモデル全体が高速化するとは限りません。変更を統合すると、CUDA グラフのキャプチャやマルチストリーム実行といった下流の依存関係との相互作用により、カーネルレベルでの性能向上が相殺されたり、場合によっては性能低下を招いたりする可能性があります。

個々のカーネルを最適化しただけでは、より高次の改善機会を見逃す恐れがあります。エンドツーエンドのトレース解析では、実際にはわずかな数のカーネルだけが大幅な改善余地を持っていることが明らかになるケースが少なくありません。より手軽に成果を得るためには、それらのカーネルを中心に計算構造を再構築することが有効です。具体的には、演算の融合や冗長処理の排除、パイプラインのバブル解消などが挙げられます。

本番環境のサービングエンジンに新しいカーネルを組み込むのは容易ではありません。スタンドアロンの PyTorch モデルを修正するのと異なり、サービングエンジンは相互接続された実行パスと依存関係を持っています。新規カーネルは正しいパスに正しく接続され、既存の計算をきれいに置換し、周囲のランタイム環境との互換性を維持する必要があります。

こうした課題を踏まえ、ベンチマークと実運用のギャップを埋めるソリューションを開発しました。このフレームワークは、モデルとサービングエンジンが与えられれば、フルワークロードのプロファイリングを行い、最適な最適化策を推論した上で、生成されたカーネルを実環境へ直接デプロイできます。

スタック：モデルレベルとカーネルレベルの最適化

最適化スタックは以下の2層に分かれます。

✕

最適化パイプラインは2つのトラックで構成されています。一つ目はモデルレベルの作業で、ワークロードのプロファイリングを行い、エンドツーエンドのレイテンシ改善につながる候補変更をランク付け・検証します。二つ目はカーネルレベルの作業で、さまざまなバリエーションにわたるカーネルをベンチマークします。これら2つの成果はエンジン統合と実運用へとつながります。点線で示されたテストループ（マイクロベンチマーク、正しさの確認、アブレーションテスト）によって、候補がアーカイブされるか、あるいは行き止まりとして記録されるかが決定されます。

• モデルレベルの最適化: フルモデルワークロードを把握し、時間の消費箇所を特定した上で、融合処理や冗長作業の排除といった変更案を提案します。
• カーネルごとの最適化: トレースから抽出された生成済みおよび性能が重要なカーネルを対象に、複数の実装を並列で探索し、最も有力な候補に対して反復改良を行います。

最初のレイヤーは、カーネルごとの改善に限定されない探索空間の拡大を可能にします。単独でカーネルを最適化するのではなく、実行グラフを再構築することで、冗長な処理の削除、中間結果の即時生成の削減、あるいは基礎となるカーネルを生成・改善する前の演算結合などを実現できます。

最適化ランニング間での学習

当フレームワークには自己改善機能も備わっています。正しさとエンドツーエンドのパフォーマンスチェックに合格したカーネルは再利用可能な候補として保持され、成功と失敗の両方の試行から得られた教訓は、負荷条件や統合に関する知見とともに、進化し続けるナレッジベースに追加されます。

これにより、各最適化イテレーションが蓄積された経験を出発点とする自己改善ループが形成され、エージェントは時間とともにより強力な候補を生成できるようになり、収束も加速します。

永続的なナレッジアーキテクチャ：成功した最適化は、パッチ、テストケース、ベンチマーク、負荷データとともにカーネルデータベースに保存され、次の最適化ループへとフィードされます。成功と失敗の両方の最適化事例が要約されてナレッジベースに取り込まれ、成功は再利用可能なパターンとして、失敗は注意すべき点や根本原因として記録されます。

結果とケーススタディ

最初の実験では、B300 GPU で SGLang を実行する Qwen-Image と FLUX.2 という拡散モデルを対象にしました。以下に示す最適化は、すべて当社のエージェントフレームワークによって特定・提案・実装されたものです。

4 つのモデル構成（FLUX.2 FP8、FLUX.2 NVFP4、Qwen-Image FP8、Qwen-Image NVFP4）における 1 ステップあたりのノイズ除去時間の中央値。左から右へ進むにつれてすべてのモデルが高速化しており、特に Qwen FP8 の改善幅が大きく、245.6 ms から 141.8 ms に短縮されました。

両モデルでの最適化

最適化 #1：事前パックされた FP8 スケール

Qwen-Image と FLUX.2 の FP8 パスでは、行列乗算前にスケールメタデータを DeepGEMM が要求する形式に変換するための起動が繰り返され、リソースが無駄になっていました。定数である重みスケールも、小さなカーネル起動の連鎖を通じて何度も再パックされていました。

これを解消するため、主要な FP8 活性化プロデューサーを直接パックされたスケールを出力するように変更し、重みとスケールのパッキング処理をモデル読み込み時に移管しました。数値計算自体は変更していないため、出力結果もビット単位で同一に保たれています。

例えば、FLUX.2 のアテンション投影（attention projection）における最適化前後の比較は以下の通りです。

Baseline

Optimized

同様に、Qwen-Image のフィードフォワード層（feed-forward layer）でも以下の結果が得られました。

Baseline

Optimized

この最適化により、Qwen-Image ではエンドツーエンドのレイテンシが 7.3% 短縮され、FLUX.2 では 6.1% の改善が見られました。これらの効果は、その後の FP8 最適化においても持続しました。

最適化 #2: QKV 投影とエピローグの融合（Fused QKV projection and epilogue）

両モデルとも、元のアテンションパスでは画像クエリ（query）、キー（key）、バリュー（value）の投影をそれぞれ独立して計算していました。これらはすべて同じ入力を使用しているにもかかわらずです。

その結果、アテンションブロック全体で活性化値の量子化（activation quantization）や GEMM 設定が繰り返される非効率な状態が生じていました。

最適化では、3 つの FP8 投影を1つのGEMMに統合し、バイアス加算、QK正規化、RoPEを融合させた上で、単一のTritonエピローグで結合された画像・テキスト注意用バッファへ書き込みます。一方、NVFP4では各投影が異なるスケールを使用するため、Q、K、VのGEMMは個別に実行されます。

✕

ベースライン

✕

最適化 #3: 正規化と量子化カーネルの融合

両方のモデルにおいて、従来の正規化処理は大きな BF16 テンソルを生成しており、その直後に続く量子化カーネルがそれを再度読み込んでいました。この問題を解決するため、これら二つの処理を単一のカーネルに融合させました。これにより、中間的な BF16 データの書き込みと読み出しという往復動作を不要にし、パフォーマンスを向上させています。

ベースライン

✕

最適化済み

Qwen-Image では、融合カーネルが元の BF16 結果と、QKV およびフィードフォワード GEMM のための事前量子化された FP8 アクティベーションの両方を出力します。これによりレイテンシが 4.3% 短縮され、パッケッドスケール最適化で利用されるプロデューサーパスが実現されます。

FLUX.2 の残差パスでは、融合カーネルが正規化された出力、更新された残差値、パックされた E2M1 値、そしてスイズルされた E4M3 スケールを 1 サイクルで出力します。これによりエンドツーエンドのレイテンシが 0.7% 改善されます。

Qwen-Image

最適化 #1: バイアスの吸収

前の最適化後、アテンションとフィードフォワードの出力投影の直後に、2 つの独立したバイアス加算が残っています。これらは Qwen の FP8 ステップ時間の約 11% を占めています。これを解消するため、各バイアスを次の融合演算（残差正規化スケールおよび残差更新）に折り込みます。その結果、レイテンシが 5.2% 短縮されました。

最適化 #2: CFG モデュレーションキャッシュ

Classifier-free guidance (CFG) では、同じタイムステップで 2 回のデノイザーパスが実行されます。各パスは異なる条件付けを使用します（片方はプロンプト cc を受け取り、もう片方は空またはネガティブなプロンプト ∅ を受け取ります）。以前のの実装では、両方のパスでタイムステップに依存する画像およびテキストのモデュレーションブランチを再計算していました：

ϵcond=F(xt,t,c),ϵuncond=F(xt,t,∅)

ϵCFG=ϵuncond+w(ϵcond−ϵuncond)

ノイジーな潜在変数 xtx_t と時間ステップ tt は、両方のパスで共有されます。画像とテキストのモジュレーションブランチは、プロンプトではなく、時間ステップ埋め込みと固定されたモデルパラメータのみを関数として依存しています：

et=Embed(t)e_t = \text{Embed}(t)

したがって：

mimage=Wimageet+bimage,mtext=Wtextet+btextm_{\text{image}} = W_{\text{image}} e_t + b_{\text{image}}, \qquad m_{\text{text}} = W_{\text{text}} e_t + b_{\text{text}}

これらのモジュレーションブランチは ete_t と固定された重みのみに依存するため、同じ時間ステップにおける条件付きパスと非条件付きパスで出力が同一になります。これによりキャッシュが可能となります。隠れ状態やアテンションなどプロンプトに依存する出力は個別に計算されます。

DiT エントリでキャッシュキーを作成（両方の CFG ブランチに同じ timestep オブジェクトが渡されます）：

1def _cfg_cache_optimization_enabled(active) -> bool:
2    return active
3
4# QwenImageTransformer2DModel.forward
5if _cfg_cache_optimization_enabled() and isinstance(timestep, torch.Tensor):
6    # Both CFG branches receive the same timestep tensor.
7		# Keep a reference to it so tensor identity can be used safely.
8		cache_key = {
9	    "timestep": timestep,
10	    "version": version_if_available(timestep),
11		}

各ブロック内で画像とテキストのモジュレーション出力をキャッシュする：

1# QwenImageTransformerBlock.forward
2cached = getattr(self, "modulation_cache", None)
3
4cache_hit = (
5    cache_key is not None
6    and cached is not None
7    and cached["timestep"] is cache_key["timestep"]
8    and cached["version"] == cache_key["version"]
9)
10
11if cache_hit:
12    # Second CFG pass: reuse the cached outputs.
13    image_modulation = cached["image_modulation"]
14    text_modulation = cached["text_modulation"]
15
16else:
17    # First CFG pass: compute the modulation outputs.
18    image_modulation = image_modulation_GEMM(timestep_embedding)
19    text_modulation = text_modulation_GEMM(timestep_embedding)
20
21    # Cache them for the second CFG pass.
22    if cache_key is not None:
23        self.modulation_cache = {
24            "timestep": cache_key["timestep"],
25            "version": cache_key["version"],
26            "image_modulation": image_modulation,
27            "text_modulation": text_modulation,
28        }

これにより、FP8 では遅延が 2.1%、NVFP4 では 3.1% 削減されます。

オプティマイゼーション #3: カーネルごとの最適化

次に、パフォーマンスクリティカルなカーネルおよび以前融合されていたカーネルに対して最適化パスを実行し、以下の改善を実現しました：

これらのカーネルごとの最適化により、FP8 では遅延が 7.6%、NVFP4 では 13.4% 削減されます。

✕

FLUX.2

最適化 #1: シングルブロックにおける QK 正規化と RoPE の統合

FLUX.2 のシングルストリームトランスフォーマーブロックでは、Python の連続性ガード（contiguity guard）が結合された GEMM ビューを拒否したため、本番環境で使われる融合された QK 正規化および RoPE カーネルは採用されていませんでした。その結果、QK RMSNorm とインタリーブされた RoPE が別々のパスとして実行され、コサインとサインのキャッシュが繰り返し結合されるという非効率な処理が発生していました。

これを置き換えたのが、トークンごとの CTA（Cooperative Thread Array）カーネルです。この新しいカーネルは、連続した 12 KB の Q/K ヘッドタイルをそれぞれ読み込み、FP32 で RMSNorm を実行し、その結果を BF16 に丸めた上で、同じパス内でインタリーブされた RoPE を適用します。これにより、コサインとサインのテンソルを直接参照できるようになり、ステップあたり 60 回のキャッシュ結合のうち 48 回を不要にすることができました。

ベースライン

最適化済み

この融合カーネルにより、2 倍の速度向上が実現し、FP8 ではエンドツーエンドのレイテンシが 2.3%、NVFP4 では 4.0% 改善されました。

最適化 #2: SwiGLU の統合と FP8/NVFP4 量子化

SwiGLU の各呼び出しでは、以前は大きな BF16 の中間値が生成され、それを別の FP8 または NVFP4 量子化カーネルが読み取って出力投影を行っていました。また、一部のシングルブロックでは、アテンション特徴量と SwiGLU の出力を結合するための追加操作も実行されていました。

この最適化により、従来の多段階処理を単一の融合カーネルに置き換え、前述の全工程を一括で実行可能になりました。

Baseline

Optimized

FP8 の場合、本番環境と完全に一致させるには、元の演算順序を維持する必要があります。具体的には、除算を用いて SiLU を計算し、BF16 に丸め込み、BF16 状態で乗算を行い、保存された BF16 の結果から FP8 スケールを導出します。

NVFP4 の場合、この融合パスは直接、下流の FP4 GEMM で必要とされるパッキング済みの E2M1 値とスイズル処理された E4M3 スケールを出力します。これにより、中間的な BF16 の書き込み・読み出しラウンドトリップや連続コピー、複数の個別カーネル起動が不要になります。

この融合カーネルにより、FP8 では遅延が 2.3% 削減され、NVFP4 では 3.8% 削減されました。

Optimization #3: Gated residual normalization

FLUX.2 の残差パスでは、従来ゲート付き残差更新と層正規化を別々の操作として実行していました。以前の生産環境スタックは FLUX.2 のゲートをサポートしていなかったため、モデルは非融合パスのまま運用されていました。

Baseline