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 , 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