From 401805dd6b6b0a425e725f4049d120942697816d Mon Sep 17 00:00:00 2001 From: Copilot Date: Sun, 29 Mar 2026 05:43:01 +0000 Subject: [PATCH] [release] bump to v0.2.4 --- docker/patch/v0.5.9/megatron.patch | 23 + docker/patch/v0.5.9/sglang.patch | 1925 +++++++++++++++++++++++++--- setup.py | 2 +- 3 files changed, 1804 insertions(+), 146 deletions(-) diff --git a/docker/patch/v0.5.9/megatron.patch b/docker/patch/v0.5.9/megatron.patch index 6d2a233949..2e6ae436b1 100644 --- a/docker/patch/v0.5.9/megatron.patch +++ b/docker/patch/v0.5.9/megatron.patch @@ -314,6 +314,29 @@ index 1fd5dcfae..c9aeef1f0 100644 ctx.v_dim, nheads, batch_size, +diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py +index 5d7b69cd3..2e0a26815 100644 +--- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py ++++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py +@@ -348,6 +348,7 @@ class MultimodalRotaryEmbedding(nn.Module): + + # shape (seq_length, bs, 1, 2 * dim) + emb = emb[..., None, :].transpose(0, 1).contiguous() ++ packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if packed_seq_params is not None and packed_seq_params.local_cp_size is not None: + if packed_seq_params.local_cp_size > 1: + # Set CP group to dynamic CP group for CP slicing +@@ -357,7 +358,9 @@ class MultimodalRotaryEmbedding(nn.Module): + cp_group = None + else: + cp_group = self.cp_group +- if cp_group is not None and cp_group.size() > 1: ++ # For THD (packed sequence) format, skip CP slicing here — it is handled ++ # per-sequence inside _apply_rotary_pos_emb_thd instead (same as RotaryEmbedding). ++ if cp_group is not None and cp_group.size() > 1 and not packed_seq: + # slice rotary_pos_emb along sequence dimension and select the parition of the current + # CP rank + emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py index 13d74aa52..060898a7a 100644 --- a/megatron/core/models/common/language_module/language_module.py diff --git a/docker/patch/v0.5.9/sglang.patch b/docker/patch/v0.5.9/sglang.patch index 8cea544bd1..8001ceb73d 100644 --- a/docker/patch/v0.5.9/sglang.patch +++ b/docker/patch/v0.5.9/sglang.patch @@ -1,5 +1,14 @@ +diff --git a/.codespellrc b/.codespellrc +index 808a344b4..a34624958 100644 +--- a/.codespellrc ++++ b/.codespellrc +@@ -1,3 +1,3 @@ + [codespell] +-ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS ++ignore-words-list = ans, als, hel, boostrap, childs, te, vas, hsa, ment, cann, thi, makro, wil, rouge, PRIS, medias + skip = *.json,*.jsonl,*.patch,*.txt diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py -index 6fbd1db82..f80ec11bb 100644 +index 6fbd1db82..4c681b58d 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -274,6 +274,7 @@ class ModelConfig: @@ -10,6 +19,32 @@ index 6fbd1db82..f80ec11bb 100644 "GlmMoeDsaForCausalLM", ]: self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" +@@ -1016,10 +1017,10 @@ class ModelConfig: + + if tf_version < required_version: + if needs_tf_v5: +- raise ValueError( +- f"Transformers version {tf_version_str} is not supported for model {self.model_path} " ++ logger.warning( ++ f"Transformers version {tf_version_str} may not be fully supported for model {self.model_path} " + f"or model type {self.hf_config.model_type}. " +- "Please upgrade transformers to >= 5.0.0." ++ "Recommended transformers >= 5.0.0, but proceeding with current version." + ) + elif not needs_tf_v5: + logger.warning( +diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py +index da4629e52..c03f98231 100644 +--- a/python/sglang/srt/disaggregation/base/conn.py ++++ b/python/sglang/srt/disaggregation/base/conn.py +@@ -17,6 +17,7 @@ class KVArgs: + kv_data_ptrs: List[int] + kv_data_lens: List[int] + kv_item_lens: List[int] ++ aux_buffer_names: List[str] + aux_data_ptrs: List[int] + aux_data_lens: List[int] + aux_item_lens: List[int] diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 67fe82ad6..2ef25c49b 100644 --- a/python/sglang/srt/disaggregation/common/conn.py @@ -91,7 +126,7 @@ index 67fe82ad6..2ef25c49b 100644 "prefill_pp_size": self.pp_size, "prefill_page_size": self.page_size, diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py -index 1d8baf002..1672de78d 100644 +index 1d8baf002..1ebb95929 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -21,6 +21,7 @@ Life cycle of a request in the decode server @@ -102,7 +137,26 @@ index 1d8baf002..1672de78d 100644 import time from collections import deque from dataclasses import dataclass -@@ -336,6 +337,16 @@ class DecodePreallocQueue: +@@ -40,8 +41,10 @@ from sglang.srt.disaggregation.utils import ( + MetadataBuffers, + ReqToMetadataIdxAllocator, + TransferBackend, ++ apply_prefill_timing_payload, + get_kv_class, + is_mla_backend, ++ is_slime_profiling_enabled, + kv_to_page_indices, + poll_and_all_reduce, + prepare_abort, +@@ -295,6 +298,7 @@ class DecodePreallocQueue: + kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( + self.metadata_buffers.get_buf_infos() + ) ++ kv_args.aux_buffer_names = self.metadata_buffers.get_aux_buffer_names() + + if hasattr(self.token_to_kv_pool, "get_state_buf_infos"): + state_data_ptrs, state_data_lens, state_item_lens = ( +@@ -336,6 +340,16 @@ class DecodePreallocQueue: ) return kv_manager @@ -119,7 +173,7 @@ index 1d8baf002..1672de78d 100644 def add(self, req: Req, is_retracted: bool = False) -> None: """Add a request to the pending queue.""" if self._check_if_req_exceed_kv_capacity(req): -@@ -440,12 +451,37 @@ class DecodePreallocQueue: +@@ -440,12 +454,37 @@ class DecodePreallocQueue: [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group ) @@ -158,7 +212,38 @@ index 1d8baf002..1672de78d 100644 elif poll == KVPoll.WaitingForInput: decode_req.waiting_for_input = True elif poll == KVPoll.Failed: -@@ -830,6 +866,13 @@ class DecodeTransferQueue: +@@ -590,6 +629,7 @@ class DecodePreallocQueue: + self.req_to_metadata_buffer_idx_allocator.alloc() + ) + assert decode_req.metadata_buffer_index is not None ++ self.metadata_buffers.clear_profiling_buf(decode_req.metadata_buffer_index) + page_indices = kv_to_page_indices(kv_indices, page_size) + decode_req.kv_receiver.init( + page_indices, decode_req.metadata_buffer_index, state_indices +@@ -751,6 +791,7 @@ class DecodeTransferQueue: + output_topk_index, + output_hidden_states, + output_bootstrap_room, ++ output_prefill_timing, + ) = self.metadata_buffers.get_buf(idx) + + # Validate bootstrap_room to detect context corruption +@@ -813,6 +854,14 @@ class DecodeTransferQueue: + output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist() + ) + ++ # Inject prefill-side PD timing forwarded from the P instance. ++ # Layout: [bootstrap_queue, forward, transfer_queue, bootstrap, ++ # alloc_waiting, transfer_speed, transfer_mb, retry_count] ++ if is_slime_profiling_enabled(): ++ apply_prefill_timing_payload( ++ decode_req.req.time_stats, output_prefill_timing ++ ) ++ + decode_req.kv_receiver.clear() + decode_req.kv_receiver = None + trace_slice_end( +@@ -830,6 +879,13 @@ class DecodeTransferQueue: [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group ) @@ -172,7 +257,7 @@ index 1d8baf002..1672de78d 100644 transferred_reqs = [] indices_to_remove = set() for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): -@@ -877,7 +920,31 @@ class DecodeTransferQueue: +@@ -877,7 +933,20 @@ class DecodeTransferQueue: KVPoll.WaitingForInput, KVPoll.Transferring, ]: @@ -190,22 +275,11 @@ index 1d8baf002..1672de78d 100644 + f"{decode_req.req.rid=} {decode_req.req.bootstrap_room=}" + ) + logger.error(error_message) -+ prepare_abort( -+ decode_req.req, -+ error_message, -+ status_code=HTTPStatus.GATEWAY_TIMEOUT, -+ ) -+ self.scheduler.stream_output( -+ [decode_req.req], decode_req.req.return_logprob -+ ) -+ release_kv_cache(decode_req.req, self.tree_cache, is_insert=False) -+ indices_to_remove.add(i) -+ if self.scheduler.enable_metrics: -+ self.scheduler.metrics_collector.increment_transfer_failed_reqs() ++ decode_req.kv_receiver.abort() else: raise ValueError(f"Unexpected poll case: {poll}") -@@ -893,6 +960,14 @@ class DecodeTransferQueue: +@@ -893,6 +962,14 @@ class DecodeTransferQueue: return transferred_reqs @@ -220,7 +294,7 @@ index 1d8baf002..1672de78d 100644 class SchedulerDisaggregationDecodeMixin: -@@ -1072,7 +1147,15 @@ class SchedulerDisaggregationDecodeMixin: +@@ -1072,7 +1149,15 @@ class SchedulerDisaggregationDecodeMixin: resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs() self.waiting_queue.extend(resumed_reqs) if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: @@ -237,10 +311,60 @@ index 1d8baf002..1672de78d 100644 return if not hasattr(self, "polling_count"): +diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py +index a2d08e0e3..ed0790604 100644 +--- a/python/sglang/srt/disaggregation/encode_server.py ++++ b/python/sglang/srt/disaggregation/encode_server.py +@@ -117,7 +117,7 @@ def _convert(data): + return data + + +-_image_grid_attrs = ["image_grid_thw", "image_grid_hws"] ++_image_grid_attrs = ["image_grid_thw", "image_grid_hws", "grid_thws"] + + + def _get_image_grid_dim(images_input): +@@ -320,7 +320,26 @@ class MMEncoder: + + try: + kwargs = {"device": self.device} if self.use_image_processor_gpu else {} +- images_input = self.image_processor(images=images, **kwargs) ++ # Some processors (e.g., KimiK25VisionProcessor) expect MediaInput ++ # dicts rather than raw PIL Images. Wrap PIL images as needed. ++ from PIL import Image as PILImage ++ ++ if ( ++ isinstance(images, (list, tuple)) ++ and images ++ and isinstance(images[0], PILImage.Image) ++ ): ++ import inspect ++ ++ sig = inspect.signature(self.image_processor.preprocess) ++ first_param = list(sig.parameters.keys())[0] ++ if first_param == "medias": ++ medias = [{"type": "image", "image": img} for img in images] ++ images_input = self.image_processor.preprocess(medias, **kwargs) ++ else: ++ images_input = self.image_processor(images=images, **kwargs) ++ else: ++ images_input = self.image_processor(images=images, **kwargs) + feature = images_input["pixel_values"] + mm_item = MultimodalDataItem.from_dict( + { diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py -index d0d4efd95..fc4eef7b9 100644 +index d0d4efd95..b3a207063 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py +@@ -30,7 +30,7 @@ from sglang.srt.disaggregation.common.utils import ( + from sglang.srt.disaggregation.mooncake.utils import ( + check_mooncake_custom_mem_pool_enabled, + ) +-from sglang.srt.disaggregation.utils import DisaggregationMode ++from sglang.srt.disaggregation.utils import DisaggregationMode, iter_aux_transfer_specs + from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine + from sglang.srt.environ import envs + from sglang.srt.server_args import ServerArgs @@ -260,6 +260,19 @@ class MooncakeKVManager(CommonKVManager): self.kv_args.state_data_ptrs, self.kv_args.state_data_lens ) @@ -261,7 +385,44 @@ index d0d4efd95..fc4eef7b9 100644 def _transfer_data(self, mooncake_session_id, transfer_blocks): if not transfer_blocks: return 0 -@@ -643,13 +656,13 @@ class MooncakeKVManager(CommonKVManager): +@@ -524,10 +537,14 @@ class MooncakeKVManager(CommonKVManager): + prefill_aux_ptrs = self.kv_args.aux_data_ptrs + prefill_aux_item_lens = self.kv_args.aux_item_lens + +- for i, dst_aux_ptr in enumerate(dst_aux_ptrs): +- length = prefill_aux_item_lens[i] +- src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index +- dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index ++ for _, src_addr, dst_addr, length in iter_aux_transfer_specs( ++ self.kv_args.aux_buffer_names, ++ prefill_aux_ptrs, ++ prefill_aux_item_lens, ++ dst_aux_ptrs, ++ prefill_aux_index, ++ req.dst_aux_index, ++ ): + transfer_blocks.append((src_addr, dst_addr, length)) + + return self._transfer_data(req.mooncake_session_id, transfer_blocks) +@@ -541,9 +558,14 @@ class MooncakeKVManager(CommonKVManager): + prefill_aux_ptrs = self.kv_args.aux_data_ptrs + prefill_aux_item_lens = self.kv_args.aux_item_lens + +- for i in range(len(prefill_aux_ptrs)): +- length = prefill_aux_item_lens[i] +- src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index ++ for i, src_addr, _, length in iter_aux_transfer_specs( ++ self.kv_args.aux_buffer_names, ++ prefill_aux_ptrs, ++ prefill_aux_item_lens, ++ dst_aux_ptrs, ++ prefill_aux_index, ++ req.dst_aux_index, ++ ): + data = AuxDataCodec.serialize_data_from_buffer(src_addr, length) + + self.send_aux_data_to_endpoint( +@@ -643,13 +665,13 @@ class MooncakeKVManager(CommonKVManager): raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet." ) @@ -281,7 +442,20 @@ index d0d4efd95..fc4eef7b9 100644 # Reuse _send_kvcache_generic interface to send extra pool data prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32) dst_state_indices = np.array(req.dst_state_indices, dtype=np.int32) -@@ -880,13 +893,43 @@ class MooncakeKVManager(CommonKVManager): +@@ -858,12 +880,6 @@ class MooncakeKVManager(CommonKVManager): + if ret != 0: + with self.session_lock: + self.session_failures[req.mooncake_session_id] += 1 +- # Failures should never happen if the session is not dead, if the session fails once, mark it as failed +- if self.session_failures[req.mooncake_session_id] >= 1: +- self.failed_sessions.add(req.mooncake_session_id) +- logger.error( +- f"Session {req.mooncake_session_id} failed." +- ) + self.record_failure( + kv_chunk.room, + f"Failed to send kv chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}", +@@ -880,13 +896,31 @@ class MooncakeKVManager(CommonKVManager): if kv_chunk.is_last: if kv_chunk.state_indices is not None: @@ -298,18 +472,6 @@ index d0d4efd95..fc4eef7b9 100644 + self.session_failures[ + req.mooncake_session_id + ] += 1 -+ if ( -+ self.session_failures[ -+ req.mooncake_session_id -+ ] -+ >= 1 -+ ): -+ self.failed_sessions.add( -+ req.mooncake_session_id -+ ) -+ logger.error( -+ f"Session {req.mooncake_session_id} failed." -+ ) + self.record_failure( + kv_chunk.room, + f"Failed to send extra state chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}", @@ -326,7 +488,7 @@ index d0d4efd95..fc4eef7b9 100644 # Only the last chunk we need to send the aux data ret = self.send_aux( -@@ -895,6 +938,21 @@ class MooncakeKVManager(CommonKVManager): +@@ -895,6 +929,11 @@ class MooncakeKVManager(CommonKVManager): target_rank_registration_info.dst_aux_ptrs, ) polls.append(True if ret == 0 else False) @@ -335,20 +497,10 @@ index d0d4efd95..fc4eef7b9 100644 + # on subsequent batch_transfer_sync calls + with self.session_lock: + self.session_failures[req.mooncake_session_id] += 1 -+ if ( -+ self.session_failures[req.mooncake_session_id] -+ >= 1 -+ ): -+ self.failed_sessions.add( -+ req.mooncake_session_id -+ ) -+ logger.error( -+ f"Session {req.mooncake_session_id} failed (send_aux)." -+ ) dst_ranks_infos.append( (req.endpoint, req.dst_port, req.room) ) -@@ -977,15 +1035,20 @@ class MooncakeKVManager(CommonKVManager): +@@ -977,15 +1016,20 @@ class MooncakeKVManager(CommonKVManager): if status == KVPoll.Success: if bootstrap_room in self.request_status: @@ -357,11 +509,11 @@ index d0d4efd95..fc4eef7b9 100644 + # between the request_status check and dict access here. expected_response_num = ( - self.required_prefill_response_num_table[bootstrap_room] -- ) -- arrived_response_num = len( -- self.prefill_response_tracker[bootstrap_room] + self.required_prefill_response_num_table.get(bootstrap_room) ) +- arrived_response_num = len( +- self.prefill_response_tracker[bootstrap_room] +- ) - if arrived_response_num == expected_response_num: - self.update_status(bootstrap_room, KVPoll.Success) + if expected_response_num is not None: @@ -376,7 +528,7 @@ index d0d4efd95..fc4eef7b9 100644 elif status == KVPoll.Failed: self.record_failure( bootstrap_room, -@@ -1266,7 +1329,10 @@ class MooncakeKVReceiver(CommonKVReceiver): +@@ -1266,7 +1310,10 @@ class MooncakeKVReceiver(CommonKVReceiver): super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank) self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) @@ -389,7 +541,7 @@ index d0d4efd95..fc4eef7b9 100644 def _register_kv_args(self): for bootstrap_info in self.bootstrap_infos: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py -index fbc801635..21fc1ce0d 100644 +index fbc801635..ade111c9f 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -20,6 +20,7 @@ Life cycle of a request in the prefill server @@ -400,7 +552,15 @@ index fbc801635..21fc1ce0d 100644 import time from collections import deque from http import HTTPStatus -@@ -276,6 +277,12 @@ class PrefillBootstrapQueue: +@@ -167,6 +168,7 @@ class PrefillBootstrapQueue: + kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( + self.metadata_buffers.get_buf_infos() + ) ++ kv_args.aux_buffer_names = self.metadata_buffers.get_aux_buffer_names() + kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device + kv_args.gpu_id = self.scheduler.gpu_id + +@@ -276,6 +278,12 @@ class PrefillBootstrapQueue: [req.disagg_kv_sender for req in self.queue], self.gloo_group ) @@ -413,7 +573,7 @@ index fbc801635..21fc1ce0d 100644 for i, (req, poll) in enumerate(zip(self.queue, polls)): if rids_to_check is not None: # if req not in reqs_info_to_check, skip -@@ -283,6 +290,27 @@ class PrefillBootstrapQueue: +@@ -283,6 +291,27 @@ class PrefillBootstrapQueue: continue if poll == KVPoll.Bootstrapping: @@ -441,7 +601,7 @@ index fbc801635..21fc1ce0d 100644 continue elif poll == KVPoll.Failed: error_message = f"Prefill bootstrap failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}" -@@ -335,6 +363,15 @@ class PrefillBootstrapQueue: +@@ -335,6 +364,15 @@ class PrefillBootstrapQueue: else: return bootstrapped_reqs, failed_reqs @@ -457,8 +617,39 @@ index fbc801635..21fc1ce0d 100644 class SchedulerDisaggregationPrefillMixin: """ -@@ -564,6 +601,13 @@ class SchedulerDisaggregationPrefillMixin: - self.attn_tp_cpu_group, +@@ -547,6 +585,18 @@ class SchedulerDisaggregationPrefillMixin: + + self.maybe_send_health_check_signal() + ++ if ( ++ self.current_scheduler_metrics_enabled ++ and hasattr(batch, "prefill_stats") ++ and batch.prefill_stats is not None ++ ): ++ can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) ++ self.log_prefill_stats( ++ prefill_stats=batch.prefill_stats, ++ can_run_cuda_graph=can_run_cuda_graph, ++ dp_cooperation_info=getattr(batch, "dp_cooperation_info", None), ++ ) ++ + def process_disagg_prefill_inflight_queue( + self: Scheduler, rids_to_check: Optional[List[str]] = None + ) -> List[Req]: +@@ -559,11 +609,24 @@ class SchedulerDisaggregationPrefillMixin: + + done_reqs = [] + ++ # When CP > 1, use the full TP gloo group so all CP ranks reach ++ # consensus; otherwise a subset may enter run_batch while others wait ++ # in recv_requests, causing a deadlock. ++ disagg_gloo_group = ( ++ self.tp_cpu_group if self.attn_cp_size > 1 else self.attn_tp_cpu_group ++ ) + polls = poll_and_all_reduce( + [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], +- self.attn_tp_cpu_group, ++ disagg_gloo_group, ) + # Transfer timeout: if a request has been in the inflight queue for too long @@ -471,7 +662,7 @@ index fbc801635..21fc1ce0d 100644 undone_reqs: List[Req] = [] # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue for req, poll in zip(self.disagg_prefill_inflight_queue, polls): -@@ -573,10 +617,35 @@ class SchedulerDisaggregationPrefillMixin: +@@ -573,10 +636,35 @@ class SchedulerDisaggregationPrefillMixin: undone_reqs.append(req) continue @@ -509,6 +700,212 @@ index fbc801635..21fc1ce0d 100644 elif poll == KVPoll.Success: # transfer done release_kv_cache(req, self.tree_cache) # unlock the tree req.finished_reason = FINISH_LENGTH(length=0) +@@ -628,9 +716,12 @@ class SchedulerDisaggregationPrefillMixin: + """ + Used by PP, get the transferred rids but **do not pop** + """ ++ disagg_gloo_group = ( ++ self.tp_cpu_group if self.attn_cp_size > 1 else self.attn_tp_cpu_group ++ ) + polls = poll_and_all_reduce( + [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue], +- self.attn_tp_cpu_group, ++ disagg_gloo_group, + ) + + transferred_rids: List[str] = [] +diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py +index 6d58f415a..84723c342 100644 +--- a/python/sglang/srt/disaggregation/utils.py ++++ b/python/sglang/srt/disaggregation/utils.py +@@ -21,6 +21,17 @@ if TYPE_CHECKING: + # Constants & Enums + ######################### + FAKE_BOOTSTRAP_HOST = "2.2.2.2" ++PREFILL_TIMING_AUX_BUFFER_NAME = "prefill_timing" ++PREFILL_TIMING_DEST_ATTRS = ( ++ ("fwd_prefill_bootstrap_queue_duration", float), ++ ("fwd_prefill_forward_duration", float), ++ ("fwd_prefill_transfer_queue_duration", float), ++ ("fwd_bootstrap_duration", float), ++ ("fwd_alloc_waiting_duration", float), ++ ("fwd_transfer_speed_gb_s", float), ++ ("fwd_transfer_total_mb", float), ++ ("fwd_prefill_retry_count", int), ++) + + + class DisaggregationMode(Enum): +@@ -139,46 +150,35 @@ class MetadataBuffers: + self.bootstrap_room = torch.zeros( + (size, 8), dtype=torch.uint64, device=device + ) ++ # Prefill-side PD timing (8 floats, padded to 16 for RDMA alignment). ++ # Layout: [bootstrap_queue, forward, transfer_queue, bootstrap, ++ # alloc_waiting, transfer_speed, transfer_mb, retry_count] ++ self.prefill_timing = torch.zeros( ++ (size, 16), dtype=torch.float32, device=device ++ ) ++ self.aux_buffers = [ ++ ("output_ids", self.output_ids), ++ ("cached_tokens", self.cached_tokens), ++ ("output_token_logprobs_val", self.output_token_logprobs_val), ++ ("output_token_logprobs_idx", self.output_token_logprobs_idx), ++ ("output_top_logprobs_val", self.output_top_logprobs_val), ++ ("output_top_logprobs_idx", self.output_top_logprobs_idx), ++ ("output_topk_p", self.output_topk_p), ++ ("output_topk_index", self.output_topk_index), ++ ("output_hidden_states", self.output_hidden_states), ++ ("bootstrap_room", self.bootstrap_room), ++ (PREFILL_TIMING_AUX_BUFFER_NAME, self.prefill_timing), ++ ] + + def get_buf_infos(self): +- ptrs = [ +- self.output_ids.data_ptr(), +- self.cached_tokens.data_ptr(), +- self.output_token_logprobs_val.data_ptr(), +- self.output_token_logprobs_idx.data_ptr(), +- self.output_top_logprobs_val.data_ptr(), +- self.output_top_logprobs_idx.data_ptr(), +- self.output_topk_p.data_ptr(), +- self.output_topk_index.data_ptr(), +- self.output_hidden_states.data_ptr(), +- self.bootstrap_room.data_ptr(), +- ] +- data_lens = [ +- self.output_ids.nbytes, +- self.cached_tokens.nbytes, +- self.output_token_logprobs_val.nbytes, +- self.output_token_logprobs_idx.nbytes, +- self.output_top_logprobs_val.nbytes, +- self.output_top_logprobs_idx.nbytes, +- self.output_topk_p.nbytes, +- self.output_topk_index.nbytes, +- self.output_hidden_states.nbytes, +- self.bootstrap_room.nbytes, +- ] +- item_lens = [ +- self.output_ids[0].nbytes, +- self.cached_tokens[0].nbytes, +- self.output_token_logprobs_val[0].nbytes, +- self.output_token_logprobs_idx[0].nbytes, +- self.output_top_logprobs_val[0].nbytes, +- self.output_top_logprobs_idx[0].nbytes, +- self.output_topk_p[0].nbytes, +- self.output_topk_index[0].nbytes, +- self.output_hidden_states[0].nbytes, +- self.bootstrap_room[0].nbytes, +- ] ++ ptrs = [buffer.data_ptr() for _, buffer in self.aux_buffers] ++ data_lens = [buffer.nbytes for _, buffer in self.aux_buffers] ++ item_lens = [buffer[0].nbytes for _, buffer in self.aux_buffers] + return ptrs, data_lens, item_lens + ++ def get_aux_buffer_names(self): ++ return [name for name, _ in self.aux_buffers] ++ + def get_buf(self, idx: int): + return ( + self.output_ids[idx], +@@ -191,8 +191,12 @@ class MetadataBuffers: + self.output_topk_index[idx], + self.output_hidden_states[idx], + self.bootstrap_room[idx], ++ self.prefill_timing[idx], + ) + ++ def clear_profiling_buf(self, idx: int): ++ self.prefill_timing[idx].zero_() ++ + def set_buf(self, req: Req): + + self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0] +@@ -237,6 +241,84 @@ class MetadataBuffers: + self.bootstrap_room[req.metadata_buffer_index, 0] = ( + req.bootstrap_room if req.bootstrap_room is not None else 0 + ) ++ # Pack prefill-side PD timing durations for transfer to decode instance. ++ # Note: set_buf is called at the START of the last KV chunk send, so ++ # completion_time and prefill_transfer_queue_entry_time are not yet set. ++ # We use time.perf_counter() as the "forward just completed" timestamp. ++ import time ++ ++ ts = req.time_stats ++ timing = self.prefill_timing[req.metadata_buffer_index] ++ self.clear_profiling_buf(req.metadata_buffer_index) ++ if not is_slime_profiling_enabled(): ++ return ++ for idx, value in enumerate( ++ build_prefill_timing_payload(ts, now=time.perf_counter()) ++ ): ++ if value > 0: ++ timing[idx] = value ++ ++ ++def is_slime_profiling_enabled() -> bool: ++ return envs.SLIME_ENABLE_PROFILING.get() ++ ++ ++def build_prefill_timing_payload(time_stats, now: float) -> tuple[float, ...]: ++ bootstrap_queue_duration = 0.0 ++ if ( ++ time_stats.prefill_bootstrap_queue_entry_time > 0 ++ and time_stats.wait_queue_entry_time > 0 ++ ): ++ bootstrap_queue_duration = ( ++ time_stats.wait_queue_entry_time ++ - time_stats.prefill_bootstrap_queue_entry_time ++ ) ++ ++ prefill_forward_duration = ( ++ now - time_stats.forward_entry_time ++ if time_stats.forward_entry_time > 0 ++ else 0.0 ++ ) ++ ++ return ( ++ bootstrap_queue_duration, ++ prefill_forward_duration, ++ 0.0, ++ max(0.0, time_stats.bootstrap_duration), ++ max(0.0, time_stats.alloc_waiting_duration), ++ max(0.0, time_stats.transfer_speed_gb_s), ++ max(0.0, time_stats.transfer_total_mb), ++ float(max(0, time_stats.prefill_retry_count)), ++ ) ++ ++ ++def apply_prefill_timing_payload(time_stats, timing) -> None: ++ for value, (attr_name, caster) in zip( ++ timing[: len(PREFILL_TIMING_DEST_ATTRS)].tolist(), ++ PREFILL_TIMING_DEST_ATTRS, ++ ): ++ if value > 0: ++ setattr(time_stats, attr_name, caster(value)) ++ ++ ++def iter_aux_transfer_specs( ++ aux_buffer_names: list[str], ++ prefill_aux_ptrs: list[int], ++ prefill_aux_item_lens: list[int], ++ dst_aux_ptrs: list[int], ++ prefill_aux_index: int, ++ dst_aux_index: int, ++): ++ profiling_enabled = is_slime_profiling_enabled() ++ for i, (buffer_name, dst_aux_ptr) in enumerate(zip(aux_buffer_names, dst_aux_ptrs)): ++ if not profiling_enabled and buffer_name == PREFILL_TIMING_AUX_BUFFER_NAME: ++ continue ++ length = prefill_aux_item_lens[i] ++ if length <= 0: ++ continue ++ src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index ++ dst_addr = dst_aux_ptr + length * dst_aux_index ++ yield i, src_addr, dst_addr, length + + + ######################### diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 8f1069c00..e47589295 100644 --- a/python/sglang/srt/distributed/parallel_state.py @@ -634,6 +1031,18 @@ index 1d6816c01..402b42e05 100644 @app.post("/update_weight_version") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): +diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py +index 8293796a2..bff34e422 100644 +--- a/python/sglang/srt/environ.py ++++ b/python/sglang/srt/environ.py +@@ -244,6 +244,7 @@ class Envs: + SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2) + SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) + SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") ++ SLIME_ENABLE_PROFILING = EnvBool(False) + + # Scheduler: others: + SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py index 1cdf65b91..4783cd18f 100644 --- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py @@ -655,7 +1064,7 @@ index 1cdf65b91..4783cd18f 100644 buf_numel_per_page: tl.constexpr, index_head_dim: tl.constexpr, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py -index ca54a931b..6c102a251 100644 +index ca54a931b..3540f77ba 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1,6 +1,7 @@ @@ -666,20 +1075,41 @@ index ca54a931b..6c102a251 100644 from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple -@@ -207,7 +208,11 @@ class Indexer(MultiPlatformOp): +@@ -201,14 +202,31 @@ class Indexer(MultiPlatformOp): + prefix=add_prefix("weights_proj", prefix), + ) + self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32) ++ server_args = get_global_server_args() ++ disable_flag = server_args.disable_indexer_rope_neox_style ++ env_raw = os.environ.get("INDEXER_ROPE_NEOX_STYLE", None) ++ if env_raw is not None: ++ env_value = env_raw == "1" ++ if disable_flag and env_value: ++ raise ValueError( ++ "Conflict: --disable-indexer-rope-neox-style is set but " ++ "INDEXER_ROPE_NEOX_STYLE='1'. " ++ "Please remove one or make them consistent." ++ ) ++ resolved_neox_style = env_value ++ elif disable_flag: ++ resolved_neox_style = False ++ else: ++ resolved_neox_style = is_neox_style ++ + self.rotary_emb = get_rope_wrapper( + rope_head_dim, + rotary_dim=rope_head_dim, max_position=max_position_embeddings, base=rope_theta, # type: ignore rope_scaling=rope_scaling, - is_neox_style=is_neox_style, -+ is_neox_style=( -+ os.environ.get("INDEXER_ROPE_NEOX_STYLE", "1") == "1" -+ if os.environ.get("INDEXER_ROPE_NEOX_STYLE", None) -+ else is_neox_style -+ ), - device=get_global_server_args().device, +- device=get_global_server_args().device, ++ is_neox_style=resolved_neox_style, ++ device=server_args.device, ) self.block_size = block_size -@@ -244,6 +249,11 @@ class Indexer(MultiPlatformOp): + self.scale_fmt = scale_fmt +@@ -244,6 +262,11 @@ class Indexer(MultiPlatformOp): x = x.to(self.weights_proj.weight.dtype) weights, _ = self.weights_proj(x) weights = weights.float() @@ -691,7 +1121,7 @@ index ca54a931b..6c102a251 100644 weights = weights * self.n_heads**-0.5 weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale return weights -@@ -982,15 +992,26 @@ class Indexer(MultiPlatformOp): +@@ -982,15 +1005,26 @@ class Indexer(MultiPlatformOp): query, key = self._get_q_k_bf16( q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch ) @@ -719,24 +1149,10 @@ index ca54a931b..6c102a251 100644 if enable_dual_stream: current_stream = torch.cuda.current_stream() diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py -index 00ef96f9b..5adaec804 100644 +index 00ef96f9b..c2c2c78fe 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py -@@ -54,7 +54,12 @@ def can_nsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"): - return False - cp_size = get_attention_cp_size() - seq_len = sum(forward_batch.extend_seq_lens_cpu) -- return is_nsa_prefill_cp_round_robin_split() and seq_len > 0 and cp_size > 1 -+ return ( -+ is_nsa_prefill_cp_round_robin_split() -+ and seq_len >= cp_size -+ and seq_len % cp_size == 0 -+ and cp_size > 1 -+ ) - - - def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]): -@@ -91,20 +96,29 @@ def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]): +@@ -91,20 +91,29 @@ def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]): def cal_padded_tokens(forward_batch: "ForwardBatch"): # Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch, # calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode. @@ -775,32 +1191,7 @@ index 00ef96f9b..5adaec804 100644 if can_nsa_prefill_cp_round_robin_split(forward_batch): tokens = ceil_div(tokens, attn_cp_size) return tokens -@@ -152,10 +166,20 @@ class NSAContextParallelMetadata: - - def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): - if is_nsa_prefill_cp_round_robin_split(): -- cur_cp_seq_len = seq_len // cp_size -- assert ( -- seq_len % cp_size == 0 -- ), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when nsa_prefill_cp_mode is round-robin-split" -+ # Use actual extend sequence length instead of (possibly padded) input_ids -+ # length to stay consistent with can_nsa_prefill_cp_round_robin_split(), -+ # which also checks sum(extend_seq_lens_cpu). When prepare_mlp_sync_batch -+ # pads input_ids to ceil_align(n, attn_cp_size), len(input_ids) can become -+ # divisible by cp_size even though the real extend length is not, causing -+ # hidden_states to be CP-split while the attention metadata is not. -+ actual_seq_len = ( -+ sum(forward_batch.extend_seq_lens_cpu) -+ if forward_batch.extend_seq_lens_cpu is not None -+ else seq_len -+ ) -+ if actual_seq_len < cp_size or actual_seq_len % cp_size != 0: -+ return False -+ cur_cp_seq_len = actual_seq_len // cp_size - else: - # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2 - # Note: (self.cp_size * 2) To achieve load balancing for seq computation, -@@ -175,10 +199,6 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): +@@ -175,10 +184,6 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): if is_nsa_prefill_cp_round_robin_split(): @@ -811,7 +1202,7 @@ index 00ef96f9b..5adaec804 100644 return nsa_cp_round_robin_split_data(input_) input_list = list( -@@ -192,11 +212,6 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): +@@ -192,11 +197,6 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor): if is_nsa_prefill_cp_round_robin_split(): @@ -1172,7 +1563,7 @@ index ebcc696ec..3b527021a 100644 self, dispatch_output: Union[DeepEPNormalDispatchOutput, DeepEPLLDispatchOutput], diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py -index de8a07ab3..5c9f4813a 100644 +index de8a07ab3..952f8a67b 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -697,6 +697,7 @@ class FusedMoE(torch.nn.Module): @@ -1183,7 +1574,31 @@ index de8a07ab3..5c9f4813a 100644 else loaded_weight ) -@@ -916,6 +917,7 @@ class FusedMoE(torch.nn.Module): +@@ -821,13 +822,16 @@ class FusedMoE(torch.nn.Module): + FusedMoeWeightScaleSupported.GROUP.value, + FusedMoeWeightScaleSupported.BLOCK.value, + ]: +- self._load_model_weight_or_group_weight_scale( +- shard_id=shard_id, +- shard_dim=shard_dim, +- loaded_weight=loaded_weight, +- expert_data=expert_data, +- tp_rank=tp_rank, +- ) ++ if getattr(param, "load_full_w2", False) and shard_id == "w2": ++ expert_data.copy_(loaded_weight) ++ else: ++ self._load_model_weight_or_group_weight_scale( ++ shard_id=shard_id, ++ shard_dim=shard_dim, ++ loaded_weight=loaded_weight, ++ expert_data=expert_data, ++ tp_rank=tp_rank, ++ ) + elif quant_method == FusedMoeWeightScaleSupported.TENSOR.value: + # INT4-FP8 (INT4 MoE Weight, FP8 Compute): Adjust FP8 per-tensor scaling number for e4m3fnuz (AMD) + if _is_hip and get_bool_env_var("SGLANG_INT4_WEIGHT"): +@@ -916,6 +920,7 @@ class FusedMoE(torch.nn.Module): "CompressedTensorsWNA16TritonMoE", ] ) @@ -1253,7 +1668,7 @@ index 00bd68755..12d5577af 100644 def get_routed_experts( diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py -index 8539639d5..e7f5d1565 100644 +index 8539639d5..d44496c2f 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py @@ -388,6 +388,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase): @@ -1292,7 +1707,7 @@ index 8539639d5..e7f5d1565 100644 output = hidden_states else: raise NotImplementedError() # triton runner was supported but it's temporarily disabled -@@ -551,10 +562,12 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): +@@ -551,10 +562,18 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): buffer = self._get_buffer() topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids topk_ids = topk_ids.to(torch.int64) @@ -1300,16 +1715,22 @@ index 8539639d5..e7f5d1565 100644 - hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1] - + self.num_experts - ) // self.num_experts -+ # Use a correctness-preserving upper bound for per-expert token count. -+ # In the worst case, every rank routes all local tokens to the same expert. -+ expected_m = min( -+ hidden_states.shape[0] * buffer.group_size, -+ self.num_max_dispatch_tokens_per_rank * buffer.group_size, -+ ) ++ if self.quant_config.get("bf16_weights", False): ++ # BF16 low-latency path slices hidden_states[:, :expected_m, :], so ++ # expected_m must remain a correctness-preserving upper bound. ++ expected_m = min( ++ hidden_states.shape[0] * buffer.group_size, ++ self.num_max_dispatch_tokens_per_rank * buffer.group_size, ++ ) ++ else: ++ expected_m = ( ++ hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1] ++ + self.num_experts ++ ) // self.num_experts hidden_states, masked_m, event, hook = self._dispatch_core( hidden_states, topk_ids, -@@ -609,7 +622,9 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): +@@ -609,7 +628,9 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase): input_global_scale = self.quant_config.get("input_global_scale", None) if input_global_scale is not None: use_nvfp4 = True @@ -1486,11 +1907,76 @@ index ae0614635..3b6a8d254 100644 # TODO: remove this when npu_mrope supports QNumHeads * QHeadSize > 4096 assert ( fused_set_kv_buffer_arg is None +diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py +index 652227860..7d3a5d0c4 100644 +--- a/python/sglang/srt/managers/detokenizer_manager.py ++++ b/python/sglang/srt/managers/detokenizer_manager.py +@@ -405,6 +405,17 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): + prefill_launch_delay=recv_obj.prefill_launch_delay, + prefill_launch_latency=recv_obj.prefill_launch_latency, + prefill_finished_ts=recv_obj.prefill_finished_ts, ++ pd_prefill_bootstrap_queue_duration=recv_obj.pd_prefill_bootstrap_queue_duration, ++ pd_prefill_forward_duration=recv_obj.pd_prefill_forward_duration, ++ pd_prefill_transfer_queue_duration=recv_obj.pd_prefill_transfer_queue_duration, ++ pd_decode_prealloc_duration=recv_obj.pd_decode_prealloc_duration, ++ pd_decode_transfer_duration=recv_obj.pd_decode_transfer_duration, ++ pd_decode_forward_duration=recv_obj.pd_decode_forward_duration, ++ pd_bootstrap_duration=recv_obj.pd_bootstrap_duration, ++ pd_alloc_waiting_duration=recv_obj.pd_alloc_waiting_duration, ++ pd_transfer_speed_gb_s=recv_obj.pd_transfer_speed_gb_s, ++ pd_transfer_total_mb=recv_obj.pd_transfer_total_mb, ++ pd_prefill_retry_count=recv_obj.pd_prefill_retry_count, + ) + + def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index ff1774567..42d27a82a 100644 +index ff1774567..f947e71d7 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -1403,6 +1403,20 @@ class UpdateWeightsFromIPCReqOutput(BaseReq): +@@ -101,6 +101,42 @@ class RequestTimingMetricsMixin: + # This marks when the prefill computation finishes. + prefill_finished_ts: Optional[List[Optional[float]]] + ++ # --- PD disaggregation timing fields --- ++ # All fields are None when profiling is disabled or not in PD disaggregation mode. ++ ++ # P instance: duration spent in bootstrap queue before entering the wait queue. ++ pd_prefill_bootstrap_queue_duration: Optional[List[Optional[float]]] ++ ++ # P instance: duration for the actual prefill forward computation. ++ pd_prefill_forward_duration: Optional[List[Optional[float]]] ++ ++ # P instance: duration spent in the KV transfer queue. ++ pd_prefill_transfer_queue_duration: Optional[List[Optional[float]]] ++ ++ # D instance: duration waiting for KV cache slot pre-allocation. ++ pd_decode_prealloc_duration: Optional[List[Optional[float]]] ++ ++ # D instance: duration waiting for the KV cache transfer to complete. ++ pd_decode_transfer_duration: Optional[List[Optional[float]]] ++ ++ # D instance: duration for the actual decode forward computation. ++ pd_decode_forward_duration: Optional[List[Optional[float]]] ++ ++ # Bootstrap handshake duration (P and D instances). ++ pd_bootstrap_duration: Optional[List[Optional[float]]] ++ ++ # KV cache allocation waiting duration (P and D instances). ++ pd_alloc_waiting_duration: Optional[List[Optional[float]]] ++ ++ # KV cache transfer speed in GB/s. ++ pd_transfer_speed_gb_s: Optional[List[Optional[float]]] ++ ++ # Total KV cache transferred in MB. ++ pd_transfer_total_mb: Optional[List[Optional[float]]] ++ ++ # Number of prefill retries (P instance only). ++ pd_prefill_retry_count: Optional[List[Optional[int]]] ++ + + @dataclass + class SpeculativeDecodingMetricsMixin: +@@ -1403,6 +1439,20 @@ class UpdateWeightsFromIPCReqOutput(BaseReq): message: str @@ -1511,7 +1997,7 @@ index ff1774567..42d27a82a 100644 @dataclass class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq): success: bool -@@ -1802,6 +1816,10 @@ class GetLoadReqOutput(BaseReq): +@@ -1802,6 +1852,10 @@ class GetLoadReqOutput(BaseReq): num_waiting_reqs: int num_tokens: int ts_tic: float @@ -1522,6 +2008,202 @@ index ff1774567..42d27a82a 100644 @dataclass +diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py +index e1236aa0f..daa598a1f 100644 +--- a/python/sglang/srt/managers/multi_tokenizer_mixin.py ++++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py +@@ -142,6 +142,39 @@ def _handle_output_by_index(output, i): + prefill_finished_ts=_extract_field_by_index( + output, "prefill_finished_ts", i + ), ++ pd_prefill_bootstrap_queue_duration=_extract_field_by_index( ++ output, "pd_prefill_bootstrap_queue_duration", i ++ ), ++ pd_prefill_forward_duration=_extract_field_by_index( ++ output, "pd_prefill_forward_duration", i ++ ), ++ pd_prefill_transfer_queue_duration=_extract_field_by_index( ++ output, "pd_prefill_transfer_queue_duration", i ++ ), ++ pd_decode_prealloc_duration=_extract_field_by_index( ++ output, "pd_decode_prealloc_duration", i ++ ), ++ pd_decode_transfer_duration=_extract_field_by_index( ++ output, "pd_decode_transfer_duration", i ++ ), ++ pd_decode_forward_duration=_extract_field_by_index( ++ output, "pd_decode_forward_duration", i ++ ), ++ pd_bootstrap_duration=_extract_field_by_index( ++ output, "pd_bootstrap_duration", i ++ ), ++ pd_alloc_waiting_duration=_extract_field_by_index( ++ output, "pd_alloc_waiting_duration", i ++ ), ++ pd_transfer_speed_gb_s=_extract_field_by_index( ++ output, "pd_transfer_speed_gb_s", i ++ ), ++ pd_transfer_total_mb=_extract_field_by_index( ++ output, "pd_transfer_total_mb", i ++ ), ++ pd_prefill_retry_count=_extract_field_by_index( ++ output, "pd_prefill_retry_count", i ++ ), + finished_reasons=_extract_field_by_index(output, "finished_reasons", i), + decoded_texts=_extract_field_by_index(output, "decoded_texts", i), + decode_ids=_extract_field_by_index(output, "decode_ids", i), +@@ -211,6 +244,50 @@ def _handle_output_by_index(output, i): + elif isinstance(output, BatchEmbeddingOutput): + new_output = BatchEmbeddingOutput( + rids=[output.rids[i]], ++ queue_time=_extract_field_by_index(output, "queue_time", i), ++ forward_entry_time=_extract_field_by_index(output, "forward_entry_time", i), ++ prefill_launch_delay=_extract_field_by_index( ++ output, "prefill_launch_delay", i ++ ), ++ prefill_launch_latency=_extract_field_by_index( ++ output, "prefill_launch_latency", i ++ ), ++ prefill_finished_ts=_extract_field_by_index( ++ output, "prefill_finished_ts", i ++ ), ++ pd_prefill_bootstrap_queue_duration=_extract_field_by_index( ++ output, "pd_prefill_bootstrap_queue_duration", i ++ ), ++ pd_prefill_forward_duration=_extract_field_by_index( ++ output, "pd_prefill_forward_duration", i ++ ), ++ pd_prefill_transfer_queue_duration=_extract_field_by_index( ++ output, "pd_prefill_transfer_queue_duration", i ++ ), ++ pd_decode_prealloc_duration=_extract_field_by_index( ++ output, "pd_decode_prealloc_duration", i ++ ), ++ pd_decode_transfer_duration=_extract_field_by_index( ++ output, "pd_decode_transfer_duration", i ++ ), ++ pd_decode_forward_duration=_extract_field_by_index( ++ output, "pd_decode_forward_duration", i ++ ), ++ pd_bootstrap_duration=_extract_field_by_index( ++ output, "pd_bootstrap_duration", i ++ ), ++ pd_alloc_waiting_duration=_extract_field_by_index( ++ output, "pd_alloc_waiting_duration", i ++ ), ++ pd_transfer_speed_gb_s=_extract_field_by_index( ++ output, "pd_transfer_speed_gb_s", i ++ ), ++ pd_transfer_total_mb=_extract_field_by_index( ++ output, "pd_transfer_total_mb", i ++ ), ++ pd_prefill_retry_count=_extract_field_by_index( ++ output, "pd_prefill_retry_count", i ++ ), + finished_reasons=_extract_field_by_index(output, "finished_reasons", i), + embeddings=_extract_field_by_index(output, "embeddings", i), + prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i), +@@ -239,6 +316,39 @@ def _handle_output_by_index(output, i): + prefill_finished_ts=_extract_field_by_index( + output, "prefill_finished_ts", i + ), ++ pd_prefill_bootstrap_queue_duration=_extract_field_by_index( ++ output, "pd_prefill_bootstrap_queue_duration", i ++ ), ++ pd_prefill_forward_duration=_extract_field_by_index( ++ output, "pd_prefill_forward_duration", i ++ ), ++ pd_prefill_transfer_queue_duration=_extract_field_by_index( ++ output, "pd_prefill_transfer_queue_duration", i ++ ), ++ pd_decode_prealloc_duration=_extract_field_by_index( ++ output, "pd_decode_prealloc_duration", i ++ ), ++ pd_decode_transfer_duration=_extract_field_by_index( ++ output, "pd_decode_transfer_duration", i ++ ), ++ pd_decode_forward_duration=_extract_field_by_index( ++ output, "pd_decode_forward_duration", i ++ ), ++ pd_bootstrap_duration=_extract_field_by_index( ++ output, "pd_bootstrap_duration", i ++ ), ++ pd_alloc_waiting_duration=_extract_field_by_index( ++ output, "pd_alloc_waiting_duration", i ++ ), ++ pd_transfer_speed_gb_s=_extract_field_by_index( ++ output, "pd_transfer_speed_gb_s", i ++ ), ++ pd_transfer_total_mb=_extract_field_by_index( ++ output, "pd_transfer_total_mb", i ++ ), ++ pd_prefill_retry_count=_extract_field_by_index( ++ output, "pd_prefill_retry_count", i ++ ), + finished_reasons=_extract_field_by_index(output, "finished_reasons", i), + output_strs=_extract_field_by_index(output, "output_strs", i), + output_ids=_extract_field_by_index(output, "output_ids", i), +@@ -524,6 +634,60 @@ def monkey_patch_uvicorn_multiprocessing(timeout: float = 10): + "uvicorn.supervisors.multiprocess not found, skipping monkey patch" + ) + ++ # Fix stdin fd issue when running under Ray (or other managed ++ # environments where stdin may not be a real terminal): ++ # ++ # Uvicorn's get_subprocess() captures sys.stdin.fileno() in the parent ++ # and passes it to spawn'd children, which call os.fdopen(stdin_fileno) ++ # to re-attach stdin. This is intended for interactive debugging (e.g. ++ # pdb attach to a child worker). ++ # ++ # In Ray Actors, sys.stdin.fileno() succeeds in the parent (returns a ++ # valid fd number), but the fd is not inheritable across spawn. The ++ # child's os.fdopen() then crashes with OSError: [Errno 9] Bad file ++ # descriptor, killing every tokenizer worker. ++ # ++ # Instead of unconditionally disabling stdin passthrough, we probe ++ # whether the fd is truly usable by dup'ing it. If os.dup() fails, ++ # the fd won't survive spawn either, so we fall back to None. In a ++ # normal terminal environment os.dup() succeeds and debugging ability ++ # is preserved. ++ try: ++ import uvicorn._subprocess as _uv_sub ++ import uvicorn.supervisors.multiprocess as _uv_mp ++ ++ def _safe_get_stdin_fileno(): ++ """Return stdin fileno only if it is genuinely usable.""" ++ try: ++ fileno = sys.stdin.fileno() ++ # Verify the fd is valid and duplicable — if it isn't, ++ # spawn'd children won't be able to reopen it either. ++ dup_fd = os.dup(fileno) ++ os.close(dup_fd) ++ return fileno ++ except (AttributeError, OSError): ++ return None ++ ++ def _patched_get_subprocess(config, target, sockets): ++ stdin_fileno = _safe_get_stdin_fileno() ++ kwargs = { ++ "config": config, ++ "target": target, ++ "sockets": sockets, ++ "stdin_fileno": stdin_fileno, ++ } ++ return _uv_sub.spawn.Process( ++ target=_uv_sub.subprocess_started, kwargs=kwargs ++ ) ++ ++ # Must patch both: the supervisor module caches its own reference ++ # to get_subprocess at import time via ++ # ``from uvicorn._subprocess import get_subprocess``. ++ _uv_sub.get_subprocess = _patched_get_subprocess ++ _uv_mp.get_subprocess = _patched_get_subprocess ++ except Exception: ++ pass ++ + + class SenderWrapper: + def __init__(self, port_args: PortArgs, send_to_scheduler: zmq.Socket): diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index c07995798..dd8ca7167 100644 --- a/python/sglang/srt/managers/schedule_batch.py @@ -1539,7 +2221,7 @@ index c07995798..dd8ca7167 100644 break diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index a9ff0ac94..a50dd5122 100644 +index a9ff0ac94..c124f43bc 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -114,6 +114,7 @@ from sglang.srt.managers.io_struct import ( @@ -1550,7 +2232,28 @@ index a9ff0ac94..a50dd5122 100644 ProfileReq, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, -@@ -1063,6 +1064,7 @@ class Scheduler( +@@ -952,6 +953,11 @@ class Scheduler( + custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(), + ) + ++ # When CP > 1, all CP ranks must agree on poll results so they ++ # enter run_batch together; use the full TP gloo group for consensus. ++ disagg_prefill_gloo_group = ( ++ self.tp_cpu_group if self.attn_cp_size > 1 else self.attn_tp_cpu_group ++ ) + self.disagg_prefill_bootstrap_queue = PrefillBootstrapQueue( + token_to_kv_pool=self.token_to_kv_pool_allocator.get_kvcache(), + draft_token_to_kv_pool=draft_token_to_kv_pool, +@@ -961,7 +967,7 @@ class Scheduler( + tp_size=self.tp_size, + gpu_id=self.gpu_id, + bootstrap_port=self.server_args.disaggregation_bootstrap_port, +- gloo_group=self.attn_tp_cpu_group, ++ gloo_group=disagg_prefill_gloo_group, + max_total_num_tokens=self.max_total_num_tokens, + decode_tp_size=self.server_args.disaggregation_decode_tp, + decode_dp_size=self.server_args.disaggregation_decode_dp, +@@ -1063,6 +1069,7 @@ class Scheduler( ), (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor), (UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc), @@ -1618,10 +2321,70 @@ index 30b2732b9..68090b161 100644 def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py -index 482bc6ca6..857cfa6a3 100644 +index 482bc6ca6..fbc486417 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py -@@ -1134,7 +1134,7 @@ class SchedulerOutputProcessorMixin: +@@ -922,6 +922,18 @@ class SchedulerOutputProcessorMixin: + prefill_launch_delays = [] + prefill_launch_latencies = [] + prefill_finished_timestamps = [] ++ profiling_enabled = envs.SLIME_ENABLE_PROFILING.get() ++ pd_prefill_bootstrap_queue_durations = [] if profiling_enabled else None ++ pd_prefill_forward_durations = [] if profiling_enabled else None ++ pd_prefill_transfer_queue_durations = [] if profiling_enabled else None ++ pd_decode_prealloc_durations = [] if profiling_enabled else None ++ pd_decode_transfer_durations = [] if profiling_enabled else None ++ pd_decode_forward_durations = [] if profiling_enabled else None ++ pd_bootstrap_durations = [] if profiling_enabled else None ++ pd_alloc_waiting_durations = [] if profiling_enabled else None ++ pd_transfer_speeds_gb_s = [] if profiling_enabled else None ++ pd_transfer_totals_mb = [] if profiling_enabled else None ++ pd_prefill_retry_counts = [] if profiling_enabled else None + + if return_logprob: + input_token_logprobs_val = [] +@@ -1037,6 +1049,40 @@ class SchedulerOutputProcessorMixin: + prefill_finished_timestamps.append( + req.time_stats.get_prefill_finished_ts() + ) ++ if profiling_enabled: ++ pd_prefill_bootstrap_queue_durations.append( ++ req.time_stats.get_pd_prefill_bootstrap_queue_duration() ++ ) ++ pd_prefill_forward_durations.append( ++ req.time_stats.get_pd_prefill_forward_duration() ++ ) ++ pd_prefill_transfer_queue_durations.append( ++ req.time_stats.get_pd_prefill_transfer_queue_duration() ++ ) ++ pd_decode_prealloc_durations.append( ++ req.time_stats.get_pd_decode_prealloc_duration() ++ ) ++ pd_decode_transfer_durations.append( ++ req.time_stats.get_pd_decode_transfer_duration() ++ ) ++ pd_decode_forward_durations.append( ++ req.time_stats.get_pd_decode_forward_duration() ++ ) ++ pd_bootstrap_durations.append( ++ req.time_stats.get_pd_bootstrap_duration() ++ ) ++ pd_alloc_waiting_durations.append( ++ req.time_stats.get_pd_alloc_waiting_duration() ++ ) ++ pd_transfer_speeds_gb_s.append( ++ req.time_stats.get_pd_transfer_speed_gb_s() ++ ) ++ pd_transfer_totals_mb.append( ++ req.time_stats.get_pd_transfer_total_mb() ++ ) ++ pd_prefill_retry_counts.append( ++ req.time_stats.get_pd_prefill_retry_count() ++ ) + + if not self.spec_algorithm.is_none(): + spec_verify_ct.append(req.spec_verify_ct) +@@ -1134,7 +1180,7 @@ class SchedulerOutputProcessorMixin: req.log_time_stats() # Send to detokenizer @@ -1630,6 +2393,102 @@ index 482bc6ca6..857cfa6a3 100644 if self.model_config.is_multimodal_gen: return self.send_to_detokenizer.send_output( +@@ -1149,6 +1195,17 @@ class SchedulerOutputProcessorMixin: + prefill_launch_delay=prefill_launch_delays, + prefill_launch_latency=prefill_launch_latencies, + prefill_finished_ts=prefill_finished_timestamps, ++ pd_prefill_bootstrap_queue_duration=pd_prefill_bootstrap_queue_durations, ++ pd_prefill_forward_duration=pd_prefill_forward_durations, ++ pd_prefill_transfer_queue_duration=pd_prefill_transfer_queue_durations, ++ pd_decode_prealloc_duration=pd_decode_prealloc_durations, ++ pd_decode_transfer_duration=pd_decode_transfer_durations, ++ pd_decode_forward_duration=pd_decode_forward_durations, ++ pd_bootstrap_duration=pd_bootstrap_durations, ++ pd_alloc_waiting_duration=pd_alloc_waiting_durations, ++ pd_transfer_speed_gb_s=pd_transfer_speeds_gb_s, ++ pd_transfer_total_mb=pd_transfer_totals_mb, ++ pd_prefill_retry_count=pd_prefill_retry_counts, + finished_reasons=finished_reasons, + decoded_texts=decoded_texts, + decode_ids=decode_ids_list, +@@ -1198,6 +1255,18 @@ class SchedulerOutputProcessorMixin: + prefill_launch_delays = [] + prefill_launch_latencies = [] + prefill_finished_timestamps = [] ++ profiling_enabled = envs.SLIME_ENABLE_PROFILING.get() ++ pd_prefill_bootstrap_queue_durations = [] if profiling_enabled else None ++ pd_prefill_forward_durations = [] if profiling_enabled else None ++ pd_prefill_transfer_queue_durations = [] if profiling_enabled else None ++ pd_decode_prealloc_durations = [] if profiling_enabled else None ++ pd_decode_transfer_durations = [] if profiling_enabled else None ++ pd_decode_forward_durations = [] if profiling_enabled else None ++ pd_bootstrap_durations = [] if profiling_enabled else None ++ pd_alloc_waiting_durations = [] if profiling_enabled else None ++ pd_transfer_speeds_gb_s = [] if profiling_enabled else None ++ pd_transfer_totals_mb = [] if profiling_enabled else None ++ pd_prefill_retry_counts = [] if profiling_enabled else None + retraction_counts = [] + for req in reqs: + if req.finished(): +@@ -1221,6 +1290,40 @@ class SchedulerOutputProcessorMixin: + prefill_finished_timestamps.append( + req.time_stats.get_prefill_finished_ts() + ) ++ if profiling_enabled: ++ pd_prefill_bootstrap_queue_durations.append( ++ req.time_stats.get_pd_prefill_bootstrap_queue_duration() ++ ) ++ pd_prefill_forward_durations.append( ++ req.time_stats.get_pd_prefill_forward_duration() ++ ) ++ pd_prefill_transfer_queue_durations.append( ++ req.time_stats.get_pd_prefill_transfer_queue_duration() ++ ) ++ pd_decode_prealloc_durations.append( ++ req.time_stats.get_pd_decode_prealloc_duration() ++ ) ++ pd_decode_transfer_durations.append( ++ req.time_stats.get_pd_decode_transfer_duration() ++ ) ++ pd_decode_forward_durations.append( ++ req.time_stats.get_pd_decode_forward_duration() ++ ) ++ pd_bootstrap_durations.append( ++ req.time_stats.get_pd_bootstrap_duration() ++ ) ++ pd_alloc_waiting_durations.append( ++ req.time_stats.get_pd_alloc_waiting_duration() ++ ) ++ pd_transfer_speeds_gb_s.append( ++ req.time_stats.get_pd_transfer_speed_gb_s() ++ ) ++ pd_transfer_totals_mb.append( ++ req.time_stats.get_pd_transfer_total_mb() ++ ) ++ pd_prefill_retry_counts.append( ++ req.time_stats.get_pd_prefill_retry_count() ++ ) + retraction_counts.append(req.retraction_count) + self.send_to_detokenizer.send_output( + BatchEmbeddingOutput( +@@ -1231,6 +1334,17 @@ class SchedulerOutputProcessorMixin: + prefill_launch_delay=prefill_launch_delays, + prefill_launch_latency=prefill_launch_latencies, + prefill_finished_ts=prefill_finished_timestamps, ++ pd_prefill_bootstrap_queue_duration=pd_prefill_bootstrap_queue_durations, ++ pd_prefill_forward_duration=pd_prefill_forward_durations, ++ pd_prefill_transfer_queue_duration=pd_prefill_transfer_queue_durations, ++ pd_decode_prealloc_duration=pd_decode_prealloc_durations, ++ pd_decode_transfer_duration=pd_decode_transfer_durations, ++ pd_decode_forward_duration=pd_decode_forward_durations, ++ pd_bootstrap_duration=pd_bootstrap_durations, ++ pd_alloc_waiting_duration=pd_alloc_waiting_durations, ++ pd_transfer_speed_gb_s=pd_transfer_speeds_gb_s, ++ pd_transfer_total_mb=pd_transfer_totals_mb, ++ pd_prefill_retry_count=pd_prefill_retry_counts, + finished_reasons=finished_reasons, + embeddings=embeddings, + prompt_tokens=prompt_tokens, diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 1a65a3c3d..f76606469 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -1958,7 +2817,7 @@ index f2ffa9909..6e4d1d460 100644 self, obj: InitWeightsSendGroupForRemoteInstanceReqInput, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 0914a5230..cce2d8a2b 100644 +index 0914a5230..33bb3844a 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -324,8 +324,12 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi @@ -1993,6 +2852,82 @@ index 0914a5230..cce2d8a2b 100644 self.is_pause_cond.notify_all() async def update_weights_from_disk( +@@ -1510,6 +1514,40 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi + self._add_metric_if_present( + recv_obj, "prefill_finished_ts", meta_info, i + ) ++ # PD disaggregation timing ++ self._add_metric_if_present( ++ recv_obj, "pd_prefill_bootstrap_queue_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_prefill_forward_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_prefill_transfer_queue_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_decode_prealloc_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_decode_transfer_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_decode_forward_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_bootstrap_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_alloc_waiting_duration", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_transfer_speed_gb_s", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_transfer_total_mb", meta_info, i ++ ) ++ self._add_metric_if_present( ++ recv_obj, "pd_prefill_retry_count", meta_info, i ++ ) + + if getattr(state.obj, "return_logprob", False): + self.convert_logprob_style( +@@ -1955,19 +1993,17 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi + if custom_labels + else self.metrics_collector.labels + ) +- if ( +- state.first_token_time == 0.0 +- and self.disaggregation_mode != DisaggregationMode.PREFILL +- ): ++ if state.first_token_time == 0.0: + state.first_token_time = state.last_time = time.time() + state.first_token_time_perf = time.perf_counter() + state.last_completion_tokens = completion_tokens +- self.metrics_collector.observe_time_to_first_token( +- labels, state.first_token_time - state.created_time +- ) ++ if self.disaggregation_mode != DisaggregationMode.PREFILL: ++ self.metrics_collector.observe_time_to_first_token( ++ labels, state.first_token_time - state.created_time ++ ) + else: + num_new_tokens = completion_tokens - state.last_completion_tokens +- if num_new_tokens: ++ if num_new_tokens > 0: + new_time = time.time() + interval = new_time - state.last_time + self.metrics_collector.observe_inter_token_latency( +@@ -1976,7 +2012,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi + num_new_tokens, + ) + state.last_time = new_time +- state.last_completion_tokens = completion_tokens ++ state.last_completion_tokens = completion_tokens + + if state.finished: + retraction_count = ( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 86b009df4..16ebd52ae 100644 --- a/python/sglang/srt/managers/tp_worker.py @@ -2299,8 +3234,224 @@ index 42b169728..8e799196a 100644 node = node.parent return delta +diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py +index 255d41ccc..f93bedb4d 100644 +--- a/python/sglang/srt/metrics/collector.py ++++ b/python/sglang/srt/metrics/collector.py +@@ -20,7 +20,10 @@ import time + from dataclasses import dataclass, field + from typing import Any, Dict, List, Optional, Union + +-from sglang.srt.disaggregation.utils import DisaggregationMode ++from sglang.srt.disaggregation.utils import ( ++ DisaggregationMode, ++ is_slime_profiling_enabled, ++) + from sglang.srt.environ import envs + from sglang.srt.metrics.utils import exponential_buckets, generate_buckets + from sglang.srt.model_executor.forward_batch_info import ForwardMode +@@ -77,6 +80,17 @@ class TimeStats: + # Number of prefill retries for this request + prefill_retry_count: int = 0 + ++ # Prefill-side durations forwarded via metadata transfer from P to D instance. ++ # Set on the decode instance after KV cache transfer completes. ++ fwd_prefill_bootstrap_queue_duration: Optional[float] = None ++ fwd_prefill_forward_duration: Optional[float] = None ++ fwd_prefill_transfer_queue_duration: Optional[float] = None ++ fwd_bootstrap_duration: Optional[float] = None ++ fwd_alloc_waiting_duration: Optional[float] = None ++ fwd_transfer_speed_gb_s: Optional[float] = None ++ fwd_transfer_total_mb: Optional[float] = None ++ fwd_prefill_retry_count: Optional[int] = None ++ + # Timestamp when prefill phase finishes, obtained from `time.time()`. + # Note that this differs from the other `_time` fields tracked by the + # `TimeStats` class, which are obtained from `time.perf_counter()`. +@@ -102,6 +116,148 @@ class TimeStats: + return self.prefill_finished_ts + return None + ++ # --- PD disaggregation timing getters --- ++ ++ def get_pd_prefill_bootstrap_queue_duration(self) -> Optional[float]: ++ """P instance: time spent in bootstrap queue before entering the wait queue.""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_bootstrap_queue_duration is not None: ++ return self.fwd_prefill_bootstrap_queue_duration ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.prefill_bootstrap_queue_entry_time > 0.0 ++ and self.wait_queue_entry_time > 0.0 ++ ): ++ return self.wait_queue_entry_time - self.prefill_bootstrap_queue_entry_time ++ return None ++ ++ def get_pd_prefill_forward_duration(self) -> Optional[float]: ++ """P instance: time for the actual prefill forward computation.""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_forward_duration is not None: ++ return self.fwd_prefill_forward_duration ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.forward_entry_time > 0.0 ++ and self.completion_time > 0.0 ++ ): ++ return self.completion_time - self.forward_entry_time ++ return None ++ ++ def get_pd_prefill_transfer_queue_duration(self) -> Optional[float]: ++ """P instance: time spent in the transfer queue (KV cache send).""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_transfer_queue_duration is not None: ++ return self.fwd_prefill_transfer_queue_duration ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.prefill_transfer_queue_entry_time > 0.0 ++ and self.completion_time > 0.0 ++ ): ++ return self.completion_time - self.prefill_transfer_queue_entry_time ++ return None ++ ++ def get_pd_decode_prealloc_duration(self) -> Optional[float]: ++ """D instance: time spent in the pre-alloc queue (waiting for KV cache slot allocation).""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.decode_prealloc_queue_entry_time > 0.0 ++ and self.decode_transfer_queue_entry_time > 0.0 ++ ): ++ return ( ++ self.decode_transfer_queue_entry_time ++ - self.decode_prealloc_queue_entry_time ++ ) ++ return None ++ ++ def get_pd_decode_transfer_duration(self) -> Optional[float]: ++ """D instance: time spent waiting for KV cache transfer to complete.""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.decode_transfer_queue_entry_time > 0.0 ++ and self.wait_queue_entry_time > 0.0 ++ ): ++ return self.wait_queue_entry_time - self.decode_transfer_queue_entry_time ++ return None ++ ++ def get_pd_decode_forward_duration(self) -> Optional[float]: ++ """D instance: time for the actual decode forward computation.""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.forward_entry_time > 0.0 ++ and self.completion_time > 0.0 ++ ): ++ return self.completion_time - self.forward_entry_time ++ return None ++ ++ def get_pd_bootstrap_duration(self) -> Optional[float]: ++ """Bootstrap handshake duration (both P and D instances).""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_bootstrap_duration is not None: ++ return self.fwd_bootstrap_duration ++ if ( ++ self.disagg_mode != DisaggregationMode.NULL ++ and self.bootstrap_duration > 0.0 ++ ): ++ return self.bootstrap_duration ++ return None ++ ++ def get_pd_alloc_waiting_duration(self) -> Optional[float]: ++ """KV cache allocation waiting duration (both P and D instances).""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_alloc_waiting_duration is not None: ++ return self.fwd_alloc_waiting_duration ++ if ( ++ self.disagg_mode != DisaggregationMode.NULL ++ and self.alloc_waiting_duration > 0.0 ++ ): ++ return self.alloc_waiting_duration ++ return None ++ ++ def get_pd_transfer_speed_gb_s(self) -> Optional[float]: ++ """KV cache transfer speed in GB/s.""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_transfer_speed_gb_s is not None: ++ return self.fwd_transfer_speed_gb_s ++ if ( ++ self.disagg_mode != DisaggregationMode.NULL ++ and self.transfer_speed_gb_s > 0.0 ++ ): ++ return self.transfer_speed_gb_s ++ return None ++ ++ def get_pd_transfer_total_mb(self) -> Optional[float]: ++ """Total KV cache transferred in MB.""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_transfer_total_mb is not None: ++ return self.fwd_transfer_total_mb ++ if self.disagg_mode != DisaggregationMode.NULL and self.transfer_total_mb > 0.0: ++ return self.transfer_total_mb ++ return None ++ ++ def get_pd_prefill_retry_count(self) -> Optional[int]: ++ """Number of prefill retries for this request.""" ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_retry_count is not None: ++ return self.fwd_prefill_retry_count ++ if self.disagg_mode == DisaggregationMode.PREFILL: ++ return self.prefill_retry_count ++ return None ++ + def convert_to_duration(self) -> str: + if self.disagg_mode == DisaggregationMode.NULL: + queue_duration = self.forward_entry_time - self.wait_queue_entry_time +diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py +index 234523532..f5d479945 100644 +--- a/python/sglang/srt/model_executor/forward_batch_info.py ++++ b/python/sglang/srt/model_executor/forward_batch_info.py +@@ -909,6 +909,28 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): + tokens_padded = (tokens + rank_size - 1) // rank_size * rank_size + self._pad_inputs_to_size(model_runner, tokens_padded, self.batch_size) + ++ def prepare_cp_padding(self, model_runner: ModelRunner): ++ """Pad input_ids and extend_num_tokens to CP size multiples. ++ ++ In the PP disagg prefill + CP path, MLP sync is skipped so ++ prepare_mlp_sync_batch never runs. This method performs the ++ subset of padding that CP collective communication requires: ++ input_ids (and related tensors) must be divisible by cp_size. ++ """ ++ attn_cp_size = get_attention_cp_size() ++ if attn_cp_size <= 1: ++ return ++ if not self.forward_mode.is_extend(): ++ return ++ ++ tokens = self.input_ids.shape[0] ++ tokens_padded = ceil_align(tokens, attn_cp_size) ++ if tokens_padded == tokens: ++ return ++ ++ self._pad_inputs_to_size(model_runner, tokens_padded, self.batch_size) ++ self.extend_num_tokens = tokens_padded ++ + def post_forward_mlp_sync_batch(self, logits_output: LogitsProcessorOutput): + + self.forward_mode = getattr(self, "_original_forward_mode", self.forward_mode) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 275775a73..e4e2fdc39 100644 +index 275775a73..f0bd3ebf8 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -395,7 +395,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): @@ -2352,7 +3503,17 @@ index 275775a73..e4e2fdc39 100644 if self.eplb_manager is not None: self.eplb_manager.on_forward_pass_end() -@@ -2664,6 +2678,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -2472,6 +2486,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): + forward_batch.prepare_mlp_sync_batch(self) + else: + forward_batch.prepare_attn_tp_scatter_input(self) ++ # In PP disagg prefill + CP, MLP sync is skipped so CP padding ++ # must be done separately to keep input_ids divisible by cp_size. ++ forward_batch.prepare_cp_padding(self) + + # Normalize num_token_non_padded to be local to this attention TP rank if needed. + if ( +@@ -2664,6 +2681,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): device=self.device, ) @@ -2415,7 +3576,7 @@ index cc673a9ca..06c430d2c 100644 return AttnForwardMethod.MHA_ONE_SHOT return AttnForwardMethod.MLA diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py -index cb13a7c67..d62111471 100644 +index cb13a7c67..d9669ce08 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -29,6 +29,7 @@ from sglang.srt.layers.attention.nsa.utils import ( @@ -2426,16 +3587,289 @@ index cb13a7c67..d62111471 100644 is_nsa_enable_prefill_cp, nsa_use_prefill_cp, prepare_input_dp_with_cp_dsa, -@@ -160,6 +161,7 @@ class DeepseekModelNextN(nn.Module): +@@ -160,15 +161,17 @@ class DeepseekModelNextN(nn.Module): if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp): hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) + positions = cp_split_and_rebuild_position(forward_batch, positions) residual = None with get_global_expert_distribution_recorder().disable_this_region(): - hidden_states, residual = self.decoder( +- hidden_states, residual = self.decoder( ++ hidden_states, residual, *rest = self.decoder( + positions, + hidden_states, + forward_batch, + residual, + zero_allocator, + ) ++ topk_indices = rest[0] if rest else None + + if not forward_batch.forward_mode.is_idle(): + if residual is not None: +diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py +index 1583dd788..a35c00f96 100644 +--- a/python/sglang/srt/models/deepseek_v2.py ++++ b/python/sglang/srt/models/deepseek_v2.py +@@ -1085,6 +1085,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + skip_rope: bool = False, ++ is_nextn: bool = False, + ) -> None: + super().__init__() + self.layer_id = layer_id +@@ -1154,6 +1155,8 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + prefix=add_prefix("kv_a_proj_with_mqa", prefix), + ) + ++ self.skip_topk = False ++ self.next_skip_topk = False + if self.use_nsa: + is_neox_style = not getattr(config, "indexer_rope_interleave", False) + self.indexer = Indexer( +@@ -1174,6 +1177,31 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + layer_id=layer_id, + alt_stream=alt_stream, + ) ++ if not is_nextn: ++ self.index_topk_freq = getattr(config, "index_topk_freq", 1) ++ self.index_topk_pattern = getattr(config, "index_topk_pattern", None) ++ self.index_skip_topk_offset = getattr( ++ config, "index_skip_topk_offset", 2 ++ ) ++ if self.index_topk_pattern is None: ++ self.skip_topk = ( ++ max(layer_id - self.index_skip_topk_offset + 1, 0) ++ % self.index_topk_freq ++ != 0 ++ ) ++ self.next_skip_topk = ( ++ max(layer_id - self.index_skip_topk_offset + 2, 0) ++ % self.index_topk_freq ++ != 0 ++ ) ++ else: ++ self.skip_topk = self.index_topk_pattern[layer_id] == "S" ++ if layer_id < len(self.index_topk_pattern) - 1: ++ self.next_skip_topk = ( ++ self.index_topk_pattern[layer_id + 1] == "S" ++ ) ++ else: ++ self.next_skip_topk = False + + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, +@@ -1362,6 +1390,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + llama_4_scaling: Optional[torch.Tensor] = None, ++ prev_topk_indices: Optional[torch.Tensor] = None, + ): + s = self.forward_prepare( + positions=positions, +@@ -1369,6 +1398,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + forward_batch=forward_batch, + zero_allocator=zero_allocator, + llama_4_scaling=llama_4_scaling, ++ prev_topk_indices=prev_topk_indices, + ) + return self.forward_core(s) + +@@ -1379,6 +1409,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + llama_4_scaling: Optional[torch.Tensor] = None, ++ prev_topk_indices: Optional[torch.Tensor] = None, + ): + if self.attn_mha.kv_b_proj is None: + self.attn_mha.kv_b_proj = self.kv_b_proj +@@ -1418,7 +1449,12 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + ) + elif attn_forward_method == AttnForwardMethod.MLA: + inner_state = self.forward_absorb_prepare( +- positions, hidden_states, forward_batch, zero_allocator, llama_4_scaling ++ positions, ++ hidden_states, ++ forward_batch, ++ zero_allocator, ++ llama_4_scaling, ++ prev_topk_indices, + ) + elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE: + inner_state = self.forward_absorb_fused_mla_rope_prepare( +@@ -1529,6 +1565,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + llama_4_scaling: Optional[torch.Tensor] = None, ++ prev_topk_indices: Optional[torch.Tensor] = None, + ): + from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode + +@@ -1620,18 +1657,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + q = self.q_b_proj(q)[0].view( + -1, self.num_local_heads, self.qk_head_dim + ) +- topk_indices = self.indexer( +- x=hidden_states, +- q_lora=q_lora, +- positions=positions, +- forward_batch=forward_batch, +- layer_id=self.layer_id, +- ) +- current_stream.wait_stream(self.alt_stream) +- else: +- k_nope = k_nope.unsqueeze(1) +- q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) +- if q_lora is not None: ++ if not self.skip_topk: + topk_indices = self.indexer( + x=hidden_states, + q_lora=q_lora, +@@ -1639,6 +1665,23 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + forward_batch=forward_batch, + layer_id=self.layer_id, + ) ++ else: ++ topk_indices = prev_topk_indices ++ current_stream.wait_stream(self.alt_stream) ++ else: ++ k_nope = k_nope.unsqueeze(1) ++ q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) ++ if q_lora is not None: ++ if not self.skip_topk: ++ topk_indices = self.indexer( ++ x=hidden_states, ++ q_lora=q_lora, ++ positions=positions, ++ forward_batch=forward_batch, ++ layer_id=self.layer_id, ++ ) ++ else: ++ topk_indices = prev_topk_indices + else: + q = self.q_proj(hidden_states)[0].view( + -1, self.num_local_heads, self.qk_head_dim +@@ -1929,8 +1972,10 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): + ).transpose(0, 1), + ) + output, _ = self.o_proj(attn_bmm_output) +- +- return output ++ if not self.next_skip_topk: ++ return output, None ++ else: ++ return output, topk_indices + + def forward_absorb_fused_mla_rope_prepare( + self, +@@ -2275,6 +2320,7 @@ class DeepseekV2DecoderLayer(nn.Module): + reduce_results=False, + prefix=add_prefix("self_attn", prefix), + alt_stream=alt_stream, ++ is_nextn=is_nextn, + ) + + self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn) +@@ -2357,6 +2403,7 @@ class DeepseekV2DecoderLayer(nn.Module): + zero_allocator: BumpAllocator, + gemm_output_zero_allocator: BumpAllocator = None, + llama_4_scaling: Optional[torch.Tensor] = None, ++ prev_topk_indices: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + quant_format = ( + "mxfp4" +@@ -2398,7 +2445,12 @@ class DeepseekV2DecoderLayer(nn.Module): + forward_batch=forward_batch, + zero_allocator=zero_allocator, + llama_4_scaling=llama_4_scaling, ++ prev_topk_indices=prev_topk_indices, + ) ++ if isinstance(hidden_states, tuple): ++ hidden_states, topk_indices = hidden_states ++ else: ++ topk_indices = None + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch +@@ -2434,7 +2486,7 @@ class DeepseekV2DecoderLayer(nn.Module): + hidden_states, residual, forward_batch + ) + +- return hidden_states, residual ++ return hidden_states, residual, topk_indices + + def op_comm_prepare_attn( + self, +@@ -2710,6 +2762,7 @@ class DeepseekV2Model(nn.Module): + elif self.first_k_dense_replace < normal_start_layer: + normal_end_layer = normal_start_layer = 0 + aux_hidden_states = [] ++ topk_indices = None + for i in range(normal_start_layer, normal_end_layer): + # NOTE: torch dynamo does not support graph break in context manager + ctx = ( +@@ -2727,7 +2780,7 @@ class DeepseekV2Model(nn.Module): + else: + aux_hidden_states.append(hidden_states + residual) + layer = self.layers[i] +- hidden_states, residual = layer( ++ hidden_states, residual, *rest = layer( + positions, + hidden_states, + forward_batch, +@@ -2735,7 +2788,9 @@ class DeepseekV2Model(nn.Module): + zero_allocator, + gemm_output_zero_allocator, + llama_4_scaling, ++ prev_topk_indices=topk_indices, + ) ++ topk_indices = rest[0] if rest else None + + if normal_end_layer != self.end_layer: + hidden_states, residual = model_forward_maybe_tbo( +diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py +index db8c1c7ce..53ffadf6d 100644 +--- a/python/sglang/srt/models/glm4_moe.py ++++ b/python/sglang/srt/models/glm4_moe.py +@@ -678,8 +678,13 @@ class Glm4MoeDecoderLayer(nn.Module): + nn.Module.__init__(self) + self.hidden_size = config.hidden_size + self.config = config +- rope_theta = getattr(config, "rope_theta", 10000) +- rope_scaling = getattr(config, "rope_scaling", None) ++ # rope_theta may be stored in rope_parameters dict (e.g. GLM-4.6V) ++ _rope_params = getattr(config, "rope_parameters", None) ++ if isinstance(_rope_params, dict) and "rope_theta" in _rope_params: ++ rope_theta = _rope_params["rope_theta"] ++ else: ++ rope_theta = getattr(config, "rope_theta", 10000) ++ rope_scaling = getattr(config, "rope_scaling", None) or _rope_params + partial_rotary_factor = getattr( + getattr(config, "rope_parameters", None), "partial_rotary_factor", None + ) or getattr(config, "partial_rotary_factor", 0.5) +@@ -773,6 +778,7 @@ class Glm4MoeDecoderLayer(nn.Module): + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], ++ **kwargs, + ) -> torch.Tensor: + + hidden_states, residual = self.layer_communicator.prepare_attn( +diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py +index 1f6e75364..546cce4ab 100644 +--- a/python/sglang/srt/models/glm4_moe_nextn.py ++++ b/python/sglang/srt/models/glm4_moe_nextn.py +@@ -103,7 +103,7 @@ class Glm4MoeModelNextN(nn.Module): + + residual = None + with get_global_expert_distribution_recorder().disable_this_region(): +- hidden_states, residual = self.decoder( ++ hidden_states, residual, *rest = self.decoder( + positions, hidden_states, forward_batch, residual + ) + diff --git a/python/sglang/srt/models/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py -index 324de18b4..c99723f49 100644 +index 324de18b4..fc72faa03 100644 --- a/python/sglang/srt/models/glm4v_moe.py +++ b/python/sglang/srt/models/glm4v_moe.py @@ -52,11 +52,31 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): @@ -2475,7 +3909,7 @@ index 324de18b4..c99723f49 100644 self.visual = Glm4vVisionModel( config.vision_config, quant_config=quant_config, -@@ -64,21 +84,6 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +@@ -64,24 +84,14 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): use_data_parallel=self.use_data_parallel, ) @@ -2496,8 +3930,17 @@ index 324de18b4..c99723f49 100644 - self.logits_processor = LogitsProcessor(config) self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True) - self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling -@@ -219,6 +224,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +- self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling ++ _rope_cfg = ( ++ getattr(self.config, "rope_scaling", None) ++ or getattr(self.config, "rope_parameters", None) ++ or {} ++ ) ++ self.is_mrope_enabled = "mrope_section" in _rope_cfg + + # For EAGLE3 support + self.capture_aux_hidden_states = False +@@ -219,6 +229,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue @@ -2509,7 +3952,7 @@ index 324de18b4..c99723f49 100644 if name not in params_dict: continue -@@ -234,6 +244,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +@@ -234,6 +249,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue @@ -2518,7 +3961,7 @@ index 324de18b4..c99723f49 100644 # Mark as expert weight regardless of whether we can process it is_expert_weight = True -@@ -265,6 +277,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): +@@ -265,6 +282,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue @@ -2644,6 +4087,154 @@ index 2cf813bce..1250c49e4 100644 def _canonicalize_weights(config, weights_in: Iterable[Tuple[str, torch.Tensor]]): weights_out_dict = dict(weights_in) +diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py +index d8399a691..0277bc671 100644 +--- a/python/sglang/srt/models/kimi_k25.py ++++ b/python/sglang/srt/models/kimi_k25.py +@@ -666,25 +666,30 @@ class KimiK25ForConditionalGeneration(nn.Module): + self.config = config + self.quant_config = quant_config + self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder +- # Create vision tower +- self.vision_tower = MoonViT3dPretrainedModel( +- config.vision_config, use_data_parallel=self.use_data_parallel +- ) +- # Create mm projector +- self.mm_projector = K2VLMultiModalProjector(config.vision_config) + +- self.language_model = DeepseekV3ForCausalLM(config.text_config, quant_config) ++ # EPD: conditionally create components based on encoder_only / language_only ++ if not getattr(self.config, "language_only", False): ++ # Create vision tower and mm projector (needed for encoder_only and normal mode) ++ self.vision_tower = MoonViT3dPretrainedModel( ++ config.vision_config, use_data_parallel=self.use_data_parallel ++ ) ++ self.mm_projector = K2VLMultiModalProjector(config.vision_config) ++ ++ if not getattr(self.config, "encoder_only", False): ++ # Create language model (needed for language_only and normal mode) ++ self.language_model = DeepseekV3ForCausalLM( ++ config.text_config, quant_config ++ ) + +- # Ensure that the dtype of the vision_tower and mm_projector matches that of the language_model. +- # This solves the dtype mismatch issue when using device_map="auto" and torch_dtype. +- if hasattr(self.language_model, "dtype"): +- target_dtype = self.language_model.dtype +- self.vision_tower = self.vision_tower.to(dtype=target_dtype) +- self.mm_projector = self.mm_projector.to(dtype=target_dtype) ++ # Ensure dtype consistency between vision and language components ++ if hasattr(self, "vision_tower") and hasattr(self.language_model, "dtype"): ++ target_dtype = self.language_model.dtype ++ self.vision_tower = self.vision_tower.to(dtype=target_dtype) ++ self.mm_projector = self.mm_projector.to(dtype=target_dtype) + + def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor: +- pixel_values = torch.cat([item.feature for item in items], dim=0).type( +- self.vision_tower.dtype ++ pixel_values = torch.cat([item.feature for item in items], dim=0).to( ++ dtype=self.vision_tower.dtype, device=self.vision_tower.device + ) + grid_thws = torch.concat([item.grid_thws for item in items], dim=0).to( + self.vision_tower.device +@@ -735,41 +740,59 @@ class KimiK25ForConditionalGeneration(nn.Module): + + return hidden_states + ++ def set_eagle3_layers_to_capture(self, layer_ids=None): ++ self.language_model.set_eagle3_layers_to_capture(layer_ids) ++ ++ def get_embed_and_head(self): ++ return self.language_model.get_embed_and_head() ++ ++ def set_embed_and_head(self, embed, head): ++ self.language_model.set_embed_and_head(embed, head) ++ + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + """Load weights for the model, separating vision and language weights""" + mapper = getattr(self, "hf_to_sglang_mapper", None) + if mapper is not None: + weights = mapper.apply(weights) + ++ is_encoder_only = getattr(self.config, "encoder_only", False) ++ is_language_only = getattr(self.config, "language_only", False) ++ + # Separate vision tower weights and language model weights + vision_weights = [] + language_weights = [] + + for name, loaded_weight in weights: + if "vision_tower" in name or "mm_projector" in name: ++ # Skip vision weights in language_only mode ++ if is_language_only: ++ continue + name = name.replace(r"wqkv.", r"attn.qkv_proj.") + name = name.replace(r"wo.", r"attn.proj.") + name = name.replace("mm_projector.proj.0", "mm_projector.linear_1") + name = name.replace("mm_projector.proj.2", "mm_projector.linear_2") + vision_weights.append((name, loaded_weight)) + else: ++ # Skip language weights in encoder_only mode ++ if is_encoder_only: ++ continue + name = name.replace("language_model.", "") + # All other weights go to language model + language_weights.append((name, loaded_weight)) + + # Load vision tower weights +- vision_state_dict = dict(vision_weights) +- params_dict = dict(self.named_parameters(remove_duplicate=False)) +- for name, loaded_weight in vision_state_dict.items(): +- if name not in params_dict: +- raise ValueError(f"Weight {name} not found in params_dict") +- param = params_dict[name] +- weight_loader = getattr(param, "weight_loader", default_weight_loader) +- # loaded_weight = self._pad_vit_attn_dummy_heads(name, loaded_weight) +- weight_loader(param, loaded_weight) ++ if not is_language_only: ++ vision_state_dict = dict(vision_weights) ++ params_dict = dict(self.named_parameters(remove_duplicate=False)) ++ for name, loaded_weight in vision_state_dict.items(): ++ if name not in params_dict: ++ raise ValueError(f"Weight {name} not found in params_dict") ++ param = params_dict[name] ++ weight_loader = getattr(param, "weight_loader", default_weight_loader) ++ weight_loader(param, loaded_weight) + + # Load language model weights +- if language_weights: ++ if not is_encoder_only and language_weights: + self.language_model.load_weights(language_weights) + + +diff --git a/python/sglang/srt/models/llama_eagle3.py b/python/sglang/srt/models/llama_eagle3.py +index 49f938a1c..8eea383bb 100644 +--- a/python/sglang/srt/models/llama_eagle3.py ++++ b/python/sglang/srt/models/llama_eagle3.py +@@ -85,6 +85,11 @@ class LlamaDecoderLayer(LlamaDecoderLayer): + embeds = self.input_layernorm(embeds) + hidden_states = self.hidden_norm(hidden_states) + ++ if embeds.dtype != hidden_states.dtype: ++ raise RuntimeError( ++ f"Eagle3 dtype mismatch: embeds.dtype={embeds.dtype}, " ++ f"hidden_states.dtype={hidden_states.dtype}" ++ ) + hidden_states = torch.cat([embeds, hidden_states], dim=-1) + # Self Attention + hidden_states = self.self_attn( +@@ -160,6 +165,11 @@ class LlamaModel(nn.Module): + + hidden_states = forward_batch.spec_info.hidden_states + if hidden_states.shape[-1] != embeds.shape[-1]: ++ if hidden_states.dtype != self.fc.weight.dtype: ++ raise RuntimeError( ++ f"Eagle3 dtype mismatch: hidden_states.dtype={hidden_states.dtype}, " ++ f"fc.weight.dtype={self.fc.weight.dtype}" ++ ) + hidden_states = self.fc(hidden_states) + + # idle batch diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index f01225487..1dad8bb8e 100644 --- a/python/sglang/srt/models/qwen3_5.py @@ -2808,6 +4399,29 @@ index 33cce6fe2..0970c4550 100644 async def process_mm_data_async( self, image_data: List[Union[str, bytes]], +diff --git a/python/sglang/srt/multimodal/processors/kimi_k25.py b/python/sglang/srt/multimodal/processors/kimi_k25.py +index d8bb9ceb3..9311a431b 100644 +--- a/python/sglang/srt/multimodal/processors/kimi_k25.py ++++ b/python/sglang/srt/multimodal/processors/kimi_k25.py +@@ -25,6 +25,18 @@ class KimiK2_5VLImageProcessor(SGLangBaseProcessor): + image_token_id=hf_config.media_placeholder_token_id, + image_token_regex=re.compile(r"(?:<\|media_pad\|>)+"), + ).build(_processor) ++ # Required by base class get_mm_data / build_input_ids for EPD mode ++ self.IM_TOKEN_ID = hf_config.media_placeholder_token_id ++ self.IM_START_TOKEN_ID = None ++ self.IM_END_TOKEN_ID = None ++ merge_kernel = getattr(hf_config.vision_config, "merge_kernel_size", [2, 2]) ++ self._spatial_merge_size = ( ++ merge_kernel[0] if isinstance(merge_kernel, (list, tuple)) else merge_kernel ++ ) ++ ++ @property ++ def spatial_merge_size(self): ++ return self._spatial_merge_size + + async def process_mm_data_async( + self, diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py index 4395654e4..f9b5ea4ab 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py @@ -2822,7 +4436,7 @@ index 4395654e4..f9b5ea4ab 100644 image_data=image_data, video_data=request_obj.video_data, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index b080aeb16..957a613fa 100644 +index b080aeb16..b0322fef4 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -580,6 +580,7 @@ class ServerArgs: @@ -2833,7 +4447,15 @@ index b080aeb16..957a613fa 100644 disable_cuda_graph_padding: bool = False enable_profile_cuda_graph: bool = False enable_cudagraph_gc: bool = False -@@ -2089,7 +2090,16 @@ class ServerArgs: +@@ -635,6 +636,7 @@ class ServerArgs: + # Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 + enable_nsa_prefill_context_parallel: bool = False + nsa_prefill_cp_mode: str = "round-robin-split" ++ disable_indexer_rope_neox_style: bool = False + enable_fused_qk_norm_rope: bool = False + enable_precise_embedding_interpolation: bool = False + +@@ -2089,7 +2091,16 @@ class ServerArgs: assert ( self.tp_size % (self.dp_size * self.attn_cp_size) == 0 ), "tp_size must be divisible by dp_size * attn_cp_size" @@ -2851,7 +4473,7 @@ index b080aeb16..957a613fa 100644 if self.moe_dp_size > 1: # The tp_size is the world size, not the real tensor parallel size -@@ -4491,6 +4501,11 @@ class ServerArgs: +@@ -4491,6 +4502,11 @@ class ServerArgs: action="store_true", help="Disable cuda graph.", ) @@ -2863,7 +4485,20 @@ index b080aeb16..957a613fa 100644 parser.add_argument( "--disable-cuda-graph-padding", action="store_true", -@@ -5636,6 +5651,54 @@ class PortArgs: +@@ -4781,6 +4797,12 @@ class ServerArgs: + help="Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: 'round-robin-split'(default), 'in-seq-split' " + "'round-robin-split' distributes tokens across ranks based on token_idx %% cp_size. It supports multi-batch prefill, fused MoE, and FP8 KV cache.", + ) ++ parser.add_argument( ++ "--disable-indexer-rope-neox-style", ++ action="store_true", ++ help="Disable NSA indexer RoPE neox style (equivalent to INDEXER_ROPE_NEOX_STYLE=0). " ++ "If the environment variable INDEXER_ROPE_NEOX_STYLE is also set and conflicts, an error is raised.", ++ ) + parser.add_argument( + "--enable-fused-qk-norm-rope", + action="store_true", +@@ -5636,6 +5658,54 @@ class PortArgs: ) if not server_args.enable_dp_attention: diff --git a/setup.py b/setup.py index ce28b0cddc..c424a3a33d 100644 --- a/setup.py +++ b/setup.py @@ -32,7 +32,7 @@ def get_tag(self): setup( author="slime Team", name="slime", - version="0.2.3", + version="0.2.4", packages=find_packages(include=["slime*", "slime_plugins*"]), include_package_data=True, install_requires=_fetch_requirements("requirements.txt"),