diff --git a/python/sglang/srt/layers/flashinfer_comm_fusion.py b/python/sglang/srt/layers/flashinfer_comm_fusion.py index 9417d573c..2411e09c0 100644 --- a/python/sglang/srt/layers/flashinfer_comm_fusion.py +++ b/python/sglang/srt/layers/flashinfer_comm_fusion.py @@ -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() diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py new file mode 100644 index 000000000..02c781fa4 --- /dev/null +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutedsl.py @@ -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, + ) diff --git a/python/sglang/srt/model_executor/cuda_graph_config.py b/python/sglang/srt/model_executor/cuda_graph_config.py new file mode 100644 index 000000000..0532e987d --- /dev/null +++ b/python/sglang/srt/model_executor/cuda_graph_config.py @@ -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 diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py new file mode 100644 index 000000000..9cc28257a --- /dev/null +++ b/python/sglang/srt/runtime_context.py @@ -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