minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[kimi k3][pd disagg] support pp prefill + dcp decode with dspark (#40045)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-09-20 00:15:28 -07:00
committed by GitHub
parent 22f02cc339
commit f4c256354c
21 changed files with 724 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/speculative_hook.py
+5 -3
View File
@@ -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):
python/sglang/srt/arg_groups/validation_hook.py
+8 -1
View File
@@ -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
python/sglang/srt/disaggregation/mooncake/conn.py
+107 -9
View File
@@ -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:
python/sglang/srt/managers/scheduler_pp_mixin.py
+15 -1
View File
@@ -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
python/sglang/srt/model_executor/model_runner.py
+7
View File
@@ -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,
python/sglang/srt/model_executor/model_runner_components/misc_utils.py
+16 -1
View File
@@ -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
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+3
View File
@@ -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
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+16 -3
View File
@@ -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
python/sglang/srt/model_executor/runner_utils/buffers.py
+9
View File
@@ -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
python/sglang/srt/models/kimi_k3.py
+30 -9
View File
@@ -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(
python/sglang/srt/models/kimi_linear.py
+18 -10
View File
@@ -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."
python/sglang/srt/speculative/draft_worker_common.py
+2
View File
@@ -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,
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+49 -7
View File
@@ -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()