From 3c5bf1f6d2a7c747aaba9bd515509bd591c8c690 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 25 Jul 2026 02:25:16 -0700 Subject: [PATCH] [Spec] Derive NGRAM grammar tree links on the host instead of reading back `retrive_next_token` (#32380) --- python/sglang/srt/speculative/ngram_worker.py | 57 +++++++++++++++---- test/registered/spec/test_spec_ngram.py | 14 ++++- 2 files changed, 58 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 6b68cd103..50f95c9e4 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -38,6 +38,34 @@ logger = logging.getLogger(__name__) USE_FULL_MASK = True +def _derive_tree_links( + mask: np.ndarray, bs: int, draft_token_num: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Host-side (retrive_next_token, retrive_next_sibling), matching what + reconstruct_indices_from_tree_mask produces on device. + + ``mask[b, i, j]`` marks node j as an ancestor of node i, so i's immediate parent + is the largest such j < i, and both links follow from the parents alone. + """ + tree = mask.reshape(bs, draft_token_num, draft_token_num) + node_order = np.arange(draft_token_num) + ancestors = tree & (node_order < node_order[:, None]) + parents = np.where(ancestors.any(-1), (ancestors * node_order).argmax(-1), -1) + + next_token = np.full((bs, draft_token_num), -1, dtype=np.int64) + next_sibling = np.full((bs, draft_token_num), -1, dtype=np.int64) + for b in range(bs): + # Descending scan, so every k > i is already recorded when i is reached. + earliest_child_of = {} + for i in reversed(range(draft_token_num)): + next_token[b, i] = earliest_child_of.get(i, -1) + parent = int(parents[b, i]) + if parent >= 0: + next_sibling[b, i] = earliest_child_of.get(parent, -1) + earliest_child_of[parent] = i + return torch.from_numpy(next_token), torch.from_numpy(next_sibling) + + class NGRAMWorker(BaseSpecWorker): def alloc_memory_pool(self, **kwargs): # The target memory pool does not exist yet when __init__ runs. @@ -74,6 +102,7 @@ class NGRAMWorker(BaseSpecWorker): # rids of the last decode batch; used to erase corpus match state for # requests that left the batch (see forward_batch_generation). self._prev_decode_rids: set = set() + self.grammar_tree_host: Optional[tuple] = None self.ngram_corpus = NgramCorpus( min_bfs_breadth=server_args.speculative_ngram_min_bfs_breadth, @@ -280,6 +309,9 @@ class NGRAMWorker(BaseSpecWorker): tree_mask.copy_(torch.from_numpy(mask), non_blocking=True) draft_tokens.copy_(torch.from_numpy(req_drafts), non_blocking=True) + # Staged for the grammar bitmask, derived after the verify launch below. + self.grammar_tree_host = (mask, req_drafts) if batch.has_grammar else None + # generate positions and some indices using tree_mask reconstruct_indices_from_tree_mask( tree_mask, @@ -380,14 +412,6 @@ class NGRAMWorker(BaseSpecWorker): accept_lens = torch.ones(bs, dtype=torch.int32, device=self.device) if batch.forward_mode.is_target_verify(): - # Prepare grammar data on CPU if needed - if batch.has_grammar: - retrieve_next_token_cpu = verify_input.retrieve_next_token.cpu() - retrieve_next_sibling_cpu = verify_input.retrieve_next_sibling.cpu() - draft_tokens_cpu = verify_input.draft_token.view( - verify_input.retrieve_next_token.shape - ).cpu() - batch_result = self.target_worker.forward_batch_generation( batch, is_verify=True ) @@ -400,8 +424,15 @@ class NGRAMWorker(BaseSpecWorker): verify_input: NgramVerifyInput = batch.spec_info vocab_mask = None if batch.has_grammar: - # Generate the logit mask for structured output. - # Overlap the CPU operations for bitmask generation with the forward pass. + # From the host tree rather than the device output: no readback to + # wait on, and deriving here keeps it under the verify forward. + mask, req_drafts = self.grammar_tree_host + retrieve_next_token_cpu, retrieve_next_sibling_cpu = _derive_tree_links( + mask, bs, self.draft_token_num + ) + draft_tokens_cpu = ( + torch.from_numpy(req_drafts).to(torch.int64).view(bs, -1) + ) vocab_mask = generate_token_bitmask( batch.reqs, verify_input, @@ -413,7 +444,11 @@ class NGRAMWorker(BaseSpecWorker): if vocab_mask is not None: assert verify_input.grammar is not None - vocab_mask = vocab_mask.to(verify_input.retrieve_next_token.device) + # non_blocking is safe: the bitmask source is pinned, and stream + # order keeps the copy ahead of apply_vocab_mask. + vocab_mask = vocab_mask.to( + verify_input.retrieve_next_token.device, non_blocking=True + ) # NOTE (sk): otherwise, this vocab mask will be the one from the previous extend stage # and will be applied to produce wrong results batch.sampling_info.vocab_mask = None diff --git a/test/registered/spec/test_spec_ngram.py b/test/registered/spec/test_spec_ngram.py index 24f360afc..4a96166d7 100644 --- a/test/registered/spec/test_spec_ngram.py +++ b/test/registered/spec/test_spec_ngram.py @@ -2,16 +2,26 @@ import unittest from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.kits.eval_accuracy_kit import GSM8KMixin +from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin +from sglang.test.kits.regex_constrained_kit import RegexConstrainedMixin from sglang.test.kits.spec_server_kits import SpecLogprobKit from sglang.test.server_fixtures.ngram_fixture import NgramServerBase # Per-commit: Paged backend only. # - FA3 base test archived to test/manual/spec/test_spec_ngram_fa3.py # - Triton + Flashinfer moved to test_spec_ngram_extra.py -register_cuda_ci(est_time=400, stage="base-b", runner_config="1-gpu-large") +register_cuda_ci(est_time=460, stage="base-b", runner_config="1-gpu-large") -class TestNgramSpeculativeDecodingPaged(NgramServerBase, GSM8KMixin, SpecLogprobKit): +class TestNgramSpeculativeDecodingPaged( + NgramServerBase, + GSM8KMixin, + SpecLogprobKit, + RegexConstrainedMixin, + JSONConstrainedMixin, +): + # Constrained mixins reuse this server; they cover the grammar verify path, + # where the bitmask is built by walking the host draft tree. attention_backend = "flashinfer" extra_args = ["--page-size", "64"]