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

sglang/srt/server_args_config_parser.py

@@ -0,0 +1,146 @@
+"""
+Configuration argument parser for command-line applications.
+Handles merging of YAML configuration files with command-line arguments.
+"""
+
+import logging
+from pathlib import Path
+from typing import Any, Dict, List
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+class ConfigArgumentMerger:
+    """Handles merging of configuration file arguments with command-line arguments."""
+
+    def __init__(self, boolean_actions: List[str] = None):
+        """Initialize with list of boolean action destinations."""
+        self.boolean_actions = boolean_actions or []
+
+    def merge_config_with_args(self, cli_args: List[str]) -> List[str]:
+        """
+        Merge configuration file arguments with command-line arguments.
+
+        Configuration arguments are inserted after the subcommand to maintain
+        proper precedence: CLI > Config > Defaults
+
+        Args:
+            cli_args: List of command-line arguments
+
+        Returns:
+            Merged argument list with config values inserted
+
+        Raises:
+            ValueError: If multiple config files specified or no config file provided
+        """
+        config_file_path = self._extract_config_file_path(cli_args)
+        if not config_file_path:
+            return cli_args
+
+        config_args = self._parse_yaml_config(config_file_path)
+        return self._insert_config_args(cli_args, config_args, config_file_path)
+
+    def _extract_config_file_path(self, args: List[str]) -> str:
+        """Extract the config file path from arguments."""
+        config_indices = [i for i, arg in enumerate(args) if arg == "--config"]
+
+        if len(config_indices) > 1:
+            raise ValueError("Multiple config files specified! Only one allowed.")
+
+        if not config_indices:
+            return None
+
+        config_index = config_indices[0]
+        if config_index == len(args) - 1:
+            raise ValueError("No config file specified after --config flag!")
+
+        return args[config_index + 1]
+
+    def _insert_config_args(
+        self, cli_args: List[str], config_args: List[str], config_file_path: str
+    ) -> List[str]:
+        """Insert configuration arguments into the CLI argument list."""
+        config_index = cli_args.index("--config")
+
+        # Split arguments around config file
+        before_config = cli_args[:config_index]
+        after_config = cli_args[config_index + 2 :]  # Skip --config and file path
+
+        # Simple merge: config args + CLI args
+        return config_args + before_config + after_config
+
+    def _parse_yaml_config(self, file_path: str) -> List[str]:
+        """
+        Parse YAML configuration file and convert to argument list.
+
+        Args:
+            file_path: Path to the YAML configuration file
+
+        Returns:
+            List of arguments in format ['--key', 'value', ...]
+
+        Raises:
+            ValueError: If file is not YAML or cannot be read
+        """
+        self._validate_yaml_file(file_path)
+
+        try:
+            with open(file_path, "r") as file:
+                config_data = yaml.safe_load(file)
+        except Exception as e:
+            logger.error(f"Failed to read config file {file_path}: {e}")
+            raise
+
+        # Handle empty files or None content
+        if config_data is None:
+            config_data = {}
+
+        if not isinstance(config_data, dict):
+            raise ValueError("Config file must contain a dictionary at root level")
+
+        return self._convert_config_to_args(config_data)
+
+    def _validate_yaml_file(self, file_path: str) -> None:
+        """Validate that the file is a YAML file."""
+        path = Path(file_path)
+        if path.suffix.lower() not in [".yaml", ".yml"]:
+            raise ValueError(f"Config file must be YAML format, got: {path.suffix}")
+
+        if not path.exists():
+            raise ValueError(f"Config file not found: {file_path}")
+
+    def _convert_config_to_args(self, config: Dict[str, Any]) -> List[str]:
+        """Convert configuration dictionary to argument list."""
+        args = []
+
+        for key, value in config.items():
+            if isinstance(value, bool):
+                self._add_boolean_arg(args, key, value)
+            elif isinstance(value, list):
+                self._add_list_arg(args, key, value)
+            else:
+                self._add_scalar_arg(args, key, value)
+
+        return args
+
+    def _add_boolean_arg(self, args: List[str], key: str, value: bool) -> None:
+        """Add boolean argument to the list."""
+        if key in self.boolean_actions:
+            # For boolean actions, always add the flag and value
+            args.extend([f"--{key}", str(value).lower()])
+        else:
+            # For regular booleans, only add flag if True
+            if value:
+                args.append(f"--{key}")
+
+    def _add_list_arg(self, args: List[str], key: str, value: List[Any]) -> None:
+        """Add list argument to the list."""
+        if value:  # Only add if list is not empty
+            args.append(f"--{key}")
+            args.extend(str(item) for item in value)
+
+    def _add_scalar_arg(self, args: List[str], key: str, value: Any) -> None:
+        """Add scalar argument to the list."""
+        args.extend([f"--{key}", str(value)])

sglang/srt/single_batch_overlap.py

