feat(prefill-probe): add learned prefill complexity router

benchmark/routing-profiles/prefill-probe-local.yaml

@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Local learned prefill-probe profile. Before serving, start the Qwen probe as
+# documented in docs/vllm-serve-hidden-state.md and export:
+#   NVIDIA_API_KEY
+#   VLLM_BASE_URL=http://127.0.0.1:8000
+#   PREFILL_ROUTER_CHECKPOINT_DIR=/absolute/path/to/exported/router
+#
+# Serve with:
+#   switchyard serve --config benchmark/routing-profiles/prefill-probe-local.yaml
+# Clients select this profile by sending model: router.
+
+endpoints:
+  inference-hub:
+    base_url: https://inference-api.nvidia.com/v1
+    api_key: ${NVIDIA_API_KEY}
+    timeout_secs: 600.0
+  vllm-local:
+    base_url: ${VLLM_BASE_URL}/v1
+
+targets:
+  strong:
+    endpoint: inference-hub
+    model: aws/anthropic/bedrock-claude-opus-4-7
+    format: anthropic
+  probe:
+    endpoint: vllm-local
+    model: Qwen/Qwen3.6-35B-A3B
+    format: openai
+  weak:
+    endpoint: inference-hub
+    model: nvidia/nvidia/nemotron-3-super-120b-long-ctx
+    format: openai
+
+profiles:
+  router:
+    type: prefill-probe
+    probe: probe
+    strong: strong
+    strong_checkpoint_head: opus-4.7
+    weak: weak
+    weak_checkpoint_head: nemotron-3-super
+    hidden_states_dir: /tmp/switchyard-hidden-states
+    checkpoint_dir: ${PREFILL_ROUTER_CHECKPOINT_DIR}
+    routing_policy:
+      type: cost-aware
+      lambda: 0.5
+      weak_cost: 0
+      strong_cost: 1

crates/switchyard-components-v2/Cargo.toml

@@ -13,14 +13,17 @@ rust-version.workspace = true
 
 [dependencies]
 async-trait = "0.1"
+half = "2"
 parking_lot = "0.12"
 rand = "0.8"
 reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls-native-roots"] }
+safetensors = "0.4"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 switchyard-components = { path = "../switchyard-components" }
 switchyard-components-v2-macros = { path = "../switchyard-components-v2-macros" }
 switchyard-core = { path = "../switchyard-core" }
+tokio = { version = "1", features = ["time"] }
 toml = "0.8"
 tracing = { version = "0.1", default-features = false, features = ["std"] }
 yaml_serde = "0.10"

crates/switchyard-components-v2/src/lib.rs

@@ -25,10 +25,12 @@ pub use profiles::{
     EndpointHealth, EndpointHealthStatus, LatencyServiceProcessedRequest, LatencyServiceProfile,
     LatencyServiceProfileConfig, LlmRoutingDecision, LlmRoutingProcessedRequest, LlmRoutingProfile,
     LlmRoutingProfileConfig, LlmRoutingTierMapping, NoopProfile, NoopProfileConfig,
-    PassthroughProfile, PassthroughProfileConfig, RandomRoutingProcessedRequest,
-    RandomRoutingProfile, RandomRoutingProfileConfig, SelectedTarget, StageRouterClassifierConfig,
-    StageRouterDecision, StageRouterDecisionSource, StageRouterPickerMode,
-    StageRouterProcessedRequest, StageRouterProfile, StageRouterProfileConfig, StageRouterTier,
+    PassthroughProfile, PassthroughProfileConfig, PrefillProbeDecision,
+    PrefillProbeProcessedRequest, PrefillProbeProfile, PrefillProbeProfileConfig,
+    PrefillProbeRoutingPolicyConfig, RandomRoutingProcessedRequest, RandomRoutingProfile,
+    RandomRoutingProfileConfig, SelectedTarget, StageRouterClassifierConfig, StageRouterDecision,
+    StageRouterDecisionSource, StageRouterPickerMode, StageRouterProcessedRequest,
+    StageRouterProfile, StageRouterProfileConfig, StageRouterTier,
 };
 pub use stats::profile_stats_accumulator;
 pub use switchyard_components_v2_macros::profile_config;

crates/switchyard-components-v2/src/profiles/mod.rs

@@ -8,6 +8,7 @@ mod llm_routing;
 mod macros;
 mod noop;
 mod passthrough;
+mod prefill_router;
 mod profile_types;
 mod random_routing;
 mod stage_router;
@@ -22,6 +23,10 @@ pub use llm_routing::{
 };
 pub use noop::{NoopProfile, NoopProfileConfig};
 pub use passthrough::{PassthroughProfile, PassthroughProfileConfig};
+pub use prefill_router::{
+    PrefillProbeDecision, PrefillProbeProcessedRequest, PrefillProbeProfile,
+    PrefillProbeProfileConfig, PrefillProbeRoutingPolicyConfig,
+};
 pub(crate) use profile_types::{parse_profile_config, ProfileConfigEntry};
 pub use random_routing::{
     RandomRoutingProcessedRequest, RandomRoutingProfile, RandomRoutingProfileConfig,