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

[Feature] Spec-Overlap supporting DP-ATTN; PD-Disaggregation; npugraph mode (#12443)

Browse Source
This commit is contained in:
Even Zhou
2025-11-15 21:51:07 +08:00
committed by GitHub
parent 0d41ddfbd0
commit 2aec8b6e1b
24 changed files with 657 additions and 366 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -866,6 +866,10 @@ class ForwardBatch:
self.spec_info.accept_length = self.spec_info.accept_length[:bs]
logits_output.next_token_logits = logits_output.next_token_logits[:bs]
logits_output.hidden_states = logits_output.hidden_states[:bs]
elif self.forward_mode.is_draft_extend_v2(): # draft extend_v2
bs = bs * self.spec_info.num_tokens_per_batch
logits_output.next_token_logits = logits_output.next_token_logits[:bs]
logits_output.hidden_states = logits_output.hidden_states[:bs]
elif self.forward_mode.is_extend() or self.forward_mode.is_idle():
logits_output.next_token_logits = logits_output.next_token_logits[:bs]
logits_output.hidden_states = logits_output.hidden_states[:bs]

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

@@ -2236,7 +2236,7 @@ class ModelRunner:
reinit_attn_backend=reinit_attn_backend,
forward_count=split_forward_count,
)
elif forward_batch.forward_mode.is_extend():
elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
ret = self.forward_extend(
forward_batch,
skip_attn_backend_init=skip_attn_backend_init,