unlock unlimited code runs for $79, limited time →
NOSSLETTER.techLIVE
back to archive

[Feature] Per-row candidate IDs for prefill token scoring (M2 of #56860)

Language: Python+365 -161feature
concepts demonstrated
NumPy arrays and C-contiguous memory layoutType annotations with Python's typing.Annotated and Pydantic schema integrationInput validation and shape checking for multi-dimensional arraysEfficient zero-copy data handling for inter-process communication
The Issue & Context

Previously, the system used a single set of candidate token IDs for scoring all prompt tokens uniformly, limiting the ability to tailor scoring per token position in causal language model prompts.

Root Cause Analysis

The lack of per-row candidate IDs meant that on-policy distillation teachers could not score each response position against the student's own top-K tokens, reducing scoring accuracy and flexibility.

Solution Strategy & Architectural Pattern

Introduced a per-row integer array input for candidate token IDs, allowing each prompt token position to have its own candidate list. This array is validated, normalized to a C-contiguous int32 ndarray, and integrated into both Python and Rust frontends with consistent validation and zero-copy data paths. Padding with -1 is handled by masking to -inf in logprob outputs.

Key Takeaway for Contributors

Using per-row candidate IDs significantly improves scoring precision in prompt token evaluation while maintaining performance through efficient array handling and consistent cross-language validation.