Skip to content

Commit d19d722

Browse files
roytmanroytman
authored
Cap max_tokens/max_completion_tokens on synthetic single-token requests (#2086)
* Add shared helper to cap max_tokens/max_completion_tokens to 1 Chat-completions clients may set max_tokens, max_completion_tokens, or neither. Several call sites that synthesize a single-output-token request (multimodal encoder fan-out, EC/KV prefill legs) only touched max_tokens, so a client-supplied max_completion_tokens survived untouched into the outgoing request. In the coordinator's prefill step this meant a cloned request body kept its original max_completion_tokens value alongside the newly capped max_tokens=1. Add request.PrimeSingleTokenRequest to pkg/common/request and use it in the sidecar's EC/mooncake/p2p connectors and the coordinator's encode/prefill steps: always cap max_tokens to 1, and only add max_completion_tokens when the original request already carried it. Sidecar connector_shared_storage.go and connector_nixlv2.go are left untouched; they unconditionally set both fields regardless of what the client sent, a stricter, separately regression-tested requirement (commit bb181d6) not covered by this helper. pkg/sidecar/proxy's requestFieldMaxTokens/MaxCompletionTokens/Stream/ StreamOptions constants now alias the pkg/common/request constants instead of duplicating the string literals. Signed-off-by: roytman <roytman@il.ibm.com> * refactor: dedupe request field-name constants into pkg/common/request pkg/sidecar/proxy declared its own requestField*/requestHeaderRequestID string constants; pkg/coordinator's connectors and steps used the matching literals directly instead of referencing them. Move the definitions into pkg/common/request/constants.go (folding the single-constant headers.go into it) and have both packages reference the shared constants. No behavior change. Signed-off-by: roytman <roytman@il.ibm.com> * add tests Signed-off-by: roytman <roytman@il.ibm.com> --------- Signed-off-by: roytman <roytman@il.ibm.com>
1 parent f3394ed commit d19d722

19 files changed

Lines changed: 418 additions & 93 deletions

‎pkg/common/request/constants.go‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
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 request
18+
19+
const (
20+
RequestIDHeaderKey = "x-request-id"
21+
22+
FieldKVTransferParams = "kv_transfer_params"
23+
FieldECTransferParams = "ec_transfer_params"
24+
FieldMaxOutputTokens = "max_output_tokens" // Used by Responses API
25+
FieldMinTokens = "min_tokens"
26+
FieldSamplingParams = "sampling_params"
27+
FieldDoRemotePrefill = "do_remote_prefill"
28+
FieldDoRemoteDecode = "do_remote_decode"
29+
FieldRemoteBlockIDs = "remote_block_ids"
30+
FieldRemoteEngineID = "remote_engine_id"
31+
FieldRemoteHost = "remote_host"
32+
FieldRemotePort = "remote_port"
33+
FieldCacheHitThreshold = "cache_hit_threshold"
34+
FieldContinueFinalMessage = "continue_final_message"
35+
FieldAddGenerationPrompt = "add_generation_prompt"
36+
)

‎pkg/common/request/headers.go‎

Lines changed: 0 additions & 21 deletions
This file was deleted.

‎pkg/common/request/tokens.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
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 request
18+
19+
const (
20+
FieldMaxTokens = "max_tokens"
21+
FieldMaxCompletionTokens = "max_completion_tokens"
22+
FieldStream = "stream"
23+
FieldStreamOptions = "stream_options"
24+
)
25+
26+
// PrimeSingleTokenRequest mutates target in place into a synthetic,
27+
// non-streaming, single-output-token chat-completions request derived from
28+
// original (which may be the same map as target). max_tokens is always
29+
// capped to 1; max_completion_tokens is only added when original already
30+
// carries it, and leaves it untached otherwise.
31+
func PrimeSingleTokenRequest(target, original map[string]any) {
32+
target[FieldMaxTokens] = 1
33+
if _, ok := original[FieldMaxCompletionTokens]; ok {
34+
target[FieldMaxCompletionTokens] = 1
35+
}
36+
37+
target[FieldStream] = false
38+
delete(target, FieldStreamOptions)
39+
}

‎pkg/common/request/tokens_test.go‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
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 request
18+
19+
import (
20+
"maps"
21+
"testing"
22+
)
23+
24+
func TestPrimeSingleTokenRequest(t *testing.T) {
25+
tests := []struct {
26+
name string
27+
original map[string]any
28+
want map[string]any
29+
}{
30+
{
31+
name: "neither field present",
32+
original: map[string]any{"model": "m"},
33+
want: map[string]any{"model": "m", FieldMaxTokens: 1, FieldStream: false},
34+
},
35+
{
36+
name: "max_tokens only",
37+
original: map[string]any{"model": "m", FieldMaxTokens: 50},
38+
want: map[string]any{"model": "m", FieldMaxTokens: 1, FieldStream: false},
39+
},
40+
{
41+
name: "max_completion_tokens only",
42+
original: map[string]any{"model": "m", FieldMaxCompletionTokens: 100},
43+
want: map[string]any{"model": "m", FieldMaxTokens: 1, FieldMaxCompletionTokens: 1, FieldStream: false},
44+
},
45+
{
46+
name: "both fields present",
47+
original: map[string]any{"model": "m", FieldMaxTokens: 50, FieldMaxCompletionTokens: 100},
48+
want: map[string]any{"model": "m", FieldMaxTokens: 1, FieldMaxCompletionTokens: 1, FieldStream: false},
49+
},
50+
{
51+
name: "stream_options is stripped and stream is forced false",
52+
original: map[string]any{"model": "m", FieldStream: true, FieldStreamOptions: map[string]any{"include_usage": true}},
53+
want: map[string]any{"model": "m", FieldMaxTokens: 1, FieldStream: false},
54+
},
55+
}
56+
57+
for _, tt := range tests {
58+
t.Run(tt.name, func(t *testing.T) {
59+
target := maps.Clone(tt.original)
60+
PrimeSingleTokenRequest(target, tt.original)
61+
62+
if len(target) != len(tt.want) {
63+
t.Fatalf("got %v, want %v", target, tt.want)
64+
}
65+
for k, want := range tt.want {
66+
if got := target[k]; got != want {
67+
t.Errorf("target[%q] = %v, want %v", k, got, want)
68+
}
69+
}
70+
})
71+
}
72+
}
73+
74+
func TestPrimeSingleTokenRequest_TargetIsOriginal(t *testing.T) {
75+
// target and original may be the same map (in-place mutation).
76+
req := map[string]any{"model": "m", FieldMaxCompletionTokens: 100}
77+
78+
PrimeSingleTokenRequest(req, req)
79+
80+
if req[FieldMaxTokens] != 1 {
81+
t.Errorf("max_tokens = %v, want 1", req[FieldMaxTokens])
82+
}
83+
if req[FieldMaxCompletionTokens] != 1 {
84+
t.Errorf("max_completion_tokens = %v, want 1", req[FieldMaxCompletionTokens])
85+
}
86+
if req[FieldStream] != false {
87+
t.Errorf("stream = %v, want false", req[FieldStream])
88+
}
89+
}
90+
91+
func TestPrimeSingleTokenRequest_TargetDistinctFromOriginal(t *testing.T) {
92+
// target starts empty; original decides which max-tokens field is set,
93+
// mirroring buildEncoderRequest-style callers that build a fresh map.
94+
original := map[string]any{"model": "m", FieldMaxTokens: 50}
95+
target := map[string]any{}
96+
97+
PrimeSingleTokenRequest(target, original)
98+
99+
if _, ok := original[FieldMaxCompletionTokens]; ok {
100+
t.Fatalf("original should not be mutated, got %v", original)
101+
}
102+
if target[FieldMaxTokens] != 1 {
103+
t.Errorf("max_tokens = %v, want 1", target[FieldMaxTokens])
104+
}
105+
if _, ok := target[FieldMaxCompletionTokens]; ok {
106+
t.Errorf("max_completion_tokens should be absent, got %v", target[FieldMaxCompletionTokens])
107+
}
108+
}

‎pkg/coordinator/connectors/kv/nixl.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"sigs.k8s.io/controller-runtime/pkg/log"
2323

2424
logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
25+
reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"
2526
"github.com/llm-d/llm-d-router/pkg/coordinator/pipeline"
2627
)
2728

