diff --git a/python/sglang/srt/speculative/dflash_disaggregation.py b/python/sglang/srt/speculative/dflash_disaggregation.py new file mode 100644 index 000000000000..a38b40faceca --- /dev/null +++ b/python/sglang/srt/speculative/dflash_disaggregation.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.managers.overlap_utils import RelayPayload +from sglang.srt.speculative.draft_worker_common import make_draft_input_v2 + +if TYPE_CHECKING: + from sglang.srt.managers.overlap_utils import FutureMap + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.server_args import ServerArgs + from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 + + +def build_dflash_disagg_draft_input( + batch: ScheduleBatch, + server_args: ServerArgs, + last_tokens_tensor: torch.Tensor, + future_map: FutureMap, +) -> "DFlashDraftInputV2": + spec_info = make_draft_input_v2( + bonus_tokens=last_tokens_tensor, + new_seq_lens=batch.seq_lens, + ) + + if batch.enable_overlap: + spec_info.future_indices = batch.req_pool_indices + # Seed the relay buf with the known seq_lens; publish's chained record + # keeps the in-flight forward's fence intact (see FutureMap.publish). + future_map.publish(spec_info.future_indices, batch.seq_lens) + future_map.stash( + spec_info.future_indices, RelayPayload.from_draft_input(spec_info) + ) + + return spec_info diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index c4cb2e013ffb..34b1db903250 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -415,6 +415,7 @@ def _fill_dp_moe_sync_metadata( batch.global_num_tokens_for_logprob, ) device = self.draft_model_runner.device + forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens forward_batch.global_num_tokens_cpu = gnt forward_batch.global_num_tokens_for_logprob_cpu = gnt_logprob forward_batch.global_num_tokens_gpu = torch.tensor(gnt, dtype=torch.int64).to( diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index c9d27b76079d..b3222ad387c7 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -176,6 +176,14 @@ def build_disagg_draft_input( return build_eagle_disagg_draft_input( batch, server_args, last_tokens_tensor, future_map ) + if self.is_dflash_family(): + from sglang.srt.speculative.dflash_disaggregation import ( + build_dflash_disagg_draft_input, + ) + + return build_dflash_disagg_draft_input( + batch, server_args, last_tokens_tensor, future_map + ) return None def need_topk(self) -> bool: