fix(sidecar): cap min_tokens on the prefill leg for chat completions (llm-d#2041)

* fix(sidecar): cap min_tokens on the prefill leg for chat completions

The prefill leg caps max_tokens at 1 to suppress output, but min_tokens
was staged only for the generate API. On the chat completions APIs a
client-supplied min_tokens above 1 reached the prefiller next to
max_tokens=1, which vLLM rejects because min_tokens may not exceed
max_tokens, so the request failed with a 400 before decode ran.

Stage min_tokens with the other chat completions token limits, and do
the same in the MoRI-IO concurrent-dispatch path, which builds its
prefill body separately. The saved-value buffer now sizes off the field
list, which a fixed two would have overflowed.

Signed-off-by: Maroon Ayoub <mayoub@redhat.com>

* Potential fix for pull request finding

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Maroon Ayoub <Maroonay@gmail.com>

---------

Signed-off-by: Maroon Ayoub <mayoub@redhat.com>
Signed-off-by: Maroon Ayoub <Maroonay@gmail.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Author: vMaroon

pkg/sidecar/proxy/connector_nixlv2.go

@@ -115,7 +115,7 @@ func (s *Server) handleNIXLV2(w http.ResponseWriter, r *http.Request, prefillPod
 		present bool
 	}
 	tokenMap, createdSamplingParams := tokenLimitMap(completionRequest, apiType)
-	var savedTokenValues [2]savedField
+	savedTokenValues := make([]savedField, len(tokenLimitFields))
 	for i, field := range tokenLimitFields {
 		if v, ok := tokenMap[field]; ok {
 			savedTokenValues[i] = savedField{field: field, val: v, present: true}
@@ -324,7 +324,7 @@ retryLoop:
 		completionRequest[requestFieldStreamOptions] = streamOptionsValue
 	}
 
-	for i := range savedTokenValues[:len(tokenLimitFields)] {
+	for i := range savedTokenValues {
 		sv := &savedTokenValues[i]
 		delete(tokenMap, sv.field)
 		if sv.present {
@@ -459,6 +459,7 @@ func (s *Server) runNIXLProtocolV2WriteParallel(
 	maxTokensValue, maxTokensOk := completionRequest[requestFieldMaxTokens]
 	maxCompletionTokensValue, maxCompletionTokensOk := completionRequest[requestFieldMaxCompletionTokens]
 	maxOutputTokensValue, maxOutputTokensOk := completionRequest[requestFieldMaxOutputTokens]
+	minTokensValue, minTokensOk := completionRequest[requestFieldMinTokens]
 
 	// Pin both legs to the same DP rank (kv_transfer_params + HTTP header).
 	dpRank := pickDPRank(uuidStr, s.config.MoRIIODPSize)
@@ -503,6 +504,7 @@ func (s *Server) runNIXLProtocolV2WriteParallel(
 	completionRequest[requestFieldMaxTokens] = 1
 	completionRequest[requestFieldMaxCompletionTokens] = 1
 	completionRequest[requestFieldMaxOutputTokens] = 1
+	completionRequest[requestFieldMinTokens] = 1
 
 	pbody, err := json.Marshal(completionRequest)
 	if err != nil {
@@ -513,7 +515,7 @@ func (s *Server) runNIXLProtocolV2WriteParallel(
 	}
 
 	// ---------- Build decode body ----------
-	// Restore the client's streaming flags and max-token caps.
+	// Restore the client's streaming flags and token-limit fields.
 	delete(completionRequest, requestFieldStream)
 	if streamOk {
 		completionRequest[requestFieldStream] = streamValue
@@ -533,6 +535,10 @@ func (s *Server) runNIXLProtocolV2WriteParallel(
 	if maxOutputTokensOk {
 		completionRequest[requestFieldMaxOutputTokens] = maxOutputTokensValue
 	}
+	delete(completionRequest, requestFieldMinTokens)
+	if minTokensOk {
+		completionRequest[requestFieldMinTokens] = minTokensValue
+	}
 
 	// Synthesise decode-leg kv_transfer_params that the serial path would
 	// otherwise read from the prefill response. do_remote_prefill must be true:

pkg/sidecar/proxy/connector_nixlv2_test.go

@@ -249,6 +249,60 @@ var _ = Describe("NIXL Connector (v2)", func() {
 		Expect(responseBody).To(ContainSubstring("data: [DONE]"))
 	})
 
+	sendChatCompletionsRequestWithBody := func(proxyBaseAddr, body string) {
+		req, err := http.NewRequest(http.MethodPost, proxyBaseAddr+ChatCompletionsPath, bytes.NewReader([]byte(body)))
+		Expect(err).ToNot(HaveOccurred())
+		req.Header.Add(routing.PrefillEndpointHeader, testInfo.prefillBackend.URL[len("http://"):])
+
+		rp, err := http.DefaultClient.Do(req)
+		Expect(err).ToNot(HaveOccurred())
+		defer rp.Body.Close()
+
+		responseBody, err := io.ReadAll(rp.Body)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(rp.StatusCode).To(Equal(http.StatusOK), string(responseBody))
+	}
+
+	It("should cap token limits in prefill and restore originals in decode", func() {
+		proxyBaseAddr := startProxy()
+
+		sendChatCompletionsRequestWithBody(proxyBaseAddr, `{
+				"model": "Qwen/Qwen2-0.5B",
+				"messages": [
+				  {"role": "user", "content": "Hello"}
+				],
+				"max_tokens": 100,
+				"min_tokens": 5
+			}`)
+
+		prefillReq := testInfo.prefillHandler.CompletionRequests[0]
+		Expect(prefillReq).To(HaveKeyWithValue(requestFieldMaxTokens, BeNumerically("==", 1)))
+		Expect(prefillReq).To(HaveKeyWithValue(requestFieldMinTokens, BeNumerically("==", 1)))
+
+		decodeReq := testInfo.decodeHandler.CompletionRequests[0]
+		Expect(decodeReq).To(HaveKeyWithValue(requestFieldMaxTokens, BeNumerically("==", 100)))
+		Expect(decodeReq).To(HaveKeyWithValue(requestFieldMinTokens, BeNumerically("==", 5)))
+	})
+
+	It("should cap prefill and drop the caps in decode when the request omits them", func() {
+		proxyBaseAddr := startProxy()
+
+		sendChatCompletionsRequestWithBody(proxyBaseAddr, `{
+				"model": "Qwen/Qwen2-0.5B",
+				"messages": [
+				  {"role": "user", "content": "Hello"}
+				]
+			}`)
+
+		prefillReq := testInfo.prefillHandler.CompletionRequests[0]
+		Expect(prefillReq).To(HaveKeyWithValue(requestFieldMaxTokens, BeNumerically("==", 1)))
+		Expect(prefillReq).To(HaveKeyWithValue(requestFieldMinTokens, BeNumerically("==", 1)))
+
+		decodeReq := testInfo.decodeHandler.CompletionRequests[0]
+		Expect(decodeReq).ToNot(HaveKey(requestFieldMaxTokens))
+		Expect(decodeReq).ToNot(HaveKey(requestFieldMinTokens))
+	})
+
 	// Messages API tests — verify /v1/messages routes through the disaggregation
 	// handler with the same token-limit fields as chat completions.
 
@@ -1018,6 +1072,30 @@ var _ = Describe("NIXL Connector (v2)", func() {
 		Expect(pRank).To(And(BeNumerically(">=", 0), BeNumerically("<", 16)))
 	})
 
+	// The concurrent-dispatch path stages its own prefill body rather than going
+	// through the serial path's token-limit handling, so it caps min_tokens too.
+	It("concurrent WRITE-mode dispatch caps min_tokens in prefill and restores it in decode", func() {
+		env := startMoRIProxy(func(c *Config) {
+			c.MoRIIOParallelDispatch = true
+		})
+		env.sendBody(`{
+				"model": "Qwen/Qwen2-0.5B",
+				"messages": [
+				  {"role": "user", "content": "Hello"}
+				],
+				"max_tokens": 100,
+				"min_tokens": 5
+			}`)
+
+		prefillReq := env.prefillHandler.GetCompletionRequests()[0]
+		Expect(prefillReq).To(HaveKeyWithValue(requestFieldMaxTokens, BeNumerically("==", 1)))
+		Expect(prefillReq).To(HaveKeyWithValue(requestFieldMinTokens, BeNumerically("==", 1)))
+
+		decodeReq := env.decodeHandler.GetCompletionRequests()[0]
+		Expect(decodeReq).To(HaveKeyWithValue(requestFieldMaxTokens, BeNumerically("==", 100)))
+		Expect(decodeReq).To(HaveKeyWithValue(requestFieldMinTokens, BeNumerically("==", 5)))
+	})
+
 	// 1P1D DP=8, serial dispatch: the prefill leg sets the DP-rank header and
 	// the decode leg's kv_transfer_params are backfilled with the same rank.
 	It("serial WRITE-mode DP=8 pins prefill and decode HTTP legs to the same DP rank", func() {
@@ -1142,15 +1220,20 @@ func startMoRIProxy(mutate func(cfg *Config)) *moriProxyEnv {
 // sequential in the serial path) by the time this returns, so the captured
 // requests / headers are safe to read afterwards.
 func (env *moriProxyEnv) send() {
-	req, err := http.NewRequest(http.MethodPost, env.baseAddr+ChatCompletionsPath, strings.NewReader(chatCompletionsRequestBody))
+	env.sendBody(chatCompletionsRequestBody)
+}
+
+// sendBody sends a caller-supplied request body.
+func (env *moriProxyEnv) sendBody(body string) {
+	req, err := http.NewRequest(http.MethodPost, env.baseAddr+ChatCompletionsPath, strings.NewReader(body))
 	Expect(err).ToNot(HaveOccurred())
 	req.Header.Add(routing.PrefillEndpointHeader, env.prefillBackend.URL[len("http://"):])
 
 	rp, err := http.DefaultClient.Do(req)
 	Expect(err).ToNot(HaveOccurred())
 	defer rp.Body.Close()
-	body, _ := io.ReadAll(rp.Body) //nolint:errcheck
-	Expect(rp.StatusCode).To(Equal(http.StatusOK), string(body))
+	responseBody, _ := io.ReadAll(rp.Body) //nolint:errcheck
+	Expect(rp.StatusCode).To(Equal(http.StatusOK), string(responseBody))
 }
 
 // kvParams returns the kv_transfer_params map of the i-th request captured by h.

pkg/sidecar/proxy/proxy.go

@@ -136,7 +136,7 @@ func (a APIType) String() string {
 // JSON request field names used for token limits in prefill/decode staging.
 // Do not mutate these slices.
 var (
-	chatCompletionTokenLimitFields = []string{requestFieldMaxTokens, requestFieldMaxCompletionTokens}
+	chatCompletionTokenLimitFields = []string{requestFieldMaxTokens, requestFieldMaxCompletionTokens, requestFieldMinTokens}
 	responsesStyleTokenLimitFields = []string{requestFieldMaxOutputTokens}
 	generateStyleTokenLimitFields  = []string{requestFieldMaxTokens, requestFieldMinTokens}
 )