Skip to content

feat(prefill-router): run batched prefill inference #539

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nachiketb-nvidia wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(prefill-router): run batched prefill inference #539

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

4 changes: 2 additions & 2 deletions crates/prefill-router/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down
6 changes: 1 addition & 5 deletions crates/prefill-router/pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
289 changes: 171 additions & 118 deletions crates/prefill-router/python/transformers_forward.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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


Expand Down Expand Up @@ -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,
Expand All @@ -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()
Loading
Loading