fix: Bump sglang version from 0.5.9 to 0.5.10

Author: moehanabi

pyproject.toml

@@ -15,7 +15,7 @@ dependencies = [
     "torch==2.9.1",
     "torchaudio==2.9.1",
     "torchvision==0.24.1",
-    "transformers==4.57.1",
+    "transformers==5.3.0",
     "qwen-vl-utils==0.0.11",
     "datasets",
     "setuptools",
@@ -25,7 +25,7 @@ dependencies = [
     "numpy",
     "accelerate",
     "pydantic",
-    "sglang==0.5.9",
+    "sglang==0.5.10",
     "openai-harmony",
     "ninja",
     "packaging",

requirements-rocm.txt

@@ -5,7 +5,7 @@ pre-commit
 torch==2.8.0+rocm6.3
 torchaudio==2.8.0+rocm6.3
 torchvision==0.23.0+rocm6.3
-transformers==4.57.1
+transformers==5.3.0
 qwen-vl-utils==0.0.11
 datasets
 setuptools
@@ -15,6 +15,6 @@ psutil
 numpy
 accelerate
 pydantic
-sglang[all]==0.5.4
+sglang[all]==0.5.10
 openai-harmony
 tensorboard

specforge/args.py

@@ -96,7 +96,7 @@ class SGLangBackendArgs:
     sglang_enable_torch_compile: bool = True
     sglang_enable_dp_attention: bool = False
     sglang_enable_dp_lm_head: bool = False
-    sglang_enable_piecewise_cuda_graph: bool = False
+    sglang_enforce_piecewise_cuda_graph: bool = False
     sglang_piecewise_cuda_graph_max_tokens: int = 4096
     sglang_piecewise_cuda_graph_tokens: List[int] = None
     sglang_ep_size: int = 1
@@ -151,9 +151,9 @@ def add_args(parser: argparse.ArgumentParser) -> None:
             help="Enable piecewise CUDA graph for SGLang backend",
         )
         parser.add_argument(
-            "--sglang-enable-piecewise-cuda-graph",
+            "--sglang-enforce-piecewise-cuda-graph",
             action="store_true",
-            help="Enable piecewise CUDA graph for SGLang backend's prefill",
+            help="Enforce piecewise CUDA graph for SGLang backend's prefill",
         )
         parser.add_argument(
             "--sglang-piecewise-cuda-graph-max-tokens",
@@ -186,7 +186,7 @@ def from_args(args: argparse.Namespace) -> "SGLangBackendArgs":
             sglang_enable_torch_compile=args.sglang_enable_torch_compile,
             sglang_enable_dp_attention=args.sglang_enable_dp_attention,
             sglang_enable_dp_lm_head=args.sglang_enable_dp_lm_head,
-            sglang_enable_piecewise_cuda_graph=args.sglang_enable_piecewise_cuda_graph,
+            sglang_enforce_piecewise_cuda_graph=args.sglang_enforce_piecewise_cuda_graph,
             sglang_piecewise_cuda_graph_max_tokens=args.sglang_piecewise_cuda_graph_max_tokens,
             sglang_piecewise_cuda_graph_tokens=args.sglang_piecewise_cuda_graph_tokens,
             sglang_ep_size=args.sglang_ep_size,
@@ -210,7 +210,7 @@ def to_kwargs(self) -> Dict[str, Any]:
             enable_torch_compile=self.sglang_enable_torch_compile,
             enable_dp_attention=self.sglang_enable_dp_attention,
             enable_dp_lm_head=self.sglang_enable_dp_lm_head,
-            enable_piecewise_cuda_graph=self.sglang_enable_piecewise_cuda_graph,
+            enforce_piecewise_cuda_graph=self.sglang_enforce_piecewise_cuda_graph,
             piecewise_cuda_graph_max_tokens=self.sglang_piecewise_cuda_graph_max_tokens,
             piecewise_cuda_graph_tokens=self.sglang_piecewise_cuda_graph_tokens,
             ep_size=self.sglang_ep_size,

specforge/modeling/draft/llama3_eagle.py

@@ -272,6 +272,17 @@ def _set_cos_sin_cache(self, seq_len, device, dtype):
             "sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False
         )
 
+    def rebuild_buffers(self, device):
+        """Rebuild non-persistent RoPE buffers corrupted by transformers 5.x meta-device init."""
+        self.inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2, device=device).float() / self.dim)
+        )
+        self._set_cos_sin_cache(
+            seq_len=self.max_position_embeddings + 20,
+            device=device,
+            dtype=torch.get_default_dtype(),
+        )
+
     @torch.compile(dynamic=True)
     def forward(self, x, seq_len=None):
         # x: [bs, num_attention_heads, seq_len, head_size]
