From e41278bef8cef3337526820dc7633163a3fb5686 Mon Sep 17 00:00:00 2001 From: michxu Date: Thu, 16 Jul 2026 18:15:08 +0000 Subject: [PATCH 01/12] feat(prefill-probe): add learned artifact inference Signed-off-by: michxu --- Cargo.lock | 11 + crates/switchyard-components-v2/Cargo.toml | 1 + .../src/profiles/mod.rs | 1 + .../src/profiles/prefill_router/artifact.rs | 1100 +++++++++++++++++ .../src/profiles/prefill_router/mod.rs | 7 + 5 files changed, 1120 insertions(+) create mode 100644 crates/switchyard-components-v2/src/profiles/prefill_router/artifact.rs create mode 100644 crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs diff --git a/Cargo.lock b/Cargo.lock index cc34de7d..3b938e56 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1182,6 +1182,16 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safetensors" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44560c11236a6130a46ce36c836a62936dc81ebf8c36a37947423571be0e55b6" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "schannel" version = "0.1.29" @@ -1378,6 +1388,7 @@ dependencies = [ "parking_lot", "rand 0.8.6", "reqwest", + "safetensors", "serde", "serde_json", "switchyard-components", diff --git a/crates/switchyard-components-v2/Cargo.toml b/crates/switchyard-components-v2/Cargo.toml index d257a065..40aee965 100644 --- a/crates/switchyard-components-v2/Cargo.toml +++ b/crates/switchyard-components-v2/Cargo.toml @@ -16,6 +16,7 @@ async-trait = "0.1" 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" } diff --git a/crates/switchyard-components-v2/src/profiles/mod.rs b/crates/switchyard-components-v2/src/profiles/mod.rs index 91ca996b..f6d80e5d 100644 --- a/crates/switchyard-components-v2/src/profiles/mod.rs +++ b/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; diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/artifact.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/artifact.rs new file mode 100644 index 00000000..f1f3666e --- /dev/null +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/artifact.rs @@ -0,0 +1,1100 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Loading, validation, and CPU inference for learned prefill-router artifacts. + +use std::collections::{BTreeMap, BTreeSet}; +use std::path::Path; + +use safetensors::{Dtype, SafeTensors}; +use serde::Deserialize; +use switchyard_core::{Result, SwitchyardError}; + +const METADATA_FILE: &str = "router.json"; +const TENSOR_FILE: &str = "router.safetensors"; +const FORMAT_VERSION: u64 = 1; +const TRAINING_MODE: &str = "single_pca_block"; +const REPRESENTATION: &str = "token_mean_per_layer_concat"; +const PCA_DIM: usize = 200; +const TRUNK_HIDDEN: [usize; 2] = [256, 128]; +const ENSEMBLE_SIZE: usize = 5; +const OUTPUT_NAMES: [&str; 4] = ["qwen-122b", "nemotron-3-super", "opus-4.7", "gpt-5.5"]; +const PROBABILITY_LINK: &str = "independent_sigmoid"; +const ENSEMBLE_REDUCTION: &str = "probability_mean"; + +/// Immutable learned-router metadata and decoded tensors owned by a profile runtime. +pub(crate) struct InferenceArtifact { + metadata: ArtifactMetadata, + tensors: BTreeMap>, +} + +impl InferenceArtifact { + /// Loads and validates an exported artifact against the configured probe model. + pub(crate) fn load(directory: impl AsRef, probe_model: &str) -> Result { + let directory = directory.as_ref(); + let metadata_path = directory.join(METADATA_FILE); + let metadata_json = std::fs::read_to_string(&metadata_path).map_err(|error| { + invalid_artifact(format!( + "failed to read {}: {error}", + metadata_path.display() + )) + })?; + let metadata: ArtifactMetadata = serde_json::from_str(&metadata_json).map_err(|error| { + invalid_artifact(format!( + "failed to parse {}: {error}", + metadata_path.display() + )) + })?; + metadata.validate(probe_model)?; + + let tensor_path = directory.join(&metadata.tensor_file); + let tensor_bytes = std::fs::read(&tensor_path).map_err(|error| { + invalid_artifact(format!("failed to read {}: {error}", tensor_path.display())) + })?; + let tensors = { + let tensors = SafeTensors::deserialize(&tensor_bytes).map_err(|error| { + invalid_artifact(format!( + "failed to parse {}: {error}", + tensor_path.display() + )) + })?; + validate_tensors(&tensors, &metadata)?; + decode_tensors(&tensors)? + }; + + Ok(Self { metadata, tensors }) + } + + /// Returns the encoder checkpoint named by the artifact. + pub(crate) fn encoder(&self) -> &str { + &self.metadata.encoder + } + + /// Returns the number of independently trained trunk members. + pub(crate) fn ensemble_size(&self) -> usize { + self.metadata.ensemble_size + } + + /// Returns the number of output heads in checkpoint order. + pub(crate) fn output_count(&self) -> usize { + self.metadata.output_names.len() + } + + /// Returns checkpoint output names in the order used by learned probabilities. + pub(crate) fn output_names(&self) -> &[String] { + &self.metadata.output_names + } + + /// Returns the number of hidden-state layers expected from the probe server. + pub(crate) fn layer_count(&self) -> usize { + self.metadata.extraction_layer_ids.len() + } + + /// Returns the hidden width expected for each extracted layer. + pub(crate) fn hidden_size(&self) -> usize { + self.metadata.hidden_size + } + + /// Returns the flattened token-mean feature dimension. + pub(crate) fn raw_feature_dim(&self) -> usize { + self.metadata.raw_feature_dim + } + + /// Returns the number of decoded artifact tensors. + pub(crate) fn tensor_count(&self) -> usize { + self.tensors.len() + } + + /// Returns the number of decoded F32 values across all artifact tensors. + pub(crate) fn tensor_value_count(&self) -> usize { + self.tensors.values().map(Vec::len).sum() + } + + /// Applies the artifact's fitted scaler followed by its PCA-200 projection. + pub(crate) fn project(&self, raw_features: &[f32]) -> Result> { + let standardized = standardize( + raw_features, + self.tensor("transform.scaler_mean")?, + self.tensor("transform.scaler_scale")?, + )?; + project_pca( + &standardized, + self.tensor("transform.pca_mean")?, + self.tensor("transform.pca_components")?, + self.metadata.pca_dim, + ) + } + + /// Runs all learned shared-trunk members and returns logits in artifact output order. + pub(crate) fn ensemble_logits(&self, pca_features: &[f32]) -> Result>> { + if pca_features.len() != self.metadata.pca_dim { + return Err(trunk_error(format!( + "input dimension {} does not match pca_dim {}", + pca_features.len(), + self.metadata.pca_dim, + ))); + } + + let mut ensemble_logits = Vec::with_capacity(self.metadata.ensemble_size); + for index in 0..self.metadata.ensemble_size { + let prefix = format!("ensemble.{index}"); + let hidden1 = dense_layer( + pca_features, + self.tensor(&format!("{prefix}.linear1.weight"))?, + self.tensor(&format!("{prefix}.linear1.bias"))?, + TRUNK_HIDDEN[0], + true, + )?; + let hidden2 = dense_layer( + &hidden1, + self.tensor(&format!("{prefix}.linear2.weight"))?, + self.tensor(&format!("{prefix}.linear2.bias"))?, + TRUNK_HIDDEN[1], + true, + )?; + let logits = dense_layer( + &hidden2, + self.tensor(&format!("{prefix}.output.weight"))?, + self.tensor(&format!("{prefix}.output.bias"))?, + self.metadata.output_names.len(), + false, + )?; + ensemble_logits.push(logits); + } + Ok(ensemble_logits) + } + + /// Applies independent sigmoid links and averages probabilities across members. + pub(crate) fn ensemble_probabilities(&self, ensemble_logits: &[Vec]) -> Result> { + if self.metadata.ensemble_size == 0 || ensemble_logits.len() != self.metadata.ensemble_size + { + return Err(trunk_error(format!( + "logit member count {} does not match ensemble_size {}", + ensemble_logits.len(), + self.metadata.ensemble_size, + ))); + } + + let output_count = self.metadata.output_names.len(); + let mut probability_sums = vec![0.0f32; output_count]; + for (member, logits) in ensemble_logits.iter().enumerate() { + if logits.len() != output_count { + return Err(trunk_error(format!( + "member {member} logit count {} does not match output count {output_count}", + logits.len(), + ))); + } + for (sum, logit) in probability_sums.iter_mut().zip(logits) { + *sum += sigmoid_probability(*logit)?; + } + } + + let member_count = self.metadata.ensemble_size as f32; + probability_sums + .iter_mut() + .for_each(|probability| *probability /= member_count); + Ok(probability_sums) + } + + fn tensor(&self, name: &str) -> Result<&[f32]> { + self.tensors + .get(name) + .map(Vec::as_slice) + .ok_or_else(|| invalid_artifact(format!("missing decoded tensor {name}"))) + } +} + +#[derive(Deserialize)] +#[cfg_attr(test, derive(serde::Serialize))] +struct ArtifactMetadata { + format_version: u64, + training_mode: String, + encoder: String, + representation: String, + extraction_layer_ids: Vec, + hidden_size: usize, + raw_feature_dim: usize, + feature_block_count: usize, + pca_dim: usize, + pca_whiten: bool, + output_names: Vec, + trunk_hidden: Vec, + ensemble_size: usize, + probability_link: String, + ensemble_reduction: String, + tensor_file: String, +} + +impl ArtifactMetadata { + fn validate(&self, probe_model: &str) -> Result<()> { + require( + self.format_version == FORMAT_VERSION, + format!( + "unsupported format_version {}; expected {FORMAT_VERSION}", + self.format_version + ), + )?; + require( + self.training_mode == TRAINING_MODE, + format!( + "training_mode must be {TRAINING_MODE}; got {}", + self.training_mode + ), + )?; + require( + self.encoder == probe_model, + format!( + "artifact encoder {} does not match probe model {probe_model}", + self.encoder + ), + )?; + require( + self.representation == REPRESENTATION, + format!( + "representation must be {REPRESENTATION}; got {}", + self.representation + ), + )?; + require( + !self.extraction_layer_ids.is_empty(), + "extraction_layer_ids must not be empty", + )?; + let expected_layer_ids = (0..self.extraction_layer_ids.len()).collect::>(); + require( + self.extraction_layer_ids == expected_layer_ids, + "extraction_layer_ids must be contiguous and ordered from zero", + )?; + require(self.hidden_size > 0, "hidden_size must be positive")?; + let expected_raw_dim = self + .extraction_layer_ids + .len() + .checked_mul(self.hidden_size) + .ok_or_else(|| invalid_artifact("raw feature dimension overflow"))?; + require( + self.raw_feature_dim == expected_raw_dim, + format!( + "raw_feature_dim {} does not equal layer count {} * hidden_size {}", + self.raw_feature_dim, + self.extraction_layer_ids.len(), + self.hidden_size + ), + )?; + require( + self.feature_block_count == 1, + format!( + "feature_block_count must be 1; got {}", + self.feature_block_count + ), + )?; + require( + self.pca_dim == PCA_DIM, + format!("pca_dim must be {PCA_DIM}; got {}", self.pca_dim), + )?; + require(!self.pca_whiten, "pca_whiten must be false")?; + require( + self.output_names + .iter() + .map(String::as_str) + .eq(OUTPUT_NAMES), + format!("output_names must be ordered as {OUTPUT_NAMES:?}"), + )?; + require( + self.trunk_hidden == TRUNK_HIDDEN, + format!("trunk_hidden must be {TRUNK_HIDDEN:?}"), + )?; + require( + self.ensemble_size == ENSEMBLE_SIZE, + format!( + "ensemble_size must be {ENSEMBLE_SIZE}; got {}", + self.ensemble_size + ), + )?; + require( + self.probability_link == PROBABILITY_LINK, + format!( + "probability_link must be {PROBABILITY_LINK}; got {}", + self.probability_link + ), + )?; + require( + self.ensemble_reduction == ENSEMBLE_REDUCTION, + format!( + "ensemble_reduction must be {ENSEMBLE_REDUCTION}; got {}", + self.ensemble_reduction + ), + )?; + require( + self.tensor_file == TENSOR_FILE, + format!( + "tensor_file must be {TENSOR_FILE}; got {}", + self.tensor_file + ), + )?; + Ok(()) + } +} + +struct TensorSpec { + name: String, + shape: Vec, +} + +fn expected_tensors(metadata: &ArtifactMetadata) -> Vec { + let mut expected = vec![ + TensorSpec { + name: "transform.scaler_mean".into(), + shape: vec![metadata.raw_feature_dim], + }, + TensorSpec { + name: "transform.scaler_scale".into(), + shape: vec![metadata.raw_feature_dim], + }, + TensorSpec { + name: "transform.pca_mean".into(), + shape: vec![metadata.raw_feature_dim], + }, + TensorSpec { + name: "transform.pca_components".into(), + shape: vec![metadata.pca_dim, metadata.raw_feature_dim], + }, + ]; + for index in 0..metadata.ensemble_size { + let prefix = format!("ensemble.{index}"); + expected.extend([ + TensorSpec { + name: format!("{prefix}.linear1.weight"), + shape: vec![TRUNK_HIDDEN[0], metadata.pca_dim], + }, + TensorSpec { + name: format!("{prefix}.linear1.bias"), + shape: vec![TRUNK_HIDDEN[0]], + }, + TensorSpec { + name: format!("{prefix}.linear2.weight"), + shape: vec![TRUNK_HIDDEN[1], TRUNK_HIDDEN[0]], + }, + TensorSpec { + name: format!("{prefix}.linear2.bias"), + shape: vec![TRUNK_HIDDEN[1]], + }, + TensorSpec { + name: format!("{prefix}.output.weight"), + shape: vec![metadata.output_names.len(), TRUNK_HIDDEN[1]], + }, + TensorSpec { + name: format!("{prefix}.output.bias"), + shape: vec![metadata.output_names.len()], + }, + ]); + } + expected +} + +fn validate_tensors(tensors: &SafeTensors<'_>, metadata: &ArtifactMetadata) -> Result<()> { + let expected = expected_tensors(metadata); + for spec in &expected { + let tensor = tensors + .tensor(&spec.name) + .map_err(|_| invalid_artifact(format!("missing tensor {}", spec.name)))?; + require( + tensor.dtype() == Dtype::F32, + format!("tensor {} must use F32", spec.name), + )?; + require( + tensor.shape() == spec.shape, + format!( + "tensor {} has shape {:?}; expected {:?}", + spec.name, + tensor.shape(), + spec.shape + ), + )?; + validate_finite_f32(&spec.name, tensor.data())?; + if spec.name == "transform.scaler_scale" { + validate_positive_f32(&spec.name, tensor.data())?; + } + } + + let expected_names = expected + .iter() + .map(|spec| spec.name.as_str()) + .collect::>(); + let actual_names = tensors + .names() + .into_iter() + .map(String::as_str) + .collect::>(); + require( + actual_names == expected_names, + "artifact contains unexpected tensors", + )?; + Ok(()) +} + +/// Decodes artifact tensors once so requests never reparse the safetensors file. +fn decode_tensors(tensors: &SafeTensors<'_>) -> Result>> { + let mut decoded = BTreeMap::new(); + for name in tensors.names() { + let tensor = tensors + .tensor(name) + .map_err(|error| invalid_artifact(format!("failed to read tensor {name}: {error}")))?; + let values = tensor + .data() + .chunks_exact(4) + .map(|bytes| f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])) + .collect(); + decoded.insert(name.clone(), values); + } + Ok(decoded) +} + +/// Applies the fitted `StandardScaler` elementwise in F32. +fn standardize(raw: &[f32], mean: &[f32], scale: &[f32]) -> Result> { + if raw.len() != mean.len() || raw.len() != scale.len() { + return Err(transform_error(format!( + "scaler dimensions do not match: raw={}, mean={}, scale={}", + raw.len(), + mean.len(), + scale.len(), + ))); + } + + raw.iter() + .zip(mean) + .zip(scale) + .map(|((value, mean), scale)| { + if !value.is_finite() { + return Err(transform_error("raw features contain non-finite values")); + } + if !scale.is_finite() || *scale <= 0.0 { + return Err(transform_error("scaler scale must be finite and positive")); + } + let standardized = (*value - *mean) / *scale; + if !standardized.is_finite() { + return Err(transform_error( + "standardization produced a non-finite value", + )); + } + Ok(standardized) + }) + .collect() +} + +/// Projects standardized features using row-major sklearn PCA components. +fn project_pca( + standardized: &[f32], + mean: &[f32], + components: &[f32], + output_dim: usize, +) -> Result> { + if standardized.is_empty() || standardized.len() != mean.len() { + return Err(transform_error(format!( + "PCA input dimensions do not match: standardized={}, mean={}", + standardized.len(), + mean.len(), + ))); + } + let expected_components = output_dim + .checked_mul(standardized.len()) + .ok_or_else(|| transform_error("PCA component dimensions overflow"))?; + if output_dim == 0 || components.len() != expected_components { + return Err(transform_error(format!( + "PCA component dimensions do not match: values={}, expected={expected_components}", + components.len(), + ))); + } + + let centered = standardized + .iter() + .zip(mean) + .map(|(value, mean)| { + let centered = *value - *mean; + if centered.is_finite() { + Ok(centered) + } else { + Err(transform_error("PCA centering produced a non-finite value")) + } + }) + .collect::>>()?; + + components + .chunks_exact(standardized.len()) + .map(|component| { + // Wide F32 dot products drift beyond the exporter's golden tolerance. + let projected = component + .iter() + .zip(¢ered) + .map(|(weight, value)| f64::from(*weight) * f64::from(*value)) + .sum::() as f32; + if projected.is_finite() { + Ok(projected) + } else { + Err(transform_error( + "PCA projection produced a non-finite value", + )) + } + }) + .collect() +} + +/// Applies a row-major dense layer and an optional ReLU activation. +fn dense_layer( + input: &[f32], + weights: &[f32], + bias: &[f32], + output_dim: usize, + relu: bool, +) -> Result> { + if input.is_empty() || output_dim == 0 || bias.len() != output_dim { + return Err(trunk_error(format!( + "dense dimensions do not match: input={}, output={output_dim}, bias={}", + input.len(), + bias.len(), + ))); + } + let expected_weights = output_dim + .checked_mul(input.len()) + .ok_or_else(|| trunk_error("dense weight dimensions overflow"))?; + if weights.len() != expected_weights { + return Err(trunk_error(format!( + "dense weight dimensions do not match: values={}, expected={expected_weights}", + weights.len(), + ))); + } + if input.iter().any(|value| !value.is_finite()) { + return Err(trunk_error("dense input contains non-finite values")); + } + + weights + .chunks_exact(input.len()) + .zip(bias) + .map(|(row, bias)| { + let value = row + .iter() + .zip(input) + .map(|(weight, input)| *weight * *input) + .sum::() + + *bias; + if !value.is_finite() { + return Err(trunk_error("dense layer produced a non-finite value")); + } + Ok(if relu { value.max(0.0) } else { value }) + }) + .collect() +} + +/// Converts one finite logit to a probability without exponential overflow. +fn sigmoid_probability(logit: f32) -> Result { + if !logit.is_finite() { + return Err(trunk_error("sigmoid input contains a non-finite value")); + } + let probability = if logit >= 0.0 { + 1.0 / (1.0 + (-logit).exp()) + } else { + let exp = logit.exp(); + exp / (1.0 + exp) + }; + if probability.is_finite() { + Ok(probability) + } else { + Err(trunk_error("sigmoid produced a non-finite value")) + } +} + +fn validate_finite_f32(name: &str, data: &[u8]) -> Result<()> { + for bytes in data.chunks_exact(4) { + let value = f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + require( + value.is_finite(), + format!("tensor {name} contains non-finite values"), + )?; + } + Ok(()) +} + +fn validate_positive_f32(name: &str, data: &[u8]) -> Result<()> { + for bytes in data.chunks_exact(4) { + let value = f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + require( + value > 0.0, + format!("tensor {name} contains non-positive values"), + )?; + } + Ok(()) +} + +fn require(condition: bool, message: impl Into) -> Result<()> { + if condition { + Ok(()) + } else { + Err(invalid_artifact(message)) + } +} + +fn invalid_artifact(message: impl Into) -> SwitchyardError { + SwitchyardError::InvalidConfig(format!( + "invalid prefill-router artifact: {}", + message.into() + )) +} + +fn transform_error(message: impl Into) -> SwitchyardError { + SwitchyardError::Other(format!( + "prefill-router feature transform error: {}", + message.into() + )) +} + +fn trunk_error(message: impl Into) -> SwitchyardError { + SwitchyardError::Other(format!( + "prefill-router trunk inference error: {}", + message.into() + )) +} + +#[cfg(test)] +mod tests { + use std::path::{Path, PathBuf}; + use std::sync::atomic::{AtomicU64, Ordering}; + + use safetensors::tensor::{serialize, TensorView}; + + use super::*; + + static NEXT_TEST_DIRECTORY: AtomicU64 = AtomicU64::new(0); + + struct TestArtifactDirectory { + path: PathBuf, + } + + impl TestArtifactDirectory { + fn create() -> Result { + let sequence = NEXT_TEST_DIRECTORY.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "switchyard-prefill-artifact-{}-{sequence}", + std::process::id() + )); + std::fs::create_dir(&path).map_err(|error| { + invalid_artifact(format!( + "failed to create test directory {}: {error}", + path.display() + )) + })?; + Ok(Self { path }) + } + + fn write_metadata(&self, metadata: &ArtifactMetadata) -> Result<()> { + let bytes = serde_json::to_vec(metadata).map_err(|error| { + invalid_artifact(format!("failed to serialize test metadata: {error}")) + })?; + self.write(METADATA_FILE, &bytes) + } + + fn write(&self, name: &str, bytes: &[u8]) -> Result<()> { + let path = self.path.join(name); + std::fs::write(&path, bytes).map_err(|error| { + invalid_artifact(format!("failed to write {}: {error}", path.display())) + }) + } + + fn path(&self) -> &Path { + &self.path + } + } + + impl Drop for TestArtifactDirectory { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.path); + } + } + + fn test_metadata() -> ArtifactMetadata { + ArtifactMetadata { + format_version: FORMAT_VERSION, + training_mode: TRAINING_MODE.into(), + encoder: "probe/model".into(), + representation: REPRESENTATION.into(), + extraction_layer_ids: vec![0, 1], + hidden_size: 2, + raw_feature_dim: 4, + feature_block_count: 1, + pca_dim: PCA_DIM, + pca_whiten: false, + output_names: OUTPUT_NAMES.iter().map(|name| (*name).into()).collect(), + trunk_hidden: TRUNK_HIDDEN.to_vec(), + ensemble_size: ENSEMBLE_SIZE, + probability_link: PROBABILITY_LINK.into(), + ensemble_reduction: ENSEMBLE_REDUCTION.into(), + tensor_file: TENSOR_FILE.into(), + } + } + + fn repeated_f32_bytes(value: f32, count: usize) -> Vec { + let mut bytes = Vec::with_capacity(count * size_of::()); + for _ in 0..count { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes + } + + fn serialize_single_tensor(name: &str, shape: Vec, values: &[f32]) -> Result> { + let bytes = values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let tensor = TensorView::new(Dtype::F32, shape, &bytes).map_err(|error| { + invalid_artifact(format!("failed to create test tensor {name}: {error}")) + })?; + serialize([(name, tensor)], &None).map_err(|error| { + invalid_artifact(format!("failed to serialize test tensor {name}: {error}")) + }) + } + + fn serialize_valid_artifact(metadata: &ArtifactMetadata) -> Result> { + let storage = expected_tensors(metadata) + .into_iter() + .map(|spec| { + let value_count = spec.shape.iter().product(); + let fill = if spec.name == "transform.scaler_scale" { + 1.0 + } else { + 0.0 + }; + (spec.name, spec.shape, repeated_f32_bytes(fill, value_count)) + }) + .collect::>(); + let tensors = storage + .iter() + .map(|(name, shape, bytes)| { + TensorView::new(Dtype::F32, shape.clone(), bytes) + .map(|tensor| (name.as_str(), tensor)) + .map_err(|error| { + invalid_artifact(format!("failed to create test tensor {name}: {error}")) + }) + }) + .collect::>>()?; + serialize(tensors, &None).map_err(|error| { + invalid_artifact(format!("failed to serialize test artifact: {error}")) + }) + } + + fn trunk_test_artifact() -> InferenceArtifact { + let metadata = test_metadata(); + let mut tensors = BTreeMap::new(); + for index in 0..ENSEMBLE_SIZE { + let prefix = format!("ensemble.{index}"); + let mut linear1_bias = vec![0.0; TRUNK_HIDDEN[0]]; + linear1_bias[0] = (index + 1) as f32; + let mut linear2_weight = vec![0.0; TRUNK_HIDDEN[1] * TRUNK_HIDDEN[0]]; + linear2_weight[0] = 1.0; + let mut output_weight = vec![0.0; OUTPUT_NAMES.len() * TRUNK_HIDDEN[1]]; + for (output, factor) in [1.0, -2.0, 3.0, -4.0].into_iter().enumerate() { + output_weight[output * TRUNK_HIDDEN[1]] = factor; + } + tensors.insert( + format!("{prefix}.linear1.weight"), + vec![0.0; TRUNK_HIDDEN[0] * PCA_DIM], + ); + tensors.insert(format!("{prefix}.linear1.bias"), linear1_bias); + tensors.insert(format!("{prefix}.linear2.weight"), linear2_weight); + tensors.insert(format!("{prefix}.linear2.bias"), vec![0.0; TRUNK_HIDDEN[1]]); + tensors.insert(format!("{prefix}.output.weight"), output_weight); + tensors.insert( + format!("{prefix}.output.bias"), + vec![0.0; OUTPUT_NAMES.len()], + ); + } + InferenceArtifact { metadata, tensors } + } + + #[test] + fn artifact_loads_decodes_and_executes_the_exported_shape() -> Result<()> { + let directory = TestArtifactDirectory::create()?; + let metadata = test_metadata(); + directory.write_metadata(&metadata)?; + directory.write(TENSOR_FILE, &serialize_valid_artifact(&metadata)?)?; + + let artifact = InferenceArtifact::load(directory.path(), "probe/model")?; + assert_eq!(artifact.encoder(), "probe/model"); + assert_eq!(artifact.layer_count(), 2); + assert_eq!(artifact.hidden_size(), 2); + assert_eq!(artifact.raw_feature_dim(), 4); + assert_eq!(artifact.ensemble_size(), ENSEMBLE_SIZE); + assert_eq!(artifact.output_count(), OUTPUT_NAMES.len()); + assert!(artifact + .output_names() + .iter() + .map(String::as_str) + .eq(OUTPUT_NAMES)); + assert_eq!(artifact.tensor_count(), 4 + ENSEMBLE_SIZE * 6); + assert!(artifact.tensor_value_count() > 0); + + let projected = artifact.project(&vec![0.0; artifact.raw_feature_dim()])?; + assert_eq!(projected, vec![0.0; PCA_DIM]); + let logits = artifact.ensemble_logits(&projected)?; + assert_eq!(logits, vec![vec![0.0; OUTPUT_NAMES.len()]; ENSEMBLE_SIZE]); + let probabilities = artifact.ensemble_probabilities(&logits)?; + assert_eq!(probabilities, vec![0.5; OUTPUT_NAMES.len()]); + Ok(()) + } + + #[test] + fn metadata_rejects_encoder_layer_and_raw_dimension_mismatches() -> Result<()> { + let metadata = test_metadata(); + let encoder_error = metadata + .validate("different/probe") + .err() + .ok_or_else(|| invalid_artifact("encoder mismatch should fail"))?; + assert!(format!("{encoder_error}").contains("does not match probe model")); + + let mut metadata = test_metadata(); + metadata.extraction_layer_ids = vec![0, 2]; + let layer_error = metadata + .validate("probe/model") + .err() + .ok_or_else(|| invalid_artifact("layer ordering mismatch should fail"))?; + assert!(format!("{layer_error}").contains("contiguous and ordered")); + + let mut metadata = test_metadata(); + metadata.raw_feature_dim += 1; + let dimension_error = metadata + .validate("probe/model") + .err() + .ok_or_else(|| invalid_artifact("raw dimension mismatch should fail"))?; + assert!(format!("{dimension_error}").contains("does not equal layer count")); + Ok(()) + } + + #[test] + fn tensor_validation_rejects_missing_and_malformed_tensors() -> Result<()> { + let metadata = test_metadata(); + let unexpected = serialize_single_tensor("unexpected", vec![1], &[0.0])?; + let tensors = SafeTensors::deserialize(&unexpected).map_err(|error| { + invalid_artifact(format!("failed to deserialize test tensor: {error}")) + })?; + let missing_error = validate_tensors(&tensors, &metadata) + .err() + .ok_or_else(|| invalid_artifact("missing tensor should fail"))?; + assert!(format!("{missing_error}").contains("missing tensor transform.scaler_mean")); + + let wrong_shape = serialize_single_tensor( + "transform.scaler_mean", + vec![metadata.raw_feature_dim - 1], + &vec![0.0; metadata.raw_feature_dim - 1], + )?; + let tensors = SafeTensors::deserialize(&wrong_shape).map_err(|error| { + invalid_artifact(format!("failed to deserialize test tensor: {error}")) + })?; + let shape_error = validate_tensors(&tensors, &metadata) + .err() + .ok_or_else(|| invalid_artifact("malformed tensor shape should fail"))?; + assert!(format!("{shape_error}").contains("has shape")); + Ok(()) + } + + #[test] + fn tensor_validation_rejects_non_finite_values() -> Result<()> { + let metadata = test_metadata(); + let mut values = vec![0.0; metadata.raw_feature_dim]; + values[0] = f32::NAN; + let encoded = serialize_single_tensor( + "transform.scaler_mean", + vec![metadata.raw_feature_dim], + &values, + )?; + let tensors = SafeTensors::deserialize(&encoded).map_err(|error| { + invalid_artifact(format!("failed to deserialize test tensor: {error}")) + })?; + let error = validate_tensors(&tensors, &metadata) + .err() + .ok_or_else(|| invalid_artifact("non-finite tensor should fail"))?; + assert!(format!("{error}").contains("contains non-finite values")); + Ok(()) + } + + #[test] + fn artifact_load_rejects_malformed_tensor_file() -> Result<()> { + let directory = TestArtifactDirectory::create()?; + directory.write_metadata(&test_metadata())?; + directory.write(TENSOR_FILE, b"not a safetensors file")?; + + let error = InferenceArtifact::load(directory.path(), "probe/model") + .err() + .ok_or_else(|| invalid_artifact("malformed tensor file should fail"))?; + assert!(format!("{error}").contains("failed to parse")); + Ok(()) + } + + #[test] + fn scaler_and_pca_match_exported_row_major_math() -> Result<()> { + let standardized = standardize(&[3.0, 6.0, 11.0], &[1.0, 2.0, 3.0], &[2.0, 2.0, 4.0])?; + assert_eq!(standardized, vec![1.0, 2.0, 2.0]); + + let projected = project_pca( + &standardized, + &[0.5, 1.0, 1.5], + &[ + 1.0, 10.0, 100.0, // PCA component 0 + -2.0, 0.5, 4.0, // PCA component 1 + ], + 2, + )?; + assert_eq!(projected, vec![60.5, 1.5]); + Ok(()) + } + + #[test] + fn scaler_rejects_dimension_mismatch() -> Result<()> { + let error = standardize(&[1.0, 2.0], &[0.0], &[1.0, 1.0]) + .err() + .ok_or_else(|| transform_error("dimension mismatch should fail"))?; + assert!(format!("{error}").contains("scaler dimensions do not match")); + Ok(()) + } + + #[test] + fn scaler_rejects_non_finite_output() -> Result<()> { + let error = standardize(&[f32::MAX], &[-f32::MAX], &[1.0]) + .err() + .ok_or_else(|| transform_error("non-finite output should fail"))?; + assert!(format!("{error}").contains("standardization produced")); + Ok(()) + } + + #[test] + fn pca_rejects_component_shape_mismatch() -> Result<()> { + let error = project_pca(&[1.0, 2.0], &[0.0, 0.0], &[1.0, 2.0, 3.0], 2) + .err() + .ok_or_else(|| transform_error("component mismatch should fail"))?; + assert!(format!("{error}").contains("PCA component dimensions do not match")); + Ok(()) + } + + #[test] + fn pca_rejects_non_finite_output() -> Result<()> { + let error = project_pca(&[f32::MAX, f32::MAX], &[0.0, 0.0], &[2.0, 2.0], 1) + .err() + .ok_or_else(|| transform_error("non-finite projection should fail"))?; + assert!(format!("{error}").contains("PCA projection produced")); + Ok(()) + } + + #[test] + fn dense_layer_applies_row_major_weights_bias_and_relu() -> Result<()> { + let output = dense_layer( + &[2.0, -1.0], + &[ + 1.0, 2.0, // output row 0 + -3.0, 4.0, // output row 1 + ], + &[1.0, 1.0], + 2, + true, + )?; + assert_eq!(output, vec![1.0, 0.0]); + Ok(()) + } + + #[test] + fn dense_layer_rejects_weight_shape_mismatch() -> Result<()> { + let error = dense_layer(&[1.0, 2.0], &[1.0, 2.0, 3.0], &[0.0, 0.0], 2, false) + .err() + .ok_or_else(|| trunk_error("weight mismatch should fail"))?; + assert!(format!("{error}").contains("dense weight dimensions do not match")); + Ok(()) + } + + #[test] + fn dense_layer_rejects_non_finite_output() -> Result<()> { + let error = dense_layer(&[f32::MAX], &[2.0], &[0.0], 1, false) + .err() + .ok_or_else(|| trunk_error("non-finite output should fail"))?; + assert!(format!("{error}").contains("dense layer produced")); + Ok(()) + } + + #[test] + fn ensemble_trunk_preserves_member_and_output_order() -> Result<()> { + let logits = trunk_test_artifact().ensemble_logits(&vec![0.0; PCA_DIM])?; + assert_eq!(logits.len(), ENSEMBLE_SIZE); + for (member, member_logits) in logits.iter().enumerate() { + let scale = (member + 1) as f32; + assert_eq!( + member_logits, + &[scale, -2.0 * scale, 3.0 * scale, -4.0 * scale] + ); + } + Ok(()) + } + + #[test] + fn ensemble_trunk_rejects_wrong_pca_dimension() -> Result<()> { + let error = trunk_test_artifact() + .ensemble_logits(&[0.0; PCA_DIM - 1]) + .err() + .ok_or_else(|| trunk_error("PCA dimension mismatch should fail"))?; + assert!(format!("{error}").contains("does not match pca_dim")); + Ok(()) + } + + #[test] + fn sigmoid_is_stable_for_extreme_logits() -> Result<()> { + assert_eq!(sigmoid_probability(-f32::MAX)?, 0.0); + assert_eq!(sigmoid_probability(0.0)?, 0.5); + assert_eq!(sigmoid_probability(f32::MAX)?, 1.0); + Ok(()) + } + + #[test] + fn ensemble_probabilities_average_members_in_output_order() -> Result<()> { + let log_three = 3.0f32.ln(); + let mut logits = vec![vec![0.0; OUTPUT_NAMES.len()]; ENSEMBLE_SIZE]; + logits[0][0] = log_three; + for member in &mut logits { + member[2] = log_three; + member[3] = -log_three; + } + + let probabilities = trunk_test_artifact().ensemble_probabilities(&logits)?; + let expected = [0.55, 0.5, 0.75, 0.25]; + for (actual, expected) in probabilities.iter().zip(expected) { + assert!((actual - expected).abs() < 1e-6); + } + assert!((probabilities.iter().sum::() - 1.0).abs() > 0.5); + Ok(()) + } + + #[test] + fn ensemble_probabilities_reject_wrong_member_count() -> Result<()> { + let logits = vec![vec![0.0; OUTPUT_NAMES.len()]; ENSEMBLE_SIZE - 1]; + let error = trunk_test_artifact() + .ensemble_probabilities(&logits) + .err() + .ok_or_else(|| trunk_error("member count mismatch should fail"))?; + assert!(format!("{error}").contains("does not match ensemble_size")); + Ok(()) + } + + #[test] + fn ensemble_probabilities_reject_wrong_output_count() -> Result<()> { + let logits = vec![vec![0.0; OUTPUT_NAMES.len() - 1]; ENSEMBLE_SIZE]; + let error = trunk_test_artifact() + .ensemble_probabilities(&logits) + .err() + .ok_or_else(|| trunk_error("output count mismatch should fail"))?; + assert!(format!("{error}").contains("does not match output count")); + Ok(()) + } + + #[test] + fn ensemble_probabilities_reject_non_finite_logits() -> Result<()> { + let mut logits = vec![vec![0.0; OUTPUT_NAMES.len()]; ENSEMBLE_SIZE]; + logits[0][0] = f32::NAN; + let error = trunk_test_artifact() + .ensemble_probabilities(&logits) + .err() + .ok_or_else(|| trunk_error("non-finite logit should fail"))?; + assert!(format!("{error}").contains("sigmoid input contains a non-finite value")); + Ok(()) + } +} diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs new file mode 100644 index 00000000..c31c5fa1 --- /dev/null +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs @@ -0,0 +1,7 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Learned prefill-router artifact inference and routing internals. + +#[allow(dead_code)] +mod artifact; From f025f7cab3d45ec3bd563b1fa6f6e82092f75d7c Mon Sep 17 00:00:00 2001 From: michxu Date: Thu, 16 Jul 2026 20:39:29 +0000 Subject: [PATCH 02/12] feat(prefill-probe): add cost-aware routing policy Signed-off-by: michxu --- .../src/profiles/prefill_router/mod.rs | 2 + .../src/profiles/prefill_router/policy.rs | 169 ++++++++++++++++++ 2 files changed, 171 insertions(+) create mode 100644 crates/switchyard-components-v2/src/profiles/prefill_router/policy.rs diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs index c31c5fa1..814e6f0a 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs @@ -5,3 +5,5 @@ #[allow(dead_code)] mod artifact; +#[allow(dead_code)] +mod policy; diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/policy.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/policy.rs new file mode 100644 index 00000000..32e22465 --- /dev/null +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/policy.rs @@ -0,0 +1,169 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Private cost-aware policy for converting learned correctness into a binary routing score. + +use switchyard_core::{Result, SwitchyardError}; + +/// Validated utility policy for the two configured completion targets. +#[derive(Debug, Clone, Copy)] +pub(crate) struct CostAwareRoutingPolicy { + lambda: f64, + normalized_weak_cost: f64, + normalized_strong_cost: f64, +} + +impl CostAwareRoutingPolicy { + /// Validates the policy and min-max normalizes costs across weak and strong. + pub(crate) fn new(lambda: f64, weak_cost: f64, strong_cost: f64) -> Result { + if !lambda.is_finite() || !(0.0..=1.0).contains(&lambda) { + return Err(SwitchyardError::InvalidConfig( + "routing_policy.lambda must be finite and in [0.0, 1.0]".into(), + )); + } + validate_cost("weak_cost", weak_cost)?; + validate_cost("strong_cost", strong_cost)?; + + let (normalized_weak_cost, normalized_strong_cost) = if weak_cost == strong_cost { + (0.0, 0.0) + } else { + let minimum = weak_cost.min(strong_cost); + let range = weak_cost.max(strong_cost) - minimum; + ( + (weak_cost - minimum) / range, + (strong_cost - minimum) / range, + ) + }; + + Ok(Self { + lambda, + normalized_weak_cost, + normalized_strong_cost, + }) + } + + /// Returns `1.0` for weak and `0.0` for strong from two mapped probabilities. + pub(crate) fn score(&self, weak_probability: f64, strong_probability: f64) -> Result { + validate_probability("weak", weak_probability)?; + validate_probability("strong", strong_probability)?; + + let cost_weight = 1.0 - self.lambda; + let weak_utility = self.lambda * weak_probability - cost_weight * self.normalized_weak_cost; + let strong_utility = + self.lambda * strong_probability - cost_weight * self.normalized_strong_cost; + let margin = weak_utility - strong_utility; + Ok(if margin >= 0.0 { 1.0 } else { 0.0 }) + } +} + +fn validate_cost(field: &str, cost: f64) -> Result<()> { + if !cost.is_finite() || cost < 0.0 { + return Err(SwitchyardError::InvalidConfig(format!( + "routing_policy.{field} must be finite and non-negative" + ))); + } + Ok(()) +} + +fn validate_probability(head: &str, probability: f64) -> Result<()> { + if !probability.is_finite() || !(0.0..=1.0).contains(&probability) { + return Err(SwitchyardError::Other(format!( + "{head} checkpoint probability must be finite and in [0.0, 1.0]" + ))); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn lambda_zero_uses_only_cost_and_selects_weak_on_tie() -> Result<()> { + let weak_cheaper = CostAwareRoutingPolicy::new(0.0, 1.0, 10.0)?; + let strong_cheaper = CostAwareRoutingPolicy::new(0.0, 10.0, 1.0)?; + let equal_cost = CostAwareRoutingPolicy::new(0.0, 3.0, 3.0)?; + + assert_eq!(weak_cheaper.score(0.0, 1.0)?, 1.0); + assert_eq!(strong_cheaper.score(1.0, 0.0)?, 0.0); + assert_eq!(equal_cost.score(0.0, 1.0)?, 1.0); + assert_eq!(equal_cost.normalized_weak_cost, 0.0); + assert_eq!(equal_cost.normalized_strong_cost, 0.0); + Ok(()) + } + + #[test] + fn lambda_one_uses_only_correctness_probabilities() -> Result<()> { + let policy = CostAwareRoutingPolicy::new(1.0, 100.0, 1.0)?; + + assert_eq!(policy.score(0.8, 0.6)?, 1.0); + assert_eq!(policy.score(0.4, 0.6)?, 0.0); + assert_eq!(policy.score(0.6, 0.6)?, 1.0); + Ok(()) + } + + #[test] + fn lambda_changes_the_route_without_another_threshold() -> Result<()> { + let balanced = CostAwareRoutingPolicy::new(0.5, 0.0, 1.0)?; + let correctness_favoring = CostAwareRoutingPolicy::new(0.9, 0.0, 1.0)?; + + assert_eq!(balanced.score(0.4, 0.8)?, 1.0); + assert_eq!(correctness_favoring.score(0.4, 0.8)?, 0.0); + Ok(()) + } + + #[test] + fn unmapped_checkpoint_outputs_do_not_affect_the_score() -> Result<()> { + let first_outputs = [0.99, 0.7, 0.4, 0.01]; + let second_outputs = [0.01, 0.7, 0.4, 0.99]; + let policy = CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?; + + assert_eq!( + policy.score(first_outputs[1], first_outputs[2])?, + policy.score(second_outputs[1], second_outputs[2])?, + ); + Ok(()) + } + + #[test] + fn invalid_policy_values_are_rejected() -> Result<()> { + let cases = [ + (f64::NAN, 0.0, 1.0, "lambda"), + (-0.1, 0.0, 1.0, "lambda"), + (1.1, 0.0, 1.0, "lambda"), + (0.5, -1.0, 1.0, "weak_cost"), + (0.5, 1.0, f64::INFINITY, "strong_cost"), + ]; + for (lambda, weak_cost, strong_cost, expected) in cases { + let error = CostAwareRoutingPolicy::new(lambda, weak_cost, strong_cost) + .err() + .ok_or_else(|| { + SwitchyardError::Other(format!( + "invalid policy value for {expected} should fail" + )) + })?; + assert!(format!("{error}").contains(expected)); + } + Ok(()) + } + + #[test] + fn invalid_checkpoint_probabilities_are_rejected() -> Result<()> { + let policy = CostAwareRoutingPolicy::new(0.5, 0.0, 1.0)?; + let cases = [ + (f64::NAN, 0.5, "weak"), + (-0.1, 0.5, "weak"), + (0.5, 1.1, "strong"), + (0.5, f64::INFINITY, "strong"), + ]; + for (weak, strong, expected) in cases { + let error = policy.score(weak, strong).err().ok_or_else(|| { + SwitchyardError::Other(format!( + "invalid checkpoint probability for {expected} should fail" + )) + })?; + assert!(format!("{error}").contains(expected)); + } + Ok(()) + } +} From b621d4f1e341558791b5f8a0ccb507c275ff283a Mon Sep 17 00:00:00 2001 From: michxu Date: Thu, 16 Jul 2026 21:07:06 +0000 Subject: [PATCH 03/12] feat(prefill-probe): add hidden-state feature extraction Signed-off-by: michxu --- Cargo.lock | 18 + crates/switchyard-components-v2/Cargo.toml | 1 + .../src/profiles/prefill_router/mod.rs | 2 + .../src/profiles/prefill_router/scorer.rs | 710 ++++++++++++++++++ 4 files changed, 731 insertions(+) create mode 100644 crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs diff --git a/Cargo.lock b/Cargo.lock index 3b938e56..dfb6f6da 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -257,6 +257,12 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + [[package]] name = "diff" version = "0.1.13" @@ -398,6 +404,17 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + [[package]] name = "hashbrown" version = "0.15.5" @@ -1385,6 +1402,7 @@ version = "0.1.0" dependencies = [ "async-trait", "futures-core", + "half", "parking_lot", "rand 0.8.6", "reqwest", diff --git a/crates/switchyard-components-v2/Cargo.toml b/crates/switchyard-components-v2/Cargo.toml index 40aee965..90048bb1 100644 --- a/crates/switchyard-components-v2/Cargo.toml +++ b/crates/switchyard-components-v2/Cargo.toml @@ -13,6 +13,7 @@ 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"] } diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs index 814e6f0a..2eaf1400 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs @@ -7,3 +7,5 @@ mod artifact; #[allow(dead_code)] mod policy; +#[allow(dead_code)] +mod scorer; diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs new file mode 100644 index 00000000..724e439f --- /dev/null +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs @@ -0,0 +1,710 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Validation and token-mean reduction for vLLM hidden-state artifacts. + +use std::fs::{File, OpenOptions}; +use std::io::Read; +use std::path::{Path, PathBuf}; + +use safetensors::{Dtype, SafeTensors}; +use switchyard_core::{Result, SwitchyardError}; + +use super::artifact::InferenceArtifact; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct HiddenStateLayout { + layer_count: usize, + hidden_size: usize, + raw_feature_dim: usize, +} + +impl HiddenStateLayout { + fn from_artifact(artifact: &InferenceArtifact) -> Self { + Self { + layer_count: artifact.layer_count(), + hidden_size: artifact.hidden_size(), + raw_feature_dim: artifact.raw_feature_dim(), + } + } +} + +/// Decodes `[tokens, layers, hidden]` data into one layer-major token-mean vector. +fn token_mean_per_layer( + data: &[u8], + dtype: Dtype, + shape: &[usize], + expected: HiddenStateLayout, +) -> Result> { + if shape.len() != 3 { + return Err(hidden_state_error( + "expected hidden_states shape [prompt_tokens, layers, hidden_size]", + )); + } + let (prompt_tokens, layer_count, hidden_size) = (shape[0], shape[1], shape[2]); + if prompt_tokens == 0 { + return Err(hidden_state_error( + "hidden_states token dimension must be non-zero", + )); + } + if layer_count == 0 || hidden_size == 0 { + return Err(hidden_state_error( + "hidden_states layer and hidden dimensions must be non-zero", + )); + } + if layer_count != expected.layer_count { + return Err(hidden_state_error(format!( + "hidden_states layer count {layer_count} does not match artifact layer count {}", + expected.layer_count, + ))); + } + if hidden_size != expected.hidden_size { + return Err(hidden_state_error(format!( + "hidden_states hidden size {hidden_size} does not match artifact hidden size {}", + expected.hidden_size, + ))); + } + + let element_size = match dtype { + Dtype::F32 => size_of::(), + Dtype::BF16 => size_of::(), + other => { + return Err(hidden_state_error(format!( + "unsupported hidden_states dtype: {other:?}" + ))) + } + }; + let features_per_token = layer_count + .checked_mul(hidden_size) + .ok_or_else(|| hidden_state_error("hidden_states shape is too large"))?; + let bytes_per_token = features_per_token + .checked_mul(element_size) + .ok_or_else(|| hidden_state_error("hidden_states byte length is too large"))?; + let expected_bytes = prompt_tokens + .checked_mul(bytes_per_token) + .ok_or_else(|| hidden_state_error("hidden_states byte length is too large"))?; + if data.len() != expected_bytes { + return Err(hidden_state_error(format!( + "hidden_states byte length {} does not match shape byte length {expected_bytes}", + data.len(), + ))); + } + + let mut pooled = vec![0.0f32; features_per_token]; + match dtype { + Dtype::F32 => { + for token in data.chunks_exact(bytes_per_token) { + for (index, bytes) in token.chunks_exact(size_of::()).enumerate() { + let value = f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + accumulate(&mut pooled[index], value)?; + } + } + } + Dtype::BF16 => { + for token in data.chunks_exact(bytes_per_token) { + for (index, bytes) in token.chunks_exact(size_of::()).enumerate() { + let value = half::bf16::from_le_bytes([bytes[0], bytes[1]]).to_f32(); + accumulate(&mut pooled[index], value)?; + } + } + } + other => { + return Err(hidden_state_error(format!( + "unsupported hidden_states dtype: {other:?}" + ))) + } + } + + let token_count = prompt_tokens as f32; + for value in &mut pooled { + *value /= token_count; + if !value.is_finite() { + return Err(hidden_state_error( + "hidden-state token mean produced a non-finite value", + )); + } + } + if pooled.len() != expected.raw_feature_dim { + return Err(hidden_state_error(format!( + "hidden-state feature length {} does not match artifact raw_feature_dim {}", + pooled.len(), + expected.raw_feature_dim, + ))); + } + Ok(pooled) +} + +fn accumulate(sum: &mut f32, value: f32) -> Result<()> { + if !value.is_finite() { + return Err(hidden_state_error( + "hidden_states contains non-finite values", + )); + } + *sum += value; + if !sum.is_finite() { + return Err(hidden_state_error( + "hidden-state token accumulation produced a non-finite value", + )); + } + Ok(()) +} + +fn validate_token_ids(tensors: &SafeTensors<'_>, prompt_tokens: usize) -> Result<()> { + if !tensors + .names() + .iter() + .any(|name| name.as_str() == "token_ids") + { + return Ok(()); + } + let token_ids = tensors + .tensor("token_ids") + .map_err(|error| hidden_state_error(format!("token_ids tensor error: {error}")))?; + if token_ids.dtype() != Dtype::I64 { + return Err(hidden_state_error(format!( + "token_ids must use I64; got {:?}", + token_ids.dtype(), + ))); + } + if token_ids.shape() != [prompt_tokens] { + return Err(hidden_state_error(format!( + "token_ids shape {:?} does not match hidden_states token count {prompt_tokens}", + token_ids.shape(), + ))); + } + for bytes in token_ids.data().chunks_exact(size_of::()) { + let token_id = i64::from_le_bytes([ + bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7], + ]); + if token_id < 0 { + return Err(hidden_state_error("token_ids contains a negative token ID")); + } + } + Ok(()) +} + +fn parse_hidden_state_features(bytes: &[u8], expected: HiddenStateLayout) -> Result> { + let tensors = SafeTensors::deserialize(bytes) + .map_err(|error| hidden_state_error(format!("safetensors parse error: {error}")))?; + let hidden_states = tensors + .tensor("hidden_states") + .map_err(|error| hidden_state_error(format!("hidden_states tensor not found: {error}")))?; + let prompt_tokens = match hidden_states.shape() { + [prompt_tokens, _, _] => *prompt_tokens, + _ => { + return Err(hidden_state_error( + "expected hidden_states shape [prompt_tokens, layers, hidden_size]", + )) + } + }; + validate_token_ids(&tensors, prompt_tokens)?; + token_mean_per_layer( + hidden_states.data(), + hidden_states.dtype(), + hidden_states.shape(), + expected, + ) +} + +fn validate_hidden_states_path(root: &Path, path: &Path) -> Result { + if !has_safetensors_extension(path) { + return Err(hidden_state_error(format!( + "hidden-state artifact must be a .safetensors file: {}", + path.display(), + ))); + } + + let root = root.canonicalize().map_err(|error| { + hidden_state_error(format!( + "hidden-states directory {} is not accessible: {error}", + root.display(), + )) + })?; + if !root.is_dir() { + return Err(hidden_state_error(format!( + "hidden-states root is not a directory: {}", + root.display(), + ))); + } + let actual = path.canonicalize().map_err(|error| { + hidden_state_error(format!( + "hidden-state artifact {} is not accessible: {error}", + path.display(), + )) + })?; + if !actual.starts_with(&root) { + return Err(hidden_state_error(format!( + "hidden-state artifact {} is outside configured directory {}", + actual.display(), + root.display(), + ))); + } + if !has_safetensors_extension(&actual) { + return Err(hidden_state_error(format!( + "canonical hidden-state artifact must be a .safetensors file: {}", + actual.display(), + ))); + } + let metadata = actual.metadata().map_err(|error| { + hidden_state_error(format!( + "hidden-state artifact metadata error for {}: {error}", + actual.display(), + )) + })?; + if !metadata.is_file() { + return Err(hidden_state_error(format!( + "hidden-state artifact is not a regular file: {}", + actual.display(), + ))); + } + Ok(actual) +} + +fn has_safetensors_extension(path: &Path) -> bool { + path.extension().and_then(|extension| extension.to_str()) == Some("safetensors") +} + +fn open_locked_artifact(path: &Path) -> Result { + let file = OpenOptions::new() + .read(true) + .write(true) + .open(path) + .map_err(|error| { + hidden_state_error(format!( + "hidden-state artifact open error for {}: {error}", + path.display(), + )) + })?; + file.lock().map_err(|error| { + hidden_state_error(format!( + "hidden-state artifact lock error for {}: {error}", + path.display(), + )) + })?; + Ok(file) +} + +/// Reads and removes one validated artifact while holding its exclusive lock. +fn read_and_cleanup_hidden_states( + root: &Path, + path: &Path, + expected: HiddenStateLayout, +) -> Result> { + let artifact_path = validate_hidden_states_path(root, path)?; + let mut artifact = open_locked_artifact(&artifact_path)?; + let mut bytes = Vec::new(); + artifact.read_to_end(&mut bytes).map_err(|error| { + hidden_state_error(format!( + "hidden-state artifact read error for {}: {error}", + artifact_path.display(), + )) + })?; + let features = parse_hidden_state_features(&bytes, expected)?; + std::fs::remove_file(&artifact_path).map_err(|error| { + hidden_state_error(format!( + "hidden-state artifact cleanup error for {}: {error}", + artifact_path.display(), + )) + })?; + Ok(features) +} + +fn hidden_state_error(message: impl Into) -> SwitchyardError { + SwitchyardError::Other(format!( + "prefill-router hidden-state error: {}", + message.into() + )) +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicU64, Ordering}; + + use safetensors::tensor::{serialize, TensorView}; + + use super::*; + + static NEXT_TEST_DIRECTORY: AtomicU64 = AtomicU64::new(0); + + struct TestDirectory { + path: PathBuf, + } + + impl TestDirectory { + fn create() -> Result { + let sequence = NEXT_TEST_DIRECTORY.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "switchyard-hidden-state-{}-{sequence}", + std::process::id() + )); + std::fs::create_dir(&path).map_err(|error| { + hidden_state_error(format!( + "failed to create test directory {}: {error}", + path.display() + )) + })?; + Ok(Self { path }) + } + + fn path(&self) -> &Path { + &self.path + } + + fn create_subdirectory(&self, name: &str) -> Result { + let path = self.path.join(name); + std::fs::create_dir(&path).map_err(|error| { + hidden_state_error(format!( + "failed to create test directory {}: {error}", + path.display() + )) + })?; + Ok(path) + } + + fn write(&self, name: &str, bytes: &[u8]) -> Result { + let path = self.path.join(name); + std::fs::write(&path, bytes).map_err(|error| { + hidden_state_error(format!("failed to write {}: {error}", path.display())) + })?; + Ok(path) + } + } + + impl Drop for TestDirectory { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.path); + } + } + + fn layout(layer_count: usize, hidden_size: usize) -> HiddenStateLayout { + HiddenStateLayout { + layer_count, + hidden_size, + raw_feature_dim: layer_count * hidden_size, + } + } + + fn f32_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + fn bf16_bytes(values: &[f32]) -> Vec { + values + .iter() + .flat_map(|value| half::bf16::from_f32(*value).to_le_bytes()) + .collect() + } + + fn i64_bytes(values: &[i64]) -> Vec { + values + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() + } + + fn serialize_f32_hidden_states( + shape: Vec, + values: &[f32], + token_ids: Option<(Dtype, Vec, Vec)>, + ) -> Result> { + let hidden_data = f32_bytes(values); + let hidden_view = TensorView::new(Dtype::F32, shape, &hidden_data).map_err(|error| { + hidden_state_error(format!( + "failed to create hidden-state test tensor: {error}" + )) + })?; + let serialized = if let Some((dtype, shape, token_data)) = token_ids.as_ref() { + let token_view = + TensorView::new(*dtype, shape.clone(), token_data).map_err(|error| { + hidden_state_error(format!("failed to create token ID test tensor: {error}")) + })?; + serialize( + [("hidden_states", hidden_view), ("token_ids", token_view)], + &None, + ) + } else { + serialize([("hidden_states", hidden_view)], &None) + }; + serialized.map_err(|error| { + hidden_state_error(format!( + "failed to serialize hidden-state test tensor: {error}" + )) + }) + } + + fn valid_artifact_bytes() -> Result> { + serialize_f32_hidden_states( + vec![2, 2, 2], + &[ + 1.0, 3.0, 5.0, 7.0, // token 0, layers 0 and 1 + 2.0, 4.0, 6.0, 8.0, // token 1, layers 0 and 1 + ], + None, + ) + } + + #[test] + fn token_mean_pooling_produces_layer_major_features() -> Result<()> { + let pooled = token_mean_per_layer( + &f32_bytes(&[ + 1.0, 3.0, 5.0, 7.0, // token 0, layers 0 and 1 + 2.0, 4.0, 6.0, 8.0, // token 1, layers 0 and 1 + ]), + Dtype::F32, + &[2, 2, 2], + layout(2, 2), + )?; + assert_eq!(pooled, vec![1.5, 3.5, 5.5, 7.5]); + Ok(()) + } + + #[test] + fn one_token_is_unchanged_by_token_mean_pooling() -> Result<()> { + let values = [1.0, 3.0, 5.0, 7.0]; + let pooled = + token_mean_per_layer(&f32_bytes(&values), Dtype::F32, &[1, 2, 2], layout(2, 2))?; + assert_eq!(pooled, values); + Ok(()) + } + + #[test] + fn token_and_layer_axes_cannot_be_swapped() -> Result<()> { + let error = + token_mean_per_layer(&f32_bytes(&[0.0; 12]), Dtype::F32, &[2, 3, 2], layout(2, 3)) + .err() + .ok_or_else(|| hidden_state_error("swapped axes should fail"))?; + assert!(format!("{error}").contains("layer count 3")); + Ok(()) + } + + #[test] + fn bf16_and_f32_token_means_agree() -> Result<()> { + let values = [1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0]; + let f32_pooled = + token_mean_per_layer(&f32_bytes(&values), Dtype::F32, &[2, 2, 2], layout(2, 2))?; + let bf16_pooled = + token_mean_per_layer(&bf16_bytes(&values), Dtype::BF16, &[2, 2, 2], layout(2, 2))?; + for (left, right) in f32_pooled.iter().zip(&bf16_pooled) { + assert!((left - right).abs() < 1e-3); + } + Ok(()) + } + + #[test] + fn token_mean_pooling_rejects_invalid_shape_and_byte_length() -> Result<()> { + let malformed = + token_mean_per_layer(&f32_bytes(&[1.0, 2.0]), Dtype::F32, &[2, 1], layout(1, 2)) + .err() + .ok_or_else(|| hidden_state_error("malformed shape should fail"))?; + assert!(format!("{malformed}").contains("expected hidden_states shape")); + + let empty = token_mean_per_layer(&[], Dtype::F32, &[0, 2, 2], layout(2, 2)) + .err() + .ok_or_else(|| hidden_state_error("empty token axis should fail"))?; + assert!(format!("{empty}").contains("token dimension must be non-zero")); + + let overflow = token_mean_per_layer( + &[], + Dtype::F32, + &[1, usize::MAX, 2], + HiddenStateLayout { + layer_count: usize::MAX, + hidden_size: 2, + raw_feature_dim: 0, + }, + ) + .err() + .ok_or_else(|| hidden_state_error("shape overflow should fail"))?; + assert!(format!("{overflow}").contains("shape is too large")); + + let wrong_length = + token_mean_per_layer(&f32_bytes(&[1.0]), Dtype::F32, &[2, 2, 2], layout(2, 2)) + .err() + .ok_or_else(|| hidden_state_error("wrong byte length should fail"))?; + assert!(format!("{wrong_length}").contains("byte length")); + Ok(()) + } + + #[test] + fn token_mean_pooling_rejects_invalid_values_and_dtype() -> Result<()> { + let non_finite = token_mean_per_layer( + &f32_bytes(&[1.0, f32::NAN]), + Dtype::F32, + &[1, 1, 2], + layout(1, 2), + ) + .err() + .ok_or_else(|| hidden_state_error("non-finite value should fail"))?; + assert!(format!("{non_finite}").contains("contains non-finite")); + + let accumulation_overflow = token_mean_per_layer( + &f32_bytes(&[f32::MAX, f32::MAX]), + Dtype::F32, + &[2, 1, 1], + layout(1, 1), + ) + .err() + .ok_or_else(|| hidden_state_error("accumulation overflow should fail"))?; + assert!(format!("{accumulation_overflow}").contains("accumulation produced")); + + let unsupported = + token_mean_per_layer(&i64_bytes(&[1]), Dtype::I64, &[1, 1, 1], layout(1, 1)) + .err() + .ok_or_else(|| hidden_state_error("unsupported dtype should fail"))?; + assert!(format!("{unsupported}").contains("unsupported hidden_states dtype")); + Ok(()) + } + + #[test] + fn final_feature_length_must_match_artifact() -> Result<()> { + let error = token_mean_per_layer( + &f32_bytes(&[1.0, 2.0]), + Dtype::F32, + &[1, 1, 2], + HiddenStateLayout { + layer_count: 1, + hidden_size: 2, + raw_feature_dim: 3, + }, + ) + .err() + .ok_or_else(|| hidden_state_error("raw feature dimension mismatch should fail"))?; + assert!(format!("{error}").contains("raw_feature_dim 3")); + Ok(()) + } + + #[test] + fn token_ids_are_validated_when_present() -> Result<()> { + let valid = serialize_f32_hidden_states( + vec![2, 1, 2], + &[1.0, 2.0, 3.0, 4.0], + Some((Dtype::I64, vec![2], i64_bytes(&[101, 102]))), + )?; + assert_eq!( + parse_hidden_state_features(&valid, layout(1, 2))?, + vec![2.0, 3.0] + ); + + let wrong_shape = serialize_f32_hidden_states( + vec![2, 1, 2], + &[1.0, 2.0, 3.0, 4.0], + Some((Dtype::I64, vec![1], i64_bytes(&[101]))), + )?; + let error = parse_hidden_state_features(&wrong_shape, layout(1, 2)) + .err() + .ok_or_else(|| hidden_state_error("token count mismatch should fail"))?; + assert!(format!("{error}").contains("token_ids shape [1]")); + + let negative = serialize_f32_hidden_states( + vec![2, 1, 2], + &[1.0, 2.0, 3.0, 4.0], + Some((Dtype::I64, vec![2], i64_bytes(&[101, -1]))), + )?; + let error = parse_hidden_state_features(&negative, layout(1, 2)) + .err() + .ok_or_else(|| hidden_state_error("negative token ID should fail"))?; + assert!(format!("{error}").contains("negative token ID")); + + let wrong_dtype = serialize_f32_hidden_states( + vec![2, 1, 2], + &[1.0, 2.0, 3.0, 4.0], + Some((Dtype::F32, vec![2], f32_bytes(&[101.0, 102.0]))), + )?; + let error = parse_hidden_state_features(&wrong_dtype, layout(1, 2)) + .err() + .ok_or_else(|| hidden_state_error("token ID dtype mismatch should fail"))?; + assert!(format!("{error}").contains("token_ids must use I64")); + Ok(()) + } + + #[test] + fn hidden_state_path_must_stay_under_configured_directory() -> Result<()> { + let directory = TestDirectory::create()?; + let root = directory.create_subdirectory("root")?; + let outside = directory.write("outside.safetensors", b"not a tensor")?; + let traversing_path = root.join("..").join("outside.safetensors"); + + let error = validate_hidden_states_path(&root, &traversing_path) + .err() + .ok_or_else(|| hidden_state_error("outside path should fail"))?; + assert!(format!("{error}").contains("outside configured directory")); + assert_eq!( + outside + .canonicalize() + .map_err(|error| hidden_state_error(error.to_string()))?, + traversing_path + .canonicalize() + .map_err(|error| hidden_state_error(error.to_string()))? + ); + Ok(()) + } + + #[test] + fn hidden_state_path_requires_extension_and_regular_file() -> Result<()> { + let directory = TestDirectory::create()?; + let wrong_extension = directory.write("hidden.bin", b"not a tensor")?; + let error = validate_hidden_states_path(directory.path(), &wrong_extension) + .err() + .ok_or_else(|| hidden_state_error("wrong extension should fail"))?; + assert!(format!("{error}").contains(".safetensors")); + + let not_a_file = directory.create_subdirectory("directory.safetensors")?; + let error = validate_hidden_states_path(directory.path(), ¬_a_file) + .err() + .ok_or_else(|| hidden_state_error("directory artifact should fail"))?; + assert!(format!("{error}").contains("not a regular file")); + Ok(()) + } + + #[test] + fn artifact_file_is_exclusively_locked() -> Result<()> { + let directory = TestDirectory::create()?; + let path = directory.write("hidden.safetensors", b"locked")?; + let first = open_locked_artifact(&path)?; + let second = OpenOptions::new() + .read(true) + .write(true) + .open(&path) + .map_err(|error| hidden_state_error(format!("second test open failed: {error}")))?; + + let error = second + .try_lock() + .err() + .ok_or_else(|| hidden_state_error("second exclusive lock should fail"))?; + assert!(matches!(error, std::fs::TryLockError::WouldBlock)); + drop(first); + second.try_lock().map_err(|error| { + hidden_state_error(format!("released lock was not reusable: {error}")) + })?; + Ok(()) + } + + #[test] + fn successful_locked_read_removes_artifact() -> Result<()> { + let directory = TestDirectory::create()?; + let path = directory.write("hidden.safetensors", &valid_artifact_bytes()?)?; + + let features = read_and_cleanup_hidden_states(directory.path(), &path, layout(2, 2))?; + + assert_eq!(features, vec![1.5, 3.5, 5.5, 7.5]); + assert!(!path.exists()); + Ok(()) + } + + #[test] + fn failed_locked_read_preserves_artifact() -> Result<()> { + let directory = TestDirectory::create()?; + let path = directory.write("hidden.safetensors", b"not a tensor")?; + + let error = read_and_cleanup_hidden_states(directory.path(), &path, layout(2, 2)) + .err() + .ok_or_else(|| hidden_state_error("malformed artifact should fail"))?; + + assert!(format!("{error}").contains("safetensors parse error")); + assert!(path.exists()); + Ok(()) + } +} From a60ba031bd4a95f1b4fd8acecf37a46fd4b10705 Mon Sep 17 00:00:00 2001 From: michxu Date: Thu, 16 Jul 2026 21:29:50 +0000 Subject: [PATCH 04/12] feat(prefill-probe): add vllm learned scorer Signed-off-by: michxu --- crates/switchyard-components-v2/Cargo.toml | 1 + .../src/profiles/prefill_router/scorer.rs | 619 +++++++++++++++++- 2 files changed, 619 insertions(+), 1 deletion(-) diff --git a/crates/switchyard-components-v2/Cargo.toml b/crates/switchyard-components-v2/Cargo.toml index 90048bb1..0f1a9f32 100644 --- a/crates/switchyard-components-v2/Cargo.toml +++ b/crates/switchyard-components-v2/Cargo.toml @@ -23,6 +23,7 @@ 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" diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs index 724e439f..ed358667 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs @@ -6,11 +6,203 @@ use std::fs::{File, OpenOptions}; use std::io::Read; use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::Duration; +use async_trait::async_trait; use safetensors::{Dtype, SafeTensors}; -use switchyard_core::{Result, SwitchyardError}; +use serde::Deserialize; +use switchyard_core::{ChatRequest, Result, SwitchyardError}; use super::artifact::InferenceArtifact; +use super::policy::CostAwareRoutingPolicy; + +const ARTIFACT_WAIT_ATTEMPTS: usize = 20; +const ARTIFACT_WAIT_INTERVAL: Duration = Duration::from_millis(50); + +/// Private scalar scorer consumed by the prefill routing profile. +#[async_trait] +pub(crate) trait ProbeScorer: Send + Sync { + /// Returns `1.0` for weak and `0.0` for strong. + async fn score(&self, request: &ChatRequest) -> Result; +} + +struct CheckpointPrediction { + probabilities: Vec, +} + +impl CheckpointPrediction { + fn mapped_probabilities( + &self, + weak_head_index: usize, + strong_head_index: usize, + ) -> Result<(f64, f64)> { + let weak = self.probabilities.get(weak_head_index).ok_or_else(|| { + probe_error(format!( + "weak checkpoint head index {weak_head_index} is outside prediction length {}", + self.probabilities.len(), + )) + })?; + let strong = self.probabilities.get(strong_head_index).ok_or_else(|| { + probe_error(format!( + "strong checkpoint head index {strong_head_index} is outside prediction length {}", + self.probabilities.len(), + )) + })?; + Ok((f64::from(*weak), f64::from(*strong))) + } +} + +struct LearnedRouting { + artifact: Arc, + weak_head_index: usize, + strong_head_index: usize, + policy: CostAwareRoutingPolicy, +} + +impl LearnedRouting { + fn score(&self, raw_features: &[f32]) -> Result { + if raw_features.len() != self.artifact.raw_feature_dim() { + return Err(probe_error(format!( + "hidden-state feature length {} does not match artifact raw_feature_dim {}", + raw_features.len(), + self.artifact.raw_feature_dim(), + ))); + } + + let projected = self.artifact.project(raw_features)?; + let logits = self.artifact.ensemble_logits(&projected)?; + let prediction = CheckpointPrediction { + probabilities: self.artifact.ensemble_probabilities(&logits)?, + }; + let (weak_probability, strong_probability) = + prediction.mapped_probabilities(self.weak_head_index, self.strong_head_index)?; + let score = self.policy.score(weak_probability, strong_probability)?; + tracing::debug!(score, "prefill router learned score"); + Ok(score) + } +} + +/// Learned scorer backed by a dedicated vLLM hidden-state probe endpoint. +pub(crate) struct HiddenStateProbeScorer { + completions_url: String, + model: String, + hidden_states_dir: PathBuf, + client: reqwest::Client, + routing: LearnedRouting, +} + +impl HiddenStateProbeScorer { + /// Builds a scorer from startup-validated artifact and policy resources. + pub(crate) fn new( + base_url: impl Into, + model: impl Into, + hidden_states_dir: impl Into, + artifact: Arc, + weak_head_index: usize, + strong_head_index: usize, + policy: CostAwareRoutingPolicy, + ) -> Self { + let base_url = base_url.into(); + Self { + completions_url: format!("{}/chat/completions", base_url.trim_end_matches('/')), + model: model.into(), + hidden_states_dir: hidden_states_dir.into(), + client: reqwest::Client::new(), + routing: LearnedRouting { + artifact, + weak_head_index, + strong_head_index, + policy, + }, + } + } +} + +#[derive(Deserialize)] +struct CompletionResponse { + kv_transfer_params: Option, +} + +#[derive(Deserialize)] +struct KvTransferParams { + hidden_states_path: String, +} + +#[async_trait] +impl ProbeScorer for HiddenStateProbeScorer { + async fn score(&self, request: &ChatRequest) -> Result { + let messages = request + .body() + .as_object() + .and_then(|body| body.get("messages")) + .cloned() + .unwrap_or_else(|| serde_json::Value::Array(Vec::new())); + let response = self + .client + .post(&self.completions_url) + .json(&serde_json::json!({ + "model": self.model, + "messages": messages, + "max_tokens": 1, + "kv_transfer_params": { + "hidden_states_path": self.hidden_states_dir, + "include_output_tokens": false, + }, + })) + .send() + .await + .map_err(|error| probe_error(format!("request failed: {error}")))?; + + if !response.status().is_success() { + return Err(probe_error(format!( + "endpoint returned HTTP {}", + response.status(), + ))); + } + let response: CompletionResponse = response + .json() + .await + .map_err(|error| probe_error(format!("response parse failed: {error}")))?; + let hidden_states_path = response + .kv_transfer_params + .ok_or_else(|| { + probe_error( + "response missing kv_transfer_params; verify ExampleHiddenStatesConnector", + ) + })? + .hidden_states_path; + let hidden_states_path = Path::new(&hidden_states_path); + wait_for_artifact(hidden_states_path).await?; + + let raw_features = read_and_cleanup_hidden_states( + &self.hidden_states_dir, + hidden_states_path, + HiddenStateLayout::from_artifact(&self.routing.artifact), + )?; + self.routing.score(&raw_features) + } +} + +async fn wait_for_artifact(path: &Path) -> Result<()> { + for attempt in 0..ARTIFACT_WAIT_ATTEMPTS { + if path.exists() { + return Ok(()); + } + if attempt + 1 < ARTIFACT_WAIT_ATTEMPTS { + tokio::time::sleep(ARTIFACT_WAIT_INTERVAL).await; + } + } + Err(probe_error(format!( + "hidden-state artifact {} did not appear after {} ms", + path.display(), + ARTIFACT_WAIT_INTERVAL.as_millis() * (ARTIFACT_WAIT_ATTEMPTS - 1) as u128, + ))) +} + +fn probe_error(message: impl Into) -> SwitchyardError { + SwitchyardError::Other(format!("prefill-router probe error: {}", message.into())) +} #[derive(Debug, Clone, Copy, PartialEq, Eq)] struct HiddenStateLayout { @@ -318,7 +510,11 @@ fn hidden_state_error(message: impl Into) -> SwitchyardError { #[cfg(test)] mod tests { + use std::io::Write; + use std::net::{SocketAddr, TcpListener, TcpStream}; use std::sync::atomic::{AtomicU64, Ordering}; + use std::sync::Mutex; + use std::thread::{self, JoinHandle}; use safetensors::tensor::{serialize, TensorView}; @@ -326,6 +522,141 @@ mod tests { static NEXT_TEST_DIRECTORY: AtomicU64 = AtomicU64::new(0); + #[derive(Clone, Debug, PartialEq)] + struct ObservedRequest { + path: String, + body: serde_json::Value, + } + + struct MockProbeServer { + addr: SocketAddr, + requests: Arc>>, + handle: Option>, + } + + impl MockProbeServer { + fn spawn(status: u16, body: serde_json::Value) -> Result { + let body = serde_json::to_vec(&body) + .map_err(|error| probe_error(format!("mock response encode failed: {error}")))?; + Self::spawn_bytes(status, body) + } + + fn spawn_bytes(status: u16, body: Vec) -> Result { + let listener = TcpListener::bind("127.0.0.1:0") + .map_err(|error| probe_error(format!("mock bind failed: {error}")))?; + let addr = listener + .local_addr() + .map_err(|error| probe_error(format!("mock local_addr failed: {error}")))?; + let requests = Arc::new(Mutex::new(Vec::new())); + let thread_requests = Arc::clone(&requests); + let handle = thread::spawn(move || { + let Ok((mut stream, _)) = listener.accept() else { + return; + }; + let Ok(request) = read_http_request(&mut stream) else { + return; + }; + if let Ok(mut requests) = thread_requests.lock() { + requests.push(request); + } + let _ = write_http_response(&mut stream, status, &body); + }); + Ok(Self { + addr, + requests, + handle: Some(handle), + }) + } + + fn base_url(&self) -> String { + format!("http://{}/v1", self.addr) + } + + fn request(&self) -> Result { + self.requests + .lock() + .map_err(|_| probe_error("mock requests mutex poisoned"))? + .first() + .cloned() + .ok_or_else(|| probe_error("mock server did not observe a request")) + } + } + + impl Drop for MockProbeServer { + fn drop(&mut self) { + let _ = TcpStream::connect(self.addr); + if let Some(handle) = self.handle.take() { + let _ = handle.join(); + } + } + } + + fn read_http_request(stream: &mut TcpStream) -> Result { + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .map_err(|error| probe_error(format!("mock read timeout failed: {error}")))?; + let mut bytes = Vec::new(); + let mut buffer = [0_u8; 1024]; + let header_end = loop { + let read = stream + .read(&mut buffer) + .map_err(|error| probe_error(format!("mock request read failed: {error}")))?; + if read == 0 { + return Err(probe_error("mock connection closed before headers")); + } + bytes.extend_from_slice(&buffer[..read]); + if let Some(header_end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { + break header_end; + } + }; + let headers = String::from_utf8_lossy(&bytes[..header_end]); + let request_line = headers + .lines() + .next() + .ok_or_else(|| probe_error("mock request line missing"))?; + let path = request_line + .split_whitespace() + .nth(1) + .ok_or_else(|| probe_error("mock request path missing"))? + .to_string(); + let content_length = headers + .lines() + .skip(1) + .filter_map(|line| line.split_once(':')) + .find(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .and_then(|(_, value)| value.trim().parse::().ok()) + .ok_or_else(|| probe_error("mock content-length missing"))?; + let body_start = header_end + 4; + while bytes.len().saturating_sub(body_start) < content_length { + let read = stream + .read(&mut buffer) + .map_err(|error| probe_error(format!("mock body read failed: {error}")))?; + if read == 0 { + return Err(probe_error("mock connection closed before body")); + } + bytes.extend_from_slice(&buffer[..read]); + } + let body = serde_json::from_slice(&bytes[body_start..body_start + content_length]) + .map_err(|error| probe_error(format!("mock request decode failed: {error}")))?; + Ok(ObservedRequest { path, body }) + } + + fn write_http_response(stream: &mut TcpStream, status: u16, body: &[u8]) -> Result<()> { + let reason = match status { + 200 => "OK", + 503 => "Service Unavailable", + _ => "Test Response", + }; + let headers = format!( + "HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len(), + ); + stream + .write_all(headers.as_bytes()) + .and_then(|()| stream.write_all(body)) + .map_err(|error| probe_error(format!("mock response write failed: {error}"))) + } + struct TestDirectory { path: PathBuf, } @@ -446,6 +777,139 @@ mod tests { ) } + fn repeated_f32_bytes(value: f32, count: usize) -> Vec { + let mut bytes = Vec::with_capacity(count * size_of::()); + for _ in 0..count { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes + } + + fn write_router_artifact( + directory: &TestDirectory, + name: &str, + output_biases: [f32; 4], + ) -> Result> { + const PCA_DIM: usize = 200; + const HIDDEN_1: usize = 256; + const HIDDEN_2: usize = 128; + const OUTPUTS: usize = 4; + const ENSEMBLE_SIZE: usize = 5; + const RAW_DIM: usize = 4; + + let artifact_dir = directory.create_subdirectory(name)?; + let metadata = serde_json::json!({ + "format_version": 1, + "training_mode": "single_pca_block", + "encoder": "probe/model", + "representation": "token_mean_per_layer_concat", + "extraction_layer_ids": [0, 1], + "hidden_size": 2, + "raw_feature_dim": RAW_DIM, + "feature_block_count": 1, + "pca_dim": PCA_DIM, + "pca_whiten": false, + "output_names": ["qwen-122b", "nemotron-3-super", "opus-4.7", "gpt-5.5"], + "trunk_hidden": [HIDDEN_1, HIDDEN_2], + "ensemble_size": ENSEMBLE_SIZE, + "probability_link": "independent_sigmoid", + "ensemble_reduction": "probability_mean", + "tensor_file": "router.safetensors", + }); + let metadata_bytes = serde_json::to_vec(&metadata) + .map_err(|error| probe_error(format!("test metadata encode failed: {error}")))?; + std::fs::write(artifact_dir.join("router.json"), metadata_bytes) + .map_err(|error| probe_error(format!("test metadata write failed: {error}")))?; + + let mut storage = vec![ + ( + "transform.scaler_mean".to_string(), + vec![RAW_DIM], + repeated_f32_bytes(0.0, RAW_DIM), + ), + ( + "transform.scaler_scale".to_string(), + vec![RAW_DIM], + repeated_f32_bytes(1.0, RAW_DIM), + ), + ( + "transform.pca_mean".to_string(), + vec![RAW_DIM], + repeated_f32_bytes(0.0, RAW_DIM), + ), + ( + "transform.pca_components".to_string(), + vec![PCA_DIM, RAW_DIM], + repeated_f32_bytes(0.0, PCA_DIM * RAW_DIM), + ), + ]; + for index in 0..ENSEMBLE_SIZE { + let prefix = format!("ensemble.{index}"); + storage.extend([ + ( + format!("{prefix}.linear1.weight"), + vec![HIDDEN_1, PCA_DIM], + repeated_f32_bytes(0.0, HIDDEN_1 * PCA_DIM), + ), + ( + format!("{prefix}.linear1.bias"), + vec![HIDDEN_1], + repeated_f32_bytes(0.0, HIDDEN_1), + ), + ( + format!("{prefix}.linear2.weight"), + vec![HIDDEN_2, HIDDEN_1], + repeated_f32_bytes(0.0, HIDDEN_2 * HIDDEN_1), + ), + ( + format!("{prefix}.linear2.bias"), + vec![HIDDEN_2], + repeated_f32_bytes(0.0, HIDDEN_2), + ), + ( + format!("{prefix}.output.weight"), + vec![OUTPUTS, HIDDEN_2], + repeated_f32_bytes(0.0, OUTPUTS * HIDDEN_2), + ), + ( + format!("{prefix}.output.bias"), + vec![OUTPUTS], + f32_bytes(&output_biases), + ), + ]); + } + let views = storage + .iter() + .map(|(name, shape, bytes)| { + TensorView::new(Dtype::F32, shape.clone(), bytes) + .map(|view| (name.as_str(), view)) + .map_err(|error| { + probe_error(format!("test router tensor {name} is invalid: {error}")) + }) + }) + .collect::>>()?; + let tensor_bytes = serialize(views, &None) + .map_err(|error| probe_error(format!("test router encode failed: {error}")))?; + std::fs::write(artifact_dir.join("router.safetensors"), tensor_bytes) + .map_err(|error| probe_error(format!("test router write failed: {error}")))?; + + Ok(Arc::new(InferenceArtifact::load( + artifact_dir, + "probe/model", + )?)) + } + + fn test_request() -> ChatRequest { + ChatRequest::openai_chat(serde_json::json!({ + "model": "client/model", + "messages": [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "Explain the failure."}, + ], + "temperature": 0.2, + })) + } + #[test] fn token_mean_pooling_produces_layer_major_features() -> Result<()> { let pooled = token_mean_per_layer( @@ -707,4 +1171,157 @@ mod tests { assert!(path.exists()); Ok(()) } + + #[tokio::test] + async fn learned_probe_forwards_exact_request_and_returns_binary_score() -> Result<()> { + let directory = TestDirectory::create()?; + let hidden_path = directory.write("probe-hidden.safetensors", &valid_artifact_bytes()?)?; + let artifact = write_router_artifact(&directory, "router", [8.0, 2.0, -2.0, -8.0])?; + let server = MockProbeServer::spawn( + 200, + serde_json::json!({ + "kv_transfer_params": { + "hidden_states_path": hidden_path.to_string_lossy(), + } + }), + )?; + let scorer = HiddenStateProbeScorer::new( + server.base_url(), + "probe/model", + directory.path(), + artifact, + 1, + 2, + CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, + ); + let request = test_request(); + + let score = scorer.score(&request).await?; + + assert_eq!(score, 1.0); + assert!(!hidden_path.exists()); + let observed = server.request()?; + assert_eq!(observed.path, "/v1/chat/completions"); + assert_eq!( + observed.body, + serde_json::json!({ + "model": "probe/model", + "messages": [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "Explain the failure."}, + ], + "max_tokens": 1, + "kv_transfer_params": { + "hidden_states_path": directory.path(), + "include_output_tokens": false, + }, + }) + ); + Ok(()) + } + + #[tokio::test] + async fn probe_response_errors_are_reported_before_artifact_processing() -> Result<()> { + let directory = TestDirectory::create()?; + let artifact = write_router_artifact(&directory, "router", [8.0, 2.0, -2.0, -8.0])?; + let request = test_request(); + + let missing_server = MockProbeServer::spawn(200, serde_json::json!({}))?; + let missing_scorer = HiddenStateProbeScorer::new( + missing_server.base_url(), + "probe/model", + directory.path(), + Arc::clone(&artifact), + 1, + 2, + CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, + ); + let error = missing_scorer + .score(&request) + .await + .err() + .ok_or_else(|| probe_error("missing kv_transfer_params should fail"))?; + assert!(format!("{error}").contains("missing kv_transfer_params")); + + let malformed_server = MockProbeServer::spawn( + 200, + serde_json::json!({"kv_transfer_params": {"hidden_states_path": 42}}), + )?; + let malformed_scorer = HiddenStateProbeScorer::new( + malformed_server.base_url(), + "probe/model", + directory.path(), + Arc::clone(&artifact), + 1, + 2, + CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, + ); + let error = malformed_scorer + .score(&request) + .await + .err() + .ok_or_else(|| probe_error("malformed kv_transfer_params should fail"))?; + assert!(format!("{error}").contains("response parse failed")); + + let unavailable_server = MockProbeServer::spawn(503, serde_json::json!({}))?; + let unavailable_scorer = HiddenStateProbeScorer::new( + unavailable_server.base_url(), + "probe/model", + directory.path(), + artifact, + 1, + 2, + CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, + ); + let error = unavailable_scorer + .score(&request) + .await + .err() + .ok_or_else(|| probe_error("non-success probe response should fail"))?; + assert!(format!("{error}").contains("endpoint returned HTTP 503")); + Ok(()) + } + + #[test] + fn learned_routing_uses_only_mapped_heads_and_preserves_score_direction() -> Result<()> { + let directory = TestDirectory::create()?; + let first = write_router_artifact(&directory, "router-a", [8.0, 2.0, -2.0, -8.0])?; + let second = write_router_artifact(&directory, "router-b", [-8.0, 2.0, -2.0, 8.0])?; + let policy = CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?; + let first_selected = LearnedRouting { + artifact: Arc::clone(&first), + weak_head_index: 1, + strong_head_index: 2, + policy, + }; + let second_selected = LearnedRouting { + artifact: second, + weak_head_index: 1, + strong_head_index: 2, + policy, + }; + let reversed = LearnedRouting { + artifact: Arc::clone(&first), + weak_head_index: 2, + strong_head_index: 1, + policy, + }; + + assert_eq!(first_selected.score(&[0.0; 4])?, 1.0); + assert_eq!(second_selected.score(&[0.0; 4])?, 1.0); + assert_eq!(reversed.score(&[0.0; 4])?, 0.0); + + let invalid = LearnedRouting { + artifact: first, + weak_head_index: 4, + strong_head_index: 1, + policy, + }; + let error = invalid + .score(&[0.0; 4]) + .err() + .ok_or_else(|| probe_error("out-of-range head index should fail"))?; + assert!(format!("{error}").contains("weak checkpoint head index 4")); + Ok(()) + } } From 4b037010e76c01bfae4d5a0ea18a933906943e40 Mon Sep 17 00:00:00 2001 From: michxu Date: Thu, 16 Jul 2026 22:03:47 +0000 Subject: [PATCH 05/12] feat(prefill-probe): add learned profile configuration Signed-off-by: michxu --- crates/switchyard-components-v2/src/lib.rs | 10 +- .../src/profiles/mod.rs | 4 + .../src/profiles/prefill_router/mod.rs | 8 +- .../src/profiles/prefill_router/profile.rs | 652 ++++++++++++++++++ .../src/profiles/profile_types.rs | 4 +- .../tests/prefill_probe_profile.rs | 467 +++++++++++++ 6 files changed, 1138 insertions(+), 7 deletions(-) create mode 100644 crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs create mode 100644 crates/switchyard-components-v2/tests/prefill_probe_profile.rs diff --git a/crates/switchyard-components-v2/src/lib.rs b/crates/switchyard-components-v2/src/lib.rs index 201b04a5..164117de 100644 --- a/crates/switchyard-components-v2/src/lib.rs +++ b/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; diff --git a/crates/switchyard-components-v2/src/profiles/mod.rs b/crates/switchyard-components-v2/src/profiles/mod.rs index f6d80e5d..bf8e4b1f 100644 --- a/crates/switchyard-components-v2/src/profiles/mod.rs +++ b/crates/switchyard-components-v2/src/profiles/mod.rs @@ -23,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, diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs index 2eaf1400..c042767b 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/mod.rs @@ -5,7 +5,11 @@ #[allow(dead_code)] mod artifact; -#[allow(dead_code)] mod policy; -#[allow(dead_code)] +mod profile; mod scorer; + +pub use profile::{ + PrefillProbeDecision, PrefillProbeProcessedRequest, PrefillProbeProfile, + PrefillProbeProfileConfig, PrefillProbeRoutingPolicyConfig, +}; diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs new file mode 100644 index 00000000..ab70c635 --- /dev/null +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs @@ -0,0 +1,652 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Profile configuration and runtime for learned prefill-probe routing. + +use std::collections::HashMap; +use std::hash::{Hash, Hasher}; +use std::sync::Arc; +use std::time::Instant; + +use async_trait::async_trait; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; +use switchyard_components::stats::usage_from_body; +use switchyard_components::StatsAccumulator; +use switchyard_core::{ChatResponse, LlmTarget, Result, SwitchyardError}; + +use crate::backend::{native_target_backend, TargetBackend}; +use crate::profile_stats_accumulator; +use crate::{ + profile_config, Profile, ProfileConfig, ProfileHooks, ProfileInput, ProfileResponse, + RoutingMetadata, +}; + +use super::artifact::InferenceArtifact; +use super::policy::CostAwareRoutingPolicy; +use super::scorer::{HiddenStateProbeScorer, ProbeScorer}; + +const LEARNED_SCORE_THRESHOLD: f64 = 0.5; +const TIER_STRONG: &str = "strong"; +const TIER_WEAK: &str = "weak"; + +/// Learned routing policy applied to the mapped weak and strong checkpoint heads. +#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "kebab-case", deny_unknown_fields)] +pub enum PrefillProbeRoutingPolicyConfig { + /// Balances predicted correctness against normalized completion cost. + CostAware { + /// Correctness weight in `[0, 1]`; this is the profile's routing knob. + lambda: f64, + /// Non-negative weak-target cost in the same units as `strong_cost`. + weak_cost: f64, + /// Non-negative strong-target cost in the same units as `weak_cost`. + strong_cost: f64, + }, +} + +/// Config for routing with learned prompt hidden-state features. +/// +/// The probe target produces hidden states but is never selected for completion +/// inference. The scorer returns `1.0` for weak and `0.0` for strong, and this +/// profile applies a fixed decision threshold of `0.5`. +#[profile_config("prefill-probe")] +pub struct PrefillProbeProfileConfig { + /// Probe target used only to produce prompt hidden states. + #[profile_target] + pub probe: LlmTarget, + /// Strong completion target selected by score `0.0` or probe failure. + #[profile_target] + pub strong: LlmTarget, + /// Artifact output head corresponding to the strong completion target. + pub strong_checkpoint_head: String, + /// Weak completion target selected by score `1.0`. + #[profile_target] + pub weak: LlmTarget, + /// Artifact output head corresponding to the weak completion target. + pub weak_checkpoint_head: String, + /// Directory shared with vLLM's `ExampleHiddenStatesConnector`. + pub hidden_states_dir: String, + /// Directory containing `router.json` and `router.safetensors`. + pub inference_artifact_dir: String, + /// Policy that maps the two selected correctness probabilities to a binary score. + pub routing_policy: PrefillProbeRoutingPolicyConfig, +} + +impl ProfileConfig for PrefillProbeProfileConfig { + type Runtime = PrefillProbeProfile; + + /// Validates configuration and builds the complete learned routing runtime. + fn build(&self) -> Result { + let policy = match self.routing_policy { + PrefillProbeRoutingPolicyConfig::CostAware { + lambda, + weak_cost, + strong_cost, + } => CostAwareRoutingPolicy::new(lambda, weak_cost, strong_cost)?, + }; + let artifact = + InferenceArtifact::load(&self.inference_artifact_dir, self.probe.model.as_str())?; + let weak_head_index = checkpoint_head_index( + &artifact, + "weak_checkpoint_head", + &self.weak_checkpoint_head, + )?; + let strong_head_index = checkpoint_head_index( + &artifact, + "strong_checkpoint_head", + &self.strong_checkpoint_head, + )?; + if weak_head_index == strong_head_index { + return Err(SwitchyardError::InvalidConfig(format!( + "weak_checkpoint_head and strong_checkpoint_head must map to distinct outputs; both map to `{}`", + self.weak_checkpoint_head, + ))); + } + + let base_url = self + .probe + .endpoint + .base_url + .clone() + .unwrap_or_else(|| "http://localhost:8000/v1".to_string()); + let artifact = Arc::new(artifact); + let scorer = HiddenStateProbeScorer::new( + base_url, + self.probe.model.as_str(), + self.hidden_states_dir.as_str(), + artifact, + weak_head_index, + strong_head_index, + policy, + ); + + Ok(PrefillProbeProfile { + strong_backend: native_target_backend(self.strong.clone())?, + weak_backend: native_target_backend(self.weak.clone())?, + score_threshold: LEARNED_SCORE_THRESHOLD, + scorer: Arc::new(scorer), + stats: profile_stats_accumulator(), + decision_cache: Arc::new(Mutex::new(HashMap::new())), + }) + } +} + +fn checkpoint_head_index( + artifact: &InferenceArtifact, + field: &str, + checkpoint_head: &str, +) -> Result { + artifact + .output_names() + .iter() + .position(|name| name == checkpoint_head) + .ok_or_else(|| { + SwitchyardError::InvalidConfig(format!( + "{field} `{checkpoint_head}` is not present in artifact output_names {:?}", + artifact.output_names(), + )) + }) +} + +/// One strong/weak decision emitted by the prefill-probe router. +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct PrefillProbeDecision { + /// ID of the selected completion target. + pub selected_target: String, + /// Model string sent to the selected completion backend. + pub selected_model: String, + /// Selected tier, either `strong` or `weak`. + pub tier: &'static str, + /// Binary weak-preference score: `0.0` for strong and `1.0` for weak. + pub score: f64, +} + +/// Request prepared for a completion backend with its routing decision. +pub struct PrefillProbeProcessedRequest { + /// Routed input with its model rewritten to the selected completion model. + pub profile_input: ProfileInput, + /// Decision used to select the completion backend. + pub decision: PrefillProbeDecision, +} + +/// Runtime for learned prefill-probe strong/weak routing. +pub struct PrefillProbeProfile { + strong_backend: TargetBackend, + weak_backend: TargetBackend, + score_threshold: f64, + scorer: Arc, + stats: StatsAccumulator, + // Successful scores are reused by the first string-valued user instruction. + decision_cache: Arc>>, +} + +impl PrefillProbeProfile { + // Hashes the first user instruction represented as a string. + fn instruction_key(request: &switchyard_core::ChatRequest) -> Option { + let messages = request.body().as_object()?.get("messages")?.as_array()?; + let content = messages.iter().find_map(|message| { + (message.get("role").and_then(serde_json::Value::as_str) == Some("user")) + .then(|| message.get("content").and_then(serde_json::Value::as_str)) + .flatten() + })?; + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + content.hash(&mut hasher); + Some(hasher.finish()) + } + + // Scores an uncached instruction and explicitly bypasses threshold routing on failure. + async fn route(&self, mut input: ProfileInput) -> Result { + let key = Self::instruction_key(&input.request); + let cached = key.and_then(|key| self.decision_cache.lock().get(&key).copied()); + let (score, strong_fallback) = if let Some(score) = cached { + (score, false) + } else { + match self.scorer.score(&input.request).await { + Ok(score) => { + if let Some(key) = key { + self.decision_cache.lock().insert(key, score); + } + (score, false) + } + Err(error) => { + tracing::warn!( + error = %error, + fallback_tier = TIER_STRONG, + score = 0.0, + "prefill probe failed; using uncached strong fallback" + ); + (0.0, true) + } + } + }; + + let (selected_backend, tier) = if strong_fallback { + (&self.strong_backend, TIER_STRONG) + } else if score >= self.score_threshold { + (&self.weak_backend, TIER_WEAK) + } else { + (&self.strong_backend, TIER_STRONG) + }; + input + .request + .set_model(selected_backend.target().model.as_str()); + + Ok(PrefillProbeProcessedRequest { + profile_input: input, + decision: PrefillProbeDecision { + selected_target: selected_backend.target().id.to_string(), + selected_model: selected_backend.target().model.to_string(), + tier, + score, + }, + }) + } + + fn backend_for(&self, decision: &PrefillProbeDecision) -> Result<&TargetBackend> { + if decision.selected_target == self.strong_backend.target().id.as_str() { + Ok(&self.strong_backend) + } else if decision.selected_target == self.weak_backend.target().id.as_str() { + Ok(&self.weak_backend) + } else { + Err(SwitchyardError::InvalidConfig(format!( + "prefill probe selected target {} that is not configured for this profile", + decision.selected_target, + ))) + } + } + + fn record_success( + &self, + decision: &PrefillProbeDecision, + response: &ChatResponse, + total_latency_ms: f64, + backend_latency_ms: f64, + ) -> Result<()> { + self.stats.record_success( + decision.selected_model.as_str(), + Some(backend_latency_ms), + Some(decision.tier), + )?; + let routing_overhead_ms = (total_latency_ms - backend_latency_ms).max(0.0); + let usage = response.body().map(usage_from_body).unwrap_or_default(); + self.stats.record_usage_after_success_attribution( + decision.selected_model.as_str(), + usage, + Some(total_latency_ms), + Some(routing_overhead_ms), + Some(decision.tier), + )?; + Ok(()) + } + + fn record_error(&self, decision: &PrefillProbeDecision) -> Result<()> { + self.stats + .record_error(decision.selected_model.as_str(), Some(decision.tier)) + } + + fn routing_metadata(&self, decision: &PrefillProbeDecision) -> RoutingMetadata { + RoutingMetadata { + selected_model: Some(decision.selected_model.clone()), + selected_tier: Some(decision.tier.to_string()), + confidence: None, + router_version: Some("prefill-probe:v1".to_string()), + tolerance: Some(LEARNED_SCORE_THRESHOLD), + rationale: Some(format!( + "binary weak-preference score {} selected {}", + decision.score, decision.tier, + )), + } + } +} + +#[async_trait] +impl ProfileHooks for PrefillProbeProfile { + type ProcessedRequest = PrefillProbeProcessedRequest; + + /// Scores the prompt and returns a request prepared for the selected backend. + async fn process(&self, input: ProfileInput) -> Result { + self.route(input).await + } + + /// Leaves the selected backend response unchanged. + async fn rprocess( + &self, + _processed: &Self::ProcessedRequest, + response: ChatResponse, + ) -> Result { + Ok(response) + } +} + +#[async_trait] +impl Profile for PrefillProbeProfile { + /// Routes one request, calls the selected completion backend, and records the outcome. + async fn run(&self, input: ProfileInput) -> Result { + let profile_started_at = Instant::now(); + let processed = self.process(input).await?; + let decision = &processed.decision; + let backend = self.backend_for(decision)?; + let backend_started_at = Instant::now(); + let response = match backend.call(&processed.profile_input.request).await { + Ok(response) => response, + Err(error) => { + self.record_error(decision)?; + return Err(error); + } + }; + let backend_latency_ms = backend_started_at.elapsed().as_secs_f64() * 1000.0; + let total_latency_ms = profile_started_at.elapsed().as_secs_f64() * 1000.0; + self.record_success(decision, &response, total_latency_ms, backend_latency_ms)?; + let response = self.rprocess(&processed, response).await?; + Ok(ProfileResponse::with_routing_metadata( + response, + self.routing_metadata(decision), + )) + } +} + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use serde_json::{json, Value}; + use switchyard_core::{BackendFormat, ChatRequest, LlmTargetId, ModelId, SwitchyardError}; + + use crate::backend::ProfileBackend; + use crate::RequestMetadata; + + use super::*; + + #[derive(Clone, Debug, PartialEq)] + struct ObservedCall { + backend: &'static str, + body: Value, + } + + struct TestBackend { + name: &'static str, + fail: bool, + calls: Arc>>, + } + + #[async_trait] + impl ProfileBackend for TestBackend { + async fn call(&self, request: &ChatRequest) -> Result { + self.calls.lock().push(ObservedCall { + backend: self.name, + body: request.body().clone(), + }); + if self.fail { + return Err(SwitchyardError::Backend(format!("{} failed", self.name))); + } + Ok(ChatResponse::openai_completion(json!({ + "served_by": self.name, + "model": request.model(), + "usage": {"prompt_tokens": 5, "completion_tokens": 3}, + }))) + } + } + + struct FixedScorer(f64); + + #[async_trait] + impl ProbeScorer for FixedScorer { + async fn score(&self, _request: &ChatRequest) -> Result { + Ok(self.0) + } + } + + struct ErrorScorer; + + #[async_trait] + impl ProbeScorer for ErrorScorer { + async fn score(&self, _request: &ChatRequest) -> Result { + Err(SwitchyardError::Other("probe unavailable".to_string())) + } + } + + struct CountingScorer { + score: f64, + calls: Arc, + } + + #[async_trait] + impl ProbeScorer for CountingScorer { + async fn score(&self, _request: &ChatRequest) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(self.score) + } + } + + struct FlakyScorer { + calls: Arc, + } + + #[async_trait] + impl ProbeScorer for FlakyScorer { + async fn score(&self, _request: &ChatRequest) -> Result { + let attempt = self.calls.fetch_add(1, Ordering::SeqCst); + if attempt == 0 { + Err(SwitchyardError::Other( + "transient probe failure".to_string(), + )) + } else { + Ok(1.0) + } + } + } + + fn target(id: &str, model: &str) -> Result { + let mut target = LlmTarget::new(LlmTargetId::new(id)?, ModelId::new(model)?); + target.format = BackendFormat::OpenAi; + Ok(target) + } + + fn profile( + scorer: Arc, + strong_fails: bool, + weak_fails: bool, + ) -> Result<(PrefillProbeProfile, Arc>>)> { + let calls = Arc::new(Mutex::new(Vec::new())); + let strong = target("strong", "frontier/model")?; + let weak = target("weak", "cheap/model")?; + Ok(( + PrefillProbeProfile { + strong_backend: TargetBackend::new( + strong, + Arc::new(TestBackend { + name: "strong-backend", + fail: strong_fails, + calls: calls.clone(), + }), + ), + weak_backend: TargetBackend::new( + weak, + Arc::new(TestBackend { + name: "weak-backend", + fail: weak_fails, + calls: calls.clone(), + }), + ), + score_threshold: LEARNED_SCORE_THRESHOLD, + scorer, + stats: StatsAccumulator::new(), + decision_cache: Arc::new(Mutex::new(HashMap::new())), + }, + calls, + )) + } + + fn input(instruction: &str) -> ProfileInput { + ProfileInput { + request: ChatRequest::openai_chat(json!({ + "model": "client/model", + "messages": [{"role": "user", "content": instruction}], + })), + metadata: RequestMetadata::default(), + } + } + + fn observed(calls: &Arc>>) -> Vec { + calls.lock().clone() + } + + #[tokio::test] + async fn binary_score_direction_and_model_rewrite_are_fixed() -> Result<()> { + let (weak_profile, weak_calls) = profile(Arc::new(FixedScorer(1.0)), false, false)?; + let weak = weak_profile.process(input("weak task")).await?; + assert_eq!(weak_profile.score_threshold, 0.5); + assert_eq!(weak.decision.tier, TIER_WEAK); + assert_eq!(weak.decision.score, 1.0); + assert_eq!(weak.profile_input.request.model(), Some("cheap/model")); + assert!(observed(&weak_calls).is_empty()); + + let (strong_profile, strong_calls) = profile(Arc::new(FixedScorer(0.0)), false, false)?; + let strong = strong_profile.process(input("strong task")).await?; + assert_eq!(strong.decision.tier, TIER_STRONG); + assert_eq!(strong.decision.score, 0.0); + assert_eq!(strong.profile_input.request.model(), Some("frontier/model")); + assert!(observed(&strong_calls).is_empty()); + Ok(()) + } + + #[tokio::test] + async fn successful_decision_is_cached_by_first_string_user_instruction() -> Result<()> { + let scorer_calls = Arc::new(AtomicUsize::new(0)); + let (profile, _calls) = profile( + Arc::new(CountingScorer { + score: 1.0, + calls: scorer_calls.clone(), + }), + false, + false, + )?; + + let first = profile.process(input("same instruction")).await?; + let second = profile.process(input("same instruction")).await?; + + assert_eq!(first.decision.tier, TIER_WEAK); + assert_eq!(second.decision.tier, TIER_WEAK); + assert_eq!(scorer_calls.load(Ordering::SeqCst), 1); + Ok(()) + } + + #[tokio::test] + async fn probe_failure_forces_uncached_strong_then_retries_and_caches_success() -> Result<()> { + let scorer_calls = Arc::new(AtomicUsize::new(0)); + let (profile, calls) = profile( + Arc::new(FlakyScorer { + calls: scorer_calls.clone(), + }), + false, + false, + )?; + + let fallback = profile.process(input("retry task")).await?; + let retry = profile.process(input("retry task")).await?; + let cached = profile.process(input("retry task")).await?; + + assert_eq!(fallback.decision.tier, TIER_STRONG); + assert_eq!(fallback.decision.score, 0.0); + assert_eq!(fallback.decision.selected_model, "frontier/model"); + assert_eq!(retry.decision.tier, TIER_WEAK); + assert_eq!(cached.decision.tier, TIER_WEAK); + assert_eq!(scorer_calls.load(Ordering::SeqCst), 2); + assert!(observed(&calls).is_empty()); + Ok(()) + } + + #[tokio::test] + async fn every_probe_failure_is_retried() -> Result<()> { + let (profile, _calls) = profile(Arc::new(ErrorScorer), false, false)?; + + let first = profile.process(input("unavailable probe")).await?; + let second = profile.process(input("unavailable probe")).await?; + + assert_eq!(first.decision.tier, TIER_STRONG); + assert_eq!(second.decision.tier, TIER_STRONG); + assert_eq!(profile.decision_cache.lock().len(), 0); + Ok(()) + } + + #[tokio::test] + async fn run_calls_selected_backend_and_returns_routing_metadata() -> Result<()> { + let (weak_profile, weak_calls) = profile(Arc::new(FixedScorer(1.0)), false, false)?; + + let response = weak_profile.run(input("route weak")).await?; + + let weak_observed = observed(&weak_calls); + assert_eq!(weak_observed.len(), 1); + assert_eq!(weak_observed[0].backend, "weak-backend"); + assert_eq!(weak_observed[0].body["model"], "cheap/model"); + let metadata = response.routing_metadata.ok_or_else(|| { + SwitchyardError::Other("routing metadata should be present".to_string()) + })?; + assert_eq!(metadata.selected_model.as_deref(), Some("cheap/model")); + assert_eq!(metadata.selected_tier.as_deref(), Some(TIER_WEAK)); + assert_eq!(metadata.confidence, None); + assert_eq!(metadata.router_version.as_deref(), Some("prefill-probe:v1")); + assert_eq!(metadata.tolerance, Some(0.5)); + assert!(metadata + .rationale + .as_deref() + .is_some_and(|rationale| rationale.contains("score 1 selected weak"))); + + let (strong_profile, strong_calls) = profile(Arc::new(FixedScorer(0.0)), false, false)?; + let response = strong_profile.run(input("route strong")).await?; + let strong_observed = observed(&strong_calls); + assert_eq!(strong_observed.len(), 1); + assert_eq!(strong_observed[0].backend, "strong-backend"); + assert_eq!(strong_observed[0].body["model"], "frontier/model"); + let metadata = response.routing_metadata.ok_or_else(|| { + SwitchyardError::Other("strong routing metadata should be present".to_string()) + })?; + assert_eq!(metadata.selected_model.as_deref(), Some("frontier/model")); + assert_eq!(metadata.selected_tier.as_deref(), Some(TIER_STRONG)); + Ok(()) + } + + #[tokio::test] + async fn run_records_usage_for_selected_tier() -> Result<()> { + let (profile, _calls) = profile(Arc::new(FixedScorer(1.0)), false, false)?; + + let _response = profile.run(input("record weak")).await?; + + let snapshot = profile.stats.snapshot()?; + assert_eq!(snapshot.total_requests, 1); + assert_eq!(snapshot.total_tokens.prompt, 5); + assert_eq!(snapshot.total_tokens.completion, 3); + let tier = snapshot.tiers.get(TIER_WEAK).ok_or_else(|| { + SwitchyardError::Other("weak tier stats should be present".to_string()) + })?; + assert_eq!(tier.calls, 1); + assert_eq!(tier.model, "cheap/model"); + Ok(()) + } + + #[tokio::test] + async fn backend_failure_is_attributed_to_selected_model_and_tier() -> Result<()> { + let (profile, calls) = profile(Arc::new(FixedScorer(1.0)), false, true)?; + + let error = profile + .run(input("weak backend fails")) + .await + .err() + .ok_or_else(|| { + SwitchyardError::Other("backend failure should be returned".to_string()) + })?; + + assert!(format!("{error}").contains("weak-backend failed")); + assert_eq!(observed(&calls).len(), 1); + let snapshot = profile.stats.snapshot()?; + assert_eq!(snapshot.total_requests, 1); + assert_eq!(snapshot.total_errors, 1); + let model = snapshot.models.get("cheap/model").ok_or_else(|| { + SwitchyardError::Other("weak model error stats should be present".to_string()) + })?; + assert_eq!(model.errors, 1); + assert_eq!(model.tier.as_deref(), Some(TIER_WEAK)); + Ok(()) + } +} diff --git a/crates/switchyard-components-v2/src/profiles/profile_types.rs b/crates/switchyard-components-v2/src/profiles/profile_types.rs index 3126a815..7ecab2f8 100644 --- a/crates/switchyard-components-v2/src/profiles/profile_types.rs +++ b/crates/switchyard-components-v2/src/profiles/profile_types.rs @@ -6,7 +6,8 @@ use super::macros::profile_types; use super::{ LatencyServiceProfileConfig, LlmRoutingProfileConfig, NoopProfileConfig, - PassthroughProfileConfig, RandomRoutingProfileConfig, StageRouterProfileConfig, + PassthroughProfileConfig, PrefillProbeProfileConfig, RandomRoutingProfileConfig, + StageRouterProfileConfig, }; profile_types! { @@ -15,5 +16,6 @@ profile_types! { RandomRoutingProfileConfig, LatencyServiceProfileConfig, LlmRoutingProfileConfig, + PrefillProbeProfileConfig, NoopProfileConfig, } diff --git a/crates/switchyard-components-v2/tests/prefill_probe_profile.rs b/crates/switchyard-components-v2/tests/prefill_probe_profile.rs new file mode 100644 index 00000000..38f6b991 --- /dev/null +++ b/crates/switchyard-components-v2/tests/prefill_probe_profile.rs @@ -0,0 +1,467 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Public configuration and startup-validation tests for the learned prefill-probe profile. + +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicU64, Ordering}; + +use safetensors::tensor::{serialize_to_file, Dtype, TensorView}; +use serde_json::{json, Value}; +use switchyard_components_v2::{ + parse_profile_config_str, PrefillProbeProfileConfig, PrefillProbeRoutingPolicyConfig, + ProfileConfig, ProfileConfigFormat, +}; +use switchyard_core::{ + BackendFormat, LlmTarget, LlmTargetId, ModelId, ProfileId, Result, SwitchyardError, +}; + +static NEXT_TEST_DIRECTORY: AtomicU64 = AtomicU64::new(0); + +#[derive(Clone, Copy)] +enum ArtifactFault { + None, + EncoderMismatch, + MissingTensor, + WrongShape, + MalformedTensorFile, +} + +struct TestTensor { + name: String, + shape: Vec, + bytes: Vec, +} + +struct TestArtifactDirectory { + path: PathBuf, +} + +impl TestArtifactDirectory { + fn create(fault: ArtifactFault) -> Result { + let sequence = NEXT_TEST_DIRECTORY.fetch_add(1, Ordering::Relaxed); + let path = std::env::temp_dir().join(format!( + "switchyard-prefill-profile-{}-{sequence}", + std::process::id(), + )); + std::fs::create_dir(&path).map_err(|error| { + SwitchyardError::Other(format!( + "failed to create test directory {}: {error}", + path.display() + )) + })?; + let directory = Self { path }; + directory.write_artifact(fault)?; + Ok(directory) + } + + fn path(&self) -> &Path { + &self.path + } + + fn write_artifact(&self, fault: ArtifactFault) -> Result<()> { + let encoder = if matches!(fault, ArtifactFault::EncoderMismatch) { + "different/probe" + } else { + "probe/model" + }; + let metadata = json!({ + "format_version": 1, + "training_mode": "single_pca_block", + "encoder": encoder, + "representation": "token_mean_per_layer_concat", + "extraction_layer_ids": [0, 1], + "hidden_size": 2, + "raw_feature_dim": 4, + "feature_block_count": 1, + "pca_dim": 200, + "pca_whiten": false, + "output_names": ["qwen-122b", "nemotron-3-super", "opus-4.7", "gpt-5.5"], + "trunk_hidden": [256, 128], + "ensemble_size": 5, + "probability_link": "independent_sigmoid", + "ensemble_reduction": "probability_mean", + "tensor_file": "router.safetensors", + }); + let metadata_bytes = serde_json::to_vec(&metadata) + .map_err(|error| SwitchyardError::Other(format!("metadata encode failed: {error}")))?; + std::fs::write(self.path.join("router.json"), metadata_bytes) + .map_err(|error| SwitchyardError::Other(format!("metadata write failed: {error}")))?; + + if matches!(fault, ArtifactFault::MalformedTensorFile) { + std::fs::write(self.path.join("router.safetensors"), b"not safetensors").map_err( + |error| SwitchyardError::Other(format!("malformed tensor write failed: {error}")), + )?; + return Ok(()); + } + + let mut tensors = test_artifact_tensors()?; + if matches!(fault, ArtifactFault::MissingTensor) { + tensors.retain(|tensor| tensor.name != "transform.scaler_mean"); + } else if matches!(fault, ArtifactFault::WrongShape) { + tensors[0] = test_tensor("transform.scaler_mean", vec![3], 0.0)?; + } + let mut views = Vec::with_capacity(tensors.len()); + for tensor in &tensors { + let view = TensorView::new(Dtype::F32, tensor.shape.clone(), &tensor.bytes).map_err( + |error| { + SwitchyardError::Other(format!("tensor {} is invalid: {error}", tensor.name)) + }, + )?; + views.push((tensor.name.as_str(), view)); + } + serialize_to_file( + views, + &Some(HashMap::new()), + &self.path.join("router.safetensors"), + ) + .map_err(|error| SwitchyardError::Other(format!("artifact encode failed: {error}")))?; + Ok(()) + } +} + +impl Drop for TestArtifactDirectory { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.path); + } +} + +fn test_tensor(name: impl Into, shape: Vec, value: f32) -> Result { + let value_count = shape + .iter() + .try_fold(1usize, |count, size| count.checked_mul(*size)) + .ok_or_else(|| SwitchyardError::Other("test tensor shape overflow".to_string()))?; + let bytes = (0..value_count).flat_map(|_| value.to_le_bytes()).collect(); + Ok(TestTensor { + name: name.into(), + shape, + bytes, + }) +} + +fn test_artifact_tensors() -> Result> { + let mut tensors = vec![ + test_tensor("transform.scaler_mean", vec![4], 0.0)?, + test_tensor("transform.scaler_scale", vec![4], 1.0)?, + test_tensor("transform.pca_mean", vec![4], 0.0)?, + test_tensor("transform.pca_components", vec![200, 4], 0.0)?, + ]; + for index in 0..5 { + let prefix = format!("ensemble.{index}"); + tensors.extend([ + test_tensor(format!("{prefix}.linear1.weight"), vec![256, 200], 0.0)?, + test_tensor(format!("{prefix}.linear1.bias"), vec![256], 0.0)?, + test_tensor(format!("{prefix}.linear2.weight"), vec![128, 256], 0.0)?, + test_tensor(format!("{prefix}.linear2.bias"), vec![128], 0.0)?, + test_tensor(format!("{prefix}.output.weight"), vec![4, 128], 0.0)?, + test_tensor(format!("{prefix}.output.bias"), vec![4], 0.0)?, + ]); + } + Ok(tensors) +} + +fn target(id: &str, model: &str) -> Result { + let mut target = LlmTarget::new(LlmTargetId::new(id)?, ModelId::new(model)?); + target.format = BackendFormat::OpenAi; + Ok(target) +} + +fn routing_policy() -> PrefillProbeRoutingPolicyConfig { + PrefillProbeRoutingPolicyConfig::CostAware { + lambda: 0.5, + weak_cost: 0.01, + strong_cost: 0.10, + } +} + +fn config_with_artifact(artifact_dir: &Path) -> Result { + Ok(PrefillProbeProfileConfig { + probe: target("probe", "probe/model")?, + strong: target("strong", "frontier/model")?, + strong_checkpoint_head: "opus-4.7".to_string(), + weak: target("weak", "cheap/model")?, + weak_checkpoint_head: "nemotron-3-super".to_string(), + hidden_states_dir: "/dev/shm/hidden_states".to_string(), + inference_artifact_dir: artifact_dir.to_string_lossy().into_owned(), + routing_policy: routing_policy(), + }) +} + +fn assert_build_error(fault: ArtifactFault, expected: &str) -> Result<()> { + let artifact = TestArtifactDirectory::create(fault)?; + let error = config_with_artifact(artifact.path())? + .build() + .err() + .ok_or_else(|| SwitchyardError::Other("invalid artifact should fail".to_string()))?; + assert!(format!("{error}").contains(expected)); + Ok(()) +} + +#[test] +fn registered_yaml_resolves_and_builds_the_profile() -> Result<()> { + let artifact = TestArtifactDirectory::create(ArtifactFault::None)?; + let yaml = format!( + r#" +targets: + probe: + model: probe/model + format: openai + base_url: http://localhost:8000/v1 + strong: + model: frontier/model + format: openai + weak: + model: cheap/model + 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: /dev/shm/hidden_states + inference_artifact_dir: {} + routing_policy: + type: cost-aware + lambda: 0.5 + weak_cost: 0.01 + strong_cost: 0.10 +"#, + artifact.path().display(), + ); + + let plan = parse_profile_config_str(&yaml, ProfileConfigFormat::Yaml)?.resolve()?; + let profile_id = ProfileId::new("router")?; + assert_eq!(plan.profile_type(&profile_id), Some("prefill-probe")); + assert_eq!(plan.target_count(), 3); + let probe = plan + .target(&LlmTargetId::new("probe")?) + .ok_or_else(|| SwitchyardError::Other("probe target should resolve".to_string()))?; + assert_eq!(probe.model.as_str(), "probe/model"); + let _profile = plan.build_profile(&profile_id)?; + Ok(()) +} + +#[test] +fn probe_can_differ_from_both_completion_targets() -> Result<()> { + let artifact = TestArtifactDirectory::create(ArtifactFault::None)?; + let config = config_with_artifact(artifact.path())?; + + assert_eq!(config.probe.model.as_str(), "probe/model"); + assert_eq!(config.strong.model.as_str(), "frontier/model"); + assert_eq!(config.weak.model.as_str(), "cheap/model"); + let _profile = config.build()?; + Ok(()) +} + +#[test] +fn config_schema_requires_every_field_and_rejects_unknown_fields() -> Result<()> { + let artifact = TestArtifactDirectory::create(ArtifactFault::None)?; + let config = config_with_artifact(artifact.path())?; + assert_eq!(PrefillProbeProfileConfig::PROFILE_TYPE, "prefill-probe"); + assert_eq!(config.profile_type(), "prefill-probe"); + let base = serde_json::to_value(&config) + .map_err(|error| SwitchyardError::Other(format!("config encode failed: {error}")))?; + + for field in [ + "probe", + "strong", + "strong_checkpoint_head", + "weak", + "weak_checkpoint_head", + "hidden_states_dir", + "inference_artifact_dir", + "routing_policy", + ] { + let mut missing = base.clone(); + missing + .as_object_mut() + .ok_or_else(|| SwitchyardError::Other("config should encode as an object".to_string()))? + .remove(field); + let error = serde_json::from_value::(missing) + .err() + .ok_or_else(|| SwitchyardError::Other(format!("missing {field} should fail")))?; + assert!(error + .to_string() + .contains(&format!("missing field `{field}`"))); + } + + let mut unknown = base; + unknown + .as_object_mut() + .ok_or_else(|| SwitchyardError::Other("config should encode as an object".to_string()))? + .insert("unexpected_field".to_string(), Value::Bool(true)); + let error = serde_json::from_value::(unknown) + .err() + .ok_or_else(|| SwitchyardError::Other("unknown profile field should fail".to_string()))?; + assert!(error + .to_string() + .contains("unknown field `unexpected_field`")); + Ok(()) +} + +#[test] +fn policy_schema_is_tagged_and_strict() -> Result<()> { + let valid = serde_json::to_value(routing_policy()) + .map_err(|error| SwitchyardError::Other(format!("policy encode failed: {error}")))?; + let parsed = serde_json::from_value::(valid.clone()) + .map_err(|error| SwitchyardError::Other(format!("policy parse failed: {error}")))?; + assert_eq!(parsed, routing_policy()); + + let mut missing_tag = valid.clone(); + missing_tag + .as_object_mut() + .ok_or_else(|| SwitchyardError::Other("policy should encode as an object".to_string()))? + .remove("type"); + let missing_tag_error = serde_json::from_value::(missing_tag) + .err() + .ok_or_else(|| SwitchyardError::Other("missing policy type should fail".to_string()))?; + assert!(missing_tag_error.to_string().contains("type")); + + let mut unknown = valid; + unknown + .as_object_mut() + .ok_or_else(|| SwitchyardError::Other("policy should encode as an object".to_string()))? + .insert("threshold".to_string(), json!(0.5)); + let unknown_error = serde_json::from_value::(unknown) + .err() + .ok_or_else(|| SwitchyardError::Other("unknown policy field should fail".to_string()))?; + assert!(unknown_error + .to_string() + .contains("unknown field `threshold`")); + Ok(()) +} + +#[test] +fn legacy_routing_knobs_are_rejected() -> Result<()> { + let artifact = TestArtifactDirectory::create(ArtifactFault::None)?; + let base = serde_json::to_value(config_with_artifact(artifact.path())?) + .map_err(|error| SwitchyardError::Other(format!("config encode failed: {error}")))?; + + for (field, value) in [ + ("confidence_threshold", json!(0.5)), + ("probe_signal", json!("entropy")), + ] { + let mut config = base.clone(); + config + .as_object_mut() + .ok_or_else(|| SwitchyardError::Other("config should encode as an object".to_string()))? + .insert(field.to_string(), value); + let error = serde_json::from_value::(config) + .err() + .ok_or_else(|| SwitchyardError::Other(format!("legacy {field} should fail")))?; + assert!(error.to_string().contains("unknown field")); + assert!(error.to_string().contains(field)); + } + Ok(()) +} + +#[test] +fn artifact_encoder_and_tensor_failures_are_reported_at_build() -> Result<()> { + assert_build_error(ArtifactFault::EncoderMismatch, "does not match probe model")?; + assert_build_error( + ArtifactFault::MissingTensor, + "missing tensor transform.scaler_mean", + )?; + assert_build_error(ArtifactFault::WrongShape, "has shape [3]; expected [4]")?; + assert_build_error(ArtifactFault::MalformedTensorFile, "failed to parse")?; + Ok(()) +} + +#[test] +fn missing_artifact_directory_is_reported_at_build() -> Result<()> { + let missing = std::env::temp_dir().join(format!( + "switchyard-missing-prefill-artifact-{}", + NEXT_TEST_DIRECTORY.fetch_add(1, Ordering::Relaxed), + )); + let error = config_with_artifact(&missing)? + .build() + .err() + .ok_or_else(|| SwitchyardError::Other("missing artifact should fail".to_string()))?; + assert!(format!("{error}").contains("failed to read")); + assert!(format!("{error}").contains("router.json")); + Ok(()) +} + +#[test] +fn checkpoint_head_mappings_are_validated_at_build() -> Result<()> { + let artifact = TestArtifactDirectory::create(ArtifactFault::None)?; + let mut unknown = config_with_artifact(artifact.path())?; + unknown.weak_checkpoint_head = "unknown-head".to_string(); + let unknown_error = unknown + .build() + .err() + .ok_or_else(|| SwitchyardError::Other("unknown head should fail".to_string()))?; + assert!(format!("{unknown_error}").contains("weak_checkpoint_head `unknown-head`")); + + let mut duplicate = config_with_artifact(artifact.path())?; + duplicate.weak_checkpoint_head = duplicate.strong_checkpoint_head.clone(); + let duplicate_error = duplicate + .build() + .err() + .ok_or_else(|| SwitchyardError::Other("duplicate head should fail".to_string()))?; + assert!(format!("{duplicate_error}").contains("must map to distinct outputs")); + Ok(()) +} + +#[test] +fn invalid_policy_values_are_reported_before_artifact_loading() -> Result<()> { + let cases = [ + ( + PrefillProbeRoutingPolicyConfig::CostAware { + lambda: 1.5, + weak_cost: 0.01, + strong_cost: 0.10, + }, + "lambda", + ), + ( + PrefillProbeRoutingPolicyConfig::CostAware { + lambda: 0.5, + weak_cost: -0.01, + strong_cost: 0.10, + }, + "weak_cost", + ), + ( + PrefillProbeRoutingPolicyConfig::CostAware { + lambda: 0.5, + weak_cost: 0.01, + strong_cost: f64::INFINITY, + }, + "strong_cost", + ), + ]; + + for (routing_policy, expected) in cases { + let config = PrefillProbeProfileConfig { + probe: target("probe", "probe/model")?, + strong: target("strong", "frontier/model")?, + strong_checkpoint_head: "opus-4.7".to_string(), + weak: target("weak", "cheap/model")?, + weak_checkpoint_head: "nemotron-3-super".to_string(), + hidden_states_dir: "/dev/shm/hidden_states".to_string(), + inference_artifact_dir: "/missing/artifact".to_string(), + routing_policy, + }; + let error = config + .build() + .err() + .ok_or_else(|| SwitchyardError::Other(format!("invalid {expected} should fail")))?; + assert!(format!("{error}").contains(expected)); + } + Ok(()) +} + +#[test] +fn loaded_artifact_is_owned_after_source_files_are_removed() -> Result<()> { + let artifact = TestArtifactDirectory::create(ArtifactFault::None)?; + let _profile = config_with_artifact(artifact.path())?.build()?; + std::fs::remove_dir_all(artifact.path()).map_err(|error| { + SwitchyardError::Other(format!("failed to remove loaded artifact: {error}")) + })?; + Ok(()) +} From 9fb64059fc698910a4c7a69ef89679abe0ccb90a Mon Sep 17 00:00:00 2001 From: michxu Date: Fri, 17 Jul 2026 17:11:20 +0000 Subject: [PATCH 06/12] docs(prefill-probe): add deployment guide and sample config Signed-off-by: michxu --- .../routing-profiles/prefill-probe-local.yaml | 50 ++++++ docs/vllm-serve-hidden-state.md | 160 +++++++++++++++--- 2 files changed, 189 insertions(+), 21 deletions(-) create mode 100644 benchmark/routing-profiles/prefill-probe-local.yaml diff --git a/benchmark/routing-profiles/prefill-probe-local.yaml b/benchmark/routing-profiles/prefill-probe-local.yaml new file mode 100644 index 00000000..86a81648 --- /dev/null +++ b/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_ARTIFACT_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 + inference_artifact_dir: ${PREFILL_ROUTER_ARTIFACT_DIR} + routing_policy: + type: cost-aware + lambda: 0.5 + weak_cost: 0 + strong_cost: 1 diff --git a/docs/vllm-serve-hidden-state.md b/docs/vllm-serve-hidden-state.md index 2e2a8112..77330dac 100644 --- a/docs/vllm-serve-hidden-state.md +++ b/docs/vllm-serve-hidden-state.md @@ -1,61 +1,157 @@ # vLLM Hidden-State Serving for Local Models -This page captures operational notes for serving a local vLLM model with hidden-state extraction enabled. The hidden-state connector writes prefill activations to `.safetensors` files and returns the actual file path in `kv_transfer_params.hidden_states_path`. +This page describes the hidden-state deployment contract for the learned +`prefill-probe` profile. The hidden-state connector writes prefill activations +to `.safetensors` files and returns the concrete file path in +`kv_transfer_params.hidden_states_path`. -Docker is not required by the protocol. Use Docker when you want a reproducible CUDA/vLLM runtime; use `vllm serve` directly when the local Python environment has a vLLM build that includes `extract_hidden_states` and `ExampleHiddenStatesConnector`. +The example profile is +`benchmark/routing-profiles/prefill-probe-local.yaml`. It separates three +roles: + +- `probe`: `Qwen/Qwen3.6-35B-A3B`, used only to extract prompt hidden states; +- `weak`: the completion target mapped to the `nemotron-3-super` artifact head; +- `strong`: the completion target mapped to the `opus-4.7` artifact head. + +The probe does not need to be one of the completion targets. Its model string +must instead match the artifact's `encoder` metadata exactly. Docker is not +required by the protocol. Use Docker for a reproducible CUDA/vLLM runtime, or +use `vllm serve` directly when the environment contains a build with +`extract_hidden_states` and `ExampleHiddenStatesConnector`. + +## Router artifact contract + +The router artifact is external deployment data. Switchyard does not package +or download it. Set `PREFILL_ROUTER_ARTIFACT_DIR` to a readable directory with: + +```text +router.json +router.safetensors +``` + +At profile startup, Switchyard validates the artifact metadata, tensor names, +and tensor shapes. The current artifact contract uses Qwen layers `0..39` in +that exact order, hidden size 2048, and one feature block. For each prompt: + +1. vLLM exports a `hidden_states` tensor shaped + `[prompt_tokens, 40, 2048]`. +2. Switchyard averages over the prompt-token dimension independently for every + layer and hidden dimension. +3. The 40 layer vectors are concatenated in layer order into one 81,920-value + vector. This is one feature block, not one block per completion target. +4. The artifact's fitted scaler and PCA transform reduce that vector to 200 + values. +5. Five learned `200 -> 256 -> 128 -> 4` trunk members produce four + correctness probabilities. Switchyard averages probabilities across the + members and reads only the two heads named by `weak_checkpoint_head` and + `strong_checkpoint_head`. + +The scaler and PCA are fitted training artifacts, not fitted online. No +checkpoint training or online learning occurs in Switchyard. + +## Routing policy + +The profile min-max normalizes the two configured costs, then computes: + +```text +weak_utility = lambda * P(weak correct) - (1 - lambda) * normalized_weak_cost +strong_utility = lambda * P(strong correct) - (1 - lambda) * normalized_strong_cost +margin = weak_utility - strong_utility +``` + +A non-negative margin becomes public score `1.0` and selects weak. A negative +margin becomes public score `0.0` and selects strong. The learned profile's +threshold is fixed at `0.5`; tune routing only with `lambda`. The sample uses +`lambda: 0.5`, `weak_cost: 0`, and `strong_cost: 1`. + +Successful decisions are cached by a hash of the first string-valued user +instruction. A cache hit reuses the selected tier without another probe. Probe, +artifact, or scoring failures route to strong and are not cached, so a later +matching request retries the probe. ## Docker launch -Pick one filesystem path for hidden states and mount it into the container. The container path used in `shared_storage_path` must be the same path clients pass as `kv_transfer_params.hidden_states_path`. +Pick one filesystem path for hidden states and mount it into the container. The +path configured in vLLM's `shared_storage_path` must be the same absolute path +as the profile's `hidden_states_dir`. Both vLLM and Switchyard need access; +vLLM must create files there, and Switchyard must read, lock, and delete them. + +The all-layer export was smoke-tested with vLLM +`0.23.1rc1.dev672+g93d8f834d`. The command below pins the tested container +digest and runs vLLM with the host user's UID and GID. Matching the host +identity is required because vLLM creates the hidden-state files with mode +`0600`; a container-owned file cannot be read or deleted by host-side +Switchyard. ```bash -export HIDDEN_STATES_DIR=/tmp/vllm-hidden-states +export HIDDEN_STATES_DIR=/tmp/switchyard-hidden-states export HF_CACHE_DIR=/tmp/vllm-hf-cache +export HOST_UID="$(id -u)" +export HOST_GID="$(id -g)" +export HOST_USER="${USER:-switchyard}" mkdir -p "${HIDDEN_STATES_DIR}" "${HF_CACHE_DIR}" -docker run -d --name vllm_qwen35 \ +docker run -d --name vllm_qwen36 \ --gpus all \ --ipc=host \ + --user "${HOST_UID}:${HOST_GID}" \ + -e USER="${HOST_USER}" \ + -e LOGNAME="${HOST_USER}" \ + -e HOME=/tmp/vllm-home \ + -e HF_HOME=/model-cache \ -p 0.0.0.0:8000:8000 \ - -v "${HF_CACHE_DIR}:/root/.cache/huggingface" \ + -v "${HF_CACHE_DIR}:/model-cache" \ -v "${HIDDEN_STATES_DIR}:${HIDDEN_STATES_DIR}" \ - vllm/vllm-openai:latest-cu129 \ + vllm/vllm-openai@sha256:3b7bb15f9f2b13f2f508d94d1900ea40b5be9f96d716ad977bcef742dac464bc \ Qwen/Qwen3.6-35B-A3B \ - --tensor-parallel-size 8 \ - --max-model-len 32768 \ + --tensor-parallel-size 2 \ + --dtype bfloat16 \ + --max-num-seqs 1 \ + --max-model-len 16384 \ + --gpu-memory-utilization 0.90 \ + --trust-remote-code \ + --language-model-only \ --reasoning-parser qwen3 \ --enable-auto-tool-choice \ --tool-call-parser hermes \ --no-enable-chunked-prefill \ - --speculative-config '{"method":"extract_hidden_states","num_speculative_tokens":1,"draft_model_config":{"hf_config":{"eagle_aux_hidden_state_layer_ids":[39]}}}' \ - --kv-transfer-config '{"kv_connector":"ExampleHiddenStatesConnector","kv_role":"kv_producer","kv_connector_extra_config":{"shared_storage_path":"/tmp/vllm-hidden-states"}}' + --speculative-config '{"method":"extract_hidden_states","num_speculative_tokens":1,"draft_model_config":{"hf_config":{"eagle_aux_hidden_state_layer_ids":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39]}}}' \ + --kv-transfer-config '{"kv_connector":"ExampleHiddenStatesConnector","kv_role":"kv_producer","kv_connector_extra_config":{"shared_storage_path":"/tmp/switchyard-hidden-states"}}' ``` Use host IPC for long-running Docker jobs. The default Docker IPC mode gives the container a private 64 MiB `/dev/shm`, which can starve vLLM's tensor-parallel shared-memory broadcast path while hidden-state extraction is enabled. -For `Qwen/Qwen3.6-35B-A3B`, layer `39` is the last hidden-state layer. To capture multiple layers, add each layer id to `eagle_aux_hidden_state_layer_ids`, for example `[0,1,2,39]`. Capturing all layers can make each probe much larger and may require a lower `--max-model-len` to leave enough KV-cache memory. +Do not shorten or reorder `eagle_aux_hidden_state_layer_ids` for this router +artifact. Switchyard validates the number and width of exported layers, but the +file format does not carry layer IDs for it to verify. A different order can +therefore produce a valid shape with incorrect features and routes. ## Direct vLLM CLI launch The direct CLI form serves the same model without Docker. There is no volume mount; `shared_storage_path` is a host path and clients must be able to read that same path. ```bash -export HIDDEN_STATES_DIR=/tmp/vllm-hidden-states +export HIDDEN_STATES_DIR=/tmp/switchyard-hidden-states mkdir -p "${HIDDEN_STATES_DIR}" vllm serve Qwen/Qwen3.6-35B-A3B \ --host 0.0.0.0 \ --port 8000 \ - --tensor-parallel-size 8 \ - --max-model-len 32768 \ + --tensor-parallel-size 2 \ + --dtype bfloat16 \ + --max-num-seqs 1 \ + --max-model-len 16384 \ + --gpu-memory-utilization 0.90 \ + --trust-remote-code \ + --language-model-only \ --reasoning-parser qwen3 \ --enable-auto-tool-choice \ --tool-call-parser hermes \ --no-enable-chunked-prefill \ - --speculative-config '{"method":"extract_hidden_states","num_speculative_tokens":1,"draft_model_config":{"hf_config":{"eagle_aux_hidden_state_layer_ids":[39]}}}' \ - --kv-transfer-config '{"kv_connector":"ExampleHiddenStatesConnector","kv_role":"kv_producer","kv_connector_extra_config":{"shared_storage_path":"/tmp/vllm-hidden-states"}}' + --speculative-config '{"method":"extract_hidden_states","num_speculative_tokens":1,"draft_model_config":{"hf_config":{"eagle_aux_hidden_state_layer_ids":[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,32,33,34,35,36,37,38,39]}}}' \ + --kv-transfer-config '{"kv_connector":"ExampleHiddenStatesConnector","kv_role":"kv_producer","kv_connector_extra_config":{"shared_storage_path":"/tmp/switchyard-hidden-states"}}' ``` Use the direct CLI only after confirming your installed vLLM accepts both `--speculative-config '{"method":"extract_hidden_states",...}'` and `--kv-transfer-config '{"kv_connector":"ExampleHiddenStatesConnector",...}'`. If those flags fail, use the known container image or install a vLLM build that contains the connector. @@ -72,7 +168,7 @@ curl http://localhost:8000/v1/chat/completions \ "messages": [{"role": "user", "content": "Return one short sentence."}], "max_tokens": 1, "kv_transfer_params": { - "hidden_states_path": "/tmp/vllm-hidden-states", + "hidden_states_path": "/tmp/switchyard-hidden-states", "include_output_tokens": false } }' @@ -81,11 +177,11 @@ curl http://localhost:8000/v1/chat/completions \ Read the path from the response rather than assuming a filename. vLLM may choose the concrete safetensors file name. ```bash -uv run python - <<'PY_INNER' +uv run --with safetensors python - <<'PY_INNER' from pathlib import Path from safetensors import safe_open -path = Path("/tmp/vllm-hidden-states") +path = Path("/tmp/switchyard-hidden-states") files = sorted(path.glob("*.safetensors"), key=lambda item: item.stat().st_mtime) if not files: raise SystemExit("no safetensors files written") @@ -97,10 +193,32 @@ with safe_open(files[-1], framework="numpy") as handle: PY_INNER ``` +The inspection command leaves the file in place. During normal routing, +Switchyard deletes a hidden-state file only after it has locked, read, and +successfully decoded it. + +## Latency and capacity implications + +Every uncached instruction adds a Qwen prefill, a 40-layer hidden-state export, +filesystem write/read, token mean pooling, scaler/PCA projection, and trunk +inference before the selected completion begins. Mean pooling does more CPU +work than selecting the last prompt token, but it does not add another model +forward pass because vLLM already exports the prompt activations. Exporting all +40 layers and moving the artifact through the filesystem are usually the larger +incremental costs. + +Artifact size grows with prompt length. For BF16 Qwen features, the +`hidden_states` payload alone is approximately +`prompt_tokens * 40 * 2048 * 2` bytes. Budget shared-storage capacity and lower +`--max-model-len` if hidden-state extraction leaves insufficient KV-cache +memory. + ## Troubleshooting - `probe response missing kv_transfer_params`: the server is not running with `ExampleHiddenStatesConnector`, or the request did not include `kv_transfer_params`. - `no safetensors files written`: check that `shared_storage_path` exists and is writable by the vLLM process. +- Artifact encoder mismatch: use the exact probe model named by `router.json`; changing probe checkpoints requires a matching exported router artifact. +- Layer-count or raw-feature-dimension mismatch: export all layers `0..39` in order and confirm the hidden width is 2048. - `No available shared memory broadcast block found` followed by `RPC call to sample_tokens timed out`: relaunch the Docker container with `--ipc=host`. -- Context-length startup errors from vLLM: lower `--max-model-len`, reduce the number of captured layers, or increase available GPU memory. +- Context-length startup errors from vLLM: lower `--max-model-len` or increase available GPU memory. Do not reduce the captured layers without exporting a matching router artifact. - Hidden-state extraction does not work with chunked prefill; keep `--no-enable-chunked-prefill` in the launch command. From 8ea2bd2a67bb0204a17a7bcdd0009c1c56250740 Mon Sep 17 00:00:00 2001 From: michxu Date: Mon, 20 Jul 2026 18:12:47 +0000 Subject: [PATCH 07/12] feat(prefill-probe): accept explicit probe input Signed-off-by: michxu --- .../src/profiles/prefill_router/profile.rs | 263 +++++++++++++++++- .../src/profiles/prefill_router/scorer.rs | 40 +-- 2 files changed, 259 insertions(+), 44 deletions(-) diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs index ab70c635..dcebcf2b 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs @@ -27,6 +27,7 @@ use super::policy::CostAwareRoutingPolicy; use super::scorer::{HiddenStateProbeScorer, ProbeScorer}; const LEARNED_SCORE_THRESHOLD: f64 = 0.5; +const PREFILL_PROBE_INPUT_FIELD: &str = "_switchyard_prefill_probe_input"; const TIER_STRONG: &str = "strong"; const TIER_WEAK: &str = "weak"; @@ -177,32 +178,60 @@ pub struct PrefillProbeProfile { score_threshold: f64, scorer: Arc, stats: StatsAccumulator, - // Successful scores are reused by the first string-valued user instruction. + // Successful scores are reused by the resolved probe input. decision_cache: Arc>>, } impl PrefillProbeProfile { - // Hashes the first user instruction represented as a string. - fn instruction_key(request: &switchyard_core::ChatRequest) -> Option { + // Returns the first string-valued user instruction. + fn first_user_instruction(request: &switchyard_core::ChatRequest) -> Option<&str> { let messages = request.body().as_object()?.get("messages")?.as_array()?; - let content = messages.iter().find_map(|message| { + messages.iter().find_map(|message| { (message.get("role").and_then(serde_json::Value::as_str) == Some("user")) .then(|| message.get("content").and_then(serde_json::Value::as_str)) .flatten() - })?; + }) + } + + // Removes the profile-local field and resolves the exact input to score and cache. + fn take_probe_input(request: &mut switchyard_core::ChatRequest) -> Result> { + let explicit = request + .body_mut() + .as_object_mut() + .and_then(|body| body.remove(PREFILL_PROBE_INPUT_FIELD)); + match explicit { + Some(serde_json::Value::String(probe_input)) if !probe_input.is_empty() => { + Ok(Some(probe_input)) + } + Some(_) => Err(SwitchyardError::InvalidRequest(format!( + "{PREFILL_PROBE_INPUT_FIELD} must be a non-empty string", + ))), + None => Ok(Self::first_user_instruction(request).map(str::to_owned)), + } + } + + // Hashes the resolved probe input for successful-decision reuse. + fn instruction_key(probe_input: &str) -> u64 { let mut hasher = std::collections::hash_map::DefaultHasher::new(); - content.hash(&mut hasher); - Some(hasher.finish()) + probe_input.hash(&mut hasher); + hasher.finish() } // Scores an uncached instruction and explicitly bypasses threshold routing on failure. async fn route(&self, mut input: ProfileInput) -> Result { - let key = Self::instruction_key(&input.request); + let probe_input = Self::take_probe_input(&mut input.request)?; + let key = probe_input.as_deref().map(Self::instruction_key); let cached = key.and_then(|key| self.decision_cache.lock().get(&key).copied()); let (score, strong_fallback) = if let Some(score) = cached { (score, false) } else { - match self.scorer.score(&input.request).await { + let scored = match probe_input.as_deref() { + Some(probe_input) => self.scorer.score(probe_input).await, + None => Err(SwitchyardError::Other( + "prefill probe request has no string-valued user instruction".to_string(), + )), + }; + match scored { Ok(score) => { if let Some(key) = key { self.decision_cache.lock().insert(key, score); @@ -392,7 +421,7 @@ mod tests { #[async_trait] impl ProbeScorer for FixedScorer { - async fn score(&self, _request: &ChatRequest) -> Result { + async fn score(&self, _probe_input: &str) -> Result { Ok(self.0) } } @@ -401,7 +430,7 @@ mod tests { #[async_trait] impl ProbeScorer for ErrorScorer { - async fn score(&self, _request: &ChatRequest) -> Result { + async fn score(&self, _probe_input: &str) -> Result { Err(SwitchyardError::Other("probe unavailable".to_string())) } } @@ -413,7 +442,7 @@ mod tests { #[async_trait] impl ProbeScorer for CountingScorer { - async fn score(&self, _request: &ChatRequest) -> Result { + async fn score(&self, _probe_input: &str) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); Ok(self.score) } @@ -425,7 +454,7 @@ mod tests { #[async_trait] impl ProbeScorer for FlakyScorer { - async fn score(&self, _request: &ChatRequest) -> Result { + async fn score(&self, _probe_input: &str) -> Result { let attempt = self.calls.fetch_add(1, Ordering::SeqCst); if attempt == 0 { Err(SwitchyardError::Other( @@ -437,6 +466,19 @@ mod tests { } } + struct RecordingScorer { + inputs: Arc>>, + score: f64, + } + + #[async_trait] + impl ProbeScorer for RecordingScorer { + async fn score(&self, probe_input: &str) -> Result { + self.inputs.lock().push(probe_input.to_string()); + Ok(self.score) + } + } + fn target(id: &str, model: &str) -> Result { let mut target = LlmTarget::new(LlmTargetId::new(id)?, ModelId::new(model)?); target.format = BackendFormat::OpenAi; @@ -488,6 +530,17 @@ mod tests { } } + fn input_with_probe_field(probe_input: Value, wrapped_instruction: &str) -> ProfileInput { + ProfileInput { + request: ChatRequest::openai_chat(json!({ + "model": "client/model", + "messages": [{"role": "user", "content": wrapped_instruction}], + (PREFILL_PROBE_INPUT_FIELD): probe_input, + })), + metadata: RequestMetadata::default(), + } + } + fn observed(calls: &Arc>>) -> Vec { calls.lock().clone() } @@ -532,6 +585,190 @@ mod tests { Ok(()) } + #[tokio::test] + async fn profile_passes_only_first_string_user_instruction_to_scorer() -> Result<()> { + let scorer_inputs = Arc::new(Mutex::new(Vec::new())); + let (profile, _calls) = profile( + Arc::new(RecordingScorer { + inputs: scorer_inputs.clone(), + score: 1.0, + }), + false, + false, + )?; + let input = ProfileInput { + request: ChatRequest::openai_chat(json!({ + "model": "client/model", + "messages": [ + {"role": "system", "content": "System wrapper"}, + {"role": "user", "content": "First instruction"}, + {"role": "assistant", "content": "Intermediate response"}, + {"role": "user", "content": "Later instruction"}, + ], + "temperature": 0.2, + })), + metadata: RequestMetadata::default(), + }; + + let processed = profile.process(input).await?; + + assert_eq!(processed.decision.tier, TIER_WEAK); + assert_eq!(&*scorer_inputs.lock(), &["First instruction"]); + Ok(()) + } + + #[tokio::test] + async fn explicit_probe_input_controls_scoring_cache_and_not_backend_body() -> Result<()> { + let scorer_inputs = Arc::new(Mutex::new(Vec::new())); + let (profile, calls) = profile( + Arc::new(RecordingScorer { + inputs: scorer_inputs.clone(), + score: 1.0, + }), + false, + false, + )?; + + let first = profile + .run(input_with_probe_field( + json!("raw instruction"), + "wrapped instruction with terminal state one", + )) + .await?; + let second = profile + .run(input_with_probe_field( + json!("raw instruction"), + "wrapped instruction with terminal state two", + )) + .await?; + + assert_eq!( + first.routing_metadata.and_then(|data| data.selected_tier), + Some(TIER_WEAK.to_string()) + ); + assert_eq!( + second.routing_metadata.and_then(|data| data.selected_tier), + Some(TIER_WEAK.to_string()) + ); + assert_eq!(&*scorer_inputs.lock(), &["raw instruction"]); + assert_eq!(profile.decision_cache.lock().len(), 1); + let observed = observed(&calls); + assert_eq!(observed.len(), 2); + assert_eq!(observed[0].backend, "weak-backend"); + assert_eq!(observed[1].backend, "weak-backend"); + assert_eq!( + observed[0].body["messages"], + json!([{"role": "user", "content": "wrapped instruction with terminal state one"}]) + ); + assert_eq!( + observed[1].body["messages"], + json!([{"role": "user", "content": "wrapped instruction with terminal state two"}]) + ); + assert!(observed + .iter() + .all(|call| call.body.get(PREFILL_PROBE_INPUT_FIELD).is_none())); + Ok(()) + } + + #[tokio::test] + async fn different_explicit_probe_inputs_create_distinct_cache_entries() -> Result<()> { + let scorer_inputs = Arc::new(Mutex::new(Vec::new())); + let (profile, _calls) = profile( + Arc::new(RecordingScorer { + inputs: scorer_inputs.clone(), + score: 1.0, + }), + false, + false, + )?; + + let _first = profile + .process(input_with_probe_field(json!("raw one"), "same wrapper")) + .await?; + let _second = profile + .process(input_with_probe_field(json!("raw two"), "same wrapper")) + .await?; + + assert_eq!(&*scorer_inputs.lock(), &["raw one", "raw two"]); + assert_eq!(profile.decision_cache.lock().len(), 2); + Ok(()) + } + + #[tokio::test] + async fn explicit_probe_field_is_removed_on_strong_and_failure_paths() -> Result<()> { + let (strong_profile, strong_calls) = profile(Arc::new(FixedScorer(0.0)), false, false)?; + let strong_response = strong_profile + .run(input_with_probe_field( + json!("raw strong"), + "wrapped strong", + )) + .await?; + assert_eq!( + strong_response + .routing_metadata + .and_then(|data| data.selected_tier), + Some(TIER_STRONG.to_string()) + ); + let strong_observed = observed(&strong_calls); + assert_eq!(strong_observed.len(), 1); + assert!(strong_observed[0] + .body + .get(PREFILL_PROBE_INPUT_FIELD) + .is_none()); + + let (fallback_profile, fallback_calls) = profile(Arc::new(ErrorScorer), false, false)?; + let fallback_response = fallback_profile + .run(input_with_probe_field( + json!("raw fallback"), + "wrapped fallback", + )) + .await?; + assert_eq!( + fallback_response + .routing_metadata + .and_then(|data| data.selected_tier), + Some(TIER_STRONG.to_string()) + ); + assert_eq!(fallback_profile.decision_cache.lock().len(), 0); + let fallback_observed = observed(&fallback_calls); + assert_eq!(fallback_observed.len(), 1); + assert!(fallback_observed[0] + .body + .get(PREFILL_PROBE_INPUT_FIELD) + .is_none()); + Ok(()) + } + + #[tokio::test] + async fn malformed_explicit_probe_input_is_rejected_without_scoring() -> Result<()> { + let scorer_calls = Arc::new(AtomicUsize::new(0)); + let (profile, calls) = profile( + Arc::new(CountingScorer { + score: 1.0, + calls: scorer_calls.clone(), + }), + false, + false, + )?; + + for invalid in [json!(""), Value::Null, json!([]), json!({}), json!(7)] { + let error = profile + .run(input_with_probe_field(invalid, "wrapped instruction")) + .await + .err() + .ok_or_else(|| { + SwitchyardError::Other("malformed explicit input should fail".to_string()) + })?; + assert!(matches!(error, SwitchyardError::InvalidRequest(_))); + assert!(format!("{error}").contains(PREFILL_PROBE_INPUT_FIELD)); + } + + assert_eq!(scorer_calls.load(Ordering::SeqCst), 0); + assert_eq!(profile.decision_cache.lock().len(), 0); + assert!(observed(&calls).is_empty()); + Ok(()) + } + #[tokio::test] async fn probe_failure_forces_uncached_strong_then_retries_and_caches_success() -> Result<()> { let scorer_calls = Arc::new(AtomicUsize::new(0)); diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs index ed358667..ce28da2c 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/scorer.rs @@ -12,7 +12,7 @@ use std::time::Duration; use async_trait::async_trait; use safetensors::{Dtype, SafeTensors}; use serde::Deserialize; -use switchyard_core::{ChatRequest, Result, SwitchyardError}; +use switchyard_core::{Result, SwitchyardError}; use super::artifact::InferenceArtifact; use super::policy::CostAwareRoutingPolicy; @@ -24,7 +24,7 @@ const ARTIFACT_WAIT_INTERVAL: Duration = Duration::from_millis(50); #[async_trait] pub(crate) trait ProbeScorer: Send + Sync { /// Returns `1.0` for weak and `0.0` for strong. - async fn score(&self, request: &ChatRequest) -> Result; + async fn score(&self, probe_input: &str) -> Result; } struct CheckpointPrediction { @@ -131,19 +131,13 @@ struct KvTransferParams { #[async_trait] impl ProbeScorer for HiddenStateProbeScorer { - async fn score(&self, request: &ChatRequest) -> Result { - let messages = request - .body() - .as_object() - .and_then(|body| body.get("messages")) - .cloned() - .unwrap_or_else(|| serde_json::Value::Array(Vec::new())); + async fn score(&self, probe_input: &str) -> Result { let response = self .client .post(&self.completions_url) .json(&serde_json::json!({ "model": self.model, - "messages": messages, + "messages": [{"role": "user", "content": probe_input}], "max_tokens": 1, "kv_transfer_params": { "hidden_states_path": self.hidden_states_dir, @@ -899,17 +893,6 @@ mod tests { )?)) } - fn test_request() -> ChatRequest { - ChatRequest::openai_chat(serde_json::json!({ - "model": "client/model", - "messages": [ - {"role": "system", "content": "Be concise."}, - {"role": "user", "content": "Explain the failure."}, - ], - "temperature": 0.2, - })) - } - #[test] fn token_mean_pooling_produces_layer_major_features() -> Result<()> { let pooled = token_mean_per_layer( @@ -1173,7 +1156,7 @@ mod tests { } #[tokio::test] - async fn learned_probe_forwards_exact_request_and_returns_binary_score() -> Result<()> { + async fn learned_probe_forwards_exact_input_and_returns_binary_score() -> Result<()> { let directory = TestDirectory::create()?; let hidden_path = directory.write("probe-hidden.safetensors", &valid_artifact_bytes()?)?; let artifact = write_router_artifact(&directory, "router", [8.0, 2.0, -2.0, -8.0])?; @@ -1194,9 +1177,7 @@ mod tests { 2, CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, ); - let request = test_request(); - - let score = scorer.score(&request).await?; + let score = scorer.score("Explain the failure.").await?; assert_eq!(score, 1.0); assert!(!hidden_path.exists()); @@ -1207,7 +1188,6 @@ mod tests { serde_json::json!({ "model": "probe/model", "messages": [ - {"role": "system", "content": "Be concise."}, {"role": "user", "content": "Explain the failure."}, ], "max_tokens": 1, @@ -1224,8 +1204,6 @@ mod tests { async fn probe_response_errors_are_reported_before_artifact_processing() -> Result<()> { let directory = TestDirectory::create()?; let artifact = write_router_artifact(&directory, "router", [8.0, 2.0, -2.0, -8.0])?; - let request = test_request(); - let missing_server = MockProbeServer::spawn(200, serde_json::json!({}))?; let missing_scorer = HiddenStateProbeScorer::new( missing_server.base_url(), @@ -1237,7 +1215,7 @@ mod tests { CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, ); let error = missing_scorer - .score(&request) + .score("Explain the failure.") .await .err() .ok_or_else(|| probe_error("missing kv_transfer_params should fail"))?; @@ -1257,7 +1235,7 @@ mod tests { CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, ); let error = malformed_scorer - .score(&request) + .score("Explain the failure.") .await .err() .ok_or_else(|| probe_error("malformed kv_transfer_params should fail"))?; @@ -1274,7 +1252,7 @@ mod tests { CostAwareRoutingPolicy::new(1.0, 0.0, 0.0)?, ); let error = unavailable_scorer - .score(&request) + .score("Explain the failure.") .await .err() .ok_or_else(|| probe_error("non-success probe response should fail"))?; From ae02bc778acb73cd4b9d00785e39aa56edb49c48 Mon Sep 17 00:00:00 2001 From: michxu Date: Mon, 20 Jul 2026 18:13:06 +0000 Subject: [PATCH 08/12] test(benchmark): inject explicit prefill probe input Signed-off-by: michxu --- benchmark/prefill_probe_terminus_2.py | 39 +++++++++ tests/test_prefill_probe_terminus_2.py | 110 +++++++++++++++++++++++++ 2 files changed, 149 insertions(+) create mode 100644 benchmark/prefill_probe_terminus_2.py create mode 100644 tests/test_prefill_probe_terminus_2.py diff --git a/benchmark/prefill_probe_terminus_2.py b/benchmark/prefill_probe_terminus_2.py new file mode 100644 index 00000000..6d1330d4 --- /dev/null +++ b/benchmark/prefill_probe_terminus_2.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Harbor Terminus 2 adapter for explicit prefill-probe input.""" + +from __future__ import annotations + +import copy + +from harbor.agents.terminus_2.terminus_2 import Terminus2 +from harbor.environments.base import BaseEnvironment +from harbor.models.agent.context import AgentContext + +PREFILL_PROBE_INPUT_FIELD = "_switchyard_prefill_probe_input" + + +class PrefillProbeTerminus2(Terminus2): + """Attach the exact task instruction to every prefill-probe LLM call.""" + + async def run( + self, + instruction: str, + environment: BaseEnvironment, + context: AgentContext, + ) -> None: + """Run stock Terminus 2 with a persistent profile-local probe input.""" + call_kwargs = dict(self._llm_call_kwargs) + configured_extra_body = call_kwargs.get("extra_body") + if configured_extra_body is None: + extra_body: dict[str, object] = {} + elif isinstance(configured_extra_body, dict): + extra_body = copy.deepcopy(configured_extra_body) + else: + raise TypeError("llm_call_kwargs.extra_body must be a dictionary") + + extra_body[PREFILL_PROBE_INPUT_FIELD] = instruction + call_kwargs["extra_body"] = extra_body + self._llm_call_kwargs = call_kwargs + await super().run(instruction, environment, context) diff --git a/tests/test_prefill_probe_terminus_2.py b/tests/test_prefill_probe_terminus_2.py new file mode 100644 index 00000000..c5d3fd0f --- /dev/null +++ b/tests/test_prefill_probe_terminus_2.py @@ -0,0 +1,110 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import copy +import importlib.util +from pathlib import Path +from types import ModuleType + +import pytest + +REPO = Path(__file__).resolve().parents[1] +ADAPTER = REPO / "benchmark" / "prefill_probe_terminus_2.py" + + +def _load_adapter_module() -> ModuleType: + spec = importlib.util.spec_from_file_location("switchyard_prefill_probe_terminus_2", ADAPTER) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +async def test_adapter_preserves_prompt_and_injects_exact_input_for_the_full_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_adapter_module() + observed: list[tuple[str, object, object, dict[str, object]]] = [] + + async def fake_run(self, instruction: str, environment: object, context: object) -> None: + observed.append( + (instruction, environment, context, copy.deepcopy(self._llm_call_kwargs)) + ) + observed.append( + (instruction, environment, context, copy.deepcopy(self._llm_call_kwargs)) + ) + + monkeypatch.setattr(module.Terminus2, "run", fake_run) + agent = object.__new__(module.PrefillProbeTerminus2) + original_kwargs = { + "extra_body": {"existing": {"nested": True}}, + "timeout": 30, + } + agent._llm_call_kwargs = original_kwargs + environment = object() + context = object() + instruction = " exact raw instruction\nwith original whitespace " + + await agent.run(instruction, environment, context) + + assert len(observed) == 2 + for forwarded_instruction, forwarded_environment, forwarded_context, call_kwargs in observed: + assert forwarded_instruction == instruction + assert forwarded_environment is environment + assert forwarded_context is context + assert call_kwargs == { + "extra_body": { + "existing": {"nested": True}, + module.PREFILL_PROBE_INPUT_FIELD: instruction, + }, + "timeout": 30, + } + assert original_kwargs == { + "extra_body": {"existing": {"nested": True}}, + "timeout": 30, + } + + +async def test_adapter_initializes_extra_body_and_refreshes_input_per_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = _load_adapter_module() + observed: list[dict[str, object]] = [] + + async def fake_run(self, instruction: str, environment: object, context: object) -> None: + observed.append(copy.deepcopy(self._llm_call_kwargs)) + + monkeypatch.setattr(module.Terminus2, "run", fake_run) + agent = object.__new__(module.PrefillProbeTerminus2) + agent._llm_call_kwargs = {"timeout": 30} + + await agent.run("first instruction", object(), object()) + await agent.run("second instruction", object(), object()) + + assert observed == [ + { + "timeout": 30, + "extra_body": {module.PREFILL_PROBE_INPUT_FIELD: "first instruction"}, + }, + { + "timeout": 30, + "extra_body": {module.PREFILL_PROBE_INPUT_FIELD: "second instruction"}, + }, + ] + + +async def test_adapter_rejects_non_mapping_extra_body(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_adapter_module() + + async def fake_run(self, instruction: str, environment: object, context: object) -> None: + raise AssertionError("stock Terminus should not run with invalid extra_body") + + monkeypatch.setattr(module.Terminus2, "run", fake_run) + agent = object.__new__(module.PrefillProbeTerminus2) + agent._llm_call_kwargs = {"extra_body": "invalid"} + + with pytest.raises(TypeError, match="extra_body must be a dictionary"): + await agent.run("instruction", object(), object()) From 7f00b7fd6a6c6f55b0d86f82712b16ed731f277b Mon Sep 17 00:00:00 2001 From: michxu Date: Mon, 20 Jul 2026 18:13:25 +0000 Subject: [PATCH 09/12] docs(prefill-probe): document explicit probe input Signed-off-by: michxu --- docs/vllm-serve-hidden-state.md | 54 ++++++++++++++++++++++++++++++--- 1 file changed, 50 insertions(+), 4 deletions(-) diff --git a/docs/vllm-serve-hidden-state.md b/docs/vllm-serve-hidden-state.md index 77330dac..35852ee1 100644 --- a/docs/vllm-serve-hidden-state.md +++ b/docs/vllm-serve-hidden-state.md @@ -64,10 +64,56 @@ margin becomes public score `0.0` and selects strong. The learned profile's threshold is fixed at `0.5`; tune routing only with `lambda`. The sample uses `lambda: 0.5`, `weak_cost: 0`, and `strong_cost: 1`. -Successful decisions are cached by a hash of the first string-valued user -instruction. A cache hit reuses the selected tier without another probe. Probe, -artifact, or scoring failures route to strong and are not cached, so a later -matching request retries the probe. +Successful decisions are cached by a hash of the resolved probe input: the +explicit benchmark input when present, otherwise the first string-valued user +instruction. A cache hit reuses the selected tier without another probe. +Probe, artifact, or scoring failures route to strong and are not cached, so a +later matching request retries the probe. + +## Explicit probe input for Terminal-Bench + +Ordinary clients do not need special request fields. When no override is +present, the profile scores and caches the first string-valued user message. +Terminal-Bench's stock Terminus 2 agent, however, places the raw task +instruction inside a larger first-user message containing its command protocol +and current terminal state. That wrapped message does not reproduce a router +checkpoint trained on the raw task instruction alone. + +For prefill-probe benchmark runs, the repository provides +`benchmark/prefill_probe_terminus_2.py`. Its `PrefillProbeTerminus2` adapter adds +the exact `Terminus2.run()` instruction as this private top-level request field: + +```json +{ + "_switchyard_prefill_probe_input": "" +} +``` + +The `prefill-probe` profile uses a non-empty string value for both scoring and +the decision-cache key. It removes the field before calling the selected +completion target, so that target still receives the unmodified Terminus +conversation. A present empty or non-string value is rejected as an invalid +request instead of silently falling back to the wrapped message. + +Run Harbor from the repository root and select the adapter by import path: + +```bash +uv run --no-sync harbor run \ + --agent-import-path benchmark.prefill_probe_terminus_2:PrefillProbeTerminus2 \ + --model openai/router \ + --path /path/to/terminal-bench-2-closed-book \ + --jobs-dir /path/to/jobs \ + --job-name prefill-probe-smoke \ + --n-tasks 1 \ + --n-concurrent 1 \ + --max-retries 0 \ + --yes +``` + +Replace the usual `--agent terminus-2` option with `--agent-import-path`; do +not pass both. Harbor gives a registered agent name precedence over a custom +import path. The adapter keeps the same explicit input in `extra_body` for all +LLM calls during the task while leaving normal prompt construction unchanged. ## Docker launch From 18bfa111e7bf3900773b252d8f99be9821f868b3 Mon Sep 17 00:00:00 2001 From: michxu Date: Wed, 22 Jul 2026 18:49:20 +0000 Subject: [PATCH 10/12] feat(prefill-probe): extract Terminus task internally Signed-off-by: michxu --- benchmark/prefill_probe_terminus_2.py | 39 ---- .../src/profiles/prefill_router/profile.rs | 207 +++++++++++------- tests/test_prefill_probe_terminus_2.py | 110 ---------- 3 files changed, 124 insertions(+), 232 deletions(-) delete mode 100644 benchmark/prefill_probe_terminus_2.py delete mode 100644 tests/test_prefill_probe_terminus_2.py diff --git a/benchmark/prefill_probe_terminus_2.py b/benchmark/prefill_probe_terminus_2.py deleted file mode 100644 index 6d1330d4..00000000 --- a/benchmark/prefill_probe_terminus_2.py +++ /dev/null @@ -1,39 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Harbor Terminus 2 adapter for explicit prefill-probe input.""" - -from __future__ import annotations - -import copy - -from harbor.agents.terminus_2.terminus_2 import Terminus2 -from harbor.environments.base import BaseEnvironment -from harbor.models.agent.context import AgentContext - -PREFILL_PROBE_INPUT_FIELD = "_switchyard_prefill_probe_input" - - -class PrefillProbeTerminus2(Terminus2): - """Attach the exact task instruction to every prefill-probe LLM call.""" - - async def run( - self, - instruction: str, - environment: BaseEnvironment, - context: AgentContext, - ) -> None: - """Run stock Terminus 2 with a persistent profile-local probe input.""" - call_kwargs = dict(self._llm_call_kwargs) - configured_extra_body = call_kwargs.get("extra_body") - if configured_extra_body is None: - extra_body: dict[str, object] = {} - elif isinstance(configured_extra_body, dict): - extra_body = copy.deepcopy(configured_extra_body) - else: - raise TypeError("llm_call_kwargs.extra_body must be a dictionary") - - extra_body[PREFILL_PROBE_INPUT_FIELD] = instruction - call_kwargs["extra_body"] = extra_body - self._llm_call_kwargs = call_kwargs - await super().run(instruction, environment, context) diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs index dcebcf2b..07a42b53 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs @@ -27,7 +27,8 @@ use super::policy::CostAwareRoutingPolicy; use super::scorer::{HiddenStateProbeScorer, ProbeScorer}; const LEARNED_SCORE_THRESHOLD: f64 = 0.5; -const PREFILL_PROBE_INPUT_FIELD: &str = "_switchyard_prefill_probe_input"; +const TERMINUS_TASK_DESCRIPTION_HEADER: &str = "Task Description:\n"; +const TERMINUS_TERMINAL_STATE_HEADER: &str = "\n\nCurrent terminal state:\n"; const TIER_STRONG: &str = "strong"; const TIER_WEAK: &str = "weak"; @@ -178,7 +179,7 @@ pub struct PrefillProbeProfile { score_threshold: f64, scorer: Arc, stats: StatsAccumulator, - // Successful scores are reused by the resolved probe input. + // Successful scores are reused by the internally resolved probe input. decision_cache: Arc>>, } @@ -193,21 +194,21 @@ impl PrefillProbeProfile { }) } - // Removes the profile-local field and resolves the exact input to score and cache. - fn take_probe_input(request: &mut switchyard_core::ChatRequest) -> Result> { - let explicit = request - .body_mut() - .as_object_mut() - .and_then(|body| body.remove(PREFILL_PROBE_INPUT_FIELD)); - match explicit { - Some(serde_json::Value::String(probe_input)) if !probe_input.is_empty() => { - Ok(Some(probe_input)) - } - Some(_) => Err(SwitchyardError::InvalidRequest(format!( - "{PREFILL_PROBE_INPUT_FIELD} must be a non-empty string", - ))), - None => Ok(Self::first_user_instruction(request).map(str::to_owned)), - } + // Extracts the task-only portion of a stock Terminus 2 first-user prompt. + fn terminus_task_instruction(instruction: &str) -> Option<&str> { + let (_, task_and_terminal) = instruction.split_once(TERMINUS_TASK_DESCRIPTION_HEADER)?; + let (task, _) = task_and_terminal.split_once(TERMINUS_TERMINAL_STATE_HEADER)?; + (!task.is_empty()).then_some(task) + } + + // Resolves the text used by both hidden-state scoring and decision caching. + fn probe_input(request: &switchyard_core::ChatRequest) -> Option { + let first_user = Self::first_user_instruction(request)?; + Some( + Self::terminus_task_instruction(first_user) + .unwrap_or(first_user) + .to_owned(), + ) } // Hashes the resolved probe input for successful-decision reuse. @@ -219,7 +220,7 @@ impl PrefillProbeProfile { // Scores an uncached instruction and explicitly bypasses threshold routing on failure. async fn route(&self, mut input: ProfileInput) -> Result { - let probe_input = Self::take_probe_input(&mut input.request)?; + let probe_input = Self::probe_input(&input.request); let key = probe_input.as_deref().map(Self::instruction_key); let cached = key.and_then(|key| self.decision_cache.lock().get(&key).copied()); let (score, strong_fallback) = if let Some(score) = cached { @@ -530,12 +531,14 @@ mod tests { } } - fn input_with_probe_field(probe_input: Value, wrapped_instruction: &str) -> ProfileInput { + fn terminus_input(preamble: &str, instruction: &str, terminal_state: &str) -> ProfileInput { + let wrapped = format!( + "{preamble}\n\n{TERMINUS_TASK_DESCRIPTION_HEADER}{instruction}{TERMINUS_TERMINAL_STATE_HEADER}{terminal_state}" + ); ProfileInput { request: ChatRequest::openai_chat(json!({ "model": "client/model", - "messages": [{"role": "user", "content": wrapped_instruction}], - (PREFILL_PROBE_INPUT_FIELD): probe_input, + "messages": [{"role": "user", "content": wrapped}], })), metadata: RequestMetadata::default(), } @@ -618,9 +621,9 @@ mod tests { } #[tokio::test] - async fn explicit_probe_input_controls_scoring_cache_and_not_backend_body() -> Result<()> { + async fn terminus_json_and_xml_wrappers_pass_exact_task_text_to_scorer() -> Result<()> { let scorer_inputs = Arc::new(Mutex::new(Vec::new())); - let (profile, calls) = profile( + let (profile, _calls) = profile( Arc::new(RecordingScorer { inputs: scorer_inputs.clone(), score: 1.0, @@ -629,49 +632,73 @@ mod tests { false, )?; - let first = profile - .run(input_with_probe_field( - json!("raw instruction"), - "wrapped instruction with terminal state one", + let _json = profile + .process(terminus_input( + "Format your response as JSON.", + " preserve JSON task whitespace ", + "$ ", )) .await?; - let second = profile - .run(input_with_probe_field( - json!("raw instruction"), - "wrapped instruction with terminal state two", + let _xml = profile + .process(terminus_input( + "Format your response as XML.", + "preserve XML task", + "project files", )) .await?; assert_eq!( - first.routing_metadata.and_then(|data| data.selected_tier), + &*scorer_inputs.lock(), + &[" preserve JSON task whitespace ", "preserve XML task"] + ); + Ok(()) + } + + #[tokio::test] + async fn terminus_task_controls_cache_and_wrapped_backend_body_is_preserved() -> Result<()> { + let scorer_inputs = Arc::new(Mutex::new(Vec::new())); + let (profile, calls) = profile( + Arc::new(RecordingScorer { + inputs: scorer_inputs.clone(), + score: 1.0, + }), + false, + false, + )?; + + let first = terminus_input("Terminus protocol", "same raw task", "state one"); + let second = terminus_input("Terminus protocol", "same raw task", "state two"); + let first_messages = first.request.body()["messages"].clone(); + let second_messages = second.request.body()["messages"].clone(); + + let first_response = profile.run(first).await?; + let second_response = profile.run(second).await?; + + assert_eq!( + first_response + .routing_metadata + .and_then(|data| data.selected_tier), Some(TIER_WEAK.to_string()) ); assert_eq!( - second.routing_metadata.and_then(|data| data.selected_tier), + second_response + .routing_metadata + .and_then(|data| data.selected_tier), Some(TIER_WEAK.to_string()) ); - assert_eq!(&*scorer_inputs.lock(), &["raw instruction"]); + assert_eq!(&*scorer_inputs.lock(), &["same raw task"]); assert_eq!(profile.decision_cache.lock().len(), 1); let observed = observed(&calls); assert_eq!(observed.len(), 2); assert_eq!(observed[0].backend, "weak-backend"); assert_eq!(observed[1].backend, "weak-backend"); - assert_eq!( - observed[0].body["messages"], - json!([{"role": "user", "content": "wrapped instruction with terminal state one"}]) - ); - assert_eq!( - observed[1].body["messages"], - json!([{"role": "user", "content": "wrapped instruction with terminal state two"}]) - ); - assert!(observed - .iter() - .all(|call| call.body.get(PREFILL_PROBE_INPUT_FIELD).is_none())); + assert_eq!(observed[0].body["messages"], first_messages); + assert_eq!(observed[1].body["messages"], second_messages); Ok(()) } #[tokio::test] - async fn different_explicit_probe_inputs_create_distinct_cache_entries() -> Result<()> { + async fn different_terminus_tasks_create_distinct_cache_entries() -> Result<()> { let scorer_inputs = Arc::new(Mutex::new(Vec::new())); let (profile, _calls) = profile( Arc::new(RecordingScorer { @@ -683,10 +710,10 @@ mod tests { )?; let _first = profile - .process(input_with_probe_field(json!("raw one"), "same wrapper")) + .process(terminus_input("Terminus protocol", "raw one", "same state")) .await?; let _second = profile - .process(input_with_probe_field(json!("raw two"), "same wrapper")) + .process(terminus_input("Terminus protocol", "raw two", "same state")) .await?; assert_eq!(&*scorer_inputs.lock(), &["raw one", "raw two"]); @@ -695,14 +722,34 @@ mod tests { } #[tokio::test] - async fn explicit_probe_field_is_removed_on_strong_and_failure_paths() -> Result<()> { + async fn malformed_terminus_envelopes_fall_back_to_full_first_user_text() -> Result<()> { + let scorer_inputs = Arc::new(Mutex::new(Vec::new())); + let (profile, _calls) = profile( + Arc::new(RecordingScorer { + inputs: scorer_inputs.clone(), + score: 1.0, + }), + false, + false, + )?; + let missing_terminal = "protocol\n\nTask Description:\nraw task without terminal marker"; + let reversed = "Current terminal state:\nstate\n\nTask Description:\nraw task"; + + let _first = profile.process(input(missing_terminal)).await?; + let _second = profile.process(input(reversed)).await?; + + assert_eq!(&*scorer_inputs.lock(), &[missing_terminal, reversed]); + Ok(()) + } + + #[tokio::test] + async fn terminus_wrapper_is_preserved_on_strong_and_probe_failure_paths() -> Result<()> { + let strong_input = terminus_input("Terminus protocol", "raw strong", "strong state"); + let strong_messages = strong_input.request.body()["messages"].clone(); let (strong_profile, strong_calls) = profile(Arc::new(FixedScorer(0.0)), false, false)?; - let strong_response = strong_profile - .run(input_with_probe_field( - json!("raw strong"), - "wrapped strong", - )) - .await?; + + let strong_response = strong_profile.run(strong_input).await?; + assert_eq!( strong_response .routing_metadata @@ -711,18 +758,14 @@ mod tests { ); let strong_observed = observed(&strong_calls); assert_eq!(strong_observed.len(), 1); - assert!(strong_observed[0] - .body - .get(PREFILL_PROBE_INPUT_FIELD) - .is_none()); + assert_eq!(strong_observed[0].body["messages"], strong_messages); + let fallback_input = terminus_input("Terminus protocol", "raw fallback", "fallback state"); + let fallback_messages = fallback_input.request.body()["messages"].clone(); let (fallback_profile, fallback_calls) = profile(Arc::new(ErrorScorer), false, false)?; - let fallback_response = fallback_profile - .run(input_with_probe_field( - json!("raw fallback"), - "wrapped fallback", - )) - .await?; + + let fallback_response = fallback_profile.run(fallback_input).await?; + assert_eq!( fallback_response .routing_metadata @@ -732,17 +775,14 @@ mod tests { assert_eq!(fallback_profile.decision_cache.lock().len(), 0); let fallback_observed = observed(&fallback_calls); assert_eq!(fallback_observed.len(), 1); - assert!(fallback_observed[0] - .body - .get(PREFILL_PROBE_INPUT_FIELD) - .is_none()); + assert_eq!(fallback_observed[0].body["messages"], fallback_messages); Ok(()) } #[tokio::test] - async fn malformed_explicit_probe_input_is_rejected_without_scoring() -> Result<()> { + async fn missing_string_user_content_is_an_uncached_strong_fallback() -> Result<()> { let scorer_calls = Arc::new(AtomicUsize::new(0)); - let (profile, calls) = profile( + let (profile, _calls) = profile( Arc::new(CountingScorer { score: 1.0, calls: scorer_calls.clone(), @@ -750,22 +790,23 @@ mod tests { false, false, )?; + let no_string_user = ProfileInput { + request: ChatRequest::openai_chat(json!({ + "model": "client/model", + "messages": [ + {"role": "system", "content": "system"}, + {"role": "user", "content": [{"type": "text", "text": "blocks"}]}, + ], + })), + metadata: RequestMetadata::default(), + }; - for invalid in [json!(""), Value::Null, json!([]), json!({}), json!(7)] { - let error = profile - .run(input_with_probe_field(invalid, "wrapped instruction")) - .await - .err() - .ok_or_else(|| { - SwitchyardError::Other("malformed explicit input should fail".to_string()) - })?; - assert!(matches!(error, SwitchyardError::InvalidRequest(_))); - assert!(format!("{error}").contains(PREFILL_PROBE_INPUT_FIELD)); - } + let processed = profile.process(no_string_user).await?; + assert_eq!(processed.decision.tier, TIER_STRONG); + assert_eq!(processed.decision.score, 0.0); assert_eq!(scorer_calls.load(Ordering::SeqCst), 0); assert_eq!(profile.decision_cache.lock().len(), 0); - assert!(observed(&calls).is_empty()); Ok(()) } diff --git a/tests/test_prefill_probe_terminus_2.py b/tests/test_prefill_probe_terminus_2.py deleted file mode 100644 index c5d3fd0f..00000000 --- a/tests/test_prefill_probe_terminus_2.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import copy -import importlib.util -from pathlib import Path -from types import ModuleType - -import pytest - -REPO = Path(__file__).resolve().parents[1] -ADAPTER = REPO / "benchmark" / "prefill_probe_terminus_2.py" - - -def _load_adapter_module() -> ModuleType: - spec = importlib.util.spec_from_file_location("switchyard_prefill_probe_terminus_2", ADAPTER) - assert spec is not None - assert spec.loader is not None - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -async def test_adapter_preserves_prompt_and_injects_exact_input_for_the_full_run( - monkeypatch: pytest.MonkeyPatch, -) -> None: - module = _load_adapter_module() - observed: list[tuple[str, object, object, dict[str, object]]] = [] - - async def fake_run(self, instruction: str, environment: object, context: object) -> None: - observed.append( - (instruction, environment, context, copy.deepcopy(self._llm_call_kwargs)) - ) - observed.append( - (instruction, environment, context, copy.deepcopy(self._llm_call_kwargs)) - ) - - monkeypatch.setattr(module.Terminus2, "run", fake_run) - agent = object.__new__(module.PrefillProbeTerminus2) - original_kwargs = { - "extra_body": {"existing": {"nested": True}}, - "timeout": 30, - } - agent._llm_call_kwargs = original_kwargs - environment = object() - context = object() - instruction = " exact raw instruction\nwith original whitespace " - - await agent.run(instruction, environment, context) - - assert len(observed) == 2 - for forwarded_instruction, forwarded_environment, forwarded_context, call_kwargs in observed: - assert forwarded_instruction == instruction - assert forwarded_environment is environment - assert forwarded_context is context - assert call_kwargs == { - "extra_body": { - "existing": {"nested": True}, - module.PREFILL_PROBE_INPUT_FIELD: instruction, - }, - "timeout": 30, - } - assert original_kwargs == { - "extra_body": {"existing": {"nested": True}}, - "timeout": 30, - } - - -async def test_adapter_initializes_extra_body_and_refreshes_input_per_run( - monkeypatch: pytest.MonkeyPatch, -) -> None: - module = _load_adapter_module() - observed: list[dict[str, object]] = [] - - async def fake_run(self, instruction: str, environment: object, context: object) -> None: - observed.append(copy.deepcopy(self._llm_call_kwargs)) - - monkeypatch.setattr(module.Terminus2, "run", fake_run) - agent = object.__new__(module.PrefillProbeTerminus2) - agent._llm_call_kwargs = {"timeout": 30} - - await agent.run("first instruction", object(), object()) - await agent.run("second instruction", object(), object()) - - assert observed == [ - { - "timeout": 30, - "extra_body": {module.PREFILL_PROBE_INPUT_FIELD: "first instruction"}, - }, - { - "timeout": 30, - "extra_body": {module.PREFILL_PROBE_INPUT_FIELD: "second instruction"}, - }, - ] - - -async def test_adapter_rejects_non_mapping_extra_body(monkeypatch: pytest.MonkeyPatch) -> None: - module = _load_adapter_module() - - async def fake_run(self, instruction: str, environment: object, context: object) -> None: - raise AssertionError("stock Terminus should not run with invalid extra_body") - - monkeypatch.setattr(module.Terminus2, "run", fake_run) - agent = object.__new__(module.PrefillProbeTerminus2) - agent._llm_call_kwargs = {"extra_body": "invalid"} - - with pytest.raises(TypeError, match="extra_body must be a dictionary"): - await agent.run("instruction", object(), object()) From 83d55a7fd07fe7752060f3af09332cfa422da5f9 Mon Sep 17 00:00:00 2001 From: michxu Date: Wed, 22 Jul 2026 18:50:13 +0000 Subject: [PATCH 11/12] docs(prefill-probe): document stock Terminus routing Signed-off-by: michxu --- docs/vllm-serve-hidden-state.md | 90 ++++++++++++++++++++------------- 1 file changed, 55 insertions(+), 35 deletions(-) diff --git a/docs/vllm-serve-hidden-state.md b/docs/vllm-serve-hidden-state.md index 35852ee1..76de60b9 100644 --- a/docs/vllm-serve-hidden-state.md +++ b/docs/vllm-serve-hidden-state.md @@ -64,56 +64,76 @@ margin becomes public score `0.0` and selects strong. The learned profile's threshold is fixed at `0.5`; tune routing only with `lambda`. The sample uses `lambda: 0.5`, `weak_cost: 0`, and `strong_cost: 1`. -Successful decisions are cached by a hash of the resolved probe input: the -explicit benchmark input when present, otherwise the first string-valued user -instruction. A cache hit reuses the selected tier without another probe. -Probe, artifact, or scoring failures route to strong and are not cached, so a -later matching request retries the probe. +Successful decisions are cached by a hash of the internally resolved probe +input. A cache hit reuses the selected tier without another probe. Probe, +artifact, or scoring failures route to strong and are not cached, so a later +matching request retries the probe. -## Explicit probe input for Terminal-Bench +## Internal task extraction for Terminal-Bench -Ordinary clients do not need special request fields. When no override is -present, the profile scores and caches the first string-valued user message. -Terminal-Bench's stock Terminus 2 agent, however, places the raw task +For ordinary requests, the profile scores and caches the first string-valued +user message. Terminal-Bench's stock Terminus 2 agent places the raw task instruction inside a larger first-user message containing its command protocol -and current terminal state. That wrapped message does not reproduce a router -checkpoint trained on the raw task instruction alone. +and current terminal state. The profile recognizes the stock envelope: -For prefill-probe benchmark runs, the repository provides -`benchmark/prefill_probe_terminus_2.py`. Its `PrefillProbeTerminus2` adapter adds -the exact `Terminus2.run()` instruction as this private top-level request field: +```text +Task Description: + -```json -{ - "_switchyard_prefill_probe_input": "" -} +Current terminal state: + ``` -The `prefill-probe` profile uses a non-empty string value for both scoring and -the decision-cache key. It removes the field before calling the selected -completion target, so that target still receives the unmodified Terminus -conversation. A present empty or non-string value is rejected as an invalid -request instead of silently falling back to the wrapped message. +When both headers occur in order and the task text is non-empty, Switchyard +uses the exact text between them for hidden-state scoring and the decision-cache +key. It does not trim or normalize that text. If the envelope is absent or +malformed, the profile falls back to the complete first string-valued user +message. If no such message exists, probing fails open to strong without +caching the failure. + +This extraction is internal to the `prefill-probe` profile. Clients send a +normal chat request, and benchmark runs use Harbor's registered `terminus-2` +agent without a custom subclass or private request field. The selected +completion target still receives the original wrapped Terminus conversation; +only its model is rewritten. + +After the probe is running, start Switchyard on an address reachable from the +benchmark task containers: + +```bash +export NVIDIA_API_KEY="..." +export VLLM_BASE_URL=http://127.0.0.1:8000 +export PREFILL_ROUTER_ARTIFACT_DIR=/absolute/path/to/exported/router + +switchyard serve \ + --config benchmark/routing-profiles/prefill-probe-local.yaml \ + --host 0.0.0.0 \ + --port 4000 +``` -Run Harbor from the repository root and select the adapter by import path: +In another terminal, use the repository's benchmark wrapper with the stock +agent. `SWITCHYARD_BENCHMARK_URL` must resolve from Docker task containers; do +not use `127.0.0.1` unless Switchyard runs in that same container: ```bash -uv run --no-sync harbor run \ - --agent-import-path benchmark.prefill_probe_terminus_2:PrefillProbeTerminus2 \ +export UPSTREAM_API_KEY=switchyard-local +export SWITCHYARD_BENCHMARK_URL=http://YOUR_DOCKER_REACHABLE_HOST:4000/v1 + +bash benchmark/run-baseline.sh \ + --harbor-path /path/to/terminal-bench-2-closed-book \ + --upstream-base-url "${SWITCHYARD_BENCHMARK_URL}" \ + --upstream-api-key-env UPSTREAM_API_KEY \ + --agent terminus-2 \ --model openai/router \ - --path /path/to/terminal-bench-2-closed-book \ - --jobs-dir /path/to/jobs \ - --job-name prefill-probe-smoke \ --n-tasks 1 \ --n-concurrent 1 \ - --max-retries 0 \ - --yes + --max-retries 0 ``` -Replace the usual `--agent terminus-2` option with `--agent-import-path`; do -not pass both. Harbor gives a registered agent name precedence over a custom -import path. The adapter keeps the same explicit input in `extra_body` for all -LLM calls during the task while leaving normal prompt construction unchanged. +In this setup, the runner treats the separately launched Switchyard server as +an OpenAI-compatible upstream. It manages the closed-book Harbor run and its +artifacts but does not manage the Switchyard lifecycle or collect Switchyard +routing statistics. ## Docker launch From 242790e86dcfa8c7f627f8bbd8c887af9161dc34 Mon Sep 17 00:00:00 2001 From: michxu Date: Wed, 22 Jul 2026 21:21:03 +0000 Subject: [PATCH 12/12] refactor(prefill-probe): rename config field to checkpoint_dir Signed-off-by: michxu --- .../routing-profiles/prefill-probe-local.yaml | 4 ++-- .../src/profiles/prefill_router/profile.rs | 1 + .../tests/prefill_probe_profile.rs | 24 +++++++++++++++++-- docs/vllm-serve-hidden-state.md | 8 +++---- 4 files changed, 29 insertions(+), 8 deletions(-) diff --git a/benchmark/routing-profiles/prefill-probe-local.yaml b/benchmark/routing-profiles/prefill-probe-local.yaml index 86a81648..03a335bf 100644 --- a/benchmark/routing-profiles/prefill-probe-local.yaml +++ b/benchmark/routing-profiles/prefill-probe-local.yaml @@ -5,7 +5,7 @@ # documented in docs/vllm-serve-hidden-state.md and export: # NVIDIA_API_KEY # VLLM_BASE_URL=http://127.0.0.1:8000 -# PREFILL_ROUTER_ARTIFACT_DIR=/absolute/path/to/exported/router +# PREFILL_ROUTER_CHECKPOINT_DIR=/absolute/path/to/exported/router # # Serve with: # switchyard serve --config benchmark/routing-profiles/prefill-probe-local.yaml @@ -42,7 +42,7 @@ profiles: weak: weak weak_checkpoint_head: nemotron-3-super hidden_states_dir: /tmp/switchyard-hidden-states - inference_artifact_dir: ${PREFILL_ROUTER_ARTIFACT_DIR} + checkpoint_dir: ${PREFILL_ROUTER_CHECKPOINT_DIR} routing_policy: type: cost-aware lambda: 0.5 diff --git a/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs index 07a42b53..fb126152 100644 --- a/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs +++ b/crates/switchyard-components-v2/src/profiles/prefill_router/profile.rs @@ -70,6 +70,7 @@ pub struct PrefillProbeProfileConfig { /// Directory shared with vLLM's `ExampleHiddenStatesConnector`. pub hidden_states_dir: String, /// Directory containing `router.json` and `router.safetensors`. + #[serde(rename = "checkpoint_dir", alias = "inference_artifact_dir")] pub inference_artifact_dir: String, /// Policy that maps the two selected correctness probabilities to a binary score. pub routing_policy: PrefillProbeRoutingPolicyConfig, diff --git a/crates/switchyard-components-v2/tests/prefill_probe_profile.rs b/crates/switchyard-components-v2/tests/prefill_probe_profile.rs index 38f6b991..f7b08358 100644 --- a/crates/switchyard-components-v2/tests/prefill_probe_profile.rs +++ b/crates/switchyard-components-v2/tests/prefill_probe_profile.rs @@ -223,7 +223,7 @@ profiles: weak: weak weak_checkpoint_head: nemotron-3-super hidden_states_dir: /dev/shm/hidden_states - inference_artifact_dir: {} + checkpoint_dir: {} routing_policy: type: cost-aware lambda: 0.5 @@ -273,7 +273,7 @@ fn config_schema_requires_every_field_and_rejects_unknown_fields() -> Result<()> "weak", "weak_checkpoint_head", "hidden_states_dir", - "inference_artifact_dir", + "checkpoint_dir", "routing_policy", ] { let mut missing = base.clone(); @@ -303,6 +303,26 @@ fn config_schema_requires_every_field_and_rejects_unknown_fields() -> Result<()> Ok(()) } +#[test] +fn legacy_inference_artifact_dir_alias_is_accepted() -> Result<()> { + let artifact = TestArtifactDirectory::create(ArtifactFault::None)?; + let config = config_with_artifact(artifact.path())?; + let mut value = serde_json::to_value(&config) + .map_err(|error| SwitchyardError::Other(format!("config encode failed: {error}")))?; + let object = value + .as_object_mut() + .ok_or_else(|| SwitchyardError::Other("config should encode as an object".to_string()))?; + let checkpoint_dir = object.remove("checkpoint_dir").ok_or_else(|| { + SwitchyardError::Other("serialized config should contain checkpoint_dir".to_string()) + })?; + object.insert("inference_artifact_dir".to_string(), checkpoint_dir); + + let parsed = serde_json::from_value::(value) + .map_err(|error| SwitchyardError::Other(format!("legacy config parse failed: {error}")))?; + assert_eq!(parsed, config); + Ok(()) +} + #[test] fn policy_schema_is_tagged_and_strict() -> Result<()> { let valid = serde_json::to_value(routing_policy()) diff --git a/docs/vllm-serve-hidden-state.md b/docs/vllm-serve-hidden-state.md index 76de60b9..64e4c39d 100644 --- a/docs/vllm-serve-hidden-state.md +++ b/docs/vllm-serve-hidden-state.md @@ -19,10 +19,10 @@ required by the protocol. Use Docker for a reproducible CUDA/vLLM runtime, or use `vllm serve` directly when the environment contains a build with `extract_hidden_states` and `ExampleHiddenStatesConnector`. -## Router artifact contract +## Router checkpoint contract -The router artifact is external deployment data. Switchyard does not package -or download it. Set `PREFILL_ROUTER_ARTIFACT_DIR` to a readable directory with: +The router checkpoint is external deployment data. Switchyard does not package +or download it. Set `PREFILL_ROUTER_CHECKPOINT_DIR` to a readable directory with: ```text router.json @@ -103,7 +103,7 @@ benchmark task containers: ```bash export NVIDIA_API_KEY="..." export VLLM_BASE_URL=http://127.0.0.1:8000 -export PREFILL_ROUTER_ARTIFACT_DIR=/absolute/path/to/exported/router +export PREFILL_ROUTER_CHECKPOINT_DIR=/absolute/path/to/exported/router switchyard serve \ --config benchmark/routing-profiles/prefill-probe-local.yaml \