sglang 0.4.4.post1__py3-none-any.whl → 0.4.4.post3__py3-none-any.whl

sglang/srt/utils.py

@@ -36,12 +36,13 @@ import tempfile
 import threading
 import time
 import warnings
+from contextlib import contextmanager
 from functools import lru_cache
 from importlib.metadata import PackageNotFoundError, version
 from importlib.util import find_spec
 from io import BytesIO
-from multiprocessing import Pool
 from multiprocessing.reduction import ForkingPickler
+from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional, Protocol, Set, Tuple, Union
 
 import numpy as np
@@ -54,13 +55,13 @@ import triton
 import zmq
 from fastapi.responses import ORJSONResponse
 from packaging import version as pkg_version
-from packaging.version import Version, parse
+from PIL import Image
 from starlette.routing import Mount
 from torch import nn
 from torch.func import functional_call
 from torch.library import Library
 from torch.profiler import ProfilerActivity, profile, record_function
-from torch.utils.cpp_extension import CUDA_HOME
+from torch.utils._contextlib import _DecoratorContextManager
 from triton.runtime.cache import (
     FileCacheManager,
     default_cache_dir,
@@ -76,6 +77,11 @@ time_infos = {}
 HIP_FP8_E4M3_FNUZ_MAX = 224.0
 
 
+def get_bool_env_var(name: str, default: str = "false") -> bool:
+    value = os.getenv(name, default)
+    return value.lower() in ("true", "1")
+
+
 # https://pytorch.org/docs/stable/notes/hip.html#checking-for-hip
 def is_hip() -> bool:
     return torch.version.hip is not None
@@ -126,6 +132,63 @@ def is_cuda_available():
     return is_cuda()
 
 
+_ENABLE_TORCH_INFERENCE_MODE = get_bool_env_var(
+    "SGLANG_ENABLE_TORCH_INFERENCE_MODE", "false"
+)
+
+
+class DynamicGradMode(_DecoratorContextManager):
+    """
+    A combination of torch.no_grad and torch.inference_mode,
+    with their behavior controlled by an environment variable. Just refer to them.
+    """
+
+    @staticmethod
+    def set_inference_mode(mode: bool):
+        if isinstance(mode, bool):
+            global _ENABLE_TORCH_INFERENCE_MODE
+
+            _ENABLE_TORCH_INFERENCE_MODE = mode
+        else:
+            logger.warning("mode is not a boolean object")
+
+    def __init__(self, mode=True):
+        if not torch._jit_internal.is_scripting():
+            super().__init__()
+        if _ENABLE_TORCH_INFERENCE_MODE:
+            self.mode = mode
+        else:
+            self.prev = False
+
+    def __new__(cls, mode_or_orig_func=True if _ENABLE_TORCH_INFERENCE_MODE else None):
+        if mode_or_orig_func is None or isinstance(mode_or_orig_func, bool):
+            return super().__new__(cls)
+        return cls()(mode_or_orig_func)
+
+    def __enter__(self) -> None:
+        if _ENABLE_TORCH_INFERENCE_MODE:
+            self._inference_mode_context = torch._C._InferenceMode(self.mode)
+            self._inference_mode_context.__enter__()
+        else:
+            self.prev = torch.is_grad_enabled()
+            torch.set_grad_enabled(False)
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        if _ENABLE_TORCH_INFERENCE_MODE:
+            self._inference_mode_context.__exit__(exc_type, exc_value, traceback)
+        else:
+            torch.set_grad_enabled(self.prev)
+
+    def clone(self) -> "DynamicGradMode":
+        r"""
+        Create a copy of this class
+        """
+        if _ENABLE_TORCH_INFERENCE_MODE:
+            return self.__class__(self.mode)
+        else:
+            return self.__class__()
+
+
 def enable_show_time_cost():
     global show_time_cost
     show_time_cost = True
@@ -198,7 +261,7 @@ def get_available_gpu_memory(device, gpu_id, distributed=False, empty_cache=True
     When distributed is True, the available memory is the minimum available memory of all GPUs.
     """
     if device == "cuda":
-        num_gpus = torch.cuda.device_count()
+        num_gpus = cuda_device_count_stateless()
         assert gpu_id < num_gpus
 
         if torch.cuda.current_device() != gpu_id:
@@ -443,17 +506,46 @@ def decode_video_base64(video_base64):
         )  # Return an empty array and size tuple if no frames were found
 
 
-def load_image(image_file: Union[str, bytes]):
-    from PIL import Image
+def load_audio(audio_file: str, sr: int = 16000, mono: bool = True) -> np.ndarray:
+    # Use soundfile here, since librosa use it under the hood,
+    # and librosa will not support audio loading in the future
+    import soundfile as sf
+    from scipy.signal import resample
+
+    # print(f"loading {audio_file}")
+    # Load audio data
+    if isinstance(audio_file, bytes):
+        audio, original_sr = sf.read(BytesIO(audio_file))
+    elif audio_file.startswith("data:"):
+        audio_file = audio_file.split(",")[1]
+        audio, original_sr = sf.read(BytesIO(base64.b64decode(audio_file)))
+    elif isinstance(audio_file, str):
+        audio, original_sr = sf.read(audio_file)
+    else:
+        raise ValueError(f"Invalid audio format: {audio_file}")
 
+    # Resample audio if the original sample rate is different from the desired sample rate
+    if original_sr != sr:
+        num_samples = int(len(audio) * float(sr) / original_sr)
+        audio = resample(audio, num_samples)
+
+    # Convert to mono if requested and audio is stereo
+    if mono and len(audio.shape) > 1:
+        audio = np.mean(audio, axis=1)
+
+    return audio
+
+
+def load_image(image_file: Union[str, bytes]) -> tuple[Image, tuple[int, int]]:
     image = image_size = None
 
     if isinstance(image_file, bytes):
         image = Image.open(BytesIO(image_file))
     elif image_file.startswith("http://") or image_file.startswith("https://"):
         timeout = int(os.getenv("REQUEST_TIMEOUT", "3"))
-        response = requests.get(image_file, timeout=timeout)
-        image = Image.open(BytesIO(response.content))
+        response = requests.get(image_file, stream=True, timeout=timeout).raw
+        image = Image.open(response)
+        response.close()
     elif image_file.lower().endswith(("png", "jpg", "jpeg", "webp", "gif")):
         image = Image.open(image_file)
     elif image_file.startswith("data:"):
@@ -471,7 +563,10 @@ def load_image(image_file: Union[str, bytes]):
 
 
 def suppress_other_loggers():
-    from vllm.logger import logger as vllm_default_logger
+    try:
+        from vllm.logger import logger as vllm_default_logger
+    except ImportError:
+        return
 
     vllm_default_logger.setLevel(logging.WARN)
     logging.getLogger("vllm.distributed.device_communicators.pynccl").setLevel(
@@ -480,6 +575,7 @@ def suppress_other_loggers():
     logging.getLogger("vllm.distributed.device_communicators.shm_broadcast").setLevel(
         logging.WARN
     )
+    logging.getLogger("vllm.config").setLevel(logging.ERROR)
 
     warnings.filterwarnings(
         "ignore", category=UserWarning, message="The given NumPy array is not writable"
@@ -527,6 +623,10 @@ def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = N
 
     if include_parent:
         try:
+            if parent_pid == os.getpid():
+                itself.kill()
+                sys.exit(0)
+
             itself.kill()
 
             # Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes),
@@ -555,11 +655,14 @@ def monkey_patch_p2p_access_check():
 
 
 def monkey_patch_vllm_gguf_config():
-    from vllm.model_executor.layers.quantization.gguf import (
-        GGUFConfig,
-        GGUFEmbeddingMethod,
-        GGUFLinearMethod,
-    )
+    try:
+        from vllm.model_executor.layers.quantization.gguf import (
+            GGUFConfig,
+            GGUFEmbeddingMethod,
+            GGUFLinearMethod,
+        )
+    except ImportError:
+        return
 
     from sglang.srt.layers.linear import LinearBase
     from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding
@@ -651,6 +754,16 @@ def prepare_model_and_tokenizer(model_path: str, tokenizer_path: str):
 
 
 def configure_logger(server_args, prefix: str = ""):
+    if SGLANG_LOGGING_CONFIG_PATH := os.getenv("SGLANG_LOGGING_CONFIG_PATH"):
+        if not os.path.exists(SGLANG_LOGGING_CONFIG_PATH):
+            raise Exception(
+                "Setting SGLANG_LOGGING_CONFIG_PATH from env with "
+                f"{SGLANG_LOGGING_CONFIG_PATH} but it does not exist!"
+            )
+        with open(SGLANG_LOGGING_CONFIG_PATH, encoding="utf-8") as file:
+            custom_config = json.loads(file.read())
+        logging.config.dictConfig(custom_config)
+        return
     format = f"[%(asctime)s{prefix}] %(message)s"
     # format = f"[%(asctime)s.%(msecs)03d{prefix}] %(message)s"
     logging.basicConfig(
@@ -774,12 +887,22 @@ def get_zmq_socket(
         buf_size = -1
 
     socket = context.socket(socket_type)
-    if socket_type == zmq.PUSH:
+
+    def set_send_opt():
         socket.setsockopt(zmq.SNDHWM, 0)
         socket.setsockopt(zmq.SNDBUF, buf_size)
-    elif socket_type == zmq.PULL:
+
+    def set_recv_opt():
         socket.setsockopt(zmq.RCVHWM, 0)
         socket.setsockopt(zmq.RCVBUF, buf_size)
+
+    if socket_type == zmq.PUSH:
+        set_send_opt()
+    elif socket_type == zmq.PULL:
+        set_recv_opt()
+    elif socket_type == zmq.DEALER:
+        set_send_opt()
+        set_recv_opt()
     else:
         raise ValueError(f"Unsupported socket type: {socket_type}")
 
@@ -910,6 +1033,13 @@ def get_amdgpu_memory_capacity():
         )
 
 
+def get_device_sm():
+    if torch.cuda.is_available():
+        major, minor = torch.cuda.get_device_capability()
+        return major * 10 + minor
+    return 0
+
+
 def get_nvgpu_memory_capacity():
     try:
         # Run nvidia-smi and capture the output
@@ -1246,11 +1376,6 @@ def set_gpu_proc_affinity(
     logger.info(f"Process {pid} gpu_id {gpu_id} is running on CPUs: {p.cpu_affinity()}")
 
 
-def get_bool_env_var(name: str, default: str = "false") -> bool:
-    value = os.getenv(name, default)
-    return value.lower() in ("true", "1")
-
-
 @lru_cache(maxsize=2)
 def disable_request_logging() -> bool:
     return get_bool_env_var("SGLANG_DISABLE_REQUEST_LOGGING")
@@ -1477,6 +1602,7 @@ def get_ip() -> str:
 def get_open_port() -> int:
     port = os.getenv("SGLANG_PORT")
     if port is not None:
+        port = int(port)
         while True:
             try:
                 with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -1505,6 +1631,38 @@ def is_valid_ipv6_address(address: str) -> bool:
         return False
 
 
+def configure_ipv6(dist_init_addr):
+    addr = dist_init_addr
+    end = addr.find("]")
+    if end == -1:
+        raise ValueError("invalid IPv6 address format: missing ']'")
+
+    host = addr[: end + 1]
+
+    # this only validates the address without brackets: we still need the below checks.
+    # if it's invalid, immediately raise an error so we know it's not formatting issues.
+    if not is_valid_ipv6_address(host[1:end]):
+        raise ValueError(f"invalid IPv6 address: {host}")
+
+    port_str = None
+    if len(addr) > end + 1:
+        if addr[end + 1] == ":":
+            port_str = addr[end + 2 :]
+        else:
+            raise ValueError("received IPv6 address format: expected ':' after ']'")
+
+    if not port_str:
+        raise ValueError(
+            "a port must be specified in IPv6 address (format: [ipv6]:port)"
+        )
+
+    try:
+        port = int(port_str)
+    except ValueError:
+        raise ValueError(f"invalid port in IPv6 address: '{port_str}'")
+    return port, host
+
+
 def rank0_print(msg: str):
     from sglang.srt.distributed import get_tensor_model_parallel_rank
 
@@ -1561,6 +1719,16 @@ def next_power_of_2(n: int):
 setattr(triton, "next_power_of_2", next_power_of_2)
 
 
+@contextmanager
+def empty_context(*args, **kwargs):
+    try:
+        # Setup code goes here
+        yield
+    finally:
+        # Cleanup code goes here
+        pass
+
+
 def add_prefix(name: str, prefix: str) -> str:
     """Add a weight path prefix to a module name.
 
@@ -1572,3 +1740,29 @@ def add_prefix(name: str, prefix: str) -> str:
         The string `prefix.name` if prefix is non-empty, otherwise just `name`.
     """
     return name if not prefix else f"{prefix}.{name}"
+
+
+def is_remote_url(url: Union[str, Path]) -> bool:
+    """
+    Check if the URL is a remote URL of the format:
+    <connector_type>://<host>:<port>/<model_name>
+    """
+    if isinstance(url, Path):
+        return False
+
+    pattern = r"(.+)://(.*)"
+    m = re.match(pattern, url)
+    return m is not None
+
+
+def parse_connector_type(url: str) -> str:
+    """
+    Parse the connector type from the URL of the format:
+    <connector_type>://<path>
+    """
+    pattern = r"(.+)://(.*)"
+    m = re.match(pattern, url)
+    if m is None:
+        return ""
+
+    return m.group(1)

sglang/test/attention/test_flashattn_backend.py

@@ -0,0 +1,312 @@
+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 MHATokenToKVPool
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+from sglang.test.test_utils import CustomTestCase
+
+
+class MockModelRunner:
+    model_config = type(
+        "ModelConfig", (), {"context_len": 2048, "is_multimodal": False}
+    )
+    sliding_window_size = None
+
+    def __init__(self, device="cuda"):
+        self.device = device
+        # Create a proper req_to_token_pool with the req_to_token attribute
+        self.req_to_token_pool = type(
+            "TokenPool",
+            (),
+            {
+                "size": 160,  # a typical max_bs * max_context_len for cuda graph decode
+                "req_to_token": torch.zeros(
+                    160, 2048, dtype=torch.int32, device=device
+                ),  # Add req_to_token attribute
+            },
+        )
+
+
+class MockReqToTokenPool:
+    def __init__(self, batch_size, seq_len, device):
+        self.req_to_token = (
+            torch.arange(batch_size * seq_len, device=device)
+            .reshape(batch_size, seq_len)
+            .to(torch.int32)
+        )
+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestFlashAttentionBackend(CustomTestCase):
+    def setUp(self):
+        """Set up test fixtures before each test method."""
+        self.model_runner = MockModelRunner()
+        self.backend = FlashAttentionBackend(self.model_runner)
+
+        # Common test parameters
+        self.batch_size = 2
+        self.seq_len = 4
+        self.num_heads = 2
+        self.head_dim = 8
+        self.device = "cuda"
+        self.dtype = torch.float16
+
+    def _create_attention_layer(self):
+        """Helper method to create an attention layer."""
+        return RadixAttention(
+            num_heads=self.num_heads,
+            head_dim=self.head_dim,
+            scaling=1.0,
+            num_kv_heads=self.num_heads,
+            layer_id=0,
+        )
+
+    def _create_kv_pool(self, size):
+        """Helper method to create a KV pool."""
+        return MHATokenToKVPool(
+            size=size,
+            page_size=1,  # only consider page=1 for unit test
+            dtype=self.dtype,
+            head_num=self.num_heads,
+            head_dim=self.head_dim,
+            layer_num=1,  # only consider layer=1 for unit test
+            device=self.device,
+            enable_memory_saver=False,
+        )
+
+    def _create_qkv_tensors(self, tokens_len):
+        """Helper method to create q, k, v tensors."""
+        return (
+            torch.randn(
+                tokens_len,
+                self.num_heads,
+                self.head_dim,
+                dtype=self.dtype,
+                device=self.device,
+            ),
+            torch.randn(
+                tokens_len,
+                self.num_heads,
+                self.head_dim,
+                dtype=self.dtype,
+                device=self.device,
+            ),
+            torch.randn(
+                tokens_len,
+                self.num_heads,
+                self.head_dim,
+                dtype=self.dtype,
+                device=self.device,
+            ),
+        )
+
+    def _verify_output(self, output, expected_shape):
+        """Helper method to verify output."""
+        self.assertEqual(
+            output.shape,
+            expected_shape,
+            f"Expected shape {expected_shape}, got {output.shape}",
+        )
+        self.assertEqual(output.dtype, self.dtype)
+        self.assertEqual(output.device.type, "cuda")
+        self.assertEqual(
+            torch.isnan(output).sum().item(), 0, "Output contains NaN values"
+        )
+
+    def test_forward_extend(self):
+        """Test the standard extend operation."""
+        # Create test inputs
+        q, k, v = self._create_qkv_tensors(self.batch_size * self.seq_len)
+
+        # Create attention layer
+        layer = self._create_attention_layer()
+
+        # Create forward batch
+        forward_batch = ForwardBatch(
+            batch_size=self.batch_size,
+            input_ids=torch.randint(
+                0, 100, (self.batch_size, self.seq_len), device=self.device
+            ),
+            out_cache_loc=torch.arange(
+                self.batch_size * self.seq_len, device=self.device
+            ),
+            seq_lens_sum=self.batch_size * self.seq_len,
+            forward_mode=ForwardMode.EXTEND,
+            req_pool_indices=torch.arange(self.batch_size, device=self.device),
+            seq_lens=torch.tensor([self.seq_len] * self.batch_size, device=self.device),
+            # 0 prefix, 4 extend
+            extend_prefix_lens=torch.tensor([0] * self.batch_size, device=self.device),
+            extend_seq_lens=torch.tensor([4] * self.batch_size, device=self.device),
+            attn_backend=self.backend,
+        )
+
+        # Add token pool and KV cache
+        forward_batch.req_to_token_pool = MockReqToTokenPool(
+            self.batch_size, self.seq_len, self.device
+        )
+        forward_batch.token_to_kv_pool = self._create_kv_pool(
+            self.batch_size * self.seq_len
+        )
+
+        # Initialize forward metadata before running the attention
+        self.backend.init_forward_metadata(forward_batch)
+
+        # Run forward_extend
+        output = self.backend.forward_extend(q, k, v, layer, forward_batch)
+
+        # Verify output
+        expected_shape = (
+            self.batch_size * self.seq_len,
+            self.num_heads * self.head_dim,
+        )
+        self._verify_output(output, expected_shape)
+
+    def test_forward_decode(self):
+        """Test the decode operation with cached tokens."""
+        # For decode, we only have one token per sequence
+        decode_len = 1
+        curr_seq_len = self.seq_len + decode_len
+
+        # Create test inputs
+        q, k, v = self._create_qkv_tensors(self.batch_size * decode_len)
+
+        # Create attention layer
+        layer = self._create_attention_layer()
+
+        # Create forward batch
+        forward_batch = ForwardBatch(
+            batch_size=self.batch_size,
+            input_ids=torch.randint(
+                0, 100, (self.batch_size, decode_len), device=self.device
+            ),
+            out_cache_loc=torch.arange(
+                self.batch_size * self.seq_len,
+                self.batch_size * curr_seq_len,
+                device=self.device,
+            ),
+            seq_lens_sum=self.batch_size * curr_seq_len,
+            forward_mode=ForwardMode.DECODE,
+            req_pool_indices=torch.arange(self.batch_size, device=self.device),
+            seq_lens=torch.tensor([curr_seq_len] * self.batch_size, device=self.device),
+            attn_backend=self.backend,
+        )
+
+        # Add token pool and KV cache
+        forward_batch.req_to_token_pool = MockReqToTokenPool(
+            self.batch_size, curr_seq_len, self.device
+        )
+        forward_batch.token_to_kv_pool = self._create_kv_pool(
+            self.batch_size * curr_seq_len
+        )
+
+        # Pre-fill KV cache
+        cache_k, cache_v, _ = self._create_qkv_tensors(self.batch_size * self.seq_len)
+        forward_batch.token_to_kv_pool.set_kv_buffer(
+            layer,
+            torch.arange(self.batch_size * self.seq_len, device=self.device),
+            cache_k,
+            cache_v,
+            layer.k_scale,
+            layer.v_scale,
+        )
+
+        # Initialize forward metadata before running the attention
+        self.backend.init_forward_metadata(forward_batch)
+
+        # Run forward_decode
+        output = self.backend.forward_decode(q, k, v, layer, forward_batch)
+
+        # Verify output
+        expected_shape = (self.batch_size, self.num_heads * self.head_dim)
+        self._verify_output(output, expected_shape)
+
+    def test_forward_extend_with_prefix(self):
+        """Test extending from cached prefix tokens."""
+        # Define prefix and extend lengths
+        prefix_len = 2
+        extend_len = 2
+        total_len = prefix_len + extend_len
+
+        # Create test inputs for the extend portion
+        q, k, v = self._create_qkv_tensors(self.batch_size * extend_len)
+
+        # Create attention layer
+        layer = self._create_attention_layer()
+
+        # Create forward batch
+        forward_batch = ForwardBatch(
+            batch_size=self.batch_size,
+            input_ids=torch.randint(
+                0, 100, (self.batch_size, extend_len), device=self.device
+            ),
+            out_cache_loc=torch.arange(
+                self.batch_size * prefix_len,
+                self.batch_size * total_len,
+                device=self.device,
+            ),
+            seq_lens_sum=self.batch_size * total_len,
+            forward_mode=ForwardMode.EXTEND,
+            req_pool_indices=torch.arange(self.batch_size, device=self.device),
+            seq_lens=torch.tensor([total_len] * self.batch_size, device=self.device),
+            extend_prefix_lens=torch.tensor(
+                [prefix_len] * self.batch_size, device=self.device
+            ),
+            extend_seq_lens=torch.tensor(
+                [extend_len] * self.batch_size, device=self.device
+            ),
+            attn_backend=self.backend,
+        )
+
+        # Add token pool and KV cache
+        forward_batch.req_to_token_pool = MockReqToTokenPool(
+            self.batch_size, total_len, self.device
+        )
+        forward_batch.token_to_kv_pool = self._create_kv_pool(
+            self.batch_size * total_len
+        )
+
+        # Pre-fill the KV cache for prefix with known values
+        cache_k = torch.ones(
+            self.batch_size * prefix_len,
+            self.num_heads,
+            self.head_dim,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        cache_v = (
+            torch.ones(
+                self.batch_size * prefix_len,
+                self.num_heads,
+                self.head_dim,
+                dtype=self.dtype,
+                device=self.device,
+            )
+            * 2
+        )
+
+        # Set the prefix KV cache
+        forward_batch.token_to_kv_pool.set_kv_buffer(
+            layer,
+            torch.arange(self.batch_size * prefix_len, device=self.device),
+            cache_k,
+            cache_v,
+            layer.k_scale,
+            layer.v_scale,
+        )
+
+        # Initialize forward metadata before running the attention
+        self.backend.init_forward_metadata(forward_batch)
+
+        # Run forward_extend
+        output = self.backend.forward_extend(q, k, v, layer, forward_batch)
+
+        # Verify output
+        expected_shape = (self.batch_size * extend_len, self.num_heads * self.head_dim)
+        self._verify_output(output, expected_shape)
+
+
+if __name__ == "__main__":
+    unittest.main()