sglang 0.5.2rc2__py3-none-any.whl → 0.5.3rc2__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, Union
+
+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,151 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import torch
+
+from sglang.srt.layers.moe import get_moe_runner_backend
+from sglang.srt.layers.moe.utils import is_sbo_enabled
+from sglang.srt.layers.quantization import deep_gemm_wrapper
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.utils import get_int_env_var
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE
+
+
+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 = -1
+
+
+@dataclass
+class DownGemmOverlapArgs:
+    num_sms: int
+    signal: torch.Tensor
+    start_event: torch.cuda.Event
+
+
+def execute_sbo(
+    forward_shared_experts: Callable[[], Any],
+    experts: "DeepEPMoE",
+    hidden_states: torch.Tensor,
+    topk_idx: torch.Tensor,
+    topk_weights: torch.Tensor,
+    forward_batch: ForwardBatch,
+    alt_stream: Optional = None,
+):
+    shared_output = None
+
+    dispatch_output = experts.dispatch(
+        hidden_states, topk_idx, topk_weights, forward_batch
+    )
+
+    combine_overlap_args, down_gemm_overlap_args, meta_overlap_args = (
+        _compute_overlap_args(dispatch_output, alt_stream)
+    )
+
+    hidden_states = experts.moe_impl(
+        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 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"]
+        ):
+            shared_output = forward_shared_experts()
+
+    hidden_states = experts.combine(
+        hidden_states,
+        dispatch_output.topk_idx,
+        dispatch_output.topk_weights,
+        forward_batch,
+        overlap_args=combine_overlap_args,
+    )
+
+    return hidden_states, shared_output
+
+
+def _compute_overlap_args(dispatch_output, alt_stream):
+    if 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_fp8
+    if isinstance(hidden_states, tuple):
+        hidden_states = hidden_states[0]
+
+    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/cpp_ngram/ngram.cpp

