feat(prefill-router): run batched prefill inference

Signed-off-by: nachiketb <nachiketb@nvidia.com>

Author: nachiketb-nvidia

Cargo.lock

@@ -4,7 +4,7 @@
 [package]
 name = "prefill-router"
 version.workspace = true
-description = "Prefill feature extraction for Switchyard routing"
+description = "Learned prefill routing for Switchyard"
 authors.workspace = true
 edition.workspace = true
 license.workspace = true
@@ -13,7 +13,7 @@ rust-version.workspace = true
 
 [dependencies]
 pyo3 = { version = "0.28.3", features = ["auto-initialize"] }
-serde_json.workspace = true
+switchyard-protocol.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
 

crates/prefill-router/Cargo.toml

@@ -6,16 +6,12 @@ name = "prefill-router-tests"
 version = "0.0.0"
 requires-python = ">=3.10"
 dependencies = [
+    "accelerate>=1.0",
     "numpy>=1.26",
     "torch>=2.0",
     "transformers>=5.0",
 ]
 
-[dependency-groups]
-dev = [
-    "model-router-toolkit @ https://github.com/NVIDIA-AI-Blueprints/llm-router/archive/8a9d3509fbde879d9795258081bab9553b458e04.tar.gz",
-]
-
 [tool.uv]
 package = false
 

crates/prefill-router/pyproject.toml

@@ -1,10 +1,11 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-"""Hugging Face Transformers prefill forward used by the Rust crate."""
+"""Complete Transformers prefill and checkpoint inference for the Rust crate."""
 
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Any
 
 
@@ -37,125 +38,148 @@ def _resolve_device(torch: Any, override: str | None) -> str:
 
 
 class TransformersForward:
-    """Lazily load a causal LM and return pooled prefill hidden states."""
+    """Run encoder extraction and learned confidence inference in one pass."""
 
     def __init__(
         self,
-        model: str,
+        checkpoint_path: str | Path,
         *,
         device: str | None = None,
         cache_dir: str | None = None,
     ) -> None:
-        self._model_path = model
+        import numpy as np
+        import torch
+
+        checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+        if checkpoint["format_version"] != 1:
+            raise ValueError("unsupported checkpoint format_version")
+
+        encoder = checkpoint["encoder"]
+        architecture = checkpoint["architecture"]
+        pipeline = checkpoint["feature_pipeline"]
+        if encoder["view"] != "task_prompt_only":
+            raise ValueError("checkpoint encoder view must be task_prompt_only")
+        if pipeline["pooling"] != "mean of independently standardized selected layers":
+            raise ValueError("unsupported checkpoint feature pooling")
+        if architecture["activation"] != "ReLU":
+            raise ValueError("checkpoint activation must be ReLU")
+        if architecture["ensemble_reduction"] != "mean(sigmoid(logits))":
+            raise ValueError("unsupported checkpoint ensemble reduction")
+
+        self._numpy = np
+        self._torch = torch
+        self._model_path = str(encoder["name"])
+        self._expected_layers = int(encoder["n_layers"])
+        self._hidden_dim = int(encoder["hidden_dim"])
+        self._models = [str(model) for model in checkpoint["models"]]
+        self._selected_layers = [int(layer) for layer in pipeline["selected_layers"]]
+        self._layer_mean = torch.stack(
+            [pipeline["layer_mean"][str(layer)].float() for layer in self._selected_layers]
+        ).numpy()
+        self._layer_std = torch.stack(
+            [pipeline["layer_std"][str(layer)].float() for layer in self._selected_layers]
+        ).numpy()
+        self._scaler_mean = pipeline["scaler_mean"].numpy()
+        self._scaler_scale = pipeline["scaler_scale"].numpy()
+        self._pca_mean = pipeline["pca_mean"].numpy()
+        self._pca_components = pipeline["pca_components"].numpy()
+        self._states = checkpoint["model_state_dicts"]
         self._cache_dir = cache_dir
         self._device_override = device
         self._model = None
         self._tokenizer = None
-        self._torch = None
-        self.n_layers = 0
-        self.hidden_dim = 0
 
-    def _ensure_loaded(self) -> str:
+        if not self._selected_layers or len(set(self._selected_layers)) != len(
+            self._selected_layers
+        ):
+            raise ValueError("checkpoint selected layers must be non-empty and unique")
+        if not self._models or not self._states:
+            raise ValueError("checkpoint must contain models and ensemble members")
+        if self._layer_mean.shape != self._layer_std.shape or self._layer_mean.shape != (
+            len(self._selected_layers),
+            self._hidden_dim,
+        ):
+            raise ValueError("checkpoint layer normalization shape is inconsistent")
+        if not bool(np.all(self._layer_std > 0)) or not bool(np.all(self._scaler_scale > 0)):
+            raise ValueError("checkpoint normalization scales must be positive")
+
+    def metadata(self) -> tuple[str, int]:
+        """Return the encoder and ordered output count consumed by Rust."""
+        return self._model_path, len(self._models)
+
+    def _ensure_loaded(self) -> None:
         if self._model is not None:
