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

Introduce ModelRunner.ps ParallelState (#31161)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-14 16:01:14 +08:00
committed by GitHub
parent 1dc48c2c3b
commit 725920915f
28 changed files with 202 additions and 276 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/benchmark/one_batch.py
+32 -6
View File
@@ -69,7 +69,9 @@ from sglang.srt.distributed.parallel_state import (
destroy_distributed_environment,
destroy_model_parallel,
)
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.entrypoints.engine import _set_envs_and_config
from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
from sglang.srt.layers.moe import initialize_moe_config
from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
@@ -300,16 +302,40 @@ def load_model(server_args, port_args, gpu_id, tp_rank):
moe_ep_rank = tp_rank // (server_args.tp_size // server_args.ep_size)
model_config = ModelConfig.from_server_args(server_args)
attn_tp_rank, attn_tp_size, attn_dp_rank, attn_dp_size = (
compute_dp_attention_world_info(
server_args.enable_dp_attention,
tp_rank,
server_args.tp_size,
server_args.dp_size,
server_args.attn_cp_size,
)
)
ps = ParallelState(
tp_rank=tp_rank,
tp_size=server_args.tp_size,
pp_rank=0,
pp_size=1,
dp_rank=None,
dp_size=server_args.dp_size,
attn_tp_rank=attn_tp_rank,
attn_tp_size=attn_tp_size,
attn_cp_rank=0,
attn_cp_size=server_args.attn_cp_size,
attn_dp_rank=attn_dp_rank,
attn_dp_size=attn_dp_size,
moe_ep_rank=moe_ep_rank,
moe_ep_size=server_args.ep_size,
moe_dp_rank=None,
moe_dp_size=server_args.moe_dp_size,
dcp_size=server_args.dcp_size,
gpu_id=gpu_id,
)
runner_kwargs = dict(
model_config=model_config,
mem_fraction_static=server_args.mem_fraction_static,
gpu_id=gpu_id,
tp_rank=tp_rank,
tp_size=server_args.tp_size,
moe_ep_rank=moe_ep_rank,
moe_ep_size=server_args.ep_size,
pp_rank=0,
pp_size=1,
ps=ps,
nccl_port=port_args.nccl_port,
server_args=server_args,
)
python/sglang/srt/distributed/bootstrap.py
+13 -10
View File
@@ -19,6 +19,7 @@ from sglang.srt.distributed import (
set_mscclpp_all_reduce,
set_torch_symm_mem_all_reduce,
)
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import initialize_dp_attention
from sglang.srt.platforms import current_platform
@@ -52,20 +53,22 @@ def init_torch_distributed(
server_args: ServerArgs,
model_config: ModelConfig,
device: str,
gpu_id: int,
tp_rank: int,
tp_size: int,
pp_rank: int,
pp_size: int,
dp_size: int,
attn_cp_size: int,
moe_ep_size: int,
moe_dp_size: int,
dcp_size: int,
ps: ParallelState,
dist_port: int,
is_draft_worker: bool,
local_omp_cpuid: Optional[List[int]],
):
gpu_id = ps.gpu_id
tp_rank = ps.tp_rank
tp_size = ps.tp_size
pp_rank = ps.pp_rank
pp_size = ps.pp_size
dp_size = ps.attn_dp_size
attn_cp_size = ps.attn_cp_size
moe_ep_size = ps.moe_ep_size
moe_dp_size = ps.moe_dp_size
dcp_size = ps.dcp_size
tic = time.perf_counter()
logger.info("Init torch distributed begin.")
python/sglang/srt/distributed/parallel_state_wrapper.py
+26
View File
@@ -20,4 +20,30 @@ class ParallelState:
moe_ep_size: int
moe_dp_rank: Optional[int]
moe_dp_size: int
dcp_size: int
gpu_id: int
@staticmethod
def trivial(**overrides: Optional[int]) -> "ParallelState":
kwargs: dict[str, Optional[int]] = dict(
tp_rank=0,
tp_size=1,
pp_rank=0,
pp_size=1,
dp_rank=0,
dp_size=1,
attn_tp_rank=0,
attn_tp_size=1,
attn_cp_rank=0,
attn_cp_size=1,
attn_dp_rank=0,
attn_dp_size=1,
moe_ep_rank=0,
moe_ep_size=1,
moe_dp_rank=0,
moe_dp_size=1,
dcp_size=1,
gpu_id=0,
)
kwargs.update(overrides)
return ParallelState(**kwargs)
python/sglang/srt/hardware_backend/mlx/tp_worker.py
+1 -7
View File
@@ -63,14 +63,8 @@ class MlxTpModelWorker(TpModelWorker):
model_config=self.model_config,
mem_fraction_static=self.server_args.mem_fraction_static,
gpu_id=self.gpu_id,
tp_rank=self.tp_rank,
tp_size=self.tp_size,
moe_ep_rank=self.moe_ep_rank,
moe_ep_size=self.ep_size,
pp_rank=self.pp_rank,
pp_size=self.pp_size,
ps=self.ps,
nccl_port=self.nccl_port,
dp_rank=self.dp_rank,
server_args=self.server_args,
is_draft_worker=self.is_draft_worker,
req_to_token_pool=self.req_to_token_pool,
python/sglang/srt/managers/scheduler.py
+3 -11
View File
@@ -387,6 +387,7 @@ class Scheduler(
moe_ep_size=server_args.ep_size,
moe_dp_rank=moe_dp_rank,
moe_dp_size=server_args.moe_dp_size,
dcp_size=server_args.dcp_size,
gpu_id=gpu_id,
)
@@ -747,12 +748,7 @@ class Scheduler(
worker_kwargs = dict(
server_args=self.server_args,
gpu_id=self.ps.gpu_id,
tp_rank=self.ps.tp_rank,
moe_ep_rank=self.ps.moe_ep_rank,
pp_rank=self.ps.pp_rank,
attn_cp_rank=self.ps.attn_cp_rank,
moe_dp_rank=self.ps.moe_dp_rank,
dp_rank=self.ps.dp_rank,
ps=self.ps,
nccl_port=self.nccl_port,
)
@@ -776,13 +772,9 @@ class Scheduler(
draft_worker_kwargs = dict(
server_args=self.server_args,
gpu_id=self.ps.gpu_id,
tp_rank=self.ps.tp_rank,
moe_ep_rank=self.ps.moe_ep_rank,
ps=self.ps,
nccl_port=self.nccl_port,
target_worker=self.tp_worker,
dp_rank=self.ps.dp_rank,
attn_cp_rank=self.ps.attn_cp_rank,
moe_dp_rank=self.ps.moe_dp_rank,
)
if self.server_args.speculative_draft_load_format is not None:
python/sglang/srt/managers/tp_worker.py
+7 -31
View File
@@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, List, Optional, Tuple
import torch
from sglang.srt.distributed import get_pp_group, get_world_group
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.io_struct import (
DestroyWeightsUpdateGroupReqInput,
GetWeightsByNameReqInput,
@@ -173,7 +174,7 @@ class BaseTpWorker(ABC):
monkey_patch_torch_reductions()
success, message = self.model_runner.weight_updater.update_weights_from_tensor(
named_tensors=MultiprocessingSerializer.deserialize(
recv_req.serialized_named_tensors[self.tp_rank]
recv_req.serialized_named_tensors[self.ps.tp_rank]
),
load_format=recv_req.load_format,
)
@@ -237,12 +238,7 @@ class TpModelWorker(BaseTpWorker):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
moe_ep_rank: int,
pp_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
dp_rank: Optional[int],
ps: ParallelState,
nccl_port: int,
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
@@ -253,21 +249,13 @@ class TpModelWorker(BaseTpWorker):
):
# Parse args
self.server_args = server_args
self.tp_size = server_args.tp_size
self.ep_size = server_args.ep_size
self.pp_size = server_args.pp_size
self.tp_rank = tp_rank
self.moe_ep_rank = moe_ep_rank
self.pp_rank = pp_rank
self.dp_rank = dp_rank
self.ps = ps
self.gpu_id = gpu_id
self.nccl_port = nccl_port
self.is_draft_worker = is_draft_worker
self.is_multi_layer_eagle = is_multi_layer_eagle
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
# Draft worker: target's resolved MemoryPoolConfig (forwarded to ModelRunner).
self.memory_pool_config = memory_pool_config
# Draft worker: target's effective context length; the draft runs at
@@ -317,7 +305,7 @@ class TpModelWorker(BaseTpWorker):
# Sync random seed across TP workers
self.random_seed = broadcast_pyobj(
[server_args.random_seed],
self.tp_size * self.pp_rank + tp_rank,
self.ps.tp_size * self.ps.pp_rank + self.ps.tp_rank,
self.world_group.cpu_group,
src=self.world_group.ranks[0],
)[0]
@@ -394,14 +382,8 @@ class TpModelWorker(BaseTpWorker):
model_config=self.model_config,
mem_fraction_static=self.server_args.mem_fraction_static,
gpu_id=self.gpu_id,
tp_rank=self.tp_rank,
tp_size=self.tp_size,
moe_ep_rank=self.moe_ep_rank,
moe_ep_size=self.ep_size,
pp_rank=self.pp_rank,
pp_size=self.pp_size,
ps=self.ps,
nccl_port=self.nccl_port,
dp_rank=self.dp_rank,
server_args=self.server_args,
is_draft_worker=self.is_draft_worker,
req_to_token_pool=self.req_to_token_pool,
@@ -420,14 +402,8 @@ class TpModelWorker(BaseTpWorker):
model_config=self.model_config,
mem_fraction_static=self.server_args.mem_fraction_static,
gpu_id=self.gpu_id,
tp_rank=self.tp_rank,
tp_size=self.tp_size,
moe_ep_rank=self.moe_ep_rank,
moe_ep_size=self.ep_size,
pp_rank=self.pp_rank,
pp_size=self.pp_size,
ps=self.ps,
nccl_port=self.nccl_port,
dp_rank=self.dp_rank,
server_args=self.server_args,
is_draft_worker=self.is_draft_worker,
req_to_token_pool=self.req_to_token_pool,
python/sglang/srt/model_executor/model_runner.py
+20 -35
View File
@@ -43,6 +43,7 @@ from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
prealloc_symmetric_memory_pool,
)
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.dllm.config import DllmConfig
from sglang.srt.elastic_ep.elastic_ep import (
ElasticEPStateManager,
@@ -242,17 +243,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
model_config: ModelConfig,
mem_fraction_static: float,
gpu_id: int,
tp_rank: int,
tp_size: int,
moe_ep_rank: int,
moe_ep_size: int,
pp_rank: int,
pp_size: int,
ps: ParallelState,
nccl_port: int,
server_args: ServerArgs,
dp_rank: Optional[int] = None,
attn_cp_rank: Optional[int] = None,
moe_dp_rank: Optional[int] = None,
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
@@ -267,22 +260,21 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.memory_pool_config = memory_pool_config
self.device = server_args.device
self.gpu_id = gpu_id
self.tp_rank = tp_rank
self.tp_size = tp_size
self.tp_rank = ps.tp_rank
self.tp_size = ps.tp_size
self.dcp_size = server_args.dcp_size
self.dcp_rank = self.tp_rank % self.dcp_size
self.moe_ep_rank = moe_ep_rank
self.moe_ep_size = moe_ep_size
self.dp_rank = dp_rank
self.attn_dp_size = (
server_args.dp_size if server_args.enable_dp_attention else 1
)
self.pp_rank = pp_rank
self.pp_size = pp_size
self.attn_cp_rank = attn_cp_rank
self.attn_cp_size = server_args.attn_cp_size
self.moe_dp_rank = moe_dp_rank
self.moe_dp_size = server_args.moe_dp_size
self.dcp_rank = ps.tp_rank % self.dcp_size
self.ps = ps
self.moe_ep_rank = ps.moe_ep_rank
self.moe_ep_size = ps.moe_ep_size
self.dp_rank = ps.dp_rank
self.attn_dp_size = ps.attn_dp_size
self.pp_rank = ps.pp_rank
self.pp_size = ps.pp_size
self.attn_cp_rank = ps.attn_cp_rank
self.attn_cp_size = ps.attn_cp_size
self.moe_dp_rank = ps.moe_dp_rank
self.moe_dp_size = ps.moe_dp_size
self.model_config = model_config
self.dist_port = nccl_port
self.server_args = server_args
@@ -356,7 +348,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
self.war_fastpath_read_done_event: Optional[torch.cuda.Event] = None
# CPU offload
set_offloader(create_offloader_from_server_args(server_args, dp_rank=dp_rank))
set_offloader(
create_offloader_from_server_args(server_args, dp_rank=self.ps.dp_rank)
)
self._weight_checker = WeightChecker(model_runner=self)
@@ -813,16 +807,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
server_args=self.server_args,
model_config=self.model_config,
device=self.device,
gpu_id=self.gpu_id,
tp_rank=self.tp_rank,
tp_size=self.tp_size,
pp_rank=self.pp_rank,
pp_size=self.pp_size,
dp_size=self.attn_dp_size,
attn_cp_size=self.attn_cp_size,
moe_ep_size=self.moe_ep_size,
moe_dp_size=self.moe_dp_size,
dcp_size=self.dcp_size,
ps=self.ps,
dist_port=self.dist_port,
is_draft_worker=self.is_draft_worker,
local_omp_cpuid=self.local_omp_cpuid if self.device == "cpu" else None,
python/sglang/srt/speculative/dflash_worker_v2.py
+10 -20
View File
@@ -1,5 +1,6 @@
import logging
import math
from dataclasses import replace
from typing import List, Optional
import torch
@@ -11,6 +12,7 @@ from sglang.kernels.ops.speculative.dflash import (
)
from sglang.srt.configs.hybrid_arch import mambaish_config
from sglang.srt.distributed import get_tp_group
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -108,21 +110,13 @@ class DFlashWorkerV2(BaseSpecWorker):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
nccl_port: int,
target_worker: TpModelWorker,
):
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.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
self.ps = ps
self.nccl_port = nccl_port
self._target_worker = target_worker
self.model_runner = target_worker.model_runner
@@ -141,11 +135,7 @@ class DFlashWorkerV2(BaseSpecWorker):
bundle = build_draft_tp_worker(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
target_model_config=target_worker.model_runner.model_config,
algo_label="DFLASH",
@@ -181,7 +171,7 @@ class DFlashWorkerV2(BaseSpecWorker):
mask_token=self._mask_token,
mask_token_id=self._mask_token_id_override,
)
if self.tp_rank == 0:
if self.ps.tp_rank == 0:
logger.info(
"Initialized DFLASH draft runner. attention_backend=%s, model=%s, block_size=%s, draft_window_size=%s, compact_cache=%s",
bundle.resolved_attention_backend,
@@ -400,7 +390,7 @@ class DFlashWorkerV2(BaseSpecWorker):
break
if fused_disable_reason is not None:
if self.tp_rank == 0:
if self.ps.tp_rank == 0:
logger.info(
"DFLASH fused KV materialization disabled: %s",
fused_disable_reason,
@@ -422,7 +412,7 @@ class DFlashWorkerV2(BaseSpecWorker):
max_position_hint=self.target_worker.model_runner.model_config.context_len
+ int(self.block_size),
)
if self.tp_rank == 0:
if self.ps.tp_rank == 0:
logger.info(
"DFLASH fused KV materialization enabled. "
"n_layers=%d, num_kv_heads=%d, head_dim=%d",
@@ -629,7 +619,7 @@ class DFlashWorkerV2(BaseSpecWorker):
if resolved_id is None:
resolved_id = tokenizer.convert_tokens_to_ids(mask_token)
if added and self.tp_rank == 0:
if added and self.ps.tp_rank == 0:
logger.info(
"Added DFLASH mask token to tokenizer. token=%s, mask_token_id=%s, tokenizer_len=%s, model_vocab_size=%s",
mask_token,
@@ -1199,7 +1189,7 @@ class DFlashWorkerV2(BaseSpecWorker):
if (
not is_dflash_sampling_verify_available()
and not self._warned_sampling_fallback
and self.tp_rank == 0
and self.ps.tp_rank == 0
):
logger.warning(
"DFLASH non-greedy verification is unavailable on this build/device; "
python/sglang/srt/speculative/draft_worker_common.py
+3 -11
View File
@@ -17,6 +17,7 @@ from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.model_executor.model_runner import ModelRunner
logger = logging.getLogger(__name__)
@@ -57,11 +58,7 @@ def build_draft_tp_worker(
*,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
nccl_port: int,
target_model_config: ModelConfig,
algo_label: str,
@@ -98,12 +95,7 @@ def build_draft_tp_worker(
draft_worker = TpModelWorker(
server_args=draft_server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
moe_ep_rank=moe_ep_rank,
pp_rank=0,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
dp_rank=dp_rank,
ps=ps,
nccl_port=nccl_port,
is_draft_worker=True,
)
python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py
+9 -19
View File
@@ -1,9 +1,11 @@
import logging
from contextlib import nullcontext
from dataclasses import replace
from typing import Optional
import torch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -63,21 +65,13 @@ class DSparkWorkerV2(BaseSpecWorker):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
nccl_port: int,
target_worker: TpModelWorker,
):
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.attn_cp_rank = attn_cp_rank
self.moe_dp_rank = moe_dp_rank
self.ps = ps
self.nccl_port = nccl_port
self._target_worker = target_worker
self.model_runner = target_worker.model_runner
@@ -100,11 +94,7 @@ class DSparkWorkerV2(BaseSpecWorker):
bundle = build_draft_tp_worker(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
target_model_config=target_worker.model_runner.model_config,
algo_label="DSPARK",
@@ -129,7 +119,7 @@ class DSparkWorkerV2(BaseSpecWorker):
self.speculative_num_draft_tokens = self.verify_num_draft_tokens
self._mask_token_id = runtime_config.mask_token_id
if self.tp_rank == 0:
if self.ps.tp_rank == 0:
logger.info(
"Initialized DSpark draft runner. attention_backend=%s, model=%s, "
"gamma=%s, verify_num_draft_tokens=%s, mask_token_id=%s, "
@@ -165,7 +155,7 @@ class DSparkWorkerV2(BaseSpecWorker):
gamma=self.gamma,
model_runner=self.model_runner,
device=self.device,
tp_rank=self.tp_rank,
tp_rank=self.ps.tp_rank,
server_args=self.server_args,
verify_num_draft_tokens=self.verify_num_draft_tokens,
)
@@ -256,7 +246,7 @@ class DSparkWorkerV2(BaseSpecWorker):
planner=self._verify_planner,
gamma=self.gamma,
verify_num_draft_tokens=self.verify_num_draft_tokens,
tp_rank=self.tp_rank,
tp_rank=self.ps.tp_rank,
device=self.device,
simulate_acc_len=self._simulate_acc_len,
)
@@ -332,7 +322,7 @@ class DSparkWorkerV2(BaseSpecWorker):
gamma=self.gamma,
max_bs=max(self.server_args.cuda_graph_config.decode.bs),
device=self.device,
tp_rank=self.tp_rank,
tp_rank=self.ps.tp_rank,
confidence_fn=(
self._verify_planner.compute_confidence_tensor
if self._verify_planner.carries_confidence
python/sglang/srt/speculative/eagle_worker_v2.py
+10 -28
View File
@@ -1,11 +1,13 @@
import contextlib
import logging
import time
from dataclasses import replace
from typing import List, Optional
import torch
from sglang.kernels.ops.speculative.eagle import fill_bonus_tokens_func
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import (
EAGLEDraftExtendNpuGraphRunner,
@@ -127,24 +129,16 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
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.ps = ps
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
@@ -173,12 +167,8 @@ class EagleDraftWorker(EagleDraftWorkerBase):
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=0, # spec workers don't support pipeline parallelism
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
# spec workers don't support pipeline parallelism
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
is_draft_worker=True,
)
@@ -1048,11 +1038,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -1061,7 +1047,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
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.tp_rank = tp_rank
self.ps = ps
self.gpu_id = gpu_id
self.device = server_args.device
self._target_worker = target_worker
@@ -1079,11 +1065,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
self._draft_worker = EagleDraftWorker(
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
ps,
nccl_port,
target_worker,
)
@@ -1762,7 +1744,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
monkey_patch_torch_reductions()
named_tensors = MultiprocessingSerializer.deserialize(
recv_req.serialized_named_tensors[self.tp_rank]
recv_req.serialized_named_tensors[self.ps.tp_rank]
)
success, message = (
self.draft_worker.draft_runner.weight_updater.update_weights_from_tensor(
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+9 -25
View File
@@ -13,11 +13,13 @@
# ==============================================================================
import logging
from typing import TYPE_CHECKING, List, Optional
from dataclasses import replace
from typing import TYPE_CHECKING, List
import torch
from sglang.kernels.ops.speculative.eagle import fill_bonus_tokens_func
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.multi_layer_eagle_draft_extend_npu_graph_runner import (
MultiLayerEagleMultiStepDraftExtendNpuGraphRunner,
@@ -92,20 +94,14 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
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.ps = ps
self.nccl_port = nccl_port
self.target_worker = target_worker
self.draft_extend_attn_backend_list = []
@@ -136,12 +132,8 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=0, # spec workers don't support pipeline parallelism
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
# spec workers don't support pipeline parallelism
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
is_draft_worker=True,
is_multi_layer_eagle=True,
@@ -665,11 +657,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -695,11 +683,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
self._draft_worker = MultiLayerEagleDraftWorker(
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
ps,
nccl_port,
target_worker,
)
python/sglang/srt/speculative/ngram_worker.py
+3 -6
View File
@@ -8,6 +8,7 @@ from sgl_kernel.speculative import reconstruct_indices_from_tree_mask
from sglang.kernels.ops.speculative.cache_locs import (
assign_extend_cache_locs_func as assign_extend_cache_locs_func,
)
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.utils.logprob import compute_spec_v2_logprobs
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
@@ -50,11 +51,7 @@ class NGRAMWorker(BaseSpecWorker):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -62,7 +59,7 @@ class NGRAMWorker(BaseSpecWorker):
self.enable_overlap = not server_args.disable_overlap_schedule
self._target_worker = target_worker
self.model_runner = target_worker.model_runner
self.tp_rank = tp_rank
self.tp_rank = ps.tp_rank
self.page_size = server_args.page_size
self.draft_token_num: int = server_args.speculative_num_draft_tokens
self.max_trie_depth: int = server_args.speculative_ngram_max_trie_depth
python/sglang/srt/speculative/standalone_worker_v2.py
+8 -26
View File
@@ -1,8 +1,10 @@
import logging
from dataclasses import replace
from typing import Optional
import torch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
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
@@ -29,24 +31,16 @@ class StandaloneDraftWorker(EagleDraftWorker):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: int,
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
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.ps = ps
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
@@ -74,12 +68,8 @@ class StandaloneDraftWorker(EagleDraftWorker):
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
tp_rank=tp_rank,
pp_rank=0, # spec workers don't support pipeline parallelism
dp_rank=dp_rank,
moe_ep_rank=moe_ep_rank,
attn_cp_rank=attn_cp_rank,
moe_dp_rank=moe_dp_rank,
# spec workers don't support pipeline parallelism
ps=replace(ps, pp_rank=0),
nccl_port=nccl_port,
is_draft_worker=True,
)
@@ -146,11 +136,7 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
self,
server_args: ServerArgs,
gpu_id: int,
tp_rank: int,
dp_rank: Optional[int],
moe_ep_rank: int,
attn_cp_rank: int,
moe_dp_rank: int,
ps: ParallelState,
nccl_port: int,
target_worker: TpModelWorker,
):
@@ -177,11 +163,7 @@ class StandaloneWorkerV2(EAGLEWorkerV2):
self._draft_worker = StandaloneDraftWorker(
server_args,
gpu_id,
tp_rank,
dp_rank,
moe_ep_rank,
attn_cp_rank,
moe_dp_rank,
ps,
nccl_port,
target_worker,
)
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+4
View File
@@ -7,6 +7,7 @@ import torch.nn.functional as F
from torch import nn
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
@@ -286,7 +287,9 @@ class TinyModelConfig:
num_key_value_heads=num_kv_heads,
head_dim=head_dim,
)
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
def get_num_attention_heads(self, tp_size: int) -> int:
assert self.num_attention_heads % tp_size == 0
@@ -322,6 +325,7 @@ class MockModelRunner(ModelRunner):
self.tp_size = 1
self.dp_size = 1
self.pp_size = 1
self.ps = ParallelState.trivial()
self.is_draft_worker = False
self.spec_algorithm = SpeculativeAlgorithm.NONE
# The runner lifecycle warms up kernels in capture() / first execute()
python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py
+4
View File
@@ -5,6 +5,7 @@ from typing import Any
import torch
from torch import nn
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, ReqToTokenPool
@@ -258,7 +259,9 @@ class TinyDSAModelConfig:
index_topk=index_topk,
num_hidden_layers=1,
)
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
class DSAMockModelRunner(ModelRunner):
@@ -308,6 +311,7 @@ class DSAMockModelRunner(ModelRunner):
self._kernel_warmed_up = True
self.dp_size = 1
self.pp_size = 1
self.ps = ParallelState.trivial()
self._server_args_override = get_context().override_server_args(
attention_backend=case.backend,
chunked_prefill_size=-1,
python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py
+4
View File
@@ -19,6 +19,7 @@ from typing import Any
import torch
from torch import nn
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.environ import envs
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.dsv4.quant_k_cache import (
@@ -282,7 +283,9 @@ class TinyDSV4ModelConfig:
num_hidden_layers=len(compression_ratios),
compress_ratios=list(compression_ratios),
)
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
class MockDSV4ModelRunner:
@@ -332,6 +335,7 @@ class MockDSV4ModelRunner:
self.tp_size = 1
self.dp_size = 1
self.pp_size = 1
self.ps = ParallelState.trivial()
self._server_args_override = get_context().override_server_args(
attention_backend=case.backend,
chunked_prefill_size=-1,
python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py
+4
View File
@@ -4,6 +4,7 @@ from types import SimpleNamespace
import torch
from torch import nn
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention import (
dual_chunk_flashattention_backend as _dual_chunk_backend,
)
@@ -286,7 +287,9 @@ class TinyDualChunkModelConfig:
dual_chunk_attention_config or DUAL_CHUNK_CONFIG
),
)
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
def get_num_attention_heads(self, tp_size: int) -> int:
assert self.num_attention_heads % tp_size == 0
@@ -323,6 +326,7 @@ class DualChunkMockModelRunner(ModelRunner):
self._kernel_warmed_up = True
self.dp_size = 1
self.pp_size = 1
self.ps = ParallelState.trivial()
self._server_args_override = get_context().override_server_args(
attention_backend=case.backend,
chunked_prefill_size=-1,
python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py
+4
View File
@@ -10,6 +10,7 @@ from sglang.srt.configs.mamba_utils import (
Mamba2StateShape,
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
HybridLinearAttnBackend,
@@ -182,7 +183,9 @@ class TinyGDNModelConfig:
self.attention_chunk_size = None
self.sliding_window_size = None
self.hf_config = SimpleNamespace(architectures=["TinyGDNForCausalLM"])
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
def get_num_kv_heads(self, tp_size: int) -> int:
assert self.num_key_value_heads % tp_size == 0
@@ -210,6 +213,7 @@ class MockGDNModelRunner(ModelRunner):
self.dtype = dtype
self.kv_cache_dtype = dtype
self.gpu_id = 0
self.ps = ParallelState.trivial()
self.canary_manager = None
self.page_size = case.page_size
self.model_config = model_config
python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py
+4
View File
@@ -10,6 +10,7 @@ from sglang.srt.configs.mamba_utils import (
Mamba2StateDType,
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
HybridLinearAttnBackend,
@@ -188,7 +189,9 @@ class TinyKDAModelConfig:
self.attention_chunk_size = None
self.sliding_window_size = None
self.hf_config = SimpleNamespace(architectures=["TinyKDAForCausalLM"])
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
def get_num_kv_heads(self, tp_size: int) -> int:
assert self.num_key_value_heads % tp_size == 0
@@ -216,6 +219,7 @@ class MockKDAModelRunner(ModelRunner):
self.dtype = dtype
self.kv_cache_dtype = dtype
self.gpu_id = 0
self.ps = ParallelState.trivial()
self.canary_manager = None
self.page_size = case.page_size
self.model_config = model_config
python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py
+4
View File
@@ -10,6 +10,7 @@ from sglang.srt.configs.mamba_utils import (
Mamba2StateShape,
)
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.attention.linear.lightning_backend import (
LightningAttentionBackend,
@@ -198,7 +199,9 @@ class TinyLightningModelConfig:
num_hidden_layers=num_hidden_layers,
linear_backend=linear_backend,
)
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
def get_num_kv_heads(self, tp_size: int) -> int:
assert self.num_key_value_heads % tp_size == 0
@@ -224,6 +227,7 @@ class MockLightningModelRunner(ModelRunner):
self.dtype = dtype
self.kv_cache_dtype = dtype
self.gpu_id = 0
self.ps = ParallelState.trivial()
self.canary_manager = None
self.page_size = case.page_size
self.model_config = model_config
python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py
+7 -4
View File
@@ -18,12 +18,14 @@ _parallel_override.__enter__()
# Provide a stub group with world_size=1 so use_symmetric_memory short-circuits.
_linear_mod.get_tp_group = lambda: SimpleNamespace(world_size=1)
from sglang.srt.configs.falcon_h1 import FalconH1Config # noqa: E402
from sglang.srt.configs.mamba_utils import ( # noqa: E402
Mamba2CacheParams,
Mamba2StateDType,
Mamba2StateShape,
)
from sglang.srt.configs.model_config import AttentionArch # noqa: E402
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention.attention_registry import ( # noqa: E402
ATTENTION_BACKENDS,
)
@@ -278,14 +280,14 @@ class TinyMamba2ModelConfig:
self.is_local_attention_model = False
self.attention_chunk_size = None
self.sliding_window_size = None
# Mamba2AttnBackend reads mamba2_config.mamba_chunk_size; expose it
# through a SimpleNamespace-as-hf_config so runner.mamba2_config returns
# something non-None with the expected attribute.
self.hf_config = SimpleNamespace(
# Mamba2AttnBackend reads mamba2_config(model_config).mamba_chunk_size; expose it
self.hf_config = FalconH1Config(
architectures=["TinyMamba2ForCausalLM"],
mamba_chunk_size=case.mamba_chunk_size,
)
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
def get_num_kv_heads(self, tp_size: int) -> int:
assert self.num_key_value_heads % tp_size == 0
@@ -310,6 +312,7 @@ class MockMamba2ModelRunner(ModelRunner):
self.dtype = dtype
self.kv_cache_dtype = dtype
self.gpu_id = 0
self.ps = ParallelState.trivial()
self.canary_manager = None
self.page_size = case.page_size
self.model_config = model_config
python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py
+4
View File
@@ -7,6 +7,7 @@ import torch.nn.functional as F
from torch import nn
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool, ReqToTokenPool
@@ -191,7 +192,9 @@ class TinyMLAModelConfig:
qk_rope_head_dim=qk_rope_head_dim,
v_head_dim=kv_lora_rank,
)
self.hf_config.get_text_config = lambda: self.hf_config
self.hf_text_config = self.hf_config
self.linear_attn_registry_result = None
def get_num_attention_heads(self, tp_size: int) -> int:
assert self.num_attention_heads % tp_size == 0
@@ -233,6 +236,7 @@ class MockMLAModelRunner(ModelRunner):
self.tp_size = 1
self.dp_size = 1
self.pp_size = 1
self.ps = ParallelState.trivial()
speculative_num_draft_tokens = (
max(case.input_lens)
if case.forward_mode.is_target_verify()