JAX トレースの調査
概要
1回のキャプチャに対する7ステップの調査: 判定、整合性確認、スイムレーン、 カーネル → HLO → Python行、ルーフライン、変更点、信頼性リスト。
すべてはユーザーが提供したキャプチャから取得: <host>.xplane.pb および/または
<host>.trace.json.gz。ワークロードを再実行したり、XLAをダンプしたり、別のプロファイラを実行したりしないでください。
トレースで回答できない部分は、「信頼できないリスト」に追加します。
目標
レポート終了後、読者は以下を理解している必要があります:
- (a) ステップを高速化するための変更点（推定値とその前提条件を付記）
- (b) このトレースが属性付けをサポートできない箇所（カウント付き）
ツール
スクリプトのパスは、このスキルのベースディレクトリ（スキル読み込み時に表示されるディレクトリ）からの相対パスです。 スクラッチディレクトリから実行してください: 両方のスクリプトは出力を現在のディレクトリに書き込みます。
B=<this skill's base directory>
python3 -P $B/../analyzing-jax-launch-overhead/scripts/skip_stats.py <capture>/<host>.trace.json.gz
python3 -P $B/scripts/attribute_kernels.py <capture dir> [--program ID] [--user-re REGEX]
skip_stats.py(標準ライブラリのみ): 判定、ストリーム別時間予算、 アイドル時間中のホストイベント、および SKIP ペーパーの指標。読み方: analyzing-jax-launch-overhead スキルを参照。attribute_kernels.py(jaxおよびxprofパッケージが必要): すべての GPU カーネル → HLO 命令 → op_name → Python行。.xplane.pbのみを読み取り、 最初に./attribution/xplane/にコピーします（XProf が HLO proto をその隣に書き込むため）。kernels.csv、ops.csv、lines.csv、how.csv、summary.txt、および プログラムの HLO テキストを./attribution/に書き込みます。- exact: トレーサー自身の
hlo_op（呼び出しごとに1つのカーネルが起動される場合）。 - name: 名前がウィンドウ内の唯一の候補であるグラフカーネル。
- order: スケジュール順にマッチングされた共有カーネル、ライブラリ GEMM、コピー、コレクティブ。
- none: 候補なし。
- ウィンドウは、プログラムのスケジュールにおいて、隣接する exact カーネル間の命令で構成され、 while ループは既知のトリップカウントで展開されます。
- 検証: MiniGPT キャプチャでは、コマンドバッファをオフにした同一プログラムの実行と 100%一致したパイプラインとして、147,840 すべてのカーネルに対して同じ命令を選択しました。
- exact: トレーサー自身の
5つのテクニック（各テクニックに対応するツール）
| 分野 | テクニック | ここでのツール |
|---|---|---|
| パフォーマンスエンジニアリング | クリティカルパス; ルーフライン | skip_stats.py からのストリーム予算; 適切なピークを使用したルーフラインのみ（トラップ参照） |
| 分散トレーシング | 相関IDによるホストイベントとGPUイベントの結合; ストラグラーの特定 | skip_stats.py → rows.csv; GPU間のステップ終了時刻のばらつき |
| コンパイラ | コンパイラパスを通じたソースの来歴 | attribute_kernels.py → lines.csv |
| フォレンジック会計 | 独立したレコード間で合計が一致する必要がある | ステップ2のチェック |
| 情報設計 | GPUストリームごとに整列したスイムレーン | 1つのステップで、すべての GPU × ストリームを1つの時間軸上に表示 |
ステップ
計画は以下の7ステップで、すべて指定されたファイルに対して実行します。新しいキャプチャでのみ回答できる質問 （ホストコールの内部動作、修正の効果、バッチスイープ）は、ステップ7の「新しいキャプチャが必要」リストに追加し、 キャプチャするかどうかはユーザーの判断に委ねます。
- 判定とホスト側原因:
skip_stats.pyを実行します。 - 整合性確認. 各チェックは成立するか、信頼性リストに追加されます:
- デバイスあたりのGPUイベント数: json = xplane;
- すべてのGPUイベントは、1つのホスト起動呼び出しに結合する;
- ステップあたりのカーネル数は一定;
- すべてのGPUが同じ作業を行う（ステップあたりのカーネル数が等しいか、差異が説明されている）;
- 各ストリームについて、計算 + コレクティブ + memcpy + アイドルのシェアがスパンに合計される。
- スイムレーン. 中央値のステップについて、すべての GPU × ストリームを1つの時間軸上に配置し、 さらに起動スレッドのホストイベントも追加します。これにより、クリティカルパスとストラグラーが表示されます: どのGPUが最後に終了するか、そのGPUが何を待っているか。
- カーネル → HLO → Python行:
attribute_kernels.pyを実行します。- GPU時間における各マッチクラスのシェアと、候補なしのカーネルを報告します。
lines.csvとops.csvを行単位および演算単位のテーブルに使用します。--user-reは、どのフレームをユーザーコードとしてカウントするかを指定します。 デフォルトでは、site-packages と jax の外側の最初のフレームです。
- ルーフライン（HLO 形状から FLOP とバイト数を取得できる演算のみ）:
attribute_kernels.pyが./attribution/に書き込む<program>(<id>).hlo.txtを使用します。 演算のデータ型のピーク（ベンダー仕様から引用し、読み取り値 とラベル付け）を使用します。 トレースはピークを測定できません。 - 変更点. クリティカルパスにおける原因のシェアでランク付けします。
各原因について、それがなくなった場合のステップ時間を推定します（例: クリティカルストリーム上の step − ホストアイドル）。
推測値、未実行 とラベル付けし、ユーザーがテストするための仮説として調整パラメータを指定します
（
donate_argnums、コマンドバッファオフ、少ないまたは大きなコレクティブ）。 - このトレースが信頼できない箇所: 各項目をカウントまたはシェアとともにリストします。その後、「新しいキャプチャが必要」リストを追加します: 未解決の質問と、それを回答するためのキャプチャ。
トラップ（ROCm MI355X キャプチャで測定）
- XProf の演算ビューはグラフカーネルを除外します。
hlo_statsは、4-GPU キャプチャで GPU カーネル時間の37.6%しかカバーしませんでした。残りの62.4%（コマンドバッファカーネル）はkernel_statsに演算