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

[refactor] Collect MoE and DP-attention runtime state into typed flag groups (#30347)

Browse Source
This commit is contained in:
Cheng Wan
2026-07-07 21:29:27 -07:00
committed by GitHub
parent be32c57598
commit b7cca0bf8f
7 changed files with 292 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/dp_attention.py
+7 -8
View File
@@ -29,6 +29,7 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory,
)
from sglang.srt.runtime_context import get_flags
from sglang.srt.utils import get_bool_env_var, is_hip
if TYPE_CHECKING:
@@ -44,8 +45,6 @@ _ATTN_DP_RANK: Optional[int] = None
_ATTN_DP_SIZE: Optional[int] = None
_LOCAL_ATTN_DP_SIZE: Optional[int] = None
_LOCAL_ATTN_DP_RANK: Optional[int] = None
_ENABLE_DP_ATTENTION_FLAG: bool = False
_DP_MAX_LEN_WITH_IDLE = False
_is_hip = is_hip()
_USE_ROCM700A_WA = _is_hip and get_bool_env_var("SGLANG_USE_ROCM700A")
@@ -77,7 +76,7 @@ class DpPaddingMode(IntEnum):
if is_extend_in_batch and dp_size > 1:
# Hybrid-SSM models materialize idle ranks via the MAX_LEN
# fabricated-row conversion; other models keep mainline SUM_LEN.
if _DP_MAX_LEN_WITH_IDLE and min(global_num_tokens) == 0:
if get_flags().dp.max_len_with_idle and min(global_num_tokens) == 0:
return DpPaddingMode.MAX_LEN
return DpPaddingMode.SUM_LEN
@@ -281,9 +280,9 @@ def initialize_dp_attention(
model_config: ModelConfig,
):
global _ATTN_DP_RANK, _ATTN_DP_SIZE
global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK, _ENABLE_DP_ATTENTION_FLAG
global _DP_MAX_LEN_WITH_IDLE
_DP_MAX_LEN_WITH_IDLE = (
global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK
dp = get_flags().dp
dp.max_len_with_idle = (
getattr(model_config.hf_config, "hybrid_override_pattern", None) is not None
)
enable_dp_attention = server_args.enable_dp_attention
@@ -291,7 +290,7 @@ def initialize_dp_attention(
moe_dense_tp_size = server_args.moe_dense_tp_size
attn_cp_size = server_args.attn_cp_size
_ENABLE_DP_ATTENTION_FLAG = enable_dp_attention
dp.enabled = enable_dp_attention
tp_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size()
@@ -321,7 +320,7 @@ def initialize_dp_attention(
def is_dp_attention_enabled() -> bool:
return _ENABLE_DP_ATTENTION_FLAG
return get_flags().dp.enabled
def is_allocation_symmetric() -> bool:
python/sglang/srt/layers/moe/utils.py
+66 -97
View File
@@ -4,7 +4,7 @@ import logging
import os
from contextlib import contextmanager
from enum import Enum, IntEnum
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING
import torch
@@ -12,7 +12,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled,
)
from sglang.srt.runtime_context import get_parallel
from sglang.srt.runtime_context import get_flags, get_parallel
from sglang.srt.utils import is_cuda, is_npu
_is_npu = is_npu()
@@ -245,122 +245,96 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype:
return DeepEPOutputDtype.FP8
MOE_A2A_BACKEND: Optional[MoeA2ABackend] = None
MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None
SPECULATIVE_MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None
SPECULATIVE_MOE_A2A_BACKEND: Optional[MoeA2ABackend] = None
DEEPEP_MODE: Optional[DeepEPMode] = None
IS_TBO_ENABLED: Optional[bool] = None
IS_SBO_ENABLED: Optional[bool] = None
TBO_TOKEN_DISTRIBUTION_THRESHOLD: Optional[float] = None
DEEPEP_CONFIG: Optional[str] = None
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER: Optional[bool] = None
MOE_QUANTIZATION: Optional[str] = None
def initialize_moe_config(server_args: ServerArgs):
global MOE_A2A_BACKEND
global MOE_RUNNER_BACKEND
global SPECULATIVE_MOE_RUNNER_BACKEND
global SPECULATIVE_MOE_A2A_BACKEND
global DEEPEP_MODE
global DEEPEP_CONFIG
global IS_TBO_ENABLED
global IS_SBO_ENABLED
global TBO_TOKEN_DISTRIBUTION_THRESHOLD
global DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER
global MOE_QUANTIZATION
MOE_A2A_BACKEND = MoeA2ABackend(server_args.moe_a2a_backend)
MOE_RUNNER_BACKEND = MoeRunnerBackend(server_args.moe_runner_backend)
SPECULATIVE_MOE_RUNNER_BACKEND = (
moe = get_flags().moe
moe.a2a_backend = MoeA2ABackend(server_args.moe_a2a_backend)
moe.runner_backend = MoeRunnerBackend(server_args.moe_runner_backend)
moe.speculative_runner_backend = (
MoeRunnerBackend(server_args.speculative_moe_runner_backend)
if server_args.speculative_moe_runner_backend is not None
else MOE_RUNNER_BACKEND
else moe.runner_backend
)
SPECULATIVE_MOE_A2A_BACKEND = (
moe.speculative_a2a_backend = (
MoeA2ABackend(server_args.speculative_moe_a2a_backend)
if server_args.speculative_moe_a2a_backend is not None
else MOE_A2A_BACKEND
else moe.a2a_backend
)
DEEPEP_MODE = DeepEPMode(server_args.deepep_mode)
DEEPEP_CONFIG = server_args.deepep_config or ""
IS_TBO_ENABLED = server_args.enable_two_batch_overlap
IS_SBO_ENABLED = server_args.enable_single_batch_overlap
if IS_SBO_ENABLED and is_cuda():
moe.deepep_mode = DeepEPMode(server_args.deepep_mode)
moe.deepep_config = server_args.deepep_config or ""
moe.tbo_enabled = server_args.enable_two_batch_overlap
moe.sbo_enabled = server_args.enable_single_batch_overlap
if moe.sbo_enabled and is_cuda():
if torch.cuda.get_device_capability()[0] == 9:
raise ValueError(
"SBO (single batch overlap) is not supported on SM90 GPUs with latest sgl-deep-gemm wheel. Please try removing --enable-single-batch-overlap argument."
)
TBO_TOKEN_DISTRIBUTION_THRESHOLD = server_args.tbo_token_distribution_threshold
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = (
server_args.disable_flashinfer_cutlass_moe_fp4_allgather
)
MOE_QUANTIZATION = server_args.quantization
moe.tbo_token_distribution_threshold = server_args.tbo_token_distribution_threshold
moe.disable_fp4_allgather = server_args.disable_flashinfer_cutlass_moe_fp4_allgather
moe.quantization = server_args.quantization
def get_moe_a2a_backend() -> MoeA2ABackend:
global MOE_A2A_BACKEND
if MOE_A2A_BACKEND is None:
MOE_A2A_BACKEND = MoeA2ABackend.NONE
return MOE_A2A_BACKEND
moe = get_flags().moe
if moe.a2a_backend is None:
moe.a2a_backend = MoeA2ABackend.NONE
return moe.a2a_backend
def get_moe_runner_backend() -> MoeRunnerBackend:
global MOE_RUNNER_BACKEND
if MOE_RUNNER_BACKEND is None:
MOE_RUNNER_BACKEND = MoeRunnerBackend.AUTO
return MOE_RUNNER_BACKEND
moe = get_flags().moe
if moe.runner_backend is None:
moe.runner_backend = MoeRunnerBackend.AUTO
return moe.runner_backend
def get_speculative_moe_runner_backend() -> MoeRunnerBackend:
global SPECULATIVE_MOE_RUNNER_BACKEND
if SPECULATIVE_MOE_RUNNER_BACKEND is None:
moe = get_flags().moe
if moe.speculative_runner_backend is None:
logger.warning(
"SPECULATIVE_MOE_RUNNER_BACKEND is not initialized, using auto backend"
)
SPECULATIVE_MOE_RUNNER_BACKEND = MoeRunnerBackend.AUTO
return SPECULATIVE_MOE_RUNNER_BACKEND
moe.speculative_runner_backend = MoeRunnerBackend.AUTO
return moe.speculative_runner_backend
def get_speculative_moe_a2a_backend() -> MoeA2ABackend:
global SPECULATIVE_MOE_A2A_BACKEND
if SPECULATIVE_MOE_A2A_BACKEND is None:
moe = get_flags().moe
if moe.speculative_a2a_backend is None:
logger.warning(
"SPECULATIVE_MOE_A2A_BACKEND is not initialized, using none backend"
)
SPECULATIVE_MOE_A2A_BACKEND = MoeA2ABackend.NONE
return SPECULATIVE_MOE_A2A_BACKEND
moe.speculative_a2a_backend = MoeA2ABackend.NONE
return moe.speculative_a2a_backend
def get_deepep_mode() -> DeepEPMode:
global DEEPEP_MODE
if DEEPEP_MODE is None:
moe = get_flags().moe
if moe.deepep_mode is None:
logger.warning("DEEPEP_MODE is not initialized, using auto mode")
DEEPEP_MODE = DeepEPMode.AUTO
return DEEPEP_MODE
moe.deepep_mode = DeepEPMode.AUTO
return moe.deepep_mode
def get_deepep_config() -> str:
global DEEPEP_CONFIG
if DEEPEP_CONFIG is None:
moe = get_flags().moe
if moe.deepep_config is None:
logger.warning("DEEPEP_CONFIG is not initialized, using default config")
DEEPEP_CONFIG = ""
return DEEPEP_CONFIG
moe.deepep_config = ""
return moe.deepep_config
def is_tbo_enabled() -> bool:
global IS_TBO_ENABLED
if IS_TBO_ENABLED is None:
IS_TBO_ENABLED = False
return IS_TBO_ENABLED
moe = get_flags().moe
if moe.tbo_enabled is None:
moe.tbo_enabled = False
return moe.tbo_enabled
def is_sbo_enabled() -> bool:
global IS_SBO_ENABLED
if IS_SBO_ENABLED is None:
IS_SBO_ENABLED = False
return IS_SBO_ENABLED
moe = get_flags().moe
if moe.sbo_enabled is None:
moe.sbo_enabled = False
return moe.sbo_enabled
def is_deepep_class_backend() -> bool:
@@ -388,13 +362,13 @@ def is_flashinfer_cutedsl_v1_path() -> bool:
def get_tbo_token_distribution_threshold() -> float:
global TBO_TOKEN_DISTRIBUTION_THRESHOLD
if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None:
moe = get_flags().moe
if moe.tbo_token_distribution_threshold is None:
logger.warning(
"TBO_TOKEN_DISTRIBUTION_THRESHOLD is not initialized, using 0.48"
)
TBO_TOKEN_DISTRIBUTION_THRESHOLD = 0.48
return TBO_TOKEN_DISTRIBUTION_THRESHOLD
moe.tbo_token_distribution_threshold = 0.48
return moe.tbo_token_distribution_threshold
def filter_moe_weight_param_global_expert(name, x, num_local_experts):
@@ -413,11 +387,11 @@ def should_use_flashinfer_cutlass_moe_fp4_allgather():
Perform FP4 quantize before all-gather for flashinfer cutlass moe to reduce communication cost for high-throughput serving.
"""
return (
not DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER
not get_flags().moe.disable_fp4_allgather
and get_moe_a2a_backend().is_none()
and get_moe_runner_backend().is_flashinfer_cutlass()
and is_dp_attention_enabled()
and MOE_QUANTIZATION == "modelopt_fp4"
and get_flags().moe.quantization == "modelopt_fp4"
and get_parallel().moe_ep_size == get_parallel().attn_dp_size
)
@@ -484,13 +458,13 @@ def speculative_moe_backend_context():
Context manager to temporarily use the speculative MoE backend for draft model operations.
This ensures that draft models in speculative decoding use the configured speculative backend.
"""
global MOE_RUNNER_BACKEND
original_backend = MOE_RUNNER_BACKEND
moe = get_flags().moe
original_backend = moe.runner_backend
try:
MOE_RUNNER_BACKEND = get_speculative_moe_runner_backend()
moe.runner_backend = get_speculative_moe_runner_backend()
yield
finally:
MOE_RUNNER_BACKEND = original_backend
moe.runner_backend = original_backend
@contextmanager
@@ -499,22 +473,17 @@ def speculative_moe_a2a_backend_context():
Context manager to temporarily use the speculative MoE A2A backend for draft model operations.
This ensures that draft models in speculative decoding use the configured speculative A2A backend.
"""
global MOE_A2A_BACKEND
global DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER
original_backend = MOE_A2A_BACKEND
original_disable_flashinfer_cutlass_moe_fp4_allgather = (
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER
)
moe = get_flags().moe
original_backend = moe.a2a_backend
original_disable_fp4_allgather = moe.disable_fp4_allgather
try:
MOE_A2A_BACKEND = get_speculative_moe_a2a_backend()
moe.a2a_backend = get_speculative_moe_a2a_backend()
# Disable FP4 allgather for spec decode since MTP layers are unquantized
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = True
moe.disable_fp4_allgather = True
yield
finally:
MOE_A2A_BACKEND = original_backend
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = (
original_disable_flashinfer_cutlass_moe_fp4_allgather
)
moe.a2a_backend = original_backend
moe.disable_fp4_allgather = original_disable_fp4_allgather
# The type of method in top-K routing, for use in torch custom op
python/sglang/srt/runtime_context.py
+39 -1
View File
@@ -270,6 +270,42 @@ class CaptureFlags(_FlagGroupBase):
enable_torch_compile: bool = False
@dataclasses.dataclass
class MoeFlags(_FlagGroupBase):
"""MoE runtime flags, materialized by ``initialize_moe_config`` (scheduler
init, after distributed setup). ``a2a_backend`` / ``runner_backend`` /
``disable_fp4_allgather`` are the ACTIVE values: the speculative contexts
in ``layers.moe.utils`` swap them around draft-model forwards. Values are
the parsed enums from ``layers.moe.utils``; ``None`` means "not
initialized yet" and the accessors fall back lazily.
"""
a2a_backend: Any = None
runner_backend: Any = None
speculative_runner_backend: Any = None
speculative_a2a_backend: Any = None
deepep_mode: Any = None
deepep_config: str | None = None
tbo_enabled: bool | None = None
sbo_enabled: bool | None = None
tbo_token_distribution_threshold: float | None = None
disable_fp4_allgather: bool | None = None
quantization: str | None = None
@dataclasses.dataclass
class DpFlags(_FlagGroupBase):
"""DP-attention runtime flags, materialized by ``initialize_dp_attention``
(after distributed setup; reads the model config). Topology values
(sizes/ranks) stay on ``layers.dp_attention`` until the parallel vertical
migrates them."""
enabled: bool = False
# Hybrid-SSM models materialize idle ranks via the MAX_LEN fabricated-row
# conversion (set when hf_config has hybrid_override_pattern).
max_len_with_idle: bool = False
@dataclasses.dataclass
class Flags(_FlagGroupBase):
"""Root of the runtime-flags tier.
@@ -277,10 +313,12 @@ class Flags(_FlagGroupBase):
Resolved configuration lives on ``server_args`` fields (materialized at
the end of ``__post_init__``) — this tier only carries genuine runtime
state whose value is not a function of the configuration alone, grouped
by lifecycle (today: ``capture``).
by lifecycle (``capture``) or subsystem (``moe`` / ``dp``).
"""
capture: CaptureFlags = dataclasses.field(default_factory=CaptureFlags)
moe: MoeFlags = dataclasses.field(default_factory=MoeFlags)
dp: DpFlags = dataclasses.field(default_factory=DpFlags)
class RuntimeContext: