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

[Spec][1/N] Decoupled speculative decoding: IPC protocol + cross-process request id + server flags (#27634)

Browse Source
This commit is contained in:
Jackey Hua
2026-06-23 17:04:45 -07:00
committed by GitHub
parent 34dd9c28ca
commit f444b5897b
4 changed files with 804 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/server_args.py
+54 -1
View File
@@ -59,6 +59,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
)
from sglang.srt.parser.reasoning_parser import ReasoningParser
from sglang.srt.platforms import current_platform
from sglang.srt.speculative.decoupled_spec_io import DecoupledSpecIpcConfig
from sglang.srt.utils.common import (
LORA_TARGET_ALL_MODULES,
SUPPORTED_LORA_TARGET_MODULES,
@@ -1558,7 +1559,36 @@ class ServerArgs:
"Path to a JSON config file for adaptive speculative decoding tuning knobs.",
] = None
# -------------------------------------------------------------------------
# Decoupled speculative decoding: draft and verify run as
# separate engines, currently connected by a ZMQ IPC mesh.
decoupled_spec_bind_endpoint: A[
Optional[str],
"ZMQ endpoint this engine binds for its inbound channel in decoupled "
"speculative decoding (verifier: result PULL; drafter: control PULL).",
] = None
decoupled_spec_connect_endpoints: A[
Optional[List[str]],
Arg(
help="Peer inbound (bind) endpoints to connect to, ordered by peer "
"rank, for decoupled speculative decoding.",
type_parser=json_list_type,
),
] = None
decoupled_spec_rank: A[
Optional[int],
"This engine's rank within its own role space (verifier-rank or "
"drafter-rank) for decoupled speculative decoding.",
] = None
decoupled_spec_role: A[
Literal["null", "verifier", "drafter"],
"Role in decoupled speculative decoding: 'null' disables it, 'verifier' "
"runs the target/verify half, 'drafter' runs the draft half.",
] = "null"
spec_trace_dir: A[
Optional[str],
"Directory to write decoupled speculative decoding trace files.",
] = None
# Speculative decoding (ngram)
# -------------------------------------------------------------------------
speculative_ngram_min_bfs_breadth: A[
@@ -7297,6 +7327,9 @@ class PortArgs:
# The ipc filename for MultiTokenizerRouter to receive inputs from TokenizerWorker processes (zmq)
tokenizer_worker_ipc_name: Optional[str]
# The ipc endpoints between verifier scheduler and drafter scheduler
decoupled_spec_ipc_config: Optional[DecoupledSpecIpcConfig]
# zmq address for load snapshot PUSH/PULL (dp-attention TCP mode only;
# empty when IPC mode derives the address from instance_id).
load_collector_ipc_name: str = ""
@@ -7325,6 +7358,24 @@ class PortArgs:
instance_id = uuid.uuid4().hex[:12]
decoupled_spec_ipc_config = None
if server_args.decoupled_spec_role != "null":
if (
server_args.decoupled_spec_bind_endpoint is None
or server_args.decoupled_spec_connect_endpoints is None
or server_args.decoupled_spec_rank is None
):
raise ValueError(
"--decoupled-spec-bind-endpoint, "
"--decoupled-spec-connect-endpoints, and "
"--decoupled-spec-rank are required for decoupled speculative decoding."
)
decoupled_spec_ipc_config = DecoupledSpecIpcConfig(
bind_endpoint=server_args.decoupled_spec_bind_endpoint,
connect_endpoints=tuple(server_args.decoupled_spec_connect_endpoints),
rank=int(server_args.decoupled_spec_rank),
)
if not server_args.enable_dp_attention:
# Normal case, use IPC within a single node
return PortArgs(
@@ -7335,6 +7386,7 @@ class PortArgs:
rpc_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}",
metrics_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}",
tokenizer_worker_ipc_name=tokenizer_worker_ipc_name,
decoupled_spec_ipc_config=decoupled_spec_ipc_config,
instance_id=instance_id,
)
else:
@@ -7405,6 +7457,7 @@ class PortArgs:
rpc_ipc_name=NetworkAddress(dist_init_host, rpc_port).to_tcp(),
metrics_ipc_name=NetworkAddress(dist_init_host, metrics_port).to_tcp(),
tokenizer_worker_ipc_name=tokenizer_worker_ipc_name,
decoupled_spec_ipc_config=decoupled_spec_ipc_config,
load_collector_ipc_name=NetworkAddress(
dist_init_host, load_collector_port
).to_tcp(),
python/sglang/srt/speculative/decoupled_spec_io.py
+384
View File
@@ -0,0 +1,384 @@
from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Optional
class DraftMeshMessageType(str, Enum):
CONTROL_BATCH = "control_batch"
TAIL_STREAM_OUTPUT_BATCH = "tail_stream_output_batch"
@dataclass(frozen=True)
class DraftReqKey:
"""Request identity on the drafter side.
The original request_id is only unique within the verifier that owns it.
src_verifier_rank keeps the drafter-side request table unambiguous when
multiple verifier ranks send work to the same drafter rank.
"""
src_verifier_rank: int
request_id: str
def build_draft_scheduler_rid(draft_key: DraftReqKey) -> str:
return f"draft:{int(draft_key.src_verifier_rank)}:{draft_key.request_id}"
def parse_draft_scheduler_rid(rid: str) -> DraftReqKey:
if rid.startswith("draft:"):
encoded = rid[len("draft:") :]
rank_text, sep, request_id = encoded.partition(":")
if sep and request_id:
return DraftReqKey(
src_verifier_rank=int(rank_text),
request_id=request_id,
)
raise ValueError(f"Invalid decoupled draft scheduler rid: {rid}")
@dataclass
class DraftSync:
"""Open or re-open a drafter request from a verifier-owned prefix.
The verifier is the source of truth for committed tokens. DraftSync gives
the drafter the prompt and already committed output prefix that it must
align to before it can emit draft tail tokens.
"""
request_id: str
src_verifier_rank: int
dst_drafter_rank: int
prompt_token_ids: list[int] = field(default_factory=list)
committed_outputs: list[int] = field(default_factory=list)
@property
def draft_key(self) -> DraftReqKey:
return DraftReqKey(
src_verifier_rank=int(self.src_verifier_rank),
request_id=self.request_id,
)
@dataclass
class VerifyCommit:
"""
Sent from verifier to drafter to commit a portion of the draft outputs.
committed_tokens is the verifier-committed contiguous output segment:
output_ids[
pre_verify_committed_len:
pre_verify_committed_len + len(committed_tokens)
].
Drafter must align its reqs to these committed tokens,
and sometimes needs to truncate tokens / reprefill.
"""
request_id: str
src_verifier_rank: int
dst_drafter_rank: int
pre_verify_committed_len: int
committed_tokens: list[int]
@property
def draft_key(self) -> DraftReqKey:
return DraftReqKey(
src_verifier_rank=int(self.src_verifier_rank),
request_id=self.request_id,
)
def validate_committed_tokens(self) -> None:
if not self.committed_tokens:
raise ValueError(
"VerifyCommit committed_tokens must be non-empty: "
f"request_id={self.request_id} "
f"pre_verify_committed_len={self.pre_verify_committed_len}"
)
if int(self.pre_verify_committed_len) < 0:
raise ValueError(
"VerifyCommit pre_verify_committed_len must be non-negative: "
f"request_id={self.request_id} "
f"pre_verify_committed_len={self.pre_verify_committed_len}"
)
@dataclass
class DraftClose:
request_id: str
src_verifier_rank: int
dst_drafter_rank: int
reason: str
@property
def draft_key(self) -> DraftReqKey:
return DraftReqKey(
src_verifier_rank=int(self.src_verifier_rank),
request_id=self.request_id,
)
@dataclass
class DraftTailStreamOutput:
"""
Drafter sends one output token to the verifier-side DraftTailBuffer.
base_committed_len records the verifier prefix length that the drafter used
as the base when this token was emitted. The verifier compares it with its
stale-base boundary before accepting the token as tail data or as
pending-prefix confirmation.
new_token_pos is the 0-based output token position for new_token. Normal
decode streams send the latest generated token.
"""
src_drafter_rank: int
dst_verifier_rank: int
request_id: str
base_committed_len: int
new_token_pos: int
new_token: int
@dataclass
class DraftTailStreamOutputBatch:
outputs: list[DraftTailStreamOutput] = field(default_factory=list)
@dataclass
class DraftControlBatch:
dst_drafter_rank: int
sync_messages: list[DraftSync] = field(default_factory=list)
verify_commit_messages: list[VerifyCommit] = field(default_factory=list)
close_messages: list[DraftClose] = field(default_factory=list)
@dataclass
class VerifierCommitSegment:
"""Contiguous VerifyCommit messages coalesced for one drafter request.
When receiving contiguous VerifyCommit messages for the same draft req,
the transport thread(TokenSync thread at drafter side) coalesces them into a single VerifierCommitSegment.
VerifierCommitSegment represents a contiguous verifier-committed token segment for drafter,
and drafter scheduler should align with these segments before emitting tail tokens
"""
draft_key: DraftReqKey
dst_drafter_rank: int
pre_verify_committed_len: int
committed_tokens: list[int] = field(default_factory=list)
@property
def end_committed_len(self) -> int:
return int(self.pre_verify_committed_len) + len(self.committed_tokens)
def append_message(self, message: VerifyCommit) -> None:
"""
It runs on TokenSyncThread under _pending_lock. That loop only
catches zmq.error.ContextTerminated, so a raise here escapes _run and
silently kills the drafter control thread. It then stops applying
ALL requests' controls while the verifier keeps pushing.
TODO: 1. peer-data violations (non-contiguous / invalid len)
should quarantine just that request (drop + add to close_keys), not
crash the thread. 2. phase 5.c will handle the drafter failure by
degrading the verifier into normal autoregressive decoding.
"""
if message.draft_key != self.draft_key:
raise RuntimeError(
"Verifier commit segment received a commit for a different "
f"request: segment_key={self.draft_key} message_key={message.draft_key}"
)
if int(message.dst_drafter_rank) != int(self.dst_drafter_rank):
raise RuntimeError(
"Verifier commit segment received a commit for a different "
"drafter rank: "
f"request_id={message.request_id} "
f"segment_drafter_rank={self.dst_drafter_rank} "
f"message_drafter_rank={message.dst_drafter_rank}"
)
message.validate_committed_tokens()
pre_verify_committed_len = int(message.pre_verify_committed_len)
if pre_verify_committed_len != self.end_committed_len:
raise RuntimeError(
"Verifier commit segment requires contiguous VerifyCommit "
"messages: "
f"request_id={message.request_id} "
f"expected_pre_verify_committed_len={self.end_committed_len} "
f"actual_pre_verify_committed_len={pre_verify_committed_len}"
)
token_ids = [int(token_id) for token_id in message.committed_tokens]
self.committed_tokens.extend(token_ids)
def extract_prefix(self, num_tokens: int) -> VerifierCommitSegment:
num_tokens = int(num_tokens)
if num_tokens <= 0:
raise ValueError(
"Verifier commit segment prefix length must be positive: "
f"request_id={self.draft_key.request_id} num_tokens={num_tokens}"
)
if num_tokens > len(self.committed_tokens):
raise ValueError(
"Verifier commit segment prefix length exceeds segment length: "
f"request_id={self.draft_key.request_id} "
f"num_tokens={num_tokens} "
f"segment_len={len(self.committed_tokens)}"
)
prefix_tokens = [
int(token_id) for token_id in self.committed_tokens[:num_tokens]
]
remaining_tokens = [
int(token_id) for token_id in self.committed_tokens[num_tokens:]
]
prefix_segment = VerifierCommitSegment(
draft_key=self.draft_key,
dst_drafter_rank=int(self.dst_drafter_rank),
pre_verify_committed_len=int(self.pre_verify_committed_len),
committed_tokens=prefix_tokens,
)
self.pre_verify_committed_len = int(self.pre_verify_committed_len) + num_tokens
self.committed_tokens = remaining_tokens
return prefix_segment
@dataclass
class DraftControlInbox:
"""Drafter-side inbox for verifier control messages.
The TokenSync thread temporarily stores incoming control messages here.
The drafter scheduler extracts and consumes them each time it finishes a decoding step.
"""
sync_messages: list[DraftSync] = field(default_factory=list)
verifier_commit_segments: dict[DraftReqKey, VerifierCommitSegment] = field(
default_factory=dict
)
close_keys: set[DraftReqKey] = field(default_factory=set)
def is_empty(self) -> bool:
return (
not self.sync_messages
and not self.verifier_commit_segments
and not self.close_keys
)
def pending_control_count(self) -> int:
return (
len(self.sync_messages)
+ len(self.verifier_commit_segments)
+ len(self.close_keys)
)
def add_control_batch_locked(self, batch: DraftControlBatch) -> None:
for message in batch.close_messages:
self.add_close_key_locked(message.draft_key)
for message in batch.sync_messages:
if message.draft_key not in self.close_keys:
self.sync_messages.append(message)
for message in batch.verify_commit_messages:
self.add_verify_commit_locked(message)
def add_close_key_locked(self, draft_key: DraftReqKey) -> None:
self.close_keys.add(draft_key)
self.verifier_commit_segments.pop(draft_key, None)
self.sync_messages = [
message for message in self.sync_messages if message.draft_key != draft_key
]
def add_verify_commit_locked(self, message: VerifyCommit) -> None:
if message.draft_key in self.close_keys:
return
segment = self.verifier_commit_segments.get(message.draft_key)
if segment is None:
segment = VerifierCommitSegment(
draft_key=message.draft_key,
dst_drafter_rank=int(message.dst_drafter_rank),
pre_verify_committed_len=int(message.pre_verify_committed_len),
)
segment.append_message(message)
self.verifier_commit_segments[message.draft_key] = segment
return
segment.append_message(message)
def extract_ready_controls_locked(
self,
consumable_commit_len: Callable[[VerifierCommitSegment], int],
) -> ReadyDraftControls:
ready_controls = ReadyDraftControls()
if self.close_keys:
ready_controls.close_keys = self.close_keys
self.close_keys = set()
if self.sync_messages:
ready_controls.sync_messages = self.sync_messages
self.sync_messages = []
for draft_key, segment in list(self.verifier_commit_segments.items()):
consumable_len = consumable_commit_len(segment)
if consumable_len <= 0:
continue
ready_controls.ready_commit_segments.append(
segment.extract_prefix(consumable_len)
)
if not segment.committed_tokens:
self.verifier_commit_segments.pop(draft_key, None)
return ready_controls
@dataclass
class ReadyDraftControls:
sync_messages: list[DraftSync] = field(default_factory=list)
close_keys: set[DraftReqKey] = field(default_factory=set)
ready_commit_segments: list[VerifierCommitSegment] = field(default_factory=list)
def is_empty(self) -> bool:
return (
not self.sync_messages
and not self.close_keys
and not self.ready_commit_segments
)
def extracted_control_count(self) -> int:
return (
len(self.sync_messages)
+ len(self.close_keys)
+ len(self.ready_commit_segments)
)
@dataclass
class DraftMeshMessage:
message_type: DraftMeshMessageType
control_batch: Optional[DraftControlBatch] = None
tail_stream_output_batch: Optional[DraftTailStreamOutputBatch] = None
@staticmethod
def from_control_batch(message: DraftControlBatch) -> DraftMeshMessage:
return DraftMeshMessage(
message_type=DraftMeshMessageType.CONTROL_BATCH,
control_batch=message,
)
@staticmethod
def from_tail_stream_output_batch(
message: DraftTailStreamOutputBatch,
) -> DraftMeshMessage:
return DraftMeshMessage(
message_type=DraftMeshMessageType.TAIL_STREAM_OUTPUT_BATCH,
tail_stream_output_batch=message,
)
@dataclass(frozen=True)
class DecoupledSpecIpcConfig:
bind_endpoint: str
connect_endpoints: tuple[str, ...]
rank: int