Skip to content

Commit 9e7f828

Browse files
revit13revit13
committed
Cap coordinator prefill/encode legs to a single token per schema
Signed-off-by: Revital Sur <eres@il.ibm.com>
1 parent d397ccd commit 9e7f828

6 files changed

Lines changed: 211 additions & 32 deletions

File tree

‎pkg/coordinator/steps/encode.go‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -202,21 +202,20 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [
202202
},
203203
},
204204
}
205-
reqcommon.PrimeSingleTokenRequest(body)
205+
capSingleTokenOutput(body, format)
206206
return body
207207
default:
208-
samplingParams := map[string]any{}
209-
reqcommon.CapSingleTokenOutput(samplingParams)
210-
return map[string]any{
208+
body := map[string]any{
211209
"model": reqCtx.Model,
212210
"token_ids": tokenIDs,
213211
"features": map[string]any{
214212
"mm_hashes": map[string][]string{ModalityImage: {entry.Hash}},
215213
"mm_placeholders": map[string][]any{ModalityImage: {map[string]any{"offset": 1, "length": entry.Placeholder.Length}}},
216214
"kwargs_data": map[string][]string{ModalityImage: {entry.KwargsData}},
217215
},
218-
reqcommon.FieldSamplingParams: samplingParams,
219216
}
217+
capSingleTokenOutput(body, format)
218+
return body
220219
}
221220
}
222221

‎pkg/coordinator/steps/encode_test.go‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -352,10 +352,10 @@ func TestEncodeStep_ChatCompletionsFormat(t *testing.T) {
352352
}
353353

354354
// TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens is a
355-
// regression test: buildEncodeBody used to always emit max_tokens=1 without
356-
// regard to which token-limit field the client used, so a reasoning-model
357-
// client's max_completion_tokens was never capped in the encoder sub-request.
358-
func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) {
355+
// The encode chat sub-request is built fresh from the request context and does
356+
// not carry the client's sampling fields, so max_completion_tokens is not
357+
// propagated and is never injected: max_tokens=1 alone caps output.
358+
func TestEncodeStep_ChatCompletionsFormat_OmitsMaxCompletionTokens(t *testing.T) {
359359
var receivedBody map[string]any
360360

361361
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -401,8 +401,11 @@ func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T)
401401
t.Fatalf("unexpected error: %v", err)
402402
}
403403

404-
if receivedBody["max_completion_tokens"] != float64(1) {
405-
t.Fatalf("expected encode sub-request max_completion_tokens capped to 1, got %v", receivedBody["max_completion_tokens"])
404+
if receivedBody["max_tokens"] != float64(1) {
405+
t.Fatalf("expected encode sub-request max_tokens capped to 1, got %v", receivedBody["max_tokens"])
406+
}
407+
if _, ok := receivedBody["max_completion_tokens"]; ok {
408+
t.Fatalf("expected encode sub-request to omit max_completion_tokens, got %v", receivedBody["max_completion_tokens"])
406409
}
407410
}
408411

@@ -539,7 +542,8 @@ func TestEncodeStep_BuildsCorrectTokenIDs(t *testing.T) {
539542

540543
// TestEncodeStep_GenerateFormat_CapsSingleToken verifies the generate-format
541544
// encoder sub-request caps output to a single token: sampling_params carries
542-
// both max_tokens=1 and min_tokens=1 (min_tokens <= max_tokens, as vLLM requires).
545+
// max_tokens=1 and strips min_tokens (it defaults to 0, keeping min_tokens <=
546+
// max_tokens).
543547
func TestEncodeStep_GenerateFormat_CapsSingleToken(t *testing.T) {
544548
var samplingParams map[string]any
545549

@@ -572,7 +576,7 @@ func TestEncodeStep_GenerateFormat_CapsSingleToken(t *testing.T) {
572576
if samplingParams["max_tokens"] != float64(1) {
573577
t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"])
574578
}
575-
if samplingParams["min_tokens"] != float64(1) {
576-
t.Fatalf("expected sampling_params.min_tokens=1, got %v", samplingParams["min_tokens"])
579+
if _, ok := samplingParams["min_tokens"]; ok {
580+
t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"])
577581
}
578582
}

‎pkg/coordinator/steps/prefill.go‎

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
137137
switch format {
138138
case gateway.FormatChatCompletions:
139139
body := maps.Clone(reqCtx.Body)
140-
reqcommon.PrimeSingleTokenRequest(body)
140+
capSingleTokenOutput(body, format)
141141
tokens := map[string]any{
142142
"token_ids": reqCtx.TokenIDs,
143143
}
@@ -166,7 +166,7 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
166166
"prompt": prompt,
167167
reqcommon.FieldKVTransferParams: kvParams,
168168
}
169-
reqcommon.CapSingleTokenOutput(body)
169+
capSingleTokenOutput(body, format)
170170
if features != nil {
171171
body["features"] = features
172172
}
@@ -176,18 +176,17 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
176176
return body, nil
177177

178178
default:
179-
samplingParams := map[string]any{
180-
"extra_args": map[string]any{
181-
reqcommon.FieldKVTransferParams: kvParams,
182-
},
183-
}
184-
reqcommon.CapSingleTokenOutput(samplingParams)
185179
body := map[string]any{
186-
"request_id": reqCtx.RequestID,
187-
"token_ids": reqCtx.TokenIDs,
188-
"model": reqCtx.Model,
189-
reqcommon.FieldSamplingParams: samplingParams,
180+
"request_id": reqCtx.RequestID,
181+
"token_ids": reqCtx.TokenIDs,
182+
"model": reqCtx.Model,
183+
reqcommon.FieldSamplingParams: map[string]any{
184+
"extra_args": map[string]any{
185+
reqcommon.FieldKVTransferParams: kvParams,
186+
},
187+
},
190188
}
189+
capSingleTokenOutput(body, format)
191190
if features != nil {
192191
body["features"] = features
193192
}

‎pkg/coordinator/steps/prefill_test.go‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -147,8 +147,8 @@ func TestPrefillStep_SendsCorrectGenerateRequest(t *testing.T) {
147147
if samplingParams["max_tokens"] != float64(1) {
148148
t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"])
149149
}
150-
if samplingParams["min_tokens"] != float64(1) {
151-
t.Fatalf("expected sampling_params.min_tokens=1, got %v", samplingParams["min_tokens"])
150+
if _, ok := samplingParams["min_tokens"]; ok {
151+
t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"])
152152
}
153153
extraArgs, ok := samplingParams["extra_args"].(map[string]any)
154154
if !ok {
@@ -221,13 +221,13 @@ func TestPrefillStep_CompletionsFormat(t *testing.T) {
221221
if prefillBody["request_id"] != "req-compl" {
222222
t.Fatalf("expected request_id, got %v", prefillBody["request_id"])
223223
}
224-
// Prefill leg caps output to a single token: both max_tokens and min_tokens
225-
// are pinned to 1 (min_tokens <= max_tokens, as vLLM requires).
224+
// Prefill leg caps output to a single token: max_tokens is pinned to 1 and
225+
// min_tokens is stripped (it defaults to 0, keeping min_tokens <= max_tokens).
226226
if prefillBody["max_tokens"] != float64(1) {
227227
t.Fatalf("expected max_tokens=1, got %v", prefillBody["max_tokens"])
228228
}
229-
if prefillBody["min_tokens"] != float64(1) {
230-
t.Fatalf("expected min_tokens=1, got %v", prefillBody["min_tokens"])
229+
if _, ok := prefillBody["min_tokens"]; ok {
230+
t.Fatalf("expected min_tokens to be stripped, got %v", prefillBody["min_tokens"])
231231
}
232232
// Completions format has top-level kv_transfer_params
233233
kvParams, ok := prefillBody["kv_transfer_params"].(map[string]any)

‎pkg/coordinator/steps/tokens.go‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/*
2+
Copyright 2026 The llm-d Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package steps
18+
19+
import (
20+
reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"
21+
"github.com/llm-d/llm-d-router/pkg/coordinator/gateway"
22+
)
23+
24+
// capSingleTokenOutput rewrites body into a single-output-token, non-streaming
25+
// request for the synthetic prefill and encode legs. max_tokens is pinned to 1
26+
// and min_tokens is stripped: min_tokens defaults to 0 in vLLM, so removing it
27+
// keeps min_tokens <= max_tokens satisfied without raising the floor above
28+
// max_tokens=1. These limits live in the sampling_params sub-map for the generate
29+
// schema (synthesized if absent) and at the top level for the OpenAI schemas.
30+
//
31+
// stream is forced false and stream_options stripped for every format, so no leg
32+
// returns a streamed response the coordinator cannot decode. max_completion_tokens
33+
// is capped to 1 only when the body already carries it (in practice only the chat
34+
// completions schema does), never added when absent.
35+
func capSingleTokenOutput(body map[string]any, format gateway.RequestFormat) {
36+
target := body
37+
if format == gateway.FormatGenerate {
38+
sp, ok := body[reqcommon.FieldSamplingParams].(map[string]any)
39+
if !ok {
40+
sp = map[string]any{}
41+
body[reqcommon.FieldSamplingParams] = sp
42+
}
43+
target = sp
44+
}
45+
46+
target[reqcommon.FieldMaxTokens] = 1
47+
delete(target, reqcommon.FieldMinTokens)
48+
49+
if _, ok := body[reqcommon.FieldMaxCompletionTokens]; ok {
50+
body[reqcommon.FieldMaxCompletionTokens] = 1
51+
}
52+
53+
body[reqcommon.FieldStream] = false
54+
delete(body, reqcommon.FieldStreamOptions)
55+
}

‎pkg/coordinator/steps/tokens_test.go‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
/*
2+
Copyright 2026 The llm-d Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package steps
18+
19+
import (
20+
"reflect"
21+
"testing"
22+
23+
"github.com/llm-d/llm-d-router/pkg/coordinator/gateway"
24+
)
25+
26+
func TestCapSingleTokenOutput(t *testing.T) {
27+
tests := []struct {
28+
name string
29+
format gateway.RequestFormat
30+
body map[string]any
31+
want map[string]any
32+
}{
33+
{
34+
name: "chat completions caps output fields and forces non-streaming",
35+
format: gateway.FormatChatCompletions,
36+
body: map[string]any{
37+
"model": "m",
38+
"max_tokens": 100,
39+
"min_tokens": 5,
40+
"max_completion_tokens": 100,
41+
"stream": true,
42+
"stream_options": map[string]any{"include_usage": true},
43+
},
44+
want: map[string]any{
45+
"model": "m",
46+
"max_tokens": 1,
47+
"max_completion_tokens": 1,
48+
"stream": false,
49+
},
50+
},
51+
{
52+
name: "max_completion_tokens is not added when the client omitted it",
53+
format: gateway.FormatChatCompletions,
54+
body: map[string]any{"model": "m"},
55+
want: map[string]any{
56+
"model": "m",
57+
"max_tokens": 1,
58+
"stream": false,
59+
},
60+
},
61+
{
62+
name: "completions caps max_tokens, strips min_tokens, forces non-streaming",
63+
format: gateway.FormatCompletions,
64+
body: map[string]any{"model": "m", "max_tokens": 100, "min_tokens": 5},
65+
want: map[string]any{"model": "m", "max_tokens": 1, "stream": false},
66+
},
67+
{
68+
name: "streaming is forced false and stream_options stripped",
69+
format: gateway.FormatCompletions,
70+
body: map[string]any{"stream": true, "stream_options": map[string]any{"include_usage": true}},
71+
want: map[string]any{"stream": false, "max_tokens": 1},
72+
},
73+
{
74+
name: "generate caps max_tokens and strips min_tokens inside sampling_params",
75+
format: gateway.FormatGenerate,
76+
body: map[string]any{
77+
"model": "m",
78+
"sampling_params": map[string]any{"max_tokens": 100, "min_tokens": 5},
79+
},
80+
want: map[string]any{
81+
"model": "m",
82+
"sampling_params": map[string]any{"max_tokens": 1},
83+
"stream": false,
84+
},
85+
},
86+
{
87+
name: "generate synthesizes sampling_params when absent",
88+
format: gateway.FormatGenerate,
89+
body: map[string]any{"model": "m"},
90+
want: map[string]any{
91+
"model": "m",
92+
"sampling_params": map[string]any{"max_tokens": 1},
93+
"stream": false,
94+
},
95+
},
96+
{
97+
name: "generate preserves other sampling_params entries",
98+
format: gateway.FormatGenerate,
99+
body: map[string]any{
100+
"sampling_params": map[string]any{
101+
"extra_args": map[string]any{"kv_transfer_params": "x"},
102+
},
103+
},
104+
want: map[string]any{
105+
"sampling_params": map[string]any{
106+
"max_tokens": 1,
107+
"extra_args": map[string]any{"kv_transfer_params": "x"},
108+
},
109+
"stream": false,
110+
},
111+
},
112+
}
113+
114+
for _, tt := range tests {
115+
t.Run(tt.name, func(t *testing.T) {
116+
capSingleTokenOutput(tt.body, tt.format)
117+
if !reflect.DeepEqual(tt.body, tt.want) {
118+
t.Fatalf("got %v, want %v", tt.body, tt.want)
119+
}
120+
})
121+
}
122+
}

0 commit comments

Comments
 (0)