@@ -1314,6 +1325,16 @@ class LlamaForCausalLMEagle3(Eagle3DraftModel):
 
     config_class = LlamaConfig
 
+    def _init_weights(self, module):
+        # Override the transformers 5.x default _init_weights which would
+        # re-randomize all Linear/Embedding weights with normal_(0, 0.02).
+        # Draft model weights come from checkpoint, not random init.
+        #
+        # For RotaryEmbedding: rebuild non-persistent buffers (inv_freq,
+        # cos_cached, sin_cached) corrupted by meta-device materialization.
+        if isinstance(module, LlamaRotaryEmbedding):
+            module.rebuild_buffers(module.inv_freq.device)
+
     def __init__(self, config, quant_config=None, attention_backend="sdpa") -> None:
         super().__init__(config)
         self.config = config
@@ -1346,6 +1367,8 @@ def __init__(self, config, quant_config=None, attention_backend="sdpa") -> None:
         self.register_buffer("t2d", t2d)
         self.register_buffer("d2t", d2t)
 
+        self.post_init()
+
     def forward(
         self,
         hidden_states: torch.Tensor,

specforge/modeling/target/custom_backend/gpt_oss.py

@@ -36,7 +36,8 @@
 from transformers.models.gpt_oss.modeling_gpt_oss import GptOssRMSNorm
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
-from transformers.utils.generic import check_model_inputs
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
 
 from specforge.distributed import get_tp_group, shard_tensor
 from specforge.layers import (
@@ -585,7 +586,8 @@ def __init__(self, config: GptOssConfig):
         # Initialize weights and apply final processing
         self.post_init()
 
-    @check_model_inputs
+    @merge_with_config_defaults
+    @capture_outputs
     @auto_docstring
     def forward(
         self,

specforge/modeling/target/custom_backend/llama.py

@@ -41,7 +41,8 @@
 )
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, logging
-from transformers.utils.generic import check_model_inputs
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
 
 from specforge.distributed import get_tp_group
 from specforge.layers import (
@@ -275,7 +276,8 @@ def __init__(self, config: LlamaConfig):
         # Initialize weights and apply final processing
         self.post_init()
 
-    @check_model_inputs
+    @merge_with_config_defaults
+    @capture_outputs
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,

specforge/modeling/target/custom_backend/llama4.py

@@ -52,7 +52,8 @@
     logging,
 )
 from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import check_model_inputs
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
 
 # [MODIFIED] Import from transformers library
 from specforge.distributed import get_tp_group, shard_tensor
@@ -431,7 +432,8 @@ def __init__(self, config: Llama4TextConfig):
         self.post_init()
 
     @can_return_tuple
-    @check_model_inputs
+    @merge_with_config_defaults
+    @capture_outputs
     @auto_docstring
     def forward(
         self,

specforge/modeling/target/custom_backend/phi3.py

@@ -43,7 +43,8 @@
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
 from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import check_model_inputs
+from transformers.utils.generic import merge_with_config_defaults
+from transformers.utils.output_capturing import capture_outputs
 
 from specforge.distributed import get_tp_group
 from specforge.layers import (
@@ -284,7 +285,8 @@ def __init__(self, config: Phi3Config):
         # Initialize weights and apply final processing
         self.post_init()
 
-    @check_model_inputs
+    @merge_with_config_defaults
+    @capture_outputs
     @auto_docstring
     def forward(
         self,

specforge/modeling/target/sglang_backend/patch.py

@@ -140,7 +140,6 @@ def initialize_model_parallel(
             "SGLANG_USE_MESSAGE_QUEUE_BROADCASTER", "true"
         ),
         group_name="tp",
-        pynccl_use_current_stream=duplicate_tp_group,
     )
 
     if duplicate_tp_group:
@@ -156,7 +155,6 @@ def initialize_model_parallel(
                 "SGLANG_USE_MESSAGE_QUEUE_BROADCASTER", "true"
             ),
             group_name="pdmux_prefill_tp",
-            pynccl_use_current_stream=True,
         )
         # NOTE: Check pynccl_comm exists before accessing it (may be None in sglang 0.5.9)
         if parallel_state._TP.pynccl_comm is not None: