spec: gate dp mlp sync with server args (#24177)
This commit is contained in:
@@ -2463,7 +2463,11 @@ class Scheduler(
|
|||||||
new_batch = self.get_new_batch_prefill()
|
new_batch = self.get_new_batch_prefill()
|
||||||
|
|
||||||
need_mlp_sync = self.require_mlp_sync
|
need_mlp_sync = self.require_mlp_sync
|
||||||
if need_mlp_sync and not self.spec_algorithm.is_none():
|
if (
|
||||||
|
need_mlp_sync
|
||||||
|
and not self.spec_algorithm.is_none()
|
||||||
|
and not self.server_args.speculative_skip_dp_mlp_sync
|
||||||
|
):
|
||||||
# NOTE: This branch makes sure prefill and decode batches will not be mixed when spec and dp-attn is enabled.
|
# NOTE: This branch makes sure prefill and decode batches will not be mixed when spec and dp-attn is enabled.
|
||||||
# Before merging the new batch into running batch:
|
# Before merging the new batch into running batch:
|
||||||
# 1. All new batches are none -> need_mlp_sync remains true (sync is needed for decode batch).
|
# 1. All new batches are none -> need_mlp_sync remains true (sync is needed for decode batch).
|
||||||
|
|||||||
@@ -526,6 +526,7 @@ class ServerArgs:
|
|||||||
speculative_draft_model_quantization: Optional[str] = None
|
speculative_draft_model_quantization: Optional[str] = None
|
||||||
speculative_adaptive: bool = False
|
speculative_adaptive: bool = False
|
||||||
speculative_adaptive_config: Optional[str] = None
|
speculative_adaptive_config: Optional[str] = None
|
||||||
|
speculative_skip_dp_mlp_sync: bool = False
|
||||||
|
|
||||||
# Speculative decoding (ngram)
|
# Speculative decoding (ngram)
|
||||||
speculative_ngram_min_bfs_breadth: int = 1
|
speculative_ngram_min_bfs_breadth: int = 1
|
||||||
@@ -3300,6 +3301,12 @@ class ServerArgs:
|
|||||||
if self.speculative_algorithm == "NEXTN":
|
if self.speculative_algorithm == "NEXTN":
|
||||||
self.speculative_algorithm = "EAGLE"
|
self.speculative_algorithm = "EAGLE"
|
||||||
|
|
||||||
|
if self.speculative_skip_dp_mlp_sync:
|
||||||
|
assert self.speculative_algorithm == "EAGLE", (
|
||||||
|
"--speculative-skip-dp-mlp-sync is only supported with "
|
||||||
|
f"speculative_algorithm == EAGLE, got {self.speculative_algorithm}."
|
||||||
|
)
|
||||||
|
|
||||||
if self.speculative_algorithm == "DFLASH":
|
if self.speculative_algorithm == "DFLASH":
|
||||||
if self.enable_dp_attention:
|
if self.enable_dp_attention:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
@@ -5497,6 +5504,13 @@ class ServerArgs:
|
|||||||
help="Path to a JSON config file for adaptive speculative decoding tuning knobs ",
|
help="Path to a JSON config file for adaptive speculative decoding tuning knobs ",
|
||||||
default=ServerArgs.speculative_adaptive_config,
|
default=ServerArgs.speculative_adaptive_config,
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--speculative-skip-dp-mlp-sync",
|
||||||
|
action="store_true",
|
||||||
|
default=ServerArgs.speculative_skip_dp_mlp_sync,
|
||||||
|
help="Skip the extra MLP sync that the scheduler performs before merging a new batch "
|
||||||
|
"when speculative decoding + DP attention are both enabled.",
|
||||||
|
)
|
||||||
|
|
||||||
# Multi-layer Eagle speculative decoding
|
# Multi-layer Eagle speculative decoding
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
|
|||||||
Reference in New Issue
Block a user