sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/constrained/reasoner_grammar_backend.py

@@ -17,7 +17,11 @@ from typing import List, Optional, Tuple
 
 import torch
 
-from .base_grammar_backend import BaseGrammarBackend, BaseGrammarObject
+from .base_grammar_backend import (
+    INVALID_GRAMMAR_OBJ,
+    BaseGrammarBackend,
+    BaseGrammarObject,
+)
 
 
 class ReasonerGrammarObject(BaseGrammarObject):
@@ -81,10 +85,9 @@ class ReasonerGrammarBackend(BaseGrammarBackend):
         self.grammar_backend = grammar_backend
         self.think_end_id = think_end_id
 
-    def _init_value_dispatch(
-        self, key: Tuple[str, str]
-    ) -> Optional[ReasonerGrammarObject]:
+    def _init_value_dispatch(self, key: Tuple[str, str]) -> Optional[BaseGrammarObject]:
         ret = self.grammar_backend._init_value_dispatch(key)
-        if ret is None:
-            return None
+        # avoid wrapping invalid grammar, so that the scheduler can detect it
+        if ret is None or ret is INVALID_GRAMMAR_OBJ:
+            return ret
         return ReasonerGrammarObject(ret, self.think_end_id)

sglang/srt/constrained/utils.py

@@ -0,0 +1,12 @@
+from typing import Dict
+
+
+def is_legacy_structural_tag(obj: Dict) -> bool:
+    # test whether an object is a legacy structural tag
+    # see `StructuralTagResponseFormat` at `sglang.srt.entrypoints.openai.protocol`
+    if obj.get("structures", None) is not None:
+        assert obj.get("triggers", None) is not None
+        return True
+    else:
+        assert obj.get("format", None) is not None
+        return False

sglang/srt/constrained/xgrammar_backend.py

@@ -34,6 +34,7 @@ from sglang.srt.constrained.base_grammar_backend import (
     BaseGrammarObject,
     GrammarStats,
 )
+from sglang.srt.constrained.utils import is_legacy_structural_tag
 from sglang.srt.utils import is_hip
 
 _is_hip = is_hip()
@@ -167,6 +168,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         tokenizer,
         vocab_size: int,
         model_eos_token_ids: Optional[List[int]] = None,
+        any_whitespace: bool = True,
     ):
         super().__init__()
 
@@ -188,6 +190,7 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info)
         self.vocab_size = vocab_size
         self.override_stop_tokens = override_stop_tokens
+        self.any_whitespace = any_whitespace
 
     def _from_context(
         self, ctx: CompiledGrammar, key_string: str, grammar_stats: GrammarStats
@@ -212,12 +215,14 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
                 # Note: This builtin JSON grammar includes *all* valid JSON (including, for example, arrays at the root)
                 ctx = self.grammar_compiler.compile_builtin_json_grammar()
             else:
-                ctx = self.grammar_compiler.compile_json_schema(schema=key_string)
+                ctx = self.grammar_compiler.compile_json_schema(
+                    schema=key_string, any_whitespace=self.any_whitespace
+                )
 
         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, GrammarStats())
+        return self._from_context(ctx, key_string, GrammarStats(dispatch_type="json"))
 
     def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -225,7 +230,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, GrammarStats())
+        return self._from_context(ctx, key_string, GrammarStats(dispatch_type="ebnf"))
 
     def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
@@ -233,26 +238,32 @@ 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, GrammarStats())
+        return self._from_context(ctx, key_string, GrammarStats(dispatch_type="regex"))
 
     def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]:
         try:
+            # TODO(dark): it's REALLY stupid to construct object from string and decode it again
             structural_tag = json.loads(key_string)
-            tags = [
-                StructuralTagItem(
-                    begin=structure["begin"],
-                    schema=json.dumps(structure["schema"]),
-                    end=structure["end"],
+            if is_legacy_structural_tag(structural_tag):
+                tags = [
+                    StructuralTagItem(
+                        begin=structure["begin"],
+                        schema=json.dumps(structure["schema"]),
+                        end=structure["end"],
+                    )
+                    for structure in structural_tag["structures"]
+                ]
+                ctx = self.grammar_compiler.compile_structural_tag(
+                    tags, structural_tag["triggers"]
                 )
-                for structure in structural_tag["structures"]
-            ]
-            ctx = self.grammar_compiler.compile_structural_tag(
-                tags, structural_tag["triggers"]
-            )
+            else:
+                ctx = self.grammar_compiler.compile_structural_tag(key_string)
         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, GrammarStats())
+        return self._from_context(
+            ctx, key_string, GrammarStats(dispatch_type="structural_tag")
+        )
 
     def reset(self):
         self.grammar_compiler.clear_cache()

sglang/srt/debug_utils/dumper.py

@@ -36,6 +36,15 @@ class _Dumper:
         self._forward_pass_id = 0
 
     def on_forward_pass_start(self):
+        """This should be called on all ranks."""
+
+        if not self._enable:
+            return
+
+        # Users may want to `dump` only on some ranks, thus determine name here
+        if self._partial_name is None:
+            self._partial_name = _get_partial_name()
+
         self._forward_pass_id += 1
         print(
             f"[Dumper] [{time.time()}] on_forward_pass_start id={self._forward_pass_id}"
@@ -48,11 +57,9 @@ class _Dumper:
         assert (
             self._forward_pass_id >= 1
         ), "Do you forget to call `dumper.on_forward_pass_start()`?"
+        assert self._partial_name is not None
         self._dump_index += 1
 
-        if self._partial_name is None:
-            self._partial_name = _get_partial_name()
-
         rank = _get_rank()
         full_kwargs = dict(
             forward_pass_id=self._forward_pass_id,

sglang/srt/disaggregation/ascend/conn.py

@@ -13,7 +13,7 @@ from sglang.srt.disaggregation.mooncake.conn import (
     MooncakeKVReceiver,
     MooncakeKVSender,
 )
-from sglang.srt.utils import get_local_ip_by_remote
+from sglang.srt.utils import get_local_ip_auto
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +21,7 @@ logger = logging.getLogger(__name__)
 class AscendKVManager(MooncakeKVManager):
     def init_engine(self):
         # TransferEngine initialized on ascend.
-        local_ip = get_local_ip_by_remote()
+        local_ip = get_local_ip_auto()
         self.engine = AscendTransferEngine(
             hostname=local_ip,
             npu_id=self.kv_args.gpu_id,

sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -1,10 +1,20 @@
 import logging
 import os
-from typing import List, Optional
+from typing import List
+
+import torch
 
 from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
 from sglang.srt.disaggregation.utils import DisaggregationMode
 
+try:
+    from mf_adapter import TransferEngine
+
+    import_error = None
+except ImportError as e:
+    import_error = e
+    pass
+
 logger = logging.getLogger(__name__)
 
 
@@ -13,12 +23,11 @@ class AscendTransferEngine(MooncakeTransferEngine):
     def __init__(
         self, hostname: str, npu_id: int, disaggregation_mode: DisaggregationMode
     ):
-        try:
-            from mf_adapter import TransferEngine
-        except ImportError as e:
-            raise ImportError(
+        if import_error is not None:
+            logger.warning(
                 "Please install mf_adapter, for details, see docs/backend/pd_disaggregation.md"
-            ) from e
+            )
+            raise import_error
 
         self.engine = TransferEngine()
         self.hostname = hostname
@@ -37,12 +46,29 @@ class AscendTransferEngine(MooncakeTransferEngine):
         self.initialize()
 
     def initialize(self) -> None:
+        from sglang.srt.layers.dp_attention import (
+            get_tensor_model_parallel_world_size,
+            get_tp_group,
+        )
+
+        transfer_protocol = self._get_transfer_protocol()
+        if transfer_protocol is None or transfer_protocol == "sdma":
+            trans_op_type = TransferEngine.TransDataOpType.SDMA
+        else:
+            trans_op_type = TransferEngine.TransDataOpType.DEVICE_RDMA
+            """with device RDMA for PD transfer"""
+            tmp_tensor = torch.zeros(1, device="npu")
+            output_tensor_list = [
+                torch.empty_like(tmp_tensor)
+                for _ in range(get_tensor_model_parallel_world_size())
+            ]
+            # Initialize hccl in advance through all_gather to avoid conflicts with rdma initialization.
+            torch.distributed.all_gather(
+                output_tensor_list, tmp_tensor, group=get_tp_group().device_group
+            )
         """Initialize the ascend transfer instance."""
         ret_value = self.engine.initialize(
-            self.store_url,
-            self.session_id,
-            self.role,
-            self.npu_id,
+            self.store_url, self.session_id, self.role, self.npu_id, trans_op_type
         )
         if ret_value != 0:
             logger.error("Ascend Transfer Engine initialization failed.")
@@ -56,3 +82,15 @@ class AscendTransferEngine(MooncakeTransferEngine):
             ret_value = -1
         if ret_value != 0:
             logger.debug(f"Ascend memory registration for ptr {ptrs} failed.")
+
+    @staticmethod
+    def _get_transfer_protocol():
+        protocol = os.getenv("ASCEND_MF_TRANSFER_PROTOCOL")
+        allowed_protocols = {"device_rdma", "sdma"}
+        if protocol and protocol.lower() in allowed_protocols:
+            return protocol.lower()
+        else:
+            logger.warning(
+                "Invalid or no transfer protocol specified, using default protocol."
+            )
+            return None

sglang/srt/disaggregation/base/conn.py

@@ -20,6 +20,10 @@ class KVArgs:
     aux_data_ptrs: List[int]
     aux_data_lens: List[int]
     aux_item_lens: List[int]
+    state_data_ptrs: List[int]
+    state_data_lens: List[int]
+    state_item_lens: List[int]
+    state_type: str  # "none", "mamba", "swa"
     ib_device: str
     ib_traffic_class: str
     gpu_id: int
@@ -76,9 +80,13 @@ class BaseKVSender(ABC):
         ...
 
     @abstractmethod
-    def send(self, kv_indices: npt.NDArray[np.int32]):
+    def send(
+        self,
+        kv_indices: npt.NDArray[np.int32],
+        state_indices: Optional[List[int]] = None,
+    ):
         """
-        Send the kv cache at the given kv indices to the decoder server
+        Send the kv cache at the given kv indices and the extra cache/state at the given indices to the decoder server
         """
         ...
 
@@ -108,9 +116,14 @@ class BaseKVReceiver(ABC):
     ): ...
 
     @abstractmethod
-    def init(self, kv_indices: npt.NDArray[np.int32], aux_index: Optional[int] = None):
+    def init(
+        self,
+        kv_indices: npt.NDArray[np.int32],
+        aux_index: Optional[int] = None,
+        state_indices: Optional[List[int]] = None,
+    ):
         """
-        Notify the prefill server about the kv indices and aux index
+        Notify the prefill server about the kv indices, aux index, and state_indices.
         """
         ...