From f99ae52ca229228706fb24821d41596861feba36 Mon Sep 17 00:00:00 2001
From: Alexazhu <alexazzjjtt@163.com>
Date: Mon, 20 Jul 2026 01:06:36 +0800
Subject: [PATCH] fix: truncate xgrammar accepted_tokens in place on rollback
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

rollback(k) rebuilt the whole accepted-token history with a slice-copy
(accepted_tokens[:-k]), which is O(output_len) per call — and the EAGLE
spec-decode tree traversal calls rollback(1) once per draft-tree node on
the scheduler hot path. The slice form also made rollback(0) clear the
entire list ([:-0] == [:0]). Use guarded in-place deletion instead:
O(k) work per call and a safe no-op for k=0.

Fixes #31711
---
 .../srt/constrained/xgrammar_backend.py       |  5 +-
 .../unit/constrained/test_xgrammar_backend.py | 66 +++++++++++++++++++
 2 files changed, 70 insertions(+), 1 deletion(-)
 create mode 100644 test/registered/unit/constrained/test_xgrammar_backend.py

diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py
index 19fdc4bbf3ac..20c97e42bfd9 100644
--- a/python/sglang/srt/constrained/xgrammar_backend.py
+++ b/python/sglang/srt/constrained/xgrammar_backend.py
@@ -92,7 +92,10 @@ def accept_token(self, token: int):
 
     def rollback(self, k: int):
         self.matcher.rollback(k)
-        self.accepted_tokens = self.accepted_tokens[:-k]
+        # Truncate in place: a slice-copy (`[:-k]`) is O(len) per call on the
+        # spec-decode hot path, and `[:-0]` would clear the whole list.
+        if k > 0:
+            del self.accepted_tokens[-k:]
 
     def is_terminated(self):
         return self.matcher.is_terminated()
diff --git a/test/registered/unit/constrained/test_xgrammar_backend.py b/test/registered/unit/constrained/test_xgrammar_backend.py
new file mode 100644
index 000000000000..4c6733f46b88
--- /dev/null
+++ b/test/registered/unit/constrained/test_xgrammar_backend.py
@@ -0,0 +1,66 @@
+"""
+Unit tests for sglang.srt.constrained.xgrammar_backend.
+
+Test Coverage:
+- XGrammarGrammar.rollback: drops exactly the last k accepted tokens,
+  k=0 is a no-op, truncation happens in place (regression for #31711 —
+  the old slice-copy was O(output_len) per call on the EAGLE spec-decode
+  hot path and `[:-0]` cleared the whole token history).
+
+Usage:
+    python -m pytest test_xgrammar_backend.py -v
+"""
+
+import unittest
+from unittest.mock import MagicMock
+
+from sglang.srt.constrained.xgrammar_backend import XGrammarGrammar
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(2.0, "base-a-test-cpu")
+
+
+def _make_grammar(tokens):
+    """Build an XGrammarGrammar with a mocked matcher and accept `tokens`."""
+    matcher = MagicMock()
+    matcher.is_terminated.return_value = False
+    matcher.accept_token.return_value = True
+    grammar = XGrammarGrammar(
+        matcher=matcher,
+        vocab_size=32000,
+        ctx=MagicMock(),
+        override_stop_tokens=None,
+        key_string="test",
+    )
+    for token in tokens:
+        grammar.accept_token(token)
+    return grammar
+
+
+class TestXGrammarGrammarRollback(unittest.TestCase):
+    """Test XGrammarGrammar.rollback token-history bookkeeping (#31711)."""
+
+    def test_rollback_drops_last_k_tokens(self):
+        grammar = _make_grammar([1, 2, 3, 4, 5])
+        grammar.rollback(2)
+        self.assertEqual(grammar.accepted_tokens, [1, 2, 3])
+        grammar.matcher.rollback.assert_called_once_with(2)
+
+    def test_rollback_zero_is_noop(self):
+        """rollback(0) must keep the history: `[:-0]` used to clear it."""
+        grammar = _make_grammar([1, 2, 3])
+        grammar.rollback(0)
+        self.assertEqual(grammar.accepted_tokens, [1, 2, 3])
+
+    def test_rollback_truncates_in_place(self):
+        """The spec-decode tree traversal calls rollback(1) per draft-tree
+        node; the history must be truncated in place, not slice-copied."""
+        grammar = _make_grammar([1, 2, 3, 4])
+        tokens_before = grammar.accepted_tokens
+        grammar.rollback(1)
+        self.assertIs(grammar.accepted_tokens, tokens_before)
+        self.assertEqual(tokens_before, [1, 2, 3])
+
+
+if __name__ == "__main__":
+    unittest.main()