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

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
 
@@ -48,6 +49,7 @@ class SeparatorStyle(IntEnum):
     DeepSeekVL2 = auto()
     QWEN2_VL_EMBED = auto()
     GEMMA3 = auto()
+    MPT = auto()
 
 
 @dataclasses.dataclass
@@ -327,6 +329,16 @@ class Conversation:
                     ret += role
             return ret
 
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
         else:
             raise ValueError(f"Invalid style: {self.sep_style}")
 
@@ -570,8 +582,11 @@ def generate_chat_conv(
                             real_content += "\n"  # for video
                         real_content += content.text
                     elif content.type == "image_url":
-                        # NOTE: Only works for llava
-                        real_content += image_token
+                        # NOTE: works for llava and intervl2_5
+                        if conv.name == "internvl-2-5":
+                            real_content = image_token + real_content
+                        else:
+                            real_content += image_token
                         conv.append_image(content.image_url.url)
                     elif content.type == "audio_url":
                         real_content += audio_token
@@ -619,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(
@@ -703,6 +732,19 @@ register_conv_template(
     )
 )
 
+register_conv_template(
+    Conversation(
+        name="internvl-2-5",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。",
+        roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
+        sep_style=SeparatorStyle.MPT,
+        sep="<|im_end|>\n",
+        stop_str=["<|im_end|>", "<|action_end|>"],
+        image_token="<image>",
+    )
+)
+
 # Reference: https://huggingface.co/docs/transformers/main/model_doc/qwen2_vl#usage-example
 register_conv_template(
     Conversation(
@@ -826,90 +868,80 @@ register_conv_template(
 
 
 @register_conv_template_matching_function
-def match_deepseek_janus_pro(model_path: str):
-    if (
-        "llama" in model_path.lower()
-        and "3.2" in model_path.lower()
-        and "vision" in model_path.lower()
-    ):
+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 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

@@ -21,6 +21,7 @@ Life cycle of a request in the decode server
 from __future__ import annotations
 
 import logging
+import os
 from collections import deque
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, Optional, Tuple
@@ -37,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,
 )
@@ -86,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
@@ -97,7 +100,9 @@ class DecodePreallocQueue:
         self.tp_size = tp_size
         self.bootstrap_port = bootstrap_port
 
-        self.num_reserved_decode_tokens = 512
+        self.num_reserved_decode_tokens = int(
+            os.environ.get("SGLANG_NUM_RESERVED_DECODE_TOKENS", "512")
+        )
 
         # Queue for requests pending pre-allocation
         self.queue: List[DecodeRequest] = []
@@ -128,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
 
@@ -506,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")

sglang/srt/disaggregation/mini_lb.py

@@ -3,10 +3,12 @@ Minimal HTTP load balancer for prefill and decode servers for testing.
 """
 
 import asyncio
+import dataclasses
+import logging
 import random
 import urllib
 from itertools import chain
-from typing import List
+from typing import List, Optional
 
 import aiohttp
 import orjson
@@ -14,11 +16,32 @@ import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import ORJSONResponse, Response, StreamingResponse
 
+from sglang.srt.disaggregation.utils import PDRegistryRequest
 
+
+def setup_logger():
+    logger = logging.getLogger("pdlb")
+    logger.setLevel(logging.INFO)
+
+    formatter = logging.Formatter(
+        "[PDLB (Python)] %(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    handler = logging.StreamHandler()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+    return logger
+
+
+logger = setup_logger()
+
+
+@dataclasses.dataclass
 class PrefillConfig:
-    def __init__(self, url: str, bootstrap_port: int):
-        self.url = url
-        self.bootstrap_port = bootstrap_port
+    url: str
+    bootstrap_port: Optional[int] = None
 
 
 class MiniLoadBalancer:
@@ -28,6 +51,10 @@ class MiniLoadBalancer:
         self.decode_servers = decode_servers
 
     def select_pair(self):
+        # TODO: return some message instead of panic
+        assert len(self.prefill_configs) > 0, "No prefill servers available"
+        assert len(self.decode_servers) > 0, "No decode servers available"
+
         prefill_config = random.choice(self.prefill_configs)
         decode_server = random.choice(self.decode_servers)
         return prefill_config.url, prefill_config.bootstrap_port, decode_server
@@ -47,7 +74,7 @@ class MiniLoadBalancer:
                 session.post(f"{decode_server}/{endpoint}", json=modified_request),
             ]
             # Wait for both responses to complete. Prefill should end first.
-            prefill_response, decode_response = await asyncio.gather(*tasks)
+            _, decode_response = await asyncio.gather(*tasks)
 
             return ORJSONResponse(
                 content=await decode_response.json(),
@@ -268,6 +295,32 @@ async def get_models():
             raise HTTPException(status_code=500, detail=str(e))
 
 
+@app.post("/register")
+async def register(obj: PDRegistryRequest):
+    if obj.mode == "prefill":
+        load_balancer.prefill_configs.append(
+            PrefillConfig(obj.registry_url, obj.bootstrap_port)
+        )
+        logger.info(
+            f"Registered prefill server: {obj.registry_url} with bootstrap port: {obj.bootstrap_port}"
+        )
+    elif obj.mode == "decode":
+        load_balancer.decode_servers.append(obj.registry_url)
+        logger.info(f"Registered decode server: {obj.registry_url}")
+    else:
+        raise HTTPException(
+            status_code=400,
+            detail="Invalid mode. Must be either PREFILL or DECODE.",
+        )
+
+    logger.info(
+        f"#Prefill servers: {len(load_balancer.prefill_configs)}, "
+        f"#Decode servers: {len(load_balancer.decode_servers)}"
+    )
+
+    return Response(status_code=200)
+
+
 def run(prefill_configs, decode_addrs, host, port):
     global load_balancer
     load_balancer = MiniLoadBalancer(prefill_configs, decode_addrs)
@@ -279,15 +332,16 @@ if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description="Mini Load Balancer Server")
     parser.add_argument(
-        "--prefill", required=True, help="Comma-separated URLs for prefill servers"
+        "--prefill", type=str, default=[], nargs="+", help="URLs for prefill servers"
     )
     parser.add_argument(
-        "--prefill-bootstrap-ports",
-        help="Comma-separated bootstrap ports for prefill servers",
-        default="8998",
+        "--decode", type=str, default=[], nargs="+", help="URLs for decode servers"
     )
     parser.add_argument(
-        "--decode", required=True, help="Comma-separated URLs for decode servers"
+        "--prefill-bootstrap-ports",
+        type=int,
+        nargs="+",
+        help="Bootstrap ports for prefill servers",
     )
     parser.add_argument(
         "--host", default="0.0.0.0", help="Host to bind the server (default: 0.0.0.0)"
@@ -297,22 +351,19 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    prefill_urls = args.prefill.split(",")
-    bootstrap_ports = [int(p) for p in args.prefill_bootstrap_ports.split(",")]
-
-    if len(bootstrap_ports) == 1:
-        bootstrap_ports = bootstrap_ports * len(prefill_urls)
+    bootstrap_ports = args.prefill_bootstrap_ports
+    if bootstrap_ports is None:
+        bootstrap_ports = [None] * len(args.prefill)
+    elif len(bootstrap_ports) == 1:
+        bootstrap_ports = bootstrap_ports * len(args.prefill)
     else:
-        if len(bootstrap_ports) != len(prefill_urls):
+        if len(bootstrap_ports) != len(args.prefill):
             raise ValueError(
                 "Number of prefill URLs must match number of bootstrap ports"
             )
-            exit(1)
-
-    prefill_configs = []
-    for url, port in zip(prefill_urls, bootstrap_ports):
-        prefill_configs.append(PrefillConfig(url, port))
 
-    decode_addrs = args.decode.split(",")
+    prefill_configs = [
+        PrefillConfig(url, port) for url, port in zip(args.prefill, bootstrap_ports)
+    ]
 
-    run(prefill_configs, decode_addrs, args.host, args.port)
+    run(prefill_configs, args.decode, args.host, args.port)