fix: truncate xgrammar accepted_tokens in place on rollback by alexzhu0 · Pull Request #31728 · sgl-project/sglang

alexzhu0 · 2026-07-19T18:09:45Z

Motivation

XGrammarGrammar.rollback() slice-copies the entire retained accepted-token history on each rollback. The speculative-decoding tree traversal calls rollback(1), so the Python bookkeeping cost grows with output history. Also, [:-0] unexpectedly clears the history for a zero-token rollback.

Modifications

Keep matcher.rollback(k) first, then truncate Python history in place with if k > 0: del self.accepted_tokens[-k:].
No matcher, model-forward, kernel, or workflow changes; do not introduce an additional negative-input exception contract.
Add seven real-xgrammar regression tests using a tiny local digit vocabulary (no tokenizer/model downloads): partial rollback, zero rollback on empty/nonempty history, list identity, full rollback, accept/rollback/accept, and oversized rollback rejection with both histories preserved.
Tests use CustomTestCase, unittest.main(), and register_cpu_ci(2.0, "base-a-test-cpu").
Merge upstream snapshot dae126d510b786a2424684c17ca4e5e70d74e8f7 without rewriting the original PR history. This snapshot includes minimum base a5f07b1241fc2c2f38058850c377a35066485ad3.
Refreshed head: 567cb90c002db4b21b6ab8a767348161c0b6fd59.

Accuracy Tests

Local environment: Linux amd64 Docker emulated on Apple Silicon, Python 3.12.14, torch 2.13.0+cpu, xgrammar 0.2.1, pytest 9.1.1. These are actual production imports and real matcher calls, not substituted modules.

The same seven regression tests against the unpatched upstream snapshot: 4 failed, 3 passed. Failures reproduce zero rollback clearing history and history reallocation.
Refreshed candidate, CPU constrained tests:
```
python -m pytest test/registered/unit/constrained \
  --ignore=test/registered/unit/constrained/test_e2e_constrained_reasoning.py \
  -q -p no:cacheprovider
```
136 passed, 3 skipped, 20 subtests passed. The three skips require CUDA. The explicitly excluded E2E file launches a GPU model server; it was not run locally.
Direct CI-style entry point: python test/registered/unit/constrained/test_xgrammar_backend.py -v — 7 passed.
pre-commit run --files python/sglang/srt/constrained/xgrammar_backend.py test/registered/unit/constrained/test_xgrammar_backend.py — passed (exact repository hooks; full registered-test tree present for registry validation).
git diff --check — passed.

Oversized rollback is intentionally tested as rejection, not successful history clearing: xgrammar 0.2.1 checks saved history before modifying matcher state (native implementation).

This is focused CPU verification, not the full upstream CUDA CI result. Maintainer-enabled CI and formal review are still required; no approval or merge is claimed.

Speed Tests and Profiling

Python-bookkeeping-only microbenchmark of the real XGrammarGrammar.rollback(1) method. The native matcher is replaced only for this timing measurement with a no-op boundary, and list construction is outside the timed interval. Median of 21 samples per size, same Linux/Python environment above:

History length	Upstream median µs/op	Refreshed candidate median µs/op
10,000	13.583	0.417
100,000	126.459	0.458
1,000,000	2701.709	2.334

This isolates removal of the retained-history copy; it is not end-to-end decoding throughput, and does not claim constant cost for arbitrary k. Emulation, scheduling and memory effects affect the timings.

Reproduction script (run with each checkout on PYTHONPATH)

import statistics
from time import perf_counter_ns

from sglang.srt.constrained.xgrammar_backend import XGrammarGrammar


class BookkeepingOnlyMatcher:
    def rollback(self, k):
        assert k == 1


for history_len in (10_000, 100_000, 1_000_000):
    grammar = XGrammarGrammar(
        matcher=BookkeepingOnlyMatcher(),
        vocab_size=11,
        ctx=None,
        override_stop_tokens=None,
        key_string="bookkeeping-only-benchmark",
    )
    samples = []
    for _ in range(21):
        grammar.accepted_tokens = list(range(history_len))
        start = perf_counter_ns()
        grammar.rollback(1)
        samples.append(perf_counter_ns() - start)
        assert len(grammar.accepted_tokens) == history_len - 1
        assert grammar.accepted_tokens[-1] == history_len - 2
    print(history_len, statistics.median(samples) / 1_000)

