sglang 0.4.6.post3__py3-none-any.whl → 0.4.6.post4__py3-none-any.whl

sglang/srt/constrained/reasoner_grammar_backend.py

@@ -13,7 +13,6 @@
 # ==============================================================================
 """The baseclass of a backend for reasoner grammar-guided constrained decoding."""
 
-from concurrent.futures import Future
 from typing import List, Optional, Tuple
 
 import torch
@@ -28,13 +27,12 @@ class ReasonerGrammarObject(BaseGrammarObject):
         self.think_end_id = think_end_id
         self.is_in_reasoning = True
 
-    @property
-    def finished(self):
-        return self.grammar.finished
+    def accept_token(self, token: int):
+        if token == self.think_end_id:
+            self.is_in_reasoning = False
 
-    @finished.setter
-    def finished(self, finished):
-        self.grammar.finished = finished
+        if not self.is_in_reasoning and token != self.think_end_id:
+            self.grammar.accept_token(token)
 
     def allocate_vocab_mask(
         self, vocab_size: int, batch_size: int, device
@@ -52,12 +50,16 @@ class ReasonerGrammarObject(BaseGrammarObject):
     def apply_vocab_mask(self):
         return self.grammar.apply_vocab_mask
 
-    def accept_token(self, token: int):
-        if token == self.think_end_id:
-            self.is_in_reasoning = False
+    def copy(self) -> BaseGrammarObject:
+        return ReasonerGrammarObject(self.grammar.copy(), self.think_end_id)
 
-        if not self.is_in_reasoning and token != self.think_end_id:
-            self.grammar.accept_token(token)
+    @property
+    def finished(self):
+        return self.grammar.finished
+
+    @finished.setter
+    def finished(self, finished):
+        self.grammar.finished = finished
 
     def try_jump_forward(self, tokenizer):
         return self.grammar.try_jump_forward(tokenizer)
@@ -72,30 +74,17 @@ class ReasonerGrammarObject(BaseGrammarObject):
             old_output_ids, new_output_ids, next_state
         )
 
-    def copy(self) -> BaseGrammarObject:
-        return ReasonerGrammarObject(self.grammar.copy(), self.think_end_id)
-
 
 class ReasonerGrammarBackend(BaseGrammarBackend):
     def __init__(self, grammar_backend: BaseGrammarBackend, think_end_id):
+        super().__init__()
         self.grammar_backend = grammar_backend
         self.think_end_id = think_end_id
 
-    def get_cached_value(self, key: Tuple[str, str]) -> Optional[ReasonerGrammarObject]:
-        grammar = self.grammar_backend.get_cached_value(key)
-        return ReasonerGrammarObject(grammar, self.think_end_id) if grammar else None
-
-    def get_future_value(self, key: Tuple[str, str]) -> Future:
-        grammar = Future()
-
-        def callback(f: Future):
-            if result := f.result():
-                grammar.set_result(ReasonerGrammarObject(result, self.think_end_id))
-            else:
-                grammar.set_result(None)
-
-        self.grammar_backend.get_future_value(key).add_done_callback(callback)
-        return grammar
-
-    def reset(self):
-        self.grammar_backend.reset()
+    def _init_value_dispatch(
+        self, key: Tuple[str, str]
+    ) -> Optional[ReasonerGrammarObject]:
+        ret = self.grammar_backend._init_value_dispatch(key)
+        if ret is None:
+            return None
+        return ReasonerGrammarObject(ret, self.think_end_id)

sglang/srt/constrained/xgrammar_backend.py

@@ -18,7 +18,6 @@ import logging
 from typing import List, Optional, Tuple, Union
 
 import torch
