chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

feat(SpecEagleV2): add standalone_worker_v2 (#12625)

Browse Source 
Co-authored-by: liugaoji.lgj <liugaoji.lgj@alibaba-inc.com>
This commit is contained in:
Gaoji Liu
2025-12-30 17:55:04 +08:00
committed by GitHub
parent b6871ba7c9
commit 7518dc3532
14 changed files with 317 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/layers/attention/triton_backend.py
View File

@@ -602,7 +602,7 @@ class TritonAttnBackend(AttentionBackend):
num_kv_splits = None
attn_logits = None
attn_lse = None
elif forward_mode.is_draft_extend():
elif forward_mode.is_draft_extend(include_v2=True):
num_tokens_per_bs = self.speculative_num_steps + 1
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[: bs + 1] = torch.arange(
@@ -749,7 +749,7 @@ class TritonAttnBackend(AttentionBackend):
seq_mask_len = self.num_draft_tokens * (seq_lens + self.num_draft_tokens)
mask_indptr = self.mask_indptr[: bs + 1]
mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
elif forward_mode.is_draft_extend():
elif forward_mode.is_draft_extend(include_v2=True):
seq_lens = seq_lens[:bs]
accept_lens = spec_info.accept_length[:bs]
qo_indptr = self.qo_indptr[: bs + 1]

14
python/sglang/srt/managers/overlap_utils.py
View File

@@ -116,7 +116,9 @@ class FutureMap:
return FutureIndices(indices=indices, interval=slice(start, end))
def resolve_future(self, model_worker_batch: ModelWorkerBatch):
if self.spec_algo.is_eagle():
if self.spec_algo.is_none():
_resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf)
else:
# TODO(lsyin): write future indices into spec_info.future_indices
draft_input: EagleDraftInput = model_worker_batch.spec_info
if draft_input is None:
@@ -129,8 +131,6 @@ class FutureMap:
draft_input.new_seq_lens = self.new_seq_lens_buf[indices]
if spec_need_hidden_states():
draft_input.hidden_states = self.hidden_states_buf[indices]
else:
_resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf)
def is_empty_slice(self, s: slice) -> bool:
start, stop, step = s.indices(self.future_buffer_len)
@@ -142,12 +142,12 @@ class FutureMap:
def store_to_map(
self, future_indices: FutureIndices, batch_result: GenerationBatchResult
):
if self.spec_algo.is_eagle():
draft_input: EagleDraftInput = batch_result.next_draft_input
self.store_to_map_for_new_batch(future_indices, draft_input)
else:
if self.spec_algo.is_none():
intv = future_indices.interval
self.token_ids_buf[intv] = batch_result.next_token_ids
else:
draft_input: EagleDraftInput = batch_result.next_draft_input
self.store_to_map_for_new_batch(future_indices, draft_input)
def store_to_map_for_new_batch(
self, future_indices: FutureIndices, draft_input: EagleDraftInput

4
python/sglang/srt/managers/schedule_batch.py
View File

@@ -1849,7 +1849,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
@property
def is_spec_v2(self):
# FIXME: finally deprecate is_spec_v2
return self.enable_overlap and self.spec_algorithm.is_eagle()
ret = self.enable_overlap and not self.spec_algorithm.is_none()
assert not ret or self.spec_algorithm.supports_spec_v2()
return ret
def prepare_for_decode(self):
self.forward_mode = ForwardMode.DECODE

6
python/sglang/srt/managers/scheduler.py
View File

@@ -499,7 +499,7 @@ class Scheduler(
# Draft workers are looked up via `SpeculativeAlgorithm` registry; new
# algorithms should register their factory instead of patching this code.
if self.spec_algorithm.is_eagle():
if self.spec_algorithm.supports_spec_v2():
draft_worker_kwargs["enable_overlap"] = self.enable_overlap
# FIXME: refactor the draft worker registration logic
@@ -852,7 +852,7 @@ class Scheduler(
if self.draft_worker is None or self.spec_algorithm.is_ngram():
draft_token_to_kv_pool = None
elif self.spec_algorithm.is_eagle() and self.enable_overlap:
elif self.spec_algorithm.supports_spec_v2() and self.enable_overlap:
if self.server_args.enable_multi_layer_eagle:
draft_runner = self.draft_worker.draft_worker.draft_runner_list[0]
else:
@@ -930,11 +930,13 @@ class Scheduler(
hidden_size=(
model_config.hidden_size
if self.spec_algorithm.is_eagle()
or self.spec_algorithm.is_standalone()
else 16 # minimal padding size for RDMA
),
hidden_states_dtype=(
model_config.dtype
if self.spec_algorithm.is_eagle()
or self.spec_algorithm.is_standalone()
else torch.float32
),
custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),

5
python/sglang/srt/model_executor/cpu_graph_runner.py
View File

@@ -606,7 +606,10 @@ class CPUGraphRunner:
def get_spec_info(self, num_tokens: int):
spec_info = None
if self.model_runner.spec_algorithm.is_eagle():
if (
self.model_runner.spec_algorithm.is_eagle()
or self.model_runner.spec_algorithm.is_standalone()
):
from sglang.srt.speculative.eagle_info import EagleVerifyInput
if self.model_runner.is_draft_worker:

2
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -375,6 +375,7 @@ class CudaGraphRunner:
cuda_graph_bs = (
max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
if self.model_runner.spec_algorithm.is_eagle()
or self.model_runner.spec_algorithm.is_standalone()
else max(forward_batch.global_num_tokens_cpu)
)
else:
@@ -777,6 +778,7 @@ class CudaGraphRunner:
max_batch_size = (
max_num_tokens / self.num_tokens_per_bs
if self.model_runner.spec_algorithm.is_eagle()
or self.model_runner.spec_algorithm.is_standalone()
else max_num_tokens
)
index = bisect.bisect_left(self.capture_bs, max_batch_size)

2
python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py
View File

@@ -303,7 +303,7 @@ class ModelRunnerKVCacheMixin:
max_num_reqs, self.server_args.max_mamba_cache_size // ratio
)
if self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone():
if not self.spec_algorithm.is_none():
if self.is_draft_worker:
self.max_total_num_tokens = self.server_args.draft_runner_cache_size
max_num_reqs = self.server_args.max_num_reqs

2
python/sglang/srt/server_args.py
View File

@@ -1998,7 +1998,7 @@ class ServerArgs:
)
if (
self.speculative_algorithm in ["EAGLE", "EAGLE3"]
self.speculative_algorithm in ["EAGLE", "EAGLE3", "STANDALONE"]
and envs.SGLANG_ENABLE_SPEC_V2.get()
):
self.disable_overlap_schedule = False

2
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
View File

@@ -141,6 +141,7 @@ class EAGLEDraftCudaGraphRunner:
cuda_graph_bs = (
max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
if self.model_runner.spec_algorithm.is_eagle()
or self.model_runner.spec_algorithm.is_standalone()
else max(forward_batch.global_num_tokens_cpu)
)
else:
@@ -328,6 +329,7 @@ class EAGLEDraftCudaGraphRunner:
max_batch_size = (
max_num_tokens // self.num_tokens_per_bs
if self.model_runner.spec_algorithm.is_eagle()
or self.model_runner.spec_algorithm.is_standalone()
else max_num_tokens
)
index = bisect.bisect_left(self.capture_bs, max_batch_size)

1
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
View File

@@ -189,6 +189,7 @@ class EAGLEDraftExtendCudaGraphRunner:
cuda_graph_bs = (
max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
if self.model_runner.spec_algorithm.is_eagle()
or self.model_runner.spec_algorithm.is_standalone()
else max(forward_batch.global_num_tokens_cpu)
)
else:

16
python/sglang/srt/speculative/eagle_worker_v2.py
View File

@@ -12,6 +12,7 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_r
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
EAGLEDraftNpuGraphRunner,
)
from sglang.srt.layers.attention.triton_backend import TritonMultiStepDraftBackend
from sglang.srt.layers.moe.utils import (
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
@@ -50,12 +51,14 @@ from sglang.srt.utils.common import (
empty_context,
fast_topk,
get_available_gpu_memory,
is_cuda,
is_npu,
next_power_of_2,
)
from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
_is_npu = is_npu()
_is_cuda = is_cuda()
logger = logging.getLogger(__name__)
@@ -251,8 +254,14 @@ class EagleDraftWorker(BaseDraftWorker):
"cuda": EAGLEDraftExtendCudaGraphRunner,
}
# Capture extend
# FIXME cuda not support draft_extend capture
if self.draft_extend_attn_backend and _is_npu:
# TODO: support draft extend cuda graph for more attention backends
if self.draft_extend_attn_backend and (
_is_npu
or (
_is_cuda
and isinstance(self.draft_attn_backend, TritonMultiStepDraftBackend)
)
):
tic = time.perf_counter()
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
logger.info(
@@ -506,6 +515,9 @@ class EagleDraftWorker(BaseDraftWorker):
self.plan_stream
)
if forward_batch.spec_info.accept_length is None:
forward_batch.spec_info.accept_length = batch_result.accept_lens
# Run draft extend batch in the main compute stream
can_cuda_graph = (
self.cuda_graph_runner_for_draft_extend

9
python/sglang/srt/speculative/spec_info.py
View File

@@ -161,6 +161,9 @@ class SpeculativeAlgorithm(metaclass=_SpeculativeAlgorithmMeta):
def is_none(self) -> bool:
return self is SpeculativeAlgorithm.NONE
def supports_spec_v2(self) -> bool:
return self.is_eagle() or self.is_eagle3() or self.is_standalone()
def is_eagle(self) -> bool:
return self._has_flag("EAGLE")
@@ -274,6 +277,12 @@ def _create_eagle_worker(**kwargs: Any) -> Any:
def _create_standalone_worker(**kwargs: Any) -> Any:
enable_overlap = kwargs.pop("enable_overlap", False)
if enable_overlap:
from sglang.srt.speculative.standalone_worker_v2 import StandaloneWorkerV2
return StandaloneWorkerV2(**kwargs)
from sglang.srt.speculative.standalone_worker import StandaloneWorker
return StandaloneWorker(**kwargs)

169
python/sglang/srt/speculative/standalone_worker_v2.py Normal file
View File

@@ -0,0 +1,169 @@
import contextlib
import logging
from typing import Optional, Tuple
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.eagle_utils import TreeMaskMode
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_utils import draft_tp_context
from sglang.srt.utils import empty_context, get_bool_env_var, is_cuda
if is_cuda():
from sgl_kernel import segment_packbits # noqa: F401
logger = logging.getLogger(__name__)
SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
def _get_plan_stream(
device: str,
) -> Tuple[any, contextlib.AbstractContextManager]:
if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get():
plan_stream = torch.get_device_module(device).Stream()
plan_stream_ctx = torch.get_device_module(device).stream(plan_stream)
return plan_stream, plan_stream_ctx
else:
return None, contextlib.nullcontext()
class StandaloneDraftWorker(EagleDraftWorker):
"""Custom EagleDraftWorker that doesn't share embeddings/lm_head with target model."""
def __init__(
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
# copy args
self.server_args = server_args
self.gpu_id = gpu_id
self.tp_rank = tp_rank
self.dp_rank = dp_rank
self.moe_ep_rank = moe_ep_rank
self.nccl_port = nccl_port
self.target_worker = target_worker
# Args for easy access
self.device = server_args.device
self.topk = server_args.speculative_eagle_topk
self.speculative_num_steps = server_args.speculative_num_steps
self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm
)
# Set constant
from sglang.srt.speculative.eagle_info import EagleDraftInput
EagleDraftInput.ALLOC_LEN_PER_DECODE = max(
self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens
)
# Do not capture cuda graph in `TpModelWorker` init,
# will capture later with init_cuda_graphs()
backup_disable_cuda_graph = server_args.disable_cuda_graph
server_args.disable_cuda_graph = True
# Share the allocator with a target worker.
# Draft and target worker own their own KV cache pools.
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
with empty_context():
# Init draft worker
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
)
# Alias for better readability
self.draft_runner = self.draft_worker.model_runner
self.init_token_map()
self.init_lm_head()
# Init attention backend and cuda graphs
self.draft_runner.server_args.disable_cuda_graph = backup_disable_cuda_graph
self.draft_tp_context = (
draft_tp_context if server_args.enable_dp_attention else empty_context
)
with self.draft_tp_context(
self.draft_runner.tp_group
), speculative_moe_backend_context():
self.init_attention_backend()
self.init_cuda_graphs()
self.tree_mask_mode = TreeMaskMode.FULL_MASK
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
def init_lm_head(self):
"""Override to prevent sharing embeddings and lm_head with target model."""
# For standalone worker, we don't share embeddings and lm_head
# The draft model uses its own embeddings and lm_head
pass
class StandaloneWorkerV2(EAGLEWorkerV2):
def __init__(
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
# Parse arguments
self.server_args = server_args
self.topk = server_args.speculative_eagle_topk
self.speculative_num_steps = server_args.speculative_num_steps
self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
self.enable_nan_detection = server_args.enable_nan_detection
self.gpu_id = gpu_id
self.device = server_args.device
self._target_worker = target_worker
self.page_size = server_args.page_size
self.speculative_algorithm = SpeculativeAlgorithm.from_string(
server_args.speculative_algorithm
)
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
target_worker.get_memory_pool()
)
# Override the context length of the draft model to be the same as the target model.
server_args.context_length = target_worker.model_runner.model_config.context_len
# Create our custom draft worker that doesn't share embeddings/lm_head
self._draft_worker = StandaloneDraftWorker(
server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
)
# Some dummy tensors
self.num_new_pages_per_topk = torch.empty(
(), dtype=torch.int64, device=self.device
)
self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device)
self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)