-            return str(self._model.device)
+            return
 
-        import torch
         from transformers import AutoModelForCausalLM, AutoTokenizer
 
-        self._torch = torch
+        torch = self._torch
         device = _resolve_device(torch, self._device_override)
-        if device == "cpu":
-            dtype = torch.float32
-        elif device == "mps":
-            dtype = torch.float16
-        else:
-            dtype = (
-                torch.bfloat16
-                if torch.cuda.get_device_capability(device)[0] >= 8
-                else torch.float16
-            )
-
+        dtype = (
+            torch.float32
+            if device == "cpu"
+            else torch.float16
+            if device == "mps"
+            else torch.bfloat16
+            if torch.cuda.get_device_capability(device)[0] >= 8
+            else torch.float16
+        )
+        model_kwargs: dict[str, Any] = {}
+        if device == "cuda":
+            model_kwargs["device_map"] = "auto"
+        elif device.startswith("cuda:"):
+            model_kwargs["device_map"] = {"": device}
         self._tokenizer = AutoTokenizer.from_pretrained(
             self._model_path,
             cache_dir=self._cache_dir,
         )
         if self._tokenizer.pad_token is None:
             self._tokenizer.pad_token = self._tokenizer.eos_token
-
-        load_kwargs: dict[str, Any] = {
-            "dtype": dtype,
-            "cache_dir": self._cache_dir,
-        }
-
-        self._model = AutoModelForCausalLM.from_pretrained(
+        causal_model = AutoModelForCausalLM.from_pretrained(
             self._model_path,
-            **load_kwargs,
+            dtype=dtype,
+            cache_dir=self._cache_dir,
+            **model_kwargs,
         )
-        if device != "cpu":
+        self._model = causal_model.base_model
+        del causal_model
+        if device == "mps":
             self._model.to(device)
         self._model.eval()
-        self.n_layers = self._model.config.num_hidden_layers
-        self.hidden_dim = self._model.config.hidden_size
-        return str(self._model.device)
+        model_config = getattr(self._model.config, "text_config", self._model.config)
+        if (
+            model_config.num_hidden_layers != self._expected_layers
+            or model_config.hidden_size != self._hidden_dim
+        ):
+            raise ValueError("loaded encoder dimensions do not match checkpoint metadata")
 
-    def extract_batch(
-        self,
-        prompts: list[str],
-        *,
-        chat_template_kwargs: dict[str, Any] | None = None,
-        extract_layers: list[int] | str = "upper_half",
-        pooling_modes: list[str] | None = None,
-        batch_size: int = 4,
-        max_length: int = 2048,
-    ) -> dict[str, Any]:
-        """Extract pooled hidden states using the blueprint's direct indexing."""
+    def forward(self, prompts: list[str], batch_size: int, max_length: int) -> bytes:
+        """Return a row-major F32 probability matrix for ordered prompts."""
         self._ensure_loaded()
+        if not prompts or any(not prompt for prompt in prompts):
+            raise ValueError("prompts must be non-empty")
+        if batch_size <= 0 or max_length <= 0:
+            raise ValueError("batch_size and max_length must be positive")
 
-        if extract_layers == "all":
-            layers = list(range(self.n_layers))
-        elif extract_layers == "upper_half":
-            layers = list(range(self.n_layers // 2, self.n_layers))
-        elif isinstance(extract_layers, list):
-            layers = [int(layer) for layer in extract_layers]
-        else:
-            raise ValueError(f"Unsupported layer selection: {extract_layers}")
-        if not layers:
-            raise ValueError("extract_layers resolved to an empty list")
-        invalid = [layer for layer in layers if layer < 0 or layer >= self.n_layers]
-        if invalid:
-            raise ValueError(
-                f"Requested indexes {invalid} are outside hidden-state range 0..{self.n_layers - 1}"
+        formatted = [
+            self._tokenizer.apply_chat_template(
+                [{"role": "user", "content": prompt}],
+                tokenize=False,
+                add_generation_prompt=True,
             )
-
-        pools = {"last", "mean"} if pooling_modes is None else set(pooling_modes)
-        unknown_pools = pools - {"last", "mean"}
-        if unknown_pools:
-            raise ValueError(f"Unknown pooling modes: {sorted(unknown_pools)}")
-        if not pools:
-            raise ValueError("At least one pooling mode is required")
-
-        template_kwargs = chat_template_kwargs or {}
-        conversations = [[{"role": "user", "content": prompt}] for prompt in prompts]
-        formatted = self._tokenizer.apply_chat_template(
-            conversations,
-            tokenize=False,
-            add_generation_prompt=True,
-            **template_kwargs,
-        )
-        all_last = {layer: [] for layer in layers} if "last" in pools else {}
-        all_mean = {layer: [] for layer in layers} if "mean" in pools else {}
-
-        for batch_start in range(0, len(formatted), batch_size):
+            for prompt in prompts
+        ]
+        predictions = []
+        for start in range(0, len(formatted), batch_size):
             inputs = self._tokenizer(
-                formatted[batch_start : batch_start + batch_size],
+                formatted[start : start + batch_size],
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
                 max_length=max_length,
             )
             input_ids = inputs["input_ids"].to(self._model.device)
             attention_mask = inputs["attention_mask"].to(self._model.device)
-
             with self._torch.inference_mode():
                 outputs = self._model(
                     input_ids=input_ids,
@@ -164,45 +188,74 @@ def extract_batch(
                     use_cache=False,
                 )
 
-            hidden_states = outputs.hidden_states
-            token_mask = attention_mask.bool()
-            token_count = token_mask.sum(dim=1, keepdim=True)
-            positions = self._torch.arange(token_mask.shape[1], device=token_mask.device).expand_as(
-                token_mask
+            layers = []
+            for layer in self._selected_layers:
+                hidden = outputs.hidden_states[layer].float()
+                # The checkpoint consumes the last real token from each selected layer.
+                token_mask = attention_mask.to(hidden.device).bool()
+                positions = self._torch.arange(hidden.shape[1], device=hidden.device)
+                last_token = (
+                    positions.expand_as(token_mask).masked_fill(~token_mask, -1).max(dim=1).values
+                )
+                pooled = hidden[
+                    self._torch.arange(hidden.shape[0], device=hidden.device), last_token
+                ]
+                layers.append(pooled.cpu().numpy())
+            predictions.append(self._predict(layers))
+
+        return self._numpy.ascontiguousarray(
+            self._numpy.concatenate(predictions, axis=0), dtype=self._numpy.float32
+        ).tobytes()
+
+    def _predict(self, layers: list[Any]) -> Any:
+        np = self._numpy
+        torch = self._torch
+        stacked = np.stack(layers)
+        if stacked.ndim != 3 or stacked.shape[0] != len(self._selected_layers):
+            raise ValueError("encoder returned invalid selected-layer features")
+        standardized = (stacked - self._layer_mean[:, None, :]) / self._layer_std[:, None, :]
+        pooled = standardized.mean(axis=0)
+        scaled = (pooled - self._scaler_mean) / self._scaler_scale
+        features = torch.from_numpy(
+            np.ascontiguousarray(
+                (scaled - self._pca_mean) @ self._pca_components.T,
+                dtype=np.float32,
             )
-            last_index = positions.masked_fill(~token_mask, -1).max(dim=1).values
-            batch_index = self._torch.arange(token_mask.shape[0], device=token_mask.device)
-
-            for layer in layers:
-                hidden = hidden_states[layer].float()
-                if "last" in pools:
-                    all_last[layer].append(hidden[batch_index, last_index].cpu())
-                if "mean" in pools:
-                    masked = hidden.masked_fill(~token_mask.unsqueeze(-1), 0)
-                    all_mean[layer].append((masked.sum(dim=1) / token_count).cpu())
-
-            del outputs, hidden_states, input_ids, attention_mask
-
-        return {
-            "hidden_last": {
-                layer: self._torch.cat(rows).contiguous().numpy().tobytes()
-                for layer, rows in all_last.items()
-            },
-            "hidden_mean": {
-                layer: self._torch.cat(rows).contiguous().numpy().tobytes()
-                for layer, rows in all_mean.items()
-            },
-            "n_layers": self.n_layers,
-            "hidden_dim": self.hidden_dim,
-        }
+        )
+
+        members = []
+        with torch.inference_mode():
+            for state in self._states:
+                logits = []
+                for index in range(len(self._models)):
+                    adapter = torch.nn.functional.relu(
+                        torch.nn.functional.linear(
+                            features,
+                            state[f"adapters.{index}.weight"],
+                            state[f"adapters.{index}.bias"],
+                        )
+                    )
+                    trunk = torch.nn.functional.relu(
+                        torch.nn.functional.linear(
+                            adapter,
+                            state["trunk.0.weight"],
+                            state["trunk.0.bias"],
+                        )
+                    )
+                    logits.append(
+                        torch.nn.functional.linear(
+                            trunk,
+                            state[f"heads.{index}.weight"],
+                            state[f"heads.{index}.bias"],
+                        )
+                    )
+                members.append(torch.sigmoid(torch.cat(logits, dim=1)))
+        return torch.stack(members).mean(dim=0).contiguous().numpy()
 
     def unload(self) -> None:
         self._model = None
         self._tokenizer = None
-        self.n_layers = 0
-        self.hidden_dim = 0
-        if self._torch is not None:
-            if self._torch.cuda.is_available():
-                self._torch.cuda.empty_cache()
-            if hasattr(self._torch, "mps") and self._torch.backends.mps.is_available():
-                self._torch.mps.empty_cache()
+        if self._torch.cuda.is_available():
+            self._torch.cuda.empty_cache()
+        if hasattr(self._torch, "mps") and self._torch.backends.mps.is_available():
+            self._torch.mps.empty_cache()