NVIDIA NeMo Switchyard の prefill router の仕組みについて調べてみた
はじめに
こんにちは、クラスメソッド製造ビジネステクノロジー部の森茂です。
NVIDIA の技術ブログで NeMo Switchyard のルーターが 4 種類に整理されていました。LLM classifier、stage router、escalation router と並んで、ひとつだけ学習ベースの prefill router が載っています。ところがリリース版にはまだ入っていなくて、手元の routes.toml に書ける route type にも見当たらない。いったい何をもって prefill と呼んでいるのか、今使っている judge と何が違うのか。気になっていた方もいるのではないでしょうか。
リポジトリを追いかけてみると、実装は止まっているのではなく、この 1 か月で 3 回作り直されている最中でした。来たときに routes.toml へ prefill_router と 1 行足せば済む話ではなく、自分のワークロードの正誤ラベルや評価セット、エンコーダの置き場所が要ります。先日の judge 記事も、判定役はいずれ weak 実行層の prefill に折り畳まれていくのではないか、と書いて締めていました（2026-08-16 時点の記事です）。
この記事では、prefill router の考え方と仕組み、生成型 judge との違い、Switchyard 側の実装の現在地、そして来たときに備えて何を用意しておくかを紹介します（理解の裏取りとして、手元で作って比べた実測も添えています）。ルーターの判定役を自分のワークロードに合わせて育てていきたい人に刺さるといいなと思っています。
ルーターの判定役は「読んで推論する」か「読み込んだ時点で当てる」か
LLM ルーターの仕事は、届いたリクエストを読んで「安いモデルで足りるか、高いモデルが要るか」を決めることです。prefill router の話に入る前に、世の中のルーターが何を見て決めているのかを軽く整理しておきます。2026 年のサーベイ（arXiv 2603.04445）は設計空間を「いつ決めるか（依頼前・実行中・応答後）」「何を見るか（依頼文の特徴・モデルのメタデータ・過去の成績）」「どう決めるか（ルール・分類器・強化学習・カスケード）」の 3 軸で整理していますが、実装として出回っているものは、おおむね次の 5 つに収まります。
| 方式 | 何を見て決めるか | 代表例 |
|---|---|---|
| ルール・メタデータ型 | トークン数やキーワード、コスト・レイテンシ・予算の実績。中身は読まない | LiteLLM の cost / latency / usage ベースの振り分け、complexity router |
| 意味類似度型 | 依頼文の埋め込みと、用途ごとに用意した例文との近さ | aurelio の semantic-router、vLLM Semantic Router（ModernBERT の意図分類） |
| 学習型（依頼文から予測） | 依頼文の特徴から「どのモデルが正解するか」を学習した分類器で当てる | RouteLLM（Chatbot Arena の選好データで学習）、Not Diamond（OpenRouter Auto の裏側）、Fireworks の FireRouter（学習済みモデルが難易度をスコアし、1〜5 の preference で閉じたモデルへの pass-through とオープンモデルへの redirect を振る。research preview）、NVIDIA LLM Router v1 / v2 |
| LLM-as-a-judge 型 | judge 役の LLM が依頼文を読み、カテゴリや確率を生成する | Switchyard の LLM classifier、OpenRouter Fusion の judge |
| カスケード・昇格型 | まず安いモデルに解かせ、品質の判定やつまずきの証拠で上位へ昇格する | FrugalGPT、Switchyard の escalation router と stage router |
prefill router はこの表の学習型の一種で、特徴量をテキスト埋め込みではなく LLM の内部状態にしたものです。Switchyard の 4 ルーターもこの構成に載っていて、LLM classifier が LLM-as-a-judge 型、stage と escalation がカスケード・昇格型、prefill router が学習型にあたります。この記事で比べるのは、自分が今使っている LLM-as-a-judge 型と、これから入ってくる学習型の 2 つです。
学習型の中身は「何から学ぶか」で分かれる
ひとくちに学習型と言っても、中身の考え方はかなり違います。分けて見るなら軸は 3 つです。
1 つ目は何から学ぶか、つまりラベルの出どころです。RouteLLM は Chatbot Arena で人がどちらの答えを好んだかという選好データから「強いモデルが勝つ確率」を学びます。Not Diamond や NVIDIA の LLM Router は候補モデルごとの評価スコアや正誤から「このモデルが解けるか」を学びます。FireRouter は依頼の難易度スコアを学んだモデルで、vLLM Semantic Router は依頼の意図や複雑さのカテゴリを学びます。選好から学べば「人が好む答えを出すモデル」が、正誤から学べば「実際に解けるモデル」が、難易度から学べば「安いモデルで足りる依頼かどうか」が選ばれることになり、同じ学習型でもルーターが知っていることが違います。
2 つ目は何を特徴量にするかです。BERT や DeBERTa を微調整する、依頼文の埋め込みで類似の過去事例を探す、モデルと依頼の埋め込みの内積で点を付ける（行列分解）、小さな LLM に分類させる、そして LLM の内部状態を読む。RouteLLM はこのうち 4 つを同じ選好データで作って比べています。特徴量の選び方は、未知の依頼への一般化のしかたと、判定 1 回のコストを決めます。
3 つ目は何を出して、どう決めるかです。二択の確率を 1 つ出してしきい値で切るもの（RouteLLM、FireRouter、Switchyard の judge）、候補ごとの確率を同時に出して「いちばん高い確率から tolerance ぶんだけ下までの中で最安」を選ぶもの（NVIDIA LLM Router）、確率とコストを足し合わせた効用で決めるもの（Switchyard の過去の prefill router PR）、カテゴリを当てて対応表で振るもの（Semantic Router）があります。
| ルーター | 何から学ぶか（ラベル） | 何を特徴量にするか | 何を出して、どう決めるか |
|---|---|---|---|
| RouteLLM（LMSYS） | Chatbot Arena の人の選好（どちらの答えを好んだか） | 類似度で重み付けした Elo、行列分解、BERT 分類器、小さな LLM の分類器の 4 通り | 強いモデルが勝つ確率 1 つ → しきい値で二択 |
| Not Diamond（OpenRouter Auto の裏側） | 候補モデルごとの応答の評価スコア | 依頼文から候補ごとに点を付けるメタモデル（中身は非公開） | 候補ごとのスコア → 品質とコストで選択 |
| FireRouter（Fireworks、research preview） | 依頼の難易度 | 学習済みのカスタムモデル（中身は非公開） | 難易度スコア 1 つ → 1〜5 の preference で二択 |
| vLLM Semantic Router | 依頼の意図・複雑さのカテゴリ | ModernBERT の分類器 | カテゴリ → 対応表で推論の有無やモデルを選択 |
| NVIDIA LLM Router v1 / v2 | 候補モデルごとの正誤（LLM judge で採点） | 埋め込みモデルの出力（v2 は CLIP で画像も）+ MLP | 候補ごとの正解確率 → tolerance で最安を選択 |
| prefill router（LLM Router v3、Switchyard で実装中） | 候補モデルごとの正誤 | エンコーダ LLM の prefill 時の内部状態 + PCA + 浅い MLP | 候補ごとの正解確率 → tolerance、または Switchyard のしきい値 |
この表で見ると、prefill router は「正誤から学び、内部状態を特徴量にし、候補ごとの確率を出す」組み合わせです。RouteLLM と違って人の好みではなく解けたかどうかを学び、FireRouter と違って二択ではなくプール全体を見積もり、LLM Router v1 / v2 と違ってテキスト埋め込みではなく LLM の中の表現を読む。後半の実測で効いてくるのは、このうち 1 つ目の軸です。どのエンコーダを使うかより、正誤ラベルがどこから来たかのほうが判定の質を決めていました。
1 つ目が、今の Switchyard が LLM classifier と呼んでいる生成型です。judge 役の LLM に判定用のプロンプトと依頼文を渡し、「weak モデルがこのタスクを一発で完遂する確率 p_solve」を JSON で生成させ、しきい値と比べて振り分けます。自分が運用している構成では、judge は DeepSeek V4 Flash を経て、先日 LoRA SFT で仕立てた Nemotron 3.5 Lightning に置き換わりました。判定根拠（crux やルール ID）が JSON に残るので、あとから「なぜ strong に振ったのか」を読めるのが持ち味です。Switchyard v0.2.0 での動かし方は以前の記事にまとめています（2026-08-12 時点の記事です）。
2 つ目が prefill router です。こちらは判定文を生成しません。LLM がプロンプトを読み込んで最初のトークンを出すまでの処理を prefill と呼び、その途中で各層を通過するたびに更新される内部ベクトルを hidden states（残差ストリーム）と呼びます。prefill router はこの hidden states を特徴量にして、小さなニューラルネットで「候補モデルそれぞれがこのタスクに正解する確率」を推定します。LLM は読み込むだけで、答えを書かない。言ってみれば、問題文を読み終えた時点の「手応え」を数値にして読み出すやり方です。
2 つを並べると、違いは「判定のために LLM に何をさせるか」に集約されます。
| 観点 | 生成型 judge（LLM classifier） | 予測型 prefill router |
|---|---|---|
| LLM の仕事 | 判定プロンプトを読み、推論し、JSON を生成する | 依頼文を読み込むだけ（prefill のみ） |
| 出力 | 根拠つきの verdict JSON（p_solve は 1 つ） | 候補モデルごとの正解確率（プール全体を同時に） |
| 1 判定のコスト | 生成ぶんの時間とトークン（思考を止めても数秒） | prefill 1 回 + MLP（数百ミリ秒、思考トークンなし） |
| 判定の揺れ | 生成なので temperature やモデル更新で変わる | 同じ入力なら同じ出力（決定論的） |
| 判定ルールの置き場 | プロンプト（capability card）に書く | 学習データのラベルに埋め込む |
| 自分の用途に合わせる | プロンプトを書き換える、または judge を事後学習する | 自分のワークロードの正誤ラベルで MLP を学習する（CPU でも回る） |
| 閉じたモデルの扱い | judge が知識で推測する | 閉じたモデルの正誤も別のエンコーダから予測できる（後述） |
| 苦手なこと | 見極めにくい難問を「解ける」と楽観する、形式崩れで暴走する | 学習分布の外で確率が狂う、会話の途中経過を読めない（入力は依頼文） |
並べてみると、生成型の弱点はどれも同じところから出ています。judge に文章を書かせている、という 1 点です。書かせるから、答えを 1 文字も返す前の判定だけで 10 秒以上待たされ、渡す会話履歴を増やせば 17 秒まで伸びる。書かせるから、学習データの JSON のキーの並び順を入れ替えただけで、18% が形式崩れで返ってこなくなる。書かせるから、judge のモデルを小さく更新しただけで、50 会話のうち 39 本を weak に振れていたのが 1 本まで落ち、ほぼ全部が高いモデル側へ倒れる。judge 記事と Switchyard 記事で順に踏んできた故障です。prefill router は、その「書かせる」を無くします。
prefill router は何を解決するのか
では、無くした結果として何が手に入るのか。論文 "LLM Router: Rethinking Routing with Prefill Activations"（arXiv 2603.20895）と NVIDIA の実装から読み取れる効能は 4 つあります。
1 つ目は判定コストの固定費化です。エンコーダの prefill は 1 回で済み、論文の測定ではルーティングに足すコストが全体の 0.76%、レイテンシは 0.79% 増でした。判定のたびに思考トークンを払う生成型とは桁が違います。毎リクエスト走る判定役にとって、ここは効いてきます。
2 つ目は候補が 2 つ以上あるプールを扱えることです。weak と strong の二択なら judge でも足りますが、プールに 5 本、9 本と並ぶと、生成型では「どれが最安で足りるか」を 1 回の判定で答えさせるのが難しくなります。prefill router は候補モデルごとの正解確率を同時に出すので、以前の基礎編で触れた tolerance の発想、つまり「いちばん高い確率から tolerance ぶんだけ下までを合格圏にして、その中で最安を選ぶ」という選び方がそのまま使えます（2026-06-21 時点の記事です）。
3 つ目は、閉じたモデルもプールに入れられることです。論文の核は Encoder-Target Decoupling と呼ばれる分離で、特徴量を出すモデル（エンコーダ）と、正誤を予測される対象モデル（ターゲット）を別物にしてよい、という立場です。Qwen3.5 の 122B をエンコーダにすると、Claude Opus や GPT のような閉じたモデルの正誤まで、対象モデル自身の内部状態より良く当てられたと報告されています。エンコーダは「他人」でよい。直感に反しますが、これがあるからクラウドとローカルが混ざったプールでも学習型ルーターが成立します。
4 つ目は、自分のワークロードに合わせられることです。学習するのは最後の小さな MLP と、その手前の標準化・PCA だけで、エンコーダは凍結したままなので学習は CPU でも回ります。必要なのは「この依頼を weak が解けたか」の正誤ラベルで、数百〜数千件あれば足ります。論文の主結果は、フロンティアモデル 11 本のプールで最良単体モデルと oracle の差の 45.6% を回収しつつ、コストを 74.3% 減らしたというものです。
解決しないことも書いておきます。学習分布から外れた問題（HLE の held-out）では精度が 11.43 ポイント落ちると論文自身が書いていて、学習型ルーターは自分のワークロードの分布で学習して初めて価値が出る道具です。また入力は依頼文であって会話ではないので、ツールの実行結果を見て「詰んでいる」と気づくような判定は守備範囲の外です。そこは生成型や escalation router の仕事として残ります。
仕組みは「どの層の何を読むか」に尽きる
内部状態から確率を読む、と言っても、LLM の各層には何千次元ものベクトルが並んでいます。prefill router の設計は、そこから何をどう取り出すかの話です。
まず層の選び方です。LLM は浅い層で語彙や構文を、深い層で意味や課題の性質を表現すると言われていて、論文では Fisher 分離度で「正解と不正解を最も分けやすい層」を選んでいます。自分の実測でも同じ傾向で、Nemotron 3.5 Lightning の 52 層のうち、最終トークンの hidden states で正解・不正解を最もよく分けたのは 27〜31 層目の中層でした。浅すぎても深すぎても落ちます。
次に pooling です。依頼文はトークン列なので、層ごとに「最終トークンの状態を取る」か「全トークンを平均する」かを選びます。研究チームの最新構成（2026 年 7 月末に llm-router v3 へ合流）は Qwen3.6-35B-A3B の全 40 層をトークン平均で連結し（81,920 次元）、PCA で 200 次元に落として共有 MLP に通す形です。以前の記事で試した既定構成は、Qwen3.5-0.8B の上半分の層から最良の 1 層を選ぶ形でした（2026-06-21 時点の記事です）。
最後が MLP です。SharedTrunkNet と名付けられた 2 層の小さなネットワークが、PCA 後の特徴量から候補モデルごとの正解確率を同時に出します。10 通りの seed で学習して検証損失の良い 5 本を平均する、という地味な作りです。学習できるパラメータはこの MLP と PCA だけなので、数千件のラベルでも過学習しにくく、CPU で数分で終わります。
Switchyard には「まだ無い」のではなく「作り直し中」だった
ここで気になるのが、Switchyard 側の実装です。リリース版の v0.2.0 にある route type は passthrough、random、llm_classifier、stage_router の 4 つで、prefill router はありません。ところが PR を時系列に並べると、実装は 3 世代作り直されている最中でした。
| PR | 期間 | 状態 | 中身 |
|---|---|---|---|
| #97 | 7/20 → 7/22 マージ | 旧アーキテクチャ（components-v2）のブランチ | prefill-probe profile。vLLM の Qwen3.6 probe から全層平均 → PCA-200 → 5 本のアンサンブル。weak / strong の 2 head をコスト込みの効用で比較 |
| #140 | 7/24 → 8/21 クローズ | v1 コンポーネントへ移植 | 同じ内容を Rust の request processor に。メンテナの指示で 3 分割 |
| #174 / #185 / #186 | 7/28 → 8/20 クローズ | 分類器・vLLM hidden-state probe・server 設定 | 「stale として閉じる、参照用に残す」 |
| #506 | 8/20 オープン | crates/prefill-router を新設 |
抽出層の契約 PrefillForward だけを定義。PCA、MLP、checkpoint、ルーティングは意図的に未実装 |
読みどころは #506 の作りです。この PR は抽出層を NVIDIA LLM Router v3 の特定コミット（8a9d3509）の PrefillExtractor とテンソル単位で完全一致させる parity test を持っています。チャットテンプレートの当て方、層の選び方、pooling、バッチとパディングの扱いまで、llm-router 側の Python 実装をそのまま正とする宣言です。つまり、自分で prefill router を作るなら、同じコミットの同じ抽出関数で特徴量を作っておけば、上流がスコアリング層を載せた時点で checkpoint を差し替えるだけで済む。この記事の検証は、その前提で組みました。
同じラベルで予測型を作り、生成型 judge と比べてみた
ここまでの理解を手元で裏取りしておきます。教師データは judge 事後学習のときのものをそのまま使います。weak モデル（DeepSeek V4 Flash-0731）の実測正誤から作った LiveCodeBench 167 問と、現行 judge の判定を蒸留した合成タスク 1,608 問、採用した judge v2 と同じ配合です。ラベルは「weak が一発で完遂する確率 p_solve」で、予測型でも同じ定義のまま当てにいきます。評価も judge 記事と同じ試験問題で、学習に使わないホールドアウト 237 件、weak が 2 回連続で落ちて strong なら通った見極め 14 問、ルーター経由で流している単発プロンプト 47 問の 3 つです。
物差しは 2 つ揃え直しました。judge 記事の decision 一致率は判定カテゴリで上乗せするしきい値の則だったので、比較表では閾値に依らない AUC を主指標にし、decision は p_solve 単独 0.80 の則で judge 側も保存済みの判定から再集計しています。また合成ホールドアウト 210 件は学習側に同じ生成バッチの兄弟を平均 6.85 件持つ構造で、文体だけで当たってしまうので、難しさを見抜けたかは実測ラベルの LCB 側と見極め 14 問で測ります。
エンコーダは性格の違う 4 本を並べました。研究チームの構成である Qwen3.6-35B-A3B だけは HF Transformers でのロードが GB10 の 128GB に収まらず、Switchyard の文書が示す vLLM の hidden-state 抽出に切り替えています（壁物語に畳みます）。
| エンコーダ | 位置づけ | 層 × hidden | 2,074 件の抽出時間 | 1 件あたり |
|---|---|---|---|---|
| Qwen3.5-0.8B | llm-router v3 の既定。CPU でも回る最小構成 | 24 × 1024 | 2 分 | 0.04 秒 |
| Qwen3.6-35B-A3B NVFP4（vLLM 経由） | 研究チームの構成。Switchyard の PR が probe に使うモデル | 40 × 2048 | 18 分 | 0.51 秒 |
| Nemotron 3.5 Lightning | weak 実行層の候補そのもの。judge を prefill に折り畳む仮説 | 52 × 2688 | 34 分 | 0.98 秒 |
| judge v2-r8（Lightning SFT 済み） | 生成型 judge に仕立てたモデルの内部状態を読む | 52 × 2688 | 34 分 | 0.98 秒 |
結果です。行は現行の生成型 judge、エンコーダ不要の床（文字 n-gram の TF-IDF + ロジスティック回帰）、そしてエンコーダ 4 本の prefill router（toolkit 既定のスイープで選んだ単層構成）に、仮説本体の Lightning 全層構成を加えています。
| 判定役 | AUC（237） | AUC LCB 27 | decision @0.80 | 凍結 14 → strong | 単発 47 問 | 単発 47 問の順位 AUC |
|---|---|---|---|---|---|---|
| 生成型 judge v2-r8 NVFP4（現行） | 0.903 | 0.909 | 85.7% | 5/14 | 44/47 | 0.900 |
| 床: TF-IDF + ロジスティック回帰 | 0.783 | 0.670 | 68.4% | 12/14 | 31/47 | 0.758 |
| prefill 0.8B（L13 last PCA-50） | 0.911 | 0.744 | 82.7% | 11/14 | 37/47 | 0.753 |
| prefill Lightning（L29 last PCA-50） | 0.918 | 0.790 | 83.5% | 12/14 | 38/47 | 0.837 |
| prefill Lightning（全層 mean PCA-200） | 0.912 | 0.875 | 74.3% | 10/14 | 31/47 | 0.755 |
| prefill judge v2-r8（L32 last PCA-50） | 0.931 | 0.824 | 86.5% | 11/14 | 38/47 | 0.797 |
| prefill Qwen3.6-35B vLLM（L29 last PCA-50） | 0.938 | 0.932 | 84.0% | 11/14 | 35/47 | 0.875 |
読み取れることは 3 つです。
学習分布の内側では、prefill router は生成型 judge に並びます。最も良かった Qwen3.6-35B の構成は AUC 0.938 で、237 件を 2,000 回リサンプルした paired bootstrap でも judge との差 +0.035 の 95% 信頼区間が [+0.002, +0.070] とかろうじて 0 を外しました。他のエンコーダは信頼区間が 0 をまたぎ、237 件では 3〜4 ポイントの差は見分けがつきません。素の Lightning も全層構成で LCB 側の AUC 0.875 と judge と同じ水準の難度信号を内部状態から取り出せていて、judge を weak の prefill に折り畳むという方向は、少なくとも学習分布の内側では成立していました。
見極め 14 問は、judge が 5/14 しか strong に振れなかったところを、prefill router は 10〜12/14 で拾います。生成型は「仕様が明確で検証器があるタスクは解ける」とルールを当てて楽観するのに対し、予測型はラベルにある実測の正誤をそのまま写すからです。
一方で、ルーター経由で日常的に流している単発プロンプト 47 問では 31〜38/47 と、常に strong と答える床の 31/47 からわずかしか離れません。順位 AUC は 0.75〜0.88 あるので、難しい順に並べる力は保っているのに、確率が 0.4〜0.6 の帯に縮んで 0.80 のしきい値では全部 strong 側に倒れている形です。学習データの分布から離れると確率の較正が崩れる、という論文の限界がそのまま出ました。学習内の交差検証で選んだしきい値や Platt 較正を当てても、この 47 問は動きません。
学習データの配合を変えると、難しさの信号がどこから来ているかも見えます。Lightning の全層構成で、合成データだけで学習すると AUC は 0.911 と混合とほぼ同じなのに、LCB の weak 側 16 問を 1 問も weak に振れず、見極め 14 問は 14/14 で strong です。コーディング問題の形を見た時点で一律に低い確率を返しているだけで、難しさを見抜いてはいません。合成データのラベルは judge の判定を蒸留したもので、教師の盲点をそのまま写しています。難度の信号は実測ラベルのある 167 問からしか出ておらず、それを 6 倍に複製して混ぜる配合が今のデータでは落としどころでした。
今の Switchyard に挿すなら「偽 judge」になる
上流の route type を待たずに今の Switchyard で使うには、どうすればよいでしょうか。v0.2.0 の capability モードは、judge に OpenAI 互換の chat completions を投げて verdict JSON を受け取り、p_solve をしきい値と比べます。つまり、この JSON を返す HTTP サーバーなら judge の席に座れます。そこで prefill router を「偽 judge」として立て、judge 用のプロンプトは捨てて最初の user メッセージだけを取り出し、エンコーダ → PCA → MLP で出した p_solve を固定のルールと判定カテゴリに添えて返すようにしました。判定カテゴリを固定するので、route 側は threshold_step = 0 にします。
[targets.prefill]
id = "prefill-judge"
llm_client = "node2_prefill" # shim の /v1/chat/completions
[routes.auto-prefill]
type = "llm_classifier"
mode = "capability"
classifier_target = "prefill"
strong_target = "strong"
weak_target = "weak"
base_threshold = 0.80
threshold_step = 0.0
この shim を評価用の Switchyard に挿し、judge 記事と同じ 87 判定を既存のハーネスで流しました。87 件のうち実運用形状の 40 件は冒頭の依頼文が全件同一で、会話の途中経過を見る judge には 40 通りの入力でも、依頼文だけを見る prefill router には 1 入力です。
| route | 判定役 | 単発 47 問 | 実運用形状 40 件（冒頭同一） | 合計 | 判定レイテンシ p50 |
|---|---|---|---|---|---|
| auto-judge-w4（本番設定） | 生成型 judge v2-r8 NVFP4（vLLM） | 42/47 | 40/40 | 82/87 | 1.93 秒 |
| auto-judge-w0（冒頭タスクのみ） | 同上 | 42/47 | 37/40 | 79/87 | 2.18 秒 |
| auto-prefill（shim） | prefill router・Lightning スイープ | 38/47 | 0/40 | 38/87 | 1.04 秒 |
| auto-prefill（shim） | prefill router・judge v2-r8 スイープ | 38/47 | 0/40 | 38/87 | 0.88 秒 |
judge の本番設定は judge 記事の 82/87 をそのまま再現し、prefill router の単発 47 問はスクリプトから直接判定したときと同じ 38/47 でした。shim が judge の席に座れることの確認です。40 件が 0/40 なのは、その 1 入力の確率が 0.80 に届かず strong に倒れたからで、judge は同じ冒頭文でも直近のツール実行結果を見て 40/40 を weak に振れています。会話の途中経過を読めるかどうかが、この 40 件の差のすべてです。
レイテンシは、shim 内のエンコーダ forward と MLP で p50 0.6 秒、ルーターから見た判定全体で 0.9〜1.0 秒です。judge の 1.9 秒より速いですが、こちらは 62GB の BF16 を HF Transformers の eager 実行で動かした値で、vLLM の NVFP4 で 21GB に収めている judge とは条件が違います。この shim はつなぎで、上流が route type を出した時点で役目を終える想定です。本番の routes.toml には入れていません。
壁物語
壁 1: Qwen3.6-35B-A3B が HF Transformers で読み込めない
BF16 の重み 70GB を HF Transformers 5.15 で device_map="auto" 読み込みしたところ、2 回とも 9 割強を読んだところで止まりました。
kernel: Out of memory: Killed process 2728499 (python3) total-vm:158829808kB, anon-rss:50454500kB
GB10 は GPU とホストが 128GB を共有しています。読み込み中のプロセスは GPU 側に 66GB を先に確保したうえで、ホスト側の anon メモリが 50GB まで膨らみ、合計が上限に当たりました。CPU 側の学習ジョブを止めて単独で再実行しても同じ位置で落ちたので、このモデルと読み込み経路の組み合わせ自体が 128GB に収まらないと判断しました。
代わりに使ったのが、Switchyard の docs/vllm-serve-hidden-state.md が示す vLLM の経路です。--speculative-config で extract_hidden_states という擬似の投機デコードを指定し、ExampleHiddenStatesConnector に prefill の hidden states を safetensors として書き出させます。NVFP4 の Qwen3.6-35B-A3B（21.8GB）に全 40 層を指定したところ vLLM 0.27.1 でそのまま動き、1 件 0.51 秒で 2,074 件を 18 分で抽出できました。踏んだ穴は 3 つ。リクエストの hidden_states_path は allow_custom_save_path が無効だと無視されて既定の shared_storage_path に書かれること、書かれたファイルが root 所有の 600 でホスト側からは読めず docker exec cat で取り出す必要があること、そして .lock ファイルが書き終わっても残るので、ロックの消滅ではなくファイルサイズの安定で完了を判定することです。なお vLLM が書き出すのは各層の出力で HF 側の添字と 1 つずれ、量子化も違うので、この 1 本だけは #506 とテンソル一致の特徴量ではありません。
壁 2: NemotronH の Mamba 層に CUDA カーネルがない
Lightning と judge v2-r8 は Mamba 混在の NemotronH です。手元の環境には mamba_ssm と causal_conv1d の aarch64 向け wheel が無く、Transformers は純 PyTorch の fallback で動きます。1 件あたり 0.98 秒、2,074 件で 34 分でした。判定用途では 1 件ずつなので許容範囲でしたが、学習データを大きくするならカーネルを入れるか、vLLM 経由で抽出したほうがよさそうです。
作った checkpoint は上流に合流できるのか
最後に、作ったものが上流に合流できる形になっているかを確かめておきます。#506 のブランチを DGX Spark に clone して crates/prefill-router の parity test をそのまま走らせると、Rust 側（PyO3 で埋め込んだ Transformers）と llm-router 8a9d3509 の PrefillExtractor の hidden states が全層で一致することを確かめるテストが通りました。
test transformers::tests::matches_the_reference_transformers_tensors_exactly ... ok
test result: ok. 1 passed; 0 failed; 0 ignored; 0 measured; 1 filtered out; finished in 18.58s
手元の抽出はその PrefillExtractor を同じコミット・同じ既定値で呼んでいるだけなので、学習に使った特徴量は上流の Rust 実装が再現する特徴量と同じものです。checkpoint も llm-router v3 の形式（feature_spec と shared_once レイアウトの version 3）で保存し、toolkit 自身のスコアリング関数で読み戻して学習時の確率と一致することを確認しています。model-router serve がそのまま読める形なので、Switchyard がこの形式を読む実装を出した時点で差し替えるだけで済む見込みです。ただし #506 は「PCA、MLP、checkpoint の読み込み、ルーティングは意図的に含めない」と明記しているので、どの形式を正とするかは次の PR で決まります。#140 の系譜では safetensors + JSON の独自形式だったので、そちらに倒れる可能性も残っています。
まとめ
prefill router は、judge に判定文を生成させる代わりに、エンコーダが依頼文を読み込んだ時点の内部状態から候補モデルごとの正解確率を当てる予測器です。生成型 judge と違って判定のたびに思考トークンを払わず、同じ入力には同じ答えを返し、プールに並ぶ複数のモデルを同時に見積もれて、エンコーダとターゲットを分けられるので閉じたモデルも対象にできます。その代わり、学習した分布の外では確率が縮み、会話の途中経過を読んで昇格を判断するような仕事は守備範囲の外です。生成型と置き換わるというより、分布の内側で毎回走る判定は予測型に、分布の外や会話の途中経過を踏まえた判断は生成型や escalation router に、という分担になりそうだというのが、作って比べてみた手触りでした。
Switchyard に入ってきたときに備えて用意しておくものは 3 つだと思っています。1 つ目は、自分のワークロードで weak が実際に解けたかどうかの正誤ラベルです。難しさの信号は実測ラベルからしか出ず、judge の判定を蒸留した合成データで数を稼いでも見極めは育ちませんでした。2 つ目は、学習に使わない凍結した評価セットで、学習分布の内側と外側を分けて測れる形にしておくことです。3 つ目はエンコーダの置き場所で、研究チームの構成である Qwen3.6-35B 級をローカルで prefill させるなら、vLLM の hidden-state 抽出を含めて動かし方を決めておく必要があります。抽出は #506 が llm-router v3 の PrefillExtractor と一致させる契約を固定したので、同じ関数で特徴量を作っておけば、上流にスコアリング層が入った時点で checkpoint の差し替えだけで載せられる見込みです。
次は、実運用で流れている依頼の形を学習データに足して単発プロンプトでの較正崩れが直るかを確かめ、#506 以降の PR でスコアリング層が入ったら、この checkpoint をそのまま載せて動かしてみる予定です。
参考リンク
- LLM Router: Rethinking Routing with Prefill Activations（arXiv 2603.20895）
- Route AI Agent Workloads Across Models with NVIDIA NeMo Switchyard — 4 ルーターの公式整理
- NVIDIA-AI-Blueprints/llm-router — v3 ブランチ。本記事はコミット
8a9d3509を使用 - NVIDIA-NeMo/Switchyard PR #506 —
crates/prefill-routerの抽出契約 - NVIDIA-NeMo/Switchyard —
docs/vllm-serve-hidden-state.mdに probe 用 vLLM の起動手順 - Dynamic Model Routing and Cascading for Efficient LLM Inference: A Survey（arXiv 2603.04445） — ルーターの設計空間を 3 軸で整理したサーベイ
- RouteLLM: An Open-Source Framework for Cost-Effective LLM Routing（LMSYS） — 選好データで学習する学習型ルーターの代表例
- FireRouter routing preferences（Fireworks AI Docs） — 学習済みスコアラー + 1〜5 の preference で二択に振る managed router（research preview）
- vLLM: Hidden State Extraction
- Rust に生まれ変わった NeMo Switchyard v0.2.0 を試してみた — capability classifier の較正と運用（2026-08-12 時点の記事です）
- Nemotron 3.5 Lightning を LoRA 事後学習で LLM ルーターの判定役に仕立ててみた — 本記事のラベルと試験問題の出どころ（2026-08-16 時点の記事です）
