sglang 0.5.2rc1__py3-none-any.whl → 0.5.2rc2__py3-none-any.whl

sglang/srt/mem_cache/storage/mooncake_store/mooncake_store.py

@@ -1,4 +1,3 @@
-import hashlib
 import json
 import logging
 import os
@@ -6,10 +5,8 @@ import uuid
 from dataclasses import dataclass
 from typing import Any, List, Optional
 
-import numpy as np
 import torch
 
-from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
 
 DEFAULT_GLOBAL_SEGMENT_SIZE = 4 * 1024 * 1024 * 1024  # 4 GiB
@@ -154,21 +151,36 @@ class MooncakeStore(HiCacheStorage):
         target_location: Optional[List[int]] = None,
         target_sizes: Optional[List[int]] = None,
     ) -> bool:
-        return self.batch_set([key], [value], [target_location], [target_sizes])
+        # Only support zero copy set for now
+        assert target_location is not None and target_sizes is not None
+        exist_result = self._batch_exist([key])
+        if exist_result[0] == 1:
+            return True
+        put_result = self._put_batch_zero_copy_impl(
+            [key], [target_location], [target_sizes]
+        )
+        return put_result[0] == 0
 
     def batch_set(
         self,
         keys: List[str],
         values: Optional[List[torch.Tensor]] = None,
-        target_location: Optional[List[int]] = None,
+        target_locations: Optional[List[int]] = None,
         target_sizes: Optional[List[int]] = None,
     ) -> bool:
-        assert len(keys) == len(target_location) == len(target_sizes)
+        # Only support zero copy set for now
+        assert target_locations is not None and target_sizes is not None
+        assert len(keys) == len(target_locations) == len(target_sizes)
+
         if len(keys) == 0:
             return False
 
         for i in range(len(keys)):
-            if keys[i] is None or target_location[i] is None or target_sizes[i] is None:
+            if (
+                keys[i] is None
+                or target_locations[i] is None
+                or target_sizes[i] is None
+            ):
                 return False
 
         exist_result = self._batch_exist(keys)
@@ -179,7 +191,7 @@ class MooncakeStore(HiCacheStorage):
         for i in range(len(keys)):
             if exist_result[i] != 1:
                 set_keys.append(keys[i])
-                set_target_locations.append(target_location[i])
+                set_target_locations.append(target_locations[i])
                 set_target_sizes.append(target_sizes[i])
                 set_indices.append(i)
         # Only set non-existing keys to storage
@@ -204,18 +216,24 @@ class MooncakeStore(HiCacheStorage):
         target_location: Optional[Any] = None,
         target_sizes: Optional[Any] = None,
     ) -> bool:
-        return self.batch_get([key], [target_location], [target_sizes]) == 1
+        assert target_location is not None and target_sizes is not None
+        get_result = self._get_batch_zero_copy_impl(
+            [key], [target_location], [target_sizes]
+        )
+        return get_result[0] >= 0
 
     def batch_get(
         self,
         keys: List[str],
-        target_location: Optional[Any] = None,
+        target_locations: Optional[Any] = None,
         target_sizes: Optional[Any] = None,
     ) -> int:
-        assert len(keys) == len(target_location) == len(target_sizes)
+        assert len(keys) == len(target_locations) == len(target_sizes)
         if len(keys) == 0:
             return 0
-        get_result = self._get_batch_zero_copy_impl(keys, target_location, target_sizes)
+        get_result = self._get_batch_zero_copy_impl(
+            keys, target_locations, target_sizes
+        )
         if self.is_mla_backend:
             key_multiplier = 1
         else:
@@ -226,7 +244,8 @@ class MooncakeStore(HiCacheStorage):
         return len(keys) // key_multiplier
 
     def exists(self, key) -> bool:
-        return self.batch_exists([key]) > 0
+        exist_result = self._batch_exist([key])
+        return exist_result[0] == 1
 
     def batch_exists(self, keys) -> int:
         if self.is_mla_backend:

sglang/srt/mem_cache/storage/mooncake_store/test_mooncake_store.py