@@ -0,0 +1,161 @@
+# Copyright 2025 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import torch
+
+from sglang.srt.layers import deep_gemm_wrapper
+from sglang.srt.layers.moe import get_moe_runner_backend
+from sglang.srt.layers.moe.topk import TopKOutput
+from sglang.srt.layers.moe.utils import is_sbo_enabled
+from sglang.srt.utils import get_int_env_var
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
+
+
+class SboFlags:
+    # TODO may have: "enable_dispatch_shared_one_stream_overlap", "enable_dispatch_gateup_gemm_two_stream_overlap", ...
+
+    @classmethod
+    def enable_combine_down_gemm_two_stream_overlap(cls):
+        return (
+            is_sbo_enabled()
+            # currently only cutedsl backend supports it
+            and get_moe_runner_backend().is_flashinfer_cutedsl()
+        )
+
+    @classmethod
+    def enable_combine_shared_two_stream_overlap(cls):
+        return is_sbo_enabled()
+
+    @classmethod
+    def fuse_shared_experts_inside_sbo(cls):
+        # TODO after antgroup's PR, should be `... or cls.enable_dispatch_shared_one_stream_overlap()`
+        return cls.enable_combine_shared_two_stream_overlap()
+
+
+@dataclass
+class CombineOverlapArgs:
+    # this "overlap" flag means overlapping with down gemm, not the general two-stream overlap
+    overlap: bool
+    stream: torch.cuda.Stream
+    wait_event: torch.cuda.Event
+    num_sms: int
+    signal: Optional[torch.Tensor] = None
+    threshold: int = 0
+
+
+@dataclass
+class DownGemmOverlapArgs:
+    num_sms: int
+    signal: torch.Tensor
+    start_event: torch.cuda.Event
+
+
+def execute_sbo(
+    forward_shared_experts: Callable[[], Any],
+    experts: FusedMoE,
+    hidden_states: torch.Tensor,
+    topk_output: TopKOutput,
+    alt_stream: Optional[torch.cuda.Stream] = None,
+    disable_sbo: bool = False,
+):
+
+    dispatch_output = experts.dispatcher.dispatch(
+        hidden_states=hidden_states, topk_output=topk_output
+    )
+
+    combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
+        _compute_overlap_args(dispatch_output, alt_stream, disable_sbo=disable_sbo)
+    )
+
+    hidden_states = experts.run_moe_core(
+        dispatch_output, down_gemm_overlap_args=down_gemm_overlap_args
+    )
+    if (e := meta_overlap_args.get("record_event_after_down")) is not None:
+        e.record()
+
+    if (not disable_sbo) and SboFlags.enable_combine_shared_two_stream_overlap():
+        # TODO reduce sm for non-deepgemm
+        with deep_gemm_wrapper.configure_deep_gemm_num_sms(
+            meta_overlap_args["compute_num_sms"]
+        ):
+            forward_shared_experts()
+
+    hidden_states = experts.dispatcher.combine(
+        hidden_states=hidden_states,
+        topk_ids=dispatch_output.topk_ids,
+        topk_weights=dispatch_output.topk_weights,
+        overlap_args=combine_overlap_args,
+    )
+
+    return hidden_states
+
+
+def _compute_overlap_args(dispatch_output, alt_stream, disable_sbo):
+    if disable_sbo or not (
+        SboFlags.enable_combine_down_gemm_two_stream_overlap()
+        or SboFlags.enable_combine_shared_two_stream_overlap()
+    ):
+        return None, None, {}
+
+    hidden_states = dispatch_output.hidden_states
+
+    num_local_experts, num_tokens_static, hidden_dim = hidden_states.shape
+
+    total_num_sms = torch.cuda.get_device_properties(
+        device="cuda"
+    ).multi_processor_count
+    communicate_num_sms = get_int_env_var("SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS", 32)
+    compute_num_sms = total_num_sms - communicate_num_sms
+
+    assert alt_stream is not None
+    combine_wait_event = torch.cuda.Event()
+    combine_overlap_args = CombineOverlapArgs(
+        overlap=False,
+        num_sms=communicate_num_sms,
+        stream=alt_stream,
+        wait_event=combine_wait_event,
+    )
+    meta_overlap_args = dict(
+        compute_num_sms=compute_num_sms,
+    )
+    down_gemm_overlap_args = None
+
+    if SboFlags.enable_combine_down_gemm_two_stream_overlap():
+        # TODO use zero_allocator to remove this `torch.zeros` call
+        # NOTE ours v2 use uint32 not int32 currently
+        combine_signal = torch.zeros(
+            num_local_experts, dtype=torch.uint32, device=hidden_states.device
+        )
+
+        down_gemm_overlap_args = DownGemmOverlapArgs(
+            signal=combine_signal,
+            start_event=combine_wait_event,
+            num_sms=compute_num_sms,
+        )
+        combine_overlap_args.overlap = True
+        combine_overlap_args.signal = combine_signal
+        combine_overlap_args.threshold = compute_num_sms
+    else:
+        meta_overlap_args |= dict(
+            record_event_after_down=combine_wait_event,
+        )
+
+    return combine_overlap_args, down_gemm_overlap_args, meta_overlap_args

sglang/srt/speculative/base_spec_worker.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.tp_worker import TpModelWorker
+
+
+class BaseDraftWorker(ABC):
+    @abstractmethod
+    def draft():
+        pass
+
+    @abstractmethod
+    def draft_extend():
+        pass
+
+
+class BaseSpecWorker(ABC):
+    @property
+    @abstractmethod
+    def target_worker(self) -> TpModelWorker:
+        pass
+
+    @property
+    @abstractmethod
+    def draft_worker(self) -> BaseDraftWorker:
+        pass
+
+    @abstractmethod
+    def clear_cache_pool(self):
+        # TODO: move this abstract method to BaseTpWorker and call through self.model_runner
+        pass