Checklist

Run pre-commit on touched files.
Add registered regression tests with real xgrammar.
Provide fresh correctness and bounded microbenchmark evidence.
Follow the contribution guide's registered-test conventions.
Full upstream CI and maintainer review pending.

Documentation update is not applicable to this internal bookkeeping fix.

Disclosure: this fix and follow-up were prepared with AI assistance. Reported results are from actual local executions; upstream approval is not implied.

CI States

Latest PR Test (Base): ❌ Run #33856594474
Latest PR Test (Extra): ❌ Run #33856593952
Latest PR Test (AMD ROCm 7.2): ❌ Run #33856594395

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 sgl-project#31711

gemini-code-assist

Code Review

This pull request optimizes the rollback method in XGrammarGrammar to truncate the accepted tokens list in place rather than using a slice-copy, which avoids O(len) overhead and fixes a bug where k=0 cleared the entire history. It also introduces a comprehensive unit test suite for this behavior. The reviewer recommended adding defensive validation to raise a ValueError if k is negative, along with corresponding unit tests to cover negative values of k and values of k larger than the history length.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-19T18:11:03Z

    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:]


It is a good practice to defensively guard against negative values of k to prevent unexpected behavior or cryptic errors from the underlying C++ matcher. Adding a check for k < 0 and raising a clear ValueError improves robustness.

Suggested change

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 rollback(self, k: int):

if k < 0:

raise ValueError(f"k must be non-negative, got {k}")

self.matcher.rollback(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:]

Checked against the repository-pinned xgrammar 0.2.1 before adding validation. Its native rollback checks the upper history bound, and a negative count currently returns without consuming history (verified by accepting one token, calling rollback(-1), then successfully rolling back that token). Introducing a Python ValueError would add a new input contract beyond this bookkeeping fix. I have kept matcher.rollback(k) first and left validation semantics at that boundary. The refreshed commit 567cb90 includes real-matcher regression tests, rather than using a mock to define dependency behavior. Source: https://github.com/mlc-ai/xgrammar/blob/v0.2.1/cpp/grammar_matcher.cc#L1023

gemini-code-assist · 2026-07-19T18:11:03Z

+        grammar = _make_grammar([1, 2, 3])
+        grammar.rollback(0)
+        self.assertEqual(grammar.accepted_tokens, [1, 2, 3])
+


To ensure the robustness of the rollback implementation, we should add unit tests covering edge cases such as negative values of k and values of k larger than the length of the accepted tokens history.

Suggested change

def test_rollback_negative_k_raises_error(self):

grammar = _make_grammar([1, 2, 3])

with self.assertRaises(ValueError):

grammar.rollback(-1)

def test_rollback_large_k_clears_history(self):

grammar = _make_grammar([1, 2, 3])

grammar.rollback(5)

self.assertEqual(grammar.accepted_tokens, [])

Expanded the tests in 567cb90 to use a real xgrammar 0.2.1 digit grammar and local vocabulary. Oversized rollback is rejected by the native matcher before state mutation, so a test expecting it to clear history would encode incorrect semantics. The new test asserts RuntimeError, unchanged Python list identity/contents, and a subsequent valid rollback to verify native history is intact. Added full/empty rollback and accept-rollback-accept coverage as well: seven real-matcher tests pass on the refreshed candidate; the same tests give 4 failures and 3 passes on unpatched main. Negative-input validation is intentionally not introduced as part of this narrow fix (see the companion reply). The PR body now has fresh CPU test, lint, and bounded benchmark evidence.

alexzhu0 · 2026-07-23T15:24:55Z

Per the PR Merge Process, pinging Merge Oncalls to start the process — the assignment bot doesn't seem to have picked this one up (no assignee, 4 days).

