diff --git a/vllm/distributed/eplb/eplb_utils.py b/vllm/distributed/eplb/eplb_utils.py index 92fffd22977..f10891d6cdf 100644 --- a/vllm/distributed/eplb/eplb_utils.py +++ b/vllm/distributed/eplb/eplb_utils.py @@ -61,25 +61,31 @@ class CpuGpuEvent: self._recorded.set() -def override_envs_for_eplb(parallel_config: ParallelConfig) -> None: +def override_envs_for_eplb( + parallel_config: ParallelConfig, + moe_backend: str | None = None, +) -> None: """ Override environment variables for EPLB when specific conditions are met. Args: parallel_config: The parallel configuration object. + moe_backend: The configured MoE backend (e.g. ``deep_gemm_mega_moe``). """ is_data_parallel = parallel_config.data_parallel_size > 1 is_eplb_enabled = parallel_config.enable_eplb async_eplb = parallel_config.eplb_config.use_async is_deepep_ll = parallel_config.all2all_backend == "deepep_low_latency" + is_mega_moe = moe_backend == "deep_gemm_mega_moe" is_nccl_based_eplb_communicator = parallel_config.eplb_config.communicator in ( "torch_nccl", "pynccl", ) - # Override NCCL_MAX_CTAS to avoid hangs when using async EPLB with the - # DeepEP low-latency backend. + # Override NCCL_MAX_CTAS to avoid hangs when EPLB's NCCL weight exchange + # contends with MoE backend's cooperative-launch on GPU SMs. # + # DeepEP low-latency: # The hang happens when two ranks interleave kernel launches differently # between NCCL collectives (used by async EPLB weight exchange) and DeepEP # low-latency (LL) kernels. DeepEP LL uses a cooperative launch and tries @@ -94,12 +100,14 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None: # Limiting NCCL occupancy via NCCL_MAX_CTAS leaves space for the DeepEP # cooperative kernel to launch and complete, breaking the deadlock. # See: https://github.com/deepseek-ai/DeepEP/issues/496 + # + # DeepGEMM Mega MoE also uses cooperative launch and will cause hang even + # with sync EPLB. if ( is_data_parallel and is_eplb_enabled - and is_deepep_ll - and async_eplb and is_nccl_based_eplb_communicator + and ((is_deepep_ll and async_eplb) or is_mega_moe) ): current_value_str = os.getenv("NCCL_MAX_CTAS") @@ -108,9 +116,10 @@ def override_envs_for_eplb(parallel_config: ParallelConfig) -> None: override_value = 8 os.environ["NCCL_MAX_CTAS"] = str(override_value) + backend = "deepep_low_latency" if is_deepep_ll else "deep_gemm_mega_moe" logger.info_once( f"EPLB: Setting NCCL_MAX_CTAS={override_value} " - "for expert parallel with NCCL-based EPLB communicator and " - "deepep_low_latency backend", + f"for expert parallel with NCCL-based EPLB communicator and " + f"cooperative MoE backend ({backend})", scope="global", ) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 13e58360c8b..33baed72795 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import typing -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, MutableSequence, Sequence from itertools import islice import regex as re @@ -15,6 +15,7 @@ from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) +from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.kernels.mhc.tilelang import ( hc_head_fused_kernel_tilelang, mhc_fused_post_pre_tilelang, @@ -23,6 +24,9 @@ from vllm.model_executor.kernels.mhc.tilelang import ( ) from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( fused_topk_bias, ) @@ -40,7 +44,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.interfaces import SupportsPP +from vllm.model_executor.models.interfaces import MixtureOfExperts, SupportsPP from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, @@ -144,6 +148,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): hidden_size: int, intermediate_size: int, prefix: str = "", + num_logical_experts: int | None = None, ): super().__init__() self.prefix = prefix @@ -156,6 +161,12 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.intermediate_size = intermediate_size self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens + self.num_logical_experts = ( + num_logical_experts if num_logical_experts is not None else num_experts + ) + + self.eplb_state = EplbLayerState() + weight_attrs = {"weight_loader": self.weight_loader} self.w13_weight = nn.Parameter( torch.zeros( @@ -206,10 +217,22 @@ class DeepseekV4MegaMoEExperts(nn.Module): self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None - def _map_global_expert_id(self, expert_id: int) -> int: - if expert_id < self.experts_start_idx or expert_id >= self.experts_end_idx: - return -1 - return expert_id - self.experts_start_idx + # Register in the static forward context so the custom-op wrapper + # can look up this module by name from within a torch.compile graph. + compilation_config = vllm_config.compilation_config + if prefix in compilation_config.static_forward_context: + raise ValueError(f"Duplicate layer name: {prefix}") + compilation_config.static_forward_context[prefix] = self + + def _map_global_expert_id(self, expert_id: int) -> list[int]: + """Return local (per-rank) slot offsets where logical expert + `expert_id` should land on this rank. + """ + physical_ids: list[int] = [] + for p in range(self.experts_start_idx, self.experts_end_idx): + if p % self.num_logical_experts == expert_id: + physical_ids.append(p - self.experts_start_idx) + return physical_ids def weight_loader( self, @@ -220,30 +243,38 @@ class DeepseekV4MegaMoEExperts(nn.Module): expert_id: int, return_success: bool = False, ) -> bool | None: - local_expert_id = self._map_global_expert_id(expert_id) - if local_expert_id == -1: + local_expert_ids = self._map_global_expert_id(expert_id) + if not local_expert_ids: return False if return_success else None - expert_data = param.data[local_expert_id] - if shard_id in ("w1", "w3"): - if "w13_" not in weight_name: - return False if return_success else None - shard_offset = 0 if shard_id == "w1" else self.intermediate_size - expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size) - elif shard_id == "w2": - if "w2_" not in weight_name: - return False if return_success else None - else: - raise ValueError(f"Unsupported expert shard id: {shard_id}") + loaded_any = False + for local_expert_id in local_expert_ids: + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + if "w13_" not in weight_name: + continue + shard_offset = 0 if shard_id == "w1" else self.intermediate_size + expert_data = expert_data.narrow( + 0, shard_offset, self.intermediate_size + ) + elif shard_id == "w2": + if "w2_" not in weight_name: + continue + else: + raise ValueError(f"Unsupported expert shard id: {shard_id}") - if expert_data.shape != loaded_weight.shape: - raise ValueError( - f"DeepSeek V4 MegaMoE expert weight shape mismatch for " - f"{weight_name}: parameter shard {tuple(expert_data.shape)} " - f"vs checkpoint {tuple(loaded_weight.shape)}" - ) - expert_data.copy_(loaded_weight) - return True if return_success else None + if expert_data.shape != loaded_weight.shape: + raise ValueError( + f"DeepSeek V4 MegaMoE expert weight shape mismatch for " + f"{weight_name}: parameter shard {tuple(expert_data.shape)} " + f"vs checkpoint {tuple(loaded_weight.shape)}" + ) + expert_data.copy_(loaded_weight) + loaded_any = True + + if return_success: + return loaded_any + return None @staticmethod def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: @@ -264,7 +295,9 @@ class DeepseekV4MegaMoEExperts(nn.Module): return self._check_runtime_supported() - import vllm.third_party.deep_gemm as deep_gemm + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() w13_scale = deep_gemm.transform_sf_into_required_layout( self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(), @@ -298,7 +331,9 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.w2_weight_scale = None def get_symm_buffer(self): - import vllm.third_party.deep_gemm as deep_gemm + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() group = get_ep_group().device_group device = torch.accelerator.current_device_index() @@ -324,6 +359,52 @@ class DeepseekV4MegaMoEExperts(nn.Module): self._symm_buffer_cache[key] = symm_buffer return symm_buffer + def set_eplb_state( + self, + moe_layer_idx: int, + expert_load_view: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + ) -> None: + self.eplb_state.set_layer_state( + moe_layer_idx, + expert_load_view, + logical_to_physical_map, + logical_replica_count, + ) + + def get_expert_weights(self) -> list[torch.Tensor]: + self.finalize_weights() + assert self._transformed_l1_weights is not None + assert self._transformed_l2_weights is not None + + def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor: + """Return a (num_local_experts, -1) view with contiguous memory layout.""" + assert t.shape[0] == self.num_local_experts + if t.is_contiguous(): + return t.view(self.num_local_experts, -1) + elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]: + # scales have shape (E, M, N) with memory layout (E, N, M) + back = torch.transpose(t, 1, 2) + assert back.is_contiguous() + return back.view(self.num_local_experts, -1) + + raise AssertionError( + f"DSv4 EPLB {name}: non-contiguous expert tensor with " + f"unexpected layout shape={tuple(t.shape)} " + f"stride={tuple(t.stride())} dtype={t.dtype}" + ) + + return [ + _to_eplb_view("l1_packed", self._transformed_l1_weights[0]), + _to_eplb_view("l1_scale", self._transformed_l1_weights[1]), + _to_eplb_view("l2_weight", self._transformed_l2_weights[0]), + _to_eplb_view("l2_scale", self._transformed_l2_weights[1]), + ] + + def update_expert_map(self) -> None: + pass + def forward( self, hidden_states: torch.Tensor, @@ -358,10 +439,27 @@ class DeepseekV4MegaMoEExperts(nn.Module): activation_clamp: float | None, fast_math: bool, ) -> None: - import vllm.third_party.deep_gemm as deep_gemm + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] + + # EPLB: map logical expert IDs to physical replicas and record load. + eplb_state = self.eplb_state + if eplb_state.logical_to_physical_map is not None: + assert eplb_state.expert_load_view is not None + assert eplb_state.logical_replica_count is not None + assert eplb_state.should_record_tensor is not None + topk_ids = eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=eplb_state.expert_load_view, + logical_to_physical_map=eplb_state.logical_to_physical_map, + logical_replica_count=eplb_state.logical_replica_count, + record_enabled=eplb_state.should_record_tensor, + ) + prepare_megamoe_inputs( hidden_states, topk_weights, @@ -493,17 +591,33 @@ class DeepseekV4MoE(nn.Module): self.ep_group = get_ep_group() self.ep_size = self.ep_group.world_size self.ep_rank = self.ep_group.rank_in_group - assert config.n_routed_experts % self.ep_size == 0 - self.n_local_experts = config.n_routed_experts // self.ep_size - self.experts_start_idx = self.ep_rank * self.n_local_experts - self.experts_end_idx = self.experts_start_idx + self.n_local_experts + eplb_config = vllm_config.parallel_config.eplb_config + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_routed_experts = config.n_routed_experts + self.n_shared_experts = config.n_shared_experts or 0 + self.n_logical_experts = self.n_routed_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + assert self.n_physical_experts % self.ep_size == 0, ( + f"n_physical_experts={self.n_physical_experts} must be divisible by " + f"ep_size={self.ep_size}. Adjust num_redundant_experts." + ) + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + self.n_local_experts = self.n_local_physical_experts + self.experts_start_idx = self.physical_expert_start + self.experts_end_idx = self.physical_expert_end self.experts = DeepseekV4MegaMoEExperts( vllm_config, - num_experts=config.n_routed_experts, - num_local_experts=self.n_local_experts, - experts_start_idx=self.experts_start_idx, + num_experts=self.n_physical_experts, + num_local_experts=self.n_local_physical_experts, + experts_start_idx=self.physical_expert_start, + num_logical_experts=self.n_logical_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, @@ -1242,7 +1356,44 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: ) -class DeepseekV4ForCausalLM(nn.Module, SupportsPP): +class DeepseekV4MixtureOfExperts(MixtureOfExperts): + moe_mlp_layers: list["DeepseekV4MoE"] + + def extract_moe_parameters(self, example_moe: "DeepseekV4MoE | None") -> None: + if example_moe is None: + self.num_moe_layers = 0 + self.num_expert_groups = 0 + self.num_logical_experts = 0 + self.num_physical_experts = 0 + self.num_local_physical_experts = 0 + self.num_routed_experts = 0 + self.num_shared_experts = 0 + self.num_redundant_experts = 0 + return + self.num_logical_experts = example_moe.n_logical_experts + self.num_physical_experts = example_moe.n_physical_experts + self.num_local_physical_experts = example_moe.n_local_physical_experts + self.num_routed_experts = example_moe.n_routed_experts + self.num_shared_experts = example_moe.n_shared_experts + self.num_redundant_experts = example_moe.n_redundant_experts + + def update_physical_experts_metadata( + self, + num_physical_experts: int, + num_local_physical_experts: int, + ) -> None: + assert self.num_local_physical_experts == num_local_physical_experts + self.num_physical_experts = num_physical_experts + self.num_local_physical_experts = num_local_physical_experts + self.num_redundant_experts = num_physical_experts - self.num_logical_experts + for moe in self.moe_mlp_layers: + moe.n_local_physical_experts = num_local_physical_experts + moe.n_physical_experts = num_physical_experts + moe.n_redundant_experts = self.num_redundant_experts + moe.experts.update_expert_map() + + +class DeepseekV4ForCausalLM(nn.Module, SupportsPP, DeepseekV4MixtureOfExperts): model_cls = DeepseekV4Model # Default mapper assumes the original FP4-expert checkpoint layout. @@ -1274,6 +1425,28 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP): self.model.make_empty_intermediate_tensors ) + self.set_moe_parameters() + + def set_moe_parameters(self) -> None: + self.expert_weights: MutableSequence[Sequence[torch.Tensor]] = [] + self.num_expert_groups = getattr(self.config, "n_group", 1) + self.num_moe_layers = self.config.num_hidden_layers + self.moe_layers: list[nn.Module] = [] + self.moe_mlp_layers: list[DeepseekV4MoE] = [] + example_moe: DeepseekV4MoE | None = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + if not isinstance(layer, DeepseekV4DecoderLayer): + continue + if isinstance(layer.ffn, DeepseekV4MoE): + example_moe = layer.ffn + self.moe_mlp_layers.append(layer.ffn) + self.moe_layers.append(layer.ffn.experts) + + self.num_moe_layers = len(self.moe_layers) + self.extract_moe_parameters(example_moe) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.embed_input_ids(input_ids) diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 6b89f5c3320..4252ce87754 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -140,6 +140,7 @@ _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _transform_sf_into_required_layout_impl: Callable[..., Any] | None = None +@functools.cache def _import_deep_gemm(): """Import the deep_gemm module. diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 121fc69f532..259cd05554c 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -1147,7 +1147,10 @@ def init_worker_distributed_environment( from vllm.model_executor.layers.batch_invariant import init_batch_invariance init_batch_invariance() - override_envs_for_eplb(parallel_config) + override_envs_for_eplb( + parallel_config, + moe_backend=getattr(vllm_config.kernel_config, "moe_backend", None), + ) set_custom_all_reduce(not parallel_config.disable_custom_all_reduce) init_method = distributed_init_method or "env://"