From 7c13fe5d182ca191764874cd1565a2e8e4b8db6d Mon Sep 17 00:00:00 2001 From: Revital Sur Date: Mon, 27 Jul 2026 14:04:47 +0300 Subject: [PATCH 1/5] Cap coordinator prefill/encode legs to a single token per schema Signed-off-by: Revital Sur --- pkg/coordinator/steps/encode.go | 7 +- pkg/coordinator/steps/encode_test.go | 56 ++++++++++-- pkg/coordinator/steps/prefill.go | 6 +- pkg/coordinator/steps/prefill_test.go | 11 +++ pkg/coordinator/steps/tokens.go | 55 ++++++++++++ pkg/coordinator/steps/tokens_test.go | 122 ++++++++++++++++++++++++++ 6 files changed, 245 insertions(+), 12 deletions(-) create mode 100644 pkg/coordinator/steps/tokens.go create mode 100644 pkg/coordinator/steps/tokens_test.go diff --git a/pkg/coordinator/steps/encode.go b/pkg/coordinator/steps/encode.go index 66c08b6395..5fba3bda1c 100644 --- a/pkg/coordinator/steps/encode.go +++ b/pkg/coordinator/steps/encode.go @@ -202,10 +202,10 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [ }, }, } - reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body) + capSingleTokenOutput(body, format) return body default: - return map[string]any{ + body := map[string]any{ "model": reqCtx.Model, "token_ids": tokenIDs, "features": map[string]any{ @@ -213,8 +213,9 @@ func (s *EncodeStep) buildEncodeBody(reqCtx *pipeline.RequestContext, tokenIDs [ "mm_placeholders": map[string][]any{ModalityImage: {map[string]any{"offset": 1, "length": entry.Placeholder.Length}}}, "kwargs_data": map[string][]string{ModalityImage: {entry.KwargsData}}, }, - reqcommon.FieldSamplingParams: map[string]any{reqcommon.FieldMaxTokens: 1}, } + capSingleTokenOutput(body, format) + return body } } diff --git a/pkg/coordinator/steps/encode_test.go b/pkg/coordinator/steps/encode_test.go index 9fdddc2411..31211fdc78 100644 --- a/pkg/coordinator/steps/encode_test.go +++ b/pkg/coordinator/steps/encode_test.go @@ -352,10 +352,10 @@ func TestEncodeStep_ChatCompletionsFormat(t *testing.T) { } // TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens is a -// regression test: buildEncodeBody used to always emit max_tokens=1 without -// regard to which token-limit field the client used, so a reasoning-model -// client's max_completion_tokens was never capped in the encoder sub-request. -func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) { +// The encode chat sub-request is built fresh from the request context and does +// not carry the client's sampling fields, so max_completion_tokens is not +// propagated and is never injected: max_tokens=1 alone caps output. +func TestEncodeStep_ChatCompletionsFormat_OmitsMaxCompletionTokens(t *testing.T) { var receivedBody map[string]any server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -401,8 +401,11 @@ func TestEncodeStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) t.Fatalf("unexpected error: %v", err) } - if receivedBody["max_completion_tokens"] != float64(1) { - t.Fatalf("expected encode sub-request max_completion_tokens capped to 1, got %v", receivedBody["max_completion_tokens"]) + if receivedBody["max_tokens"] != float64(1) { + t.Fatalf("expected encode sub-request max_tokens capped to 1, got %v", receivedBody["max_tokens"]) + } + if _, ok := receivedBody["max_completion_tokens"]; ok { + t.Fatalf("expected encode sub-request to omit max_completion_tokens, got %v", receivedBody["max_completion_tokens"]) } } @@ -536,3 +539,44 @@ func TestEncodeStep_BuildsCorrectTokenIDs(t *testing.T) { } } } + +// TestEncodeStep_GenerateFormat_CapsSingleToken verifies the generate-format +// encoder sub-request caps output to a single token: sampling_params carries +// max_tokens=1 and strips min_tokens (it defaults to 0, keeping min_tokens <= +// max_tokens). +func TestEncodeStep_GenerateFormat_CapsSingleToken(t *testing.T) { + var samplingParams map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var parsed map[string]any + _ = json.Unmarshal(body, &parsed) + samplingParams, _ = parsed["sampling_params"].(map[string]any) + _ = json.NewEncoder(w).Encode(map[string]any{ + "ec_transfer_params": map[string]any{"h1": map[string]any{"peer_port": 5501}}, + }) + })) + defer server.Close() + + gwClient := gateway.New(config.GatewayConfig{Address: server.URL}) + step, _ := NewEncodeStep(gwClient, map[string]any{"use_openai_format": false}) + + reqCtx := &pipeline.RequestContext{ + RequestID: "req-gen-cap", + Model: "test", + TokenIDs: []int{1, 32000, 32000, 32000, 2345}, + MultimodalEntries: []pipeline.MultimodalEntry{ + {Index: 0, Hash: "h1", KwargsData: "dGVzdA==", Placeholder: pipeline.PlaceholderRange{Offset: 1, Length: 3}}, + }, + } + + if err := step.Execute(context.Background(), reqCtx); err != nil { + t.Fatalf("unexpected error: %v", err) + } + if samplingParams["max_tokens"] != float64(1) { + t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"]) + } + if _, ok := samplingParams["min_tokens"]; ok { + t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"]) + } +} diff --git a/pkg/coordinator/steps/prefill.go b/pkg/coordinator/steps/prefill.go index 9ede0e2e44..bcb54d59a9 100644 --- a/pkg/coordinator/steps/prefill.go +++ b/pkg/coordinator/steps/prefill.go @@ -137,7 +137,7 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req switch format { case gateway.FormatChatCompletions: body := maps.Clone(reqCtx.Body) - reqcommon.PrimeSingleTokenRequest(body, reqCtx.Body) + capSingleTokenOutput(body, format) tokens := map[string]any{ "token_ids": reqCtx.TokenIDs, } @@ -164,9 +164,9 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req "request_id": reqCtx.RequestID, "model": reqCtx.Model, "prompt": prompt, - reqcommon.FieldMaxTokens: 1, reqcommon.FieldKVTransferParams: kvParams, } + capSingleTokenOutput(body, format) if features != nil { body["features"] = features } @@ -181,12 +181,12 @@ func (s *PrefillStep) buildPrefillBody(ctx context.Context, reqCtx *pipeline.Req "token_ids": reqCtx.TokenIDs, "model": reqCtx.Model, reqcommon.FieldSamplingParams: map[string]any{ - reqcommon.FieldMaxTokens: 1, "extra_args": map[string]any{ reqcommon.FieldKVTransferParams: kvParams, }, }, } + capSingleTokenOutput(body, format) if features != nil { body["features"] = features } diff --git a/pkg/coordinator/steps/prefill_test.go b/pkg/coordinator/steps/prefill_test.go index 281596859f..3a67c19a13 100644 --- a/pkg/coordinator/steps/prefill_test.go +++ b/pkg/coordinator/steps/prefill_test.go @@ -147,6 +147,9 @@ func TestPrefillStep_SendsCorrectGenerateRequest(t *testing.T) { if samplingParams["max_tokens"] != float64(1) { t.Fatalf("expected sampling_params.max_tokens=1, got %v", samplingParams["max_tokens"]) } + if _, ok := samplingParams["min_tokens"]; ok { + t.Fatalf("expected sampling_params.min_tokens to be stripped, got %v", samplingParams["min_tokens"]) + } extraArgs, ok := samplingParams["extra_args"].(map[string]any) if !ok { t.Fatal("expected sampling_params.extra_args in generate format") @@ -218,6 +221,14 @@ func TestPrefillStep_CompletionsFormat(t *testing.T) { if prefillBody["request_id"] != "req-compl" { t.Fatalf("expected request_id, got %v", prefillBody["request_id"]) } + // Prefill leg caps output to a single token: max_tokens is pinned to 1 and + // min_tokens is stripped (it defaults to 0, keeping min_tokens <= max_tokens). + if prefillBody["max_tokens"] != float64(1) { + t.Fatalf("expected max_tokens=1, got %v", prefillBody["max_tokens"]) + } + if _, ok := prefillBody["min_tokens"]; ok { + t.Fatalf("expected min_tokens to be stripped, got %v", prefillBody["min_tokens"]) + } // Completions format has top-level kv_transfer_params kvParams, ok := prefillBody["kv_transfer_params"].(map[string]any) if !ok { diff --git a/pkg/coordinator/steps/tokens.go b/pkg/coordinator/steps/tokens.go new file mode 100644 index 0000000000..34de78fa9e --- /dev/null +++ b/pkg/coordinator/steps/tokens.go @@ -0,0 +1,55 @@ +/* +Copyright 2026 The llm-d Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package steps + +import ( + reqcommon "github.com/llm-d/llm-d-router/pkg/common/request" + "github.com/llm-d/llm-d-router/pkg/coordinator/gateway" +) + +// capSingleTokenOutput rewrites body into a single-output-token, non-streaming +// request for the synthetic prefill and encode legs. max_tokens is pinned to 1 +// and min_tokens is stripped: min_tokens defaults to 0 in vLLM, so removing it +// keeps min_tokens <= max_tokens satisfied without raising the floor above +// max_tokens=1. These limits live in the sampling_params sub-map for the generate +// schema (synthesized if absent) and at the top level for the OpenAI schemas. +// +// stream is forced false and stream_options stripped for every format, so no leg +// returns a streamed response the coordinator cannot decode. max_completion_tokens +// is capped to 1 only when the body already carries it (in practice only the chat +// completions schema does), never added when absent. +func capSingleTokenOutput(body map[string]any, format gateway.RequestFormat) { + target := body + if format == gateway.FormatGenerate { + sp, ok := body[reqcommon.FieldSamplingParams].(map[string]any) + if !ok { + sp = map[string]any{} + body[reqcommon.FieldSamplingParams] = sp + } + target = sp + } + + target[reqcommon.FieldMaxTokens] = 1 + delete(target, reqcommon.FieldMinTokens) + + if _, ok := body[reqcommon.FieldMaxCompletionTokens]; ok { + body[reqcommon.FieldMaxCompletionTokens] = 1 + } + + body[reqcommon.FieldStream] = false + delete(body, reqcommon.FieldStreamOptions) +} diff --git a/pkg/coordinator/steps/tokens_test.go b/pkg/coordinator/steps/tokens_test.go new file mode 100644 index 0000000000..a09d1335a5 --- /dev/null +++ b/pkg/coordinator/steps/tokens_test.go @@ -0,0 +1,122 @@ +/* +Copyright 2026 The llm-d Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package steps + +import ( + "reflect" + "testing" + + "github.com/llm-d/llm-d-router/pkg/coordinator/gateway" +) + +func TestCapSingleTokenOutput(t *testing.T) { + tests := []struct { + name string + format gateway.RequestFormat + body map[string]any + want map[string]any + }{ + { + name: "chat completions caps output fields and forces non-streaming", + format: gateway.FormatChatCompletions, + body: map[string]any{ + "model": "m", + "max_tokens": 100, + "min_tokens": 5, + "max_completion_tokens": 100, + "stream": true, + "stream_options": map[string]any{"include_usage": true}, + }, + want: map[string]any{ + "model": "m", + "max_tokens": 1, + "max_completion_tokens": 1, + "stream": false, + }, + }, + { + name: "max_completion_tokens is not added when the client omitted it", + format: gateway.FormatChatCompletions, + body: map[string]any{"model": "m"}, + want: map[string]any{ + "model": "m", + "max_tokens": 1, + "stream": false, + }, + }, + { + name: "completions caps max_tokens, strips min_tokens, forces non-streaming", + format: gateway.FormatCompletions, + body: map[string]any{"model": "m", "max_tokens": 100, "min_tokens": 5}, + want: map[string]any{"model": "m", "max_tokens": 1, "stream": false}, + }, + { + name: "streaming is forced false and stream_options stripped", + format: gateway.FormatCompletions, + body: map[string]any{"stream": true, "stream_options": map[string]any{"include_usage": true}}, + want: map[string]any{"stream": false, "max_tokens": 1}, + }, + { + name: "generate caps max_tokens and strips min_tokens inside sampling_params", + format: gateway.FormatGenerate, + body: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 100, "min_tokens": 5}, + }, + want: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 1}, + "stream": false, + }, + }, + { + name: "generate synthesizes sampling_params when absent", + format: gateway.FormatGenerate, + body: map[string]any{"model": "m"}, + want: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 1}, + "stream": false, + }, + }, + { + name: "generate preserves other sampling_params entries", + format: gateway.FormatGenerate, + body: map[string]any{ + "sampling_params": map[string]any{ + "extra_args": map[string]any{"kv_transfer_params": "x"}, + }, + }, + want: map[string]any{ + "sampling_params": map[string]any{ + "max_tokens": 1, + "extra_args": map[string]any{"kv_transfer_params": "x"}, + }, + "stream": false, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + capSingleTokenOutput(tt.body, tt.format) + if !reflect.DeepEqual(tt.body, tt.want) { + t.Fatalf("got %v, want %v", tt.body, tt.want) + } + }) + } +} From 27fc9e20d2f7c32c0757876b502697be16286de1 Mon Sep 17 00:00:00 2001 From: Revital Sur Date: Mon, 27 Jul 2026 22:53:40 +0300 Subject: [PATCH 2/5] Add e2e test. Signed-off-by: Revital Sur --- pkg/coordinator/gateway/client.go | 11 ++- pkg/coordinator/gateway/client_test.go | 12 +-- pkg/coordinator/steps/decode_proxy.go | 8 ++ pkg/coordinator/steps/prefill_test.go | 50 ++++++++++ pkg/coordinator/steps/tokens.go | 19 ++-- .../e2e/coordinator/coordinator_test.go | 92 ++++++++++++++----- 6 files changed, 150 insertions(+), 42 deletions(-) diff --git a/pkg/coordinator/gateway/client.go b/pkg/coordinator/gateway/client.go index a07b28dca5..d4b7835854 100644 --- a/pkg/coordinator/gateway/client.go +++ b/pkg/coordinator/gateway/client.go @@ -92,7 +92,7 @@ func (c *Client) Request(ctx context.Context, method, path string, body []byte, logger := log.FromContext(ctx).WithName("gateway") if body != nil { if v := logger.V(logutil.TRACE); v.Enabled() { - v.Info("request body", "method", method, "path", path, "headers", httplog.RedactedHeaders(req.Header), "body", redactBody(body)) + v.Info("request body", "method", method, "path", path, "headers", httplog.RedactedHeaders(req.Header), "body", RedactBody(body)) } } @@ -109,7 +109,7 @@ func (c *Client) Request(ctx context.Context, method, path string, body []byte, if err != nil { return nil, fmt.Errorf("reading response from gateway: %w", err) } - v.Info("response body", "status", resp.StatusCode, "body", redactBody(respBody)) + v.Info("response body", "status", resp.StatusCode, "body", RedactBody(respBody)) resp.Body = io.NopCloser(bytes.NewReader(respBody)) } @@ -146,10 +146,11 @@ const ( maxRedactDepth = 32 ) -// redactBody parses JSON and redacts oversized string values (see +// RedactBody parses JSON and redacts oversized string values (see // maxRedactStringLen); a body that is not valid JSON is returned verbatim, or -// truncated at maxRedactRawBodyLen. -func redactBody(data []byte) any { +// truncated at maxRedactRawBodyLen. It is exported for the decode reverse-proxy +// path, which bypasses Client.Request and so redacts its own body before logging. +func RedactBody(data []byte) any { var v any if err := json.Unmarshal(data, &v); err != nil { if len(data) > maxRedactRawBodyLen { diff --git a/pkg/coordinator/gateway/client_test.go b/pkg/coordinator/gateway/client_test.go index cf4a418f38..1bf557d929 100644 --- a/pkg/coordinator/gateway/client_test.go +++ b/pkg/coordinator/gateway/client_test.go @@ -147,9 +147,9 @@ func TestRedactBody(t *testing.T) { t.Fatal(err) } - got, ok := redactBody(body).(map[string]any) + got, ok := RedactBody(body).(map[string]any) if !ok { - t.Fatalf("expected map[string]any, got %T", redactBody(body)) + t.Fatalf("expected map[string]any, got %T", RedactBody(body)) } if got["model"] != "m" { t.Errorf("short value changed: %v", got["model"]) @@ -160,23 +160,23 @@ func TestRedactBody(t *testing.T) { }) t.Run("non-JSON under max length returned verbatim", func(t *testing.T) { - if got := redactBody([]byte("not json")); got != "not json" { + if got := RedactBody([]byte("not json")); got != "not json" { t.Errorf("redactBody = %v, want verbatim string", got) } }) t.Run("non-JSON exactly max length returned verbatim", func(t *testing.T) { raw := strings.Repeat("q", maxRedactRawBodyLen) - if got := redactBody([]byte(raw)); got != raw { + if got := RedactBody([]byte(raw)); got != raw { t.Errorf("redactBody = %v, want verbatim (boundary is > not >=)", got) } }) t.Run("non-JSON over max length truncated", func(t *testing.T) { raw := strings.Repeat("q", maxRedactRawBodyLen+50) - got, ok := redactBody([]byte(raw)).(string) + got, ok := RedactBody([]byte(raw)).(string) if !ok { - t.Fatalf("expected string, got %T", redactBody([]byte(raw))) + t.Fatalf("expected string, got %T", RedactBody([]byte(raw))) } if want := raw[:maxRedactRawBodyLen] + "..."; got != want { t.Errorf("redactBody truncation = %q, want %q", got, want) diff --git a/pkg/coordinator/steps/decode_proxy.go b/pkg/coordinator/steps/decode_proxy.go index e6783c5b4e..e7c42866f1 100644 --- a/pkg/coordinator/steps/decode_proxy.go +++ b/pkg/coordinator/steps/decode_proxy.go @@ -78,6 +78,14 @@ func newDecodeProxyRequest(ctx context.Context, logger logr.Logger, step string, v.Info("request body", "method", "POST", "path", reqCtx.OriginalPath, "bodyLen", len(bodyBytes), "headers", httplog.RedactedHeaders(proxyReq.Header)) } + // The decode leg is proxied straight to the gateway and never passes through + // 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() { + v.Info("request body", "method", "POST", "path", reqCtx.OriginalPath, "headers", httplog.RedactedHeaders(proxyReq.Header), "body", gateway.RedactBody(bodyBytes)) + } + return proxyReq, nil } diff --git a/pkg/coordinator/steps/prefill_test.go b/pkg/coordinator/steps/prefill_test.go index 3a67c19a13..f6fa4a03d5 100644 --- a/pkg/coordinator/steps/prefill_test.go +++ b/pkg/coordinator/steps/prefill_test.go @@ -477,6 +477,56 @@ func TestPrefillStep_ChatCompletionsFormat_CapsMaxCompletionTokens(t *testing.T) } } +// TestPrefillStep_ChatCompletionsFormat_StripsClientMinTokens is a regression +// test for the one coordinator path where a client-supplied min_tokens survives +// into the capped body: chat-completions clones reqCtx.Body, so a client +// min_tokens > 1 would leave min_tokens > max_tokens=1 and vLLM rejects the leg. +// The generate and completions legs build fresh bodies that never carry it. +func TestPrefillStep_ChatCompletionsFormat_StripsClientMinTokens(t *testing.T) { + var prefillBody map[string]any + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + _ = json.Unmarshal(body, &prefillBody) + _ = json.NewEncoder(w).Encode(map[string]any{ + "kv_transfer_params": map[string]any{"block_id": "block-5"}, + }) + })) + defer server.Close() + + gwClient := gateway.New(config.GatewayConfig{Address: server.URL}) + step, err := NewPrefillStep(gwClient, map[string]any{}) + if err != nil { + t.Fatal(err) + } + + reqCtx := &pipeline.RequestContext{ + RequestID: "req-chat-min-tokens", + OriginalPath: gateway.PathChatCompletions, + Model: "test-model", + Body: map[string]any{ + "model": "test-model", + "min_tokens": 50, + "max_tokens": 200, + "messages": []any{ + map[string]any{"role": "user", "content": "hello"}, + }, + }, + KVTransferParams: make(map[string]any), + } + + if err := step.Execute(context.Background(), reqCtx); err != nil { + t.Fatalf("unexpected error: %v", err) + } + + if prefillBody["max_tokens"] != float64(1) { + t.Fatalf("expected prefill request max_tokens=1, got %v", prefillBody["max_tokens"]) + } + if _, ok := prefillBody["min_tokens"]; ok { + t.Fatalf("expected client min_tokens to be stripped, got %v", prefillBody["min_tokens"]) + } +} + // TestSharedStorage_OmitsECTransferParams_InPrefillBody verifies that the // ec-shared-storage EC connector never emits an ec_transfer_params field on // the prefill body, in every prefill wire format (chat-completions, diff --git a/pkg/coordinator/steps/tokens.go b/pkg/coordinator/steps/tokens.go index 34de78fa9e..da74c0344e 100644 --- a/pkg/coordinator/steps/tokens.go +++ b/pkg/coordinator/steps/tokens.go @@ -22,16 +22,11 @@ import ( ) // capSingleTokenOutput rewrites body into a single-output-token, non-streaming -// request for the synthetic prefill and encode legs. max_tokens is pinned to 1 -// and min_tokens is stripped: min_tokens defaults to 0 in vLLM, so removing it -// keeps min_tokens <= max_tokens satisfied without raising the floor above -// max_tokens=1. These limits live in the sampling_params sub-map for the generate -// schema (synthesized if absent) and at the top level for the OpenAI schemas. +// request for the synthetic prefill and encode legs. // -// stream is forced false and stream_options stripped for every format, so no leg -// returns a streamed response the coordinator cannot decode. max_completion_tokens -// is capped to 1 only when the body already carries it (in practice only the chat -// completions schema does), never added when absent. +// 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) { target := body if format == gateway.FormatGenerate { @@ -44,12 +39,18 @@ func capSingleTokenOutput(body map[string]any, format gateway.RequestFormat) { } target[reqcommon.FieldMaxTokens] = 1 + // Strip rather than clamp min_tokens: it defaults to 0 in vLLM, so removing it + // keeps min_tokens <= max_tokens=1 without raising the floor above the cap. delete(target, reqcommon.FieldMinTokens) if _, ok := body[reqcommon.FieldMaxCompletionTokens]; ok { body[reqcommon.FieldMaxCompletionTokens] = 1 } + // TODO: max_output_tokens is another client-supplied output cap (Responses + // API) that a client can send instead of max_tokens/max_completion_tokens; it + // should be capped to 1 here as well so the synthetic legs stay single-token. + body[reqcommon.FieldStream] = false delete(body, reqcommon.FieldStreamOptions) } diff --git a/test/coordinator/e2e/coordinator/coordinator_test.go b/test/coordinator/e2e/coordinator/coordinator_test.go index 495fcfe28d..ba1317d060 100644 --- a/test/coordinator/e2e/coordinator/coordinator_test.go +++ b/test/coordinator/e2e/coordinator/coordinator_test.go @@ -50,35 +50,42 @@ var _ = ginkgo.Describe("Coordinator pipeline", func() { runCoordinatorPipeline([]byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":"hello"}]}`, modelName, - )), textOnlySteps, 0) + )), textOnlySteps, 0, 0, 0) + }) + + ginkgo.It("forwards the client min_tokens/max_tokens to decode and caps them on prefill and encode", func() { + runCoordinatorPipeline([]byte(fmt.Sprintf( + `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"min_tokens":3,"max_tokens":5}`, + modelName, inlineImageDataURI, + )), allSteps, 1, 3, 5) }) ginkgo.It("routes a multimodal image chat completion end-to-end", func() { runCoordinatorPipeline([]byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"max_tokens":150}`, modelName, testImageURL, - )), allSteps, 1) + )), allSteps, 1, 0, 0) }) ginkgo.It("routes a multimodal chat completion with two images end-to-end", func() { runCoordinatorPipeline([]byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"image_url","image_url":{"url":%q},"uuid":"image-1"},{"type":"text","text":"What is in these two images?"}]}],"max_tokens":150}`, modelName, testImageURL, testImageURL2, - )), allSteps, 2) + )), allSteps, 2, 0, 0) }) ginkgo.It("routes a multimodal chat completion with an inline base64 image end-to-end", func() { runCoordinatorPipeline([]byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"max_tokens":150}`, modelName, inlineImageDataURI, - )), allSteps, 1) + )), allSteps, 1, 0, 0) }) ginkgo.It("routes a multimodal chat completion with one inline and one remote image end-to-end", func() { runCoordinatorPipeline([]byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"image_url","image_url":{"url":%q},"uuid":"image-1"},{"type":"text","text":"Describe what you see in both images."}]}],"max_tokens":150}`, modelName, inlineImageDataURI, testImageURL, - )), allSteps, 2) + )), allSteps, 2, 0, 0) }) }) @@ -92,8 +99,11 @@ const inlineImageDataURI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAEAAAA // given chat-completion body, asserts a 200 with a non-empty body, verifies // that the coordinator logs show all expected pipeline steps completed, then // tears the workload down. expectedImages is the number of images in the -// request; when > 0 the encoder log assertions are also verified. -func runCoordinatorPipeline(body []byte, expectedSteps []string, expectedImages int) { +// request; when > 0 the encoder log assertions are also verified. When +// wantMaxTokens > 0 the token-limit contract is additionally asserted: decode +// forwards the client's min_tokens/max_tokens unchanged, while prefill (and +// encode, for multimodal requests) cap max_tokens to 1 and strip min_tokens. +func runCoordinatorPipeline(body []byte, expectedSteps []string, expectedImages, wantMinTokens, wantMaxTokens int) { nsName := getNamespace() var ( coordinator []string @@ -158,30 +168,30 @@ func runCoordinatorPipeline(body []byte, expectedSteps []string, expectedImages gomega.Expect(resp.StatusCode).To(gomega.Equal(http.StatusOK), "coordinator returned non-200: body=%s", string(raw)) gomega.Expect(raw).NotTo(gomega.BeEmpty(), "coordinator returned empty body") - verifyCoordinatorSteps(nsName, expectedSteps, expectedImages, true, true) + + logs := fetchCoordinatorLogs(nsName) + verifyCoordinatorSteps(logs, expectedSteps, expectedImages, true, true) + if wantMaxTokens > 0 { + // Prefill is always capped; encode is capped only when the request carries + // images, since it fires one sub-request per image. + capLegs := []string{"prefill"} + if expectedImages > 0 { + capLegs = append(capLegs, "encode") + } + verifyTokenLimits(logs, wantMinTokens, wantMaxTokens, capLegs) + } } -// verifyCoordinatorSteps fetches the coordinator pod logs and asserts that -// every expected step has a "step complete" log entry. When kvNIXL is set it +// verifyCoordinatorSteps asserts that the coordinator logs show every expected +// step has a "step complete" log entry. When kvNIXL is set it // also asserts kv_transfer_params on the prefill and decode legs, and when // expectedImages > 0 it asserts the encode step completed all image // sub-requests, plus, when ecNIXL is set, the ec_transfer_params total via the // "merged encode response" marker. The kv/ec params surface in the logs only // for the NIXL connectors, so those checks are gated accordingly. -func verifyCoordinatorSteps(nsName string, expectedSteps []string, expectedImages int, kvNIXL, ecNIXL bool) { +func verifyCoordinatorSteps(logs string, expectedSteps []string, expectedImages int, kvNIXL, ecNIXL bool) { ginkgo.By("Verifying coordinator logs contain all pipeline steps") - args := []string{"logs", "deployment/llm-d-coordinator", - "-c", "coordinator", "--namespace=" + nsName} - if k8sContext != "" { - args = append(args, "--context="+k8sContext) - } - - out, err := exec.Command("kubectl", args...).CombinedOutput() - gomega.Expect(err).ShouldNot(gomega.HaveOccurred(), - "failed to fetch coordinator logs: %s", string(out)) - - logs := string(out) for _, step := range expectedSteps { stepField := `"step":"` + step + `"` gomega.Expect(logHasLine(logs, `"msg":"step complete"`, stepField)).To(gomega.BeTrue(), @@ -233,6 +243,44 @@ func verifyCoordinatorSteps(nsName string, expectedSteps []string, expectedImage } } +// fetchCoordinatorLogs returns the coordinator container logs for the current spec. +func fetchCoordinatorLogs(nsName string) string { + args := []string{"logs", "deployment/llm-d-coordinator", + "-c", "coordinator", "--namespace=" + nsName} + if k8sContext != "" { + args = append(args, "--context="+k8sContext) + } + + out, err := exec.Command("kubectl", args...).CombinedOutput() + gomega.Expect(err).ShouldNot(gomega.HaveOccurred(), + "failed to fetch coordinator logs: %s", string(out)) + return string(out) +} + +// verifyTokenLimits asserts the pipeline's token-limit contract: the decode leg +// forwards the client's min_tokens/max_tokens unchanged, while the synthetic +// prefill and encode legs (capLegs) cap output to a single token (max_tokens=1) +// and strip min_tokens. Leg request bodies surface only at TRACE, so this relies +// on the coordinator running at log_level 5. +func verifyTokenLimits(logs string, wantMin, wantMax int, capLegs []string) { + ginkgo.By("Verifying decode leg forwards the client min_tokens/max_tokens") + gomega.Expect(logHasLine(logs, `"msg":"request body"`, `"epp-phase":"decode"`, + fmt.Sprintf(`"min_tokens":%d`, wantMin), fmt.Sprintf(`"max_tokens":%d`, wantMax))).To(gomega.BeTrue(), + "coordinator logs have no decode request body carrying min_tokens=%d and max_tokens=%d", wantMin, wantMax) + + for _, phase := range capLegs { + phaseField := `"epp-phase":"` + phase + `"` + + ginkgo.By("Verifying " + phase + " leg caps max_tokens to 1") + gomega.Expect(logHasLine(logs, `"msg":"request body"`, phaseField, `"max_tokens":1`)).To(gomega.BeTrue(), + "coordinator logs have no %s request body carrying max_tokens=1", phase) + + ginkgo.By("Verifying " + phase + " leg strips min_tokens") + gomega.Expect(logHasLine(logs, `"msg":"request body"`, phaseField, `"min_tokens"`)).To(gomega.BeFalse(), + "%s request body must not carry min_tokens", phase) + } +} + // logHasLine reports whether any single line in logs contains all of substrs. func logHasLine(logs string, substrs ...string) bool { for _, line := range strings.Split(logs, "\n") { From 6933d43a1848e69509619dfc512623c3753b3bd3 Mon Sep 17 00:00:00 2001 From: Revital Sur Date: Tue, 28 Jul 2026 13:52:11 +0300 Subject: [PATCH 3/5] Address review comments. Signed-off-by: Revital Sur --- pkg/coordinator/steps/tokens.go | 56 ------------ pkg/coordinator/steps/tokens_test.go | 122 --------------------------- pkg/coordinator/steps/utils.go | 35 ++++++++ pkg/coordinator/steps/utils_test.go | 101 ++++++++++++++++++++++ 4 files changed, 136 insertions(+), 178 deletions(-) delete mode 100644 pkg/coordinator/steps/tokens.go delete mode 100644 pkg/coordinator/steps/tokens_test.go diff --git a/pkg/coordinator/steps/tokens.go b/pkg/coordinator/steps/tokens.go deleted file mode 100644 index da74c0344e..0000000000 --- a/pkg/coordinator/steps/tokens.go +++ /dev/null @@ -1,56 +0,0 @@ -/* -Copyright 2026 The llm-d Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -package steps - -import ( - reqcommon "github.com/llm-d/llm-d-router/pkg/common/request" - "github.com/llm-d/llm-d-router/pkg/coordinator/gateway" -) - -// 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 -// 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) { - target := body - if format == gateway.FormatGenerate { - sp, ok := body[reqcommon.FieldSamplingParams].(map[string]any) - if !ok { - sp = map[string]any{} - body[reqcommon.FieldSamplingParams] = sp - } - target = sp - } - - target[reqcommon.FieldMaxTokens] = 1 - // Strip rather than clamp min_tokens: it defaults to 0 in vLLM, so removing it - // keeps min_tokens <= max_tokens=1 without raising the floor above the cap. - delete(target, reqcommon.FieldMinTokens) - - if _, ok := body[reqcommon.FieldMaxCompletionTokens]; ok { - body[reqcommon.FieldMaxCompletionTokens] = 1 - } - - // TODO: max_output_tokens is another client-supplied output cap (Responses - // API) that a client can send instead of max_tokens/max_completion_tokens; it - // should be capped to 1 here as well so the synthetic legs stay single-token. - - body[reqcommon.FieldStream] = false - delete(body, reqcommon.FieldStreamOptions) -} diff --git a/pkg/coordinator/steps/tokens_test.go b/pkg/coordinator/steps/tokens_test.go deleted file mode 100644 index a09d1335a5..0000000000 --- a/pkg/coordinator/steps/tokens_test.go +++ /dev/null @@ -1,122 +0,0 @@ -/* -Copyright 2026 The llm-d Authors. - -Licensed under the Apache License, Version 2.0 (the "License"); -you may not use this file except in compliance with the License. -You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. -*/ - -package steps - -import ( - "reflect" - "testing" - - "github.com/llm-d/llm-d-router/pkg/coordinator/gateway" -) - -func TestCapSingleTokenOutput(t *testing.T) { - tests := []struct { - name string - format gateway.RequestFormat - body map[string]any - want map[string]any - }{ - { - name: "chat completions caps output fields and forces non-streaming", - format: gateway.FormatChatCompletions, - body: map[string]any{ - "model": "m", - "max_tokens": 100, - "min_tokens": 5, - "max_completion_tokens": 100, - "stream": true, - "stream_options": map[string]any{"include_usage": true}, - }, - want: map[string]any{ - "model": "m", - "max_tokens": 1, - "max_completion_tokens": 1, - "stream": false, - }, - }, - { - name: "max_completion_tokens is not added when the client omitted it", - format: gateway.FormatChatCompletions, - body: map[string]any{"model": "m"}, - want: map[string]any{ - "model": "m", - "max_tokens": 1, - "stream": false, - }, - }, - { - name: "completions caps max_tokens, strips min_tokens, forces non-streaming", - format: gateway.FormatCompletions, - body: map[string]any{"model": "m", "max_tokens": 100, "min_tokens": 5}, - want: map[string]any{"model": "m", "max_tokens": 1, "stream": false}, - }, - { - name: "streaming is forced false and stream_options stripped", - format: gateway.FormatCompletions, - body: map[string]any{"stream": true, "stream_options": map[string]any{"include_usage": true}}, - want: map[string]any{"stream": false, "max_tokens": 1}, - }, - { - name: "generate caps max_tokens and strips min_tokens inside sampling_params", - format: gateway.FormatGenerate, - body: map[string]any{ - "model": "m", - "sampling_params": map[string]any{"max_tokens": 100, "min_tokens": 5}, - }, - want: map[string]any{ - "model": "m", - "sampling_params": map[string]any{"max_tokens": 1}, - "stream": false, - }, - }, - { - name: "generate synthesizes sampling_params when absent", - format: gateway.FormatGenerate, - body: map[string]any{"model": "m"}, - want: map[string]any{ - "model": "m", - "sampling_params": map[string]any{"max_tokens": 1}, - "stream": false, - }, - }, - { - name: "generate preserves other sampling_params entries", - format: gateway.FormatGenerate, - body: map[string]any{ - "sampling_params": map[string]any{ - "extra_args": map[string]any{"kv_transfer_params": "x"}, - }, - }, - want: map[string]any{ - "sampling_params": map[string]any{ - "max_tokens": 1, - "extra_args": map[string]any{"kv_transfer_params": "x"}, - }, - "stream": false, - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - capSingleTokenOutput(tt.body, tt.format) - if !reflect.DeepEqual(tt.body, tt.want) { - t.Fatalf("got %v, want %v", tt.body, tt.want) - } - }) - } -} diff --git a/pkg/coordinator/steps/utils.go b/pkg/coordinator/steps/utils.go index f8e196349d..c10007d852 100644 --- a/pkg/coordinator/steps/utils.go +++ b/pkg/coordinator/steps/utils.go @@ -23,6 +23,7 @@ import ( "github.com/go-logr/logr" logutil "github.com/llm-d/llm-d-router/pkg/common/observability/logging" + reqcommon "github.com/llm-d/llm-d-router/pkg/common/request" "github.com/llm-d/llm-d-router/pkg/coordinator/gateway" "github.com/llm-d/llm-d-router/pkg/coordinator/pipeline" @@ -71,6 +72,40 @@ func resolveFormat(useOpenAIFormat bool, path string) gateway.RequestFormat { return detected } +// 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 +// 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) { + target := body + if format == gateway.FormatGenerate { + sp, ok := body[reqcommon.FieldSamplingParams].(map[string]any) + if !ok { + sp = map[string]any{} + body[reqcommon.FieldSamplingParams] = sp + } + target = sp + } + + target[reqcommon.FieldMaxTokens] = 1 + // Strip rather than clamp min_tokens: it defaults to 0 in vLLM, so removing it + // keeps min_tokens <= max_tokens=1 without raising the floor above the cap. + delete(target, reqcommon.FieldMinTokens) + + if _, ok := body[reqcommon.FieldMaxCompletionTokens]; ok { + body[reqcommon.FieldMaxCompletionTokens] = 1 + } + + // TODO: max_output_tokens is another client-supplied output cap (Responses + // API) that a client can send instead of max_tokens/max_completion_tokens; it + // should be capped to 1 here as well so the synthetic legs stay single-token. + + body[reqcommon.FieldStream] = false + delete(body, reqcommon.FieldStreamOptions) +} + // buildMMFeatures builds the multimodal features map (mm_hashes, mm_placeholders, // and optionally kwargs_data) from the request's multimodal entries. It returns // nil when there are no entries. diff --git a/pkg/coordinator/steps/utils_test.go b/pkg/coordinator/steps/utils_test.go index 2b3c65ca56..b70395ef93 100644 --- a/pkg/coordinator/steps/utils_test.go +++ b/pkg/coordinator/steps/utils_test.go @@ -17,8 +17,11 @@ limitations under the License. package steps import ( + "reflect" "strings" "testing" + + "github.com/llm-d/llm-d-router/pkg/coordinator/gateway" ) func TestReadErrorBody_CapsOversizedBody(t *testing.T) { @@ -34,3 +37,101 @@ func TestReadErrorBody_ReturnsSmallBodyVerbatim(t *testing.T) { t.Fatalf("expected %q, got %q", "overloaded", string(body)) } } + +func TestCapSingleTokenOutput(t *testing.T) { + tests := []struct { + name string + format gateway.RequestFormat + body map[string]any + want map[string]any + }{ + { + name: "chat completions caps output fields and forces non-streaming", + format: gateway.FormatChatCompletions, + body: map[string]any{ + "model": "m", + "max_tokens": 100, + "min_tokens": 5, + "max_completion_tokens": 100, + "stream": true, + "stream_options": map[string]any{"include_usage": true}, + }, + want: map[string]any{ + "model": "m", + "max_tokens": 1, + "max_completion_tokens": 1, + "stream": false, + }, + }, + { + name: "max_completion_tokens is not added when the client omitted it", + format: gateway.FormatChatCompletions, + body: map[string]any{"model": "m"}, + want: map[string]any{ + "model": "m", + "max_tokens": 1, + "stream": false, + }, + }, + { + name: "completions caps max_tokens, strips min_tokens, forces non-streaming", + format: gateway.FormatCompletions, + body: map[string]any{"model": "m", "max_tokens": 100, "min_tokens": 5}, + want: map[string]any{"model": "m", "max_tokens": 1, "stream": false}, + }, + { + name: "streaming is forced false and stream_options stripped", + format: gateway.FormatCompletions, + body: map[string]any{"stream": true, "stream_options": map[string]any{"include_usage": true}}, + want: map[string]any{"stream": false, "max_tokens": 1}, + }, + { + name: "generate caps max_tokens and strips min_tokens inside sampling_params", + format: gateway.FormatGenerate, + body: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 100, "min_tokens": 5}, + }, + want: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 1}, + "stream": false, + }, + }, + { + name: "generate synthesizes sampling_params when absent", + format: gateway.FormatGenerate, + body: map[string]any{"model": "m"}, + want: map[string]any{ + "model": "m", + "sampling_params": map[string]any{"max_tokens": 1}, + "stream": false, + }, + }, + { + name: "generate preserves other sampling_params entries", + format: gateway.FormatGenerate, + body: map[string]any{ + "sampling_params": map[string]any{ + "extra_args": map[string]any{"kv_transfer_params": "x"}, + }, + }, + want: map[string]any{ + "sampling_params": map[string]any{ + "max_tokens": 1, + "extra_args": map[string]any{"kv_transfer_params": "x"}, + }, + "stream": false, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + capSingleTokenOutput(tt.body, tt.format) + if !reflect.DeepEqual(tt.body, tt.want) { + t.Fatalf("got %v, want %v", tt.body, tt.want) + } + }) + } +} From bb7803195a34450e246f38d084dc2a34b734ee86 Mon Sep 17 00:00:00 2001 From: Revital Sur Date: Tue, 28 Jul 2026 10:34:53 +0300 Subject: [PATCH 4/5] Fix e2e test. Signed-off-by: Revital Sur --- test/coordinator/e2e/coordinator/coordinator_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/coordinator/e2e/coordinator/coordinator_test.go b/test/coordinator/e2e/coordinator/coordinator_test.go index ba1317d060..8b106e3dbf 100644 --- a/test/coordinator/e2e/coordinator/coordinator_test.go +++ b/test/coordinator/e2e/coordinator/coordinator_test.go @@ -264,12 +264,12 @@ func fetchCoordinatorLogs(nsName string) string { // on the coordinator running at log_level 5. func verifyTokenLimits(logs string, wantMin, wantMax int, capLegs []string) { ginkgo.By("Verifying decode leg forwards the client min_tokens/max_tokens") - gomega.Expect(logHasLine(logs, `"msg":"request body"`, `"epp-phase":"decode"`, + gomega.Expect(logHasLine(logs, `"msg":"request body"`, `"epp-profile":"decode"`, fmt.Sprintf(`"min_tokens":%d`, wantMin), fmt.Sprintf(`"max_tokens":%d`, wantMax))).To(gomega.BeTrue(), "coordinator logs have no decode request body carrying min_tokens=%d and max_tokens=%d", wantMin, wantMax) for _, phase := range capLegs { - phaseField := `"epp-phase":"` + phase + `"` + phaseField := `"epp-profile":"` + phase + `"` ginkgo.By("Verifying " + phase + " leg caps max_tokens to 1") gomega.Expect(logHasLine(logs, `"msg":"request body"`, phaseField, `"max_tokens":1`)).To(gomega.BeTrue(), From 0471c32c06ff4364baea6c8945da5d2e2cef8db2 Mon Sep 17 00:00:00 2001 From: Revital Sur Date: Thu, 30 Jul 2026 11:57:38 +0300 Subject: [PATCH 5/5] Add e2e test. Signed-off-by: Revital Sur --- .../coordinator/e2e/coordinator/configs_test.go | 10 ++++++++++ .../e2e/coordinator/coordinator_test.go | 17 +++++++++++++++-- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/test/coordinator/e2e/coordinator/configs_test.go b/test/coordinator/e2e/coordinator/configs_test.go index f030b0f290..9c71b22ecd 100644 --- a/test/coordinator/e2e/coordinator/configs_test.go +++ b/test/coordinator/e2e/coordinator/configs_test.go @@ -16,6 +16,8 @@ limitations under the License. package coordinate2e +import "strings" + // coordinatorConfigNIXL is the coordinator pipeline config for the e-p-d-pools topology. // ${NAMESPACE} and ${VLLM_RENDER_PORT} are substituted by createCoordinator before the ConfigMap is built. const coordinatorConfigNIXL = `log_level: 5 @@ -55,6 +57,14 @@ pipeline: - type: decode ` +// coordinatorConfigNIXLGenerate is coordinatorConfigNIXL with OpenAI passthrough +// disabled. With use_openai_format: false a chat/completions request collapses to +// the native generate wire format on the encode and prefill legs, which build a +// fresh sampling_params carrying only max_tokens: 1. This exercises a different +// min_tokens-stripping path than the OpenAI clone-and-cap path. +var coordinatorConfigNIXLGenerate = strings.Replace( + coordinatorConfigNIXL, "use_openai_format: true", "use_openai_format: false", 1) + // eppConfig is the scheduling config for the single EPP that serves all three // phases. Each request runs exactly one scheduling profile, named by its // EPP-Profile header value (see header-profile-handler); the role filters diff --git a/test/coordinator/e2e/coordinator/coordinator_test.go b/test/coordinator/e2e/coordinator/coordinator_test.go index 7ede89bc26..8ff1edc374 100644 --- a/test/coordinator/e2e/coordinator/coordinator_test.go +++ b/test/coordinator/e2e/coordinator/coordinator_test.go @@ -61,6 +61,13 @@ var _ = ginkgo.Describe("Coordinator pipeline", func() { )), allSteps, 1, 3, 5) }) + ginkgo.It("forwards min_tokens/max_tokens to decode and caps them on prefill and encode with OpenAI passthrough disabled", func() { + runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( + `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"min_tokens":3,"max_tokens":5}`, + modelName, inlineImageDataURI, + )), allSteps, 1, 3, 5, coordinatorConfigNIXLGenerate) + }) + ginkgo.It("routes a multimodal image chat completion end-to-end", func() { runCoordinatorPipeline(gateway.PathChatCompletions, []byte(fmt.Sprintf( `{"model":%q,"messages":[{"role":"user","content":[{"type":"image_url","image_url":{"url":%q},"uuid":"image-0"},{"type":"text","text":"Describe what you see."}]}],"max_tokens":150}`, @@ -105,7 +112,9 @@ const inlineImageDataURI = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAEAAAA // contract is additionally asserted: decode forwards the client's // min_tokens/max_tokens unchanged, while prefill (and encode, for multimodal // requests) cap max_tokens to 1 and strip min_tokens. -func runCoordinatorPipeline(path string, body []byte, expectedSteps []string, expectedImages, wantMinTokens, wantMaxTokens int) { +// coordinatorConfig, when supplied, overrides the default coordinatorConfigNIXL +// pipeline config for the coordinator deployment. +func runCoordinatorPipeline(path string, body []byte, expectedSteps []string, expectedImages, wantMinTokens, wantMaxTokens int, coordinatorConfig ...string) { nsName := getNamespace() var ( coordinator []string @@ -151,7 +160,11 @@ func runCoordinatorPipeline(path string, body []byte, expectedSteps []string, ex gomega.Expect(prefillPods).Should(gomega.HaveLen(prefillReplicas)) gomega.Expect(decodePods).Should(gomega.HaveLen(decodeReplicas)) - coordinator = createCoordinator(coordinatorConfigNIXL) + cfg := coordinatorConfigNIXL + if len(coordinatorConfig) > 0 { + cfg = coordinatorConfig[0] + } + coordinator = createCoordinator(cfg) req, err := http.NewRequest(http.MethodPost, gatewayBaseURL()+path,