fix(sglang): support min_tokens on tokenizer-free decode workers

.github/workflows/pr.yaml

@@ -940,6 +940,7 @@ jobs:
     uses: $/.github/workflows/shared-test.yml
     with:
       test_suite_name: triton
+      framework: triton
       test_type: CPU Test
       amd_runner: prod-tester-amd-v2
       target_tag_plain: ${{ needs.triton-build.outputs.target_tag_plain }}

components/src/dynamo/sglang/request_handlers/llm/decode_handler.py

@@ -361,12 +361,25 @@ def _build_sampling_params(self, request: Dict[str, Any]) -> Dict[str, Any]:
             _plain = stop_conditions.get("stop_token_ids") or []
             _merged = list(set(_hidden).union(_plain))
             stop_token_ids = _merged if _merged else None
+            # A tokenizer-free SGLang rejects positive min_new_tokens. Let
+            # Dynamo's decoder enforce the floor and keep the engine running so
+            # it cannot stop before Dynamo reaches it.
+            min_tokens = stop_conditions.get("min_tokens")
+            hold_engine_open = bool(
+                min_tokens
+                and min_tokens > 0
+                and self.config.server_args.skip_tokenizer_init
+            )
+            if hold_engine_open:
+                stop_token_ids = None
 
             param_mapping = {
                 "n": sampling_opts.get("n"),
                 "max_new_tokens": stop_conditions.get("max_tokens"),
-                "min_new_tokens": stop_conditions.get("min_tokens"),
-                "ignore_eos": stop_conditions.get("ignore_eos"),
+                "min_new_tokens": None if hold_engine_open else min_tokens,
+                "ignore_eos": True
+                if hold_engine_open
+                else stop_conditions.get("ignore_eos"),
                 "stop_token_ids": stop_token_ids,
                 **_sampling_option_params(sampling_opts),
                 **self._get_guided_decoding_params(

components/src/dynamo/sglang/tests/test_sglang_decode_handler.py

@@ -242,12 +242,20 @@ def test_openai_stop_sampling_params_maps_token_id_stop_array():
     }
 
 
-def _new_decode_handler(*, use_sglang_tokenizer: bool = False, enable_rl: bool = False):
+def _new_decode_handler(
+    *,
+    use_sglang_tokenizer: bool = False,
+    skip_tokenizer_init: bool = False,
+    enable_rl: bool = False,
+):
     handler = DecodeWorkerHandler.__new__(DecodeWorkerHandler)
     handler.shutdown_event = None
     handler.use_sglang_tokenizer = use_sglang_tokenizer
     handler.config = SimpleNamespace(
-        server_args=SimpleNamespace(served_model_name="test-model"),
+        server_args=SimpleNamespace(
+            served_model_name="test-model",
+            skip_tokenizer_init=skip_tokenizer_init,
+        ),
         dynamo_args=SimpleNamespace(enable_rl=enable_rl),
     )
     handler._first_token_source = None
@@ -813,23 +821,55 @@ def test_guided_decoding_params_are_accepted_by_sglang(guided_decoding):
     SamplingParams(**sampling_params)
 
 
-def test_build_sampling_params_maps_min_tokens_for_token_requests():
-    handler = _new_decode_handler(use_sglang_tokenizer=False)
+def test_build_sampling_params_holds_tokenizer_free_engine_open_for_min_tokens():
+    handler = _new_decode_handler(
+        use_sglang_tokenizer=False,
+        skip_tokenizer_init=True,
+    )
 
     sampling_params = handler._build_sampling_params(
         {
             "sampling_options": {},
             "stop_conditions": {
                 "max_tokens": 64,
-                "min_tokens": 64,
-                "ignore_eos": True,
+                "min_tokens": 8,
+                "stop_token_ids": [7],
+                "stop_token_ids_hidden": [151643],
             },
         }
     )
 
-    assert sampling_params["min_new_tokens"] == 64
-    assert sampling_params["max_new_tokens"] == 64
+    assert "min_new_tokens" not in sampling_params
+    assert "stop_token_ids" not in sampling_params
     assert sampling_params["ignore_eos"] is True
+    assert sampling_params["max_new_tokens"] == 64
+
+    from sglang.srt.sampling.sampling_params import SamplingParams
+
+    SamplingParams(**sampling_params).normalize(tokenizer=None)
+
+
+def test_build_sampling_params_forwards_min_tokens_when_engine_has_tokenizer():
+    handler = _new_decode_handler(
+        use_sglang_tokenizer=False,
+        skip_tokenizer_init=False,
+    )
+
+    sampling_params = handler._build_sampling_params(
+        {
+            "sampling_options": {},
+            "stop_conditions": {
+                "max_tokens": 64,
+                "min_tokens": 8,
+                "stop_token_ids": [7],
+            },
+        }
+    )
+
+    assert sampling_params["min_new_tokens"] == 8
+    assert sampling_params["stop_token_ids"] == [7]
+    assert "ignore_eos" not in sampling_params
+    assert sampling_params["max_new_tokens"] == 64
 
 
 def test_build_sampling_params_maps_min_tokens_for_sglang_tokenizer_requests():
@@ -844,7 +884,10 @@ def test_build_sampling_params_maps_min_tokens_for_sglang_tokenizer_requests():
 
 
 def test_build_sampling_params_omits_min_new_tokens_when_min_tokens_absent():
-    handler = _new_decode_handler(use_sglang_tokenizer=False)
+    handler = _new_decode_handler(
+        use_sglang_tokenizer=False,
+        skip_tokenizer_init=True,
+    )
 
     sampling_params = handler._build_sampling_params(
         {"sampling_options": {}, "stop_conditions": {"max_tokens": 8}}
@@ -854,7 +897,10 @@ def test_build_sampling_params_omits_min_new_tokens_when_min_tokens_absent():
 
 
 def test_build_sampling_params_forwards_explicit_zero_min_tokens():
-    handler = _new_decode_handler(use_sglang_tokenizer=False)
+    handler = _new_decode_handler(
+        use_sglang_tokenizer=False,
+        skip_tokenizer_init=True,
+    )
 
     sampling_params = handler._build_sampling_params(
         {"sampling_options": {}, "stop_conditions": {"max_tokens": 8, "min_tokens": 0}}

lib/llm/src/backend.rs

@@ -253,7 +253,22 @@ impl
                 });
             }
 
-            match state.stream.next().await {
+            let output = loop {
+                let output = state.stream.next().await;
+                let choice_already_finished = output.as_ref().is_some_and(|output| {
+                    !output.is_event()
+                        && output.data.as_ref().is_some_and(|data| {
+                            state
+                                .finished_choices
+                                .contains(&data.index.unwrap_or_default())
+                        })
+                });
+                if !choice_already_finished {
+                    break output;
+                }
+            };
+
+            match output {
                 Some(output) => {
                     // move to state.process_output
                     // handle any error conditions / unwraps here
@@ -486,6 +501,20 @@ impl
                         data.stop_reason = stop_reason.or(data.stop_reason);
                     }
                     data.text = text;
+
+                    // A tokenizer-free engine can return tokens past a local
+                    // stop in the same chunk. Keep its per-token fields aligned
+                    // with the prefix the decoder consumed.
+                    let consumed = tokens.len();
+                    if consumed < data.token_ids.len() {
+                        data.token_ids.truncate(consumed);
+                        if let Some(log_probs) = data.log_probs.as_mut() {
+                            log_probs.truncate(consumed);
+                        }
+                        if let Some(top_logprobs) = data.top_logprobs.as_mut() {
+                            top_logprobs.truncate(consumed);
+                        }
+                    }
                     data.tokens = Some(tokens);
                     // Snapshot of whatever this choice's decoder is still withholding as a
                     // possible hidden-stop-sequence prefix after this step -- `None` once
@@ -1147,6 +1176,7 @@ mod tests {
                 [102] => " Okay",
                 [103] => "<|user|>",
                 [101, 103] => "Okay<|user|>",
+                [101, 103, 102] => "Okay<|user|> Okay",
                 _ => anyhow::bail!("unexpected token IDs: {token_ids:?}"),
             };
             Ok(traits::DecodeResult::Complete(token.to_string()))
@@ -1169,20 +1199,26 @@ mod tests {
             &self,
             request: SingleIn<PreprocessedRequest>,
         ) -> Result<ManyOut<Annotated<LLMEngineOutput>>, Error> {
-            let output = LLMEngineOutput {
-                // SGLang's output_ids_through_stop() includes the matched stop
-                // position. The frontend must hide its text without mutating
-                // this raw token sequence.
-                token_ids: vec![101, 103],
-                finish_reason: Some(FinishReason::Stop),
-                index: Some(0),
-                ..Default::default()
+            let output = |index, token_ids: Vec<TokenIdType>| {
+                let len = token_ids.len();
+                Annotated::from_data(LLMEngineOutput {
+                    token_ids,
+                    log_probs: Some(vec![-0.1; len]),
+                    top_logprobs: Some(vec![vec![]; len]),
+                    index: Some(index),
+                    ..Default::default()
+                })
             };
+            let outputs = vec![
+                // Choice 0 stops inside this chunk; its suffix must be removed.
+                output(0, vec![101, 103, 102]),
+                // It keeps generating while choice 1 is unfinished; drop this chunk.
+                output(0, vec![102]),
+                output(1, vec![101, 103]),
+            ];
 
             Ok(ResponseStream::new(
-                Box::pin(futures::stream::once(async move {
-                    Annotated::from_data(output)
-                })),
+                Box::pin(futures::stream::iter(outputs)),
                 request.context(),
             ))
         }
@@ -1329,7 +1365,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn test_sglang_hidden_stop_text_is_suppressed_without_mutating_tito_ids() {
+    async fn test_sglang_local_stop_discards_all_suffix_tokens() {
         use crate::protocols::common::extensions::NvExt;
 
         let tokenizer: Arc<dyn traits::Tokenizer> = Arc::new(CandidateDecoder);
@@ -1341,27 +1377,45 @@ mod tests {
                 stop_token_ids_hidden: Some(vec![103]),
                 ..Default::default()
             })
-            .sampling_options(SamplingOptions::default())
+            .sampling_options(SamplingOptions {
+                n: Some(2),
+                ..Default::default()
+            })
             .output_options(OutputOptions::default())
             .build()
             .expect("valid preprocessed request");
         let engine: ServerStreamingEngine<PreprocessedRequest, Annotated<LLMEngineOutput>> =
             Arc::new(SyntheticSglangStopEngine);
 
-        let mut stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine)
+        let stream = Operator::generate(backend.as_ref(), SingleIn::new(request), engine)
             .await
             .expect("backend generation succeeds");
-        let output = stream
-            .next()
-            .await
-            .expect("backend emits a response")
+        let mut outputs = stream.collect::<Vec<_>>().await;
+        assert_eq!(
+            outputs.len(),
+            2,
+            "chunks after a local stop must be dropped"
+        );
+        let output = outputs
+            .remove(0)
+            .data
+            .expect("response contains backend output");
+        let second = outputs
+            .remove(0)
             .data
             .expect("response contains backend output");
 
+        assert_eq!(output.index, Some(0));
         assert_eq!(output.text.as_deref(), Some("Okay"));
         assert_eq!(output.tokens, vec![Some("Okay".to_string()), None]);
         assert_eq!(output.token_ids, vec![101, 103]);
+        assert_eq!(output.log_probs, Some(vec![-0.1; 2]));
+        assert_eq!(output.top_logprobs.as_ref().map(Vec::len), Some(2));
         assert_eq!(output.finish_reason, Some(FinishReason::Stop));
+        assert_eq!(second.index, Some(1));
+        assert_eq!(second.text.as_deref(), Some("Okay"));
+        assert_eq!(second.token_ids, vec![101, 103]);
+        assert_eq!(second.finish_reason, Some(FinishReason::Stop));
 
         let nvext = NvExt::builder()
             .extra_fields(vec!["completion_token_ids".to_string()])