@@ -0,0 +1,161 @@
+import logging
+import uuid
+
+import torch
+from mooncake_store import MooncakeStore
+
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorageConfig
+
+logging.basicConfig(
+    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+def generate_batch_query_keys(kv_num: int, config: HiCacheStorageConfig):
+    keys = []
+    for _ in range(kv_num):
+        key = "test_" + str(uuid.uuid4())
+        keys.append(key)
+    set_keys = []
+    for key in keys:
+        if config.is_mla_model:
+            set_keys.append(key + "_k")
+        else:
+            set_keys.append(key + f"_{config.tp_rank}_k")
+            set_keys.append(key + f"_{config.tp_rank}_v")
+    get_keys = set_keys
+    exist_keys = keys
+    return set_keys, get_keys, exist_keys
+
+
+def test_single_operation():
+    """Test the set API with a single key-value pair."""
+    print("=" * 100)
+    print("Testing single operation")
+
+    buffer_size = 1024 * 1024 * 16  # 16MB
+    value_elements = 1024
+    store = MooncakeStore()
+    buffer = torch.randn(buffer_size, dtype=torch.float32)
+    store.register_buffer(buffer)
+    value_size = value_elements * buffer.element_size()
+
+    key = str(uuid.uuid4())
+    set_slice = buffer[:value_elements]
+    get_slice = buffer[value_elements : 2 * value_elements]
+    set_location = set_slice.data_ptr()
+    get_location = get_slice.data_ptr()
+
+    # Test set operation
+    result = store.set(key, target_location=set_location, target_sizes=value_size)
+    assert result is True, f"❌set operation failed for key: {key}"
+
+    # Test exists operation
+    assert store.exists(key), f"❌key {key} should exist after set operation"
+
+    # Test get operation
+    result = store.get(key, target_location=get_location, target_sizes=value_size)
+    assert result is True, f"❌get operation failed for key: {key}"
+
+    # Compare the data using proper tensor indices
+    assert torch.allclose(
+        set_slice, get_slice, atol=1e-6
+    ), f"❌get operation failed for key: {key}"
+
+    logger.info(f"✅ Single operation passed")
+
+
+def test_batch_operation(config: HiCacheStorageConfig):
+    """Test the batch set/get APIs with multiple key-value pairs."""
+    print("=" * 100)
+    print(f"Testing batch operation with config: {config}")
+
+    buffer_size = 1024 * 1024 * 16  # 16MB
+    value_elements = 256
+    kv_num = 13
+    store = MooncakeStore(config)
+    buffer = torch.randn(buffer_size, dtype=torch.float32)
+    store.register_buffer(buffer)
+    value_size = value_elements * buffer.element_size()
+
+    set_keys, get_keys, exist_keys = generate_batch_query_keys(kv_num, config)
+    set_slices = [
+        buffer[i * value_elements : (i + 1) * value_elements]
+        for i in range(len(set_keys))
+    ]
+    set_locations = [set_slice.data_ptr() for set_slice in set_slices]
+    target_sizes = [value_size for _ in range(len(set_keys))]
+
+    # Test batch set operation
+    result = store.batch_set(
+        set_keys, target_locations=set_locations, target_sizes=target_sizes
+    )
+    assert result is True, f"❌batch set operation failed"
+
+    # Test batch exists operation
+    assert store.batch_exists(
+        exist_keys
+    ), f"❌keys should exist after batch set operation"
+
+    # Test batch get operation
+    get_slices = [
+        buffer[
+            (len(set_keys) + i)
+            * value_elements : (len(set_keys) + i + 1)
+            * value_elements
+        ]
+        for i in range(len(get_keys))
+    ]
+    get_locations = [get_slice.data_ptr() for get_slice in get_slices]
+    result = store.batch_get(
+        get_keys, target_locations=get_locations, target_sizes=target_sizes
+    )
+    assert result == kv_num, f"❌batch get operation failed"
+    for i in range(len(get_keys)):
+        assert torch.allclose(
+            set_slices[i], get_slices[i], atol=1e-6
+        ), f"❌batch get operation failed for key: {get_keys[i]}"
+
+    logger.info(f"✅ Batch operation passed")
+
+
+if __name__ == "__main__":
+    test_single_operation()
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=False,
+            tp_rank=0,
+            tp_size=1,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=True,
+            tp_rank=0,
+            tp_size=1,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=False,
+            tp_rank=1,
+            tp_size=4,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    test_batch_operation(
+        HiCacheStorageConfig(
+            is_mla_model=True,
+            tp_rank=3,
+            tp_size=8,
+            model_name=None,
+            is_page_first_layout=True,
+        )
+    )
+    logger.info(f"✅ All tests passed")

sglang/srt/metrics/collector.py

@@ -18,6 +18,8 @@ from dataclasses import dataclass
 from enum import Enum
 from typing import Dict, List, Optional, Union
 
+from sglang.srt.metrics.utils import generate_buckets
+from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import get_bool_env_var
 
 SGLANG_TEST_REQUEST_TIME_STATS = get_bool_env_var("SGLANG_TEST_REQUEST_TIME_STATS")
@@ -309,6 +311,7 @@ class SchedulerMetricsCollector:
 class TokenizerMetricsCollector:
     def __init__(
         self,
+        server_args: ServerArgs,
         labels: Dict[str, str],
         bucket_time_to_first_token: Optional[List[float]] = None,
         bucket_inter_token_latency: Optional[List[float]] = None,
@@ -334,7 +337,7 @@ class TokenizerMetricsCollector:
         )
 
         if collect_tokens_histogram:
-            bucket_prompt_tokens = [
+            default_bucket_prompt_tokens = [
                 100,
                 300,
                 500,
@@ -363,9 +366,11 @@ class TokenizerMetricsCollector:
                 name="sglang:prompt_tokens_histogram",
                 documentation="Histogram of prompt token length.",
                 labelnames=labels.keys(),
-                buckets=bucket_prompt_tokens,
+                buckets=generate_buckets(
+                    server_args.prompt_tokens_buckets, default_bucket_prompt_tokens
+                ),
             )
-            bucket_generation_tokens = [
+            default_bucket_generation_tokens = [
                 100,
                 300,
                 500,
@@ -390,7 +395,10 @@ class TokenizerMetricsCollector:
                 name="sglang:generation_tokens_histogram",
                 documentation="Histogram of generation token length.",
                 labelnames=labels.keys(),
-                buckets=bucket_generation_tokens,
+                buckets=generate_buckets(
+                    server_args.generation_tokens_buckets,
+                    default_bucket_generation_tokens,
+                ),
             )
 
         self.cached_tokens_total = Counter(

sglang/srt/metrics/utils.py

@@ -0,0 +1,48 @@
+# Copyright 2023-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.
+# ==============================================================================
+"""Utilities for Prometheus Metrics."""
+import math
+from typing import List
+
+
+def two_sides_exponential_buckets(
+    middle: float, base: float, count: int
+) -> List[float]:
+    buckets = []
+    half_count = math.ceil(count / 2)
+    distance = 1
+    buckets.append(middle)
+    for i in range(half_count):
+        distance *= base
+        buckets.append(middle + distance)
+        buckets.append(max(0, middle - distance))
+    return sorted(set(buckets))
+
+
+def generate_buckets(
+    buckets_rule: List[str], default_buckets: List[float]
+) -> List[float]:
+    if not buckets_rule:
+        buckets_rule = ["default"]
+
+    assert len(buckets_rule) > 0
+    rule = buckets_rule[0]
+    if rule == "tse":
+        middle, base, count = buckets_rule[1:]
+        assert float(base) > 1.0, "Base must be greater than 1.0"
+        return two_sides_exponential_buckets(float(middle), float(base), int(count))
+    if rule == "default":
+        return sorted(set(default_buckets))
+    assert rule == "customer"
+    return sorted(set([float(x) for x in buckets_rule[1:]]))

sglang/srt/model_executor/forward_batch_info.py

@@ -516,24 +516,23 @@ class ForwardBatch:
         for batch_idx in range(batch_size):
             mm_input = batch.multimodal_inputs[batch_idx]
             if self.forward_mode.is_decode():
-                mrope_position_deltas = (
-                    [0]
-                    if mm_input is None
-                    else flatten_nested_list(mm_input.mrope_position_delta.tolist())
-                )
-                next_input_positions = []
-                for mrope_position_delta in mrope_position_deltas:
-                    # batched deltas needs to be processed separately
-                    # Convert list of lists to tensor with shape [3, seq_len]
-                    next_input_positions += [
-                        MRotaryEmbedding.get_next_input_positions(
-                            mrope_position_delta,
-                            int(self.seq_lens[batch_idx]) - 1,
-                            int(self.seq_lens[batch_idx]),
-                        )
-                    ]
                 # 3 * N
-                mrope_positions_list[batch_idx] = torch.cat(next_input_positions, dim=1)
+                if mm_input is None:
+                    mrope_positions_list[batch_idx] = torch.full(
+                        (3, 1),
+                        self.seq_lens[batch_idx] - 1,
+                        dtype=torch.int64,
+                        device=model_runner.device,
+                    )
+                else:
+                    mrope_position_deltas = mm_input.mrope_position_delta.flatten().to(
+                        model_runner.device, non_blocking=True
+                    )
+                    mrope_positions_list[batch_idx] = (
+                        (mrope_position_deltas + self.seq_lens[batch_idx] - 1)
+                        .unsqueeze(0)
+                        .repeat(3, 1)
+                    )
             elif self.forward_mode.is_extend():
                 extend_seq_len, extend_prefix_len = (
                     batch.extend_seq_lens[batch_idx],

sglang/srt/model_executor/model_runner.py

@@ -1655,7 +1655,7 @@ class ModelRunner:
 
     def apply_torch_tp(self):
         logger.info(f"Enabling torch tensor parallelism on {self.tp_size} devices.")
-        from sglang.srt.model_parallel import tensor_parallel
+        from sglang.srt.layers.model_parallel import tensor_parallel
 
         device_mesh = torch.distributed.init_device_mesh(self.device, (self.tp_size,))
         tensor_parallel(self.model, device_mesh)