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:
@@ -1,93 +1,81 @@
|
||||
import contextlib
|
||||
import inspect
|
||||
import logging
|
||||
import platform
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch.distributed import ProcessGroup
|
||||
|
||||
from sglang.srt.distributed import (
|
||||
get_attn_tensor_model_parallel_rank,
|
||||
get_attn_tensor_model_parallel_world_size,
|
||||
get_moe_expert_parallel_rank,
|
||||
get_moe_expert_parallel_world_size,
|
||||
get_moe_tensor_parallel_rank,
|
||||
get_moe_tensor_parallel_world_size,
|
||||
get_attn_tp_group,
|
||||
get_moe_ep_group,
|
||||
get_moe_tp_group,
|
||||
get_tp_group,
|
||||
)
|
||||
from sglang.srt.distributed.parallel_state import in_the_same_node_as
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.utils import (
|
||||
ceil_align,
|
||||
get_cuda_driver_bindings,
|
||||
is_flashinfer_available,
|
||||
is_sm90_supported,
|
||||
is_sm100_supported,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import is_flashinfer_available
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# FlashInfer allreduce fusion: set when flashinfer is available (see block below)
|
||||
_flashinfer_comm = None
|
||||
_workspace_manager = None
|
||||
_TorchDistBackend = None
|
||||
_mnnvl_comm_backend = None
|
||||
_create_allreduce_fusion_workspace = None
|
||||
_flashinfer_allreduce_unavailable = False
|
||||
_posix_transport_override_logged = False
|
||||
_flashinfer_create_workspace_supports_group = False
|
||||
_flashinfer_create_workspace_supports_comm_backend = False
|
||||
_flashinfer_allreduce_supports_trigger_completion = False
|
||||
_mnnvl_non_blackwell_fallback_logged = False
|
||||
|
||||
|
||||
def _should_force_posix_fd_transport() -> bool:
|
||||
force_posix_env = envs.SGLANG_FLASHINFER_FORCE_POSIX_FD_TRANSPORT.get()
|
||||
if force_posix_env is not None:
|
||||
return force_posix_env
|
||||
def _mnnvl_supported(is_multi_node: bool) -> bool:
|
||||
"""Whether the mnnvl backend is usable on the current system.
|
||||
|
||||
machine = platform.machine().lower()
|
||||
if machine not in ("aarch64", "arm64"):
|
||||
return False
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
return False
|
||||
|
||||
try:
|
||||
major, _minor = torch.cuda.get_device_capability(torch.cuda.current_device())
|
||||
except Exception as e:
|
||||
logger.debug("Failed to get CUDA device capability: %s", e)
|
||||
return False
|
||||
|
||||
return major == 10
|
||||
mnnvl runs on Blackwell (SM10x) for both single- and multi-node, and on
|
||||
SM90 for single-node only. Multi-node mnnvl on non-Blackwell is not
|
||||
supported and must fall back to trtllm.
|
||||
"""
|
||||
if is_sm100_supported():
|
||||
return True
|
||||
return is_sm90_supported() and not is_multi_node
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _flashinfer_posix_fd_transport_override_if_needed():
|
||||
# TODO(mmangkad): Remove this temporary override once the
|
||||
# FlashInfer unified allreduce-fusion transport issue on
|
||||
# GB200/GB300 platforms is fixed and verified resolved.
|
||||
global _posix_transport_override_logged
|
||||
def _resolve_backend(backend: str, is_multi_node: bool = False) -> str:
|
||||
"""Resolve the requested FlashInfer allreduce fusion backend."""
|
||||
global _mnnvl_non_blackwell_fallback_logged
|
||||
|
||||
if not _should_force_posix_fd_transport():
|
||||
yield
|
||||
return
|
||||
if backend == "auto":
|
||||
# Prefer mnnvl wherever it is supported (any Blackwell system, or SM90
|
||||
# single-node); fall back to trtllm otherwise.
|
||||
return "mnnvl" if _mnnvl_supported(is_multi_node) else "trtllm"
|
||||
|
||||
try:
|
||||
import flashinfer.comm.mnnvl as flashinfer_mnnvl
|
||||
except Exception as e:
|
||||
logger.debug(
|
||||
"Failed to import flashinfer.comm.mnnvl for transport override: %s", e
|
||||
)
|
||||
yield
|
||||
return
|
||||
if backend == "mnnvl" and not _mnnvl_supported(is_multi_node):
|
||||
if not _mnnvl_non_blackwell_fallback_logged:
|
||||
logger.info(
|
||||
"FlashInfer allreduce fusion: forcing trtllm backend "
|
||||
"(mnnvl requires a Blackwell system, or SM90 single-node)."
|
||||
)
|
||||
_mnnvl_non_blackwell_fallback_logged = True
|
||||
return "trtllm"
|
||||
return backend
|
||||
|
||||
original_checker = getattr(flashinfer_mnnvl, "is_mnnvl_fabric_supported", None)
|
||||
if original_checker is None:
|
||||
yield
|
||||
return
|
||||
|
||||
if not _posix_transport_override_logged:
|
||||
logger.warning(
|
||||
"Applying FlashInfer transport workaround: forcing PosixFD "
|
||||
"symmetric-memory handle exchange on aarch64 + sm10x to avoid "
|
||||
"known data corruption with Fabric handle exchange on GB systems. "
|
||||
"Set SGLANG_FLASHINFER_FORCE_POSIX_FD_TRANSPORT=0 to disable."
|
||||
)
|
||||
_posix_transport_override_logged = True
|
||||
|
||||
def _always_disable_fabric(_device_idx: int) -> bool:
|
||||
return False
|
||||
|
||||
flashinfer_mnnvl.is_mnnvl_fabric_supported = _always_disable_fabric
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
flashinfer_mnnvl.is_mnnvl_fabric_supported = original_checker
|
||||
def resolve_flashinfer_allreduce_fusion_backend(server_args) -> Optional[str]:
|
||||
backend = getattr(server_args, "flashinfer_allreduce_fusion_backend", None)
|
||||
if backend is None:
|
||||
return None
|
||||
is_multi_node = getattr(server_args, "nnodes", 1) > 1
|
||||
return _resolve_backend(backend, is_multi_node)
|
||||
|
||||
|
||||
if is_flashinfer_available():
|
||||
@@ -98,33 +86,309 @@ if is_flashinfer_available():
|
||||
comm, "create_allreduce_fusion_workspace"
|
||||
):
|
||||
_flashinfer_comm = comm
|
||||
_create_allreduce_fusion_workspace = comm.create_allreduce_fusion_workspace
|
||||
workspace_params = inspect.signature(
|
||||
comm.create_allreduce_fusion_workspace
|
||||
).parameters
|
||||
allreduce_params = inspect.signature(comm.allreduce_fusion).parameters
|
||||
_flashinfer_create_workspace_supports_group = "group" in workspace_params
|
||||
_flashinfer_create_workspace_supports_comm_backend = (
|
||||
"comm_backend" in workspace_params
|
||||
)
|
||||
_flashinfer_allreduce_supports_trigger_completion = (
|
||||
"trigger_completion_at_end" in allreduce_params
|
||||
)
|
||||
else:
|
||||
_flashinfer_allreduce_unavailable = True
|
||||
logger.warning(
|
||||
"flashinfer.comm unified allreduce_fusion API is not available, "
|
||||
"falling back to standard implementation"
|
||||
)
|
||||
except ImportError:
|
||||
except (ImportError, AttributeError) as e:
|
||||
_flashinfer_allreduce_unavailable = True
|
||||
logger.warning(
|
||||
"flashinfer.comm is not available, falling back to standard "
|
||||
"implementation"
|
||||
"flashinfer.comm allreduce_fusion API is not available (%s), "
|
||||
"falling back to standard implementation",
|
||||
e,
|
||||
)
|
||||
|
||||
try:
|
||||
from flashinfer.comm.mnnvl import TorchDistBackend
|
||||
|
||||
class _FixedTorchDistBackend(TorchDistBackend):
|
||||
"""Workaround for FlashInfer TorchDistBackend issues.
|
||||
|
||||
1. bcast fix: TorchDistBackend.bcast passes the in-group rank
|
||||
directly as `src` to broadcast_object_list, which expects a
|
||||
global rank.
|
||||
2. Graph-capture fix: initialize with NCCL device_group (so
|
||||
the backend derives correct device_idx / GPU mapping), but
|
||||
broadcast via GLOO cpu_group (to avoid NCCL collectives
|
||||
that interfere with CUDA graph capture).
|
||||
"""
|
||||
|
||||
def __init__(self, device_group, cpu_group):
|
||||
super().__init__(group=device_group)
|
||||
self._cpu_group = cpu_group
|
||||
|
||||
def bcast(self, data, root):
|
||||
import torch.distributed as dist
|
||||
|
||||
group_ranks = dist.get_process_group_ranks(self._cpu_group)
|
||||
global_root = group_ranks[root]
|
||||
object_list = [data]
|
||||
dist.broadcast_object_list(
|
||||
object_list, src=global_root, group=self._cpu_group
|
||||
)
|
||||
return object_list[0]
|
||||
|
||||
_TorchDistBackend = _FixedTorchDistBackend
|
||||
except ImportError:
|
||||
logger.debug(
|
||||
"flashinfer.comm.mnnvl.TorchDistBackend is not available, "
|
||||
"allreduce fusion will use the default process group"
|
||||
)
|
||||
|
||||
try:
|
||||
from flashinfer.comm.mnnvl import CommBackend
|
||||
|
||||
class TorchDistributedCommBackend(CommBackend):
|
||||
"""
|
||||
Use torch distributed instead of MPI to set up flashinfer MNNVL
|
||||
workspaces during initialization.
|
||||
"""
|
||||
|
||||
def __init__(self, group: ProcessGroup):
|
||||
self._group = group
|
||||
|
||||
def Get_rank(self) -> int:
|
||||
return self._group.rank()
|
||||
|
||||
def Get_size(self) -> int:
|
||||
return self._group.size()
|
||||
|
||||
def allgather(self, data: int):
|
||||
gathered = [None] * self.Get_size()
|
||||
dist.all_gather_object(gathered, data, group=self._group)
|
||||
return gathered
|
||||
|
||||
def bcast(self, data, root: int = 0):
|
||||
"""Broadcast a picklable Python object from root to all ranks."""
|
||||
obj_list = [data]
|
||||
dist.broadcast_object_list(obj_list, src=root, group=self._group)
|
||||
return obj_list[0]
|
||||
|
||||
def barrier(self):
|
||||
dist.barrier(group=self._group)
|
||||
|
||||
def Split(self, color: int, key: int):
|
||||
# No need to split; we already use the proper group.
|
||||
return self._group
|
||||
|
||||
_mnnvl_comm_backend = TorchDistributedCommBackend
|
||||
except ImportError:
|
||||
_mnnvl_comm_backend = None
|
||||
|
||||
|
||||
# FlashInfer allreduce fusion backend support matrix for
|
||||
# --flashinfer-allreduce-fusion-backend:
|
||||
#
|
||||
# Backend | SM103 | SM100 | SM90 | Single-Node | Multi-Node |
|
||||
# --------- | ----- | ----- | ----------- | ----------- | ---------- |
|
||||
# trtllm | Yes | Yes | Yes | Yes | No |
|
||||
# mnnvl | Yes | Yes | Single-node | Yes | Blackwell |
|
||||
#
|
||||
# mnnvl runs on any Blackwell GPU (SM10x) for both single- and multi-node, and
|
||||
# on SM90 for single-node only. auto resolves to mnnvl wherever it is supported
|
||||
# and to trtllm otherwise. An explicit mnnvl request on an unsupported
|
||||
# configuration (e.g. SM90 multi-node) falls back to trtllm (see
|
||||
# _resolve_backend).
|
||||
|
||||
|
||||
def is_flashinfer_allreduce_unavailable() -> bool:
|
||||
return _flashinfer_allreduce_unavailable
|
||||
|
||||
|
||||
def _make_flashinfer_workspace_allocation_prop(cuda_driver):
|
||||
from flashinfer.comm.mnnvl import is_mnnvl_fabric_supported
|
||||
|
||||
handle_types = cuda_driver.CUmemAllocationHandleType
|
||||
if is_mnnvl_fabric_supported(torch.cuda.current_device()):
|
||||
handle_type = handle_types.CU_MEM_HANDLE_TYPE_FABRIC
|
||||
else:
|
||||
handle_type = handle_types.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR
|
||||
|
||||
prop = cuda_driver.CUmemAllocationProp()
|
||||
prop.requestedHandleTypes = handle_type
|
||||
prop.type = cuda_driver.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
|
||||
prop.location = cuda_driver.CUmemLocation()
|
||||
prop.location.type = cuda_driver.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
|
||||
prop.location.id = torch.cuda.current_device()
|
||||
prop.allocFlags.gpuDirectRDMACapable = 1
|
||||
return prop
|
||||
|
||||
|
||||
def _flashinfer_trtllm_workspace_allocation_sizes(
|
||||
cuda_driver,
|
||||
prop,
|
||||
world_size: int,
|
||||
max_token_num: int,
|
||||
hidden_dim: int,
|
||||
dtype: torch.dtype,
|
||||
) -> list[int]:
|
||||
"""Mirror FlashInfer TRTLLM SymmDeviceMemory local allocation sizes."""
|
||||
elem_size = 4 if dtype == torch.float32 else 2
|
||||
buffer_size = world_size * max_token_num * hidden_dim * 2
|
||||
flag_size = world_size * 256 * 4
|
||||
|
||||
max_comm_size = 2147483647 & ~((1 << 21) - 1)
|
||||
lamport_comm_size = min(
|
||||
world_size * max_token_num * hidden_dim * elem_size,
|
||||
max_comm_size,
|
||||
)
|
||||
lamport_buffer_size = lamport_comm_size * 3
|
||||
|
||||
# trtllm_create_ipc_workspace_for_all_reduce_fusion rounds each logical
|
||||
# buffer to 2 MiB before passing it to SymmDeviceMemory.
|
||||
buffer_sizes = (
|
||||
ceil_align(size, 1 << 21)
|
||||
for size in (buffer_size, flag_size, lamport_buffer_size)
|
||||
)
|
||||
|
||||
signal_pad_size = 2048
|
||||
allocation_sizes = []
|
||||
for buffer_size in buffer_sizes:
|
||||
err, alloc_granularity = cuda_driver.cuMemGetAllocationGranularity(
|
||||
prop,
|
||||
cuda_driver.CUmemAllocationGranularity_flags.CU_MEM_ALLOC_GRANULARITY_RECOMMENDED,
|
||||
)
|
||||
if err != cuda_driver.CUresult.CUDA_SUCCESS:
|
||||
raise RuntimeError(
|
||||
"cuMemGetAllocationGranularity failed for FlashInfer "
|
||||
f"workspace preflight: {err}"
|
||||
)
|
||||
|
||||
allocation_size = ceil_align(buffer_size + signal_pad_size, alloc_granularity)
|
||||
|
||||
mc_prop = cuda_driver.CUmulticastObjectProp()
|
||||
mc_prop.numDevices = world_size
|
||||
mc_prop.size = allocation_size
|
||||
mc_prop.handleTypes = prop.requestedHandleTypes
|
||||
|
||||
err, mc_granularity = cuda_driver.cuMulticastGetGranularity(
|
||||
mc_prop,
|
||||
cuda_driver.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED,
|
||||
)
|
||||
if err != cuda_driver.CUresult.CUDA_SUCCESS:
|
||||
raise RuntimeError(
|
||||
"cuMulticastGetGranularity failed for FlashInfer "
|
||||
f"workspace preflight: {err}"
|
||||
)
|
||||
|
||||
allocation_size = ceil_align(allocation_size, mc_granularity)
|
||||
allocation_sizes.append(allocation_size)
|
||||
return allocation_sizes
|
||||
|
||||
|
||||
def _probe_cumem_create_sequence(cuda_driver, allocation_sizes, prop) -> bool:
|
||||
handles = []
|
||||
try:
|
||||
for allocation_size in allocation_sizes:
|
||||
err, handle = cuda_driver.cuMemCreate(allocation_size, prop, 0)
|
||||
if err != cuda_driver.CUresult.CUDA_SUCCESS:
|
||||
return False
|
||||
handles.append(handle)
|
||||
return True
|
||||
finally:
|
||||
for handle in reversed(handles):
|
||||
cuda_driver.cuMemRelease(handle)
|
||||
|
||||
|
||||
def _preflight_check_workspace_memory(
|
||||
world_size: int,
|
||||
max_token_num: int,
|
||||
hidden_dim: int,
|
||||
dtype: torch.dtype,
|
||||
cpu_group: Optional["torch.distributed.ProcessGroup"] = None,
|
||||
) -> bool:
|
||||
"""Collectively decide whether to enter FlashInfer workspace creation.
|
||||
|
||||
FlashInfer TRTLLM workspaces allocate several SymmDeviceMemory buffers and
|
||||
then exchange handles across ranks. If one rank fails local cuMemCreate and
|
||||
exits while peers enter handle exchange, peers can hang until the watchdog
|
||||
aborts. Probe the same handle type and allocation sequence first, then vote
|
||||
on a CPU group so all ranks proceed or skip together.
|
||||
"""
|
||||
import torch.distributed as dist
|
||||
|
||||
group = cpu_group
|
||||
if group is None:
|
||||
tp_group = get_tp_group()
|
||||
if tp_group.world_size <= 1:
|
||||
return True
|
||||
group = tp_group.cpu_group
|
||||
|
||||
allocation_sizes = []
|
||||
try:
|
||||
cuda_driver = get_cuda_driver_bindings()
|
||||
prop = _make_flashinfer_workspace_allocation_prop(cuda_driver)
|
||||
allocation_sizes = _flashinfer_trtllm_workspace_allocation_sizes(
|
||||
cuda_driver,
|
||||
prop,
|
||||
world_size,
|
||||
max_token_num,
|
||||
hidden_dim,
|
||||
dtype,
|
||||
)
|
||||
local_ok = _probe_cumem_create_sequence(cuda_driver, allocation_sizes, prop)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"FlashInfer workspace preflight probe failed (%s). "
|
||||
"Skipping allreduce fusion.",
|
||||
e,
|
||||
)
|
||||
local_ok = False
|
||||
|
||||
flag = torch.tensor([1 if local_ok else 0], dtype=torch.int32)
|
||||
dist.all_reduce(flag, op=dist.ReduceOp.BAND, group=group)
|
||||
|
||||
logger.debug(
|
||||
"FlashInfer workspace preflight [rank %s]: probe=%.2f GB, "
|
||||
"local_probe=%s, vote=%s",
|
||||
dist.get_rank(group=group),
|
||||
sum(allocation_sizes) / 1e9,
|
||||
"OK" if local_ok else "FAIL",
|
||||
"PROCEED" if flag.item() == 1 else "SKIP",
|
||||
)
|
||||
if flag.item() == 0:
|
||||
logger.warning(
|
||||
"FlashInfer workspace preflight: cuMemCreate probe failed on at "
|
||||
"least one rank. Skipping allreduce fusion to avoid cross-rank "
|
||||
"desync inside the flashinfer collective."
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class FlashInferWorkspaceManager:
|
||||
"""
|
||||
Manages FlashInfer's unified allreduce workspace.
|
||||
Supports trtllm and mnnvl backends via create_allreduce_fusion_workspace().
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self.workspace = None
|
||||
self.world_size = None
|
||||
self.rank = None
|
||||
self.group = None
|
||||
self.max_token_num = None
|
||||
self.hidden_dim = None
|
||||
self.dtype = None
|
||||
self.initialized = False
|
||||
# Track max sizes ever requested so the workspace only grows (fewer recreates)
|
||||
self._max_token_num_seen: Optional[int] = None
|
||||
self._max_hidden_dim_seen: Optional[int] = None
|
||||
self._logged_init = False
|
||||
|
||||
def initialize(
|
||||
self,
|
||||
@@ -132,52 +396,134 @@ class FlashInferWorkspaceManager:
|
||||
rank: int,
|
||||
max_token_num: int,
|
||||
hidden_dim: int,
|
||||
dtype: torch.dtype,
|
||||
backend: str = "auto",
|
||||
group: Optional[ProcessGroup] = None,
|
||||
use_fp32_lamport: bool = False,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
use_oneshot: Optional[bool] = None,
|
||||
device_group: Optional["torch.distributed.ProcessGroup"] = None,
|
||||
cpu_group: Optional["torch.distributed.ProcessGroup"] = None,
|
||||
):
|
||||
"""Initialize workspace"""
|
||||
if _flashinfer_comm is None:
|
||||
"""Initialize workspace using FlashInfer's unified API."""
|
||||
global _flashinfer_allreduce_unavailable
|
||||
|
||||
# Track the high-water mark so allocations only grow
|
||||
self._max_token_num_seen = max(max_token_num, self._max_token_num_seen or 0)
|
||||
self._max_hidden_dim_seen = max(hidden_dim, self._max_hidden_dim_seen or 0)
|
||||
|
||||
# Reuse existing workspace if it already covers this problem size
|
||||
if (
|
||||
self.initialized
|
||||
and self.world_size == world_size
|
||||
and self.is_buffer_size_sufficient(
|
||||
token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype or torch.bfloat16,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
):
|
||||
return
|
||||
|
||||
# Same world_size but buffer too small: free old workspace before creating new
|
||||
if self.initialized and self.world_size == world_size:
|
||||
self.cleanup()
|
||||
|
||||
if _flashinfer_comm is None or _create_allreduce_fusion_workspace is None:
|
||||
logger.warning(
|
||||
"FlashInfer comm not available, skipping workspace initialization"
|
||||
)
|
||||
return
|
||||
|
||||
self.cleanup()
|
||||
|
||||
if not _preflight_check_workspace_memory(
|
||||
world_size=world_size,
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
cpu_group=cpu_group,
|
||||
):
|
||||
_flashinfer_allreduce_unavailable = True
|
||||
self.workspace = None
|
||||
self.initialized = False
|
||||
return
|
||||
|
||||
# Determine GPUs per node for MNNVL topology detection
|
||||
gpus_per_node = None
|
||||
node_pg = cpu_group if cpu_group is not None else group
|
||||
if node_pg is not None:
|
||||
gpus_per_node = sum(in_the_same_node_as(node_pg, source_rank=0))
|
||||
comm_backend = None
|
||||
if (
|
||||
_TorchDistBackend is not None
|
||||
and device_group is not None
|
||||
and cpu_group is not None
|
||||
):
|
||||
comm_backend = _TorchDistBackend(
|
||||
device_group=device_group, cpu_group=cpu_group
|
||||
)
|
||||
elif _mnnvl_comm_backend is not None and group is not None:
|
||||
comm_backend = _mnnvl_comm_backend(group)
|
||||
|
||||
try:
|
||||
with _flashinfer_posix_fd_transport_override_if_needed():
|
||||
self.workspace = _flashinfer_comm.create_allreduce_fusion_workspace(
|
||||
backend="trtllm",
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
force_oneshot_support=bool(use_oneshot),
|
||||
alloc_token_num = max(max_token_num, self._max_token_num_seen or 0)
|
||||
alloc_hidden_dim = max(hidden_dim, self._max_hidden_dim_seen or 0)
|
||||
create_kw = dict(
|
||||
backend=backend,
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=alloc_token_num,
|
||||
hidden_dim=alloc_hidden_dim,
|
||||
dtype=dtype or torch.bfloat16,
|
||||
gpus_per_node=gpus_per_node,
|
||||
)
|
||||
if (
|
||||
_flashinfer_create_workspace_supports_comm_backend
|
||||
and comm_backend is not None
|
||||
):
|
||||
create_kw["comm_backend"] = comm_backend
|
||||
if _flashinfer_create_workspace_supports_group:
|
||||
# Pin the symmetric-memory rendezvous to the actual
|
||||
# subgroup. Without this, flashinfer >=0.6.10 falls back
|
||||
# to WORLD and TP/EP/CP subgroup peers get addressed
|
||||
# incorrectly (kernel hangs in cuda-graph warmup).
|
||||
create_kw["group"] = device_group
|
||||
if use_oneshot is not None:
|
||||
create_kw["force_oneshot_support"] = bool(use_oneshot)
|
||||
if use_fp32_lamport:
|
||||
create_kw["use_fp32_lamport"] = True
|
||||
self.workspace = _create_allreduce_fusion_workspace(**create_kw)
|
||||
self.world_size = world_size
|
||||
self.rank = rank
|
||||
self.group = (device_group, cpu_group)
|
||||
self.max_token_num = alloc_token_num
|
||||
self.hidden_dim = alloc_hidden_dim
|
||||
self.dtype = dtype or torch.bfloat16
|
||||
self.initialized = True
|
||||
|
||||
backend_name = getattr(self.workspace, "backend", "unknown")
|
||||
if not self._logged_init:
|
||||
logger.info(
|
||||
f"FlashInfer AllReduce Fusion enabled and workspace initialized: "
|
||||
f"backend={backend_name}, rank={rank}, world_size={world_size}, "
|
||||
f"max_token_num={self.max_token_num}, hidden_dim={self.hidden_dim}"
|
||||
)
|
||||
self._logged_init = True
|
||||
else:
|
||||
logger.debug(
|
||||
f"FlashInfer workspace re-initialized: backend={backend_name}, "
|
||||
f"rank={rank}, world_size={world_size}"
|
||||
)
|
||||
except Exception as e:
|
||||
global _flashinfer_allreduce_unavailable
|
||||
_flashinfer_allreduce_unavailable = True
|
||||
logger.warning(
|
||||
f"Failed to initialize FlashInfer workspace: {e}. "
|
||||
f"Failed to initialize FlashInfer workspace (backend={backend}): {e}. "
|
||||
"Disabling flashinfer allreduce fusion permanently."
|
||||
)
|
||||
self.workspace = None
|
||||
self.initialized = False
|
||||
return
|
||||
|
||||
self.world_size = world_size
|
||||
self.rank = rank
|
||||
self.max_token_num = max_token_num
|
||||
self.hidden_dim = hidden_dim
|
||||
self.dtype = dtype
|
||||
self.initialized = True
|
||||
|
||||
backend = getattr(self.workspace, "backend", "unknown")
|
||||
logger.info(
|
||||
f"FlashInfer workspace initialized for rank {rank}, "
|
||||
f"world_size {world_size}, backend {backend}"
|
||||
)
|
||||
|
||||
def is_buffer_size_sufficient(
|
||||
self,
|
||||
token_num: int,
|
||||
@@ -197,13 +543,23 @@ class FlashInferWorkspaceManager:
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"FlashInfer workspace size check failed: {e}")
|
||||
# Fallback: some backends may not implement is_buffer_size_sufficient;
|
||||
# reuse if within our allocated dimensions.
|
||||
if (
|
||||
self.max_token_num is not None
|
||||
and self.hidden_dim is not None
|
||||
and token_num <= self.max_token_num
|
||||
and hidden_dim <= self.hidden_dim
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
def cleanup(self):
|
||||
"""Clean up workspace"""
|
||||
"""Clean up workspace."""
|
||||
if self.workspace is not None:
|
||||
try:
|
||||
self.workspace.destroy()
|
||||
if hasattr(self.workspace, "destroy"):
|
||||
self.workspace.destroy()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to cleanup FlashInfer workspace: {e}")
|
||||
finally:
|
||||
@@ -211,23 +567,64 @@ class FlashInferWorkspaceManager:
|
||||
self.initialized = False
|
||||
self.world_size = None
|
||||
self.rank = None
|
||||
self.group = None
|
||||
self.max_token_num = None
|
||||
self.hidden_dim = None
|
||||
self.dtype = None
|
||||
self._logged_init = False
|
||||
|
||||
|
||||
_workspace_manager = FlashInferWorkspaceManager()
|
||||
_attn_tp_workspace_manager = FlashInferWorkspaceManager()
|
||||
_moe_tp_workspace_manager = FlashInferWorkspaceManager()
|
||||
|
||||
|
||||
def _get_workspace_manager(use_attn_tp_group: bool) -> FlashInferWorkspaceManager:
|
||||
return (
|
||||
_attn_tp_workspace_manager if use_attn_tp_group else _moe_tp_workspace_manager
|
||||
)
|
||||
|
||||
|
||||
def _sync_allreduce_unavailable_across_tp():
|
||||
"""Synchronize _flashinfer_allreduce_unavailable across all TP ranks.
|
||||
|
||||
If workspace initialization fails on any rank, all ranks must agree to
|
||||
disable fusion. Otherwise ranks diverge during CUDA graph capture: some
|
||||
use FlashInfer fusion (skipping custom allreduce), others fall back to
|
||||
standard allreduce (calling register_buffer collectives), causing a hang
|
||||
in register_graph_buffers.
|
||||
"""
|
||||
global _flashinfer_allreduce_unavailable
|
||||
try:
|
||||
import torch.distributed as dist
|
||||
|
||||
tp_group = get_tp_group()
|
||||
if tp_group.world_size <= 1:
|
||||
return
|
||||
flag = torch.tensor(
|
||||
[1 if _flashinfer_allreduce_unavailable else 0],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=tp_group.cpu_group)
|
||||
if flag.item() > 0 and not _flashinfer_allreduce_unavailable:
|
||||
_flashinfer_allreduce_unavailable = True
|
||||
logger.warning(
|
||||
"FlashInfer allreduce fusion disabled globally because "
|
||||
"workspace initialization failed on at least one rank."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.debug(f"Failed to sync flashinfer unavailable flag: {e}")
|
||||
|
||||
|
||||
def ensure_workspace_initialized(
|
||||
max_token_num: int = 2048,
|
||||
hidden_dim: int = 4096,
|
||||
dtype: torch.dtype = torch.float16,
|
||||
use_fp32_lamport: bool = False,
|
||||
dtype: Optional[torch.dtype] = None,
|
||||
token_num: Optional[int] = None,
|
||||
use_oneshot: Optional[bool] = None,
|
||||
use_attn_tp_group: bool = True,
|
||||
):
|
||||
"""Ensure workspace is initialized"""
|
||||
"""Ensure workspace is initialized."""
|
||||
if _flashinfer_allreduce_unavailable:
|
||||
return False
|
||||
|
||||
@@ -235,45 +632,67 @@ def ensure_workspace_initialized(
|
||||
return False
|
||||
|
||||
if use_attn_tp_group:
|
||||
world_size = get_attn_tensor_model_parallel_world_size()
|
||||
rank = get_attn_tensor_model_parallel_rank()
|
||||
world_size = get_parallel().attn_tp_size
|
||||
rank = get_parallel().attn_tp_rank
|
||||
coordinator = get_attn_tp_group()
|
||||
else:
|
||||
# If MoE expert parallel world size > 1, use expert parallel group
|
||||
# Otherwise, use tensor parallel group
|
||||
# The two values cannot be larger than 1 at the same time
|
||||
if get_moe_expert_parallel_world_size() > 1:
|
||||
world_size = get_moe_expert_parallel_world_size()
|
||||
rank = get_moe_expert_parallel_rank()
|
||||
if get_parallel().moe_ep_size > 1:
|
||||
world_size = get_parallel().moe_ep_size
|
||||
rank = get_parallel().moe_ep_rank
|
||||
coordinator = get_moe_ep_group()
|
||||
else:
|
||||
world_size = get_moe_tensor_parallel_world_size()
|
||||
rank = get_moe_tensor_parallel_rank()
|
||||
world_size = get_parallel().moe_tp_size
|
||||
rank = get_parallel().moe_tp_rank
|
||||
coordinator = get_moe_tp_group()
|
||||
|
||||
# Always pass the coordinator's groups: flashinfer >=0.6.10 reads the
|
||||
# rendezvous group from `group=...` (falling back to WORLD when None),
|
||||
# so leaving it None silently rendezvouses on WORLD and the kernel ends
|
||||
# up addressing the wrong peers in TP/EP/CP subgroup setups.
|
||||
device_group = coordinator.device_group
|
||||
cpu_group = coordinator.cpu_group
|
||||
|
||||
if world_size <= 1:
|
||||
return False
|
||||
|
||||
workspace_manager = _get_workspace_manager(use_attn_tp_group)
|
||||
token_num = token_num or max_token_num
|
||||
group_key = (device_group, cpu_group)
|
||||
effective_dtype = dtype or torch.bfloat16
|
||||
server_args = get_global_server_args()
|
||||
backend = resolve_flashinfer_allreduce_fusion_backend(server_args)
|
||||
if backend is None:
|
||||
return False
|
||||
|
||||
if (
|
||||
not _workspace_manager.initialized
|
||||
or _workspace_manager.world_size != world_size
|
||||
or _workspace_manager.rank != rank
|
||||
or not _workspace_manager.is_buffer_size_sufficient(
|
||||
not workspace_manager.initialized
|
||||
or workspace_manager.world_size != world_size
|
||||
or workspace_manager.rank != rank
|
||||
or workspace_manager.group != group_key
|
||||
or not workspace_manager.is_buffer_size_sufficient(
|
||||
token_num=token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
dtype=effective_dtype,
|
||||
use_oneshot=use_oneshot,
|
||||
)
|
||||
):
|
||||
_workspace_manager.initialize(
|
||||
workspace_manager.initialize(
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
backend=backend,
|
||||
group=cpu_group,
|
||||
use_fp32_lamport=use_fp32_lamport,
|
||||
dtype=dtype,
|
||||
use_oneshot=use_oneshot,
|
||||
device_group=device_group,
|
||||
cpu_group=cpu_group,
|
||||
)
|
||||
|
||||
return _workspace_manager.initialized
|
||||
_sync_allreduce_unavailable_across_tp()
|
||||
|
||||
return workspace_manager.initialized
|
||||
|
||||
|
||||
def fake_flashinfer_allreduce_residual_rmsnorm(
|
||||
@@ -308,7 +727,9 @@ def flashinfer_allreduce_residual_rmsnorm(
|
||||
use_attn_tp_group: bool = True,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Use FlashInfer's fused allreduce + residual + RMS norm operation
|
||||
Use FlashInfer's unified fused allreduce + residual + RMS norm operation.
|
||||
Automatically selects between trtllm and mnnvl backends based on topology
|
||||
and hardware (controlled by --flashinfer-allreduce-fusion-backend).
|
||||
|
||||
Args:
|
||||
input_tensor: Input tensor that needs allreduce
|
||||
@@ -331,15 +752,12 @@ def flashinfer_allreduce_residual_rmsnorm(
|
||||
return None, None
|
||||
|
||||
if use_attn_tp_group:
|
||||
world_size = get_attn_tensor_model_parallel_world_size()
|
||||
world_size = get_parallel().attn_tp_size
|
||||
else:
|
||||
# If MoE expert parallel world size > 1, use expert parallel group
|
||||
# Otherwise, use tensor parallel group
|
||||
# The two values cannot be larger than 1 at the same time
|
||||
if get_moe_expert_parallel_world_size() > 1:
|
||||
world_size = get_moe_expert_parallel_world_size()
|
||||
if get_parallel().moe_ep_size > 1:
|
||||
world_size = get_parallel().moe_ep_size
|
||||
else:
|
||||
world_size = get_moe_tensor_parallel_world_size()
|
||||
world_size = get_parallel().moe_tp_size
|
||||
|
||||
if world_size <= 1:
|
||||
logger.debug("Single GPU, no need for allreduce fusion")
|
||||
@@ -357,6 +775,7 @@ def flashinfer_allreduce_residual_rmsnorm(
|
||||
if not ensure_workspace_initialized(
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=input_tensor.shape[-1],
|
||||
use_fp32_lamport=(input_tensor.dtype == torch.float32),
|
||||
dtype=input_tensor.dtype,
|
||||
token_num=input_tensor.shape[0],
|
||||
use_oneshot=use_oneshot,
|
||||
@@ -365,12 +784,17 @@ def flashinfer_allreduce_residual_rmsnorm(
|
||||
logger.debug("FlashInfer workspace not available")
|
||||
return None, None
|
||||
|
||||
workspace_manager = _get_workspace_manager(use_attn_tp_group)
|
||||
if workspace_manager.workspace is None:
|
||||
logger.debug("FlashInfer workspace is None")
|
||||
return None, None
|
||||
|
||||
residual_out = torch.empty_like(residual)
|
||||
norm_out = torch.empty_like(input_tensor)
|
||||
|
||||
_flashinfer_comm.allreduce_fusion(
|
||||
kwargs = dict(
|
||||
input=input_tensor,
|
||||
workspace=_workspace_manager.workspace,
|
||||
workspace=workspace_manager.workspace,
|
||||
pattern=_flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm,
|
||||
launch_with_pdl=True,
|
||||
residual_out=residual_out,
|
||||
@@ -381,11 +805,53 @@ def flashinfer_allreduce_residual_rmsnorm(
|
||||
use_oneshot=use_oneshot,
|
||||
fp32_acc=fp32_acc,
|
||||
)
|
||||
if _flashinfer_allreduce_supports_trigger_completion:
|
||||
kwargs["trigger_completion_at_end"] = trigger_completion_at_end
|
||||
_flashinfer_comm.allreduce_fusion(**kwargs)
|
||||
|
||||
return norm_out, residual_out
|
||||
|
||||
|
||||
def pre_initialize_workspaces(
|
||||
max_token_num: int,
|
||||
hidden_dim: int,
|
||||
dtype: torch.dtype,
|
||||
use_oneshot: Optional[bool] = None,
|
||||
):
|
||||
"""Pre-initialize flashinfer workspaces before CUDA graph capture.
|
||||
|
||||
This must be called before graph capture to avoid collective operations
|
||||
(broadcasts, barriers) inside the graph capture context, which can
|
||||
deadlock with custom_all_reduce.register_graph_buffers.
|
||||
"""
|
||||
if _flashinfer_allreduce_unavailable or _flashinfer_comm is None:
|
||||
return
|
||||
|
||||
# Initialize MoE workspace
|
||||
ensure_workspace_initialized(
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
use_oneshot=use_oneshot,
|
||||
use_attn_tp_group=False,
|
||||
)
|
||||
|
||||
# Initialize attention workspace
|
||||
ensure_workspace_initialized(
|
||||
max_token_num=max_token_num,
|
||||
hidden_dim=hidden_dim,
|
||||
dtype=dtype,
|
||||
use_oneshot=use_oneshot,
|
||||
use_attn_tp_group=True,
|
||||
)
|
||||
|
||||
|
||||
def cleanup_flashinfer_workspace():
|
||||
global _workspace_manager
|
||||
if _workspace_manager is not None:
|
||||
_workspace_manager.cleanup()
|
||||
global _attn_tp_workspace_manager, _moe_tp_workspace_manager
|
||||
if _attn_tp_workspace_manager is not None:
|
||||
_attn_tp_workspace_manager.cleanup()
|
||||
if (
|
||||
_moe_tp_workspace_manager is not None
|
||||
and _moe_tp_workspace_manager is not _attn_tp_workspace_manager
|
||||
):
|
||||
_moe_tp_workspace_manager.cleanup()
|
||||
|
||||
522
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py
Normal file
522
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py
Normal file
@@ -0,0 +1,522 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.moe.moe_runner.base import (
|
||||
MoeQuantInfo,
|
||||
MoeRunnerConfig,
|
||||
register_fused_func,
|
||||
)
|
||||
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
|
||||
from sglang.srt.utils.common import log_info_on_rank0, print_warning_once
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.batch_overlap.single_batch_overlap import DownGemmOverlapArgs
|
||||
from sglang.srt.layers.moe.token_dispatcher import (
|
||||
DeepEPLLCombineInput,
|
||||
DeepEPLLDispatchOutput,
|
||||
StandardCombineInput,
|
||||
StandardDispatchOutput,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
|
||||
FlashinferCombineInput,
|
||||
FlashinferDispatchOutput,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_FP4_SF_VEC_SIZE = 16
|
||||
_cutedsl_logged_scalarize: set = set()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Weight / scale preparation utilities (called from modelopt_quant.py during
|
||||
# process_weights_after_loading and lazy wrapper init)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def interleave_w13_halves(
|
||||
tensor: torch.Tensor, group_size: int = 64, dim: int = 1
|
||||
) -> torch.Tensor:
|
||||
"""Interleave the two logical W13 halves for CuteDSL's SwiGLU GEMM1 layout.
|
||||
|
||||
The caller is responsible for loading W13 in the expected two-half order.
|
||||
This helper only rewrites the first and second halves into alternating
|
||||
`group_size` chunks along `dim`.
|
||||
"""
|
||||
if tensor.shape[dim] % 2 != 0:
|
||||
raise ValueError(
|
||||
"Expected even size on interleave dimension for W13 half split."
|
||||
)
|
||||
split = tensor.shape[dim] // 2
|
||||
if split % group_size != 0:
|
||||
raise ValueError(
|
||||
f"Expected split dim divisible by group_size={group_size}, got {split}."
|
||||
)
|
||||
first_half = tensor.narrow(dim, 0, split)
|
||||
second_half = tensor.narrow(dim, split, split)
|
||||
first_half_groups = first_half.split(group_size, dim=dim)
|
||||
second_half_groups = second_half.split(group_size, dim=dim)
|
||||
interleaved = [
|
||||
item for pair in zip(first_half_groups, second_half_groups) for item in pair
|
||||
]
|
||||
return torch.cat(interleaved, dim=dim)
|
||||
|
||||
|
||||
def cutedsl_quant_scale_to_scalar(
|
||||
quant_scale: torch.Tensor,
|
||||
*,
|
||||
name: str,
|
||||
) -> torch.Tensor:
|
||||
"""Reduce per-expert quant-domain scale vector to a single scalar.
|
||||
|
||||
The quant domain is the reciprocal of the raw checkpoint scale:
|
||||
quant_scale = 1 / raw_scale
|
||||
|
||||
Returns min(quant_scale) = 1/max(raw_scale), which is the TRTLLM CuteDSL
|
||||
convention for global scalar activation scales (see TRTLLM quantization.py
|
||||
lines 2137-2141: fc2_input_scale = tmp_fc2_input_scale.max().reciprocal()).
|
||||
|
||||
If quant_scale is already scalar (numel==1), returns it unchanged.
|
||||
"""
|
||||
quant_scale = quant_scale.to(torch.float32)
|
||||
if quant_scale.numel() == 0:
|
||||
print_warning_once(
|
||||
f"CuteDSL got empty {name}; using 1.0 fallback.",
|
||||
)
|
||||
return torch.ones(1, device=quant_scale.device, dtype=torch.float32)
|
||||
if quant_scale.numel() == 1:
|
||||
return quant_scale.reshape(1)
|
||||
if name not in _cutedsl_logged_scalarize:
|
||||
log_info_on_rank0(
|
||||
logger,
|
||||
f"CuteDSL: reducing per-expert {name} to scalar via "
|
||||
"min(quant_scale) = 1/max(raw_scale), matching TRTLLM convention.",
|
||||
)
|
||||
_cutedsl_logged_scalarize.add(name)
|
||||
return quant_scale.min().reshape(1)
|
||||
|
||||
|
||||
def resolve_cutedsl_standard_scales(
|
||||
layer: torch.nn.Module,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Resolve standard-path CuteDSL scales (baseline: scalar fc2/w13 input scales).
|
||||
|
||||
Returns (w1_alpha, fc2_input_scale, w2_alpha, used_input_scale).
|
||||
used_input_scale is the scalarized w13 input scale for FP4 quantize and GEMM1.
|
||||
"""
|
||||
|
||||
def _to_fp32_tensor(x: torch.Tensor | float, ref: torch.Tensor) -> torch.Tensor:
|
||||
if not isinstance(x, torch.Tensor):
|
||||
x = torch.tensor(x, device=ref.device)
|
||||
return x.to(device=ref.device, dtype=torch.float32)
|
||||
|
||||
def _align_scale_to_alpha(
|
||||
scale: torch.Tensor, alpha: torch.Tensor, scale_name: str
|
||||
) -> torch.Tensor:
|
||||
scale = scale.to(device=alpha.device, dtype=torch.float32)
|
||||
alpha = alpha.to(torch.float32)
|
||||
if scale.ndim == 0:
|
||||
return scale
|
||||
# Gated weight scales may be (num_experts, 2) with separate gate/up
|
||||
# columns. Collapse to 1D by taking the first column (gate == up for
|
||||
# well-formed checkpoints; mismatch is warned in process_weights_after_loading).
|
||||
if scale.ndim == 2 and scale.shape[1] <= 2:
|
||||
scale = scale[:, 0]
|
||||
if scale.numel() == alpha.numel():
|
||||
return scale
|
||||
if scale.numel() == 1:
|
||||
return scale.reshape(())
|
||||
|
||||
# Some EP setups may carry global-per-expert scale vectors while alphas are
|
||||
# local-per-expert vectors. Slice to this rank's local expert range.
|
||||
num_local_experts = getattr(layer, "num_local_experts", None)
|
||||
num_experts = getattr(layer, "num_experts", None)
|
||||
moe_ep_rank = getattr(layer, "moe_ep_rank", 0)
|
||||
if (
|
||||
num_local_experts is not None
|
||||
and num_experts is not None
|
||||
and scale.numel() == num_experts
|
||||
and alpha.numel() == num_local_experts
|
||||
):
|
||||
start = moe_ep_rank * num_local_experts
|
||||
end = start + num_local_experts
|
||||
return scale[start:end]
|
||||
|
||||
raise ValueError(
|
||||
f"Unable to align {scale_name} shape={tuple(scale.shape)} "
|
||||
f"to alpha shape={tuple(alpha.shape)} for CuteDSL standard scale resolution."
|
||||
)
|
||||
|
||||
def _resolve_w1_alpha_from_scalar_input_scale(
|
||||
used_input_scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Resolve GEMM1 alpha consistent with scalarized activation quant scale.
|
||||
|
||||
CuteDSL pre-quantizes x with a single scalar (used_input_scale), but
|
||||
g1_alphas was derived with per-expert activation scales:
|
||||
g1_alphas[e] = (1/w13_isq[e]) * w13_ws2[e]
|
||||
Correct alpha for scalar quantization:
|
||||
w1_alpha[e] = w13_ws2[e] / used_input_scale
|
||||
= g1_alphas[e] * w13_isq[e] / used_input_scale
|
||||
When w13_isq is already scalar, this is a no-op (ratio = 1).
|
||||
"""
|
||||
eps = 1e-12
|
||||
scalar = torch.clamp(used_input_scale.to(torch.float32).reshape(()), min=eps)
|
||||
|
||||
if hasattr(layer, "w13_weight_scale_2"):
|
||||
w13_weight_scale_2 = _align_scale_to_alpha(
|
||||
layer.w13_weight_scale_2, layer.g1_alphas, "w13_weight_scale_2"
|
||||
)
|
||||
return w13_weight_scale_2.to(torch.float32) / scalar
|
||||
|
||||
w13_isq = _align_scale_to_alpha(
|
||||
layer.w13_input_scale_quant, layer.g1_alphas, "w13_input_scale_quant"
|
||||
)
|
||||
w13_isq = torch.clamp(_to_fp32_tensor(w13_isq, layer.g1_alphas), min=eps)
|
||||
return (layer.g1_alphas.to(torch.float32) * w13_isq / scalar).to(torch.float32)
|
||||
|
||||
def _resolve_w2_alpha_from_scalar_fc2_input_scale(
|
||||
fc2_input_scale: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Resolve GEMM2 alpha consistent with scalarized FC2 input scale.
|
||||
|
||||
CuteDSL standard path uses a scalar global scale for GEMM1 FP4 output
|
||||
quantization (`fc2_input_scale`). GEMM2 alpha must use the same scalar
|
||||
convention: alpha2 = w2_weight_scale_2 / fc2_input_scale.
|
||||
"""
|
||||
eps = 1e-12
|
||||
fc2_input_scale = fc2_input_scale.to(torch.float32)
|
||||
fc2_scalar = torch.clamp(fc2_input_scale.reshape(-1)[:1], min=eps).reshape(())
|
||||
|
||||
if hasattr(layer, "w2_weight_scale_2"):
|
||||
w2_weight_scale_2 = _align_scale_to_alpha(
|
||||
layer.w2_weight_scale_2, layer.g2_alphas, "w2_weight_scale_2"
|
||||
)
|
||||
w2_weight_scale_2 = w2_weight_scale_2.to(torch.float32)
|
||||
return w2_weight_scale_2 / fc2_scalar
|
||||
|
||||
w2_q_for_w2 = _align_scale_to_alpha(
|
||||
layer.w2_input_scale_quant, layer.g2_alphas, "w2_input_scale_quant"
|
||||
)
|
||||
w2_q_for_w2 = torch.clamp(
|
||||
_to_fp32_tensor(w2_q_for_w2, layer.g2_alphas), min=eps
|
||||
)
|
||||
w2_weight_scale_2 = layer.g2_alphas.to(torch.float32) * w2_q_for_w2
|
||||
return w2_weight_scale_2 / fc2_scalar
|
||||
|
||||
fc2_input_scale = cutedsl_quant_scale_to_scalar(
|
||||
layer.w2_input_scale_quant,
|
||||
name="w2_input_scale_quant",
|
||||
)
|
||||
w2_alpha = _resolve_w2_alpha_from_scalar_fc2_input_scale(fc2_input_scale)
|
||||
used_input_scale = cutedsl_quant_scale_to_scalar(
|
||||
layer.w13_input_scale_quant,
|
||||
name="w13_input_scale_quant",
|
||||
)
|
||||
w1_alpha = _resolve_w1_alpha_from_scalar_input_scale(used_input_scale)
|
||||
return w1_alpha, fc2_input_scale, w2_alpha, used_input_scale
|
||||
|
||||
|
||||
def ensure_cutedsl_wrapper(layer: torch.nn.Module) -> None:
|
||||
"""Lazily create CuteDslMoEWrapper and resolve scales on first forward.
|
||||
|
||||
The wrapper is created lazily (not in __init__ / create_weights) because
|
||||
it depends on final weight shapes and EP configuration. The wrapper's
|
||||
CUDA-graph buffers are allocated inside CuteDslMoEWrapper.__init__, which
|
||||
typically runs during the autotune dummy forward under inference_mode().
|
||||
We wrap the creation in inference_mode(False) so that those pre-allocated
|
||||
buffers are normal tensors -- inference tensors cannot be inplace-updated
|
||||
during later CUDA graph capture, which runs outside inference_mode.
|
||||
"""
|
||||
if getattr(layer, "_cutedsl_wrapper", None) is not None:
|
||||
return
|
||||
|
||||
try:
|
||||
from flashinfer import CuteDslMoEWrapper
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"flashinfer_cutedsl backend requires FlashInfer with CuteDSL support. "
|
||||
"Install with: pip install flashinfer"
|
||||
) from e
|
||||
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
|
||||
assert layer.intermediate_size_per_partition > 0, (
|
||||
f"CuteDSL MoE: intermediate_size_per_partition must be > 0, "
|
||||
f"got {layer.intermediate_size_per_partition}. Check EP/TP configuration."
|
||||
)
|
||||
|
||||
server_args = get_global_server_args()
|
||||
# CuteDSL wrapper preallocates CG buffers used by any captured graph
|
||||
# that routes through this MoE — decode and prefill alike.
|
||||
use_cuda_graph = not cuda_graph_fully_disabled()
|
||||
|
||||
# Size the wrapper's CUDA-graph buffers for the largest number of tokens a
|
||||
# single forward can route through this layer.
|
||||
dispatcher = getattr(layer, "dispatcher", None)
|
||||
if hasattr(dispatcher, "max_num_tokens"):
|
||||
# A2A path: bounded by the dispatcher's own workspace limit.
|
||||
max_num_tokens = dispatcher.max_num_tokens * getattr(dispatcher, "ep_size", 1)
|
||||
else:
|
||||
# Standard allgather path: the MoE sees up to dp_size local forwards
|
||||
# gathered together, so scale the per-rank forward bound by dp_size.
|
||||
max_num_tokens = server_args.dp_size * server_args.cutedsl_moe_max_num_tokens()
|
||||
top_k = layer.top_k if layer.top_k is not None else layer.moe_runner_config.top_k
|
||||
# inference_mode(False) ensures the wrapper's pre-allocated CUDA-graph
|
||||
# buffers are normal tensors. This call typically happens inside
|
||||
# _dummy_run which runs under inference_mode(); inference tensors cannot
|
||||
# be inplace-updated during later CUDA graph capture (which runs outside
|
||||
# inference_mode), so we must opt out here.
|
||||
with torch.inference_mode(False):
|
||||
layer._cutedsl_wrapper = CuteDslMoEWrapper(
|
||||
num_experts=layer.num_experts,
|
||||
top_k=top_k,
|
||||
hidden_size=layer.hidden_size,
|
||||
intermediate_size=layer.intermediate_size_per_partition,
|
||||
use_cuda_graph=use_cuda_graph,
|
||||
max_num_tokens=max_num_tokens,
|
||||
num_local_experts=layer.num_local_experts,
|
||||
local_expert_offset=layer.moe_ep_rank * layer.num_local_experts,
|
||||
output_dtype=layer.moe_runner_config.params_dtype,
|
||||
device=str(layer.w13_weight.device),
|
||||
)
|
||||
|
||||
w1_alpha, fc2_input_scale, w2_alpha, used_input_scale = (
|
||||
resolve_cutedsl_standard_scales(layer)
|
||||
)
|
||||
layer._cutedsl_scales = (w1_alpha, fc2_input_scale, w2_alpha)
|
||||
layer._cutedsl_input_scale = used_input_scale
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataclass + fused function for moe_runner dispatch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class CuteDslFp4MoeQuantInfo(MoeQuantInfo):
|
||||
"""Quantization payload for FlashInfer CuteDSL FP4 MoE kernels.
|
||||
|
||||
Shared by the two CuteDSL runner entries:
|
||||
|
||||
* "v2" standard path (a2a=none/flashinfer): consumed by the
|
||||
@register_fused_func("none", "flashinfer_cutedsl") entry, which
|
||||
drives CuteDslMoEWrapper.run. Weights are [Up, Gate]
|
||||
interleaved with MMA-layout blockscales. wrapper is set;
|
||||
w*_scale are scalarized.
|
||||
|
||||
* "v1" DeepEP low-latency path (a2a=deepep): consumed by the
|
||||
@register_fused_func("deepep", "flashinfer_cutedsl") entry,
|
||||
which drives flashinfer_cutedsl_moe_masked. Weights are
|
||||
[Gate, Up] non-interleaved with swizzled blockscales.
|
||||
wrapper is None; w*_scale are per-expert.
|
||||
"""
|
||||
|
||||
# FP4 packed weights (uint8)
|
||||
w13_weight: torch.Tensor
|
||||
w2_weight: torch.Tensor
|
||||
|
||||
# Block-scale factors (MMA layout for v2, swizzled for v1)
|
||||
w13_weight_sf: torch.Tensor
|
||||
w2_weight_sf: torch.Tensor
|
||||
|
||||
# Per-expert GEMM dequant alphas (scalarized for v2, per-expert for v1)
|
||||
w1_alpha: torch.Tensor
|
||||
w2_alpha: torch.Tensor
|
||||
|
||||
# Activation quant scales (1 / raw_input_scale).
|
||||
# - a1_scale: quantizes hidden_states before GEMM1
|
||||
# - a2_scale: quantizes GEMM1 output before GEMM2 (a.k.a. fc2 input)
|
||||
a1_scale: torch.Tensor
|
||||
a2_scale: torch.Tensor
|
||||
|
||||
# v2 only: lazily-created CuteDslMoEWrapper (None on the v1 path).
|
||||
wrapper: Optional[Any] = None
|
||||
|
||||
# v1 only: True when DeepEP pre-quantizes activations to NVFP4.
|
||||
use_nvfp4_dispatch: bool = False
|
||||
|
||||
# v1 only: SBO down-GEMM overlap args.
|
||||
down_gemm_overlap_args: Optional[DownGemmOverlapArgs] = None
|
||||
|
||||
|
||||
@register_fused_func("none", "flashinfer_cutedsl")
|
||||
def fused_experts_none_to_flashinfer_cutedsl_fp4(
|
||||
dispatch_output: StandardDispatchOutput,
|
||||
quant_info: CuteDslFp4MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> StandardCombineInput:
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardCombineInput
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
|
||||
|
||||
assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
|
||||
assert quant_info.wrapper is not None, "CuteDSL v2 path requires CuteDslMoEWrapper."
|
||||
|
||||
hidden_states = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
assert TopKOutputChecker.format_is_standard(topk_output)
|
||||
|
||||
topk_ids = topk_output.topk_ids
|
||||
topk_weights = topk_output.topk_weights
|
||||
if topk_ids.dtype != torch.int32:
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
|
||||
x_fp4, x_sf = fp4_quantize(
|
||||
hidden_states,
|
||||
quant_info.a1_scale,
|
||||
sf_vec_size=_FP4_SF_VEC_SIZE,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
|
||||
output = quant_info.wrapper.run(
|
||||
x=x_fp4,
|
||||
x_sf=x_sf,
|
||||
token_selected_experts=topk_ids,
|
||||
token_final_scales=topk_weights,
|
||||
w1_weight=quant_info.w13_weight,
|
||||
w1_weight_sf=quant_info.w13_weight_sf,
|
||||
w1_alpha=quant_info.w1_alpha,
|
||||
fc2_input_scale=quant_info.a2_scale,
|
||||
w2_weight=quant_info.w2_weight,
|
||||
w2_weight_sf=quant_info.w2_weight_sf,
|
||||
w2_alpha=quant_info.w2_alpha,
|
||||
)
|
||||
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
@register_fused_func("flashinfer", "flashinfer_cutedsl")
|
||||
def fused_experts_flashinfer_to_flashinfer_cutedsl_fp4(
|
||||
dispatch_output: FlashinferDispatchOutput,
|
||||
quant_info: CuteDslFp4MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> FlashinferCombineInput:
|
||||
"""CuteDSL fused func for flashinfer alltoall dispatcher.
|
||||
|
||||
Two cases depending on whether the dispatcher did FP4 quantization:
|
||||
- bf16 input (SGLANG_MOE_NVFP4_DISPATCH=0): quantize with cutedsl's scale
|
||||
- FP4 input (SGLANG_MOE_NVFP4_DISPATCH=1): pass through (same fp4_quantize params)
|
||||
"""
|
||||
from sglang.srt.layers.moe.token_dispatcher.flashinfer import (
|
||||
FlashinferCombineInput,
|
||||
)
|
||||
from sglang.srt.layers.moe.topk import TopKOutputChecker
|
||||
from sglang.srt.layers.quantization.fp4_utils import fp4_quantize
|
||||
|
||||
assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
|
||||
assert quant_info.wrapper is not None, "CuteDSL v2 path requires CuteDslMoEWrapper."
|
||||
|
||||
hidden_states = dispatch_output.hidden_states
|
||||
x_sf = dispatch_output.hidden_states_scale
|
||||
topk_output = dispatch_output.topk_output
|
||||
assert TopKOutputChecker.format_is_standard(topk_output)
|
||||
|
||||
topk_ids = topk_output.topk_ids
|
||||
topk_weights = topk_output.topk_weights
|
||||
if topk_ids.dtype != torch.int32:
|
||||
topk_ids = topk_ids.to(torch.int32)
|
||||
|
||||
if x_sf is not None:
|
||||
# NVFP4 dispatch, inputs are already quantized.
|
||||
x_fp4 = hidden_states
|
||||
else:
|
||||
x_fp4, x_sf = fp4_quantize(
|
||||
hidden_states,
|
||||
quant_info.a1_scale,
|
||||
sf_vec_size=_FP4_SF_VEC_SIZE,
|
||||
is_sf_swizzled_layout=False,
|
||||
)
|
||||
|
||||
output = quant_info.wrapper.run(
|
||||
x=x_fp4,
|
||||
x_sf=x_sf,
|
||||
token_selected_experts=topk_ids,
|
||||
token_final_scales=topk_weights,
|
||||
w1_weight=quant_info.w13_weight,
|
||||
w1_weight_sf=quant_info.w13_weight_sf,
|
||||
w1_alpha=quant_info.w1_alpha,
|
||||
fc2_input_scale=quant_info.a2_scale,
|
||||
w2_weight=quant_info.w2_weight,
|
||||
w2_weight_sf=quant_info.w2_weight_sf,
|
||||
w2_alpha=quant_info.w2_alpha,
|
||||
)
|
||||
|
||||
# Note: output contains routed expert results; shared_expert is handled separately
|
||||
|
||||
# Write into pre-allocated workspace buffer if available
|
||||
if dispatch_output.moe_output is not None:
|
||||
dispatch_output.moe_output.copy_(output)
|
||||
output = dispatch_output.moe_output
|
||||
|
||||
return FlashinferCombineInput(hidden_states=output)
|
||||
|
||||
|
||||
@register_fused_func("deepep", "flashinfer_cutedsl")
|
||||
def fused_experts_deepep_to_flashinfer_cutedsl_fp4(
|
||||
dispatch_output: DeepEPLLDispatchOutput,
|
||||
quant_info: CuteDslFp4MoeQuantInfo,
|
||||
runner_config: MoeRunnerConfig,
|
||||
) -> DeepEPLLCombineInput:
|
||||
from sglang.srt.layers.moe.flashinfer_cutedsl_moe import (
|
||||
flashinfer_cutedsl_moe_masked,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPLLCombineInput
|
||||
|
||||
assert runner_config.activation == "silu", "Only silu is supported for CuteDSL MoE."
|
||||
assert (
|
||||
not runner_config.apply_router_weight_on_input
|
||||
), "apply_router_weight_on_input is not supported for Flashinfer"
|
||||
|
||||
hidden_states, hidden_states_scale, _, _, masked_m, _ = dispatch_output
|
||||
|
||||
# flashinfer_cutedsl_moe_masked reinterprets scales as float8_e4m3fn.
|
||||
# Same-dtype .view is a no-op; only wider dtypes (e.g. int32-packed
|
||||
# UE8M0) need stride(-1)==1.
|
||||
if (
|
||||
quant_info.use_nvfp4_dispatch
|
||||
and hidden_states_scale is not None
|
||||
and hidden_states_scale.element_size() != 1
|
||||
and hidden_states_scale.stride(-1) != 1
|
||||
):
|
||||
raise AssertionError(
|
||||
f"NVFP4 dispatch scale has stride(-1)={hidden_states_scale.stride(-1)}, "
|
||||
f"dtype={hidden_states_scale.dtype}; .view(float8_e4m3fn) requires stride(-1)==1. "
|
||||
"Try SGLANG_MOE_NVFP4_DISPATCH=0 or check DeepEP version."
|
||||
)
|
||||
|
||||
overlap = quant_info.down_gemm_overlap_args
|
||||
output = flashinfer_cutedsl_moe_masked(
|
||||
hidden_states=(hidden_states, hidden_states_scale),
|
||||
input_global_scale=(
|
||||
None if quant_info.use_nvfp4_dispatch else quant_info.a1_scale
|
||||
),
|
||||
w1=quant_info.w13_weight,
|
||||
w1_blockscale=quant_info.w13_weight_sf,
|
||||
w1_alpha=quant_info.w1_alpha,
|
||||
w2=quant_info.w2_weight,
|
||||
a2_global_scale=quant_info.a2_scale,
|
||||
w2_blockscale=quant_info.w2_weight_sf,
|
||||
w2_alpha=quant_info.w2_alpha,
|
||||
masked_m=masked_m,
|
||||
**(
|
||||
dict(
|
||||
down_sm_count=overlap.num_sms,
|
||||
down_signals=overlap.signal,
|
||||
down_start_event=overlap.start_event,
|
||||
)
|
||||
if overlap is not None
|
||||
else {}
|
||||
),
|
||||
)
|
||||
|
||||
return DeepEPLLCombineInput(
|
||||
hidden_states=output,
|
||||
topk_ids=dispatch_output.topk_ids,
|
||||
topk_weights=dispatch_output.topk_weights,
|
||||
)
|
||||
224
python/sglang/srt/model_executor/cuda_graph_config.py
Normal file
224
python/sglang/srt/model_executor/cuda_graph_config.py
Normal file
@@ -0,0 +1,224 @@
|
||||
# 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.
|
||||
# ==============================================================================
|
||||
"""Phase / backend identifiers, the canonical default for
|
||||
cuda_graph_config, and the --cuda-graph-config JSON CLI parser.
|
||||
|
||||
Module-level imports are pure stdlib — no torch / sglang.srt deps — so
|
||||
ServerArgs can import everything here without pulling in backend
|
||||
classes. check_cuda_graph_backend lazy-imports get_global_server_args
|
||||
inside the function body to preserve that invariant.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import dataclasses
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
|
||||
class Phase:
|
||||
"""The two phases of model forward."""
|
||||
|
||||
DECODE = "decode"
|
||||
PREFILL = "prefill"
|
||||
ALL = (DECODE, PREFILL)
|
||||
|
||||
|
||||
class Backend:
|
||||
"""CUDA graph capture backends a phase can use."""
|
||||
|
||||
FULL = "full"
|
||||
BREAKABLE = "breakable"
|
||||
TC_PIECEWISE = "tc_piecewise"
|
||||
DISABLED = "disabled"
|
||||
ALL = (FULL, BREAKABLE, TC_PIECEWISE, DISABLED)
|
||||
|
||||
|
||||
ALLOWED_BACKENDS_PER_PHASE = {
|
||||
Phase.DECODE: (
|
||||
Backend.FULL,
|
||||
Backend.BREAKABLE,
|
||||
Backend.TC_PIECEWISE,
|
||||
Backend.DISABLED,
|
||||
),
|
||||
# full is rejected for prefill — full CUDA graph capture only
|
||||
# fits fixed-shape and prefill is variable-shape. Use breakable
|
||||
# or tc_piecewise for prefill.
|
||||
Phase.PREFILL: (Backend.BREAKABLE, Backend.TC_PIECEWISE, Backend.DISABLED),
|
||||
}
|
||||
|
||||
# Per-phase settings schema. Keys other than backend are runner-level
|
||||
# (read by any backend in that phase); tc_compiler is the lone
|
||||
# backend-specific knob (only meaningful when backend == tc_piecewise).
|
||||
# For prefill, bs carries the captured shape size (token count for
|
||||
# tc_piecewise, request count for breakable) — one shape knob per phase.
|
||||
ALLOWED_KEYS_PER_PHASE = {
|
||||
Phase.DECODE: ("backend", "max_bs", "bs", "tc_compiler"),
|
||||
Phase.PREFILL: ("backend", "max_bs", "bs", "tc_compiler"),
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class PhaseConfig:
|
||||
"""Per-phase CUDA graph settings."""
|
||||
|
||||
backend: str = Backend.DISABLED
|
||||
max_bs: Optional[int] = None
|
||||
bs: Optional[List[int]] = None
|
||||
# Only meaningful when backend == tc_piecewise; ignored otherwise.
|
||||
tc_compiler: str = "eager"
|
||||
|
||||
|
||||
@dataclass
|
||||
class CudaGraphConfig:
|
||||
"""Top-level CUDA graph config: one PhaseConfig per phase."""
|
||||
|
||||
decode: PhaseConfig = field(
|
||||
default_factory=lambda: PhaseConfig(backend=Backend.FULL)
|
||||
)
|
||||
prefill: PhaseConfig = field(
|
||||
default_factory=lambda: PhaseConfig(backend=Backend.TC_PIECEWISE)
|
||||
)
|
||||
|
||||
def __getitem__(self, phase: str) -> PhaseConfig:
|
||||
"""Phase-string lookup; kept for migration ergonomics."""
|
||||
if phase not in Phase.ALL:
|
||||
raise KeyError(phase)
|
||||
return getattr(self, phase)
|
||||
|
||||
def to_dict(self) -> Dict[str, Dict[str, Any]]:
|
||||
# Diff-only, not asdict: the parser locks every (phase, key) it sees,
|
||||
# so emitting defaults would lock fields the caller never set.
|
||||
baseline = default_cuda_graph_config()
|
||||
return {
|
||||
Phase.DECODE: _diff_phase(self.decode, baseline.decode),
|
||||
Phase.PREFILL: _diff_phase(self.prefill, baseline.prefill),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, raw: Optional[Dict[str, Dict[str, Any]]]) -> "CudaGraphConfig":
|
||||
"""Build from a (partial) dict of overrides, defaults fill the rest.
|
||||
Unknown phases / keys are silently dropped — the JSON-input
|
||||
validator (parse_cuda_graph_config_arg) rejects them upstream."""
|
||||
cfg = cls()
|
||||
if not raw:
|
||||
return cfg
|
||||
for phase, phase_settings in raw.items():
|
||||
if phase not in Phase.ALL or not isinstance(phase_settings, dict):
|
||||
continue
|
||||
phase_cfg = getattr(cfg, phase)
|
||||
allowed = ALLOWED_KEYS_PER_PHASE[phase]
|
||||
for key, value in phase_settings.items():
|
||||
if key in allowed:
|
||||
setattr(phase_cfg, key, value)
|
||||
return cfg
|
||||
|
||||
|
||||
def default_cuda_graph_config() -> CudaGraphConfig:
|
||||
"""Fresh CudaGraphConfig populated with canonical defaults."""
|
||||
return CudaGraphConfig()
|
||||
|
||||
|
||||
def _diff_phase(actual: PhaseConfig, baseline: PhaseConfig) -> Dict[str, Any]:
|
||||
"""Return only fields whose value differs from the per-phase default."""
|
||||
return {
|
||||
f.name: getattr(actual, f.name)
|
||||
for f in dataclasses.fields(actual)
|
||||
if getattr(actual, f.name) != getattr(baseline, f.name)
|
||||
}
|
||||
|
||||
|
||||
def check_cuda_graph_backend(phase: str, backend: str) -> bool:
|
||||
"""True if cuda_graph_config[phase].backend == backend on the
|
||||
global server args. Returns False if the global server args have not
|
||||
been initialized yet (e.g. unit tests, early startup)."""
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
|
||||
try:
|
||||
server_args = get_global_server_args()
|
||||
except ValueError:
|
||||
return False
|
||||
cfg = server_args.cuda_graph_config
|
||||
if cfg is None or phase not in Phase.ALL:
|
||||
return False
|
||||
return getattr(cfg, phase).backend == backend
|
||||
|
||||
|
||||
def cuda_graph_fully_disabled() -> bool:
|
||||
"""True iff cuda_graph_config has Backend.DISABLED on every phase.
|
||||
|
||||
Use at sites that ask the legacy server_args.disable_cuda_graph
|
||||
question ("no CG anywhere globally") — e.g., preallocating buffers
|
||||
that any captured graph would otherwise reuse, or one-shot init
|
||||
that's a no-op when CG is completely off.
|
||||
"""
|
||||
return check_cuda_graph_backend(
|
||||
Phase.DECODE, Backend.DISABLED
|
||||
) and check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED)
|
||||
|
||||
|
||||
def parse_cuda_graph_config_arg(raw: str) -> Dict[str, Dict[str, Any]]:
|
||||
"""argparse type for --cuda-graph-config: parse JSON dict of
|
||||
phase → settings dict. Each phase's settings dict is itself validated
|
||||
against ALLOWED_KEYS_PER_PHASE. Returns a plain dict — the
|
||||
precedence pipeline in ServerArgs converts to CudaGraphConfig
|
||||
after merging."""
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--cuda-graph-config must be JSON: {e}")
|
||||
if not isinstance(parsed, dict):
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"--cuda-graph-config must be a JSON object, got {type(parsed).__name__}"
|
||||
)
|
||||
|
||||
result: Dict[str, Dict[str, Any]] = {}
|
||||
for phase, phase_settings in parsed.items():
|
||||
phase = str(phase)
|
||||
if phase not in Phase.ALL:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"--cuda-graph-config: unknown phase '{phase}', expected one of {Phase.ALL}"
|
||||
)
|
||||
if not isinstance(phase_settings, dict):
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"--cuda-graph-config['{phase}'] must be a JSON object, got "
|
||||
f"{type(phase_settings).__name__}"
|
||||
)
|
||||
allowed = ALLOWED_KEYS_PER_PHASE[phase]
|
||||
result[phase] = {}
|
||||
for key, value in phase_settings.items():
|
||||
if key not in allowed:
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"--cuda-graph-config['{phase}']: unknown key '{key}', expected one of {allowed}"
|
||||
)
|
||||
result[phase][key] = value
|
||||
return result
|
||||
|
||||
|
||||
def explicit_keys_in(
|
||||
settings: Optional[Dict[str, Dict[str, Any]]],
|
||||
) -> set:
|
||||
"""Return the set of (phase, key) tuples present in settings
|
||||
(the raw dict form, as it arrives from CLI/SDK). Used by ServerArgs
|
||||
to track keys the user explicitly set so the auto-disable cascade can
|
||||
skip them."""
|
||||
out: set = set()
|
||||
if not settings:
|
||||
return out
|
||||
for phase, phase_settings in settings.items():
|
||||
if not isinstance(phase_settings, dict):
|
||||
continue
|
||||
for key in phase_settings.keys():
|
||||
out.add((phase, key))
|
||||
return out
|
||||
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