diff --git a/pyproject.toml b/pyproject.toml index 9d8fc8cff..01c30ef82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,7 +26,7 @@ dependencies = [ "accelerate", "pydantic", "pyyaml", - "sglang==0.5.9", + "sglang==0.5.14", "openai-harmony", "ninja", "packaging", diff --git a/specforge/inference/target_engine/sglang_backend/capture.py b/specforge/inference/target_engine/sglang_backend/capture.py index 2e5901e86..982e40aa2 100644 --- a/specforge/inference/target_engine/sglang_backend/capture.py +++ b/specforge/inference/target_engine/sglang_backend/capture.py @@ -36,6 +36,7 @@ from __future__ import annotations +from array import array from typing import List, Optional, Tuple import sglang.srt.managers.mm_utils as mm_utils @@ -55,8 +56,9 @@ ScheduleBatch, ) -# prepare_mlp_sync_batch_raw is a module-level function, not a Scheduler method -from sglang.srt.managers.scheduler_dp_attn_mixin import prepare_mlp_sync_batch_raw +# 0.5.14 relocated prepare_mlp_sync_batch_raw to scheduler_components.dp_attn +# (module-level function, not a Scheduler method). +from sglang.srt.managers.scheduler_components.dp_attn import prepare_mlp_sync_batch_raw from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.radix_cache import RadixCache from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch @@ -119,15 +121,18 @@ def build( original eagle3 path). """ tp_size = dist.get_world_size(get_tp_group()) - # NOTE: sglang 0.5.9 requires dtype to be non-None. If torch_dtype is None, - # use "auto" to let sglang decide the dtype. + # 0.5.13+ requires a non-None dtype; "auto" lets sglang decide. dtype_arg = torch_dtype if torch_dtype is not None else "auto" + # We prefill the whole batch in one forward, so disable chunked prefill: + # 0.5.14's eager runner otherwise caps per-call buffers at chunked_prefill_size + # (default 8192) and errors when our token count exceeds it. server_args = ServerArgs( model_path=pretrained_model_name_or_path, trust_remote_code=trust_remote_code, dtype=dtype_arg, enable_return_hidden_states=True, disable_cuda_graph=True, # we use piecewise cuda graph for prefill instead + chunked_prefill_size=-1, tp_size=tp_size, pp_size=1, **kwargs, @@ -152,6 +157,12 @@ def build( nccl_port=None, is_draft_worker=False, ) + # 0.5.14 moved post-load setup out of ModelRunner.initialize() (now + # weights-only). We drive the runner directly, so run it here: pools for + # forward(), attention backend, and the (eager) runner. + model_runner.alloc_memory_pool() + model_runner.init_attention_backends() + model_runner.init_cuda_graphs() if wrap_eagle3_logits: wrap_eagle3_logits_processors_in_module( model_runner.model, return_full_logits=return_full_logits @@ -256,8 +267,17 @@ def _forward_extend(self, reqs): ) batch.prepare_for_extend() self._maybe_prepare_mlp_sync_batch(batch) - model_worker_batch = batch.get_model_worker_batch() - forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner) + # 0.5.13+ stages input_ids on pinned CPU and leaves batch.input_ids=None; + # the scheduler normally does this H2D copy, but we bypass it. + if getattr(batch, "prefill_input_ids_cpu", None) is not None: + batch.input_ids = batch.prefill_input_ids_cpu.to( + batch.device, non_blocking=True + ) + batch.prefill_input_ids_cpu = None + # 0.5.13+ dropped ModelWorkerBatch; ForwardBatch.init_new reads + # capture_hidden_mode off the ScheduleBatch, so set it first. + batch.capture_hidden_mode = CaptureHiddenMode.FULL + forward_batch = ForwardBatch.init_new(batch, self.model_runner) forward_batch.capture_hidden_mode = CaptureHiddenMode.FULL output = self.model_runner.forward(forward_batch) return output.logits_output if hasattr(output, "logits_output") else output @@ -291,8 +311,9 @@ def _forward_eagle3_reqs( self._set_eagle3_logits_flags( return_last_hidden_states, return_logits, shard_returns ) - eagle3_output = self._forward_extend(reqs) + # Capture lengths before the forward: 0.5.13+ releases origin_input_ids in it. input_lens = [len(req.origin_input_ids) for req in reqs] + eagle3_output = self._forward_extend(reqs) logits = eagle3_output.logits aux_hidden_states = eagle3_output.aux_hidden_states @@ -370,8 +391,11 @@ def extend_eagle3( origin_input_ids=input_id_.view(-1).tolist(), sampling_params=sampling_params, ) - req.fill_ids = req.origin_input_ids - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + # 0.5.13+ replaced Req.fill_ids with full_untruncated_fill_ids + int + # fill_len (set by PrefillAdder, which we bypass; no prefix reuse). + req.full_untruncated_fill_ids = array("q", req.origin_input_ids) + req.fill_len = len(req.full_untruncated_fill_ids) + req.extend_input_len = req.fill_len - len(req.prefix_indices) req.logprob_start_len = len(req.origin_input_ids) - 1 data_cache.append([input_id_, attention_mask_, loss_mask_]) reqs.append(req) @@ -542,8 +566,11 @@ def extend_eagle3_vlm( origin_input_ids=input_id_list, sampling_params=sampling_params, ) - req.fill_ids = req.origin_input_ids - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + # 0.5.13+ replaced Req.fill_ids with full_untruncated_fill_ids + int + # fill_len (set by PrefillAdder, which we bypass; no prefix reuse). + req.full_untruncated_fill_ids = array("q", req.origin_input_ids) + req.fill_len = len(req.full_untruncated_fill_ids) + req.extend_input_len = req.fill_len - len(req.prefix_indices) req.logprob_start_len = len(req.origin_input_ids) - 1 req.multimodal_inputs = mm_inputs data_cache.append([input_id_, attention_mask_, loss_mask_]) @@ -595,13 +622,17 @@ def extend_dflash( origin_input_ids=curr_ids.view(-1).tolist(), sampling_params=sampling_params, ) - req.fill_ids = req.origin_input_ids - req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + # 0.5.13+ replaced Req.fill_ids with full_untruncated_fill_ids + int + # fill_len (set by PrefillAdder, which we bypass; no prefix reuse). + req.full_untruncated_fill_ids = array("q", req.origin_input_ids) + req.fill_len = len(req.full_untruncated_fill_ids) + req.extend_input_len = req.fill_len - len(req.prefix_indices) data_cache.append((curr_ids, curr_attn, curr_loss)) reqs.append(req) - output = self._forward_extend(reqs) + # Capture lengths before the forward: 0.5.13+ releases origin_input_ids in it. input_lens = [len(req.origin_input_ids) for req in reqs] + output = self._forward_extend(reqs) if ( hasattr(output, "aux_hidden_states") and output.aux_hidden_states is not None diff --git a/specforge/inference/target_engine/sglang_backend/patch.py b/specforge/inference/target_engine/sglang_backend/patch.py index 1ec608f8e..137e3bd23 100644 --- a/specforge/inference/target_engine/sglang_backend/patch.py +++ b/specforge/inference/target_engine/sglang_backend/patch.py @@ -131,7 +131,7 @@ def initialize_model_parallel( group_ranks.append(ranks) # message queue broadcaster is only used in tensor model parallel group - # NOTE: torch_compile parameter was removed in sglang 0.5.9 + # NOTE: torch_compile removed in 0.5.9; pynccl_use_current_stream removed in 0.5.13 parallel_state._TP = init_model_parallel_group( group_ranks, parallel_state._WORLD.local_rank, @@ -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: @@ -355,11 +353,14 @@ def initialize_dp_attention( dp_attention._ENABLE_DP_ATTENTION_FLAG = enable_dp_attention - # NOTE: Added attn_cp_size parameter for sglang 0.5.9 + # NOTE: attn_cp_size added in 0.5.9. 0.5.13: compute_dp_attention_world_info + # returns a 4-tuple; attn-tp rank/size are no longer module globals, so we keep + # only _ATTN_DP_RANK (mirrors sglang's initialize_dp_attention). ( - dp_attention._ATTN_TP_RANK, - dp_attention._ATTN_TP_SIZE, + _, + _, dp_attention._ATTN_DP_RANK, + _, ) = compute_dp_attention_world_info( enable_dp_attention, tp_rank, tp_size, dp_size, attn_cp_size ) diff --git a/specforge/inference/target_engine/sglang_backend/utils.py b/specforge/inference/target_engine/sglang_backend/utils.py index 8af3bce92..816c8f99e 100644 --- a/specforge/inference/target_engine/sglang_backend/utils.py +++ b/specforge/inference/target_engine/sglang_backend/utils.py @@ -53,17 +53,21 @@ def replaced_logits_processor_forward_for_eagle3( Updated for sglang 0.5.9: - Added hidden_states_before_norm parameter for compatibility """ + # 0.5.13 moved the multi-item delimiter off the LogitsProcessor onto the + # ForwardBatch (multi_item_delimiter_indices); read it before the conversion. + multi_item_delimiter_indices = None if isinstance(logits_metadata, ForwardBatch): + multi_item_delimiter_indices = logits_metadata.multi_item_delimiter_indices logits_metadata = LogitsMetadata.from_forward_batch(logits_metadata) # Multi-item scoring only for prefill-only requests. - if self.multi_item_delimiter is not None and logits_metadata.is_prefill_only: + if multi_item_delimiter_indices is not None and logits_metadata.is_prefill_only: return self.compute_logprobs_for_multi_item_scoring( input_ids, hidden_states, lm_head, logits_metadata, - self.multi_item_delimiter, + multi_item_delimiter_indices, ) # Diffusion LLM only.