Skip to content

Commit 6c2f855

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 f56f3bd commit 6c2f855

6 files changed

Lines changed: 245 additions & 12 deletions

File tree

‎pkg/coordinator/steps/encode.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -202,19 +202,20 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [
202202
},
203203
},
204204
}
205-
reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body)
205+
capSingleTokenOutput(body, format)
206206
return body
207207
default:
208-
return map[string]any{
208+
body := map[string]any{
209209
"model": reqCtx.Model,
210210
"token_ids": tokenIDs,
211211
"features": map[string]any{
212212
"mm_hashes": map[string][]string{ModalityImage: {entry.Hash}},
213213
"mm_placeholders": map[string][]any{ModalityImage: {map[string]any{"offset": 1, "length": entry.Placeholder.Length}}},
214214
"kwargs_data": map[string][]string{ModalityImage: {entry.KwargsData}},
215215
},
216-
reqcommon.FieldSamplingParams: map[string]any{reqcommon.FieldMaxTokens: 1},
217216
}
217+
capSingleTokenOutput(body, format)
218+
return body
218219
}
219220
}
220221

‎pkg/coordinator/steps/encode_test.go‎

Lines changed: 50 additions & 6 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

@@ -536,3 +539,44 @@ func TestEncodeStep_BuildsCorrectTokenIDs(t *testing.T) {
536539
}
537540
}
538541
}
542+
543+
// TestEncodeStep_GenerateFormat_CapsSingleToken verifies the generate-format
544+
// encoder sub-request caps output to a single token: sampling_params carries
545+
// max_tokens=1 and strips min_tokens (it defaults to 0, keeping min_tokens <=
546+
// max_tokens).
547+
func TestEncodeStep_GenerateFormat_CapsSingleToken(t *testing.T) {
548+
var samplingParams map[string]any
549+
550+
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
551+
body, _ := io.ReadAll(r.Body)
552+
var parsed map[string]any
553+
_ = json.Unmarshal(body, &parsed)
554+
samplingParams, _ = parsed["sampling_params"].(map[string]any)
555+
_ = json.NewEncoder(w).Encode(map[string]any{
556+
"ec_transfer_params": map[string]any{"h1": map[string]any{"peer_port": 5501}},
557+
})
558+
}))
559+
defer server.Close()
560+
561+
gwClient := gateway.New(config.GatewayConfig{Address: server.URL})
562+
step, _ := NewEncodeStep(gwClient, map[string]any{"use_openai_format": false})
563+
564+
reqCtx := &pipeline.RequestContext{
565+
RequestID: "req-gen-cap",
566+
Model: "test",
567+
TokenIDs: []int{1, 32000, 32000, 32000, 2345},
568+
MultimodalEntries: []pipeline.MultimodalEntry{
569+
{Index: 0, Hash: "h1", KwargsData: "dGVzdA==", Placeholder: pipeline.PlaceholderRange{Offset: 1, Length: 3}},
570+
},
571+
}
572+
573+
if err := step.Execute(context.Background(), reqCtx); err != nil {
574+
t.Fatalf("unexpected error: %v", err)
575+
}
576+
if samplingParams["max_tokens"] != float64(1) {
577+
t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"])
578+
}
579+
if _, ok := samplingParams["min_tokens"]; ok {
580+
t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"])
581+
}
582+
}

‎pkg/coordinator/steps/prefill.go‎

Lines changed: 3 additions & 3 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, reqCtx.Body)
140+
capSingleTokenOutput(body, format)
141141
tokens := map[string]any{
142142
"token_ids": reqCtx.TokenIDs,
143143
}
@@ -164,9 +164,9 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
164164
"request_id": reqCtx.RequestID,
165165
"model": reqCtx.Model,
166166
"prompt": prompt,
167-
reqcommon.FieldMaxTokens: 1,
168167
reqcommon.FieldKVTransferParams: kvParams,
169168
}
169+
capSingleTokenOutput(body, format)
170170
if features != nil {
171171
body["features"] = features
172172
}
@@ -181,12 +181,12 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req
181181
"token_ids": reqCtx.TokenIDs,
182182
"model": reqCtx.Model,
183183
reqcommon.FieldSamplingParams: map[string]any{
184-
reqcommon.FieldMaxTokens: 1,
185184
"extra_args": map[string]any{
186185
reqcommon.FieldKVTransferParams: kvParams,
187186
},
188187
},
189188
}
189+
capSingleTokenOutput(body, format)
190190
if features != nil {
191191
body["features"] = features
192192
}

‎pkg/coordinator/steps/prefill_test.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,9 @@ 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 _, ok := samplingParams["min_tokens"]; ok {
151+
t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"])
152+
}
150153
extraArgs, ok := samplingParams["extra_args"].(map[string]any)
151154
if !ok {
152155
t.Fatal("expected sampling_params.extra_args in generate format")
@@ -218,6 +221,14 @@ func TestPrefillStep_CompletionsFormat(t *testing.T) {
218221
if prefillBody["request_id"] != "req-compl" {
219222
t.Fatalf("expected request_id, got %v", prefillBody["request_id"])
220223
}
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).
226+
if prefillBody["max_tokens"] != float64(1) {
227+
t.Fatalf("expected max_tokens=1, got %v", prefillBody["max_tokens"])
228+
}
229+
if _, ok := prefillBody["min_tokens"]; ok {
230+
t.Fatalf("expected min_tokens to be stripped, got %v", prefillBody["min_tokens"])
231+
}
221232
// Completions format has top-level kv_transfer_params
222233
kvParams, ok := prefillBody["kv_transfer_params"].(map[string]any)
223234
if !ok {

‎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)