I looked into how the prefill router in NVIDIA NeMo Switchyard works
This page has been translated by machine translation. View original
Introduction
Hello, I'm Morishige from Classmethod's Manufacturing Business Technology Division.
In NVIDIA's technical blog, NeMo Switchyard's routers were organized into 4 types. Alongside LLM classifier, stage router, and escalation router, there's one learning-based prefill router listed. However, it hasn't made it into the release version yet, and you can't find it among the route types you can write in your local routes.toml. Just what does "prefill" mean here, and how is it different from the judge I'm currently using? Some of you may have been wondering about this.
When I tracked down the repository, I found that implementation hadn't stalled — it had been rebuilt 3 times over the past month. When it arrives, it won't be as simple as adding one line prefill_router to routes.toml; you'll need correctness labels for your own workload, an evaluation set, and a place to store your encoder. In my previous judge article, I closed by writing that the judge role would eventually be folded into the prefill of the weak execution layer (this is an article as of 2026-08-16).
This article introduces the concept and mechanics of the prefill router, its differences from the generative judge, the current state of implementation on the Switchyard side, and what to prepare for when it arrives (with actual measurements I made and compared at home as verification). I hope it resonates with those who want to cultivate the router's judgment role to fit their own workload.
Does a router's judgment role "read and reason" or "get it right at the point of reading"?
The job of an LLM router is to read an incoming request and decide "is a cheap model enough, or does an expensive one need to be called?" Before getting into the prefill router discussion, let me briefly organize what the various routers out there are looking at to make their decisions. A 2026 survey (arXiv 2603.04445) organizes the design space along 3 axes: "when to decide (before request, during execution, after response)," "what to look at (features of the request text, model metadata, past performance)," and "how to decide (rules, classifiers, reinforcement learning, cascade)." But implementations that are actually circulating generally fall into the following 5 categories.
| Method | What it looks at to decide | Representative examples |
|---|---|---|
| Rule / metadata-based | Token count, keywords, cost/latency/budget track record. Doesn't read the content | LiteLLM's cost / latency / usage-based routing, complexity router |
| Semantic similarity-based | Closeness between the request embedding and example sentences prepared for each use case | aurelio's semantic-router, vLLM Semantic Router (ModernBERT intent classification) |
| Learning-based (predict from request) | A classifier trained to predict "which model will answer correctly" from the features of the request | RouteLLM (trained on Chatbot Arena preference data), Not Diamond (the backend of OpenRouter Auto), Fireworks' FireRouter (a trained model scores difficulty, routing between closed model pass-through and open model redirect via 1–5 preferences. research preview), NVIDIA LLM Router v1 / v2 |
| LLM-as-a-judge-based | A judge LLM reads the request and generates a category or probability | Switchyard's LLM classifier, OpenRouter Fusion's judge |
| Cascade / escalation-based | First lets a cheap model solve it, then escalates to a higher tier based on quality judgment or signs of struggle | FrugalGPT, Switchyard's escalation router and stage router |
The prefill router is a type of learning-based approach in this table, one that uses the LLM's internal state as features instead of text embeddings. Switchyard's 4 routers also fit this framework, with the LLM classifier being the LLM-as-a-judge type, stage and escalation being cascade/escalation types, and the prefill router being the learning type. What I'm comparing in this article is the LLM-as-a-judge type I currently use and the learning type that is coming in.
Learning-based approaches differ by "what they learn from"
Even under the umbrella of "learning-based," the underlying thinking differs quite a bit. If you break it down, there are 3 axes.
The first is what to learn from — that is, where the labels come from. RouteLLM learns from preference data where humans chose which answer they preferred on Chatbot Arena, learning "the probability that the strong model wins." Not Diamond and NVIDIA's LLM Router learn "whether this model can solve it" from evaluation scores or correctness for each candidate model. FireRouter uses a model trained on request difficulty scores, and vLLM Semantic Router learns the category of intent or complexity of the request. Learn from preferences and you select "the model that produces answers humans prefer"; learn from correctness and you select "the model that actually gets it right"; learn from difficulty and you select "whether the request can be handled by a cheap model" — even within the same learning-based type, what the router knows differs.
The second axis is what to use as features. Fine-tuning BERT or DeBERTa, searching for similar past cases using request embeddings, scoring by the dot product of model and request embeddings (matrix factorization), having a small LLM classify, and reading the LLM's internal state. RouteLLM builds and compares all four of these using the same preference data. The choice of features determines how the system generalizes to unknown requests and the cost per judgment call.
The third axis is what it outputs and how it decides. Some output a single binary probability and cut at a threshold (RouteLLM, FireRouter, Switchyard's judge), some simultaneously output probabilities for each candidate and select "the cheapest within tolerance below the highest probability" (NVIDIA LLM Router), some decide based on a utility combining probability and cost (Switchyard's previous prefill router PR), and some match a category to a lookup table (Semantic Router).
| Router | What it learns from (labels) | What it uses as features | What it outputs and how it decides |
|---|---|---|---|
| RouteLLM (LMSYS) | Human preferences from Chatbot Arena (which answer they preferred) | 4 variants: Elo weighted by similarity, matrix factorization, BERT classifier, small LLM classifier | Single probability that strong model wins → binary decision at threshold |
| Not Diamond (behind OpenRouter Auto) | Evaluation scores of responses per candidate model | Meta-model that scores each candidate from the request (internals not public) | Score per candidate → selection by quality and cost |
| FireRouter (Fireworks, research preview) | Request difficulty | Custom trained model (internals not public) | Single difficulty score → binary decision via 1–5 preference |
| vLLM Semantic Router | Category of request intent / complexity | ModernBERT classifier | Category → select model or inference presence via lookup table |
| NVIDIA LLM Router v1 / v2 | Correctness per candidate model (scored by LLM judge) | Embedding model output (v2 also uses CLIP for images) + MLP | Correct probability per candidate → select cheapest via tolerance |
| prefill router (LLM Router v3, being implemented in Switchyard) | Correctness per candidate model | Internal state of encoder LLM during prefill + PCA + shallow MLP | Correct probability per candidate → tolerance, or Switchyard's threshold |
Looking at this table, the prefill router is the combination of "learns from correctness, uses internal state as features, outputs per-candidate probabilities." Unlike RouteLLM, it learns from whether a question was solved rather than human preferences; unlike FireRouter, it estimates the whole pool rather than a binary choice; and unlike LLM Router v1/v2, it reads the LLM's internal representations rather than text embeddings. What matters in the measurements in the latter half is the first axis. It was where the correctness labels came from, rather than which encoder was used, that determined the quality of judgment.
The first is what today's Switchyard calls the generative LLM classifier. It passes a judgment prompt and the request text to a judge LLM, has it generate "the probability p_solve that a weak model will complete this task in one shot" as JSON, and compares it against a threshold to route. In the configuration I'm running, the judge went through DeepSeek V4 Flash and was then replaced by the LoRA SFT-trained Nemotron 3.5 Lightning I set up recently. Its strength is that judgment rationale (crux, rule ID, etc.) remains in the JSON, so you can read back "why it was routed to strong" afterward. The way to run it in Switchyard v0.2.0 is summarized in a previous article (this is an article as of 2026-08-12).
The second is the prefill router. This one does not generate judgment text. The process of an LLM reading a prompt and producing the first token is called prefill, and the internal vectors that are updated as they pass through each layer during that process are called hidden states (residual stream). The prefill router uses these hidden states as features to estimate, with a small neural network, "the probability that each candidate model will answer this task correctly." The LLM only reads — it doesn't write the answer. In other words, it's a method of reading out numerically the "sense of difficulty" at the point of finishing reading the problem.
Placing the two side by side, the difference boils down to "what you make the LLM do for judgment."
| Aspect | Generative judge (LLM classifier) | Predictive prefill router |
|---|---|---|
| LLM's job | Reads judgment prompt, reasons, generates verdict JSON | Only reads the request text (prefill only) |
| Output | Verdict JSON with rationale (single p_solve) | Correct probability per candidate model (whole pool simultaneously) |
| Cost per judgment | Time and tokens for generation (takes seconds even without reasoning) | 1 prefill pass + MLP (hundreds of milliseconds, no reasoning tokens) |
| Judgment variability | Varies with temperature or model updates since it generates | Same input produces same output (deterministic) |
| Where judgment rules live | Written in the prompt (capability card) | Embedded in the training data labels |
| Adapting to your use case | Rewrite the prompt, or post-train the judge | Train MLP on correctness labels from your workload (runs on CPU) |
| Handling closed models | Judge infers from knowledge | Can predict correctness of closed models too from a different encoder (described later) |
| Weaknesses | Optimistic about difficult questions that are hard to assess, runs amok on format errors | Probabilities go wrong outside training distribution, can't read conversation context (input is request text) |
Laid side by side, all the weaknesses of the generative type stem from the same place. The single point that you're making the judge write text. Because it writes, you wait over 10 seconds for just the judgment before a single character of the answer comes back, and it stretches to 17 seconds when you increase the conversation history passed in. Because it writes, just reordering the JSON keys in the training data causes 18% to come back with malformed output. Because it writes, just a small update to the judge model causes the rate of routing to weak — which had been 39 out of 50 conversations — to drop to just 1, with almost everything falling to the expensive model side. These are failures I encountered one by one in the judge article and Switchyard article. The prefill router eliminates that "writing."
What does the prefill router solve?
So what do you gain from eliminating it? There are 4 benefits I can read from the paper "LLM Router: Rethinking Routing with Prefill Activations" (arXiv 2603.20895) and NVIDIA's implementation.
The first is making judgment cost a fixed overhead. The encoder's prefill runs only once, and in the paper's measurements, the added cost for routing was 0.76% of the total, with latency increasing only 0.79%. This is orders of magnitude different from the generative type that pays reasoning tokens on every judgment. For a judgment role that runs on every request, this matters.
The second is the ability to handle a pool of more than 2 candidates. If it's a binary choice between weak and strong, a judge is sufficient, but once the pool grows to 5 or 9 models, it becomes difficult to make a generative type answer "which is the cheapest that's sufficient" in a single judgment. The prefill router outputs the correct probability for each candidate model simultaneously, so the tolerance concept touched on in the previous foundational article — "set as the passing range everything within tolerance below the highest probability, and choose the cheapest within that" — can be used as-is (this is an article as of 2026-06-21).
The third is that closed models can be included in the pool. The core of the paper is a separation called Encoder-Target Decoupling, which takes the position that the model that outputs features (encoder) and the target model whose correctness is being predicted (target) can be different entities. It reports that using Qwen3.5 122B as the encoder can predict the correctness of closed models like Claude Opus and GPT better than the target model's own internal state. The encoder is allowed to be an "outsider." Counter-intuitive as it may seem, this is what makes a learning-based router viable for pools that mix cloud and local models.
The fourth is the ability to adapt to your own workload. Only the final small MLP and the normalization and PCA before it are trained, while the encoder stays frozen, so training runs even on CPU. What you need is correctness labels of "whether weak solved this request," and a few hundred to a few thousand examples are sufficient. The main result of the paper is recovering 45.6% of the gap between the best single model and oracle on a pool of 11 frontier models while reducing cost by 74.3%.
Let me also write what it doesn't solve. The paper itself states that accuracy drops 11.43 points on out-of-distribution problems (HLE held-out), and a learning-based router is a tool that only delivers value when trained on the distribution of your own workload. Also, the input is the request text rather than the conversation, so judgments like noticing "it's stuck" from tool execution results are outside its scope. That remains the job of the generative type or the escalation router.
The mechanism comes down to "which layer to read, and what"
Even saying "read from internal state to get probabilities," each layer of an LLM has vectors of thousands of dimensions. The design of the prefill router is about what to extract from there and how.
First is how to select layers. LLMs are said to represent vocabulary and syntax in shallow layers and meaning and task nature in deeper layers, and the paper selects "the layers that best separate correct from incorrect" using Fisher discriminant ratio. My own measurements showed the same tendency: among the 52 layers of Nemotron 3.5 Lightning, the layers that best separated correct and incorrect answers using the last token's hidden states were middle layers 27–31. Both too shallow and too deep led to worse performance.
Next is pooling. Since the request is a token sequence, you choose per layer whether to "take the state of the last token" or "average all tokens." The research team's latest configuration (merged into llm-router v3 at the end of July 2026) concatenates all 40 layers of Qwen3.6-35B-A3B with token mean (81,920 dimensions), reduces to 200 dimensions with PCA, and passes through a shared MLP. The default configuration I tried in the previous article selected the best single layer from the upper half of Qwen3.5-0.8B's layers (this is an article as of 2026-06-21).
Last is the MLP. A small 2-layer network called SharedTrunkNet outputs the correct probability for each candidate model simultaneously from the post-PCA features. It's an understated design that trains 10 seeds and averages the 5 with the best validation loss. Since only this MLP and PCA have trainable parameters, it's resistant to overfitting even with a few thousand labels and finishes in minutes on CPU.
Switchyard doesn't "not have it yet" — it's being rebuilt
What's curious here is the state of implementation on the Switchyard side. The route types available in the released v0.2.0 are passthrough, random, llm_classifier, and stage_router — no prefill router. However, when you line up the PRs chronologically, you find that implementation is in the middle of being rebuilt for the 3rd time.
| PR | Period | Status | Contents |
|---|---|---|---|
| #97 | 7/20 → merged 7/22 | Branch for old architecture (components-v2) | prefill-probe profile. All-layer mean from vLLM Qwen3.6 probe → PCA-200 → 5-model ensemble. Compares weak / strong 2 heads with cost-based utility |
| #140 | 7/24 → closed 8/21 | Port to v1 components | Same content into Rust request processor. Split into 3 on maintainer's direction |
| #174 / #185 / #186 | 7/28 → closed 8/20 | Classifier, vLLM hidden-state probe, server settings | "Close as stale, keep for reference" |
| #506 | Opened 8/20 | New crates/prefill-router |
Defines only the extraction layer contract PrefillForward. PCA, MLP, checkpoint, and routing intentionally not implemented |
The noteworthy part is the design of #506. This PR has a parity test that makes the extraction layer match tensor-by-tensor with PrefillExtractor from a specific commit (8a9d3509) of NVIDIA LLM Router v3. It's a declaration that the Python implementation on the llm-router side is the ground truth, covering how to apply the chat template, which layers to select, pooling, and how to handle batches and padding. In other words, if you're building your own prefill router, if you prepare features using the same extraction function from the same commit, when the upstream adds the scoring layer, you'll only need to swap the checkpoint. This article's verification was built on that premise.
I built a predictive type with the same labels and compared it to the generative judge
Let me verify my understanding so far with hands-on measurements. For training data, I use the same data from the judge post-training. It's the same composition as judge v2: 167 LiveCodeBench questions created from actual correctness labels of the weak model (DeepSeek V4 Flash-0731), and 1,608 synthetic task questions distilled from the current judge's verdicts. Labels are the probability p_solve that "weak completes it in one shot," and I aim to predict with the same definition for the predictive type. Evaluation also uses the same test problems as in the judge article: a held-out 237 cases not used in training, 14 triage questions where weak failed twice in a row but strong passed, and 47 single-shot prompts flowing through the router.
I prepared two metrics. Since the decision agreement rate in the judge article used the threshold rule with per-category adjustments, I use AUC as the primary metric independent of threshold in the comparison table, and for decision I re-aggregate from saved judge judgments using the p_solve-only 0.80 rule. Also, the 210 synthetic held-out cases have an average of 6.85 siblings in the training set from the same generation batch, and they can be matched on writing style alone, so I measure the ability to discern difficulty in the LCB side with actual labels and the 14 triage questions.
I lined up 4 encoders with different characteristics. The research team's configuration with Qwen3.6-35B-A3B alone couldn't be loaded with HF Transformers within the GB10's 128GB, so I switched to the vLLM hidden-state extraction shown in the Switchyard documentation (the struggle story is folded below).
| Encoder | Role | Layers × hidden | Extraction time for 2,074 cases | Per case |
|---|---|---|---|---|
| Qwen3.5-0.8B | llm-router v3 default. Minimal configuration that runs on CPU | 24 × 1024 | 2 min | 0.04 sec |
| Qwen3.6-35B-A3B NVFP4 (via vLLM) | Research team's configuration. Model the Switchyard PR uses for probing | 40 × 2048 | 18 min | 0.51 sec |
| Nemotron 3.5 Lightning | The weak execution layer candidate itself. Hypothesis of folding judge into prefill | 52 × 2688 | 34 min | 0.98 sec |
| judge v2-r8 (Lightning SFT) | Reading the internal state of the model set up as generative judge | 52 × 2688 | 34 min | 0.98 sec |
Here are the results. Rows are the current generative judge, the floor requiring no encoder (character n-gram TF-IDF + logistic regression), prefill routers for 4 encoders (single-layer configuration selected by toolkit default sweep), plus the Lightning full-layer configuration that is the main hypothesis.
| Judge | AUC (237) | AUC LCB 27 | decision @0.80 | Triage 14 → strong | Single 47 | Single 47 rank AUC |
|---|---|---|---|---|---|---|
| Generative judge v2-r8 NVFP4 (current) | 0.903 | 0.909 | 85.7% | 5/14 | 44/47 | 0.900 |
| Floor: TF-IDF + logistic regression | 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 (all layers 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 |
There are 3 things to read from this.
Within the training distribution, the prefill router matches the generative judge. The best configuration using Qwen3.6-35B achieved AUC 0.938, and in a paired bootstrap with 2,000 resamples of the 237 cases, the 95% confidence interval of the difference from the judge of +0.035 was [+0.002, +0.070], barely excluding 0. Other encoders had confidence intervals straddling 0, and with 237 cases, a 3–4 point difference is indistinguishable. Even plain Lightning achieved AUC 0.875 on the LCB side with the full-layer configuration, extracting a difficulty signal from internal state at the same level as the judge, and the direction of folding the judge into the weak prefill was at least viable within the training distribution.
For the 14 triage questions, where the judge only managed to route 5/14 to strong, the prefill router catches 10–12/14. This is because the generative type applies a rule — "tasks with clear specs and a verifier are solvable" — and is optimistic, while the predictive type directly copies the actual correctness labels from training data.
On the other hand, for the 47 single-shot prompts flowing through the router daily, performance is 31–38/47 — barely separating from the floor of always answering strong at 31/47. Rank AUC is 0.75–0.88, so the ability to rank by difficulty is preserved, but probabilities are compressed into the 0.4–0.6 band, causing everything to fall to the strong side at the 0.80 threshold. This directly shows the limitation the paper mentions: calibration collapses outside the training data distribution. Even applying cross-validation-selected thresholds or Platt scaling doesn't move these 47 questions.
Changing the training data composition also reveals where the difficulty signal comes from. When training the Lightning full-layer configuration on synthetic data alone, AUC is 0.911 — nearly the same as mixed — yet it routes not a single one of the 16 LCB weak questions to weak, and all 14 triage questions go to strong. It's simply returning a uniformly low probability whenever it sees the form of a coding problem, without discerning difficulty. The labels in the synthetic data are distilled from the judge's verdicts, directly copying the teacher's blind spots. The difficulty signal only comes from the 167 questions with actual labels, and the configuration of duplicating those 6 times and mixing them in was the sweet spot with the current data.
Inserting into the current Switchyard as a "fake judge"
How can we use this with the current Switchyard without waiting for the upstream route type? The v0.2.0 capability mode sends an OpenAI-compatible chat completions request to a judge, receives a verdict JSON, and compares p_solve against a threshold. This means any HTTP server that returns this JSON can sit in the judge's seat. So we set up the prefill router as a "fake judge," discarding the judge's prompt, extracting only the first user message, and returning the p_solve produced by encoder → PCA → MLP along with a fixed rule and judgment category. Since the judgment category is fixed, the route side sets threshold_step = 0.
[targets.prefill]
id = "prefill-judge"
llm_client = "node2_prefill" # shim's /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
We inserted this shim into an evaluation Switchyard and ran the same 87 judgments from the judge article through the existing harness. Of the 87 cases, the 40 production-shaped cases all have an identical opening request, so while the judge sees 40 different inputs (observing mid-conversation progress), the prefill router sees only 1 input.
| Route | Judge role | 47 single queries | 40 production-shaped cases (identical opening) | Total | Judgment latency p50 |
|---|---|---|---|---|---|
| auto-judge-w4 (production config) | Generative judge v2-r8 NVFP4 (vLLM) | 42/47 | 40/40 | 82/87 | 1.93 s |
| auto-judge-w0 (opening task only) | Same as above | 42/47 | 37/40 | 79/87 | 2.18 s |
| auto-prefill (shim) | Prefill router · Lightning sweep | 38/47 | 0/40 | 38/87 | 1.04 s |
| auto-prefill (shim) | Prefill router · judge v2-r8 sweep | 38/47 | 0/40 | 38/87 | 0.88 s |
The production judge configuration reproduced the 82/87 from the judge article exactly, and the prefill router's 47 single queries matched the same 38/47 obtained when judging directly from the script. This confirms that the shim can sit in the judge's seat. The 40 cases scoring 0/40 is because the probability for that 1 input did not reach 0.80 and fell to strong, while the judge, seeing the same opening text, could observe the most recent tool execution results and route all 40/40 to weak. Whether or not one can read the mid-conversation progress is the entire difference for those 40 cases.
For latency, the encoder forward and MLP inside the shim give p50 0.6 s, and the total judgment as seen from the router is 0.9–1.0 s. This is faster than the judge's 1.9 s, but this figure is for running 62 GB of BF16 under HF Transformers eager execution, which is a different condition from the judge running in NVFP4 at 21 GB under vLLM. This shim is a temporary measure, expected to serve its purpose until the upstream provides a route type. It has not been added to the production routes.toml.
Wall stories
Wall 1: Qwen3.6-35B-A3B cannot be loaded with HF Transformers
When loading the 70 GB BF16 weights with HF Transformers 5.15 using device_map="auto", the process stopped at over 90% completion both times.
kernel: Out of memory: Killed process 2728499 (python3) total-vm:158829808kB, anon-rss:50454500kB
The GB10 shares 128 GB between GPU and host. During loading, the process first reserved 66 GB on the GPU side, and then the host-side anon memory grew to 50 GB, hitting the total limit. Even after stopping the CPU-side training job and re-running alone, it crashed at the same point, so we concluded that this combination of model and loading path simply cannot fit within 128 GB.
The alternative we used was the vLLM path shown in Switchyard's docs/vllm-serve-hidden-state.md. By specifying extract_hidden_states as a pseudo-speculative decoding via --speculative-config and having ExampleHiddenStatesConnector write out prefill hidden states as safetensors, the NVFP4 Qwen3.6-35B-A3B (21.8 GB) worked as-is under vLLM 0.27.1, extracting 2,074 cases at 0.51 s each in 18 minutes total. There were three pitfalls: the request's hidden_states_path is ignored if allow_custom_save_path is disabled and is instead written to the default shared_storage_path; the written files are root-owned at 600 and cannot be read from the host side, requiring docker exec cat to retrieve them; and .lock files remain even after writing completes, so completion must be determined by file size stabilization rather than lock disappearance. Note that what vLLM writes out is the output of each layer, which is offset by one index from the HF side and also uses a different quantization, so this one path does not have feature tensors that match #506.
Wall 2: No CUDA kernel for NemotronH's Mamba layers
Lightning and judge v2-r8 are NemotronH with mixed Mamba layers. Our local environment had no aarch64 wheels for mamba_ssm and causal_conv1d, so Transformers ran on the pure PyTorch fallback. That was 0.98 s per case, and 34 minutes for 2,074 cases. For judgment purposes with one case at a time this was acceptable, but if the training data grows larger, it would be better to install the kernels or extract via vLLM.
Can the checkpoint we built be merged upstream?
Finally, let's verify that what we built is in a form that can be merged upstream. Cloning the #506 branch on a DGX Spark and running the parity test in crates/prefill-router as-is, the test confirming that the Rust side (Transformers embedded via PyO3) and the PrefillExtractor of llm-router 8a9d3509 produce identical hidden states across all layers passed.
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
Our local extraction simply calls that PrefillExtractor with the same commit and the same defaults, so the features used for training are the same features that the upstream Rust implementation reproduces. The checkpoint is also saved in llm-router v3 format (feature_spec and shared_once layout, version 3), and we confirmed that reading it back with the toolkit's own scoring function reproduces the training-time probabilities. Since model-router serve can read it as-is, the plan is that once Switchyard ships an implementation that reads this format, it will be a simple drop-in replacement. However, since #506 explicitly states "PCA, MLP, checkpoint loading, and routing are intentionally not included," which format becomes canonical will be decided in the next PR. Given that the #140 lineage used a custom safetensors + JSON format, the possibility of reverting to that remains.
Summary
The prefill router is a predictor that, instead of having a judge generate a judgment, estimates the probability of correctness for each candidate model from the internal state at the time the encoder reads the request. Unlike a generative judge, it pays no thinking tokens per judgment, returns the same answer for the same input, can estimate multiple pooled models simultaneously, and can target closed models since the encoder and target are separate. On the other hand, its probabilities shrink outside the training distribution, and tasks such as judging based on mid-conversation progress are outside its scope. Rather than replacing the generative type, the feel from building and comparing is that the division of labor will likely be: predictions within the distribution that run every time go to the predictive type, while judgments outside the distribution or based on mid-conversation progress go to the generative type or escalation router.
I think there are three things to prepare in anticipation of Switchyard incorporating this. The first is correctness labels indicating whether weak actually solved the problem in your own workload. The difficulty signal can only come from empirically measured labels; inflating the count with synthetic data distilled from judge judgments did not cultivate the ability to discriminate. The second is a frozen evaluation set not used in training, structured so that inside and outside the training distribution can be measured separately. The third is where to place the encoder: if you want to run prefill locally with the Qwen3.6-35B-class that the research team uses, you need to decide on the operational approach including vLLM hidden-state extraction. Since #506 has fixed the contract that extraction aligns with llm-router v3's PrefillExtractor, building features with the same function means that once a scoring layer arrives upstream, you should be able to deploy just by swapping in the checkpoint.
Next, we plan to add production-traffic request shapes to the training data to see if the calibration drift on single-prompt queries can be fixed, and once a scoring layer arrives in PRs following #506, to try loading this checkpoint as-is and running it.
Reference links
- LLM Router: Rethinking Routing with Prefill Activations (arXiv 2603.20895)
- Route AI Agent Workloads Across Models with NVIDIA NeMo Switchyard — Official overview of 4 routers
- NVIDIA-AI-Blueprints/llm-router — v3 branch. This article uses commit
8a9d3509 - NVIDIA-NeMo/Switchyard PR #506 — Extraction contract for
crates/prefill-router - NVIDIA-NeMo/Switchyard —
docs/vllm-serve-hidden-state.mdcontains vLLM startup instructions for probing - Dynamic Model Routing and Cascading for Efficient LLM Inference: A Survey (arXiv 2603.04445) — A survey organizing the router design space along 3 axes
- RouteLLM: An Open-Source Framework for Cost-Effective LLM Routing (LMSYS) — A representative example of a trained router that learns from preference data
- FireRouter routing preferences (Fireworks AI Docs) — A managed router (research preview) that routes binary choices using a trained scorer + preferences 1–5
- vLLM: Hidden State Extraction
- Trying out NeMo Switchyard v0.2.0 reborn in Rust — Calibration and operation of the capability classifier (article as of 2026-08-12)
- Building Nemotron 3.5 Lightning into an LLM Router judge via LoRA post-training — The source of labels and test questions used in this article (article as of 2026-08-16)
