sglang 0.4.9.post5__py3-none-any.whl → 0.4.10__py3-none-any.whl

sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py

@@ -0,0 +1,278 @@
+import atexit
+import concurrent.futures
+import json
+import logging
+import os
+import signal
+import threading
+from collections import OrderedDict
+from functools import wraps
+from typing import List, Optional
+
+import torch
+
+from sglang.srt.mem_cache.hicache_storage import HiCacheStorage
+from sglang.srt.mem_cache.storage.hf3fs.client_hf3fs import Hf3fsClient
+
+logger = logging.getLogger(__name__)
+
+
+class AtomicCounter:
+    def __init__(self, n: int):
+        assert n > 0
+        self.n = n
+        self._value = 0
+        self._lock = threading.Lock()
+
+    def next(self) -> int:
+        with self._lock:
+            current = self._value
+            self._value = (current + 1) % self.n
+            return current
+
+
+def synchronized():
+    def _decorator(func):
+        @wraps(func)
+        def wrapper(self, *args, **kwargs):
+            with self.lock:
+                return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return _decorator
+
+
+class HiCacheHF3FS(HiCacheStorage):
+    default_env_var: str = "SGLANG_HICACHE_HF3FS_CONFIG_PATH"
+
+    def __init__(
+        self,
+        file_path: str,
+        file_size: int,
+        numjobs: int,
+        bytes_per_page: int,
+        entries: int,
+        dtype: torch.dtype,
+    ):
+        self.file_path = file_path
+        self.file_size = file_size
+        self.numjobs = numjobs
+        self.bytes_per_page = bytes_per_page
+        self.entries = entries
+        self.dtype = dtype
+
+        self.numel = self.bytes_per_page // self.dtype.itemsize
+
+        self.num_pages = self.file_size // self.bytes_per_page
+
+        logger.info(
+            "HiCacheHF3FS "
+            f"file_path = {self.file_path}, "
+            f"file_size = {self.file_size/(2**30):.2f} GB, "
+            f"numjobs = {self.numjobs}, "
+            f"bytes_per_page = {self.bytes_per_page/(2**20):.2f} MB, "
+            f"entries = {self.entries}, "
+            f"num_pages = {self.num_pages}"
+        )
+
+        self.ac = AtomicCounter(self.numjobs)
+        self.clients = [
+            Hf3fsClient(
+                self.file_path, self.file_size, self.bytes_per_page, self.entries
+            )
+            for _ in range(numjobs)
+        ]
+        self.executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=self.numjobs, thread_name_prefix="HiCacheHF3FS"
+        )
+
+        # Implemented a preliminary single-file page_hash -> file_offset index as interim storage.
+        # Future iterations may adopt a global KVCache manager to coordinate external cache instances
+        # through centralized metadata orchestration.
+        self.lock = threading.RLock()
+        self.free_pages = list(range(self.num_pages))
+        self.key_to_index = OrderedDict()
+
+        atexit.register(self.close)
+
+        signal.signal(signal.SIGINT, lambda sig, frame: self.close())
+        signal.signal(signal.SIGTERM, lambda sig, frame: self.close())
+        signal.signal(signal.SIGQUIT, lambda sig, frame: self.close())
+
+    @staticmethod
+    def from_env_config(
+        rank: int, bytes_per_page: int, dtype: torch.dtype
+    ) -> "HiCacheHF3FS":
+        config_path = os.getenv(HiCacheHF3FS.default_env_var)
+        if not config_path:
+            return HiCacheHF3FS(
+                file_path=f"/data/hicache.{rank}.bin",
+                file_size=1 << 40,
+                numjobs=16,
+                bytes_per_page=bytes_per_page,
+                entries=8,
+                dtype=dtype,
+            )
+
+        try:
+            with open(config_path, "r") as f:
+                config = json.load(f)
+        except Exception as e:
+            raise RuntimeError(f"Failed to load config from {config_path}: {str(e)}")
+
+        required_keys = {
+            "file_path_prefix",
+            "file_size",
+            "numjobs",
+            "entries",
+        }
+        missing_keys = required_keys - set(config.keys())
+        if missing_keys:
+            raise ValueError(f"Missing required keys in config: {missing_keys}")
+
+        return HiCacheHF3FS(
+            file_path=f"{config['file_path_prefix']}.{rank}.bin",
+            file_size=int(config["file_size"]),
+            numjobs=int(config["numjobs"]),
+            bytes_per_page=bytes_per_page,
+            entries=int(config["entries"]),
+            dtype=dtype,
+        )
+
+    def get(
+        self, key: str, target_location: Optional[torch.Tensor] = None
+    ) -> torch.Tensor | None:
+        return self.batch_get([key], target_location)[0]
+
+    @synchronized()
+    def batch_get(
+        self,
+        keys: List[str],
+        target_locations: Optional[List[torch.Tensor]] = None,
+    ) -> List[torch.Tensor | None]:
+        batch_indices, file_offsets = [], []
+        for i, key in enumerate(keys):
+            if key not in self.key_to_index:
+                continue
+            batch_indices.append(i)
+            file_offsets.append(self.key_to_index[key] * self.bytes_per_page)
+            self.key_to_index.move_to_end(key)
+        # TODO: target_locations
+        file_results = [
+            torch.empty(self.numel, dtype=self.dtype) for _ in range(len(batch_indices))
+        ]
+
+        futures = [
+            self.executor.submit(
+                self.clients[self.ac.next()].batch_read,
+                file_offsets[i : i + self.entries],
+                file_results[i : i + self.entries],
+            )
+            for i in range(0, len(batch_indices), self.entries)
+        ]
+        read_results = [result for future in futures for result in future.result()]
+
+        results = [None] * len(keys)
+        for batch_index, file_result, read_result in zip(
+            batch_indices, file_results, read_results
+        ):
+            if read_result == self.bytes_per_page:
+                results[batch_index] = file_result
+            else:
+                logger.error(f"HiCacheHF3FS get {keys[batch_index]} failed")
+
+        return results
+
+    def set(self, key: str, value: torch.Tensor) -> bool:
+        return self.batch_set([key], [value])
+
+    def batch_set(self, keys: List[str], values: List[torch.Tensor]) -> bool:
+        indices = self.get_batch_set_indices(keys)
+        batch_indices, file_offsets, file_values = [], [], []
+        for i, (value, (is_written, index)) in enumerate(zip(values, indices)):
+            if is_written or index == -1:
+                continue
+            batch_indices.append(i)
+            file_offsets.append(index * self.bytes_per_page)
+            file_values.append(value.contiguous())
+
+        futures = [
+            self.executor.submit(
+                self.clients[self.ac.next()].batch_write,
+                file_offsets[i : i + self.entries],
+                file_values[i : i + self.entries],
+            )
+            for i in range(0, len(batch_indices), self.entries)
+        ]
+        write_results = [
+            result == self.bytes_per_page
+            for future in futures
+            for result in future.result()
+        ]
+
+        results = [index[0] for index in indices]
+        for batch_index, write_result in zip(batch_indices, write_results):
+            key = keys[batch_index]
+            index = indices[batch_index][1]
+            if write_result:
+                self.key_to_index[key] = index
+                self.key_to_index.move_to_end(key)
+            else:
+                logger.error(f"HiCacheHF3FS set {key} failed")
+                self.free_pages.append(index)
+            results[batch_index] = write_result
+        return all(results)
+
+    @synchronized()
+    def get_batch_set_indices(self, keys: List[str]) -> list:
+        ionum = len(keys)
+        # results: tuples of (is_written: bool, page_idx: int)
+        # - is_written: True = hit (no I/O), False = write (miss)
+        # - page_idx: page storing data
+        results = [None] * min(ionum, self.num_pages)
+        if ionum > self.num_pages:
+            results.extend([(False, -1)] * (ionum - self.num_pages))
+
+        new_keys = []
+        for batch_index, key in enumerate(keys[: self.num_pages]):
+            if key in self.key_to_index:
+                results[batch_index] = (True, self.key_to_index[key])
+                self.key_to_index.move_to_end(key)
+            else:
+                new_keys.append((batch_index, key))
+
+        for batch_index, _ in new_keys:
+            index = (
+                self.free_pages.pop()
+                if len(self.free_pages) > 0
+                else self.key_to_index.popitem(last=False)[1]
+            )
+            results[batch_index] = (False, index)
+
+        return results
+
+    @synchronized()
+    def delete(self, key: str) -> None:
+        if key not in self.key_to_index:
+            return
+        index = self.key_to_index.pop(key)
+        self.free_pages.append(index)
+
+    @synchronized()
+    def exists(self, key: str) -> bool:
+        return key in self.key_to_index
+
+    @synchronized()
+    def clear(self) -> None:
+        self.free_pages = list(range(self.num_pages))
+        self.key_to_index.clear()
+
+    def close(self) -> None:
+        try:
+            for c in self.clients:
+                c.close()
+            self.executor.shutdown(wait=True)
+        except Exception as e:
+            logger.error(f"close HiCacheHF3FS: {e}")
+        logger.info("close HiCacheHF3FS")

sglang/srt/mem_cache/storage/hf3fs/test_hf3fs_utils.py

@@ -0,0 +1,43 @@
+import multiprocessing.shared_memory
+from pathlib import Path
+
+import pytest
+import torch
+from torch.utils.cpp_extension import load
+from tqdm import tqdm
+
+root = Path(__file__).parent.resolve()
+hf3fs_utils = load(
+    name="hf3fs_utils", sources=[f"{root}/hf3fs_utils.cpp"], verbose=True
+)
+
+
+def test_rw_shm():
+    numel = 8 << 20
+    dtype = torch.bfloat16
+    page_num = 128
+    page_bytes = numel * dtype.itemsize
+    shm = multiprocessing.shared_memory.SharedMemory(
+        size=page_num * page_bytes, create=True
+    )
+    tshm = torch.frombuffer(shm.buf, dtype=torch.uint8)
+    a = [
+        torch.randn(numel, dtype=dtype)
+        for _ in tqdm(range(page_num), desc="prepare input")
+    ]
+    b = [
+        torch.empty(numel, dtype=dtype)
+        for _ in tqdm(range(page_num), desc="prepare output")
+    ]
+    hf3fs_utils.write_shm(a, tshm)
+    hf3fs_utils.read_shm(tshm, b)
+    for _a, _b in tqdm(zip(a, b), desc="assert_close"):
+        torch.testing.assert_close(_a, _b)
+
+    del tshm
+    shm.close()
+    shm.unlink()
+
+
+if __name__ == "__main__":
+    pytest.main([__file__])

sglang/srt/model_executor/model_runner.py

@@ -157,6 +157,8 @@ class ModelRunner:
         gpu_id: int,
         tp_rank: int,
         tp_size: int,
+        moe_ep_rank: int,
+        moe_ep_size: int,
         pp_rank: int,
         pp_size: int,
         nccl_port: int,
@@ -175,6 +177,8 @@ class ModelRunner:
             logger.addFilter(RankZeroFilter(tp_rank == 0))
         self.tp_rank = tp_rank
         self.tp_size = tp_size
+        self.moe_ep_rank = moe_ep_rank
+        self.moe_ep_size = moe_ep_size
         self.dp_size = server_args.dp_size
         self.pp_rank = pp_rank
         self.pp_size = pp_size
@@ -285,11 +289,21 @@ class ModelRunner:
             if architectures and not any("Llama4" in arch for arch in architectures):
                 self.is_hybrid = self.model_config.is_hybrid = True
 
-        self.start_layer = getattr(self.model, "start_layer", 0)
-        self.end_layer = getattr(
-            self.model, "end_layer", self.model_config.num_hidden_layers
+        # For MTP models like DeepSeek-V3 or GLM-4.5, the MTP layer(s) are used separately as draft
+        # models for speculative decoding. In those cases, `num_nextn_predict_layers` is used to
+        # determine the number of layers.
+        model_has_mtp_layers = self.model_config.num_nextn_predict_layers is not None
+        model_num_layers = (
+            self.model_config.num_nextn_predict_layers
+            if self.is_draft_worker and model_has_mtp_layers
+            else self.model_config.num_hidden_layers
         )
+        self.start_layer = getattr(self.model, "start_layer", 0)
+        self.end_layer = getattr(self.model, "end_layer", model_num_layers)
         self.num_effective_layers = self.end_layer - self.start_layer
+        assert (not model_has_mtp_layers) or (
+            self.num_effective_layers == model_num_layers
+        ), "PP is not compatible with MTP models."
 
         # Apply torchao quantization
         torchao_applied = getattr(self.model, "torchao_applied", False)
@@ -539,6 +553,7 @@ class ModelRunner:
             initialize_model_parallel(
                 tensor_model_parallel_size=self.tp_size,
                 pipeline_model_parallel_size=self.pp_size,
+                expert_model_parallel_size=self.moe_ep_size,
                 duplicate_tp_group=self.server_args.enable_pdmux,
             )
             initialize_dp_attention(
@@ -1178,11 +1193,7 @@ class ModelRunner:
                 dtype=self.kv_cache_dtype,
                 kv_lora_rank=self.model_config.kv_lora_rank,
                 qk_rope_head_dim=self.model_config.qk_rope_head_dim,
-                layer_num=(
-                    self.model_config.num_hidden_layers
-                    if not self.is_draft_worker
-                    else self.model_config.hf_config.num_nextn_predict_layers
-                ),  # PP is not compatible with mla backend
+                layer_num=self.num_effective_layers,
                 device=self.device,
                 enable_memory_saver=self.server_args.enable_memory_saver,
                 start_layer=self.start_layer,
@@ -1195,11 +1206,7 @@ class ModelRunner:
                 dtype=self.kv_cache_dtype,
                 kv_lora_rank=self.model_config.kv_lora_rank,
                 qk_rope_head_dim=self.model_config.qk_rope_head_dim,
-                layer_num=(
-                    self.model_config.num_hidden_layers
-                    if not self.is_draft_worker
-                    else self.model_config.hf_config.num_nextn_predict_layers
-                ),  # PP is not compatible with mla backend
+                layer_num=self.num_effective_layers,
                 device=self.device,
                 enable_memory_saver=self.server_args.enable_memory_saver,
                 start_layer=self.start_layer,