@hnyls2002 @Qiaolin-Yu (speculative decoding — the hot caller is traverse_tree's per-node rollback(1)), cc @ch-wan (recent constrained/ work).

One-line summary: XGrammarGrammar.rollback() slice-copies the whole accepted-token history per call — O(output_len) on the spec-decode scheduler path — replaced with guarded in-place del (constant-time, ~17,700× at 1M-token history; also fixes rollback(0) silently clearing the history). CPU-only reproducible benchmark + regression tests in the PR body. Could someone tag run-ci when convenient? Thanks!

alexzhu0 · 2026-09-04T08:37:53Z

Update after the branch refresh (2026-09-04): head is now 567cb90. Upstream main was merged without rewriting history; the required minimum base is now an ancestor, and GitHub reports no merge conflict.

Fresh local evidence is in the PR body: 7 real-xgrammar regression tests pass; the same tests fail 4/7 on unpatched main. The CPU constrained suite passes 136 tests (3 CUDA skips; GPU model E2E explicitly not run locally), and all applicable pre-commit hooks pass. The existing reviewer requests target the constrained CODEOWNERS. Both inline bot suggestions now have replies grounded in the pinned matcher implementation.

The newly triggered base CI gate still exits with Missing required label run-ci; base-a-test-cpu is skipped. This is a permission gate, not a regression-test failure. Full upstream tests and formal review remain pending.

Could a maintainer enable run-ci for the refreshed branch and review this small bookkeeping change? No workflow or matcher-validation changes were made. Thanks!

alexzhu0 · 2026-09-08T07:08:22Z

@DarkSharpness @hnyls2002 @JustinTong0323 — a brief follow-up on the September 4 refresh: could one of you review this rollback bookkeeping fix and enable run-ci when convenient?

Head 567cb90c0 has seven passing real-xgrammar regression tests, 136 passing local CPU constrained tests (3 CUDA skips; GPU E2E not run locally), and passing upstream lint. The base CI gate still requires the run-ci label. The PR body includes the reproduction and validation details. Thanks!

alexzhu0 requested review from DarkSharpness, JustinTong0323 and hnyls2002 as code owners July 19, 2026 18:09

gemini-code-assist Bot reviewed Jul 19, 2026

View reviewed changes

benmyles mentioned this pull request Jul 27, 2026

Harden DeepSeek V4 DSpark DP2 serving shiftup-ai/sglang#1

Draft

21 tasks

Merge upstream main and validate xgrammar rollback with real matcher

567cb90

Navigation Menu

Uh oh!

fix: truncate xgrammar accepted_tokens in place on rollback - #31728

fix: truncate xgrammar accepted_tokens in place on rollback#31728
alexzhu0 wants to merge 2 commits into
sgl-project:mainfrom
alexzhu0:fix/xgrammar-rollback-inplace-truncate

alexzhu0 commented Jul 19, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

gemini-code-assist Bot Jul 19, 2026

Uh oh!

alexzhu0 Sep 4, 2026

Uh oh!

gemini-code-assist Bot Jul 19, 2026

Uh oh!

alexzhu0 Sep 4, 2026

Uh oh!

alexzhu0 commented Jul 23, 2026

Uh oh!

alexzhu0 commented Sep 4, 2026 •

edited

Loading

Uh oh!

alexzhu0 commented Sep 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

+    def test_rollback_negative_k_raises_error(self):
+        grammar = _make_grammar([1, 2, 3])
+        with self.assertRaises(ValueError):
+            grammar.rollback(-1)
+    def test_rollback_large_k_clears_history(self):
+        grammar = _make_grammar([1, 2, 3])
+        grammar.rollback(5)
+        self.assertEqual(grammar.accepted_tokens, [])

Uh oh!

Conversation

alexzhu0 commented Jul 19, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 19, 2026

Choose a reason for hiding this comment

Uh oh!

alexzhu0 Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 19, 2026

Choose a reason for hiding this comment

Uh oh!

alexzhu0 Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

alexzhu0 commented Jul 23, 2026

Uh oh!

alexzhu0 commented Sep 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

alexzhu0 commented Sep 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

alexzhu0 commented Jul 19, 2026 •

edited by github-actions Bot

Loading

alexzhu0 commented Sep 4, 2026 •

edited

Loading