Skip to content

feat(server): OpenAI-compatible logprobs for chat and legacy completions - #224

Open
Artemowka22 wants to merge 5 commits into
FlashML-org:mainfrom
Artemowka22:feat/openai-logprobs
Open

Artemowka22 wants to merge 5 commits into
FlashML-org:mainfrom
Artemowka22:feat/openai-logprobs

Conversation

@Artemowka22

Copy link
Copy Markdown

Summary

Implements logprobs for sampled tokens across the whole pipeline; today /v1/completions rejects the field ("logprobs is not supported") and /v1/chat/completions silently swallows it (extra="allow").

  • chat: logprobs: true + top_logprobs: 0..20 → each choice carries logprobs.content[] entries {token, logprob, bytes, top_logprobs[]}, streaming and non-streaming.
  • completions (legacy): logprobs: 0..5{tokens, token_logprobs, top_logprobs, text_offset}. echo together with logprobs stays rejected (prompt logprobs need prefill logits and are out of scope here).

Design

  • Reported values are raw model logprobs: log_softmax over the pre-temperature logits, so temperature/top-k/top-p do not change what is reported (matches vLLM's default) and greedy eval harnesses get the true model distribution.
  • Zero cost when off: the sampler computes nothing unless some request in the batch asked; per-step cost when on is one log_softmax + topk + a small D2H copy that rides the existing copy_done_event.
  • Engine → scheduler → detokenizer → API plumbing via optional fields (DetokenizeMsg.chosen_logprob/top_ids/top_logprobs, UserReply.logprobs), all defaulting to None — wire-compatible with older peers.
  • Token strings for top-k alternatives come from the detokenizer worker (tokenizer.decode([id])); the bytes field carries UTF-8 so clients can reassemble partial-UTF-8 pieces, same trade-off OpenAI documents.
  • Stop-string trimming can drop the visible text of final tokens; logprob entries still cover every sampled token.

Why

Any evaluation gate worth trusting (teacher-forced agreement/KL against a reference checkpoint, perplexity tracking of quantized variants) needs token logprobs from the OpenAI endpoint; with the radix prefix cache, per-position 1-token continuation calls make teacher-forced scoring practical without echo support.

Test plan

  • CPU-only unit tests for the sampler math (tests/engine/test_sample_logprobs.py) and the entry builder (tests/tokenizer/test_logprobs_entry.py)
  • FakeState API tests for both endpoints, streaming and not, plus the validation matrix (tests/server/test_logprobs_api.py) — no GPU, no weights, no network
  • full server/engine/tokenizer suites pass

Adds the engine half of OpenAI logprobs support. SamplingParams gains
logprobs/top_logprobs; the sampler computes log_softmax over the PRE-temperature
logits (raw model distribution, so temperature/top-k/top-p do not change reported
values), gathers the chosen token and batch-max top-k, and ships CPU copies
covered by the existing copy_done_event. The scheduler attaches per-request
values (cut to each request's own top_logprobs) to DetokenizeMsg; the detokenizer
builds a neutral entry (token text via single-id decode + UTF-8 bytes, so clients
can reassemble partial-UTF-8 pieces) onto UserReply.logprobs.

Zero cost when off: no row asked -> no mask tensor, no log_softmax, ForwardOutput
carries None. Message fields default to None, so old and new peers interoperate.
Stop-string trimming can hide final visible text; entries still cover every
sampled token.
API half of logprobs support. Chat: logprobs + top_logprobs (0..20) yield
choice.logprobs.content entries {token, logprob, bytes, top_logprobs[]},
streaming and non-streaming. Completions: the legacy integer field (0..5) yields
{tokens, token_logprobs, top_logprobs, text_offset} with absolute offsets across
the stream; echo+logprobs stays rejected (prompt logprobs need prefill logits
and are out of scope here).

The protocol-neutral event layer carries entries on ContentDelta (a list -- parser
buffering can release several tokens' text in one delta); entries always
accumulate on GenResult for the non-streaming path, and reasoning/tool-call
buffering carries pending entries onto the next content delta. Formatting lives
in server/logprobs.py; wire compatibility follows the OpenAI shapes.
@HaileyStorm

Copy link
Copy Markdown

Fresh compatibility check from current main (58f4b9ec0e166205c4dfd0c6ec184ea83b5957e6): git merge-tree --write-tree origin/main origin/pr224 completed without conflicts. CPU-only focused tests on PR head ea7ae6d6d718ecd2c648474145aab8cf72f5148a passed: tests/engine/test_sample_logprobs.py, tests/server/test_logprobs_api.py, and tests/tokenizer/test_logprobs_entry.py — 12 passed in 3.63s. I was about to build the same surface for Flash-Next/ECS, found this PR during duplicate checking, and am avoiding a competing implementation. The raw pre-temperature semantics and no-work-when-disabled boundary are exactly the useful general primitive.

@HaileyStorm

Copy link
Copy Markdown

RTX 5090 / torch 2.11 cu130 microbenchmark on Qwen3.8 vocab size 248,320: the PR implementation shape (log_softmax(float logits) + chosen gather + top-20) measured 97.7 us/call at BS1, 109.1 us BS2, 115.8 us BS4 (30 warmups, 300 timed iterations, CUDA events). A split topk(logits) + logsumexp variant was not better at BS1/2 (112.6/113.6 us for top-20; 112.0 us at BS4). So the current straightforward implementation is a reasonable fast default; no speculative kernel rewrite recommended without end-to-end evidence.

@HaileyStorm

Copy link
Copy Markdown

One correctness issue before merge: streaming semantic parsing currently misattributes hidden-token logprobs to later visible content. In _generate_events_impl, every ack.logprobs is appended to pending_logprobs, but _content_delta() drains the entire pending list only when visible content is emitted. test_reasoning_logprob_is_carried_to_next_content_delta explicitly expects the logprob for thought inside <think>...</think> to be attached to the later visible answer delta. That breaks token/content alignment, can expose hidden reasoning token strings, and differs from non-streaming (which retains every sampled-token entry) and from the end-of-stream path (which drops undrained entries). I recommend fail-closing chat logprobs when semantic reasoning/tool parsing can hide/reclassify tokens, or adding token-aware routing so each public logprob entry is emitted only with its corresponding public content token. Please do not carry hidden entries onto the next visible delta. The raw legacy-completions path does not have this semantic-layer ambiguity.

Review finding on FlashML-org#224 (HaileyStorm): streaming chat carried logprob entries
of hidden reasoning tokens onto the next visible content delta -- pending
entries accumulated before the reasoning parser classified the text, and
_content_delta() drained the whole list into whatever visible chunk came
next. That broke token<->content alignment and leaked hidden reasoning token
strings through logprobs entries; non-streaming meanwhile kept every sampled
entry and end-of-stream dropped undrained ones -- three behaviors for one
surface.

Close it at both layers:

- /v1/chat/completions rejects logprobs=true with a 400 when the server runs
  a reasoning parser or the request enables tool parsing: the semantic layer
  can hide or reclassify tokens, so entries cannot be aligned 1:1 with
  visible content tokens. tool_choice="none" keeps logprobs available; the
  raw /v1/completions path (no semantic layer) is unchanged.

- _generate_events_impl now collects entries only on the passthrough path
  (no reasoning parser, no tool parsing), so a hidden-token entry can never
  ride a later visible delta even for a caller that skips API validation.

Replaces test_reasoning_logprob_is_carried_to_next_content_delta (which
enshrined the carry) with fail-close coverage for both conflict cases, the
tool_choice="none" pass-through, and a no-carry guard test.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@Artemowka22

Copy link
Copy Markdown
Author

Confirmed and fixed — thanks for the precise diagnosis (and for the perf check and for deferring your own implementation).

You were right on all three counts: entries were appended before the reasoning parser classified the ack's text, _content_delta() drained the whole pending list into the next visible chunk, and test_reasoning_logprob_is_carried_to_next_content_delta enshrined exactly that behavior. The new commit closes it at both layers, taking your fail-close option:

  • /v1/chat/completions now rejects logprobs: true with a 400 up front when the server runs a reasoning parser or the request enables tool parsing (tool_choice: "none" keeps logprobs available; the raw /v1/completions path is untouched).
  • _generate_events_impl collects entries only on the passthrough path (no reasoning parser, no tool parsing), so a hidden-token entry can never ride a later visible delta even for a caller that skips API validation.

The carry test is replaced by a no-carry guard test plus fail-close coverage for both conflict cases and the tool_choice: "none" pass-through.

Exact token-aware routing (each public entry emitted only with its own public content token) would need the reasoning-parser interface to report per-token consumption; happy to do that as a follow-up if maintainers want logprobs and reasoning to coexist on chat.

gdevenyi added a commit to gdevenyi/FreeToken that referenced this pull request Sep 5, 2026
…gprobs for chat and legacy completions

Upstream FlashML-org#224 at 855650d, merged onto deploy/chatdnp for the PR sweep.
Conflicts: engine.py keeps FlashML-org#231's stats readout before the logprobs-aware return;
openai_api.py keeps the vision `images` argument and FlashML-org#222's disconnect-watching drain
with the logprobs entries added; generation.py keeps FlashML-org#266's marker filter and routes every
content delta through FlashML-org#224's _content_delta so the logprobs entries ride the filtered text.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
gdevenyi added a commit to gdevenyi/FreeToken that referenced this pull request Sep 5, 2026
…ollows ForwardOutput (FlashML-org#224)

The donate and hit-admission barriers of PR FlashML-org#287 call torch.cuda.synchronize on the
cache manager's device; the scheduler unit tests run that path on a CPU device, so guard
the call on device.type. The drained-forward fake in test_abort_inflight_prefill built a
bare tuple, which PR FlashML-org#224's scheduler no longer unpacks: build a ForwardOutput instead.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
@gdevenyi

gdevenyi commented Sep 5, 2026

Copy link
Copy Markdown

Tried on 2 x RTX 6000 Ada (sm_89) serving Qwen3.8-Flash-Next (RadixArk NVFP4) at TP=2, offload backend, fp8 KV pool of 8 x 262,144 tokens, merged onto my deploy branch (main af71ba4 + #385/#386/#389/#392/#354 and ten other open PRs), tests run on the box, then put in production.

Merged with conflicts against #231 (engine.py), #222 (openai_api.py) and #266 (generation.py, the content deltas now go through _content_delta after the marker filter). Findings:

  • /v1/chat/completions with logprobs: true answers 400 here: logprobs on /v1/chat/completions are not supported when the server runs a reasoning parser. Every Qwen3.8 deployment runs one, so on this model only the legacy /v1/completions logprobs are usable (they work: tokens, token_logprobs, top_logprobs).
  • The scheduler now reads ForwardOutput attributes, which breaks the tuple fake in tests/scheduler/test_abort_inflight_prefill.py (3 tests, 'tuple' object has no attribute 'next_tokens_cpu'); I changed the fake to build a ForwardOutput.
  • tests/engine/test_sample_logprobs.py 5 passed on the GPU; serving numbers unchanged with logprobs off.

The logprobs change moved the drain from positional unpacking to named
ForwardOutput attributes (next_tokens_cpu, copy_done_event, the *_cpu
logprob columns), which the bare-tuple fake in the in-flight-abort
scheduler tests predates: three of them died with 'tuple' object has no
attribute 'next_tokens_cpu' on a GPU run. Caught by gdevenyi on the PR
thread; a macOS run cannot see it (the file already fails on the missing
kernel packages there, identically on both sides of the series).

A namespace mirroring the drain-facing fields keeps the test off the
engine import chain, whose kernel packages a scheduler unit test should
not require.

Assisted-by: Claude
@Artemowka22

Copy link
Copy Markdown
Author

Thanks for running it in production — both points addressed:

The scheduler test fake: fixed in a new commit. _as_last_data now hands _process_last_data an object with ForwardOutput's drain-facing fields (next_tokens_cpu, copy_done_event, the *_cpu logprob columns set to None). We went with a namespace mirroring the fields rather than constructing the real ForwardOutput, only because importing it pulls the whole engine chain (attention/moe/kernel packages) into a scheduler unit test that today runs without them — the shape is otherwise exactly your fix, and credit for catching it on a GPU run; on macOS that file already fails on the missing kernel packages, identically on both sides of the series, so our gate could not see it.

Chat logprobs 400 on Qwen3.8 deployments: confirmed, and it is the intended trade-off rather than an oversight — the fail-close option was chosen in the earlier review round (over silently misaligned entries, which is what the previous behavior produced: hidden reasoning-token logprobs riding the next visible content delta). While a reasoning parser is routing tokens out of visible content there is no 1:1 alignment to report honestly, and /v1/completions keeps full-fidelity logprobs as you saw. The real fix for chat — reasoning-aware routing, so entries follow tokens through the parser and visible-content entries stay aligned — needs the reasoning-parser interface to expose per-token classification, which is more than this PR should carry; we would like to do that as a follow-up once this lands. If serving Qwen3.8 with chat logprobs is a pressing need on your deployment, a per-request escape hatch (e.g. logprobs honored when the request disables thinking) could be discussed there too.

Two conflicts:

- python/freetoken/server/openai_api.py: upstream threads the configurable
  default output budget (default_max_tokens, FlashML-org#411) into the resolve_sampling
  call that this branch had extracted into a variable to attach the logprobs
  fields. Kept the variable; its construction now passes default_max_tokens.
- python/freetoken/tokenizer/server.py: import conflict between upstream's
  get_mm_processor (image input, FlashML-org#454) and this branch's build_logprobs_entry.
  Kept both.

Assisted-by: Claude
MT-z added a commit to MT-z/FreeToken that referenced this pull request Sep 16, 2026
…with FlashML-org#393)

Artemowka22's FlashML-org#224, open since 2026-08-26. FlashML-org#393 declines logprobs on both routes and its
own compatibility table points here ("400 (see FlashML-org#224 for the sampled-token logprobs)"), so
the two are complementary by design -- and both rewrite the same sampling and message
plumbing, which is where the work was.

Before taking it, measured what this branch actually did with a logprobs request:

    /v1/completions   logprobs=5   400 "logprobs is not supported"        declared
    /v1/chat/...      logprobs=true, top_logprobs=5   200, no logprobs    silent

`ChatCompletionRequest` carries `extra="allow"` and declares neither field, so they were
swallowed before any validator saw them. FlashML-org#393's table says this route returns 400; it
returned 200 and dropped them. Same on pr393 alone, so it is upstream, not this stack.

Merge: 9 files, 20 conflict blocks, all from FlashML-org#393's `n` sampling (many uids per request)
meeting FlashML-org#224's single-uid shape. Resolved by keeping FlashML-org#393's structure and hanging the
logprobs off it -- the multi-uid `choices` loop, `_completion_chunk`, `_resolve_sampling`
(which also fills min_p, penalties, logit_bias, stop_token_ids) all stay. FlashML-org#224's rejections
of `echo`, `suffix` and `logit_bias` were dropped: FlashML-org#393 implements all three, and those
lines are older than it. Its `echo` + `logprobs` rejection is kept -- that one is real, the
prompt logits are not there.

Two failures had to be split apart before either could be fixed:

  - `tests/scheduler/test_abort_inflight_prefill.py`, 4 tests: **pr224 fails these on its
    own base too**, so it is the PR's regression, not the merge. `ForwardOutput` grows from
    3 fields to 6, `_process_last_data` switches to reading it by name, and the upstream
    test (there since 3af9d90) hands it a bare 3-tuple. Fixed by having the test build a
    real `ForwardOutput`.
  - `tests/engine/test_sample_logprobs.py`, 1 test: this one is the merge. FlashML-org#393 adds the
    `needs_logits_processing` property to SamplingParams and FlashML-org#224's new stub is a
    SimpleNamespace without it. Fixed in the stub, not by loosening `_plan` -- a getattr
    default there would hide a real type mismatch. Third instance of this pattern today
    (FlashML-org#354's `k_scale`, FlashML-org#198's direct `num_page_override`).

Verified on the wire, Ornith-1.5-35B-A3B-NVFP4:

    /v1/completions logprobs=3   ' Paris' -0.7205, top3 [' Paris' -0.72, ' a' -2.16,
                                 '\n' -2.60], text_offset [0,6,7,8]; streaming carries
                                 per-chunk logprobs
    /v1/chat/...    with the default qwen3 reasoning parser: 400 naming the reason
                    ("reasoning tokens are hidden from message content ... use
                    /v1/completions"); with --reasoning-parser off: full entries with
                    token/logprob/bytes/top_logprobs, streaming too

FlashML-org#393's side survives: `n=2` returns 2 choices with correct usage, `echo` echoes, `suffix`
and `logit_bias` are accepted. FlashML-org#224's own guards fire: `echo`+`logprobs` and `logprobs=9`
both 400 with their reasons.

Default path, no logprobs requested, 3 single-stream decodes of 399 tokens:

    try/all   91.7 / 95.2 / 94.1 tok/s   median 94.1
    try/224   92.5 / 95.3 / 94.4 tok/s   median 94.4

Full suite: 11 failed, 1849 passed, 60 skipped -- the same 11 as try/all, plus the 13
tests the PR adds.

Assisted-by: Claude Opus 5
MT-z added a commit to MT-z/FreeToken that referenced this pull request Sep 16, 2026
…he stubs rather than the code

Two stub-vs-contract mismatches surfaced by the FlashML-org#224 merge, both fixed on the test side so
the production code keeps saying what it means:

`ForwardOutput` grew from 3 fields to 6 and `_process_last_data` now reads it by name, but
`tests/scheduler/test_abort_inflight_prefill.py` (upstream since 3af9d90) handed it a bare
3-tuple. It builds a real `ForwardOutput` now. pr224 fails these four on its own base too,
so this is the PR's regression rather than the merge's.

`tests/engine/test_sample_logprobs.py` builds SamplingParams as a SimpleNamespace without
`needs_logits_processing`, the property FlashML-org#393 added and `_plan` reads. Added to the stub. A
getattr default in `_plan` would have hidden a real type mismatch instead.

Assisted-by: Claude Opus 5
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.

3 participants