[refactor] Collect MoE and DP-attention runtime state into typed flag groups (#30347)
This commit is contained in:
Cheng Wan
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -420,9 +420,9 @@ class TestAiterAllreduceFusionGate(CustomTestCase):
|
||||
stack.enter_context(
|
||||
mock.patch.object(comm, "get_global_server_args", lambda: server_args)
|
||||
)
|
||||
stack.enter_context(
|
||||
mock.patch.object(comm, "is_dp_attention_enabled", lambda: dp_attention)
|
||||
)
|
||||
from sglang.srt.runtime_context import get_flags
|
||||
|
||||
stack.enter_context(get_flags().dp.override(enabled=dp_attention))
|
||||
stack.enter_context(
|
||||
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
|
||||
|
||||
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"):
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
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):
|
||||
"""Publish installs the resolved server_args and seeds the capture tier."""
|
||||
|
||||
|
||||
Reference in New Issue
Block a user