-import xgrammar
 from xgrammar import (
     CompiledGrammar,
     GrammarCompiler,
@@ -35,7 +34,6 @@ from sglang.srt.constrained.base_grammar_backend import (
 from sglang.srt.constrained.triton_ops.bitmask_ops import (
     apply_token_bitmask_inplace_triton,
 )
-from sglang.srt.utils import get_bool_env_var
 
 logger = logging.getLogger(__name__)
 
@@ -51,49 +49,35 @@ class XGrammarGrammar(BaseGrammarObject):
         vocab_size: int,
         ctx: CompiledGrammar,
         override_stop_tokens: Optional[Union[List[int], int]],
+        key_string: Optional[str] = None,  # TODO (sk): for debugging, remove later
     ) -> None:
-        super().__init__()
         self.matcher = matcher
         self.vocab_size = vocab_size
         self.ctx = ctx
         self.override_stop_tokens = override_stop_tokens
         self.finished = False
-
-        from xgrammar.kernels.apply_token_bitmask_inplace_cpu import (
-            apply_token_bitmask_inplace_cpu,
-        )
-
-        self.apply_vocab_mask_cpu = apply_token_bitmask_inplace_cpu
+        self.accepted_tokens = []
+        self.key_string = key_string
 
     def accept_token(self, token: int):
-        assert self.matcher.accept_token(token)
-
-    def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
-        s = self.matcher.find_jump_forward_string()
-        if s:
-            return [], s
-        return None
-
-    def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]:
-        _, data = helper
-        return data, -1
-
-    def jump_and_retokenize(
-        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
-    ):
-        k = 0
-        for i, old_id in enumerate(old_output_ids):
-            if old_id == new_output_ids[i]:
-                k = i + 1
+        if not self.is_terminated():
+            accepted = self.matcher.accept_token(token)
+            if not accepted:
+                # log for debugging
+                raise ValueError(
+                    f"Tokens not accepted: {token}\n"
+                    f"Accepted tokens: {self.accepted_tokens}\n"
+                    f"Key string: {self.key_string}"
+                )
             else:
-                break
+                self.accepted_tokens.append(token)
 
-        # rollback to the last token that is the same
-        if k < len(old_output_ids):
-            self.matcher.rollback(len(old_output_ids) - k)
+    def rollback(self, k: int):
+        self.matcher.rollback(k)
+        self.accepted_tokens = self.accepted_tokens[:-k]
 
-        for i in range(k, len(new_output_ids)):
-            assert self.matcher.accept_token(new_output_ids[i])
+    def is_terminated(self):
+        return self.matcher.is_terminated()
 
     def allocate_vocab_mask(
         self, vocab_size: int, batch_size: int, device
@@ -122,9 +106,43 @@ class XGrammarGrammar(BaseGrammarObject):
             override_stop_tokens=self.override_stop_tokens,
         )
         return XGrammarGrammar(
-            matcher, self.vocab_size, self.ctx, self.override_stop_tokens
+            matcher,
+            self.vocab_size,
+            self.ctx,
+            self.override_stop_tokens,
+            self.key_string,
         )
 
+    def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
+        s = self.matcher.find_jump_forward_string()
+        if s:
+            return [], s
+        return None
+
+    def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]:
+        _, data = helper
+        return data, -1
+
+    def jump_and_retokenize(
+        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
+    ):
+        k = 0
+        for i, old_id in enumerate(old_output_ids):
+            if old_id == new_output_ids[i]:
+                k = i + 1
+            else:
+                break
+
+        # rollback to the last token that is the same
+        if k < len(old_output_ids):
+            self.matcher.rollback(len(old_output_ids) - k)
+
+        for i in range(k, len(new_output_ids)):
+            assert self.matcher.accept_token(new_output_ids[i])
+
+    def __repr__(self):
+        return f"XGrammarGrammar({self.key_string=}, {self.accepted_tokens=})"
+
 
 class XGrammarGrammarBackend(BaseGrammarBackend):
     def __init__(
@@ -143,9 +161,15 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         self.vocab_size = vocab_size
         self.override_stop_tokens = override_stop_tokens
 
-    def _from_context(self, ctx: CompiledGrammar) -> XGrammarGrammar:
-        matcher = GrammarMatcher(ctx, max_rollback_tokens=MAX_ROLLBACK_TOKENS)
-        return XGrammarGrammar(matcher, self.vocab_size, ctx, self.override_stop_tokens)
+    def _from_context(self, ctx: CompiledGrammar, key_string: str) -> 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
+        )
 
     def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -157,7 +181,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         except RuntimeError as e:
             logging.warning(f"Skip invalid json_schema: json_schema={key_string}, {e=}")
             return None
-        return self._from_context(ctx)
+        return self._from_context(ctx, key_string)
 
     def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -165,7 +189,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         except RuntimeError as e:
             logging.warning(f"Skip invalid ebnf: ebnf={key_string}, {e=}")
             return None
