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

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
+7 -8
View File
@@ -29,6 +29,7 @@ from sglang.srt.distributed import (
from sglang.srt.distributed.device_communicators.pynccl_allocator import ( from sglang.srt.distributed.device_communicators.pynccl_allocator import (
use_symmetric_memory, use_symmetric_memory,
) )
from sglang.srt.runtime_context import get_flags
from sglang.srt.utils import get_bool_env_var, is_hip from sglang.srt.utils import get_bool_env_var, is_hip
if TYPE_CHECKING: if TYPE_CHECKING:
@@ -44,8 +45,6 @@ _ATTN_DP_RANK: Optional[int] = None
_ATTN_DP_SIZE: Optional[int] = None _ATTN_DP_SIZE: Optional[int] = None
_LOCAL_ATTN_DP_SIZE: Optional[int] = None _LOCAL_ATTN_DP_SIZE: Optional[int] = None
_LOCAL_ATTN_DP_RANK: 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() _is_hip = is_hip()
_USE_ROCM700A_WA = _is_hip and get_bool_env_var("SGLANG_USE_ROCM700A") _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: if is_extend_in_batch and dp_size > 1:
# Hybrid-SSM models materialize idle ranks via the MAX_LEN # Hybrid-SSM models materialize idle ranks via the MAX_LEN
# fabricated-row conversion; other models keep mainline SUM_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.MAX_LEN
return DpPaddingMode.SUM_LEN return DpPaddingMode.SUM_LEN
@@ -281,9 +280,9 @@ def initialize_dp_attention(
model_config: ModelConfig, model_config: ModelConfig,
): ):
global _ATTN_DP_RANK, _ATTN_DP_SIZE global _ATTN_DP_RANK, _ATTN_DP_SIZE
global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK, _ENABLE_DP_ATTENTION_FLAG global _LOCAL_ATTN_DP_SIZE, _LOCAL_ATTN_DP_RANK
global _DP_MAX_LEN_WITH_IDLE dp = get_flags().dp
_DP_MAX_LEN_WITH_IDLE = ( dp.max_len_with_idle = (
getattr(model_config.hf_config, "hybrid_override_pattern", None) is not None getattr(model_config.hf_config, "hybrid_override_pattern", None) is not None
) )
enable_dp_attention = server_args.enable_dp_attention 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 moe_dense_tp_size = server_args.moe_dense_tp_size
attn_cp_size = server_args.attn_cp_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_rank = get_tensor_model_parallel_rank()
tp_size = get_tensor_model_parallel_world_size() tp_size = get_tensor_model_parallel_world_size()
@@ -321,7 +320,7 @@ def initialize_dp_attention(
def is_dp_attention_enabled() -> bool: def is_dp_attention_enabled() -> bool:
return _ENABLE_DP_ATTENTION_FLAG return get_flags().dp.enabled
def is_allocation_symmetric() -> bool: def is_allocation_symmetric() -> bool:
+66 -97
View File
@@ -4,7 +4,7 @@ import logging
import os import os
from contextlib import contextmanager from contextlib import contextmanager
from enum import Enum, IntEnum from enum import Enum, IntEnum
from typing import TYPE_CHECKING, Optional from typing import TYPE_CHECKING
import torch import torch
@@ -12,7 +12,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
is_dp_attention_enabled, 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 from sglang.srt.utils import is_cuda, is_npu
_is_npu = is_npu() _is_npu = is_npu()
@@ -245,122 +245,96 @@ def get_deepep_output_dtype(self) -> DeepEPOutputDtype:
return DeepEPOutputDtype.FP8 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): def initialize_moe_config(server_args: ServerArgs):
global MOE_A2A_BACKEND moe = get_flags().moe
global MOE_RUNNER_BACKEND moe.a2a_backend = MoeA2ABackend(server_args.moe_a2a_backend)
global SPECULATIVE_MOE_RUNNER_BACKEND moe.runner_backend = MoeRunnerBackend(server_args.moe_runner_backend)
global SPECULATIVE_MOE_A2A_BACKEND moe.speculative_runner_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 = (
MoeRunnerBackend(server_args.speculative_moe_runner_backend) MoeRunnerBackend(server_args.speculative_moe_runner_backend)
if server_args.speculative_moe_runner_backend is not None 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) MoeA2ABackend(server_args.speculative_moe_a2a_backend)
if server_args.speculative_moe_a2a_backend is not None 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) moe.deepep_mode = DeepEPMode(server_args.deepep_mode)
DEEPEP_CONFIG = server_args.deepep_config or "" moe.deepep_config = server_args.deepep_config or ""
IS_TBO_ENABLED = server_args.enable_two_batch_overlap moe.tbo_enabled = server_args.enable_two_batch_overlap
IS_SBO_ENABLED = server_args.enable_single_batch_overlap moe.sbo_enabled = server_args.enable_single_batch_overlap
if IS_SBO_ENABLED and is_cuda(): if moe.sbo_enabled and is_cuda():
if torch.cuda.get_device_capability()[0] == 9: if torch.cuda.get_device_capability()[0] == 9:
raise ValueError( 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." "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 moe.tbo_token_distribution_threshold = server_args.tbo_token_distribution_threshold
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = ( moe.disable_fp4_allgather = server_args.disable_flashinfer_cutlass_moe_fp4_allgather
server_args.disable_flashinfer_cutlass_moe_fp4_allgather moe.quantization = server_args.quantization
)
MOE_QUANTIZATION = server_args.quantization
def get_moe_a2a_backend() -> MoeA2ABackend: def get_moe_a2a_backend() -> MoeA2ABackend:
global MOE_A2A_BACKEND moe = get_flags().moe
if MOE_A2A_BACKEND is None: if moe.a2a_backend is None:
MOE_A2A_BACKEND = MoeA2ABackend.NONE moe.a2a_backend = MoeA2ABackend.NONE
return MOE_A2A_BACKEND return moe.a2a_backend
def get_moe_runner_backend() -> MoeRunnerBackend: def get_moe_runner_backend() -> MoeRunnerBackend:
global MOE_RUNNER_BACKEND moe = get_flags().moe
if MOE_RUNNER_BACKEND is None: if moe.runner_backend is None:
MOE_RUNNER_BACKEND = MoeRunnerBackend.AUTO moe.runner_backend = MoeRunnerBackend.AUTO
return MOE_RUNNER_BACKEND return moe.runner_backend
def get_speculative_moe_runner_backend() -> MoeRunnerBackend: def get_speculative_moe_runner_backend() -> MoeRunnerBackend:
global SPECULATIVE_MOE_RUNNER_BACKEND moe = get_flags().moe
if SPECULATIVE_MOE_RUNNER_BACKEND is None: if moe.speculative_runner_backend is None:
logger.warning( logger.warning(
"SPECULATIVE_MOE_RUNNER_BACKEND is not initialized, using auto backend" "SPECULATIVE_MOE_RUNNER_BACKEND is not initialized, using auto backend"
) )
SPECULATIVE_MOE_RUNNER_BACKEND = MoeRunnerBackend.AUTO moe.speculative_runner_backend = MoeRunnerBackend.AUTO
return SPECULATIVE_MOE_RUNNER_BACKEND return moe.speculative_runner_backend
def get_speculative_moe_a2a_backend() -> MoeA2ABackend: def get_speculative_moe_a2a_backend() -> MoeA2ABackend:
global SPECULATIVE_MOE_A2A_BACKEND moe = get_flags().moe
if SPECULATIVE_MOE_A2A_BACKEND is None: if moe.speculative_a2a_backend is None:
logger.warning( logger.warning(
"SPECULATIVE_MOE_A2A_BACKEND is not initialized, using none backend" "SPECULATIVE_MOE_A2A_BACKEND is not initialized, using none backend"
) )
SPECULATIVE_MOE_A2A_BACKEND = MoeA2ABackend.NONE moe.speculative_a2a_backend = MoeA2ABackend.NONE
return SPECULATIVE_MOE_A2A_BACKEND return moe.speculative_a2a_backend
def get_deepep_mode() -> DeepEPMode: def get_deepep_mode() -> DeepEPMode:
global DEEPEP_MODE moe = get_flags().moe
if DEEPEP_MODE is None: if moe.deepep_mode is None:
logger.warning("DEEPEP_MODE is not initialized, using auto mode") logger.warning("DEEPEP_MODE is not initialized, using auto mode")
DEEPEP_MODE = DeepEPMode.AUTO moe.deepep_mode = DeepEPMode.AUTO
return DEEPEP_MODE return moe.deepep_mode
def get_deepep_config() -> str: def get_deepep_config() -> str:
global DEEPEP_CONFIG moe = get_flags().moe
if DEEPEP_CONFIG is None: if moe.deepep_config is None:
logger.warning("DEEPEP_CONFIG is not initialized, using default config") logger.warning("DEEPEP_CONFIG is not initialized, using default config")
DEEPEP_CONFIG = "" moe.deepep_config = ""
return DEEPEP_CONFIG return moe.deepep_config
def is_tbo_enabled() -> bool: def is_tbo_enabled() -> bool:
global IS_TBO_ENABLED moe = get_flags().moe
if IS_TBO_ENABLED is None: if moe.tbo_enabled is None:
IS_TBO_ENABLED = False moe.tbo_enabled = False
return IS_TBO_ENABLED return moe.tbo_enabled
def is_sbo_enabled() -> bool: def is_sbo_enabled() -> bool:
global IS_SBO_ENABLED moe = get_flags().moe
if IS_SBO_ENABLED is None: if moe.sbo_enabled is None:
IS_SBO_ENABLED = False moe.sbo_enabled = False
return IS_SBO_ENABLED return moe.sbo_enabled
def is_deepep_class_backend() -> bool: def is_deepep_class_backend() -> bool:
@@ -388,13 +362,13 @@ def is_flashinfer_cutedsl_v1_path() -> bool:
def get_tbo_token_distribution_threshold() -> float: def get_tbo_token_distribution_threshold() -> float:
global TBO_TOKEN_DISTRIBUTION_THRESHOLD moe = get_flags().moe
if TBO_TOKEN_DISTRIBUTION_THRESHOLD is None: if moe.tbo_token_distribution_threshold is None:
logger.warning( logger.warning(
"TBO_TOKEN_DISTRIBUTION_THRESHOLD is not initialized, using 0.48" "TBO_TOKEN_DISTRIBUTION_THRESHOLD is not initialized, using 0.48"
) )
TBO_TOKEN_DISTRIBUTION_THRESHOLD = 0.48 moe.tbo_token_distribution_threshold = 0.48
return TBO_TOKEN_DISTRIBUTION_THRESHOLD return moe.tbo_token_distribution_threshold
def filter_moe_weight_param_global_expert(name, x, num_local_experts): 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. Perform FP4 quantize before all-gather for flashinfer cutlass moe to reduce communication cost for high-throughput serving.
""" """
return ( 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_a2a_backend().is_none()
and get_moe_runner_backend().is_flashinfer_cutlass() and get_moe_runner_backend().is_flashinfer_cutlass()
and is_dp_attention_enabled() 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 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. 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. This ensures that draft models in speculative decoding use the configured speculative backend.
""" """
global MOE_RUNNER_BACKEND moe = get_flags().moe
original_backend = MOE_RUNNER_BACKEND original_backend = moe.runner_backend
try: try:
MOE_RUNNER_BACKEND = get_speculative_moe_runner_backend() moe.runner_backend = get_speculative_moe_runner_backend()
yield yield
finally: finally:
MOE_RUNNER_BACKEND = original_backend moe.runner_backend = original_backend
@contextmanager @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. 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. This ensures that draft models in speculative decoding use the configured speculative A2A backend.
""" """
global MOE_A2A_BACKEND moe = get_flags().moe
global DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER original_backend = moe.a2a_backend
original_backend = MOE_A2A_BACKEND original_disable_fp4_allgather = moe.disable_fp4_allgather
original_disable_flashinfer_cutlass_moe_fp4_allgather = (
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER
)
try: 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 FP4 allgather for spec decode since MTP layers are unquantized
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = True moe.disable_fp4_allgather = True
yield yield
finally: finally:
MOE_A2A_BACKEND = original_backend moe.a2a_backend = original_backend
DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = ( moe.disable_fp4_allgather = original_disable_fp4_allgather
original_disable_flashinfer_cutlass_moe_fp4_allgather
)
# The type of method in top-K routing, for use in torch custom op # The type of method in top-K routing, for use in torch custom op
+39 -1
View File
@@ -270,6 +270,42 @@ class CaptureFlags(_FlagGroupBase):
enable_torch_compile: bool = False 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 @dataclasses.dataclass
class Flags(_FlagGroupBase): class Flags(_FlagGroupBase):
"""Root of the runtime-flags tier. """Root of the runtime-flags tier.
@@ -277,10 +313,12 @@ class Flags(_FlagGroupBase):
Resolved configuration lives on ``server_args`` fields (materialized at Resolved configuration lives on ``server_args`` fields (materialized at
the end of ``__post_init__``) — this tier only carries genuine runtime the end of ``__post_init__``) — this tier only carries genuine runtime
state whose value is not a function of the configuration alone, grouped 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) capture: CaptureFlags = dataclasses.field(default_factory=CaptureFlags)
moe: MoeFlags = dataclasses.field(default_factory=MoeFlags)
dp: DpFlags = dataclasses.field(default_factory=DpFlags)
class RuntimeContext: class RuntimeContext:
@@ -420,9 +420,9 @@ class TestAiterAllreduceFusionGate(CustomTestCase):
stack.enter_context( stack.enter_context(
mock.patch.object(comm, "get_global_server_args", lambda: server_args) mock.patch.object(comm, "get_global_server_args", lambda: server_args)
) )
stack.enter_context( from sglang.srt.runtime_context import get_flags
mock.patch.object(comm, "is_dp_attention_enabled", lambda: dp_attention)
) stack.enter_context(get_flags().dp.override(enabled=dp_attention))
stack.enter_context( stack.enter_context(
mock.patch.object(comm, "get_moe_a2a_backend", lambda: a2a_backend) mock.patch.object(comm, "get_moe_a2a_backend", lambda: a2a_backend)
) )
@@ -38,7 +38,9 @@ def _mock_global_server_args(backend="pytorch"):
device_group = None device_group = None
sampler_mod.get_tp_group = lambda: _DummyTPGroup() sampler_mod.get_tp_group = lambda: _DummyTPGroup()
sampler_mod.is_dp_attention_enabled = lambda: False from sglang.srt.runtime_context import get_flags
get_flags().dp.enabled = False
def _make_sampling_info(batch_size, vocab_size, device="cuda"): def _make_sampling_info(batch_size, vocab_size, device="cuda"):
@@ -0,0 +1,72 @@
"""Ratchet guard: module-level runtime state in the flag-owning layers may
only shrink.
Runtime flags live on ``get_flags()`` groups (``moe`` / ``dp`` / ``capture``),
where they get lifecycle reset, typo-safe writes, and the transactional
test-override primitive. A new module-level global written through a
``global`` statement in these modules recreates the pattern this replaced:
state with ad-hoc lifecycle that leaks across unit-test teardowns and cannot
be overridden scoped.
The pin lists the survivors by name: the DP-attention topology values (owned
by the parallel vertical) and the TBO comm stream (a resource, owned by the
resources vertical). Migrating one of them must shrink its pin; adding a name
fails the ratchet.
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import ast
import unittest
from pathlib import Path
import sglang.srt
from sglang.test.test_utils import CustomTestCase
_SRT_ROOT = Path(next(iter(sglang.srt.__path__)))
_PINNED_GLOBALS = {
"layers/moe/utils.py": frozenset(),
"layers/dp_attention.py": frozenset(
{
# DP-attention topology (parallel vertical scope).
"_ATTN_DP_RANK",
"_ATTN_DP_SIZE",
"_LOCAL_ATTN_DP_SIZE",
"_LOCAL_ATTN_DP_RANK",
# Comm stream resource (resources vertical scope).
"_DP_TBO_COMM_STREAM",
}
),
}
class TestModuleStateRatchet(CustomTestCase):
def test_global_statements_match_the_pins(self):
for rel, pinned in _PINNED_GLOBALS.items():
tree = ast.parse((_SRT_ROOT / rel).read_text())
declared = {
name
for node in ast.walk(tree)
if isinstance(node, ast.Global)
for name in node.names
}
grown = declared - pinned
self.assertFalse(
grown,
f"{rel} declares new module-level runtime state {sorted(grown)}; "
"put runtime flags on a get_flags() group instead "
"(see runtime_context.MoeFlags / DpFlags).",
)
shrunk = pinned - declared
self.assertFalse(
shrunk,
f"{rel} no longer declares {sorted(shrunk)}; "
"shrink the pin in this file to lock in the progress.",
)
if __name__ == "__main__":
unittest.main()
@@ -265,6 +265,108 @@ class _FakeResolvedArgs:
_resolved_overrides: list = dataclasses.field(default_factory=list) _resolved_overrides: list = dataclasses.field(default_factory=list)
class TestMoeFlagsGroup(_IsolatedServerArgs):
"""flags.moe: materialized by initialize_moe_config; the ACTIVE backends
swap under the speculative contexts and restore on exit."""
def _init(self, **kw):
from types import SimpleNamespace
from sglang.srt.layers.moe.utils import initialize_moe_config
defaults = dict(
moe_a2a_backend="none",
moe_runner_backend="auto",
speculative_moe_runner_backend=None,
speculative_moe_a2a_backend=None,
deepep_mode="auto",
deepep_config=None,
enable_two_batch_overlap=False,
enable_single_batch_overlap=False,
tbo_token_distribution_threshold=0.48,
disable_flashinfer_cutlass_moe_fp4_allgather=False,
quantization=None,
)
defaults.update(kw)
initialize_moe_config(SimpleNamespace(**defaults))
def test_lazy_defaults_before_initialize(self):
from sglang.srt.layers.moe.utils import (
get_moe_a2a_backend,
get_moe_runner_backend,
is_tbo_enabled,
)
reset_context()
self.assertTrue(get_moe_a2a_backend().is_none())
self.assertEqual(get_moe_runner_backend().name, "AUTO")
self.assertFalse(is_tbo_enabled())
def test_initialize_materializes_group(self):
from sglang.srt.layers.moe.utils import get_moe_a2a_backend, is_tbo_enabled
self._init(moe_a2a_backend="deepep", enable_two_batch_overlap=True)
self.assertTrue(get_moe_a2a_backend().is_deepep())
self.assertTrue(is_tbo_enabled())
self.assertEqual(get_flags().moe.deepep_config, "")
def test_speculative_swap_and_restore(self):
from sglang.srt.layers.moe.utils import (
get_moe_a2a_backend,
get_moe_runner_backend,
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
self._init(
moe_a2a_backend="deepep",
moe_runner_backend="triton",
speculative_moe_runner_backend="auto",
speculative_moe_a2a_backend="none",
)
with speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
self.assertEqual(get_moe_runner_backend().name, "AUTO")
self.assertTrue(get_moe_a2a_backend().is_none())
# MTP layers are unquantized: fp4 allgather is forced off
self.assertTrue(get_flags().moe.disable_fp4_allgather)
self.assertEqual(get_moe_runner_backend().name, "TRITON")
self.assertTrue(get_moe_a2a_backend().is_deepep())
self.assertFalse(get_flags().moe.disable_fp4_allgather)
def test_swap_restores_on_exception(self):
from sglang.srt.layers.moe.utils import (
get_moe_runner_backend,
speculative_moe_backend_context,
)
self._init(moe_runner_backend="triton", speculative_moe_runner_backend="auto")
with self.assertRaises(RuntimeError):
with speculative_moe_backend_context():
raise RuntimeError("boom")
self.assertEqual(get_moe_runner_backend().name, "TRITON")
class TestDpFlagsGroup(_IsolatedServerArgs):
"""flags.dp: the DP-attention runtime flags; is_dp_attention_enabled is a
thin shim over the group leaf."""
def test_shim_reads_the_leaf(self):
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
reset_context()
self.assertFalse(is_dp_attention_enabled())
get_flags().dp.enabled = True
self.assertTrue(is_dp_attention_enabled())
def test_scoped_override_forces_the_predicate(self):
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
reset_context()
with get_flags().dp.override(enabled=True):
self.assertTrue(is_dp_attention_enabled())
self.assertFalse(is_dp_attention_enabled())
class TestPublishLifecycle(_IsolatedServerArgs): class TestPublishLifecycle(_IsolatedServerArgs):
"""Publish installs the resolved server_args and seeds the capture tier.""" """Publish installs the resolved server_args and seeds the capture tier."""