Skip to content

Add requested generation logprob token IDs - #35

Merged
lukealonso merged 1 commit into
NVIDIA:mainfrom
yatesdr:feat/logprob-token-ids
Oct 23, 2026
Merged

lukealonso merged 1 commit into
NVIDIA:mainfrom
yatesdr:feat/logprob-token-ids

Conversation

@yatesdr

@yatesdr yatesdr commented Oct 23, 2026

Copy link
Copy Markdown
Collaborator

Addresses #26 (the scoring primitive; the decision routes follow in separate PRs).

Problem

Decision and distillation clients need the logprob of specific tokens at each generated position: label tokens for classification, or a teacher's top set for KLD. Today Flare returns only the natural top 20, so a label outside it has no value at all.

Change

logprob_token_ids: [id, ...] on Chat, Completions and Responses, matching vLLM's field:

  • With generation logprobs enabled, the listed IDs replace the natural alternatives at every generated position, in caller order. The sampled token is appended when absent.
  • At most 128 IDs. They must be integers within the vocabulary; duplicates are dropped and null or empty means omitted. Messages rejects the field.
  • Values are the raw FP32 logit minus the existing log-sum-exp, before sampling transforms. They are never renormalized.
  • The existing finishing kernel gathers the requested IDs with its existing normalization. There is no new launch, full-vocabulary copy or graph family.
  • Readback copies only the requested width. Requests without IDs copy exactly what they do today.
  • Natural top-k, prompt scores and echo behavior are unchanged.
  • Older packages keep their natural limits and reject nonempty IDs before admission, asking for a re-export.

Tests

Helios 893, integration 31, protocol 47 and Flare 134 pass on CPU, plus the Python component tests. An earlier revision passed CUDA, TP2/TP4 replay and 148 live API checks on Qwen3.8, with requested values within 9.0e-7 of an FP64 reference. DeepSeek V4.1 and this revision's GPU runs are still to do.

Gather up to 128 IDs through raw-logprob reporting and copy the requested frame width. Preserve older package layouts and prompt-scoring validation.

CPU: release build, qualification check and extension install passed. Helios 893 passed, 4 Engram ENOMEM failures, 22 ignored; integration 31 passed, 12 ignored; Flare 134 passed, 1 ignored; protocols 47 passed; Python 10 passed, 1 skipped. CLI clippy passed; Helios has 3 existing range errors; full formatting has existing differences. GPU and service checks remain for the test round.

Design-Signed-Off-By: yatesdr; design-review=6aeb
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants