sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

sglang/srt/utils/hf_transformers_utils.py

@@ -16,9 +16,10 @@
 import contextlib
 import json
 import os
+import tempfile
 import warnings
 from pathlib import Path
-from typing import Any, Dict, Optional, Type, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 import torch
 from huggingface_hub import snapshot_download
@@ -45,27 +46,37 @@ from sglang.srt.configs import (
     KimiVLConfig,
     LongcatFlashConfig,
     MultiModalityConfig,
+    NemotronHConfig,
+    Olmo3Config,
     Qwen3NextConfig,
     Step3VLConfig,
 )
+from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config
 from sglang.srt.configs.internvl import InternVLChatConfig
 from sglang.srt.connector import create_remote_connector
 from sglang.srt.utils import is_remote_url, logger, lru_cache_frozenset
 
-_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
-    ChatGLMConfig.model_type: ChatGLMConfig,
-    DbrxConfig.model_type: DbrxConfig,
-    ExaoneConfig.model_type: ExaoneConfig,
-    DeepseekVL2Config.model_type: DeepseekVL2Config,
-    MultiModalityConfig.model_type: MultiModalityConfig,
-    KimiVLConfig.model_type: KimiVLConfig,
-    InternVLChatConfig.model_type: InternVLChatConfig,
-    Step3VLConfig.model_type: Step3VLConfig,
-    LongcatFlashConfig.model_type: LongcatFlashConfig,
-    Qwen3NextConfig.model_type: Qwen3NextConfig,
-    FalconH1Config.model_type: FalconH1Config,
-    DotsVLMConfig.model_type: DotsVLMConfig,
-    DotsOCRConfig.model_type: DotsOCRConfig,
+_CONFIG_REGISTRY: List[Type[PretrainedConfig]] = [
+    ChatGLMConfig,
+    DbrxConfig,
+    ExaoneConfig,
+    DeepseekVL2Config,
+    MultiModalityConfig,
+    KimiVLConfig,
+    InternVLChatConfig,
+    Step3VLConfig,
+    LongcatFlashConfig,
+    Olmo3Config,
+    Qwen3NextConfig,
+    FalconH1Config,
+    DotsVLMConfig,
+    DotsOCRConfig,
+    NemotronHConfig,
+    DeepseekVLV2Config,
+]
+
+_CONFIG_REGISTRY = {
+    config_cls.model_type: config_cls for config_cls in _CONFIG_REGISTRY
 }
 
 for name, cls in _CONFIG_REGISTRY.items():
@@ -106,6 +117,12 @@ def get_hf_text_config(config: PretrainedConfig):
         # if transformers config doesn't align with this assumption.
         assert hasattr(config.text_config, "num_attention_heads")
         return config.text_config
+
+    if hasattr(config, "llm_config"):
+        # PointsV1.5 Chat Model
+        assert hasattr(config.llm_config, "num_attention_heads")
+        return config.llm_config
+
     if hasattr(config, "language_config"):
         return config.language_config
     if hasattr(config, "thinker_config"):
@@ -143,7 +160,7 @@ def _load_deepseek_v32_model(
     config_json["architectures"] = ["DeepseekV3ForCausalLM"]
     config_json["model_type"] = "deepseek_v3"
 
-    tmp_path = os.path.join(local_path, "_tmp_config_folder")
+    tmp_path = os.path.join(tempfile.gettempdir(), "_tmp_config_folder")
     os.makedirs(tmp_path, exist_ok=True)
 
     unique_path = os.path.join(tmp_path, f"deepseek_v32_{os.getpid()}")
@@ -180,6 +197,15 @@ def get_config(
         config = AutoConfig.from_pretrained(
             model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
         )
+        if (
+            getattr(config, "auto_map", None) is not None
+            and config.auto_map.get("AutoModel")
+            == "modeling_deepseekocr.DeepseekOCRForCausalLM"
+        ):
+            config.model_type = "deepseek-ocr"
+            # TODO: Remove this workaround when AutoConfig correctly identifies deepseek-ocr.
+            # Hugging Face's AutoConfig currently misidentifies it as deepseekvl2.
+
     except ValueError as e:
         if not "deepseek_v32" in str(e):
             raise e
@@ -202,7 +228,8 @@ def get_config(
             "intermediate_size": 4304,
             "model_type": "siglip_vision_model",
             "num_attention_heads": 16,
-            "num_hidden_layers": 26,  # Model is originally 27-layer, we only need the first 26 layers for feature extraction.
+            "num_hidden_layers": 26,
+            # Model is originally 27-layer, we only need the first 26 layers for feature extraction.
             "patch_size": 14,
         }
         config.vision_config = SiglipVisionConfig(**vision_config)

sglang/srt/{host_shared_memory.py → utils/host_shared_memory.py}

@@ -1,5 +1,4 @@
 import logging
-import os
 from dataclasses import dataclass
 from multiprocessing import shared_memory
 from pathlib import Path

sglang/srt/{offloader.py → utils/offloader.py}

@@ -11,14 +11,14 @@ from sglang.srt.distributed.naive_distributed import (
     get_naive_distributed,
     set_naive_distributed,
 )
-from sglang.srt.host_shared_memory import (
+from sglang.srt.layers.parameter import ModelWeightParameter
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import MultiprocessingSerializer, is_pin_memory_available
+from sglang.srt.utils.host_shared_memory import (
     HostSharedMemoryManager,
     get_host_shared_memory_manager,
     set_host_shared_memory_manager,
 )
-from sglang.srt.layers.parameter import ModelWeightParameter
-from sglang.srt.server_args import ServerArgs
-from sglang.srt.utils import MultiprocessingSerializer, is_pin_memory_available
 
 logger = logging.getLogger(__name__)
 

sglang/srt/utils/profile_merger.py

@@ -0,0 +1,199 @@
+"""Merge Chrome trace files from multiple ranks (TP, DP, PP, EP) into a single trace."""
+
+import glob
+import gzip
+import json
+import logging
+import os
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class ProfileMerger:
+    """Merge profile traces from all parallelism types: TP, DP, PP, EP."""
+
+    def __init__(self, output_dir: str, profile_id: str):
+        self.output_dir = output_dir
+        self.profile_id = profile_id
+        self.merged_trace_path = os.path.join(
+            output_dir, f"merged-{profile_id}.trace.json.gz"
+        )
+
+        # Rank types in priority order (used for sorting and labeling)
+        self.rank_types = ["tp", "dp", "pp", "ep"]
+
+        # Sort index multipliers: DP (highest) > EP > PP > TP (lowest)
+        # These ensure proper visual ordering in trace viewer
+        self.sort_index_multipliers = {
+            "dp_rank": 100_000_000,
+            "ep_rank": 1_000_000,
+            "pp_rank": 10_000,
+            "tp_rank": 100,
+        }
+
+        # PID threshold for sort_index updates (only update for system PIDs < 1000)
+        self.pid_sort_index_threshold = 1000
+
+    def merge_chrome_traces(self) -> str:
+        """Merge Chrome traces from all ranks into a single trace.
+
+        Returns:
+            Path to merged trace file.
+
+        Raises:
+            ValueError: If no trace files found.
+        """
+        trace_files = self._discover_trace_files()
+        if not trace_files:
+            raise ValueError(f"No trace files found for profile_id: {self.profile_id}")
+
+        logger.info(f"Found {len(trace_files)} trace files to merge")
+
+        merged_trace = {"traceEvents": []}
+        all_device_properties = []
+
+        for trace_file in sorted(trace_files, key=self._get_rank_sort_key):
+            rank_info = self._extract_rank_info(trace_file)
+            logger.info(f"Processing {trace_file} with rank info: {rank_info}")
+
+            output = self._handle_file(trace_file, rank_info)
+
+            merged_trace["traceEvents"].extend(output["traceEvents"])
+
+            if "deviceProperties" in output:
+                all_device_properties.extend(output["deviceProperties"])
+                del output["deviceProperties"]
+
+            for key, value in output.items():
+                if key != "traceEvents" and key not in merged_trace:
+                    merged_trace[key] = value
+
+        if all_device_properties:
+            merged_trace["deviceProperties"] = all_device_properties
+
+        with gzip.open(self.merged_trace_path, "wb") as f:
+            f.write(json.dumps(merged_trace).encode("utf-8"))
+
+        logger.info(f"Merged profile saved to: {self.merged_trace_path}")
+        logger.info(f"Total events merged: {len(merged_trace['traceEvents'])}")
+
+        return self.merged_trace_path
+
+    def _discover_trace_files(self) -> List[str]:
+        """Discover trace files matching profile_id (supports TP/DP/PP/EP formats)."""
+        patterns = [f"{self.profile_id}*.trace.json.gz"]
+
+        trace_files = []
+        for pattern in patterns:
+            search_pattern = os.path.join(self.output_dir, pattern)
+            trace_files.extend(glob.glob(search_pattern))
+
+        trace_files = [
+            f
+            for f in trace_files
+            if not f.endswith(f"merged-{self.profile_id}.trace.json.gz")
+            and not f.endswith("-memory.pickle")
+            and "TP-" in f
+        ]
+        trace_files = list(set(trace_files))
+        return trace_files
+
+    def _extract_rank_info(self, filename: str) -> Dict[str, int]:
+        """Extract rank info (TP/DP/PP/EP) from filename."""
+        basename = os.path.basename(filename)
+        rank_info = {}
+
+        for rank_type in self.rank_types:
+            match = re.search(rf"{rank_type.upper()}-(\d+)", basename)
+            if match:
+                rank_info[f"{rank_type}_rank"] = int(match.group(1))
+
+        return rank_info
+
+    def _create_rank_label(self, rank_info: Dict[str, int]) -> str:
+        parts = []
+        for rank_type in self.rank_types:
+            rank_key = f"{rank_type}_rank"
+            if rank_key in rank_info:
+                parts.append(f"{rank_type.upper()}{rank_info[rank_key]:02d}")
+
+        return f"[{'-'.join(parts)}]" if parts else "[Unknown]"
+
+    def _handle_file(self, path: str, rank_info: Dict[str, int]) -> Dict[str, Any]:
+        logger.info(f"Processing file: {path}")
+
+        try:
+            with gzip.open(path, "rt", encoding="utf-8") as f:
+                trace = json.load(f)
+
+            output = {
+                key: value for key, value in trace.items() if key != "traceEvents"
+            }
+            output["traceEvents"] = self._process_events(
+                trace.get("traceEvents", []), rank_info
+            )
+            return output
+
+        except Exception as e:
+            logger.error(f"Failed to process trace file {path}: {e}")
+            return {"traceEvents": []}
+
+    def _process_events(
+        self, events: List[Dict], rank_info: Dict[str, int]
+    ) -> List[Dict]:
+        """Process events: update sort_index and add rank labels to PIDs."""
+        rank_label = self._create_rank_label(rank_info)
+
+        for event in events:
+            if event.get("name") == "process_sort_index":
+                pid = self._maybe_cast_int(event.get("pid"))
+                if pid is not None and pid < self.pid_sort_index_threshold:
+                    event["args"]["sort_index"] = self._calculate_sort_index(
+                        rank_info, pid
+                    )
+
+            event["pid"] = f"{rank_label} {event['pid']}"
+
+        return events
+
+    def _calculate_sort_index(self, rank_info: Dict[str, int], pid: int) -> int:
+        sort_index = pid
+        for rank_type, multiplier in self.sort_index_multipliers.items():
+            sort_index += rank_info.get(rank_type, 0) * multiplier
+        return sort_index
+
+    def _get_rank_sort_key(self, path: str) -> Tuple[int, int, int, int]:
+        rank_info = self._extract_rank_info(path)
+        return tuple(
+            rank_info.get(f"{rank_type}_rank", 0)
+            for rank_type in ["dp", "ep", "pp", "tp"]
+        )
+
+    def _maybe_cast_int(self, x) -> Optional[int]:
+        try:
+            return int(x)
+        except (ValueError, TypeError):
+            return None
+
+    def get_merge_summary(self) -> Dict[str, Any]:
+        if not os.path.exists(self.merged_trace_path):
+            return {"error": "Merged trace file not found"}
+
+        try:
+            with gzip.open(self.merged_trace_path, "rt") as f:
+                merged_data = json.load(f)
+
+            trace_files = self._discover_trace_files()
+
+            return {
+                "merged_file": self.merged_trace_path,
+                "total_events": len(merged_data.get("traceEvents", [])),
+                "total_files": len(trace_files),
+                "source_files": [os.path.basename(f) for f in trace_files],
+                "profile_id": self.profile_id,
+                "device_properties_count": len(merged_data.get("deviceProperties", [])),
+            }
+        except Exception as e:
+            return {"error": f"Failed to read merged trace: {str(e)}"}

sglang/test/attention/test_flashattn_backend.py

@@ -66,7 +66,7 @@ class MockModelRunner:
             enable_memory_saver=False,
         )
         # Required by torch native backend
-        self.server_args = ServerArgs(model_path="fake_model_path")
+        self.server_args = ServerArgs(model_path="dummy")
 
 
 @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")

sglang/test/attention/test_flashattn_mla_backend.py

@@ -4,7 +4,6 @@ import torch
 
 from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
-from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode

sglang/test/attention/test_prefix_chunk_info.py

@@ -2,8 +2,6 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
-from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.test.test_utils import CustomTestCase