B300 NVFP4 port P1 stage-2: C2 flashinfer_cutedsl MoE runner + C3 flashinfer comm-fusion
Complete the flashinfer port ("port全"): ADD flashinfer_cutedsl.py (C2, deepep-paired NVFP4 MoE
runner; selectable via --moe-runner-backend flashinfer_cutedsl, already a MoeRunnerBackend enum) and
REPLACE flashinfer_comm_fusion.py to HEAD (C3, allreduce fusion). Both pull self-contained infra
modules target lacked: ADD model_executor/cuda_graph_config.py (cuda_graph_fully_disabled, 0 sglang
top-level imports) and runtime_context.py (get_parallel, lazy imports). comm_fusion consumers
(communicator.py is_flashinfer_allreduce_unavailable, layernorm.py flashinfer_allreduce_residual_rmsnorm)
verified present in HEAD.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
226
python/sglang/srt/runtime_context.py
Normal file
226
python/sglang/srt/runtime_context.py
Normal file
@@ -0,0 +1,226 @@
|
||||
# Copyright 2023-2026 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""A single structured accessor for process-static parallel-topology state.
|
||||
|
||||
``get_parallel()`` returns a ``ParallelContext`` whose attributes — tp / pp /
|
||||
moe / attn size and rank, plus the process-group handles — each delegate live to
|
||||
the canonical getter in ``distributed.parallel_state`` / ``layers.dp_attention``.
|
||||
Returned values are exactly what those getters return; this is a read-through
|
||||
wrapper, not a cache. It gives call-sites one import and one naming scheme in
|
||||
place of a dozen free functions, plus a test-only ``override()`` hook to force a
|
||||
topology without monkeypatching the underlying getters.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import Any
|
||||
|
||||
|
||||
# Imported lazily so this module has no import-time dependencies: any module can
|
||||
# import get_parallel at module level without risking an import cycle.
|
||||
def _ps():
|
||||
from sglang.srt.distributed import parallel_state
|
||||
|
||||
return parallel_state
|
||||
|
||||
|
||||
def _dp():
|
||||
from sglang.srt.layers import dp_attention
|
||||
|
||||
return dp_attention
|
||||
|
||||
|
||||
_PARALLEL_FIELDS = frozenset(
|
||||
{
|
||||
"world_size",
|
||||
"world_rank",
|
||||
"tp_size",
|
||||
"tp_rank",
|
||||
"pp_size",
|
||||
"pp_rank",
|
||||
"moe_ep_size",
|
||||
"moe_ep_rank",
|
||||
"moe_dp_size",
|
||||
"moe_dp_rank",
|
||||
"moe_tp_size",
|
||||
"moe_tp_rank",
|
||||
"attn_tp_size",
|
||||
"attn_tp_rank",
|
||||
"attn_cp_size",
|
||||
"attn_cp_rank",
|
||||
"attn_dp_size",
|
||||
"attn_dp_rank",
|
||||
"world_group",
|
||||
"tp_group",
|
||||
"pp_group",
|
||||
"moe_ep_group",
|
||||
"moe_dp_group",
|
||||
"moe_tp_group",
|
||||
"attn_tp_group",
|
||||
"attn_cp_group",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class ParallelContext:
|
||||
"""Parallel-topology namespace; the only instance state is ``_overrides``."""
|
||||
|
||||
__slots__ = ("_overrides",)
|
||||
|
||||
def __init__(self):
|
||||
self._overrides = {}
|
||||
|
||||
def _v(self, name, getter):
|
||||
overrides = self._overrides
|
||||
return overrides[name] if name in overrides else getter()
|
||||
|
||||
@contextmanager
|
||||
def override(self, **kwargs):
|
||||
"""Temporarily force parallel values, restoring on exit. Validates keys and
|
||||
supports nesting."""
|
||||
unknown = set(kwargs) - _PARALLEL_FIELDS
|
||||
if unknown:
|
||||
raise ValueError(f"unknown parallel field(s): {sorted(unknown)}")
|
||||
saved = dict(self._overrides)
|
||||
self._overrides.update(kwargs)
|
||||
try:
|
||||
yield self
|
||||
finally:
|
||||
self._overrides = saved
|
||||
|
||||
@property
|
||||
def world_size(self) -> int:
|
||||
return self._v("world_size", _ps().get_world_size)
|
||||
|
||||
@property
|
||||
def world_rank(self) -> int:
|
||||
return self._v("world_rank", _ps().get_world_rank)
|
||||
|
||||
@property
|
||||
def tp_size(self) -> int:
|
||||
return self._v("tp_size", _ps().get_tensor_model_parallel_world_size)
|
||||
|
||||
@property
|
||||
def tp_rank(self) -> int:
|
||||
return self._v("tp_rank", _ps().get_tensor_model_parallel_rank)
|
||||
|
||||
@property
|
||||
def pp_size(self) -> int:
|
||||
return self._v("pp_size", _ps().get_pipeline_model_parallel_world_size)
|
||||
|
||||
@property
|
||||
def pp_rank(self) -> int:
|
||||
return self._v("pp_rank", _ps().get_pipeline_model_parallel_rank)
|
||||
|
||||
@property
|
||||
def moe_ep_size(self) -> int:
|
||||
return self._v("moe_ep_size", _ps().get_moe_expert_parallel_world_size)
|
||||
|
||||
@property
|
||||
def moe_ep_rank(self) -> int:
|
||||
return self._v("moe_ep_rank", _ps().get_moe_expert_parallel_rank)
|
||||
|
||||
@property
|
||||
def moe_dp_size(self) -> int:
|
||||
return self._v("moe_dp_size", _ps().get_moe_data_parallel_world_size)
|
||||
|
||||
@property
|
||||
def moe_dp_rank(self) -> int:
|
||||
return self._v("moe_dp_rank", _ps().get_moe_data_parallel_rank)
|
||||
|
||||
@property
|
||||
def moe_tp_size(self) -> int:
|
||||
return self._v("moe_tp_size", _ps().get_moe_tensor_parallel_world_size)
|
||||
|
||||
@property
|
||||
def moe_tp_rank(self) -> int:
|
||||
return self._v("moe_tp_rank", _ps().get_moe_tensor_parallel_rank)
|
||||
|
||||
@property
|
||||
def attn_tp_size(self) -> int:
|
||||
return self._v("attn_tp_size", _ps().get_attn_tensor_model_parallel_world_size)
|
||||
|
||||
@property
|
||||
def attn_tp_rank(self) -> int:
|
||||
return self._v("attn_tp_rank", _ps().get_attn_tensor_model_parallel_rank)
|
||||
|
||||
@property
|
||||
def attn_cp_size(self) -> int:
|
||||
return self._v("attn_cp_size", _ps().get_attn_context_model_parallel_world_size)
|
||||
|
||||
@property
|
||||
def attn_cp_rank(self) -> int:
|
||||
return self._v("attn_cp_rank", _ps().get_attn_context_model_parallel_rank)
|
||||
|
||||
@property
|
||||
def attn_dp_size(self) -> int:
|
||||
return self._v("attn_dp_size", _dp().get_attention_dp_size)
|
||||
|
||||
@property
|
||||
def attn_dp_rank(self) -> int:
|
||||
return self._v("attn_dp_rank", _dp().get_attention_dp_rank)
|
||||
|
||||
@property
|
||||
def world_group(self) -> Any:
|
||||
return self._v("world_group", _ps().get_world_group)
|
||||
|
||||
@property
|
||||
def tp_group(self) -> Any:
|
||||
return self._v("tp_group", _ps().get_tp_group)
|
||||
|
||||
@property
|
||||
def pp_group(self) -> Any:
|
||||
return self._v("pp_group", _ps().get_pp_group)
|
||||
|
||||
@property
|
||||
def moe_ep_group(self) -> Any:
|
||||
return self._v("moe_ep_group", _ps().get_moe_ep_group)
|
||||
|
||||
@property
|
||||
def moe_dp_group(self) -> Any:
|
||||
return self._v("moe_dp_group", _ps().get_moe_dp_group)
|
||||
|
||||
@property
|
||||
def moe_tp_group(self) -> Any:
|
||||
return self._v("moe_tp_group", _ps().get_moe_tp_group)
|
||||
|
||||
@property
|
||||
def attn_tp_group(self) -> Any:
|
||||
return self._v("attn_tp_group", _ps().get_attn_tp_group)
|
||||
|
||||
@property
|
||||
def attn_cp_group(self) -> Any:
|
||||
return self._v("attn_cp_group", _ps().get_attn_cp_group)
|
||||
|
||||
|
||||
class RuntimeContext:
|
||||
"""Container for the structured runtime accessors; currently exposes ``parallel``."""
|
||||
|
||||
__slots__ = ("parallel",)
|
||||
|
||||
def __init__(self, parallel: ParallelContext):
|
||||
self.parallel = parallel
|
||||
|
||||
|
||||
_PARALLEL = ParallelContext()
|
||||
_CONTEXT = RuntimeContext(parallel=_PARALLEL)
|
||||
|
||||
|
||||
def get_context() -> RuntimeContext:
|
||||
return _CONTEXT
|
||||
|
||||
|
||||
def get_parallel() -> ParallelContext:
|
||||
return _PARALLEL
|
||||
Reference in New Issue
Block a user