From 5f8b3819027fa4fe8a1d9c9b62402d60f50bc0f2 Mon Sep 17 00:00:00 2001
From: David <12414531+DavidBellamy@users.noreply.github.com>
Date: Fri, 5 Jun 2026 08:37:03 -0700
Subject: [PATCH] Add ForcedSequenceLogitProcessor (teacher-forcing custom
 logit processor)

---
 .../srt/sampling/custom_logit_processor.py    | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/python/sglang/srt/sampling/custom_logit_processor.py b/python/sglang/srt/sampling/custom_logit_processor.py
index d58a5f6cf149..ac0639c7d2f6 100644
--- a/python/sglang/srt/sampling/custom_logit_processor.py
+++ b/python/sglang/srt/sampling/custom_logit_processor.py
@@ -200,3 +200,31 @@ def __call__(
             logits[batch_idx, indices] = -float("inf")
 
         return logits
+
+
+class ForcedSequenceLogitProcessor(CustomLogitProcessor):
+    """Teacher-forcing: pin each decode step to a per-request recorded token sequence.
+
+    Per-request custom_params: {"forced_output_ids": [int, ...]}. At decode step t
+    (t = tokens already generated for the req), forces the next token to
+    forced_output_ids[t] by masking every other logit. Lets two servers replay an
+    IDENTICAL recorded trajectory so the only variable is a server flag (e.g.
+    --no-cache-thoughts); the reasoning parser still sees the forced </think>, so the
+    split fires as in a live generation. Past the end of the list it samples normally.
+    """
+
+    def __call__(self, logits, custom_param_list=None):
+        if not custom_param_list:
+            return logits
+        for i, params in enumerate(custom_param_list):
+            if not params:
+                continue
+            forced = params.get("forced_output_ids")
+            if not forced:
+                continue
+            req = params.get("__req__")
+            t = len(req.output_ids) if req is not None else 0
+            if t < len(forced):
+                logits[i, :] = float("-inf")
+                logits[i, forced[t]] = 0.0
+        return logits