sglang 0.5.3rc0__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/model_loader/weight_utils.py

@@ -8,7 +8,7 @@ import hashlib
 import json
 import logging
 import os
-import queue
+import re
 import tempfile
 from collections import defaultdict
 from typing import (
@@ -38,7 +38,8 @@ from sglang.srt.distributed import get_tensor_model_parallel_rank
 from sglang.srt.layers.dp_attention import get_attention_tp_rank
 from sglang.srt.layers.quantization import QuantizationConfig, get_quantization_config
 from sglang.srt.layers.quantization.modelopt_quant import ModelOptFp4Config
-from sglang.srt.utils import print_warning_once
+from sglang.srt.utils import find_local_repo_dir, print_warning_once
+from sglang.utils import is_in_ci
 
 logger = logging.getLogger(__name__)
 
@@ -236,6 +237,149 @@ def get_quant_config(
     return quant_cls.from_config(config)
 
 
+def find_local_hf_snapshot_dir(
+    model_name_or_path: str,
+    cache_dir: Optional[str],
+    allow_patterns: List[str],
+    revision: Optional[str] = None,
+) -> Optional[str]:
+    """If the weights are already local, skip downloading and returns the path."""
+    if os.path.isdir(model_name_or_path):
+        return None
+
+    found_local_snapshot_dir = None
+
+    # Check custom cache_dir (if provided)
+    if cache_dir:
+        try:
+            repo_folder = os.path.join(
+                cache_dir,
+                huggingface_hub.constants.REPO_ID_SEPARATOR.join(
+                    ["models", *model_name_or_path.split("/")]
+                ),
+            )
+            rev_to_use = revision
+            if not rev_to_use:
+                ref_main = os.path.join(repo_folder, "refs", "main")
+                if os.path.isfile(ref_main):
+                    with open(ref_main) as f:
+                        rev_to_use = f.read().strip()
+            if rev_to_use:
+                rev_dir = os.path.join(repo_folder, "snapshots", rev_to_use)
+                if os.path.isdir(rev_dir):
+                    found_local_snapshot_dir = rev_dir
+        except Exception as e:
+            logger.warning(
+                "Failed to find local snapshot in custom cache_dir %s: %s",
+                cache_dir,
+                e,
+            )
+
+    # Check default HF cache as well
+    if not found_local_snapshot_dir:
+        try:
+            rev_dir = find_local_repo_dir(model_name_or_path, revision)
+            if rev_dir and os.path.isdir(rev_dir):
+                found_local_snapshot_dir = rev_dir
+        except Exception as e:
+            logger.warning("Failed to find local snapshot in default HF cache: %s", e)
+
+    # if any incomplete file exists, force re-download by returning None
+    if found_local_snapshot_dir:
+        repo_folder = os.path.abspath(
+            os.path.join(found_local_snapshot_dir, "..", "..")
+        )
+        blobs_dir = os.path.join(repo_folder, "blobs")
+        if os.path.isdir(blobs_dir) and glob.glob(
+            os.path.join(blobs_dir, "*.incomplete")
+        ):
+            logger.info(
+                "Found .incomplete files in %s for %s. "
+                "Considering local snapshot incomplete.",
+                blobs_dir,
+                model_name_or_path,
+            )
+            return None
+
+    # if local snapshot exists, validate it contains at least one weight file
+    # matching allow_patterns before skipping download.
+    if found_local_snapshot_dir is None:
+        return None
+
+    local_weight_files: List[str] = []
+    try:
+        for pattern in allow_patterns:
+            matched_files = glob.glob(os.path.join(found_local_snapshot_dir, pattern))
+            for f in matched_files:
+                # os.path.exists returns False for broken symlinks.
+                if not os.path.exists(f):
+                    continue
+                local_weight_files.append(f)
+    except Exception as e:
+        logger.warning(
+            "Failed to scan local snapshot %s with patterns %s: %s",
+            found_local_snapshot_dir,
+            allow_patterns,
+            e,
+        )
+        local_weight_files = []
+
+    # After we have a list of valid files, check for sharded model completeness.
+    # Check if all safetensors with name model-{i}-of-{n}.safetensors exists
+    checked_sharded_model = False
+    for f in local_weight_files:
+        if checked_sharded_model:
+            break
+        base_name = os.path.basename(f)
+        # Regex for files like model-00001-of-00009.safetensors
+        match = re.match(r"(.*?)-([0-9]+)-of-([0-9]+)\.(.*)", base_name)
+        if match:
+            prefix = match.group(1)
+            shard_id_str = match.group(2)
+            total_shards_str = match.group(3)
+            suffix = match.group(4)
+            total_shards = int(total_shards_str)
+
+            # Check if all shards are present
+            missing_shards = []
+            for i in range(1, total_shards + 1):
+                # Reconstruct shard name, preserving padding of original shard id
+                shard_name = (
+                    f"{prefix}-{i:0{len(shard_id_str)}d}-of-{total_shards_str}.{suffix}"
+                )
+                expected_path = os.path.join(found_local_snapshot_dir, shard_name)
+                # os.path.exists returns False for broken symlinks, which is desired.
+                if not os.path.exists(expected_path):
+                    missing_shards.append(shard_name)
+
+            if missing_shards:
+                logger.info(
+                    "Found incomplete sharded model %s. Missing shards: %s. "
+                    "Will attempt download.",
+                    model_name_or_path,
+                    missing_shards,
+                )
+                return None
+
+            # If we found and verified one set of shards, we are done.
+            checked_sharded_model = True
+
+    if len(local_weight_files) > 0:
+        logger.info(
+            "Found local HF snapshot for %s at %s; skipping download.",
+            model_name_or_path,
+            found_local_snapshot_dir,
+        )
+        return found_local_snapshot_dir
+    else:
+        logger.info(
+            "Local HF snapshot at %s has no files matching %s; will attempt download.",
+            found_local_snapshot_dir,
+            allow_patterns,
+        )
+    return None
+
+
 def download_weights_from_hf(
     model_name_or_path: str,
     cache_dir: Optional[str],
@@ -260,6 +404,16 @@ def download_weights_from_hf(
     Returns:
         str: The path to the downloaded model weights.
     """
+
+    if is_in_ci():
+        # If the weights are already local, skip downloading and returns the path.
+        # This is used to skip too-many Huggingface API calls in CI.
+        path = find_local_hf_snapshot_dir(
+            model_name_or_path, cache_dir, allow_patterns, revision
+        )
+        if path is not None:
+            return path
+
     if not huggingface_hub.constants.HF_HUB_OFFLINE:
         # Before we download we look at that is available:
         fs = HfFileSystem()

sglang/srt/models/bailing_moe.py

@@ -45,12 +45,12 @@ from sglang.srt.layers.dp_attention import (
     get_attention_dp_size,
     get_attention_tp_rank,
     get_attention_tp_size,
+    is_dp_attention_enabled,
 )
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
-    ReplicatedLinear,
     RowParallelLinear,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -72,6 +72,10 @@ from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_loader.weight_utils import default_weight_loader
+from sglang.srt.models.utils import (
+    create_fused_set_kv_buffer_arg,
+    enable_fused_set_kv_buffer,
+)
 from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers
 
 LoraConfig = None
@@ -555,8 +559,27 @@ class BailingMoEAttention(nn.Module):
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
         if self.use_qk_norm:
             q, k = self._apply_qk_norm(q, k)
-        q, k = self.rotary_emb(positions, q, k)
-        context_layer = self.attn(q, k, v, forward_batch)
+        q, k = self.rotary_emb(
+            positions,
+            q,
+            k,
+            fused_set_kv_buffer_arg=(
+                create_fused_set_kv_buffer_arg(
+                    value=v,
+                    layer=self.attn,
+                    forward_batch=forward_batch,
+                )
+                if enable_fused_set_kv_buffer(forward_batch)
+                else None
+            ),
+        )
+        context_layer = self.attn(
+            q,
+            k,
+            v,
+            forward_batch,
+            save_kv_cache=not enable_fused_set_kv_buffer(forward_batch),
+        )
         attn_output, _ = self.dense(context_layer)
         return attn_output
 
@@ -702,7 +725,7 @@ class BailingMoEModel(nn.Module):
                 self.embed_dim,
                 quant_config=quant_config,
                 prefix=add_prefix("word_embeddings", prefix),
-                use_attn_tp_group=global_server_args_dict["enable_dp_lm_head"],
+                enable_tp=not is_dp_attention_enabled(),
             )
         else:
             self.word_embeddings = PPMissingLayer()

sglang/srt/models/deepseek_nextn.py

@@ -33,11 +33,14 @@ from sglang.srt.layers.vocab_parallel_embedding import (
 from sglang.srt.managers.schedule_batch import global_server_args_dict
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.models.deepseek_v2 import DeepseekV2DecoderLayer, DeepseekV3ForCausalLM
-from sglang.srt.utils import BumpAllocator, add_prefix
+from sglang.srt.utils import BumpAllocator, add_prefix, is_cuda
 
 logger = logging.getLogger(__name__)
 
 
+_is_cuda = is_cuda()
+
+
 class DeepseekModelNextN(nn.Module):
     def __init__(
         self,
@@ -66,12 +69,14 @@ class DeepseekModelNextN(nn.Module):
 
         self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
 
+        self.alt_stream = torch.cuda.Stream() if _is_cuda else None
         self.decoder = DeepseekV2DecoderLayer(
             config,
             0,
             quant_config=quant_config,
             is_nextn=True,
             prefix=add_prefix("decoder", prefix),
+            alt_stream=self.alt_stream,
         )
 
         self.shared_head = nn.Module()