sglang 0.5.1.post3__py3-none-any.whl → 0.5.2rc0__py3-none-any.whl

sglang/bench_one_batch.py

@@ -61,6 +61,7 @@ from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.distributed.parallel_state import destroy_distributed_environment
 from sglang.srt.entrypoints.engine import _set_envs_and_config
 from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.managers.scheduler import Scheduler
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -509,6 +510,8 @@ def latency_test(
     bench_args,
     tp_rank,
 ):
+    initialize_moe_config(server_args)
+
     # Set CPU affinity
     if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
         set_gpu_proc_affinity(server_args.tp_size, server_args.nnodes, tp_rank)

sglang/srt/configs/__init__.py

@@ -5,6 +5,7 @@ from sglang.srt.configs.exaone import ExaoneConfig
 from sglang.srt.configs.janus_pro import MultiModalityConfig
 from sglang.srt.configs.kimi_vl import KimiVLConfig
 from sglang.srt.configs.kimi_vl_moonvit import MoonViTConfig
+from sglang.srt.configs.longcat_flash import LongcatFlashConfig
 from sglang.srt.configs.step3_vl import (
     Step3TextConfig,
     Step3VisionEncoderConfig,
@@ -16,6 +17,7 @@ __all__ = [
     "ChatGLMConfig",
     "DbrxConfig",
     "DeepseekVL2Config",
+    "LongcatFlashConfig",
     "MultiModalityConfig",
     "KimiVLConfig",
     "MoonViTConfig",

sglang/srt/configs/longcat_flash.py

@@ -0,0 +1,104 @@
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+FLASH_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+class LongcatFlashConfig(PretrainedConfig):
+    model_type = "longcat_flash"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=131072,
+        hidden_size=6144,
+        intermediate_size=None,
+        ffn_hidden_size=12288,
+        expert_ffn_hidden_size=2048,
+        num_layers=28,
+        num_hidden_layers=None,
+        num_attention_heads=64,
+        ep_size=1,
+        kv_lora_rank=512,
+        q_lora_rank=1536,
+        qk_rope_head_dim=128,
+        qk_nope_head_dim=128,
+        v_head_dim=128,
+        n_routed_experts=512,
+        moe_topk=12,
+        norm_topk_prob=False,
+        max_position_embeddings=131072,
+        rms_norm_eps=1e-05,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        pretraining_tp=1,
+        tie_word_embeddings=False,
+        rope_theta=10000000.0,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mla_scale_q_lora=True,
+        mla_scale_kv_lora=True,
+        torch_dtype="bfloat16",
+        params_dtype="bfloat16",
+        rounter_params_dtype="float32",
+        router_bias=False,
+        topk_method=None,
+        routed_scaling_factor=6.0,
+        zero_expert_num=256,
+        zero_expert_type="identity",
+        nextn_use_scmoe=False,
+        num_nextn_predict_layers=1,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            torch_dtype=torch_dtype,
+            params_dtype=params_dtype,
+            rounter_params_dtype=rounter_params_dtype,
+            topk_method=topk_method,
+            router_bias=router_bias,
+            nextn_use_scmoe=nextn_use_scmoe,
+            num_nextn_predict_layers=num_nextn_predict_layers,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = (
+            num_hidden_layers if num_hidden_layers is not None else num_layers
+        )
+        self.intermediate_size = (
+            intermediate_size if intermediate_size is not None else ffn_hidden_size
+        )
+        self.moe_intermediate_size = expert_ffn_hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.ep_size = ep_size
+        self.kv_lora_rank = kv_lora_rank
+        self.q_lora_rank = q_lora_rank
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.n_routed_experts = n_routed_experts
+        self.moe_topk = moe_topk
+        self.norm_topk_prob = norm_topk_prob
+        self.rms_norm_eps = rms_norm_eps
+        self.pretraining_tp = pretraining_tp
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mla_scale_q_lora = mla_scale_q_lora
+        self.mla_scale_kv_lora = mla_scale_kv_lora
+        self.zero_expert_num = zero_expert_num
+        self.zero_expert_type = zero_expert_type
+        self.routed_scaling_factor = routed_scaling_factor
+        self.hidden_act = "silu"

sglang/srt/configs/model_config.py

@@ -132,6 +132,13 @@ class ModelConfig:
         if is_draft_model and self.hf_config.architectures[0] == "Glm4MoeForCausalLM":
             self.hf_config.architectures[0] = "Glm4MoeForCausalLMNextN"
 
+        if (
+            is_draft_model
+            and self.hf_config.architectures[0] == "LongcatFlashForCausalLM"
+        ):
+            self.hf_config.architectures[0] = "LongcatFlashForCausalLMNextN"
+            self.hf_config.num_hidden_layers = self.hf_config.num_nextn_predict_layers
+
         if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM":
             self.hf_config.architectures[0] = "MiMoMTP"
         if (
@@ -199,6 +206,8 @@ class ModelConfig:
             "DeepseekV2ForCausalLM" in self.hf_config.architectures
             or "DeepseekV3ForCausalLM" in self.hf_config.architectures
             or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures
+            or "LongcatFlashForCausalLM" in self.hf_config.architectures
+            or "LongcatFlashForCausalLMNextN" in self.hf_config.architectures
         ):
             self.head_dim = 256
             self.attention_arch = AttentionArch.MLA
@@ -270,6 +279,9 @@ class ModelConfig:
             self.num_key_value_heads = self.num_attention_heads
         self.hidden_size = self.hf_text_config.hidden_size
         self.num_hidden_layers = self.hf_text_config.num_hidden_layers
+        self.num_attention_layers = self.num_hidden_layers
+        if "LongcatFlashForCausalLM" in self.hf_config.architectures:
+            self.num_attention_layers = self.num_hidden_layers * 2
         self.num_nextn_predict_layers = getattr(
             self.hf_text_config, "num_nextn_predict_layers", None
         )

sglang/srt/connector/__init__.py

@@ -20,7 +20,7 @@ class ConnectorType(str, enum.Enum):
     KV = "KV"
 
 
-def create_remote_connector(url, device="cpu") -> BaseConnector:
+def create_remote_connector(url, **kwargs) -> BaseConnector:
     connector_type = parse_connector_type(url)
     if connector_type == "redis":
         return RedisConnector(url)

sglang/srt/connector/base_connector.py

@@ -20,9 +20,8 @@ class BaseConnector(ABC):
     <connector_type://<host>:<port>/<model_name>/files/<filename>
     """
 
-    def __init__(self, url: str, device: torch.device = "cpu"):
+    def __init__(self, url: str):
         self.url = url
-        self.device = device
         self.closed = False
         self.local_dir = tempfile.mkdtemp()
         for sig in (signal.SIGINT, signal.SIGTERM):

sglang/srt/connector/redis.py

@@ -15,10 +15,10 @@ logger = logging.getLogger(__name__)
 
 class RedisConnector(BaseKVConnector):
 
-    def __init__(self, url: str, device: torch.device = "cpu"):
+    def __init__(self, url: str):
         import redis
 
-        super().__init__(url, device)
+        super().__init__(url)
         parsed_url = urlparse(url)
         self.connection = redis.Redis(host=parsed_url.hostname, port=parsed_url.port)
         self.model_name = parsed_url.path.lstrip("/")

sglang/srt/connector/serde/__init__.py

@@ -15,7 +15,7 @@ def create_serde(serde_type: str) -> Tuple[Serializer, Deserializer]:
 
     if serde_type == "safe":
         s = SafeSerializer()
-        d = SafeDeserializer(torch.uint8)
+        d = SafeDeserializer()
     else:
         raise ValueError(f"Unknown serde type: {serde_type}")
 

sglang/srt/connector/serde/safe_serde.py

@@ -19,11 +19,12 @@ class SafeSerializer(Serializer):
 
 class SafeDeserializer(Deserializer):
 
-    def __init__(self, dtype):
-        super().__init__(dtype)
+    def __init__(self):
+        # TODO: dtype options
+        super().__init__(torch.float32)
 
     def from_bytes_normal(self, b: Union[bytearray, bytes]) -> torch.Tensor:
-        return load(bytes(b))["tensor_bytes"].to(dtype=self.dtype)
+        return load(bytes(b))["tensor_bytes"]
 
     def from_bytes(self, b: Union[bytearray, bytes]) -> torch.Tensor:
         return self.from_bytes_normal(b)

sglang/srt/disaggregation/ascend/conn.py

@@ -1,6 +1,12 @@
+import concurrent.futures
 import logging
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
 
 from sglang.srt.disaggregation.ascend.transfer_engine import AscendTransferEngine
+from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
 from sglang.srt.disaggregation.mooncake.conn import (
     MooncakeKVBootstrapServer,
     MooncakeKVManager,
@@ -29,6 +35,75 @@ class AscendKVManager(MooncakeKVManager):
             self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
         )
 
+    def send_kvcache(
+        self,
+        mooncake_session_id: str,
+        prefill_kv_indices: npt.NDArray[np.int32],
+        dst_kv_ptrs: list[int],
+        dst_kv_indices: npt.NDArray[np.int32],
+        executor: concurrent.futures.ThreadPoolExecutor,
+    ):
+        # Group by indices
+        prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
+            prefill_kv_indices, dst_kv_indices
+        )
+
+        num_layers = len(self.kv_args.kv_data_ptrs)
+        layers_params = [
+            (
+                self.kv_args.kv_data_ptrs[layer_id],
+                dst_kv_ptrs[layer_id],
+                self.kv_args.kv_item_lens[layer_id],
+            )
+            for layer_id in range(num_layers)
+        ]
+
+        def set_transfer_blocks(
+            src_ptr: int, dst_ptr: int, item_len: int
+        ) -> List[Tuple[int, int, int]]:
+            transfer_blocks = []
+            for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
+                src_addr = src_ptr + int(prefill_index[0]) * item_len
+                dst_addr = dst_ptr + int(decode_index[0]) * item_len
+                length = item_len * len(prefill_index)
+                transfer_blocks.append((src_addr, dst_addr, length))
+            return transfer_blocks
+
+        # Worker function for processing a single layer
+        def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
+            transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        # Worker function for processing all layers in a batch
+        def process_layers(layers_params: List[Tuple[int, int, int]]) -> int:
+            transfer_blocks = []
+            for src_ptr, dst_ptr, item_len in layers_params:
+                transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len))
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        if self.enable_custom_mem_pool:
+            futures = [
+                executor.submit(
+                    process_layer,
+                    src_ptr,
+                    dst_ptr,
+                    item_len,
+                )
+                for (src_ptr, dst_ptr, item_len) in layers_params
+            ]
+            for future in concurrent.futures.as_completed(futures):
+                status = future.result()
+                if status != 0:
+                    for f in futures:
+                        f.cancel()
+                    return status
+        else:
+            # Combining all layers' params in one batch transfer is more efficient
+            # compared to using multiple threads
+            return process_layers(layers_params)
+
+        return 0
+
 
 class AscendKVSender(MooncakeKVSender):
     pass

sglang/srt/disaggregation/launch_lb.py

@@ -6,7 +6,6 @@ from sglang.srt.disaggregation.mini_lb import PrefillConfig, run
 
 @dataclasses.dataclass
 class LBArgs:
-    rust_lb: bool = False
     host: str = "0.0.0.0"
     port: int = 8000
     policy: str = "random"
@@ -17,11 +16,6 @@ class LBArgs:
 
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
-        parser.add_argument(
-            "--rust-lb",
-            action="store_true",
-            help="Deprecated, please use SGLang Router instead, this argument will have no effect.",
-        )
         parser.add_argument(
             "--host",
             type=str,
@@ -92,7 +86,6 @@ class LBArgs:
         ]
 
         return cls(
-            rust_lb=args.rust_lb,
             host=args.host,
             port=args.port,
             policy=args.policy,
@@ -102,12 +95,6 @@ class LBArgs:
             timeout=args.timeout,
         )
 
-    def __post_init__(self):
-        if not self.rust_lb:
-            assert (
-                self.policy == "random"
-            ), "Only random policy is supported for Python load balancer"
-
 
 def main():
     parser = argparse.ArgumentParser(

sglang/srt/disaggregation/mini_lb.py

@@ -7,6 +7,7 @@ import dataclasses
 import logging
 import random
 import urllib
+from http import HTTPStatus
 from itertools import chain
 from typing import List, Optional
 
@@ -262,14 +263,38 @@ async def get_server_info():
 
 @app.get("/get_model_info")
 async def get_model_info():
-    # Dummy model information
-    model_info = {
-        "model_path": "/path/to/dummy/model",
-        "tokenizer_path": "/path/to/dummy/tokenizer",
-        "is_generation": True,
-        "preferred_sampling_params": {"temperature": 0.7, "max_new_tokens": 128},
-    }
-    return ORJSONResponse(content=model_info)
+    global load_balancer
+
+    if not load_balancer or not load_balancer.prefill_servers:
+        raise HTTPException(
+            status_code=HTTPStatus.SERVICE_UNAVAILABLE,
+            detail="There is no server registered",
+        )
+
+    target_server_url = load_balancer.prefill_servers[0]
+    endpoint_url = f"{target_server_url}/get_model_info"
+
+    async with aiohttp.ClientSession() as session:
+        try:
+            async with session.get(endpoint_url) as response:
+                if response.status != 200:
+                    error_text = await response.text()
+                    raise HTTPException(
+                        status_code=HTTPStatus.BAD_GATEWAY,
+                        detail=(
+                            f"Failed to get model info from {target_server_url}"
+                            f"Status: {response.status}, Response: {error_text}"
+                        ),
+                    )
+
+                model_info_json = await response.json()
+                return ORJSONResponse(content=model_info_json)
+
+        except aiohttp.ClientError as e:
+            raise HTTPException(
+                status_code=HTTPStatus.SERVICE_UNAVAILABLE,
+                detail=f"Failed to get model info from backend",
+            )
 
 
 @app.post("/generate")

sglang/srt/disaggregation/prefill.py

@@ -567,7 +567,7 @@ class SchedulerDisaggregationPrefillMixin:
                 # Move the chunked request out of the batch so that we can merge
                 # only finished requests to running_batch.
                 self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
-                self.tree_cache.cache_unfinished_req(self.chunked_req)
+                self.tree_cache.cache_unfinished_req(self.chunked_req, chunked=True)
                 if self.enable_overlap:
                     # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
                     self.chunked_req.tmp_end_idx = min(

sglang/srt/distributed/parallel_state.py

@@ -52,6 +52,8 @@ from sglang.srt.utils import (
 
 _is_npu = is_npu()
 
+IS_ONE_DEVICE_PER_PROCESS = get_bool_env_var("SGLANG_ONE_DEVICE_PER_PROCESS")
+
 
 @dataclass
 class GraphCaptureContext:
@@ -223,10 +225,12 @@ class GroupCoordinator:
         use_message_queue_broadcaster: bool = False,
         group_name: Optional[str] = None,
     ):
+        # Set group info
         group_name = group_name or "anonymous"
         self.unique_name = _get_unique_name(group_name)
         _register_group(self)
 
+        # Set rank info
         self.rank = torch.distributed.get_rank()
         self.local_rank = local_rank
         self.device_group = None
@@ -250,14 +254,16 @@ class GroupCoordinator:
         assert self.cpu_group is not None
         assert self.device_group is not None
 
+        device_id = 0 if IS_ONE_DEVICE_PER_PROCESS else local_rank
         if is_cuda_alike():
-            self.device = torch.device(f"cuda:{local_rank}")
+            self.device = torch.device(f"cuda:{device_id}")
         elif _is_npu:
-            self.device = torch.device(f"npu:{local_rank}")
+            self.device = torch.device(f"npu:{device_id}")
         else:
             self.device = torch.device("cpu")
         self.device_module = torch.get_device_module(self.device)
 
+        # Import communicators
         self.use_pynccl = use_pynccl
         self.use_pymscclpp = use_pymscclpp
         self.use_custom_allreduce = use_custom_allreduce
@@ -270,6 +276,9 @@ class GroupCoordinator:
         from sglang.srt.distributed.device_communicators.custom_all_reduce import (
             CustomAllreduce,
         )
+        from sglang.srt.distributed.device_communicators.pymscclpp import (
+            PyMscclppCommunicator,
+        )
         from sglang.srt.distributed.device_communicators.pynccl import (
             PyNcclCommunicator,
         )
@@ -287,10 +296,6 @@ class GroupCoordinator:
                 device=self.device,
             )
 
-        from sglang.srt.distributed.device_communicators.pymscclpp import (
-            PyMscclppCommunicator,
-        )
-
         self.pymscclpp_comm: Optional[PyMscclppCommunicator] = None
         if use_pymscclpp and self.world_size > 1:
             self.pymscclpp_comm = PyMscclppCommunicator(
@@ -325,30 +330,30 @@ class GroupCoordinator:
                 except Exception as e:
                     logger.warning(f"Failed to initialize QuickAllReduce: {e}")
 
+        # Create communicator for other hardware backends
         from sglang.srt.distributed.device_communicators.hpu_communicator import (
             HpuCommunicator,
         )
+        from sglang.srt.distributed.device_communicators.npu_communicator import (
+            NpuCommunicator,
+        )
+        from sglang.srt.distributed.device_communicators.xpu_communicator import (
+            XpuCommunicator,
+        )
 
         self.hpu_communicator: Optional[HpuCommunicator] = None
         if use_hpu_communicator and self.world_size > 1:
             self.hpu_communicator = HpuCommunicator(group=self.device_group)
 
-        from sglang.srt.distributed.device_communicators.xpu_communicator import (
-            XpuCommunicator,
-        )
-
         self.xpu_communicator: Optional[XpuCommunicator] = None
         if use_xpu_communicator and self.world_size > 1:
             self.xpu_communicator = XpuCommunicator(group=self.device_group)
 
-        from sglang.srt.distributed.device_communicators.npu_communicator import (
-            NpuCommunicator,
-        )
-
         self.npu_communicator: Optional[NpuCommunicator] = None
         if use_npu_communicator and self.world_size > 1:
             self.npu_communicator = NpuCommunicator(group=self.device_group)
 
+        # Create message queue
         from sglang.srt.distributed.device_communicators.shm_broadcast import (
             MessageQueue,
         )
@@ -848,6 +853,11 @@ class GroupCoordinator:
         )
         return obj_list
 
+    def all_gather_object(self, obj: Any) -> List[Any]:
+        objs = [None] * self.world_size
+        torch.distributed.all_gather_object(objs, obj, group=self.cpu_group)
+        return objs
+
     def send_object(self, obj: Any, dst: int) -> None:
         """Send the input object list to the destination rank."""
         """NOTE: `dst` is the local rank of the destination rank."""

sglang/srt/entrypoints/engine.py

@@ -60,6 +60,7 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromTensorReqInput,
 )
+from sglang.srt.managers.multi_tokenizer_mixin import MultiTokenizerRouter
 from sglang.srt.managers.scheduler import run_scheduler_process
 from sglang.srt.managers.template_manager import TemplateManager
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
@@ -672,7 +673,7 @@ def _set_envs_and_config(server_args: ServerArgs):
     if server_args.attention_backend == "flashinfer":
         assert_pkg_version(
             "flashinfer_python",
-            "0.2.14.post1",
+            "0.3.0",
             "Please uninstall the old version and "
             "reinstall the latest version by following the instructions "
             "at https://docs.flashinfer.ai/installation.html.",
@@ -680,7 +681,7 @@ def _set_envs_and_config(server_args: ServerArgs):
     if _is_cuda and not get_bool_env_var("SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK"):
         assert_pkg_version(
             "sgl-kernel",
-            "0.3.7",
+            "0.3.7.post1",
             "Please reinstall the latest version with `pip install sgl-kernel --force-reinstall`",
         )
 
@@ -814,18 +815,24 @@ def _launch_subprocesses(
         ),
     )
     detoken_proc.start()
+    if server_args.tokenizer_worker_num > 1:
+        # Launch multi-tokenizer router
+        tokenizer_manager = MultiTokenizerRouter(server_args, port_args)
 
-    # Launch tokenizer process
-    tokenizer_manager = TokenizerManager(server_args, port_args)
+        # Initialize templates
+        template_manager = None
+    else:
+        # Launch tokenizer process
+        tokenizer_manager = TokenizerManager(server_args, port_args)
 
-    # Initialize templates
-    template_manager = TemplateManager()
-    template_manager.initialize_templates(
-        tokenizer_manager=tokenizer_manager,
-        model_path=server_args.model_path,
-        chat_template=server_args.chat_template,
-        completion_template=server_args.completion_template,
-    )
+        # Initialize templates
+        template_manager = TemplateManager()
+        template_manager.initialize_templates(
+            tokenizer_manager=tokenizer_manager,
+            model_path=server_args.model_path,
+            chat_template=server_args.chat_template,
+            completion_template=server_args.completion_template,
+        )
 
     # Wait for the model to finish loading
     scheduler_infos = []