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

refactor context parallel state (#17213)

Browse Source 
Co-authored-by: Shunkang <182541032+Shunkangz@users.noreply.github.co>
This commit is contained in:
dongjiyingdjy
2026-02-13 23:18:17 +08:00
committed by GitHub
parent 0012d6a4eb
commit 8b4c364960
27 changed files with 847 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -84,6 +84,8 @@ class EAGLEWorker(TpModelWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -144,6 +146,8 @@ class EAGLEWorker(TpModelWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,

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

@@ -86,6 +86,8 @@ class EagleDraftWorker(BaseDraftWorker):
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -97,6 +99,8 @@ class EagleDraftWorker(BaseDraftWorker):
self.moe_ep_rank = moe_ep_rank
self.nccl_port = nccl_port
self.target_worker = target_worker
self.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
# Args for easy access
self.device = server_args.device
@@ -134,6 +138,8 @@ class EagleDraftWorker(BaseDraftWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
@@ -582,6 +588,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -608,7 +616,15 @@ class EAGLEWorkerV2(BaseSpecWorker):
server_args.context_length = target_worker.model_runner.model_config.context_len
self._draft_worker = EagleDraftWorker(
server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
nccl_port,
target_worker,
)
# Some dummy tensors

4
python/sglang/srt/speculative/multi_layer_eagle_worker.py
View File

@@ -76,6 +76,8 @@ class MultiLayerEagleWorker(TpModelWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -135,6 +137,8 @@ class MultiLayerEagleWorker(TpModelWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,

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

@@ -70,6 +70,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -117,6 +119,8 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
@@ -532,6 +536,8 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -557,7 +563,15 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
server_args.context_length = target_worker.model_runner.model_config.context_len
self._draft_worker = MultiLayerEagleDraftWorker(
server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
nccl_port,
target_worker,
)
# Some dummy tensors

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

@@ -30,6 +30,8 @@ class NGRAMWorker:
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):

4
python/sglang/srt/speculative/standalone_worker.py
View File

@@ -30,6 +30,8 @@ class StandaloneWorker(EAGLEWorker):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -79,6 +81,8 @@ class StandaloneWorker(EAGLEWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,

18
python/sglang/srt/speculative/standalone_worker_v2.py
View File

@@ -42,6 +42,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -53,6 +55,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
self.moe_ep_rank = moe_ep_rank
self.nccl_port = nccl_port
self.target_worker = target_worker
self.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
# Args for easy access
self.device = server_args.device
@@ -89,6 +93,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
pp_rank=0, # FIXME
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
nccl_port=nccl_port,
is_draft_worker=True,
req_to_token_pool=self.req_to_token_pool,
@@ -131,6 +137,8 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -157,7 +165,15 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
# 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
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
nccl_port,
target_worker,
)
# Some dummy tensors