diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index fd50cc8f0..76964243d 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -601,9 +601,11 @@ def _handle_dspark(server_args: ServerArgs) -> None: ) if cfg.pp_size != 1: - raise ValueError( - "Currently DSpark speculative decoding only supports pp_size == 1." - ) + if cfg.disaggregation_mode != "prefill": + raise ValueError( + "DSpark pipeline parallelism requires PD prefill; " + "decode and non-disaggregated serving require pp_size == 1." + ) if cfg.speculative_draft_model_path is None: if _target_checkpoint_bundles_dspark_draft(server_args): diff --git a/python/sglang/srt/arg_groups/validation_hook.py b/python/sglang/srt/arg_groups/validation_hook.py index 335cb0222..8a1e15c95 100644 --- a/python/sglang/srt/arg_groups/validation_hook.py +++ b/python/sglang/srt/arg_groups/validation_hook.py @@ -57,7 +57,14 @@ def check_pipeline_parallel_compat( assert cfg.disable_overlap_schedule, ( "Pipeline parallelism is not compatible with overlap schedule" ) - if cfg.speculative_algorithm is not None: + if cfg.speculative_algorithm == "DSPARK": + assert cfg.disaggregation_mode == "prefill", ( + "Pipeline parallel DSPARK requires disaggregation-mode=prefill" + ) + assert not envs.SGLANG_ENABLE_PP_SPEC.get(), ( + "SGLANG_ENABLE_PP_SPEC does not support DSPARK PD prefill" + ) + elif cfg.speculative_algorithm is not None: assert ( cfg.speculative_algorithm.upper() == "EAGLE" and not cfg.enable_multi_layer_eagle diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 6c863d93f..25cbf4c80 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -154,6 +154,7 @@ class KVArgsRegisterInfo: dst_dcp_rank: int = 0 requires_dcp_relayout: bool = False dcp_token_item_lens: Optional[List[int]] = None + dst_kv_item_lens: List[int] = dataclasses.field(default_factory=list) staging_base_ptr: int = 0 staging_total_size: int = 0 staging: Optional[StagingRegisterInfo] = None @@ -201,6 +202,11 @@ class KVArgsRegisterInfo: dst_dcp_rank=( int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0 ), + dst_kv_item_lens=( + list(struct.unpack(f"{len(msg[19]) // 8}Q", msg[19])) + if len(msg) > 19 and msg[19] + else [] + ), # Note: always put the staging field at the final staging=StagingRegisterInfo.from_zmq_fields(msg, 14, slot_ids_index=18), ) @@ -1090,11 +1096,16 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): executor: concurrent.futures.ThreadPoolExecutor, dst_layer_ids: List[int], pack_buffer=None, + dst_kv_item_lens: Optional[List[int]] = None, + dst_tp_rank: int = 0, + dst_attn_tp_size: Optional[int] = None, ) -> int: if num_kv_tokens is None: raise ValueError("PD DCP transfer requires num_kv_tokens") physical_page_size = self.kv_args.page_size + if dst_kv_item_lens and len(dst_kv_item_lens) != len(dst_kv_ptrs): + raise ValueError("PD DCP destination KV lengths must match its buffers") src_layer_ids = self.kv_args.kv_layer_ids if src_layer_ids or dst_layer_ids: dst_indices = resolve_dcp_dst_entry_indices( @@ -1105,11 +1116,17 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): ) src_kv_ptrs = self.kv_args.kv_data_ptrs dst_kv_ptrs = [dst_kv_ptrs[j] for j in dst_indices] + if dst_kv_item_lens: + dst_kv_item_lens = [dst_kv_item_lens[j] for j in dst_indices] else: src_kv_ptrs, dst_kv_ptrs, _ = self.get_mla_kv_ptrs_with_pp( self.kv_args.kv_data_ptrs, dst_kv_ptrs, ) + if dst_kv_item_lens: + _, dst_kv_item_lens, _ = self.get_mla_kv_ptrs_with_pp( + self.kv_args.kv_item_lens, dst_kv_item_lens + ) num_draft = self.kv_args.num_draft_entries num_target = len(src_kv_ptrs) - num_draft @@ -1155,20 +1172,87 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): ) for entry in range(num_target) ] + sliced_draft_params = [] if num_draft > 0 and plan.draft_src_token_indices.size: + if not dst_kv_item_lens and dst_attn_tp_size not in ( + None, + self.attn_tp_size, + ): + raise ValueError( + "PD DCP with different draft TP sizes requires destination KV lengths" + ) draft_groups = group_concurrent_contiguous( plan.draft_src_token_indices, plan.draft_dst_token_indices, ) - layers_params += [ - ( - src_kv_ptrs[num_target + entry], - dst_kv_ptrs[num_target + entry], - dcp_token_item_lens[num_target + entry], - draft_groups, + for entry in range(num_target, num_target + num_draft): + src_width = dcp_token_item_lens[entry] + dst_width = src_width + if dst_kv_item_lens: + dst_width, remainder = divmod( + dst_kv_item_lens[entry], physical_page_size * dst_dcp_size + ) + if remainder or dst_width <= 0: + raise ValueError("Invalid PD DCP draft destination token width") + if src_width == dst_width: + layers_params.append( + ( + src_kv_ptrs[entry], + dst_kv_ptrs[entry], + src_width, + draft_groups, + ) + ) + continue + if self.is_mla_backend: + raise ValueError( + "PD DCP draft head slicing is unsupported for pure MLA: " + "dummy prefill senders may omit draft head shards" + ) + copy_width = min(src_width, dst_width) + if max(src_width, dst_width) % copy_width: + raise ValueError("PD DCP draft KV head shards must divide evenly") + if dst_attn_tp_size is None: + raise ValueError( + "PD DCP draft head slicing requires destination TP size" + ) + src_span = src_width * self.attn_tp_size + dst_span = dst_width * dst_attn_tp_size + src_rank = (self.kv_args.engine_rank % self.attn_tp_size) // max( + 1, src_span // dst_span ) - for entry in range(num_draft) - ] + dst_rank = dst_tp_rank // max(1, dst_span // src_span) + src_offset = (dst_rank * dst_width) % src_width + dst_offset = (src_rank * src_width) % dst_width + sliced_draft_params.append( + ( + src_kv_ptrs[entry] + src_offset, + dst_kv_ptrs[entry] + dst_offset, + src_width, + dst_width, + copy_width, + ) + ) + + def process_sliced_draft(params) -> int: + batch_size = self.max_transfer_batch_indices + if batch_size <= 0: + batch_size = 4096 + for start in range(0, plan.draft_src_token_indices.size, batch_size): + src_indices = plan.draft_src_token_indices[start : start + batch_size] + dst_indices = plan.draft_dst_token_indices[start : start + batch_size] + blocks = [] + for src_ptr, dst_ptr, src_width, dst_width, copy_width in params: + src_addrs = src_ptr + src_indices * src_width + dst_addrs = dst_ptr + dst_indices * dst_width + blocks.extend( + (int(src), int(dst), copy_width) + for src, dst in zip(src_addrs, dst_addrs) + ) + ret = self._transfer_data(mooncake_session_id, blocks) + if ret != 0: + return ret + return 0 def set_transfer_blocks( src_ptr: int, dst_ptr: int, token_item_len: int, groups @@ -1196,12 +1280,19 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): executor.submit(process_layer, *layer_params) for layer_params in layers_params ] + futures.extend( + executor.submit(process_sliced_draft, [params]) + for params in sliced_draft_params + ) return self._await_transfer_futures(futures) transfer_blocks = [] for layer_params in layers_params: transfer_blocks.extend(set_transfer_blocks(*layer_params)) - return self._transfer_data(mooncake_session_id, transfer_blocks) + ret = self._transfer_data(mooncake_session_id, transfer_blocks) + if ret != 0 or not sliced_draft_params: + return ret + return process_sliced_draft(sliced_draft_params) def send_kvcache_slice( self, @@ -2114,6 +2205,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): target_rank_registration_info.dst_kv_layer_ids ), pack_buffer=pack_buffer, + dst_kv_item_lens=target_rank_registration_info.dst_kv_item_lens, + dst_tp_rank=target_rank_registration_info.dst_tp_rank, + dst_attn_tp_size=target_rank_registration_info.dst_attn_tp_size, ) elif ( self.is_mla_backend @@ -2810,6 +2904,10 @@ class MooncakeKVReceiver(MooncakeFailureExceptionMixin, CommonKVReceiver): dst_dcp_size, dst_dcp_rank, packed_staging_slot_layer_ids, + struct.pack( + f"{len(self.kv_mgr.kv_args.kv_item_lens)}Q", + *self.kv_mgr.kv_args.kv_item_lens, + ), ] ) except zmq.ZMQError: diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index 50918355d..f8229a634 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -886,7 +886,11 @@ class SchedulerPPMixin: # Draft extend runs only on the last stage, but every rank needs its relayed # output to fill PD auxiliary buffers. draft_input = result.next_draft_input - if draft_input is not None and draft_input.topk_p is not None: + if ( + draft_input is not None + and not batch.spec_algorithm.is_dspark() + and draft_input.topk_p is not None + ): tensor_dict["draft_topk_p"] = draft_input.topk_p.contiguous() tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous() tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous() @@ -1138,6 +1142,16 @@ class SchedulerPPMixin: dsa_topk_indices=pp_outputs.tensors.get("draft_dsa_topk_indices"), ) batch.spec_info = next_draft_input + elif batch.spec_algorithm.is_dspark(): + from sglang.srt.speculative.dspark_components.dspark_draft import ( + make_next_draft_input, + ) + + next_draft_input = make_next_draft_input( + bonus_tokens=next_token_ids, + new_seq_lens=batch.seq_lens, + ) + batch.spec_info = next_draft_input if self._pp_spec_relay: # Gated single-instance PP+spec: the sampled first token roots diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index c3b51276f..7c61d494d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -799,6 +799,13 @@ class ModelRunner: enable_batch_invariant_mode() + def get_pp_proxy_dspark_hidden_size(self) -> int: + return misc_utils.resolve_pp_proxy_dspark_hidden_size( + model=self.model, + pp_size=self.ps.pp_size, + pp_rank=self.ps.pp_rank, + ) + def get_pp_proxy_topk_size(self) -> Optional[int]: return misc_utils.resolve_pp_proxy_topk_size( model_config=self.model_config, diff --git a/python/sglang/srt/model_executor/model_runner_components/misc_utils.py b/python/sglang/srt/model_executor/model_runner_components/misc_utils.py index 7ce3a09f5..026ca0434 100644 --- a/python/sglang/srt/model_executor/model_runner_components/misc_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/misc_utils.py @@ -1,7 +1,7 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable from sglang.srt.configs.model_config import ( dsa_layer_skips_topk, @@ -90,3 +90,18 @@ def resolve_pp_proxy_residual_num_blocks( if block_size is None: return None return (start_layer + block_size - 1) // block_size + + +@runtime_checkable +class _SupportsDSparkPPProxy(Protocol): + def get_pp_proxy_dspark_hidden_size(self) -> int: ... + + +def resolve_pp_proxy_dspark_hidden_size( + *, model: Any, pp_size: int, pp_rank: int +) -> int: + if pp_size <= 1 or pp_rank == 0: + return 0 + if isinstance(model, _SupportsDSparkPPProxy): + return model.get_pp_proxy_dspark_hidden_size() + return 0 diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 63f24d61f..b5dda425f 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -418,6 +418,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): pp_proxy_residual_num_blocks=( self.model_runner.get_pp_proxy_residual_num_blocks() ), + pp_proxy_dspark_hidden_size=( + self.model_runner.get_pp_proxy_dspark_hidden_size() + ), ) self.buffers.share_buffers() # FB-shared slot registry adopting DecodeInputBuffers storage (same diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 3d0c2010a..17bfceb8b 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -389,6 +389,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): pp_proxy_residual_num_blocks=( self.model_runner.get_pp_proxy_residual_num_blocks() ), + pp_proxy_dspark_hidden_size=( + self.model_runner.get_pp_proxy_dspark_hidden_size() + ), ) self.buffers.share_buffers() # Token-axis FB-shared slot registry adopting PrefillInputBuffers @@ -599,8 +602,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): f"unsupported for this model architecture." ) from exc params = list(inspect.signature(self.layer_model.forward).parameters) - self._input_embeds_arg_idx = ( - params.index("input_embeds") if "input_embeds" in params else None + self._input_embeds_arg_idx = next( + ( + params.index(name) + for name in ("input_embeds", "inputs_embeds") + if name in params + ), + None, ) # --- aiter chip info pre-warming (AMD) ------------------------- @@ -1930,6 +1938,8 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): """A text-only batch would otherwise replay the captured input_embeds.""" ie_idx = self._input_embeds_arg_idx ie = layer_kwargs.get("input_embeds") + if ie is None: + ie = layer_kwargs.get("inputs_embeds") if ie is None and ie_idx is not None and len(args) > ie_idx: ie = args[ie_idx] if ie is None: @@ -1968,7 +1978,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # text-only batches they are get_input_embeddings()(input_ids). # Copy them into the slot before replay so the graph sees the # current request's embeddings (mirrors main's BCG closure). - if self.buffer_registry.has_slot("input_embeds"): + if ( + self.model_runner.pp_group.is_first_rank + and self.buffer_registry.has_slot("input_embeds") + ): self._fill_input_embeds_slot(args, layer_kwargs, static_num_tokens) hs = self.backend.replay(shape_key, static_forward_batch, **kwargs) return _slice_output_rows(hs, raw_num_tokens) if full_path else hs diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 47f7f772b..3a7469180 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/python/sglang/srt/model_executor/runner_utils/buffers.py @@ -67,6 +67,7 @@ def _allocate_pp_proxy_tensors( hc_hidden_size: Optional[int] = None, pp_proxy_topk_size: Optional[int] = None, pp_proxy_residual_num_blocks: Optional[int] = None, + pp_proxy_dspark_hidden_size: int = 0, ) -> Dict[str, torch.Tensor]: """Allocate the stable buffers consumed by an incoming PP proxy.""" is_mhc = hc_hidden_size is not None @@ -87,6 +88,10 @@ def _allocate_pp_proxy_tensors( pp_proxy_tensors["topk_indices"] = torch.zeros( (max_num_tokens, pp_proxy_topk_size), dtype=torch.int32 ) + if pp_proxy_dspark_hidden_size: + pp_proxy_tensors["dspark_hidden_states"] = torch.zeros( + (max_num_tokens, pp_proxy_dspark_hidden_size), dtype=dtype + ) return pp_proxy_tensors @@ -136,6 +141,7 @@ class DecodeInputBuffers(ForwardInputBuffers): hc_hidden_size: Optional[int] = None, pp_proxy_topk_size: Optional[int] = None, pp_proxy_residual_num_blocks: Optional[int] = None, + pp_proxy_dspark_hidden_size: int = 0, ) -> DecodeInputBuffers: with torch.device(device): input_ids = torch.zeros((max_num_token,), dtype=torch.int64) @@ -173,6 +179,7 @@ class DecodeInputBuffers(ForwardInputBuffers): hc_hidden_size=hc_hidden_size, pp_proxy_topk_size=pp_proxy_topk_size, pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks, + pp_proxy_dspark_hidden_size=pp_proxy_dspark_hidden_size, ) if pp_size > 1 else None @@ -275,6 +282,7 @@ class PrefillInputBuffers(ForwardInputBuffers): hc_hidden_size: Optional[int] = None, pp_proxy_topk_size: Optional[int] = None, pp_proxy_residual_num_blocks: Optional[int] = None, + pp_proxy_dspark_hidden_size: int = 0, ) -> PrefillInputBuffers: with torch.device(device): input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64) @@ -311,6 +319,7 @@ class PrefillInputBuffers(ForwardInputBuffers): hc_hidden_size=hc_hidden_size, pp_proxy_topk_size=pp_proxy_topk_size, pp_proxy_residual_num_blocks=pp_proxy_residual_num_blocks, + pp_proxy_dspark_hidden_size=pp_proxy_dspark_hidden_size, ) if pp_size > 1 and not is_first_pp_rank else None diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 578f443ee..b931ff98e 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -3047,6 +3047,18 @@ class KimiK3LinearModel(nn.Module): ) sp_sharded = False aux_hidden_states = [] + if ( + self.dspark_layers_to_capture is not None + and not self.pp_group.is_first_rank + ): + if "dspark_hidden_states" in pp_proxy_tensors.tensors: + aux_hidden_states.append(pp_proxy_tensors["dspark_hidden_states"]) + if self.start_layer - 1 in self.dspark_layers_to_capture: + aux_hidden_states.append( + self._dspark_capture_stream( + self.start_layer - 1, hidden_states, residual, attn_res + ) + ) for i in range(self.start_layer, self.end_layer): if sp_sharded and not self.layers[i]._sp_moe: hidden_states = _sp_all_gather_rows(hidden_states) @@ -3065,6 +3077,7 @@ class KimiK3LinearModel(nn.Module): if ( self.dspark_layers_to_capture is not None and i in self.dspark_layers_to_capture + and (i + 1 < self.end_layer or self.pp_group.is_last_rank) ): aux_hidden_states.append( self._dspark_capture_stream(i, hidden_states, residual, attn_res) @@ -3078,9 +3091,12 @@ class KimiK3LinearModel(nn.Module): # full stream head (bit-identical to the fused fold). hidden_states = residual + hidden_states residual = attn_res.block_residual # raw bank across ranks - return PPProxyTensors( - {"hidden_states": hidden_states, "residual": residual} - ) + proxy_tensors = {"hidden_states": hidden_states, "residual": residual} + if aux_hidden_states: + proxy_tensors["dspark_hidden_states"] = torch.cat( + aux_hidden_states, dim=-1 + ) + return PPProxyTensors(proxy_tensors) if hidden_states.shape[0] != 0: if attn_res is not None: @@ -3204,13 +3220,13 @@ class KimiK3LinearForCausalLM(nn.Module): def get_input_embeddings(self): return self.model.embed_tokens + def get_pp_proxy_dspark_hidden_size(self) -> int: + layers = self.model.dspark_layers_to_capture or [] + return self.config.hidden_size * sum( + layer < self.model.start_layer - 1 for layer in layers + ) + def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None: - if self.pp_group.world_size > 1: - # Capture layers living on non-last PP ranks would be silently - # skipped (the flag is only set on the last rank). - raise NotImplementedError("DSPARK aux hidden capture requires PP=1.") - if not self.pp_group.is_last_rank: - return if layer_ids is None: raise ValueError( "DSPARK requires explicit layer_ids for aux hidden capture." @@ -3667,6 +3683,11 @@ class KimiK3ForConditionalGeneration(nn.Module): raise AttributeError("lm_head is not available in encoder-only mode") return self.language_model.lm_head + def get_pp_proxy_dspark_hidden_size(self) -> int: + if self.language_model is None: + return 0 + return self.language_model.get_pp_proxy_dspark_hidden_size() + def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None: if self.language_model is None: raise AttributeError( diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 87fdb0ae4..fd4537c4a 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -717,6 +717,12 @@ class KimiLinearModel(nn.Module): device=device, ) aux_hidden_states = [] + if ( + self.dspark_layers_to_capture is not None + and not self.pp_group.is_first_rank + and "dspark_hidden_states" in pp_proxy_tensors.tensors + ): + aux_hidden_states.append(pp_proxy_tensors["dspark_hidden_states"]) for i in range(self.start_layer, self.end_layer): ctx = get_global_expert_distribution_recorder().with_current_layer(i) with ctx: @@ -737,12 +743,12 @@ class KimiLinearModel(nn.Module): ) if not self.pp_group.is_last_rank: - return PPProxyTensors( - { - "hidden_states": hidden_states, - "residual": residual, - } - ) + proxy_tensors = {"hidden_states": hidden_states, "residual": residual} + if aux_hidden_states: + proxy_tensors["dspark_hidden_states"] = torch.cat( + aux_hidden_states, dim=-1 + ) + return PPProxyTensors(proxy_tensors) else: if hidden_states.shape[0] != 0: if residual is None: @@ -787,11 +793,13 @@ class KimiLinearForCausalLM(nn.Module): def get_input_embeddings(self): return self.model.embed_tokens + def get_pp_proxy_dspark_hidden_size(self) -> int: + layers = self.model.dspark_layers_to_capture or [] + return self.config.hidden_size * sum( + layer < self.model.start_layer for layer in layers + ) + def set_dspark_layers_to_capture(self, layer_ids: list[int]) -> None: - if self.pp_group.world_size > 1: - raise NotImplementedError("DSPARK aux hidden capture requires PP=1.") - if not self.pp_group.is_last_rank: - return if layer_ids is None: raise ValueError( "DSPARK requires explicit layer_ids for aux hidden capture." diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index c3e197981..52f7f4547 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -69,6 +69,7 @@ def build_draft_tp_worker( algo_label: str, attention_backend_override: Optional[str] = None, draft_worker_cls: type[TpModelWorker] = TpModelWorker, + random_seed: Optional[int] = None, ) -> DraftWorkerBundle: # An override names a draft-specific backend the caller has already # validated (e.g. a self-drafting architecture); it skips the generic @@ -90,6 +91,7 @@ def build_draft_tp_worker( ps=ps, nccl_port=nccl_port, is_draft_worker=True, + random_seed=random_seed, # The draft runs at absolute target positions. context_length=target_model_config.context_len, draft_attention_backend=draft_backend, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 1f4038761..2d81867bc 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -9,6 +9,7 @@ from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( is_unified_kv_triton, ) from sglang.srt.configs.hybrid_arch import mambaish_config +from sglang.srt.distributed import get_pp_group from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.environ import envs from sglang.srt.layers.logprob_processor import compute_spec_logprobs @@ -20,6 +21,7 @@ from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardMode, + PPProxyTensors, compute_position, ) from sglang.srt.runtime_context import ( @@ -126,6 +128,8 @@ def _configure_target_hidden_projection( class DSparkWorkerV2(BaseSpecWorker): + """Non-last PP stages run only the target; draft state belongs to the last stage.""" + def __init__( self, server_args: ServerArgs, @@ -145,6 +149,10 @@ class DSparkWorkerV2(BaseSpecWorker): self.model_runner = target_worker.model_runner self.page_size = get_schedule().page_size self.device = target_worker.device + self._draft_worker = None + self._hosts_draft = get_pp_group().is_last_rank + if not self._hosts_draft: + return self._draft_is_moe = draft_is_deepseek_v4() self._draft_dp_context_enabled = ( @@ -178,6 +186,7 @@ class DSparkWorkerV2(BaseSpecWorker): DSV4_DRAFT_ATTENTION_BACKEND if self._draft_is_moe else None ), draft_worker_cls=draft_worker_cls, + random_seed=target_worker.random_seed, ) self._draft_worker = bundle.draft_worker self.draft_model_runner = bundle.draft_model_runner @@ -397,10 +406,12 @@ class DSparkWorkerV2(BaseSpecWorker): @property def carries_confidence(self) -> bool: - return self._verify_planner.carries_confidence + return self._hosts_draft and self._verify_planner.carries_confidence @property def spec_v2_attn_backends(self) -> tuple: + if not self._hosts_draft: + return super().spec_v2_attn_backends return ( self._target_worker.model_runner.attn_backend, self.draft_model_runner.attn_backend, @@ -422,6 +433,8 @@ class DSparkWorkerV2(BaseSpecWorker): req_to_token_pool=None, token_to_kv_pool_allocator=None, ): + if not self._hosts_draft: + return self._draft_worker.alloc_memory_pool( memory_pool_config=memory_pool_config, req_to_token_pool=req_to_token_pool, @@ -429,6 +442,8 @@ class DSparkWorkerV2(BaseSpecWorker): ) def init_attention_backends(self): + if not self._hosts_draft: + return with draft_pp_context(), self._draft_context(): self._draft_worker.init_attention_backends() self._target_hidden_projection_enabled = _configure_target_hidden_projection( @@ -449,6 +464,8 @@ class DSparkWorkerV2(BaseSpecWorker): ) def init_cuda_graphs(self): + if not self._hosts_draft: + return capture_decode_cuda_graph = self._decode_graph_allowed available_mem = self._tp_sync.available_memory_gb( SpecTpSyncSite.DSPARK_MEM, @@ -510,19 +527,29 @@ class DSparkWorkerV2(BaseSpecWorker): pass def set_dspark_forced_budget_frac(self, frac: Optional[float]) -> None: + if not self._hosts_draft: + return self._forced_budget_frac = frac self._verify_planner.set_forced_budget_frac(frac) def dump_info_records(self) -> Optional[dict]: + if not self._hosts_draft: + return None return self._observers.dump_info_records() def clear_info_records(self) -> None: + if not self._hosts_draft: + return self._observers.clear_info_records() def block_accept_estimate_log_suffix(self) -> Optional[str]: + if not self._hosts_draft: + return None return self._observers.block_accept_estimate_log_suffix() def note_request_finished(self, *, rid: str, natural_stop: bool) -> None: + if not self._hosts_draft: + return self._observers.note_request_finished(rid=rid, natural_stop=natural_stop) def forward_batch_generation( @@ -531,19 +558,30 @@ class DSparkWorkerV2(BaseSpecWorker): on_publish=None, grammar_barrier=None, *, - pp_proxy_tensors=None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> GenerationBatchResult: - # The non-overlap scheduler passes this keyword even when PP=1. - assert pp_proxy_tensors is None, "DSpark does not support pipeline parallelism" + if not self._hosts_draft: + batch_output = self.target_worker.forward_batch_generation( + batch, + pp_proxy_tensors=pp_proxy_tensors, + capture_hidden_mode=CaptureHiddenMode.FULL, + ) + batch_output.new_seq_lens = batch.seq_lens + if on_publish is not None: + on_publish(batch_output.new_seq_lens) + return batch_output if batch.forward_mode.is_extend() or batch.is_extend_in_batch: self._verify_planner.note_non_decode_step() self._observers.note_prefill_step() - return self._forward_prefill(batch, on_publish) + return self._forward_prefill(batch, on_publish, pp_proxy_tensors) return self._forward_decode(batch, on_publish, grammar_barrier) def _forward_prefill( - self, batch: ScheduleBatch, on_publish + self, + batch: ScheduleBatch, + on_publish, + pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> GenerationBatchResult: if batch.forward_mode.is_idle(): if get_parallel().enable_dp_attention: @@ -553,7 +591,9 @@ class DSparkWorkerV2(BaseSpecWorker): return self._decode_idle_result(on_publish=on_publish) batch_output = self.target_worker.forward_batch_generation( - batch, capture_hidden_mode=CaptureHiddenMode.FULL + batch, + pp_proxy_tensors=pp_proxy_tensors, + capture_hidden_mode=CaptureHiddenMode.FULL, ) # BCG replay skips model-side Python, so re-evaluate the same pure predicate. target_hidden_is_projected = ( @@ -1010,4 +1050,6 @@ class DSparkWorkerV2(BaseSpecWorker): ) def get_confidence_budget_prepare(self): + if not self._hosts_draft: + return None return self._verify_planner.confidence_budget_prepare() diff --git a/test/registered/kernel/attention/test_kda_gate_beta_cumsum.py b/test/registered/kernels/ops/attention/test_kda_gate_beta_cumsum.py similarity index 100% rename from test/registered/kernel/attention/test_kda_gate_beta_cumsum.py rename to test/registered/kernels/ops/attention/test_kda_gate_beta_cumsum.py diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index ea7247e45..0a9bc35d4 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -96,6 +96,9 @@ class TestDisaggregationWire(unittest.TestCase): self.assertEqual(info.staging_total_size, 4096) self.assertEqual(info.dst_dcp_size, 4) self.assertEqual(info.dst_dcp_rank, 2) + self.assertEqual(info.dst_kv_item_lens, []) + info = KVArgsRegisterInfo.from_zmq(msg + [b"", struct.pack("Q", 128)]) + self.assertEqual(info.dst_kv_item_lens, [128]) def test_int_lists_roundtrip(self): cases = [ diff --git a/test/registered/unit/disaggregation/test_mooncake_transfer_batching.py b/test/registered/unit/disaggregation/test_mooncake_transfer_batching.py index ec528b7f9..c9cb9e5c6 100644 --- a/test/registered/unit/disaggregation/test_mooncake_transfer_batching.py +++ b/test/registered/unit/disaggregation/test_mooncake_transfer_batching.py @@ -130,5 +130,161 @@ class TestMooncakeTransferBatching(unittest.TestCase): ) +class TestDcpDraftHeadTransfer(unittest.TestCase): + def test_transfers_draft_heads_to_logical_destination_rows(self): + for src_tp, dst_tp in ((4, 8), (8, 4), (8, 8), (4, 32), (32, 4)): + for custom_pool in (False, True): + for batch_size in (0, 37): + with self.subTest( + src_tp=src_tp, + dst_tp=dst_tp, + custom_pool=custom_pool, + batch_size=batch_size, + ): + self._check_transfer(src_tp, dst_tp, custom_pool, batch_size) + + def test_rejects_pure_mla_with_unequal_draft_head_widths(self): + for src_tp, dst_tp in ((4, 8), (8, 4)): + with self.subTest(src_tp=src_tp, dst_tp=dst_tp): + with self.assertRaisesRegex(ValueError, "dummy prefill senders"): + self._check_transfer(src_tp, dst_tp, False, 37, pure_mla=True) + + def test_sliced_draft_stops_after_failed_batch(self): + self._check_transfer(4, 8, False, 37, fail_draft=True) + + def _check_transfer( + self, src_tp, dst_tp, custom_pool, batch_size, fail_draft=False, pure_mla=False + ): + page_size, tokens, heads, head_bytes = 64, 249, 16, 4 + src_width, dst_width = ( + max(1, heads // src_tp) * head_bytes, + max(1, heads // dst_tp) * head_bytes, + ) + src_pages = np.array([1, 3, 4, 7], dtype=np.int32) + logical = np.arange(tokens) + src_rows = src_pages[logical // page_size] * page_size + logical % page_size + expected = ( + np.arange(tokens * heads * head_bytes, dtype=np.int64) + .reshape(tokens, heads, head_bytes) + .astype(np.uint8) + ) + for dst_rank in range(dst_tp): + dst_buffers = { + base: np.zeros(16384 * max(8, dst_width), dtype=np.uint8) + for base in (1000000, 2000000, 3000000, 4000000) + } + source_ranks = ( + range(dst_rank * src_tp // dst_tp, (dst_rank + 1) * src_tp // dst_tp) + if src_tp >= dst_tp + else [dst_rank * src_tp // dst_tp] + ) + for src_rank in source_ranks: + src_head_start = (src_rank // max(1, src_tp // heads)) * max( + 1, heads // src_tp + ) + source = np.zeros(1024 * src_width, dtype=np.uint8) + source.reshape(-1, src_width)[src_rows] = expected[ + :, src_head_start : src_head_start + max(1, heads // src_tp) + ].reshape(tokens, src_width) + target = np.zeros(1024 * 8, dtype=np.uint8) + target.reshape(-1, 8)[src_rows] = ( + np.arange(tokens * 8).reshape(tokens, 8).astype(np.uint8) + ) + src_buffers = {10000: target, 100000: source, 200000: source} + + failed_batches = [] + + def transfer( + session, blocks, src_buffers=src_buffers, dst_buffers=dst_buffers + ): + draft_blocks = [block for block in blocks if block[1] >= 3000000] + if fail_draft and draft_blocks: + failed_batches.append(draft_blocks) + return 17 + if batch_size and src_width != dst_width: + self.assertLessEqual( + len(draft_blocks), batch_size * (1 if custom_pool else 2) + ) + for src, dst, size in blocks: + src_base = max(base for base in src_buffers if base <= src) + dst_base = max(base for base in dst_buffers if base <= dst) + dst_buffers[dst_base][ + dst - dst_base : dst - dst_base + size + ] = src_buffers[src_base][ + src - src_base : src - src_base + size + ] + return 0 + + manager = SimpleNamespace( + is_mla_backend=pure_mla, + kv_args=SimpleNamespace( + page_size=page_size, + kv_layer_ids=[47, 93, 93], + kv_data_ptrs=[10000, 100000, 200000], + num_draft_entries=2, + engine_rank=src_rank + 2 * src_tp, + ), + attn_tp_size=src_tp, + max_transfer_batch_indices=batch_size, + enable_custom_mem_pool=custom_pool, + _transfer_data=transfer, + _await_transfer_futures=lambda futures: max( + f.result() for f in futures + ), + ) + with concurrent.futures.ThreadPoolExecutor() as executor: + result = MooncakeKVManager.send_kvcache_dcp( + manager, + "session", + src_pages, + [1000000, 2000000, 3000000, 4000000], + np.array([2], dtype=np.int32), + dcp_token_item_lens=[8, src_width, src_width], + dst_dcp_size=dst_tp, + dst_dcp_rank=dst_rank, + src_page_offset=0, + decode_prefix_len=0, + num_kv_tokens=tokens, + executor=executor, + dst_layer_ids=[3, 47, 93, 93], + dst_kv_item_lens=[ + page_size * 8, + page_size * 8, + page_size * dst_tp * dst_width, + page_size * dst_tp * dst_width, + ], + dst_tp_rank=dst_rank, + dst_attn_tp_size=dst_tp, + ) + if fail_draft: + self.assertEqual(result, 17) + self.assertEqual(len(failed_batches), 1) + return + self.assertEqual(result, 0) + dst_head_start = (dst_rank // max(1, dst_tp // heads)) * max( + 1, heads // dst_tp + ) + for base in (3000000, 4000000): + actual = dst_buffers[base].reshape(-1, dst_width)[ + 2 * page_size * dst_tp + logical + ] + np.testing.assert_array_equal( + actual, + expected[ + :, + dst_head_start : dst_head_start + max(1, heads // dst_tp), + ].reshape(tokens, dst_width), + ) + owned = np.arange(dst_rank, tokens, dst_tp) + actual_target = dst_buffers[2000000].reshape(-1, 8)[ + 2 * page_size + owned // dst_tp + ] + np.testing.assert_array_equal( + actual_target, + np.arange(tokens * 8).reshape(tokens, 8).astype(np.uint8)[owned], + ) + self.assertFalse(dst_buffers[1000000].any()) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_auxiliary_output.py b/test/registered/unit/managers/test_auxiliary_output.py index 0799b41dc..4265ea61e 100644 --- a/test/registered/unit/managers/test_auxiliary_output.py +++ b/test/registered/unit/managers/test_auxiliary_output.py @@ -660,6 +660,7 @@ def test_pipeline_parallel_auxiliary_output_round_trip(): next_token_ids=torch.tensor([7]), ) batch = SimpleNamespace( + spec_algorithm=SpeculativeAlgorithm.NONE, return_logprob=False, req_pool_indices=torch.tensor([3]), input_ids=torch.tensor([5]), @@ -705,6 +706,7 @@ def test_pipeline_parallel_dsa_seed_round_trip(dsa_topk_indices): next_draft_input=draft_input, ) batch = SimpleNamespace( + spec_algorithm=SpeculativeAlgorithm.EAGLE3, return_logprob=False, req_pool_indices=torch.tensor([3]), input_ids=torch.tensor([5]), @@ -745,6 +747,7 @@ def test_pipeline_parallel_auxiliary_output_stays_packed_before_first_rank(): next_token_ids=torch.tensor([7]), ) batch = SimpleNamespace( + spec_algorithm=SpeculativeAlgorithm.NONE, return_logprob=False, req_pool_indices=torch.tensor([3]), input_ids=torch.tensor([5]), diff --git a/test/registered/unit/managers/test_pp_cp_rank_offsets.py b/test/registered/unit/managers/test_pp_cp_rank_offsets.py index ff6cfb8d1..06ac85e50 100644 --- a/test/registered/unit/managers/test_pp_cp_rank_offsets.py +++ b/test/registered/unit/managers/test_pp_cp_rank_offsets.py @@ -2,8 +2,11 @@ import unittest from types import SimpleNamespace from unittest.mock import patch +import torch + from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import ( + CustomTestCase, enter_scope, maybe_stub_sgl_kernel, published_topology, @@ -239,5 +242,56 @@ class TestPPCPRankOffsets(unittest.TestCase): ) +class TestDSparkPPOutput(CustomTestCase): + def test_output_ring_rebinds_dspark_state_on_each_stage(self): + from sglang.srt.model_executor.forward_batch_info import PPProxyTensors + from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 + from sglang.srt.speculative.dspark_components.dspark_draft import ( + make_next_draft_input, + ) + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + payloads = [] + scheduler = SimpleNamespace( + _pp_spec_relay=False, + pp_group=SimpleNamespace(is_first_rank=False), + future_map=SimpleNamespace( + stash=lambda indices, value: payloads.append(value) + ), + ) + tokens = torch.tensor([13, 29]) + batch = SimpleNamespace( + return_logprob=False, + req_pool_indices=torch.tensor([0, 1]), + seq_lens=torch.tensor([8, 15]), + spec_algorithm=SpeculativeAlgorithm.DSPARK, + spec_info=object(), + ) + wire = SchedulerPPMixin._pp_prepare_tensor_dict( + scheduler, + SimpleNamespace( + next_token_ids=tokens, + next_draft_input=make_next_draft_input( + bonus_tokens=tokens, new_seq_lens=batch.seq_lens + ), + logits_output=None, + ), + batch, + ) + self.assertNotIn("draft_topk_p", wire) + result = SchedulerPPMixin._pp_prep_batch_result( + scheduler, + batch, + SimpleNamespace(can_run_cuda_graph=False), + PPProxyTensors(wire), + ) + self.assertIsInstance(result.next_draft_input, DFlashDraftInputV2) + self.assertIs(batch.spec_info, result.next_draft_input) + torch.testing.assert_close(batch.spec_info.bonus_tokens, tokens) + torch.testing.assert_close(batch.spec_info.new_seq_lens, batch.seq_lens) + torch.testing.assert_close(payloads[0].bonus_tokens, tokens) + self.assertEqual(payloads[0].hidden_states.numel(), 0) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py index 9e3764810..3d4ef6b89 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py +++ b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py @@ -11,6 +11,9 @@ from sglang.srt.model_executor.cuda_graph_buffer_registry import ( build_prefill_registry, ) from sglang.srt.model_executor.forward_batch_info import PPProxyTensors +from sglang.srt.model_executor.model_runner_components.misc_utils import ( + resolve_pp_proxy_dspark_hidden_size, +) from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import ( PrefillCudaGraphRunner, _build_layer_model_forward_kwargs, @@ -52,6 +55,25 @@ def _make_pp_buffers_and_registry(): class TestPrefillCudaGraphRunnerHelpers(CustomTestCase): + def test_dspark_proxy_width_requires_receiving_stage_and_model_support(self): + class Model: + def get_pp_proxy_dspark_hidden_size(self): + return 16 + + for model, pp_size, pp_rank, expected in ( + (Model(), 1, 0, 0), + (Model(), 2, 0, 0), + (Model(), 2, 1, 16), + (object(), 2, 1, 0), + ): + with self.subTest(pp_size=pp_size, pp_rank=pp_rank, expected=expected): + self.assertEqual( + resolve_pp_proxy_dspark_hidden_size( + model=model, pp_size=pp_size, pp_rank=pp_rank + ), + expected, + ) + def test_pp_proxy_stable_buffers_accept_full_and_hidden_only_contracts(self): buffers, registry = _make_pp_buffers_and_registry() full_proxy = PPProxyTensors( @@ -170,6 +192,7 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase): pp_size=2, is_first_pp_rank=False, pp_proxy_residual_num_blocks=3, + pp_proxy_dspark_hidden_size=16, ) self.assertEqual( @@ -177,9 +200,151 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase): key: tuple(value.shape) for key, value in buffers.pp_proxy_tensors.items() }, - {"hidden_states": (16, 8), "residual": (16, 3, 8)}, + { + "hidden_states": (16, 8), + "residual": (16, 3, 8), + "dspark_hidden_states": (16, 16), + }, ) + def test_dspark_proxy_width_respects_deferred_k3_boundary_capture(self): + from sglang.srt.models.kimi_k3 import ( + KimiK3ForConditionalGeneration, + KimiK3LinearForCausalLM, + ) + from sglang.srt.models.kimi_linear import KimiLinearForCausalLM + + for start, k3_count, linear_count in [ + (0, 0, 0), + (8, 0, 1), + (24, 1, 2), + (52, 2, 3), + ]: + model = SimpleNamespace( + config=SimpleNamespace(hidden_size=8), + model=SimpleNamespace( + start_layer=start, dspark_layers_to_capture=[7, 23, 51] + ), + ) + self.assertEqual( + KimiK3LinearForCausalLM.get_pp_proxy_dspark_hidden_size(model), + k3_count * 8, + ) + self.assertEqual( + KimiLinearForCausalLM.get_pp_proxy_dspark_hidden_size(model), + linear_count * 8, + ) + model.get_pp_proxy_dspark_hidden_size = lambda: ( + KimiK3LinearForCausalLM.get_pp_proxy_dspark_hidden_size(model) + ) + self.assertEqual( + KimiK3ForConditionalGeneration.get_pp_proxy_dspark_hidden_size( + SimpleNamespace(language_model=model) + ), + k3_count * 8, + ) + + def test_body_replay_uses_plural_embeds_and_skips_embedding_on_later_pp_stage(self): + for first_rank in (True, False): + for positional in (True, False): + with self.subTest(first_rank=first_rank, positional=positional): + runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner) + runner._is_full_backend = False + runner._input_embeds_arg_idx = 3 + backing = torch.zeros(4, 8) + supplied = torch.full((4, 8), 7.0) if first_rank else None + runner.buffer_registry = SimpleNamespace( + has_slot=lambda name: name == "input_embeds", + get_slot=lambda _: SimpleNamespace( + slice_for=lambda *args: backing + ), + ) + layer = SimpleNamespace(forward=lambda *args, **kwargs: None) + original = layer.forward + runner.layer_model = layer + sentinel = object() + runner.backend = SimpleNamespace( + replay=lambda *args, **kwargs: sentinel + ) + runner._prefill_forward_context = lambda *args, **kwargs: ( + nullcontext() + ) + + def outer_forward(ids, positions, batch): + if positional: + return layer.forward(None, positions, batch, supplied) + return layer.forward( + None, positions, batch, inputs_embeds=supplied + ) + + runner.model_runner = SimpleNamespace( + pp_group=SimpleNamespace(is_first_rank=first_rank), + model=SimpleNamespace(forward=outer_forward), + ) + batch = SimpleNamespace( + input_ids=torch.arange(4), + positions=torch.arange(4), + mm_input_embeds=None, + ) + result = runner._execute_body_capture(batch, batch, 4, 4, None) + self.assertIs(result, sentinel) + self.assertIs(layer.forward, original) + if first_rank: + torch.testing.assert_close(backing, supplied) + else: + self.assertEqual(torch.count_nonzero(backing).item(), 0) + + def test_dspark_proxy_replay_updates_features_and_clears_padding(self): + buffers = PrefillInputBuffers.create( + device=torch.device("cpu"), + max_bs=1, + max_num_tokens=8, + cache_loc_dtype=torch.int64, + is_multimodal=False, + hidden_size=4, + dtype=torch.float32, + enable_mamba_track=False, + pp_size=2, + pp_proxy_dspark_hidden_size=8, + ) + registry = build_prefill_registry( + device=torch.device("cpu"), + max_bs=1, + max_num_token=8, + cache_loc_dtype=torch.int64, + share_pool=False, + source=buffers, + ) + runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner) + runner.buffers = buffers + runner.model_runner = SimpleNamespace( + pp_group=SimpleNamespace(is_first_rank=False) + ) + captured = runner._capture_pp_proxy_tensors(8)["dspark_hidden_states"] + ptr = captured.data_ptr() + for count, value in [(7, 2.0), (3, 5.0)]: + tokens = torch.arange(count) + proxy = PPProxyTensors( + { + "hidden_states": torch.zeros(count, 4), + "residual": torch.zeros(count, 4), + "dspark_hidden_states": torch.full((count, 8), value), + } + ) + registry.fill_from( + SimpleNamespace( + input_ids=tokens, positions=tokens, out_cache_loc=tokens + ), + raw_bs=1, + padded_bs=1, + raw_num_tokens=count, + padded_num_tokens=8, + pp_proxy_tensors=proxy, + ) + self.assertEqual(captured.data_ptr(), ptr) + torch.testing.assert_close(captured[:count], proxy["dspark_hidden_states"]) + self.assertEqual(torch.count_nonzero(captured[count:]).item(), 0) + def test_pipeline_proxy_output_is_supported(self): runner = PrefillCudaGraphRunner.__new__(PrefillCudaGraphRunner) runner.raw_num_tokens = 3 diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index e82fbddc0..08a67e2e4 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -2501,6 +2501,28 @@ class TestPipelineParallelCompat(CustomTestCase): def test_no_speculative_decoding_is_fine(self): check_pipeline_parallel_compat(self._cfg()) + def test_dspark_pd_prefill_does_not_require_eagle_architecture(self): + check_pipeline_parallel_compat(self._cfg(speculative_algorithm="DSPARK")) + + def test_dspark_is_rejected_outside_pd_prefill(self): + for mode in ("decode", "null"): + with self.subTest(mode=mode): + with self.assertRaisesRegex(AssertionError, "DSPARK.*prefill"): + check_pipeline_parallel_compat( + self._cfg( + speculative_algorithm="DSPARK", disaggregation_mode=mode + ) + ) + + def test_dspark_rejects_eagle_pp_relay(self): + with patch.object( + validation_hook.envs.SGLANG_ENABLE_PP_SPEC, "get", return_value=True + ): + with self.assertRaisesRegex(AssertionError, "SGLANG_ENABLE_PP_SPEC"): + check_pipeline_parallel_compat( + self._cfg(speculative_algorithm="DSPARK") + ) + def test_eagle_is_allowed_on_prefill(self): check_pipeline_parallel_compat( self._cfg(speculative_algorithm="EAGLE"), diff --git a/test/registered/unit/spec/test_dspark_target_hidden_projection.py b/test/registered/unit/spec/test_dspark_target_hidden_projection.py index ae00228ba..ddb55aee4 100644 --- a/test/registered/unit/spec/test_dspark_target_hidden_projection.py +++ b/test/registered/unit/spec/test_dspark_target_hidden_projection.py @@ -9,6 +9,7 @@ from sglang.srt.models.dspark import DSparkDraftMixin from sglang.srt.speculative.dspark_components.dspark_kv_inject import ( TargetHiddenKvInjector, ) +from sglang.srt.speculative.dspark_components.dspark_worker_v2 import DSparkWorkerV2 from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -37,6 +38,40 @@ class _Attention: class DSparkTargetHiddenProjectionTest(CustomTestCase): + def test_nonfinal_prefill_stage_only_forwards_target_proxies(self) -> None: + proxy = object() + target = SimpleNamespace( + model_runner=SimpleNamespace(attn_backend=object(), spec_algorithm=None), + device="cpu", + forward_batch_generation=lambda batch, *, pp_proxy_tensors, capture_hidden_mode: ( + SimpleNamespace(pp_hidden_states_proxy_tensors=pp_proxy_tensors) + ), + ) + with ( + mock.patch( + "sglang.srt.speculative.dspark_components.dspark_worker_v2.get_pp_group", + return_value=SimpleNamespace(is_last_rank=False), + ), + mock.patch( + "sglang.srt.speculative.dspark_components.dspark_worker_v2.get_schedule", + return_value=SimpleNamespace(page_size=1), + ), + ): + worker = DSparkWorkerV2(None, 0, None, 0, target) + worker.alloc_memory_pool() + worker.init_attention_backends() + worker.init_cuda_graphs() + batch = SimpleNamespace(seq_lens=torch.tensor([8])) + result = worker.forward_batch_generation(batch, pp_proxy_tensors=proxy) + self.assertIs(result.pp_hidden_states_proxy_tensors, proxy) + self.assertIs(result.new_seq_lens, batch.seq_lens) + self.assertIsNone(worker.get_confidence_budget_prepare()) + self.assertIsNone(worker.primary_draft_kv_pool) + self.assertEqual(worker.preloaded_weights_bytes, 0) + self.assertEqual( + worker.spec_v2_attn_backends, (target.model_runner.attn_backend,) + ) + def test_single_aux_hidden_state_is_returned_without_copy(self) -> None: hidden_states = torch.empty(2, 3)