Let the multimodal generator enable thinking and fit its output budget

CHANGELOG.rst

@@ -60,6 +60,7 @@ Changelog
 - LiLiCorr now trains offline and in streaming mode, not only online: the target logits its distractor penalty reads are reconstructed from the captured final hidden state, the same path self-logit distillation uses.
 - Add ``dflash_lilicorr_w_cal`` (default ``0.0``, off): an optional LiLiCorr proposal-calibration term, a KL divergence from the target's distribution over each slot's candidates to the reranker's candidate distribution, usable in place of the distractor penalty.
 - ``examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py`` ``--api_mode openai`` works against any OpenAI-compatible server, including vLLM, and gains ``--media_inline`` to send media as base64 ``data:`` URIs. Without ``--media_url_base`` or ``--media_inline`` the media path is still passed through verbatim, which only servers that read their own filesystem resolve (SGLang native does; vLLM needs ``--allowed-local-media-path``); that case now warns.
+- ``server_generate_vlm_sglang.py`` gains ``--enable_thinking``, which sends ``chat_template_kwargs`` so a reasoning target actually reasons; without it the template's guard leaves thinking off and nothing in the output records which mode produced it. Requests that exceed the model's context window now retry with a halved output budget down to ``--min_max_tokens`` instead of being dropped, since the requests that overflow are the longest-context ones.
 
 *Megatron Framework (M-LM / M-Bridge)*
 

examples/speculative_decoding/distributed_generate/server_generate_vlm_sglang.py

@@ -357,6 +357,16 @@ def _prompt_with_vision_token(prompt: str, media_type: str, token_format: str) -
     return f"{QWEN_IMAGE_TOKEN}\n{prompt}"
 
 
+def _is_context_overflow(body: str) -> bool:
+    """Whether a non-200 body reports the request exceeding the model's context window.
+
+    Matched on the message rather than the status code: a 400 covers many unrelated
+    problems, and retrying those with a smaller budget would just burn requests.
+    """
+    lowered = body.lower()
+    return "maximum context length" in lowered or "longer than the maximum" in lowered
+
+
 def _openai_generate_one(
     request: dict[str, Any],
     url: str,
@@ -368,6 +378,8 @@ def _openai_generate_one(
     media_root: str | None,
     input_root: str | None,
     media_inline: bool = False,
+    enable_thinking: bool = False,
+    min_max_tokens: int = 256,
 ) -> tuple[str, str | None]:
     try:
         import requests
@@ -412,20 +424,38 @@ def _openai_generate_one(
         "temperature": temperature,
         "max_tokens": max_tokens,
     }
-
+    if enable_thinking:
+        # Reasoning models gate their <think> prefill on a chat-template variable, and the
+        # usual guard is ``enable_thinking is defined and enable_thinking`` -- so leaving it
+        # unset is indistinguishable from setting it false, and the target answers without
+        # reasoning. Nothing in the response says which mode produced it.
+        payload["chat_template_kwargs"] = {"enable_thinking": True}
+
+    budget = max_tokens
     try:
-        response = requests.post(url, json=payload, timeout=timeout)
-        if response.status_code != 200:
-            return "", f"HTTP {response.status_code}: {response.text[:1000]}"
-        data = response.json()
-        choices = data.get("choices") or []
-        if not choices:
-            return "", f"No choices in OpenAI response: {data}"
-        message = choices[0].get("message") or {}
-        answer = _coerce_text(message.get("content"))
-        if not answer:
-            return "", f"Could not extract assistant content from OpenAI response: {data}"
-        return answer, None
+        while True:
+            payload["max_tokens"] = budget
+            response = requests.post(url, json=payload, timeout=timeout)
+            if response.status_code != 200:
+                body = response.text[:1000]
+                # The server REJECTS an over-budget request, it does not truncate: prompt
+                # tokens plus max_tokens must fit the context window or it is a 400.
+                # Returning the error here would drop the sample, and the samples that
+                # overflow are the longest-context ones -- long videos -- which is exactly
+                # the data a multimodal draft model needs to have seen. Shrink and retry.
+                if _is_context_overflow(body) and budget > min_max_tokens:
+                    budget = max(min_max_tokens, budget // 2)
+                    continue
+                return "", f"HTTP {response.status_code}: {body}"
+            data = response.json()
+            choices = data.get("choices") or []
+            if not choices:
+                return "", f"No choices in OpenAI response: {data}"
+            message = choices[0].get("message") or {}
+            answer = _coerce_text(message.get("content"))
+            if not answer:
+                return "", f"Could not extract assistant content from OpenAI response: {data}"
+            return answer, None
     except Exception:
         return "", traceback.format_exc()
 
@@ -507,6 +537,20 @@ def main() -> None:
         help="Send media as base64 data: URIs instead of paths/URLs. Self-contained, but "
         "inflates every request body -- prefer --media_url_base for video.",
     )
+    parser.add_argument(
+        "--enable_thinking",
+        action="store_true",
+        help="Ask the chat template to prefill the reasoning block, for targets that gate "
+        "it on an enable_thinking template variable. Reasoning traces are far longer than "
+        "direct answers, so raise --max_tokens accordingly.",
+    )
+    parser.add_argument(
+        "--min_max_tokens",
+        type=int,
+        default=256,
+        help="Floor when shrinking the output budget to fit the context window. Below this "
+        "a completion is too stunted to be worth keeping.",
+    )
     parser.add_argument(
         "--api_mode",
         choices=("openai", "native"),
@@ -602,6 +646,8 @@ def main() -> None:
                         args.media_root,
                         args.input_root,
                         args.media_inline,
+                        args.enable_thinking,
+                        args.min_max_tokens,
                     ): meta
                     for request_args, meta in zip(batch_args, batch_meta)
                 }

