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

[Spec] Fix Dspark and Dflash state divergence across TP rank (#33614)

Browse Source 
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
Rabinovich
2026-08-29 17:27:43 -07:00
committed by GitHub
co-authored by hnyls2002 Liangsheng Yin
parent a1fe4e30a9
commit f60bc73c58
8 changed files with 356 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+5
View File
@@ -1205,6 +1205,11 @@ class Envs:
# Saves the per-step draft forward, but the draft KV goes stale: an upshift
# back to steps>0 starts from a cold draft state (low accept until it recovers).
SGLANG_SPEC_SKIP_ZERO_STEP_DRAFT_EXTEND = EnvBool(False)
# Which speculative decisions rank 0 broadcasts to its TP group; narrowing
# it under live traffic isolates where ranks actually diverge. Comma
# separated presets ("all", "rng", "init", "off"), or SpecTpSyncSite slugs
# and numbers, each negatable with a leading "-": "all,-dspark-plan,-6".
SGLANG_SPEC_TP_SYNC = EnvStr("all")
# Kill-switch for the draft-extend cuda graph. Draft extend then always runs
# eager. Escape hatch for setups where the capture's memory pool costs more
# than the graph saves (e.g. DeepEP MoE workspace captured at full dispatch
python/sglang/srt/speculative/dflash_worker_v2.py
+108 -68
View File
@@ -62,6 +62,7 @@ from sglang.srt.speculative.draft_worker_common import (
)
from sglang.srt.speculative.dspark_components.dspark_draft import resolve_greedy_mask
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
from sglang.srt.speculative.spec_utils import (
SIMULATE_ACC_LEN,
SIMULATE_ACC_METHOD,
@@ -70,7 +71,7 @@ from sglang.srt.speculative.spec_utils import (
assign_req_to_token_pool_func,
build_grammar_vocab_mask,
)
from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu
from sglang.srt.utils import is_cuda, is_hip, is_npu
_is_npu = is_npu()
@@ -302,6 +303,7 @@ class DFlashWorkerV2(BaseSpecWorker):
self._warned_sampling_fallback = False
self._draft_probs_buf = None
self._logged_first_verify = False
self._tp_sync = SpecTpSync(get_tp_group())
bundle = build_draft_tp_worker(
server_args=server_args,
@@ -465,7 +467,12 @@ class DFlashWorkerV2(BaseSpecWorker):
get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
)
if is_cuda() and capture_decode_cuda_graph:
available_mem = get_available_gpu_memory(self.device, self.gpu_id)
available_mem = self._tp_sync.available_memory_gb(
SpecTpSyncSite.DFLASH_MEM,
self.device,
self.gpu_id,
group=get_tp_group(),
)
if available_mem < 1.0:
capture_decode_cuda_graph = False
logger.warning(
@@ -1619,6 +1626,89 @@ class DFlashWorkerV2(BaseSpecWorker):
self._new_seq_lens_bufs[slot][:bs],
)
def _accept_block(
self,
*,
candidates: torch.Tensor,
next_token_logits: torch.Tensor,
sampling_info,
draft_input,
prefix_lens: torch.Tensor,
bs: int,
):
new_seq_lens = None
target_predict = None
if self._selector_sample is not None:
selector_candidate_ids, selector_q_rows = self._selector_sample
accept_len, bonus = self._selector_sampling_accept(
candidates=candidates,
next_token_logits=next_token_logits,
candidate_ids=selector_candidate_ids,
q_rows=selector_q_rows,
sampling_info=sampling_info,
draft_input=draft_input,
)
self._tp_sync.sync(SpecTpSyncSite.DFLASH_SELECTOR, accept_len)
self._tp_sync.sync(SpecTpSyncSite.DFLASH_SELECTOR, bonus)
out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
elif (
not _is_all_greedy(sampling_info) and is_dflash_sampling_verify_available()
):
accept_len, bonus = compute_dflash_sampling_correct_drafts_and_bonus(
candidates=candidates,
next_token_logits=next_token_logits,
sampling_info=sampling_info,
max_top_k=draft_input.max_top_k,
uniform_top_k_value=draft_input.uniform_top_k_value,
)
self._tp_sync.sync(SpecTpSyncSite.DFLASH_ACCEPT_SAMPLE, accept_len)
self._tp_sync.sync(SpecTpSyncSite.DFLASH_ACCEPT_SAMPLE, bonus)
out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
else:
target_predict = torch.argmax(next_token_logits, dim=-1).view(
bs, int(self.block_size)
)
self._tp_sync.sync(SpecTpSyncSite.DFLASH_ACCEPT_GREEDY, target_predict)
if self._use_triton_accept_bonus:
try:
(
accept_len,
commit_lens,
bonus,
out_tokens,
new_seq_lens,
) = self._next_accept_bonus_buffers(bs)
_compute_dflash_accept_bonus_triton_unchecked(
candidates=candidates,
target_top1=target_predict,
accept_lens_out=accept_len,
commit_lens_out=commit_lens,
bonus_ids_out=bonus,
out_tokens_out=out_tokens,
prefix_lens=prefix_lens,
new_seq_lens_out=new_seq_lens,
)
except Exception as e:
self._use_triton_accept_bonus = False
logger.warning(
"DFLASH Triton accept/bonus failed; falling back to eager path: %s",
e,
)
accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
candidates=candidates,
target_predict=target_predict,
)
out_tokens, commit_lens = _commit_accept(
candidates, accept_len, bonus
)
else:
accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
candidates=candidates,
target_predict=target_predict,
)
out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
return accept_len, commit_lens, bonus, out_tokens, new_seq_lens, target_predict
def _validate_phase1_sampling_support(self, batch: ScheduleBatch) -> None:
sampling_info = batch.sampling_info
# A selector draft carries its own q and verifies through accept_sampling, so
@@ -1678,6 +1768,7 @@ class DFlashWorkerV2(BaseSpecWorker):
batch_output.logits_output,
batch_output.next_token_ids,
)
self._tp_sync.sync(SpecTpSyncSite.DFLASH_TARGET, next_token_ids)
batch_output.new_seq_lens = batch.seq_lens
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
@@ -2029,72 +2120,21 @@ class DFlashWorkerV2(BaseSpecWorker):
grammar_mask.apply(logits_output.next_token_logits)
candidates = draft_tokens
new_seq_lens = None
target_predict = None
if self._selector_sample is not None:
selector_candidate_ids, selector_q_rows = self._selector_sample
accept_len, bonus = self._selector_sampling_accept(
candidates=candidates,
next_token_logits=logits_output.next_token_logits,
candidate_ids=selector_candidate_ids,
q_rows=selector_q_rows,
sampling_info=sampling_info,
draft_input=draft_input,
)
out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
elif (
not _is_all_greedy(sampling_info) and is_dflash_sampling_verify_available()
):
accept_len, bonus = compute_dflash_sampling_correct_drafts_and_bonus(
candidates=candidates,
next_token_logits=logits_output.next_token_logits,
sampling_info=sampling_info,
max_top_k=draft_input.max_top_k,
uniform_top_k_value=draft_input.uniform_top_k_value,
)
out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
else:
target_predict = torch.argmax(logits_output.next_token_logits, dim=-1).view(
bs, int(self.block_size)
)
if self._use_triton_accept_bonus:
try:
(
accept_len,
commit_lens,
bonus,
out_tokens,
new_seq_lens,
) = self._next_accept_bonus_buffers(bs)
_compute_dflash_accept_bonus_triton_unchecked(
candidates=candidates,
target_top1=target_predict,
accept_lens_out=accept_len,
commit_lens_out=commit_lens,
bonus_ids_out=bonus,
out_tokens_out=out_tokens,
prefix_lens=prefix_lens,
new_seq_lens_out=new_seq_lens,
)
except Exception as e:
self._use_triton_accept_bonus = False
logger.warning(
"DFLASH Triton accept/bonus failed; falling back to eager path: %s",
e,
)
accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
candidates=candidates,
target_predict=target_predict,
)
out_tokens, commit_lens = _commit_accept(
candidates, accept_len, bonus
)
else:
accept_len, bonus = compute_dflash_correct_drafts_and_bonus(
candidates=candidates,
target_predict=target_predict,
)
out_tokens, commit_lens = _commit_accept(candidates, accept_len, bonus)
(
accept_len,
commit_lens,
bonus,
out_tokens,
new_seq_lens,
target_predict,
) = self._accept_block(
candidates=candidates,
next_token_logits=logits_output.next_token_logits,
sampling_info=sampling_info,
draft_input=draft_input,
prefix_lens=prefix_lens,
bs=bs,
)
if SIMULATE_ACC_LEN > 0:
if SIMULATE_ACC_TOKEN_MODE not in ("fixed", "real-draft-token"):
python/sglang/srt/speculative/dspark_components/dspark_draft.py
+20 -7
View File
@@ -27,6 +27,7 @@ from sglang.srt.speculative.spec_info import (
SpeculativeAlgorithm,
spec_scale_global_num_tokens,
)
from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
from sglang.srt.speculative.spec_utils import draft_tp_context
from sglang.srt.utils.common import is_pin_memory_available
from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect
@@ -134,6 +135,7 @@ def sample_draft_block(
sampling_info,
markov_head,
device: torch.device,
tp_sync: SpecTpSync,
) -> DraftBlockResult:
bs = base_logits.shape[0]
greedy_mask = resolve_greedy_mask(bs=bs, sampling_info=sampling_info, device=device)
@@ -151,7 +153,9 @@ def sample_draft_block(
def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
expect(_DRAFT_STEP_LOGITS, step_logits, msg=f"step {step_idx}")
return torch.argmax(step_logits, dim=-1)
return tp_sync.sync(
SpecTpSyncSite.DSPARK_DRAFT_GREEDY, torch.argmax(step_logits, dim=-1)
)
else:
@@ -161,11 +165,14 @@ def sample_draft_block(
exp_noise = torch.empty(
step_logits.shape, dtype=torch.float32, device=step_logits.device
).exponential_(1)
return SampleStepTokens.execute(
step_logits=step_logits,
temperatures=temperatures,
greedy_mask=greedy_mask,
exp_noise=exp_noise,
return tp_sync.sync(
SpecTpSyncSite.DSPARK_DRAFT_SAMPLE,
SampleStepTokens.execute(
step_logits=step_logits,
temperatures=temperatures,
greedy_mask=greedy_mask,
exp_noise=exp_noise,
),
)
else:
probs = torch.softmax(
@@ -174,7 +181,10 @@ def sample_draft_block(
probs = expect(_DRAFT_PROBS, probs)
argmax_tokens = torch.argmax(step_logits, dim=-1)
sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
return torch.where(greedy_mask, argmax_tokens, sampled_tokens)
return tp_sync.sync(
SpecTpSyncSite.DSPARK_DRAFT_MULTINOMIAL,
torch.where(greedy_mask, argmax_tokens, sampled_tokens),
)
draft_tokens, corrected_logits = markov_head.sample_block(
base_logits,
@@ -199,6 +209,7 @@ class DraftBlockProposer:
gamma: int,
mask_token_id: int,
draft_block_spec_info,
tp_sync: SpecTpSync,
dp_moe_sync: bool = False,
) -> None:
self.draft_model = draft_model
@@ -208,6 +219,7 @@ class DraftBlockProposer:
self.query_token_num = self.gamma if self.sample_from_anchor else self.gamma + 1
self._mask_token_id = mask_token_id
self._draft_block_spec_info = draft_block_spec_info
self._tp_sync = tp_sync
self._draft_sampler = None
self._dp_moe_sync = dp_moe_sync
# Persistent (bs, gamma) mask-token buffer: only column 0 (the bonus
@@ -310,6 +322,7 @@ class DraftBlockProposer:
sampling_info=sampling_info,
markov_head=self.draft_model.markov_head,
device=device,
tp_sync=self._tp_sync,
)
proposal_block_ids = (
draft_block_ids
python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py
+33 -15
View File
@@ -13,7 +13,7 @@ from sglang.srt.models.dspark import VanillaMarkov
from sglang.srt.speculative.dspark_components.dspark_draft import (
select_draft_hidden_without_anchor,
)
from sglang.srt.utils import get_available_gpu_memory
from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
logger = logging.getLogger(__name__)
@@ -46,6 +46,7 @@ class DsparkDraftSampler:
gamma,
max_bs,
device,
tp_sync: SpecTpSync,
confidence_fn=None,
out=None,
folded_sampling: bool = True,
@@ -73,6 +74,7 @@ class DsparkDraftSampler:
else None
)
self.folded_sampling = folded_sampling
self._tp_sync = tp_sync
self.temperatures = None
self.greedy_mask = None
self.exp_noise = None
@@ -144,15 +146,23 @@ class DsparkDraftSampler:
# In-graph philox noise: each replay advances the generator
# and redraws.
noise = self.exp_noise[:bs].exponential_()
return SampleStepTokens.execute(
step_logits=step_logits,
temperatures=self.temperatures[:bs],
greedy_mask=self.greedy_mask[:bs],
exp_noise=noise,
return self._tp_sync.sync(
SpecTpSyncSite.DSPARK_GRAPH_SAMPLE,
SampleStepTokens.execute(
step_logits=step_logits,
temperatures=self.temperatures[:bs],
greedy_mask=self.greedy_mask[:bs],
exp_noise=noise,
),
)
else:
sampler = greedy_step_sampler
def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
return self._tp_sync.sync(
SpecTpSyncSite.DSPARK_GRAPH_GREEDY,
greedy_step_sampler(step_logits, step_idx),
)
draft_tokens, corrected_logits = self.markov_head.sample_block(
base_logits,
@@ -177,9 +187,12 @@ class DsparkDraftSampler:
self.confidence_out[:bs].copy_(confidence)
def _resolve_folded_sampling(*, model, gamma, max_bs, device, tp_rank) -> bool:
def _resolve_folded_sampling(
*, model, gamma, max_bs, device, tp_rank, available_memory_gb: float
) -> bool:
"""The sampling buffers are baked into the captured draft graph, so AUTO
must decide before capture from a free-memory probe."""
must decide before capture from a free-memory probe. ``available_memory_gb``
is the group minimum, so every rank folds identically."""
mode = envs.SGLANG_DSPARK_FOLDED_SAMPLING.get()
if mode == DsparkFoldedSampling.OFF:
return False
@@ -189,10 +202,7 @@ def _resolve_folded_sampling(*, model, gamma, max_bs, device, tp_rank) -> bool:
noise_bytes = max_bs * vocab * 4
logits_bytes = max_bs * gamma * vocab * _base_logits_dtype(model).itemsize
need_gb = (noise_bytes + logits_bytes) / (1 << 30)
available_gb = get_available_gpu_memory(
device, torch.get_device_module().current_device()
)
if available_gb - need_gb >= _CAPTURE_HEADROOM_GB:
if available_memory_gb - need_gb >= _CAPTURE_HEADROOM_GB:
return True
if tp_rank == 0:
logger.warning(
@@ -201,7 +211,7 @@ def _resolve_folded_sampling(*, model, gamma, max_bs, device, tp_rank) -> bool:
"the eager proposal path. Set SGLANG_DSPARK_FOLDED_SAMPLING=%d "
"to force.",
need_gb,
available_gb,
available_memory_gb,
int(DsparkFoldedSampling.FORCE),
)
return False
@@ -214,6 +224,8 @@ def maybe_build_draft_sampler(
max_bs: int,
device,
tp_rank: int,
tp_sync: SpecTpSync,
available_memory_gb: float,
confidence_fn=None,
out=None,
) -> Optional[DsparkDraftSampler]:
@@ -232,7 +244,12 @@ def maybe_build_draft_sampler(
if getattr(draft_model, "markov_head", None) is None:
return _eager("no markov head")
folded_sampling = _resolve_folded_sampling(
model=draft_model, gamma=gamma, max_bs=max_bs, device=device, tp_rank=tp_rank
model=draft_model,
gamma=gamma,
max_bs=max_bs,
device=device,
tp_rank=tp_rank,
available_memory_gb=available_memory_gb,
)
if tp_rank == 0:
logger.info(
@@ -244,6 +261,7 @@ def maybe_build_draft_sampler(
gamma=gamma,
max_bs=max_bs,
device=device,
tp_sync=tp_sync,
confidence_fn=confidence_fn,
out=out,
folded_sampling=folded_sampling,
python/sglang/srt/speculative/dspark_components/dspark_planner.py
+9 -13
View File
@@ -40,6 +40,7 @@ from sglang.srt.speculative.ragged_verify import (
read_ragged_verify_mode,
round_up_grid,
)
from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
from sglang.srt.utils.common import require_mlp_tp_gather
from sglang.srt.utils.invariants import (
Bucket,
@@ -76,12 +77,14 @@ class DSparkVerifyPlanner:
device,
tp_rank: int,
verify_num_draft_tokens: int,
tp_sync: SpecTpSync,
) -> None:
self.draft_model = draft_model
self.gamma = gamma
self.model_runner = model_runner
self.device = device
self.verify_num_draft_tokens = verify_num_draft_tokens
self._tp_sync = tp_sync
self._align_verify_tokens_to_graph_tier = (
get_spec().speculative_dspark_align_verify_tokens_to_graph_tier
)
@@ -368,11 +371,15 @@ class DSparkVerifyPlanner:
return None
if not get_schedule().disable_overlap_schedule:
return draft_input.verify_token_budget
return self.compute_budget_sync(
# No collective: the budget derives only from the broadcast draft tokens
# (via confidence), replicated req_generation, and the static sps table.
draft_input.verify_token_budget = self.compute_budget_sync(
confidence=confidence,
prefix_lens=prefix_lens,
req_pool_indices=req_pool_indices,
)
return draft_input.verify_token_budget
def confidence_budget_prepare(self):
if not self.schedules_verify_budget:
@@ -574,6 +581,7 @@ class DSparkVerifyPlanner:
budget=budget,
cfg=self._schedule_cfg,
).to(device=device, dtype=torch.int32)
self._tp_sync.sync(SpecTpSyncSite.DSPARK_PLAN, verify_lens)
if resolve_level() >= InvariantCheckLevel.WARN:
verify_lens_64 = verify_lens.to(torch.int64)
@@ -593,12 +601,6 @@ class DSparkVerifyPlanner:
verify_lens=verify_lens,
)
broadcast_group, group_size = verify_lens_broadcast_group(
tp_size=get_parallel().tp_size
)
if group_size > 1:
broadcast_group.broadcast(verify_lens, src=0)
return verify_lens
def _log_verify_lens_decision(
@@ -751,12 +753,6 @@ def uniform_ragged_layout(
)
def verify_lens_broadcast_group(*, tp_size: int) -> tuple:
if is_dp_attention_enabled():
return get_parallel().attn_tp_group, get_parallel().attn_tp_size
return get_tp_group(), tp_size
def verify_layout_grid(
*,
verify_lens_cpu: list[int],
python/sglang/srt/speculative/dspark_components/dspark_verify.py
+17
View File
@@ -41,6 +41,7 @@ from sglang.srt.speculative.dspark_components.dspark_planner import (
apply_logits_adjustments_strided,
)
from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
from sglang.srt.speculative.spec_utils import (
SIMULATE_ACC_METHOD,
sample_simulated_acc_len,
@@ -86,6 +87,7 @@ class TargetVerifyExecutor:
verify_num_draft_tokens: int,
model_runner,
kv_injector: TargetHiddenKvInjector,
tp_sync: SpecTpSync,
verify_epilogue=None,
simulate_acc_len: float = 0.0,
) -> None:
@@ -94,6 +96,7 @@ class TargetVerifyExecutor:
self.verify_num_draft_tokens = verify_num_draft_tokens
self.model_runner = model_runner
self.kv_injector = kv_injector
self._tp_sync = tp_sync
self.verify_epilogue = verify_epilogue
self._verify_backend_self_adds_seq_lens_cache: Optional[bool] = None
self._simulate_acc_len = float(simulate_acc_len)
@@ -138,6 +141,15 @@ class TargetVerifyExecutor:
bs=bs, dtype=correct_len.dtype, device=correct_len.device
)
site = (
SpecTpSyncSite.DSPARK_ACCEPT_GREEDY
if sampling_info is None or sampling_info.is_all_greedy
else SpecTpSyncSite.DSPARK_ACCEPT_SAMPLE
)
self._tp_sync.sync(site, correct_len)
self._tp_sync.sync(site, bonus)
self._tp_sync.sync(site, cap_trim_lens)
finalized = FinalizeAcceptLens.execute(
correct_len=correct_len,
cap_trim_lens=cap_trim_lens,
@@ -487,12 +499,14 @@ class DsparkVerifyEpilogue:
max_bs: int,
verify_num_draft_tokens: int,
device,
tp_sync: SpecTpSync,
commit_ctx: Optional[CommitInjectCtx] = None,
) -> None:
self.max_bs = int(max_bs)
self.stride = int(verify_num_draft_tokens)
self.gamma = self.stride - 1
self.commit_ctx = commit_ctx
self._tp_sync = tp_sync
self.inject_gate_buf = torch.zeros((1,), dtype=torch.int32, device=device)
self.verify_lens_buf = torch.zeros(
(self.max_bs,), dtype=torch.int64, device=device
@@ -637,6 +651,9 @@ class DsparkVerifyEpilogue:
verify_num_draft_tokens=self.stride,
cutoff_verify_lens=verify_lens,
)
self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, correct_len)
self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, bonus)
self._tp_sync.sync(SpecTpSyncSite.DSPARK_ACCEPT_GRAPH, cap_trim_lens)
finalized = finalize_accept_lens_triton(
correct_len=correct_len,
cap_trim_lens=cap_trim_lens,
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+35 -6
View File
@@ -16,6 +16,7 @@ from sglang.srt.lora.layers import unwrap_lora_layer
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.model_executor.cuda_graph_config import Backend
from sglang.srt.model_executor.forward_batch_info import (
CaptureHiddenMode,
compute_position,
@@ -68,6 +69,7 @@ from sglang.srt.speculative.dspark_components.dspark_verify import (
TargetVerifyExecutor,
verify_logits_adjustments_are_noop,
)
from sglang.srt.speculative.spec_tp_sync import SpecTpSync, SpecTpSyncSite
from sglang.srt.speculative.spec_utils import (
GrammarTree,
build_grammar_vocab_mask,
@@ -75,8 +77,8 @@ from sglang.srt.speculative.spec_utils import (
prepare_mamba_track_for_verify,
)
from sglang.srt.utils import (
get_available_gpu_memory,
is_cuda,
is_cuda_alike,
is_npu,
is_pin_memory_available,
)
@@ -114,7 +116,8 @@ class DSparkWorkerV2(BaseSpecWorker):
)
self._is_pd_prefill = get_disagg().disaggregation_mode == "prefill"
self._decode_graph_allowed = (
not get_exec().graph.disable_cuda_graph and not self._is_pd_prefill
get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED
and not self._is_pd_prefill
)
if (
get_parallel().enable_dp_attention
@@ -172,6 +175,18 @@ class DSparkWorkerV2(BaseSpecWorker):
self.speculative_num_draft_tokens = self.verify_num_draft_tokens
self._mask_token_id = runtime_config.mask_token_id
parallel = get_parallel()
self._tp_sync = SpecTpSync(
parallel.attn_tp_group
if parallel.enable_dp_attention
else parallel.tp_group
)
self._draft_graph_group = (
parallel.attn_tp_group
if self._draft_dp_context_enabled
else parallel.tp_group
)
if self.ps.tp_rank == 0:
logger.info(
"Initialized DSpark draft runner. attention_backend=%s, model=%s, "
@@ -220,6 +235,7 @@ class DSparkWorkerV2(BaseSpecWorker):
device=self.device,
tp_rank=self.ps.tp_rank,
verify_num_draft_tokens=self.verify_num_draft_tokens,
tp_sync=self._tp_sync,
)
if (
get_parallel().enable_dp_attention
@@ -248,6 +264,7 @@ class DSparkWorkerV2(BaseSpecWorker):
gamma=self.gamma,
mask_token_id=self._mask_token_id,
draft_block_spec_info=self._draft_block_spec_info,
tp_sync=self._tp_sync,
dp_moe_sync=self._draft_is_moe and get_parallel().enable_dp_attention,
)
self._verify_epilogue = None
@@ -260,6 +277,7 @@ class DSparkWorkerV2(BaseSpecWorker):
max_bs=max(get_exec().graph.cuda_graph_config.decode.bs),
verify_num_draft_tokens=self.verify_num_draft_tokens,
device=self.device,
tp_sync=self._tp_sync,
commit_ctx=CommitInjectCtx(
draft_model=self.draft_model,
block_pos_offsets=self._block_pos_offsets,
@@ -298,6 +316,7 @@ class DSparkWorkerV2(BaseSpecWorker):
verify_num_draft_tokens=self.verify_num_draft_tokens,
model_runner=self.model_runner,
kv_injector=self._kv_injector,
tp_sync=self._tp_sync,
verify_epilogue=self._verify_epilogue,
simulate_acc_len=self._simulate_acc_len,
)
@@ -367,8 +386,13 @@ class DSparkWorkerV2(BaseSpecWorker):
def init_cuda_graphs(self):
capture_decode_cuda_graph = self._decode_graph_allowed
if is_cuda() and capture_decode_cuda_graph:
available_mem = get_available_gpu_memory(self.device, self.gpu_id)
available_mem = self._tp_sync.available_memory_gb(
SpecTpSyncSite.DSPARK_MEM,
self.device,
self.gpu_id,
group=self._draft_graph_group,
)
if is_cuda_alike() and capture_decode_cuda_graph:
if available_mem < 1.0:
capture_decode_cuda_graph = False
logger.warning(
@@ -385,7 +409,9 @@ class DSparkWorkerV2(BaseSpecWorker):
# from being the intended precision fallback, skipping the
# unused hook avoids paying for two proposal computations.
if envs.SGLANG_DSPARK_FOLDED_PROPOSAL.get():
self._draft_sampler = self._maybe_build_draft_sampler()
self._draft_sampler = self._maybe_build_draft_sampler(
available_memory_gb=available_mem
)
if self._draft_sampler is not None:
self.draft_model_runner.capture_tail_hooks.append(
make_draft_sampler_capture_hook(self._draft_sampler)
@@ -395,13 +421,15 @@ class DSparkWorkerV2(BaseSpecWorker):
capture_decode_cuda_graph=capture_decode_cuda_graph
)
def _maybe_build_draft_sampler(self):
def _maybe_build_draft_sampler(self, *, available_memory_gb: float):
return maybe_build_draft_sampler(
draft_model=self.draft_model,
gamma=self.gamma,
max_bs=max(get_exec().graph.cuda_graph_config.decode.bs),
device=self.device,
tp_rank=self.ps.tp_rank,
tp_sync=self._tp_sync,
available_memory_gb=available_memory_gb,
confidence_fn=(
self._verify_planner.compute_confidence_tensor
if self._verify_planner.carries_confidence
@@ -461,6 +489,7 @@ class DSparkWorkerV2(BaseSpecWorker):
)
logits_output = batch_output.logits_output
next_token_ids = batch_output.next_token_ids
self._tp_sync.sync(SpecTpSyncSite.DSPARK_TARGET, next_token_ids)
batch_output.new_seq_lens = batch.seq_lens
if on_publish is not None:
on_publish(batch_output.new_seq_lens)
python/sglang/srt/speculative/spec_tp_sync.py
+129
View File
@@ -0,0 +1,129 @@
from __future__ import annotations
import logging
from enum import IntEnum
import torch
from sglang.srt.environ import envs
from sglang.srt.utils import get_available_gpu_memory
logger = logging.getLogger(__name__)
class SpecTpSyncSite(IntEnum):
"""Every place a speculative step broadcasts a decision from rank 0.
Number and slug are the stable handles ``SGLANG_SPEC_TP_SYNC`` selects by.
"""
# -- DSpark --
DSPARK_MEM = 1 # gates both graph capture and folded sampling
DSPARK_DRAFT_GREEDY = 2
DSPARK_DRAFT_SAMPLE = 3
DSPARK_DRAFT_MULTINOMIAL = 4
DSPARK_GRAPH_SAMPLE = 5 # in-graph philox, redrawn per replay
DSPARK_GRAPH_GREEDY = 6
DSPARK_PLAN = 7
DSPARK_ACCEPT_GREEDY = 8
DSPARK_ACCEPT_SAMPLE = 9
DSPARK_ACCEPT_GRAPH = 10
DSPARK_TARGET = 11
# -- DFlash --
DFLASH_MEM = 12
DFLASH_SELECTOR = 13
DFLASH_ACCEPT_SAMPLE = 14
DFLASH_ACCEPT_GREEDY = 15
DFLASH_TARGET = 16
@property
def slug(self) -> str:
return self.name.lower().replace("_", "-")
_ALL = frozenset(SpecTpSyncSite)
_INIT = frozenset({SpecTpSyncSite.DSPARK_MEM, SpecTpSyncSite.DFLASH_MEM})
# Sites that draw from the RNG, so they can differ under identical logits.
_RNG = frozenset(
{
SpecTpSyncSite.DSPARK_DRAFT_SAMPLE,
SpecTpSyncSite.DSPARK_DRAFT_MULTINOMIAL,
SpecTpSyncSite.DSPARK_GRAPH_SAMPLE,
SpecTpSyncSite.DSPARK_ACCEPT_SAMPLE,
SpecTpSyncSite.DSPARK_TARGET,
SpecTpSyncSite.DFLASH_SELECTOR,
SpecTpSyncSite.DFLASH_ACCEPT_SAMPLE,
SpecTpSyncSite.DFLASH_TARGET,
}
)
_PRESETS = {
"all": _ALL,
"off": frozenset(),
"none": frozenset(),
# The one input measured to differ across ranks.
"init": _INIT,
"rng": _INIT | _RNG,
}
_BY_SLUG = {site.slug: site for site in SpecTpSyncSite}
_BY_NUMBER = {str(int(site)): site for site in SpecTpSyncSite}
def _resolve(name: str) -> frozenset[SpecTpSyncSite]:
if name in _PRESETS:
return _PRESETS[name]
site = _BY_SLUG.get(name) or _BY_NUMBER.get(name)
if site is None:
raise ValueError(
f"SGLANG_SPEC_TP_SYNC: unknown token {name!r}. "
f"Presets: {sorted(_PRESETS)}. "
f"Sites: {[(int(s), s.slug) for s in SpecTpSyncSite]}."
)
return frozenset({site})
def parse_spec_tp_sync(spec: str) -> frozenset[SpecTpSyncSite]:
"""Parse ``SGLANG_SPEC_TP_SYNC``; see its comment in environ.py for the syntax."""
sites: frozenset[SpecTpSyncSite] = frozenset()
for token in spec.replace(" ", "").replace("_", "-").lower().split(","):
if not token:
continue
negate = token.startswith("-")
value = _resolve(token[1:] if negate else token)
sites = sites - value if negate else sites | value
return sites
class SpecTpSync:
"""Broadcasts a speculative decision from rank 0 to its TP group."""
def __init__(self, tp_group) -> None:
self._tp_group = tp_group
# Parsed even on a single rank so a typo fails on every deployment.
sites = parse_spec_tp_sync(envs.SGLANG_SPEC_TP_SYNC.get())
self._sites = sites if tp_group.world_size > 1 else frozenset()
if sites != _ALL and tp_group.world_size > 1 and tp_group.rank_in_group == 0:
logger.warning(
"Speculative TP sync limited to %s.",
[f"{int(s)}:{s.slug}" for s in sorted(sites)] or "no site",
)
def enabled(self, site: SpecTpSyncSite) -> bool:
return site in self._sites
def sync(self, site: SpecTpSyncSite, values: torch.Tensor) -> torch.Tensor:
if site in self._sites:
self._tp_group.broadcast(values, src=0)
return values
def available_memory_gb(self, site: SpecTpSyncSite, device, gpu_id, *, group):
"""Free GPU memory, reduced to the group minimum when ``site`` is on."""
distributed = self.enabled(site) and group.world_size > 1
return get_available_gpu_memory(
device,
gpu_id,
distributed=distributed,
cpu_group=group.cpu_group if distributed else None,
)