Skip to content

Cap coordinator prefill/encode legs to a single token per schema - #2202

Open
revit13 wants to merge 4 commits into
llm-d:mainfrom
revit13:tokens1
Open

Cap coordinator prefill/encode legs to a single token per schema#2202
revit13 wants to merge 4 commits into
llm-d:mainfrom
revit13:tokens1

Conversation

@revit13

@revit13 revit13 commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

What type of PR is this?
/kind bug

The coordinator's synthetic prefill and encode legs only need a single output
token, but the request bodies they build did not consistently constrain the
output length across the three request schemas. This adds a single
coordinator-local helper, capSingleTokenOutput, that the prefill and encode
steps call for every schema:

  • Pins max_tokens to 1.
  • Strips min_tokens. It defaults to 0 in vLLM, so removing it keeps the
    min_tokens <= max_tokens constraint satisfied without raising the floor
    above max_tokens = 1.
  • Caps max_completion_tokens to 1 only when the client already sent it
    (in practice only the chat completions schema), never adding it otherwise.
  • Forces stream to false and strips stream_options, so no leg returns a
    streamed response the coordinator cannot decode.

The token limits live in the sampling_params sub-map for the generate schema
(synthesized when absent) and at the top level for the OpenAI completions and
chat completions schemas, so the helper targets the correct location per
schema.

This is the coordinator-side counterpart to #2041, which capped the single-token
constraint on the sidecar prefill leg for chat completions. This change is
deliberately scoped to the coordinator only; the equivalent sidecar connector
paths are handled separately.

Verified end-to-end against the coordinator together with #2077, running real vLLM instances on the following endpoints:

  • generate endpoint
  • chat completion endpoint
  • completion endpoint

Fixes #

Release note (write NONE if no user-facing change):

NONE

@revit13
revit13 requested review from a team, roytman and shmuelk as code owners July 27, 2026 11:13
@revit13
revit13 requested review from ahg-g and elevran July 27, 2026 11:13
@github-actions github-actions Bot added size/L Denotes a PR that changes 100-499 lines, ignoring generated files. kind/cleanup labels Jul 27, 2026
@revit13
revit13 marked this pull request as draft July 27, 2026 11:13
@github-actions github-actions Bot added kind/bug Categorizes issue or PR as related to a bug. and removed kind/cleanup kind/bug Categorizes issue or PR as related to a bug. labels Jul 27, 2026
@revit13
revit13 marked this pull request as ready for review July 28, 2026 04:55
// gateway.Client.Request, which is what logs full request bodies at TRACE for
// the prefill and encode legs. Log the redacted body here so the decode leg's
// fields (min_tokens, max_tokens, ...) are observable at the same verbosity.
if v := logger.V(logutil.TRACE); v.Enabled() {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

do we need it? you changed the RedactBody method visibility only for that.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, this is intentional. The decode leg is proxied straight to the gateway and never passes through gateway.Client.Request, which is what logs the full request body at TRACE for the prefill and encode legs. Without this it's the only leg whose body is invisible at TRACE.
The concrete motivation is seeing the decode leg's min_tokens/max_tokens (and the other sampling fields). Those sit either top-level or nested under sampling_params depending on format, so logging the redacted body captures them regardless of location without format-specific extraction here.

Comment thread pkg/coordinator/steps/tokens.go Outdated
// Intentionally distinct from the sidecar's reqcommon.PrimeSingleTokenRequest
// for now.
// TODO: unify the two into one shared single-token helper in a future refactor.
func capSingleTokenOutput(body map[string]any, format gateway.RequestFormat) {

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it is more suitable to be part of utils.go and do not create two separate files.

Comment thread pkg/coordinator/steps/tokens.go Outdated
// capSingleTokenOutput rewrites body into a single-output-token, non-streaming
// request for the synthetic prefill and encode legs.
//
// Intentionally distinct from the sidecar's reqcommon.PrimeSingleTokenRequest

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why cannot we use reqcommon.PrimeSingleTokenRequest maybe to extend it?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tried having capSingleTokenOutput call PrimeSingleTokenRequest, but it complicates more than it saves. Prime always writes max_tokens at the top level, whereas the generate format needs it nested under sampling_params. So the reuse ends up setting the top-level max_tokens via Prime and then deleting it to re-nest, which is harder to follow than just placing it correctly once. Prime also doesn't strip min_tokens, which we need here.

Comment thread pkg/coordinator/steps/tokens.go Outdated
target = sp
}

target[reqcommon.FieldMaxTokens] = 1

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in the case of generate, should FieldMaxTokens be part of sampling-parameters?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes. For the generate format the sampling-time fields live under sampling_params, which is why the function targets that map. It's consistent with how the generate body is built elsewhere, e.g. buildPrefillBody's generate branch nests kv_transfer_params under sampling_params.extra_args, and the "generate preserves other sampling_params entries" test asserts the resulting sampling_params: {max_tokens: 1, extra_args: {...}} shape. Top-level max_tokens is the chat/completions shape, so the function branches on format to place it correctly.

revit13 added 4 commits July 28, 2026 13:25
Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
Signed-off-by: Revital Sur <eres@il.ibm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

kind/bug Categorizes issue or PR as related to a bug. size/L Denotes a PR that changes 100-499 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants