Cap coordinator prefill/encode legs to a single token per schema - #2202
Cap coordinator prefill/encode legs to a single token per schema#2202revit13 wants to merge 4 commits into
Conversation
| // 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() { |
There was a problem hiding this comment.
do we need it? you changed the RedactBody method visibility only for that.
There was a problem hiding this comment.
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.
| // 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) { |
There was a problem hiding this comment.
it is more suitable to be part of utils.go and do not create two separate files.
| // 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 |
There was a problem hiding this comment.
why cannot we use reqcommon.PrimeSingleTokenRequest maybe to extend it?
There was a problem hiding this comment.
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.
| target = sp | ||
| } | ||
|
|
||
| target[reqcommon.FieldMaxTokens] = 1 |
There was a problem hiding this comment.
in the case of generate, should FieldMaxTokens be part of sampling-parameters?
There was a problem hiding this comment.
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.
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>
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 encodesteps call for every schema:
max_tokensto 1.min_tokens. It defaults to 0 in vLLM, so removing it keeps themin_tokens <= max_tokensconstraint satisfied without raising the floorabove
max_tokens = 1.max_completion_tokensto 1 only when the client already sent it(in practice only the chat completions schema), never adding it otherwise.
streamto false and stripsstream_options, so no leg returns astreamed response the coordinator cannot decode.
The token limits live in the
sampling_paramssub-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:
Fixes #
Release note (write
NONEif no user-facing change):