@@ -35,12 +36,12 @@ func (nixlKV) Name() string { return NIXL }
3536

3637
func (nixlKV) PreparePrefillKVParams(ctx context.Context, _ *pipeline.RequestContext) map[string]any {
3738
params := map[string]any{
38-
"do_remote_decode": true,
39-
"do_remote_prefill": false,
40-
"remote_engine_id": nil,
41-
"remote_block_ids": nil,
42-
"remote_host": nil,
43-
"remote_port": nil,
39+
reqcommon.FieldDoRemoteDecode: true,
40+
reqcommon.FieldDoRemotePrefill: false,
41+
reqcommon.FieldRemoteEngineID: nil,
42+
reqcommon.FieldRemoteBlockIDs: nil,
43+
reqcommon.FieldRemoteHost: nil,
44+
reqcommon.FieldRemotePort: nil,
4445
}
4546
log.FromContext(ctx).WithName(loggerName).V(logutil.TRACE).Info("preparing prefill kv params", "params", params)
4647
return params
@@ -51,8 +52,8 @@ func (nixlKV) PrepareDecodeKVParams(ctx context.Context, reqCtx *pipeline.Reques
5152
for k, v := range reqCtx.KVTransferParams {
5253
out[k] = v
5354
}
54-
out["do_remote_decode"] = false
55-
out["do_remote_prefill"] = true
55+
out[reqcommon.FieldDoRemoteDecode] = false
56+
out[reqcommon.FieldDoRemotePrefill] = true
5657
log.FromContext(ctx).WithName(loggerName).V(logutil.TRACE).Info("preparing decode kv params", "params", out)
5758
return out
5859
}

‎pkg/coordinator/connectors/kv/sglang.go‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import (
2727
"sigs.k8s.io/controller-runtime/pkg/log"
2828

2929
logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
30+
reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"
3031
"github.com/llm-d/llm-d-router/pkg/coordinator/pipeline"
3132
)
3233

@@ -91,10 +92,10 @@ func (sglangKV) Name() string { return SGLang }
9192

9293
func (sglangKV) PreparePrefillKVParams(ctx context.Context, _ *pipeline.RequestContext) map[string]any {
9394
params := map[string]any{
94-
"do_remote_decode": true,
95-
"do_remote_prefill": false,
96-
fieldBootstrapPort: resolveSGLangBootstrapPort(ctx),
97-
fieldBootstrapRoom: uuid.NewString(),
95+
reqcommon.FieldDoRemoteDecode: true,
96+
reqcommon.FieldDoRemotePrefill: false,
97+
fieldBootstrapPort: resolveSGLangBootstrapPort(ctx),
98+
fieldBootstrapRoom: uuid.NewString(),
9899
}
99100
log.FromContext(ctx).WithName(loggerName).V(logutil.TRACE).Info("preparing prefill kv params", "params", params)
100101
return params
@@ -105,8 +106,8 @@ func (sglangKV) PrepareDecodeKVParams(ctx context.Context, reqCtx *pipeline.Requ
105106
for k, v := range reqCtx.KVTransferParams {
106107
out[k] = v
107108
}
108-
out["do_remote_decode"] = false
109-
out["do_remote_prefill"] = true
109+
out[reqcommon.FieldDoRemoteDecode] = false
110+
out[reqcommon.FieldDoRemotePrefill] = true
110111
log.FromContext(ctx).WithName(loggerName).V(logutil.TRACE).Info("preparing decode kv params", "params", out)
111112
return out
112113
}

‎pkg/coordinator/connectors/kv/shared_storage.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"sigs.k8s.io/controller-runtime/pkg/log"
2323

2424
logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
25+
reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"
2526
"github.com/llm-d/llm-d-router/pkg/coordinator/pipeline"
2627
)
2728

@@ -33,13 +34,13 @@ type sharedStorageKV struct{}
3334
func (sharedStorageKV) Name() string { return SharedStorage }
3435

3536
func (sharedStorageKV) PreparePrefillKVParams(ctx context.Context, _ *pipeline.RequestContext) map[string]any {
36-
params := map[string]any{"do_remote_decode": true, "do_remote_prefill": false}
37+
params := map[string]any{reqcommon.FieldDoRemoteDecode: true, reqcommon.FieldDoRemotePrefill: false}
3738
log.FromContext(ctx).WithName(loggerName).V(logutil.TRACE).Info("preparing prefill kv params", "params", params)
3839
return params
3940
}
4041

4142
func (sharedStorageKV) PrepareDecodeKVParams(ctx context.Context, _ *pipeline.RequestContext) map[string]any {
42-
params := map[string]any{"do_remote_decode": false, "do_remote_prefill": true}
43+
params := map[string]any{reqcommon.FieldDoRemoteDecode: false, reqcommon.FieldDoRemotePrefill: true}
4344
log.FromContext(ctx).WithName(loggerName).V(logutil.TRACE).Info("preparing decode kv params", "params", params)
4445
return params
4546
}

‎pkg/coordinator/server/handlers.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ func (s *Server) handleInference(w http.ResponseWriter, r *http.Request) {
6666
}
6767
parseDuration := time.Since(parseStart)
6868

69-
stream, _ := parsed["stream"].(bool)
69+
stream, _ := parsed[reqcommon.FieldStream].(bool)
7070
model, _ := parsed["model"].(string)
7171

7272
requestID := r.Header.Get(reqcommon.RequestIDHeaderKey)

‎pkg/coordinator/steps/decode.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"sigs.k8s.io/controller-runtime/pkg/log"
2525

2626
logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging"
27+
reqcommon "github.com/llm-d/llm-d-router/pkg/common/request"
2728

2829
"github.com/llm-d/llm-d-router/pkg/coordinator/connectors/kv"
2930
"github.com/llm-d/llm-d-router/pkg/coordinator/gateway"
@@ -87,7 +88,7 @@ func (s *DecodeStep) Execute(ctx context.Context, reqCtx *pipeline.RequestContex
8788
// maps.Clone would still share. This is sound only while the pipeline runs steps
8889
// sequentially; if it ever goes concurrent, decode must copy like the others.
8990
func (s *DecodeStep) prepareDecodeBody(ctx context.Context, reqCtx *pipeline.RequestContext) {
90-
reqCtx.Body["kv_transfer_params"] = s.kv.PrepareDecodeKVParams(ctx, reqCtx)
91+
reqCtx.Body[reqcommon.FieldKVTransferParams] = s.kv.PrepareDecodeKVParams(ctx, reqCtx)
9192
s.injectUUIDs(reqCtx)
9293

9394
format := resolveFormat(s.useOpenAIFormat, reqCtx.OriginalPath)

‎pkg/coordinator/steps/encode.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,8 +201,8 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [
201201
"mm_placeholders": map[string][]any{ModalityImage: {map[string]any{"offset": 1, "length": entry.Placeholder.Length}}},
202202
},
203203
},
204-
"max_tokens": 1,
205204
}
205+
reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body)
206206
return body
207207
default:
208208
return map[string]any{
@@ -213,7 +213,7 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [
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-
"sampling_params": map[string]any{"max_tokens": 1},
216+
reqcommon.FieldSamplingParams: map[string]any{reqcommon.FieldMaxTokens: 1},
217217
}
218218
}
219219
}

0 commit comments

Comments
 (0)