sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc0__py3-none-any.whl

sglang/srt/constrained/base_grammar_backend.py

@@ -14,8 +14,9 @@
 """The baseclass of a backend for grammar-guided constrained decoding."""
 
 import logging
+import time
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from threading import Event
 from typing import Dict, List, Optional, Tuple
 
@@ -26,10 +27,22 @@ from sglang.srt.server_args import ServerArgs
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class GrammarStats:
+    compilation_time: Optional[float] = None
+    schema_count: Optional[int] = None
+    ebnf_size: Optional[int] = None
+    is_cache_hit: bool = False
+    is_grammar_aborted: bool = False
+    tree_traversal_time: List[float] = field(default_factory=list)
+
+
 class BaseGrammarObject:
 
     def __init__(self):
         self._finished = False
+        self.grammar_stats = None
+        self.current_token = None
 
     def accept_token(self, token: int) -> None:
         """
@@ -137,19 +150,26 @@ class BaseGrammarBackend:
         return self._not_supported("structural_tag", key_string)
 
     def _init_value_dispatch(self, key: Tuple[str, str]) -> Optional[BaseGrammarObject]:
+        s = time.perf_counter()
         key_type, key_string = key
         if key_type == "json":
-            return self.dispatch_json(key_string)
+            grammar = self.dispatch_json(key_string)
         elif key_type == "regex":
-            return self.dispatch_regex(key_string)
+            grammar = self.dispatch_regex(key_string)
         elif key_type == "ebnf":
-            return self.dispatch_ebnf(key_string)
+            grammar = self.dispatch_ebnf(key_string)
         elif key_type == "structural_tag":
-            return self.dispatch_structural_tag(key_string)
+            grammar = self.dispatch_structural_tag(key_string)
         elif key_type == "structural_pattern":
-            return self.dispatch_structural_pattern(key_string)
+            grammar = self.dispatch_structural_pattern(key_string)
+        elif key_type == "structural_pattern_v2":
+            grammar = self.dispatch_structural_pattern_v2(key_string)
         else:
-            return self.dispatch_fallback(key_type, key_string)
+            grammar = self.dispatch_fallback(key_type, key_string)
+
+        if grammar is not None and grammar.grammar_stats is not None:
+            grammar.grammar_stats.compilation_time = time.perf_counter() - s
+        return grammar
 
     def get_cached_or_future_value(
         self, key: Tuple[str, str]
@@ -167,20 +187,36 @@ class BaseGrammarBackend:
         self.cache.clear()
 
 
+GRAMMAR_BACKEND_REGISTRY = {}
+
+
+def register_grammar_backend(name, init_func):
+    GRAMMAR_BACKEND_REGISTRY[name] = init_func
+
+
 def create_grammar_backend(
     server_args: ServerArgs,
     tokenizer,
     vocab_size: int,
     eos_token_ids: Optional[set] = None,
 ) -> Optional[BaseGrammarBackend]:
-    if server_args.grammar_backend == "outlines":
+    name = server_args.grammar_backend
+
+    # Custom grammar backend has the highest priority
+    if name in GRAMMAR_BACKEND_REGISTRY:
+        return GRAMMAR_BACKEND_REGISTRY[name](
+            server_args, tokenizer, vocab_size, eos_token_ids
+        )
+
+    # Default grammar backends
+    if name == "outlines":
         from sglang.srt.constrained.outlines_backend import OutlinesGrammarBackend
 
         grammar_backend = OutlinesGrammarBackend(
             tokenizer,
             whitespace_pattern=server_args.constrained_json_whitespace_pattern,
         )
-    elif server_args.grammar_backend == "xgrammar":
+    elif name == "xgrammar":
         from sglang.srt.constrained.xgrammar_backend import XGrammarGrammarBackend
 
         # Convert Set[int] to List[int] if needed
@@ -189,17 +225,17 @@ def create_grammar_backend(
         grammar_backend = XGrammarGrammarBackend(
             tokenizer, vocab_size=vocab_size, model_eos_token_ids=eos_list
         )
-    elif server_args.grammar_backend == "llguidance":
+    elif name == "llguidance":
         from sglang.srt.constrained.llguidance_backend import GuidanceBackend
 
         grammar_backend = GuidanceBackend(
             tokenizer=tokenizer,
             whitespace_pattern=server_args.constrained_json_whitespace_pattern,
         )
-    elif server_args.grammar_backend == "none":
+    elif name == "none":
         return None
     else:
-        raise ValueError(f"Invalid grammar backend: {server_args.grammar_backend}")
+        raise ValueError(f"Invalid grammar backend: {name}")
 
     if server_args.reasoning_parser and hasattr(tokenizer, "think_end_id"):
         from sglang.srt.constrained.reasoner_grammar_backend import (

sglang/srt/constrained/llguidance_backend.py

@@ -48,7 +48,6 @@ class GuidanceGrammar(BaseGrammarObject):
             self.serialized_grammar,
             log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")),
         )
-        self.finished = False
         self.bitmask = None
 
     def accept_token(self, token: int):

sglang/srt/constrained/outlines_backend.py

@@ -49,7 +49,6 @@ class OutlinesGrammar(BaseGrammarObject):
         self.guide = guide
         self.jump_forward_map = jump_forward_map
         self.state = 0
-        self.finished = False
 
     def accept_token(self, token: int):
         self.state = self.guide.get_next_state(self.state, token)

sglang/srt/constrained/xgrammar_backend.py

@@ -13,6 +13,7 @@
 # ==============================================================================
 """Constrained decoding with xgrammar backend."""
 
+import dataclasses
 import json
 import logging
 from typing import List, Optional, Tuple, Union
@@ -31,6 +32,7 @@ from sglang.srt.constrained.base_grammar_backend import (
     INVALID_GRAMMAR_OBJ,
     BaseGrammarBackend,
     BaseGrammarObject,
+    GrammarStats,
 )
 from sglang.srt.utils import is_hip
 
@@ -41,9 +43,9 @@ else:
     from sglang.srt.constrained.triton_ops.bitmask_ops import (
         apply_token_bitmask_inplace_triton,
     )
-logger = logging.getLogger(__name__)
 
 
+logger = logging.getLogger(__name__)
 MAX_ROLLBACK_TOKENS = 200
 
 
@@ -56,17 +58,20 @@ class XGrammarGrammar(BaseGrammarObject):
         ctx: CompiledGrammar,
         override_stop_tokens: Optional[Union[List[int], int]],
         key_string: Optional[str] = None,  # TODO (sk): for debugging, remove later
+        grammar_stats: Optional[GrammarStats] = GrammarStats(),
     ) -> None:
+        super().__init__()
         self.matcher = matcher
         self.vocab_size = vocab_size
         self.ctx = ctx
         self.override_stop_tokens = override_stop_tokens
-        self.finished = False
         self.accepted_tokens = []
         self.key_string = key_string
+        self.grammar_stats = grammar_stats
 
     def accept_token(self, token: int):
         if not self.is_terminated():
+            self.current_token = token
             accepted = self.matcher.accept_token(token)
             if not accepted:
                 # log for debugging
@@ -120,6 +125,9 @@ class XGrammarGrammar(BaseGrammarObject):
             self.ctx,
             self.override_stop_tokens,
             self.key_string,
+            dataclasses.replace(
+                self.grammar_stats, is_cache_hit=True, tree_traversal_time=[]
+            ),
         )
 
     def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
@@ -150,7 +158,7 @@ class XGrammarGrammar(BaseGrammarObject):
             assert self.matcher.accept_token(new_output_ids[i])
 
     def __repr__(self):
-        return f"XGrammarGrammar({self.key_string=}, {self.accepted_tokens=})"
+        return f"XGrammarGrammar({self.key_string=}, {self.accepted_tokens=}, {self.current_token=})"
 
 
 class XGrammarGrammarBackend(BaseGrammarBackend):
@@ -165,6 +173,10 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         if hasattr(tokenizer, "init_xgrammar"):
             # For special tokenizer
             tokenizer_info, override_stop_tokens = tokenizer.init_xgrammar()
+
+            if tokenizer_info is None:
+                # Not supported tokenizer
+                return
         else:
             # Create TokenizerInfo with model's EOS tokens as the authoritative stop tokens
             # This ensures consistency between what the model considers EOS and what XGrammar uses
@@ -177,14 +189,21 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         self.vocab_size = vocab_size
         self.override_stop_tokens = override_stop_tokens
 
-    def _from_context(self, ctx: CompiledGrammar, key_string: str) -> XGrammarGrammar:
+    def _from_context(
+        self, ctx: CompiledGrammar, key_string: str, grammar_stats: GrammarStats
+    ) -> XGrammarGrammar:
         matcher = GrammarMatcher(
             ctx,
             max_rollback_tokens=MAX_ROLLBACK_TOKENS,
             override_stop_tokens=self.override_stop_tokens,
         )
         return XGrammarGrammar(
-            matcher, self.vocab_size, ctx, self.override_stop_tokens, key_string
+            matcher,
+            self.vocab_size,
+            ctx,
+            self.override_stop_tokens,
+            key_string,
+            grammar_stats,
         )
 
     def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]:
@@ -198,7 +217,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         except (RuntimeError, json.decoder.JSONDecodeError) as e:
             logging.error(f"Hit invalid json_schema: {key_string=}, {e=}")
             return INVALID_GRAMMAR_OBJ
-        return self._from_context(ctx, key_string)
+        return self._from_context(ctx, key_string, GrammarStats())
 
     def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -206,7 +225,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         except RuntimeError as e:
             logging.error(f"Hit invalid ebnf: {key_string=}, {e=}")
             return INVALID_GRAMMAR_OBJ
-        return self._from_context(ctx, key_string)
+        return self._from_context(ctx, key_string, GrammarStats())
 
     def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -214,7 +233,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         except RuntimeError as e:
             logging.error(f"Hit invalid regex: {key_string=}, {e=}")
             return INVALID_GRAMMAR_OBJ
-        return self._from_context(ctx, key_string)
+        return self._from_context(ctx, key_string, GrammarStats())
 
     def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -233,7 +252,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         except (RuntimeError, json.decoder.JSONDecodeError) as e:
             logging.error(f"Hit invalid structural_tag: {key_string=}, {e=}")
             return INVALID_GRAMMAR_OBJ
-        return self._from_context(ctx, key_string)
+        return self._from_context(ctx, key_string, GrammarStats())
 
     def reset(self):
         self.grammar_compiler.clear_cache()

sglang/srt/custom_op.py

@@ -1,12 +1,20 @@
 from torch import nn
 
-from sglang.srt.utils import cpu_has_amx_support, is_cpu, is_cuda, is_hip, is_npu
+from sglang.srt.utils import (
+    cpu_has_amx_support,
+    is_cpu,
+    is_cuda,
+    is_hip,
+    is_npu,
+    is_xpu,
+)
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
 _is_cpu = is_cpu()
 _is_cpu_amx_available = cpu_has_amx_support()
 _is_npu = is_npu()
+_is_xpu = is_xpu()
 
 
 class CustomOp(nn.Module):
@@ -88,5 +96,7 @@ class CustomOp(nn.Module):
             return self.forward_cpu
         elif _is_npu:
             return self.forward_npu
+        elif _is_xpu:
+            return self.forward_xpu
         else:
             return self.forward_native

sglang/srt/debug_utils/dump_comparator.py

@@ -1,11 +1,11 @@
 import argparse
 import functools
-import re
 from pathlib import Path
 
 import polars as pl
 import torch
 
+from sglang.srt.debug_utils.dump_loader import find_row, read_meta
 from sglang.srt.debug_utils.dumper import get_truncated_value
 
 
@@ -26,66 +26,77 @@ def main(args):
     print("df_baseline", df_baseline)
 
     for row in df_target.iter_rows(named=True):
-        rows_baseline = df_baseline.filter(
-            (
-                pl.col("forward_pass_id")
-                == row["forward_pass_id"] - args.start_id + args.baseline_start_id
-            )
-            & functools.reduce(
-                lambda a, b: a & b,
-                [
-                    pl.col(col) == row[col]
-                    for col in row.keys()
-                    if col not in ["forward_pass_id", "dump_index", "filename"]
-                ],
-            )
+        path_target = Path(args.target_path) / row["filename"]
+
+        row_baseline = find_row(
+            df_baseline,
+            conditions=dict(
+                forward_pass_id=row["forward_pass_id"]
+                - args.start_id
+                + args.baseline_start_id,
+                **{
+                    k: v
+                    for k, v in row.items()
+                    if k not in ["forward_pass_id", "dump_index", "filename"]
+                },
+            ),
         )
-        assert len(rows_baseline) == 1, f"{rows_baseline=}"
-        row_baseline = rows_baseline.to_dicts()[0]
+
+        if row_baseline is None:
+            print(f"Skip: target={str(path_target)} since no baseline")
+            x_target = _load_object(path_target)
+            if x_target is not None:
+                print(f"x_target(sample)={get_truncated_value(x_target)}")
+            continue
 
         path_baseline = Path(args.baseline_path) / row_baseline["filename"]
-        path_target = Path(args.target_path) / row["filename"]
         print(f"Check: target={str(path_target)} baseline={str(path_baseline)}")
-        check_tensor_pair(path_baseline=path_baseline, path_target=path_target)
+        check_tensor_pair(
+            path_baseline=path_baseline, path_target=path_target, name=row["name"]
+        )
         print()
 
 
-def read_meta(directory):
-    directory = Path(directory)
-    assert directory.is_dir(), f"{directory=} should be a directory"
-
-    rows = []
-    for p in directory.glob("*.pt"):
-        full_kwargs = {}
-        for kv in p.stem.split("___"):
-            k, v = kv.split("=")
-            full_kwargs[k] = v
-        rows.append(
-            {
-                "filename": str(p.name),
-                **full_kwargs,
-            }
-        )
+def check_tensor_pair(path_baseline, path_target, name=""):
+    x_baseline = _load_object(path_baseline)
+    x_target = _load_object(path_target)
 
-    df = pl.DataFrame(rows)
-    df = df.with_columns(
-        pl.col("forward_pass_id").cast(int),
-        pl.col("rank").cast(int),
+    print(
+        f"Raw "
+        f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
+        f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
     )
-    return df
-
 
-def check_tensor_pair(path_baseline, path_target):
-    x_baseline = torch.load(path_baseline, weights_only=True)
-    x_target = torch.load(path_target, weights_only=True)
+    x_baseline, x_target = _comparison_preprocessor(x_baseline, x_target, name=name)
+    x_baseline = _try_unify_shape(x_baseline, target_shape=x_target.shape)
 
     print(
+        f"After preprocessor "
         f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
         f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
     )
 
+    x_target = x_target.float()
+    x_baseline = x_baseline.float()
+
+    for name, fn in (
+        ("mean", torch.mean),
+        ("std", torch.std),
+        ("min", torch.min),
+        ("max", torch.max),
+        ("p1", functools.partial(torch.quantile, q=0.01)),
+        ("p5", functools.partial(torch.quantile, q=0.05)),
+        ("p95", functools.partial(torch.quantile, q=0.95)),
+        ("p99", functools.partial(torch.quantile, q=0.99)),
+    ):
+        value_baseline = fn(x_baseline).item()
+        value_target = fn(x_target).item()
+        print(
+            f"[{name}] {value_baseline :.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})"
+        )
+
     if x_baseline.shape != x_target.shape:
-        print(f"❌ Shape mismatch")
+        print(f"⚠️ Shape mismatch")
         return
 
     raw_abs_diff = (x_target - x_baseline).abs()
@@ -112,6 +123,19 @@ def check_tensor_pair(path_baseline, path_target):
         print(f"x_target(sample)={get_truncated_value(x_target)}")
 
 
+def _try_unify_shape(x: torch.Tensor, target_shape):
+    x_shape = x.shape
+    num_dim_to_remove = len(x_shape) - len(target_shape)
+    if (x_shape[num_dim_to_remove:] == target_shape) and all(
+        val == 1 for val in x_shape[:num_dim_to_remove]
+    ):
+        out = functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x)
+        print(f"Unify shape: {x_shape} -> {out.shape} (to match {target_shape})")
+        return out
+
+    return x
+
+
 # Copied from DeepGEMM
 def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
     x, y = x.double(), y.double()
@@ -120,6 +144,19 @@ def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
     return 1 - sim
 
 
+def _comparison_preprocessor(x_baseline, x_target, name):
+    # can insert arbitrary adhoc postprocessing logic here
+    return x_baseline, x_target
+
+
+def _load_object(path):
+    x = torch.load(path, weights_only=False)
+    if not isinstance(x, torch.Tensor):
+        print(f"Skip load {path} since {type(x)=} is not a Tensor")
+        return None
+    return x.cuda()
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--baseline-path", type=str)

sglang/srt/debug_utils/dump_loader.py

@@ -0,0 +1,97 @@
+import functools
+import os
+from pathlib import Path
+from typing import Any, Dict
+
+import polars as pl
+import torch
+
+
+class DumpLoader:
+    def __init__(self):
+        directory = os.environ.get("SGLANG_DUMP_LOADER_DIR")
+
+        self._enable = directory is not None
+        if self._enable:
+            self._directory = Path(directory)
+            self._df = read_meta(directory)
+
+    @property
+    def enable(self):
+        return self._enable
+
+    def load(self, name, **kwargs):
+        assert self._enable, "Please call DumpLoader.load only when it is enabled"
+
+        from sglang.srt.debug_utils.dumper import dumper
+
+        forward_pass_id = dumper._forward_pass_id
+        conditions = dict(name=name, forward_pass_id=forward_pass_id, **kwargs)
+        row = find_row(self._df, conditions=conditions)
+        assert (
+            row is not None
+        ), f"DumpLoader cannot find row given query {name=} {kwargs=} {self._directory=}"
+
+        path = self._directory / row["filename"]
+        output = torch.load(path, weights_only=False)
+
+        print(
+            f"[DumpLoader] load from {path=} (query: {name=} {kwargs=}, output: {type(output)})"
+        )
+        return output
+
+
+def read_meta(directory):
+    directory = Path(directory)
+    assert directory.is_dir(), f"{directory=} should be a directory"
+
+    rows = []
+    for p in directory.glob("*.pt"):
+        full_kwargs = {}
+        for kv in p.stem.split("___"):
+            k, v = kv.split("=")
+            full_kwargs[k] = v
+        rows.append(
+            {
+                "filename": str(p.name),
+                **full_kwargs,
+            }
+        )
+
+    df = pl.DataFrame(rows)
+    df = df.with_columns(
+        pl.col("forward_pass_id").cast(int),
+        pl.col("rank").cast(int),
+        pl.col("dump_index").cast(int),
+    )
+    return df
+
+
+def find_row(df, conditions: Dict[str, Any]):
+    df_sub = df.filter(
+        functools.reduce(
+            lambda a, b: a & b,
+            [
+                pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col])
+                for col in conditions.keys()
+            ],
+        )
+    )
+    assert len(df_sub) <= 1
+    return df_sub.to_dicts()[0] if len(df_sub) > 0 else None
+
+
+def _cast_to_polars_dtype(value, target_dtype):
+    if target_dtype in (pl.Int64, pl.Int32, pl.UInt64, pl.UInt32):
+        return int(value)
+    elif target_dtype in (pl.Float64, pl.Float32):
+        return float(value)
+    elif target_dtype == pl.Boolean:
+        return bool(value)
+    elif target_dtype == pl.String:
+        return str(value)
+    else:
+        return value
+
+
+dump_loader = DumpLoader()

sglang/srt/debug_utils/dumper.py

@@ -53,7 +53,7 @@ class _Dumper:
         if self._partial_name is None:
             self._partial_name = _get_partial_name()
 
-        rank = dist.get_rank()
+        rank = _get_rank()
         full_kwargs = dict(
             forward_pass_id=self._forward_pass_id,
             rank=rank,
@@ -80,12 +80,20 @@ class _Dumper:
 
 
 def _get_partial_name():
-    rank = dist.get_rank()
+    rank = _get_rank()
     object_list = [str(time.time()) if rank == 0 else None]
-    dist.broadcast_object_list(object_list, device="cuda")
+    if dist.is_initialized():
+        dist.broadcast_object_list(object_list, device="cuda")
     return object_list[0]
 
 
+def _get_rank():
+    if dist.is_initialized():
+        return dist.get_rank()
+    else:
+        return 0
+
+
 def get_truncated_value(value):
     if value is None:
         return None