@@ -0,0 +1,374 @@
+#include "ngram.h"
+
+#include <algorithm>
+#include <cstring>
+#include <limits>
+#include <queue>
+#include <vector>
+
+namespace ngram {
+
+struct Node {
+  std::unordered_map<int32_t, int32_t> next;
+};
+
+Ngram::Result fillResult(int last_token, int draft_token_num, std::vector<Node>& tree, int root) {
+  Ngram::Result info;
+  std::vector<int32_t> prevs;
+  info.token.reserve(draft_token_num);
+  prevs.reserve(draft_token_num);
+  std::queue<std::tuple<int32_t, int32_t, int32_t>> queue;
+  info.token.emplace_back(last_token);
+  prevs.emplace_back(-1);
+
+  for (auto [token, next] : tree[root].next) {
+    queue.emplace(token, next, 0);
+  }
+  while (queue.size()) {
+    auto [token, next, prev] = queue.front();
+    queue.pop();
+    info.token.emplace_back(token);
+    prevs.emplace_back(prev);
+    for (auto [t, n] : tree[next].next) {
+      queue.emplace(t, n, info.token.size() - 1);
+    }
+  }
+
+  // zero padding to length
+  while (info.token.size() < draft_token_num) {
+    info.token.emplace_back(0);
+    prevs.emplace_back(0);
+  }
+
+  int n = info.token.size();
+  info.mask.resize(n * n, 0);
+  info.mask[0] = 1;
+  for (int i = 0; i < n; ++i) {
+    if (prevs[i] != -1) {
+      memcpy(&info.mask[i * n], &info.mask[prevs[i] * n], prevs[i] + 1);
+    }
+    info.mask[i * n + i] = 1;
+  }
+
+  return info;
+}
+
+Ngram::Ngram(size_t capacity, const Param& param) {
+  param_ = param;
+  nodes_.resize(capacity);
+  for (auto& node : nodes_) {
+    node_pool_.emplace_back(&node);
+  }
+  free_node_count_ = node_pool_.size();
+  root_ = getNode();
+
+  if (!(param_.branch_length > 1)) {
+    throw std::runtime_error(
+        "param_.branch_length must be greater than 1, current value: " + std::to_string(param_.branch_length));
+  }
+  if (!(param_.min_match_window_size > 0)) {
+    throw std::runtime_error(
+        "min_match_window_size must be greater than 0, current value: " + std::to_string(param_.min_match_window_size));
+  }
+  if (!(param_.min_match_window_size <= param_.max_match_window_size)) {
+    throw std::runtime_error(
+        "min_match_window_size must be less than or equal to max_match_window_size, current min_match_window_size: " +
+        std::to_string(param_.min_match_window_size) +
+        ", max_match_window_size: " + std::to_string(param_.max_match_window_size));
+  }
+  if (!(param_.max_match_window_size < param_.branch_length)) {
+    throw std::runtime_error(
+        "max_match_window_size must be less than branch_length, current max_match_window_size: " +
+        std::to_string(param_.max_match_window_size) + ", branch_length: " + std::to_string(param_.branch_length));
+  }
+  if (!(param_.min_bfs_breadth > 0)) {
+    throw std::runtime_error(
+        "min_bfs_breadth must be greater than 0, current value: " + std::to_string(param_.min_bfs_breadth));
+  }
+  if (!(param_.min_bfs_breadth <= param_.max_bfs_breadth)) {
+    throw std::runtime_error(
+        "min_bfs_breadth must be less than or equal to max_bfs_breadth, current min_bfs_breadth: " +
+        std::to_string(param_.min_bfs_breadth) + ", max_bfs_breadth: " + std::to_string(param_.max_bfs_breadth));
+  }
+  if (!(param_.draft_token_num > 0)) {
+    throw std::runtime_error(
+        "draft_token_num must be greater than 0, current value: " + std::to_string(param_.draft_token_num));
+  }
+  for (auto config : param_.batch_draft_token_num) {
+    if (config != std::numeric_limits<decltype(config)>::max()) {
+      if (!(config <= param_.draft_token_num)) {
+        throw std::runtime_error(
+            "batch_draft_token_num config value " + std::to_string(config) +
+            " must be less than or equal to draft_token_num: " + std::to_string(param_.draft_token_num));
+      }
+    }
+  }
+  for (auto config : param_.batch_min_match_window_size) {
+    if (config != std::numeric_limits<decltype(config)>::max()) {
+      if (!(config >= param_.min_match_window_size)) {
+        throw std::runtime_error(
+            "batch_min_match_window_size config value " + std::to_string(config) +
+            " must be greater than or equal to min_match_window_size: " + std::to_string(param_.min_match_window_size));
+      }
+      if (!(config <= param_.max_match_window_size)) {
+        throw std::runtime_error(
+            "batch_min_match_window_size config value " + std::to_string(config) +
+            " must be less than or equal to max_match_window_size: " + std::to_string(param_.max_match_window_size));
+      }
+    }
+  }
+
+  quit_flag_ = false;
+  insert_worker_ = std::thread(&Ngram::insert, this);
+}
+
+Ngram::~Ngram() {
+  quit_flag_ = true;
+  insert_queue_.close();
+  insert_worker_.join();
+}
+
+std::vector<std::pair<TrieNode*, int32_t>> Ngram::match(const std::vector<int32_t>& tokens, size_t batch_size) const {
+  auto draft_token_num = param_.get_draft_token_num(batch_size);
+  auto min_match_window_size = param_.get_min_match_window_size(batch_size);
+  auto max_match_window_size = param_.max_match_window_size;
+  std::vector<std::pair<TrieNode*, int32_t>> result;
+  result.reserve(param_.max_match_window_size - param_.min_match_window_size);
+  for (int32_t match_window_size = std::min(tokens.size(), param_.max_match_window_size);
+       match_window_size >= param_.min_match_window_size;
+       --match_window_size) {
+    auto start = tokens.data() + tokens.size() - match_window_size;
+    auto end = start + match_window_size;
+    auto cursor = root_;
+    while (start != end) {
+      auto iter = cursor->child.find(*start);
+      if (iter == cursor->child.end()) {
+        cursor = nullptr;
+        break;
+      }
+      ++start;
+      cursor = iter->second;
+    }
+    if (cursor) {
+      result.emplace_back(std::make_pair(cursor, match_window_size));
+    }
+  }
+  return result;
+}
+
+void Ngram::squeeze(size_t count) {
+  if (!(node_pool_.size() >= free_node_count_ + count)) {
+    throw std::runtime_error(
+        "Insufficient node size to release required nodes. "
+        "available to release: " +
+        std::to_string(node_pool_.size() - free_node_count_) + ", required to release: " + std::to_string(count));
+  }
+  while (count--) {
+    auto last = global_lru_.back();
+    global_lru_.pop_back();
+
+    if (!last->child.empty()) {
+      throw std::runtime_error("The node to be released still has child nodes and cannot be released. ");
+    }
+
+    last->parent->lru.erase(last->parent_lru_pos);
+    last->parent->sorted_children.erase(last);
+    last->parent->child.erase(last->token);
+
+    node_pool_[free_node_count_++] = last;
+  }
+}
+
+void Ngram::synchronize() const {
+  while (!insert_queue_.empty()) {
+    std::this_thread::sleep_for(std::chrono::microseconds(10));
+  }
+}
+
+void Ngram::insert() {
+  while (!quit_flag_) {
+    std::vector<int32_t> data;
+    if (!insert_queue_.dequeue(data)) {
+      continue;
+    }
+    const auto* token = data.data();
+    size_t size = data.size();
+    std::unique_lock<std::mutex> lock(mutex_);
+
+    for (size_t i = 0; i + param_.min_match_window_size < size; ++i) {
+      auto start = token + i;
+      auto end = start + std::min(size - i, param_.branch_length);
+
+      if (end - start > free_node_count_) {
+        squeeze(end - start - free_node_count_);
+      }
+
+      TrieNode* cursor = root_;
+      path_.clear();
+      while (start != end) {
+        auto token = *start;
+        auto iter = cursor->child.find(token);
+        if (iter == cursor->child.end()) {
+          iter = cursor->child.insert({token, getNode()}).first;
+          auto node = iter->second;
+
+          cursor->lru.emplace_front(node);
+          global_lru_.emplace_back(node);
+
+          node->token = token;
+          node->parent = cursor;
+          node->parent_lru_pos = cursor->lru.begin();
+          node->global_lru_pos = --global_lru_.end();
+          node->freq = 1;
+          cursor->sorted_children.insert(node);
+        } else {
+          auto node = iter->second;
+          cursor->sorted_children.erase(node);
+          node->freq++;
+          cursor->sorted_children.insert(node);
+          cursor->lru.splice(cursor->lru.begin(), cursor->lru, node->parent_lru_pos);
+        }
+        cursor = iter->second;
+        path_.emplace_back(cursor);
+        ++start;
+      }
+
+      for (auto it = path_.rbegin(); it != path_.rend(); ++it) {
+        TrieNode* node = *it;
+        global_lru_.splice(global_lru_.begin(), global_lru_, node->global_lru_pos);
+      }
+    }
+  }
+}
+
+void Ngram::asyncInsert(std::vector<std::vector<int32_t>>&& tokens) {
+  for (auto&& token : tokens) {
+    insert_queue_.enqueue(std::move(token));
+  }
+}
+
+Ngram::Result Ngram::matchBFS(const std::vector<int32_t>& tokens, size_t batch_size) const {
+  std::vector<std::pair<TrieNode*, int32_t>> nodes = match(tokens, batch_size);
+
+  double bfs_breadth_scale = double(param_.max_bfs_breadth - param_.min_bfs_breadth) /
+                             (param_.max_match_window_size - param_.min_match_window_size + 1);
+
+  auto draft_token_num = param_.get_draft_token_num(batch_size);
+  std::vector<Node> tree(draft_token_num + 1);
+  int root = 0;
+  int cursor = 1;
+
+  for (auto [node, depth] : nodes) {
+    std::queue<std::tuple<int32_t, double, const TrieNode*>> queue;  // parent, bfs_breadth, node
+    queue.push({root, (param_.max_match_window_size - depth) * bfs_breadth_scale + param_.min_bfs_breadth, node});
+    while (queue.size() && cursor <= draft_token_num) {
+      auto front = queue.front();
+      queue.pop();
+
+      auto parent = std::get<0>(front);
+      auto cur_breadth = std::get<1>(front);
+      auto iter = std::get<2>(front)->lru.begin();
+
+      auto breadth = std::max(1, int32_t(cur_breadth));
+      for (int i = 0; i < breadth && iter != std::get<2>(front)->lru.end() && cursor <= draft_token_num; ++i, ++iter) {
+        auto token = (*iter)->token;
+        auto pos = -1;
+        if (auto tit = tree[parent].next.find(token); tit != tree[parent].next.end()) {
+          pos = tit->second;
+        } else {
+          pos = tree[parent].next.insert(std::make_pair(token, cursor++)).first->second;
+        }
+        queue.emplace(pos, cur_breadth - bfs_breadth_scale, *iter);
+      }
+    }
+  }
+
+  return fillResult(tokens.back(), draft_token_num + 1, tree, root);
+}
+
+Ngram::Result Ngram::matchProb(const std::vector<int32_t>& tokens, size_t batch_size) const {
+  std::vector<std::pair<TrieNode*, int32_t>> nodes = match(tokens, batch_size);
+  auto draft_token_num = param_.get_draft_token_num(batch_size);
+
+  struct CompareByLastDouble {
+    bool operator()(
+        const std::tuple<double, const TrieNode*, double>& a,  // parent_pos,  node, final_prob
+        const std::tuple<double, const TrieNode*, double>& b) const {
+      return std::get<2>(a) < std::get<2>(b);
+    }
+  };
+
+  std::priority_queue<
+      std::tuple<double, const TrieNode*, double>,
+      std::vector<std::tuple<double, const TrieNode*, double>>,
+      CompareByLastDouble>
+      heap;
+
+  std::vector<Node> tree(draft_token_num + 1);
+
+  int root = 0;
+  int cursor = 1;
+  int top_k = param_.max_bfs_breadth;
+
+  auto addToHeap = [&heap, &top_k](int parent, const TrieNode* trie_node, double prob) -> void {
+    double sum_freq = 0.0;
+    int count = 0;
+    std::list<std::pair<TrieNode*, int32_t>> topk_children;
+    for (auto* child : trie_node->sorted_children) {
+      sum_freq += static_cast<double>(child->freq);
+      topk_children.emplace_back(child, child->freq);
+      if (++count >= top_k) break;
+    }
+    if (sum_freq <= 0) sum_freq = 1.0;
+    for (const auto& [child, freq] : topk_children) {
+      double norm_freq = static_cast<double>(freq) / sum_freq * prob;
+      heap.emplace(parent, child, norm_freq);
+    }
+  };
+
+  for (auto [node, _] : nodes) {
+    addToHeap(root, node, 1.0);
+
+    while (!heap.empty() && cursor <= draft_token_num) {
+      auto [parent, trie_node, prob] = heap.top();  // parent_pos, node, final_prob
+      heap.pop();
+      auto token = trie_node->token;
+      int pos = -1;
+      auto tit = tree[parent].next.find(token);
+      if (tit != tree[parent].next.end()) {
+        pos = tit->second;
+      } else {
+        pos = cursor++;
+        tree[parent].next[token] = pos;
+      }
+      addToHeap(pos, trie_node, prob);
+    }
+  }
+
+  return fillResult(tokens.back(), draft_token_num + 1, tree, root);
+}
+
+Ngram::Result Ngram::batchMatch(const std::vector<std::vector<int32_t>>& tokens) const {
+  std::unique_lock<std::mutex> lock(mutex_);
+  Result merged_result;
+  auto match_func = param_.match_type == "BFS" ? &Ngram::matchBFS : &Ngram::matchProb;
+  for (const auto& tks : tokens) {
+    Result res = (this->*match_func)(tks, tokens.size());
+    merged_result.token.insert(merged_result.token.end(), res.token.begin(), res.token.end());
+    merged_result.mask.insert(merged_result.mask.end(), res.mask.begin(), res.mask.end());
+  }
+  return merged_result;
+}
+
+void Ngram::Result::truncate(size_t n) {
+  if (n < token.size()) {
+    int full_n = token.size();
+    for (int i = 1; i < n; ++i) {
+      memcpy(&mask[i * n], &mask[i * full_n], sizeof(mask[0]) * n);
+    }
+    token.resize(n);
+    mask.resize(n * n);
+  }
+}
+
+}  // namespace ngram