-        return self._from_context(ctx)
+        return self._from_context(ctx, key_string)
 
     def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -173,7 +197,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         except RuntimeError as e:
             logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
             return None
-        return self._from_context(ctx)
+        return self._from_context(ctx, key_string)
 
     def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -190,9 +214,11 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
                 tags, structural_tag["triggers"]
             )
         except RuntimeError as e:
-            logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
+            logging.warning(
+                f"Skip invalid structural_tag: structural_tag={key_string}, {e=}"
+            )
             return None
-        return self._from_context(ctx)
+        return self._from_context(ctx, key_string)
 
     def reset(self):
         if self.grammar_compiler:

sglang/srt/conversation.py

@@ -16,6 +16,7 @@
 # Adapted from
 # https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
 import dataclasses
+import re
 from enum import IntEnum, auto
 from typing import Callable, Dict, List, Optional, Tuple, Union
 
@@ -633,6 +634,20 @@ register_conv_template(
     )
 )
 
+# reference: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/blob/main/chat_template.json
+register_conv_template(
+    Conversation(
+        name="mistral",
+        system_template="[SYSTEM_PROMPT]\n{system_message}\n[/SYSTEM_PROMPT]\n\n",
+        roles=("[INST]", "[/INST]"),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=" ",
+        sep2=" </s><s>",
+        stop_str=["[INST]", "[/INST]", "[SYSTEM_PROMPT]", "[/SYSTEM_PROMPT]"],
+        image_token="[IMG]",
+    )
+)
+
 # reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/main/chat_template.json
 register_conv_template(
     Conversation(
@@ -852,91 +867,81 @@ register_conv_template(
 )
 
 
+@register_conv_template_matching_function
+def match_internvl(model_path: str):
+    if re.search(r"internvl2_5", model_path, re.IGNORECASE):
+        return "internvl-2-5"
+
+
 @register_conv_template_matching_function
 def match_llama_3_vision(model_path: str):
-    if (
-        "llama" in model_path.lower()
-        and "3.2" in model_path.lower()
-        and "vision" in model_path.lower()
-    ):
+    if re.search(r"llama.*3\.2.*vision", model_path, re.IGNORECASE):
         return "llama_3_vision"
 
 
 @register_conv_template_matching_function
 def match_deepseek_janus_pro(model_path: str):
-    if "janus" in model_path.lower():
+    if re.search(r"janus", model_path, re.IGNORECASE):
         return "janus-pro"
 
 
 @register_conv_template_matching_function
 def match_vicuna(model_path: str):
-    if "vicuna" in model_path.lower():
-        return "vicuna_v1.1"
-    if "llava-v1.5" in model_path.lower():
-        return "vicuna_v1.1"
-    if "llava-next-video-7b" in model_path.lower():
+    if re.search(r"vicuna|llava-v1\.5|llava-next-video-7b", model_path, re.IGNORECASE):
         return "vicuna_v1.1"
 
 
 @register_conv_template_matching_function
 def match_llama2_chat(model_path: str):
-    model_path = model_path.lower()
-    if "llama-2" in model_path and "chat" in model_path:
-        return "llama-2"
-    if (
-        "mistral" in model_path or "mixtral" in model_path
-    ) and "instruct" in model_path:
-        return "llama-2"
-    if "codellama" in model_path and "instruct" in model_path:
+    if re.search(
+        r"llama-2.*chat|codellama.*instruct",
+        model_path,
+        re.IGNORECASE,
+    ):
         return "llama-2"
 
 
+@register_conv_template_matching_function
+def match_mistral(model_path: str):
+    if re.search(r"pixtral|(mistral|mixtral).*instruct", model_path, re.IGNORECASE):
+        return "mistral"
+
+
 @register_conv_template_matching_function
 def match_deepseek_vl(model_path: str):
-    model_path = model_path.lower()
-    if "deepseek" in model_path and "vl2" in model_path:
+    if re.search(r"deepseek.*vl2", model_path, re.IGNORECASE):
         return "deepseek-vl2"
 
 
 @register_conv_template_matching_function
-def match_chat_ml(model_path: str):
-    # import pdb;pdb.set_trace()
-    model_path = model_path.lower()
-    # Now the suffix for qwen2 chat model is "instruct"
-    if "gme" in model_path and "qwen" in model_path and "vl" in model_path:
+def match_qwen_chat_ml(model_path: str):
+    if re.search(r"gme.*qwen.*vl", model_path, re.IGNORECASE):
         return "gme-qwen2-vl"
-    if "qwen" in model_path and "vl" in model_path:
+    if re.search(r"qwen.*vl", model_path, re.IGNORECASE):
         return "qwen2-vl"
-    if (
-        "llava-v1.6-34b" in model_path
-        or "llava-v1.6-yi-34b" in model_path
-        or "llava-next-video-34b" in model_path
-        or "llava-onevision-qwen2" in model_path
+    if re.search(
+        r"llava-v1\.6-34b|llava-v1\.6-yi-34b|llava-next-video-34b|llava-onevision-qwen2",
+        model_path,
+        re.IGNORECASE,
     ):
         return "chatml-llava"
 
 
 @register_conv_template_matching_function
-def match_gemma_it(model_path: str):
-    model_path = model_path.lower()
-    if "gemma" in model_path and "it" in model_path:
-        return "gemma-it"
-    if "gemma-3" in model_path and "1b" not in model_path:
-        # gemma-3-1b-it is completion model
+def match_gemma3_instruct(model_path: str):
+    if re.search(r"gemma-3.*it", model_path, re.IGNORECASE):
         return "gemma-it"
 
 
 @register_conv_template_matching_function
 def match_openbmb_minicpm(model_path: str):
-    model_path = model_path.lower()
-    if "minicpm-v" in model_path:
+    if re.search(r"minicpm-v", model_path, re.IGNORECASE):
         return "minicpmv"
-    elif "minicpm-o" in model_path:
+    elif re.search(r"minicpm-o", model_path, re.IGNORECASE):
         return "minicpmo"
 
 
 @register_conv_template_matching_function
 def match_moonshot_kimivl(model_path: str):
-    model_path = model_path.lower()
-    if "kimi" in model_path and "vl" in model_path:
+    if re.search(r"kimi.*vl", model_path, re.IGNORECASE):
         return "kimi-vl"

sglang/srt/disaggregation/base/conn.py

@@ -37,6 +37,7 @@ class BaseKVManager(ABC):
         args: KVArgs,
         disaggregation_mode: DisaggregationMode,
         server_args: ServerArgs,
+        is_mla_backend: Optional[bool] = False,
     ): ...
 
 

sglang/srt/disaggregation/decode.py

@@ -38,6 +38,7 @@ from sglang.srt.disaggregation.utils import (
     ReqToMetadataIdxAllocator,
     TransferBackend,
     get_kv_class,
+    is_mla_backend,
     kv_to_page_indices,
     poll_and_all_reduce,
 )
@@ -87,6 +88,7 @@ class DecodePreallocQueue:
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
         self.token_to_kv_pool = token_to_kv_pool_allocator.get_kvcache()
+        self.is_mla_backend = is_mla_backend(self.token_to_kv_pool)
         self.aux_dtype = aux_dtype
         self.metadata_buffers = metadata_buffers
         self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
@@ -131,7 +133,10 @@ class DecodePreallocQueue:
         kv_args.gpu_id = self.scheduler.gpu_id
         kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
         kv_manager = kv_manager_class(
-            kv_args, DisaggregationMode.DECODE, self.scheduler.server_args
+            kv_args,
+            DisaggregationMode.DECODE,
+            self.scheduler.server_args,
+            self.is_mla_backend,
         )
         return kv_manager
 
@@ -509,7 +514,7 @@ class SchedulerDisaggregationDecodeMixin:
     def event_loop_overlap_disagg_decode(self: Scheduler):
         result_queue = deque()
         self.last_batch: Optional[ScheduleBatch] = None
-        self.last_batch_in_queue = False  # last batch is modifed in-place, so we need another variable to track if it's extend
+        self.last_batch_in_queue = False  # last batch is modified in-place, so we need another variable to track if it's extend
 
         while True:
             recv_reqs = self.recv_requests()

sglang/srt/disaggregation/fake/conn.py

@@ -54,7 +54,7 @@ class FakeKVSender(BaseKVSender):
             logger.info(f"FakeKVSender send success")
         else:
             self.has_sent = False
-            logger.info(f"FakeKVSender send fake transfering")
+            logger.info(f"FakeKVSender send fake transferring")
 
     def failure_exception(self):
         raise Exception("Fake KVSender Exception")