tests/examples/speculative_decoding/test_server_generate_vlm_sglang.py

@@ -16,6 +16,7 @@
 """Tests for the multimodal SGLang generation client."""
 
 import base64
+import copy
 import importlib.util
 from pathlib import Path
 
@@ -173,3 +174,145 @@ def test_openai_media_value_prefers_remote_urls_over_inlining(tmp_path):
     assert (
         server_generate_vlm_sglang._as_openai_media_value(remote, None, None, None, True) == remote
     )
+
+
+def test_context_overflow_is_recognised_only_from_the_message():
+    """A 400 covers many problems; only the context message should trigger a retry."""
+
+    assert server_generate_vlm_sglang._is_context_overflow(
+        "This model's maximum context length is 32768 tokens. However, you requested 40000"
+    )
+    assert server_generate_vlm_sglang._is_context_overflow(
+        "Input is longer than the maximum supported length"
+    )
+    assert not server_generate_vlm_sglang._is_context_overflow("Invalid value for 'top_k'")
+    assert not server_generate_vlm_sglang._is_context_overflow("")
+
+
+def _fake_requests(monkeypatch, responses):
+    """Install a requests stub that returns ``responses`` in order, recording payloads."""
+
+    sent = []
+
+    class _Resp:
+        def __init__(self, status, body):
+            self.status_code = status
+            self.text = body
+            self._body = body
+
+        def json(self):
+            return {"choices": [{"message": {"content": "ok"}}]}
+
+    queue = list(responses)
+
+    class _Requests:
+        @staticmethod
+        def post(url, json=None, timeout=None):
+            # Snapshot: the caller mutates one payload dict across retries, so recording
+            # the reference would make every entry show the final budget.
+            sent.append(copy.deepcopy(json))
+            return _Resp(*queue.pop(0))
+
+    monkeypatch.setitem(__import__("sys").modules, "requests", _Requests)
+    return sent
+
+
+def test_enable_thinking_is_passed_as_chat_template_kwargs(monkeypatch):
+    """Without this the target answers in non-thinking mode and nothing says so."""
+
+    sent = _fake_requests(monkeypatch, [(200, "")])
+    request = {"prompt": "hi", "image_path": "https://example.com/a.jpg"}
+    answer, err = server_generate_vlm_sglang._openai_generate_one(
+        request,
+        "http://x/v1/chat/completions",
+        "m",
+        0.7,
+        2048,
+        60,
+        None,
+        None,
+        None,
+        False,
+        enable_thinking=True,
+    )
+
+    assert err is None and answer == "ok"
+    assert sent[0]["chat_template_kwargs"] == {"enable_thinking": True}
+
+
+def test_thinking_is_absent_unless_requested(monkeypatch):
+    sent = _fake_requests(monkeypatch, [(200, "")])
+    request = {"prompt": "hi", "image_path": "https://example.com/a.jpg"}
+    server_generate_vlm_sglang._openai_generate_one(
+        request, "http://x/v1/chat/completions", "m", 0.7, 2048, 60, None, None, None, False
+    )
+
+    assert "chat_template_kwargs" not in sent[0]
+
+
+def test_budget_shrinks_until_the_request_fits(monkeypatch):
+    """An over-budget request is retried smaller, not dropped.
+
+    The samples that overflow are the longest-context ones -- long videos -- so dropping
+    them biases the dataset against exactly the data a multimodal drafter needs.
+    """
+
+    overflow = (400, "This model's maximum context length is 32768 tokens.")
+    sent = _fake_requests(monkeypatch, [overflow, overflow, (200, "")])
+    request = {"prompt": "hi", "video_path": "https://example.com/a.mp4"}
+
+    answer, err = server_generate_vlm_sglang._openai_generate_one(
+        request,
+        "http://x/v1/chat/completions",
+        "m",
+        0.7,
+        8192,
+        60,
+        None,
+        None,
+        None,
+        False,
+        enable_thinking=False,
+        min_max_tokens=256,
+    )
+
+    assert err is None and answer == "ok"
+    assert [p["max_tokens"] for p in sent] == [8192, 4096, 2048]
+
+
+def test_budget_shrinking_stops_at_the_floor(monkeypatch):
+    """Below the floor a completion is too stunted to keep, so report instead of looping."""
+
+    overflow = (400, "This model's maximum context length is 32768 tokens.")
+    sent = _fake_requests(monkeypatch, [overflow, overflow])
+    request = {"prompt": "hi", "video_path": "https://example.com/a.mp4"}
+
+    answer, err = server_generate_vlm_sglang._openai_generate_one(
+        request,
+        "http://x/v1/chat/completions",
+        "m",
+        0.7,
+        512,
+        60,
+        None,
+        None,
+        None,
+        False,
+        enable_thinking=False,
+        min_max_tokens=256,
+    )
+
+    assert answer == "" and "maximum context length" in err
+    assert [p["max_tokens"] for p in sent] == [512, 256]
+
+
+def test_non_context_errors_are_not_retried(monkeypatch):
+    sent = _fake_requests(monkeypatch, [(400, "Invalid value for 'top_k'")])
+    request = {"prompt": "hi", "image_path": "https://example.com/a.jpg"}
+
+    answer, err = server_generate_vlm_sglang._openai_generate_one(
+        request, "http://x/v1/chat/completions", "m", 0.7, 8192, 60, None, None, None, False
+    )
+
+    assert answer == "" and "top_k" in err
+    assert len(sent) == 1