diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index b5b8b3c8e7..177d00dc4d 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -794,17 +794,6 @@ MTP 多预测参数 增加此值允许更多预测，但确保模型与指定的步数兼容。 目前 deepseekv3/r1 模型仅支持 1 步 -DeepSeek 冗余专家参数 ---------------------- - -.. option:: --ep_redundancy_expert_config_path - - 冗余专家配置的路径。可用于 deepseekv3 模型。 - -.. option:: --auto_update_redundancy_expert - - 是否通过在线专家使用计数器为 deepseekv3 模型更新冗余专家。 - 监控和日志参数 -------------- diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index dfcd5db5fa..5e7557cff0 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -825,17 +825,6 @@ MTP Multi-Prediction Parameters Increasing this value allows more predictions, but ensure the model is compatible with the specified number of steps. Currently deepseekv3/r1 models only support 1 step -DeepSeek Redundant Expert Parameters ------------------------------------- - -.. option:: --ep_redundancy_expert_config_path - - Path to redundant expert configuration. Can be used for deepseekv3 models. - -.. option:: --auto_update_redundancy_expert - - Whether to update redundant experts for deepseekv3 models through online expert usage counters. - Monitoring and Logging Parameters --------------------------------- diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 34a1c9c71c..18b5d1c171 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -28,7 +28,7 @@ from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed from lightllm.utils.log_utils import init_logger from lightllm.utils.dist_utils import get_dp_world_size -from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory +from lightllm.utils.profile_max_tokens import profile_model_memory from lightllm.utils.envs_utils import ( get_env_start_args, get_llm_data_type, @@ -109,7 +109,7 @@ def __init__(self, kvargs): self._init_quant() enable_weight_cpu_backup = self.args.enable_weight_cpu_backup - with profile_mtp_weight_memory(self): + with profile_model_memory(self): with self.torch_memory_saver.region(tag=MemoryTag.WEIGHT, enable_cpu_backup=enable_weight_cpu_backup): self._init_weights() with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): @@ -148,6 +148,7 @@ def __init__(self, kvargs): def _init_config(self): with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file: self.config = json.load(json_file) + self.config["is_mtp_draft_model"] = self.is_mtp_draft_model # rename keys repair_config(self.config, same_names=["num_attention_heads", "n_head"]) repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"]) @@ -631,7 +632,6 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_balance.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_balance.py new file mode 100644 index 0000000000..70435146c5 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_balance.py @@ -0,0 +1,14 @@ +from dataclasses import dataclass + + +@dataclass(slots=True) +class PrefillEPBalanceCounters: + """Cumulative CPU loads for one EP MoE layer's completed prefill dispatches.""" + + route_load: int = 0 + compute_load: int = 0 + + def accumulate(self, route_load: int, compute_load: int): + """Accumulate exact route and alignment-expanded compute loads for one prefill dispatch.""" + self.route_load += route_load + self.compute_load += compute_load diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py deleted file mode 100644 index ed7932e56e..0000000000 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py +++ /dev/null @@ -1,199 +0,0 @@ -import numpy as np -import torch -from .fused_moe_weight import FusedMoeWeight -from lightllm.common.quantization.quantize_method import WeightPack -from lightllm.utils.log_utils import init_logger -from typing import Dict - -logger = init_logger(__name__) - - -class FusedMoeWeightEPAutoRedundancy: - def __init__( - self, - ep_fused_moe_weight: FusedMoeWeight, - ) -> None: - super().__init__() - self._ep_w = ep_fused_moe_weight - self.redundancy_expert_num = self._ep_w.redundancy_expert_num - - def clear_counter(self): - self._ep_w.routed_expert_counter_tensor.fill_(0) - return - - def prepare_redundancy_experts( - self, - ): - expert_counter = self._ep_w.routed_expert_counter_tensor.detach().cpu().numpy() - logger.info( - f"layer_index {self._ep_w.layer_num_} global_rank {self._ep_w.global_rank_}" - f" expert_counter: {expert_counter}" - ) - self._ep_w.routed_expert_counter_tensor.fill_(0) - ep_n_routed_experts = self._ep_w.n_routed_experts // self._ep_w.global_world_size - start_expert_id = ep_n_routed_experts * self._ep_w.global_rank_ - no_redundancy_expert_ids = list(range(start_expert_id, start_expert_id + ep_n_routed_experts)) - - # 统计 0 rank 上的全局 topk 冗余信息，帮助导出一份全局可用的静态使用的冗余专家静态配置。 - if self._ep_w.global_rank_ == 0: - # int(e) for serialization, int64 can not be serialized by json.dump. - topk_redundancy_expert_ids = list(int(e) for e in np.argsort(expert_counter)[-self.redundancy_expert_num :]) - else: - topk_redundancy_expert_ids = None - - # 不要选中当前已经存在的非冗余专家作为冗余专家 - expert_counter[no_redundancy_expert_ids] = 0 - - self.redundancy_expert_ids = list(np.argsort(expert_counter)[-self.redundancy_expert_num :]) - logger.info( - f"layer_index {self._ep_w.layer_num_} global_rank {self._ep_w.global_rank_}" - f" new select redundancy_expert_ids : {self.redundancy_expert_ids}" - ) - - # 准备加载过度变量。 - self.experts_up_projs = [None] * self.redundancy_expert_num - self.experts_gate_projs = [None] * self.redundancy_expert_num - self.experts_up_proj_scales = [None] * self.redundancy_expert_num - self.experts_gate_proj_scales = [None] * self.redundancy_expert_num - self.w2_list = [None] * self.redundancy_expert_num - self.w2_scale_list = [None] * self.redundancy_expert_num - self.w13 = [None, None] # weight, weight_scale - self.w2 = [None, None] # weight, weight_scale - return topk_redundancy_expert_ids - - def load_hf_weights(self, weights): - # 加载冗余专家的权重参数 - for i, redundant_expert_id in enumerate(self.redundancy_expert_ids): - i_experts = redundant_expert_id - w1_weight = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w1_weight_name}.weight" - w2_weight = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w2_weight_name}.weight" - w3_weight = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w3_weight_name}.weight" - if w1_weight in weights: - self.experts_gate_projs[i] = weights[w1_weight] - if w3_weight in weights: - self.experts_up_projs[i] = weights[w3_weight] - if w2_weight in weights: - self.w2_list[i] = weights[w2_weight] - - self._load_weight_scale(weights) - self._fuse() - - def _fuse(self): - self._fuse_weight_scale() - - with self._ep_w.lock: - if ( - hasattr(self, "experts_up_projs") - and None not in self.experts_up_projs - and None not in self.experts_gate_projs - and None not in self.w2_list - ): - gate_out_dim, gate_in_dim = self.experts_gate_projs[0].shape - up_out_dim, up_in_dim = self.experts_up_projs[0].shape - assert gate_in_dim == up_in_dim - dtype = self.experts_gate_projs[0].dtype - total_expert_num = self.redundancy_expert_num - - w13 = torch.empty((total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu") - - for i_experts in range(self.redundancy_expert_num): - w13[i_experts, 0:gate_out_dim:, :] = self.experts_gate_projs[i_experts] - w13[i_experts, gate_out_dim:, :] = self.experts_up_projs[i_experts] - - inter_shape, hidden_size = self.w2_list[0].shape[0], self.w2_list[0].shape[1] - w2 = torch._utils._flatten_dense_tensors(self.w2_list).view(len(self.w2_list), inter_shape, hidden_size) - if self._ep_w.quant_method._check_weight_need_quanted(weight=w13): - w13_pack, _ = self._ep_w.quant_method.create_moe_weight( - out_dims=[gate_out_dim + up_out_dim], - in_dim=1, - dtype=self._ep_w.data_type_, - device_id=self._ep_w.device_id_, - num_experts=self.redundancy_expert_num, - ) - self._ep_w.quant_method.quantize(w13, w13_pack) - w2_pack, _ = self._ep_w.quant_method.create_moe_weight( - out_dims=[inter_shape], - in_dim=hidden_size, - dtype=self._ep_w.data_type_, - device_id=self._ep_w.device_id_, - num_experts=self.redundancy_expert_num, - ) - self._ep_w.quant_method.quantize(w2, w2_pack) - - self.w13[0] = w13_pack.weight - self.w13[1] = w13_pack.weight_scale - self.w2[0] = w2_pack.weight - self.w2[1] = w2_pack.weight_scale - else: - self.w13[0] = w13 - self.w2[0] = w2 - delattr(self, "w2_list") - delattr(self, "experts_up_projs") - delattr(self, "experts_gate_projs") - - def _fuse_weight_scale(self): - with self._ep_w.lock: - if ( - hasattr(self, "experts_up_proj_scales") - and None not in self.experts_up_proj_scales - and None not in self.experts_gate_proj_scales - and None not in self.w2_scale_list - ): - gate_out_dim, gate_in_dim = self.experts_gate_proj_scales[0].shape - up_out_dim, up_in_dim = self.experts_up_proj_scales[0].shape - assert gate_in_dim == up_in_dim - dtype = self.experts_gate_proj_scales[0].dtype - total_expert_num = self.redundancy_expert_num - w13_scale = torch.empty( - (total_expert_num, gate_out_dim + up_out_dim, gate_in_dim), dtype=dtype, device="cpu" - ) - for i_experts in range(self.redundancy_expert_num): - w13_scale[i_experts, 0:gate_out_dim:, :] = self.experts_gate_proj_scales[i_experts] - w13_scale[i_experts, gate_out_dim:, :] = self.experts_up_proj_scales[i_experts] - - inter_shape, hidden_size = self.w2_scale_list[0].shape[0], self.w2_scale_list[0].shape[1] - w2_scale = torch._utils._flatten_dense_tensors(self.w2_scale_list).view( - len(self.w2_scale_list), inter_shape, hidden_size - ) - self.w13[1] = w13_scale - self.w2[1] = w2_scale - delattr(self, "w2_scale_list") - delattr(self, "experts_up_proj_scales") - delattr(self, "experts_gate_proj_scales") - - def _load_weight_scale(self, weights: Dict[str, torch.Tensor]) -> None: - # 加载冗余专家的scale参数 - for i, redundant_expert_id in enumerate(self.redundancy_expert_ids): - i_experts = redundant_expert_id - weight_scale_suffix = self._ep_w.quant_method.weight_scale_suffix - w1_scale = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w1_weight_name}.{weight_scale_suffix}" - w2_scale = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w2_weight_name}.{weight_scale_suffix}" - w3_scale = f"{self._ep_w.weight_prefix}.{i_experts}.{self._ep_w.w3_weight_name}.{weight_scale_suffix}" - if w1_scale in weights: - self.experts_gate_proj_scales[i] = weights[w1_scale] - if w3_scale in weights: - self.experts_up_proj_scales[i] = weights[w3_scale] - if w2_scale in weights: - self.w2_scale_list[i] = weights[w2_scale] - - def commit(self): - for weight_pack, source in [(self._ep_w.w13, self.w13), (self._ep_w.w2, self.w2)]: - for index, dest_tensor in enumerate([weight_pack.weight, weight_pack.weight_scale]): - if dest_tensor is not None: - assert isinstance( - dest_tensor, torch.Tensor - ), f"dest_tensor should be a torch.Tensor, but got {type(dest_tensor)}" - if index == 1 and dest_tensor.dtype == torch.int32 and source[index].dtype == torch.float32: - self._ep_w.quant_method.load_weight_scale( - source[index], - WeightPack( - weight=weight_pack.weight[-self.redundancy_expert_num :], - weight_scale=dest_tensor[-self.redundancy_expert_num :], - ), - ) - else: - dest_tensor[-self.redundancy_expert_num :, :, :] = source[index][:, :, :] - - self._ep_w.redundancy_expert_ids_tensor.copy_( - torch.tensor(self.redundancy_expert_ids, dtype=torch.int64, device="cpu") - ) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_placement.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_placement.py new file mode 100644 index 0000000000..a62a5c3d08 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_placement.py @@ -0,0 +1,455 @@ +from typing import Dict, Tuple +from math import isqrt +import torch + + +def get_max_reachable_replicas(placement: torch.Tensor, num_logical_experts: int, node_world_size: int) -> int: + """Bound hash sampling by the node holdings preserved by the full planner. + + A node needs one slot per held expert. Its other slots bound how many + copies of any one expert it can own. Only experts absent from a source + node need a global-fallback histogram. + """ + ranks, capacity = placement.shape + if ranks % node_world_size: + raise ValueError("EP ranks must form complete nodes") + holdings = torch.zeros((ranks // node_world_size, num_logical_experts), dtype=torch.bool) + limits = [] + for node in range(len(holdings)): + ids = placement[node * node_world_size : (node + 1) * node_world_size].unique() + holdings[node, ids] = True + limits.append(min(node_world_size, node_world_size * capacity - len(ids) + 1)) + limits = torch.tensor(limits, dtype=torch.int64) + maximum = int(limits.max()) + absent_somewhere = ~holdings.all(0) + if absent_somewhere.any(): + global_limits = (holdings * limits[:, None]).sum(0) + maximum = max(maximum, int(global_limits[absent_somewhere].max())) + return maximum + + +def build_initial_redundant_expert_ids( + num_logical_experts: int, + num_ranks: int, + num_redundant_experts_per_rank: int, +) -> torch.Tensor: + """Build a deterministic initial placement without local duplicates.""" + assert num_logical_experts % num_ranks == 0 + num_experts_per_rank = num_logical_experts // num_ranks + assert 0 <= num_redundant_experts_per_rank <= num_logical_experts - num_experts_per_rank + + # 初始化结果确定，不依赖随机数。 + # 每个 rank 不会复制自己原本拥有的 expert。 + # 同一个 rank 的冗余槽位不会重复。 + # 最后一个 rank 通过取模自然回绕。 + rank_offsets = torch.arange(1, num_ranks + 1, dtype=torch.int64)[:, None] * num_experts_per_rank + expert_offsets = torch.arange(num_redundant_experts_per_rank, dtype=torch.int64) + return (rank_offsets + expert_offsets) % num_logical_experts + + +def expand_redundant_placement(placement: torch.Tensor, num_logical_experts: int) -> torch.Tensor: + """Prepend the initial contiguous primary rows to a redundant-only layout.""" + num_ranks = placement.shape[-2] + assert num_logical_experts % num_ranks == 0 + primary = torch.arange(num_logical_experts, device=placement.device, dtype=placement.dtype) + primary = primary.reshape(num_ranks, -1).expand(*placement.shape[:-2], num_ranks, -1) + return torch.cat((primary, placement), dim=-1) + + +def validate_physical_placement(placement: torch.Tensor, num_logical_experts: int) -> None: + """Reject layouts that cannot be represented by one replica per rank.""" + if placement.ndim != 3 or placement.dtype not in (torch.int32, torch.int64): + raise ValueError("physical placement must be integer [layers, ranks, slots]") + if min(placement.shape) <= 0 or not torch.all((placement >= 0) & (placement < num_logical_experts)): + raise ValueError("physical placement has empty dimensions or invalid expert IDs") + if num_logical_experts % placement.shape[1] != 0: + raise ValueError("logical expert count must be divisible by the EP world size") + ordered = placement.sort(dim=-1).values + if torch.any(ordered[..., 1:] == ordered[..., :-1]): + raise ValueError("physical placement contains duplicate experts on a rank") + if not torch.all(_expert_locations(placement, num_logical_experts).any(dim=-1)): + raise ValueError("physical placement must cover every logical expert") + + +def build_logical_to_physical_map( + physical_expert_ids: torch.Tensor, # 完整物理布局，shape 为 [num_ranks, num_physical_experts_per_rank]。 + num_logical_experts: int, # 逻辑 expert 的总数。 + source_rank: int | None = None, # 可选的全局源 rank；传入时优先选择同节点副本。 + node_world_size: int | None = None, # 单个节点包含的 rank 数；source_rank 非空时必填。 + replica_rotation_rank: int | None = None, +) -> Tuple[ + torch.Tensor, torch.Tensor +]: # logical_to_physical [num_logical_experts, num_ranks], replica_counts [num_logical_experts] + """构建单层完整物理布局的逻辑 expert 到副本映射。""" + + logical_to_physical, replica_counts = _build_layer_maps( + physical_expert_ids.unsqueeze(0), + num_logical_experts, + source_rank=source_rank, + node_world_size=node_world_size, + replica_rotation_rank=replica_rotation_rank, + ) + return logical_to_physical.squeeze(0), replica_counts.squeeze(0) + + +def build_logical_to_physical_maps_for_layers( + physical_expert_ids_by_layer: torch.Tensor, # [num_layers, num_ranks, num_physical_experts_per_rank] + num_logical_experts: int, # 逻辑 expert 的总数。 + source_rank: int | None = None, # 可选的全局源 rank；传入时优先选择同节点副本。 + node_world_size: int | None = None, # 单个节点包含的 rank 数；source_rank 非空时必填。 + replica_rotation_rank: int | None = None, +) -> Tuple[ + torch.Tensor, # logical_to_physical, shape [num_layers, num_logical_experts, num_ranks] + torch.Tensor, # replica_counts, shape [num_layers, num_logical_experts] +]: + """为调用方传入的多个指定层构建逻辑 expert 到物理副本的 CPU int32 映射。 + + 第一维是层，不是请求或 token 的 batch，也不会自动处理模型中的其他层。 + 输入包含全部物理行，副本按 rank-major、slot-major 的稳定顺序排列。 + source_rank 决定本地节点；replica_rotation_rank 独立控制轮转，默认使用 source_rank。 + """ + return _build_layer_maps( + physical_expert_ids_by_layer, + num_logical_experts, + source_rank=source_rank, + node_world_size=node_world_size, + replica_rotation_rank=replica_rotation_rank, + ) + + +def select_improving_placements( + expert_load: torch.Tensor, + current_placement: torch.Tensor, + candidate_placement: torch.Tensor, + *, + rebalance_gain_threshold: float, + expert_alignment: int | None = None, + node_world_size: int | None = None, + replica_hash_load: torch.Tensor | None = None, +) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, float | int], torch.Tensor, torch.Tensor]: + """Select better layers and return current/final rank loads without re-estimation.""" + if not 0.0 <= rebalance_gain_threshold <= 1.0: + raise ValueError("rebalance_gain_threshold must be between 0.0 and 1.0") + assert current_placement.shape == candidate_placement.shape + current_rank_load = _estimate_rank_load( + expert_load, current_placement, expert_alignment, node_world_size, replica_hash_load=replica_hash_load + ) + candidate_rank_load = _estimate_rank_load( + expert_load, candidate_placement, expert_alignment, node_world_size, replica_hash_load=replica_hash_load + ) + if expert_load.ndim == 2: + current_critical = current_rank_load.max(dim=1).values + candidate_critical = candidate_rank_load.max(dim=1).values + else: + current_critical = current_rank_load.max(dim=2).values.sum(dim=0) + candidate_critical = candidate_rank_load.max(dim=2).values.sum(dim=0) + # Each changed layer must reduce its own critical load. All selected + # changes must then collectively meet the configured model-level + # critical-load reduction threshold, avoiding low-gain migrations. + improved = candidate_critical < current_critical + selected = current_placement.clone() + selected[improved] = candidate_placement[improved] + if current_rank_load.ndim == 2: + selected_rank_load = torch.where(improved[:, None], candidate_rank_load, current_rank_load) + else: + selected_rank_load = torch.where(improved[None, :, None], candidate_rank_load, current_rank_load) + model_current_critical = current_critical.sum() + if expert_load.ndim == 2: + model_current_mean = current_rank_load.mean(dim=1).sum() + model_selected_critical = selected_rank_load.max(dim=1).values.sum() + model_selected_mean = selected_rank_load.mean(dim=1).sum() + else: + model_current_mean = current_rank_load.mean(dim=2).sum() + model_selected_critical = selected_rank_load.max(dim=2).values.sum() + model_selected_mean = selected_rank_load.mean(dim=2).sum() + model_ratio = model_current_critical / model_current_mean.clamp_min(1.0) + candidate_model_ratio = model_selected_critical / model_selected_mean.clamp_min(1.0) + candidate_rebalance_gain = (model_current_critical - model_selected_critical) / model_current_critical.clamp_min( + 1.0 + ) + metrics = { + "model_imbalance_ratio": float(model_ratio.item()), + "candidate_model_imbalance_ratio": float(candidate_model_ratio.item()), + "candidate_rebalance_gain": float(candidate_rebalance_gain.item()), + "candidate_changed_layer_count": int(improved.sum().item()), + } + if candidate_rebalance_gain >= rebalance_gain_threshold: + return selected, improved, metrics, current_rank_load, selected_rank_load + return ( + current_placement.clone(), + torch.zeros_like(improved), + metrics, + current_rank_load, + current_rank_load, + ) + + +def _select_source_node_replicas( + logical_to_physical: torch.Tensor, + replica_counts: torch.Tensor, + *, + source_rank: int, + node_world_size: int, + num_physical_experts_per_rank: int, + replica_positions: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """按来源节点筛选全局候选副本，并将结果压缩为连续前缀。 + + ``logical_to_physical[layer, logical_expert]`` 的前 ``replica_counts`` 个位置 + 是有效副本，尾部 ``-1`` 表示没有副本；正常输入和输出都不会在有效前缀中 + 出现 ``-1``。连续的 ``node_world_size`` 个 rank 构成一个节点，来源节点为 + ``source_rank // node_world_size``。对每个逻辑 expert，若来源节点有副本， + 则只保留该节点的全部副本，远程副本（包括远程主副本）全部排除；否则保留 + 所有全局有效副本作为回退，避免候选为空。筛选可能选中原有效前缀中不连续的 + 位置，因此按原相对顺序将选中副本复制到新映射的连续前缀，其余位置填 ``-1``， + 并返回新的有效数量；不会原地修改输入。此函数只筛选和压缩候选集合，不做 + 负载优化、最终副本选择或按 source rank 轮转，轮转由后续函数完成。 + + 输出： + - ``compact_maps_by_layer``：与 ``logical_to_physical`` 同 shape + ``[num_layers, num_logical_experts, num_ranks]``，dtype/device 相同；每行 + 为筛选后副本的连续有效前缀，尾部为 ``-1``。 + - ``selected_counts_by_layer``：输入同 device 的 ``int32`` Tensor，shape 为 + ``[num_layers, num_logical_experts]``；每个值是对应输出行的有效前缀长度。 + """ + num_layers, num_logical_experts, _max_replicas = logical_to_physical.shape + source_node = source_rank // node_world_size + output_positions = replica_positions.view(1, 1, -1) + valid = output_positions < replica_counts.unsqueeze(-1) + local = valid & ( + torch.div( + logical_to_physical, + num_physical_experts_per_rank * node_world_size, + rounding_mode="floor", + ) + == source_node + ) + selected = torch.where(local.any(dim=2, keepdim=True), local, valid) + selected_counts_by_layer = selected.sum(dim=2, dtype=torch.int32) + + compact_maps_by_layer = torch.full_like(logical_to_physical, -1) + selected_positions = selected.cumsum(dim=2) - 1 + layers = torch.arange(num_layers, dtype=torch.int64).view(-1, 1, 1).expand_as(selected) + experts = torch.arange(num_logical_experts, dtype=torch.int64).view(1, -1, 1).expand_as(selected) + compact_maps_by_layer[layers[selected], experts[selected], selected_positions[selected]] = logical_to_physical[ + selected + ] + return compact_maps_by_layer, selected_counts_by_layer + + +def _rotate_selected_replicas( + compact_maps_by_layer: torch.Tensor, + selected_count_by_layer: torch.Tensor, + *, + source_rank: int, + replica_positions: torch.Tensor, +) -> torch.Tensor: + """根据 source_rank 循环调整每个逻辑 expert 的候选副本顺序，让不同源 rank 优先使用不同副本，同时保持候选副本集合和副本数量不变""" + output_positions = replica_positions.view(1, 1, -1) + selected_count64_by_layer = selected_count_by_layer.to(torch.int64).unsqueeze(-1) + rotation_by_layer = source_rank % selected_count64_by_layer + source_positions_by_layer = (output_positions + rotation_by_layer) % selected_count64_by_layer + maps_by_layer = compact_maps_by_layer.gather(2, source_positions_by_layer) + maps_by_layer.masked_fill_(output_positions >= selected_count64_by_layer, -1) + return maps_by_layer + + +def _build_layer_maps( + physical_expert_ids_by_layer: torch.Tensor, + num_logical_experts: int, # 逻辑 expert 的总数。 + source_rank: int | None = None, # 可选的全局源 rank；传入时优先选择同节点副本。 + node_world_size: int | None = None, # 单个节点包含的 rank 数；source_rank 非空时必填。 + replica_rotation_rank: int | None = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + """构建调用方指定层的映射；全局构建、节点筛选和轮转均在此完成。""" + if physical_expert_ids_by_layer.ndim != 3: + raise ValueError("physical_expert_ids_by_layer must be [layers, ranks, num_physical_experts_per_rank]") + num_ranks, num_physical_experts_per_rank = physical_expert_ids_by_layer.shape[1:] + assert num_logical_experts % num_ranks == 0 + validate_physical_placement(physical_expert_ids_by_layer, num_logical_experts) + physical = physical_expert_ids_by_layer.cpu().to(torch.int64) + layers = physical.shape[0] + logical_to_physical = torch.full((layers, num_logical_experts, num_ranks), -1, dtype=torch.int32) + replica_counts = torch.zeros((layers, num_logical_experts), dtype=torch.int32) + layer_ids = torch.arange(layers)[:, None] + for rank in range(num_ranks): + experts = physical[:, rank] + positions = replica_counts[layer_ids, experts].to(torch.int64) + logical_to_physical[layer_ids, experts, positions] = rank * num_physical_experts_per_rank + torch.arange( + num_physical_experts_per_rank, dtype=torch.int32 + ) + replica_counts[layer_ids, experts] += 1 + if source_rank is None: + if replica_rotation_rank is not None: + raise ValueError("replica_rotation_rank requires source_rank") + return logical_to_physical, replica_counts + + assert node_world_size is not None and 0 < node_world_size <= num_ranks + assert num_ranks % node_world_size == 0 and 0 <= source_rank < num_ranks + rotation_rank = source_rank if replica_rotation_rank is None else replica_rotation_rank + if not 0 <= rotation_rank < num_ranks: + raise ValueError("replica_rotation_rank must be an EP rank") + replica_positions = torch.arange(num_ranks, dtype=torch.int64) + compact_maps_by_layer, selected_counts_by_layer = _select_source_node_replicas( + logical_to_physical, + replica_counts, + source_rank=source_rank, + node_world_size=node_world_size, + num_physical_experts_per_rank=num_physical_experts_per_rank, + replica_positions=replica_positions, + ) + return ( + _rotate_selected_replicas( + compact_maps_by_layer, + selected_counts_by_layer, + source_rank=rotation_rank, + replica_positions=replica_positions, + ), + selected_counts_by_layer, + ) + + +def _estimate_rank_load( + expert_load: torch.Tensor, + physical_expert_ids: torch.Tensor, + expert_alignment: int | None = None, + node_world_size: int | None = None, + *, + replica_hash_load: torch.Tensor | None = None, +) -> torch.Tensor: + """Estimate runtime source-node-local routing load per physical expert. + + ``expert_load`` accepts the historic ``[layers, experts]`` and + ``[samples, layers, experts]`` forms, which are both one source node, and + the distributed ``[samples, layers, source_nodes, experts]`` form. Source + loads are kept separate until they are assigned to physical replicas, then + combined before applying the per-expert alignment used by DeepEP. + """ + source_load, squeeze_sample, node_world_size = _as_source_node_load( + expert_load, physical_expert_ids.shape[1], node_world_size + ) + num_samples, num_layers, num_nodes, num_logical_experts = source_load.shape + assert physical_expert_ids.ndim == 3 and physical_expert_ids.shape[0] == num_layers + num_ranks, num_physical_experts_per_rank = physical_expert_ids.shape[1:] + assert num_logical_experts % num_ranks == 0 + if expert_alignment is not None: + assert expert_alignment > 0 + + rank_load = _expert_rank_load_all( + source_load, + _expert_locations(physical_expert_ids, num_logical_experts), + num_nodes, + node_world_size, + expert_alignment, + replica_hash_load=replica_hash_load, + ).sum(dim=2) + return rank_load.squeeze(0) if squeeze_sample else rank_load + + +def _as_source_node_load( + expert_load: torch.Tensor, num_ranks: int, node_world_size: int | None +) -> Tuple[torch.Tensor, bool, int]: + """Normalize load to ``[samples, layers, source_nodes, experts]``.""" + assert expert_load.ndim in (2, 3, 4) + squeeze_sample = expert_load.ndim == 2 + if expert_load.ndim == 2: + source_load = expert_load.unsqueeze(0).unsqueeze(2) + elif expert_load.ndim == 3: + source_load = expert_load.unsqueeze(2) + else: + source_load = expert_load + # Historic 2D/3D loads represent one source node containing every rank. + if expert_load.ndim < 4: + return source_load, squeeze_sample, num_ranks + return source_load, squeeze_sample, _resolve_node_world_size(source_load, num_ranks, node_world_size) + + +def _resolve_node_world_size(expert_load: torch.Tensor, num_ranks: int, node_world_size: int | None) -> int: + """Validate production [samples, layers, source_nodes, experts] planner loads.""" + assert expert_load.ndim == 4 + num_nodes = expert_load.shape[2] + if node_world_size is None: + assert num_ranks % num_nodes == 0 + node_world_size = num_ranks // num_nodes + assert 0 < node_world_size <= num_ranks and num_ranks % node_world_size == 0 + assert num_nodes == num_ranks // node_world_size + return node_world_size + + +def _expert_locations(physical_expert_ids: torch.Tensor, num_logical_experts: int) -> torch.Tensor: + """Return ``[layer, logical expert, rank]`` physical-copy occupancy.""" + num_layers, num_ranks, num_physical_experts_per_rank = physical_expert_ids.shape + assert num_logical_experts % num_ranks == 0 + locations = torch.zeros( + (num_layers, num_logical_experts, num_ranks), + dtype=torch.bool, + device=physical_expert_ids.device, + ) + layers = torch.arange(num_layers, device=locations.device)[:, None] + ranks = torch.arange(num_ranks, device=locations.device).repeat_interleave(num_physical_experts_per_rank)[None, :] + physical_ids = physical_expert_ids.reshape(num_layers, -1) + valid = physical_ids >= 0 + if torch.any(valid): + expanded_layers = layers.expand_as(physical_ids) + expanded_ranks = ranks.expand_as(physical_ids) + locations[ + expanded_layers[valid], + physical_ids[valid], + expanded_ranks[valid], + ] = True + return locations + + +def _source_route(slots: torch.Tensor, num_nodes: int, node_world_size: int) -> torch.Tensor: + """Route each source node to its local copies, or all copies as fallback.""" + num_ranks = slots.shape[-1] + assert num_ranks % node_world_size == 0 and num_nodes == num_ranks // node_world_size + rank_nodes = torch.arange(num_ranks, device=slots.device) // node_world_size + source_nodes = torch.arange(num_nodes, device=slots.device) + copies = slots.unsqueeze(-3).expand(*slots.shape[:-2], num_nodes, *slots.shape[-2:]) + rank_node_shape = (1,) * slots.ndim + (num_ranks,) + source_node_shape = (1,) * (slots.ndim - 2) + (num_nodes, 1, 1) + local = copies & (rank_nodes.reshape(rank_node_shape) == source_nodes.reshape(source_node_shape)) + selected = torch.where(local.any(dim=-1, keepdim=True), local, copies) + return selected.to(torch.float64) / selected.sum(dim=-1, keepdim=True) + + +def _expert_rank_load_all( + source_load: torch.Tensor, + locations: torch.Tensor, + num_nodes: int, + node_world_size: int, + expert_alignment: int | None, + *, + replica_hash_load: torch.Tensor | None = None, +) -> torch.Tensor: + """Return aligned ``[samples, layers, expert, rank]`` contributions.""" + if replica_hash_load is None: + route = _source_route(locations, num_nodes, node_world_size) + physical_load = torch.einsum("slne,lner->sler", source_load.to(torch.float64), route) + else: + if replica_hash_load.shape[:-1] != source_load.shape: + raise ValueError("replica hash samples must match logical source loads") + ranks = locations.shape[-1] + channels = replica_hash_load.shape[-1] + max_replicas = (isqrt(1 + 8 * channels) - 1) // 2 + if max_replicas * (max_replicas + 1) // 2 != channels: + raise ValueError("replica hash samples must contain complete replica buckets") + # Select exactly the same rank-major candidate prefix as runtime maps. + owners = locations[:, None].expand(-1, num_nodes, -1, -1) + rank_nodes = torch.arange(ranks, device=locations.device) // node_world_size + source_nodes = torch.arange(num_nodes, device=locations.device) + local = owners & (rank_nodes[None, None, None, :] == source_nodes[None, :, None, None]) + selected = torch.where(local.any(-1, keepdim=True), local, owners) + copies = selected.sum(-1, keepdim=True) + if int(copies.max()) > max_replicas: + raise ValueError("candidate exceeds sampled reachable replica counts") + ordinal = selected.cumsum(-1) - 1 + channels = (copies * (copies - 1) // 2 + ordinal).clamp_min(0) + per_source = torch.gather( + replica_hash_load.to(torch.float64), -1, channels[None].expand(source_load.shape[0], -1, -1, -1, -1) + ) + physical_load = (per_source * selected[None]).sum(dim=2) + if expert_alignment is not None: + physical_load = torch.ceil(physical_load / expert_alignment) * expert_alignment + return physical_load diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_planner.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_planner.py new file mode 100644 index 0000000000..2f18e80965 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_planner.py @@ -0,0 +1,313 @@ +"""Bounded CPU candidate search, scored with the runtime's node-local routing. + +The score is samplewise aligned critical token load, not a prediction of latency. +Full layouts preserve each node's expert set; cross-node placement is deliberately +left to a future planner with a calibrated communication cost model. +""" + +import torch + +from .eplb_placement import ( + _estimate_rank_load, + _expert_locations, + _expert_rank_load_all, + _resolve_node_world_size, + validate_physical_placement, +) + + +def _critical_load(load, placement, alignment, node_world_size, replica_hash_load=None): + return ( + _estimate_rank_load(load, placement, alignment, node_world_size, replica_hash_load=replica_hash_load) + .max(dim=-1) + .values.sum(dim=0) + ) + + +def refine_placement_candidates( + expert_load: torch.Tensor, + current: torch.Tensor, + candidates: list[torch.Tensor], + *, + expert_alignment: int | None = None, + node_world_size: int | None = None, + stickiness: float = 0.0, + search_steps: int = 4, + candidate_budget: int = 256, + replica_hash_load: torch.Tensor | None = None, +) -> torch.Tensor: + """Keep the best layer of each seed and its bounded local search result. + + Current placement is always included and wins ties. Each search step scores + at most ``candidate_budget`` replacements/swaps per layer. Only the two + affected experts are re-evaluated for each move, bounding temporary memory. + Stickiness is a minimum cumulative improvement over a search seed, in + aligned-token units; it must not block a series of individually small moves. + """ + if search_steps < 0 or candidate_budget <= 0 or not 0 <= stickiness <= 1: + raise ValueError("invalid EPLB search budget or stickiness") + node_world_size = _resolve_node_world_size(expert_load, current.shape[1], node_world_size) + load = expert_load.to(device="cpu", dtype=torch.float64) + best = current.cpu().clone() + if replica_hash_load is not None: + replica_hash_load = replica_hash_load.to(device="cpu", dtype=torch.float64) + score = _critical_load(load, best, expert_alignment, node_world_size, replica_hash_load) + for seed in [current, *candidates]: + if seed.shape != current.shape: + raise ValueError("candidate capacity differs from the current placement") + seed = seed.cpu().clone() + validate_physical_placement(seed, load.shape[-1]) + # Preserve node holdings, including when a caller supplies a seed. + for node_start in range(0, current.shape[1], node_world_size): + for layer in range(current.shape[0]): + if set(seed[layer, node_start : node_start + node_world_size].flatten().tolist()) != set( + current[layer, node_start : node_start + node_world_size].flatten().tolist() + ): + raise ValueError("full EPLB candidates must preserve each node's expert set") + for proposal in ( + seed, + _local_search( + load, + seed, + expert_alignment, + node_world_size, + stickiness, + search_steps, + candidate_budget, + replica_hash_load=replica_hash_load, + ), + ): + proposal_score = _critical_load(load, proposal, expert_alignment, node_world_size, replica_hash_load) + improved = proposal_score < score + best[improved] = proposal[improved] + score[improved] = proposal_score[improved] + return best + + +def _local_search(load, placement, alignment, node_world_size, stickiness, steps, budget, replica_hash_load=None): + result = placement.clone() + experts = load.shape[-1] + for layer in range(placement.shape[0]): + source = load[:, layer : layer + 1] + source_hash = replica_hash_load[:, layer : layer + 1] if replica_hash_load is not None else None + margin = stickiness * float(source.sum()) / experts + initial_score = _critical_load(source, placement[layer : layer + 1], alignment, node_world_size, source_hash) + for _ in range(steps): + rows = result[layer] + locations = _expert_locations(rows[None], experts) + contributions = _expert_rank_load_all( + source, locations, source.shape[2], node_world_size, alignment, replica_hash_load=source_hash + )[:, 0] + rank_load = contributions.sum(dim=1) + proposals = _propose_moves(rows, locations[0], contributions, node_world_size, budget) + if not proposals: + break + # Each move changes two distinct experts. Build only their routing + # contributions, never a [candidates, all experts] load tensor. + proposal_tensor = torch.tensor(proposals, dtype=torch.long, device=rows.device) + old_ids = proposal_tensor[:, 2] + new_ids = proposal_tensor[:, 3] + ids = torch.stack((old_ids, new_ids), dim=1) + new_locations = locations[0, ids].clone() + proposal_indexes = torch.arange(len(proposals), device=rows.device) + ranks = proposal_tensor[:, 0] + new_locations[proposal_indexes, 0, ranks] = False + new_locations[proposal_indexes, 1, ranks] = True + has_other_rank = proposal_tensor[:, 4] >= 0 + other_indexes = proposal_indexes[has_other_rank] + other_ranks = proposal_tensor[has_other_rank, 4] + new_locations[other_indexes, 1, other_ranks] = False + new_locations[other_indexes, 0, other_ranks] = True + move_load = source[:, 0, :, ids].permute(0, 2, 1, 3) + move_hash = source_hash[:, 0, :, ids, :].permute(0, 2, 1, 3, 4) if source_hash is not None else None + next_contributions = _expert_rank_load_all( + move_load, new_locations, source.shape[2], node_world_size, alignment, replica_hash_load=move_hash + ).sum(dim=2) + previous = contributions[:, ids].sum(dim=2) + scores = (rank_load[:, None] - previous + next_contributions).max(dim=-1).values.sum(dim=0) + winner = int(scores.argmin()) + if float(scores[winner]) >= float(rank_load.max(dim=-1).values.sum()): + break + rank, slot, old, new, other_rank, other_slot = proposals[winner] + rows[rank, slot] = new + if other_rank >= 0: + rows[other_rank, other_slot] = old + final_score = _critical_load(source, result[layer : layer + 1], alignment, node_world_size, source_hash) + if final_score + margin >= initial_score: + result[layer].copy_(placement[layer]) + return result + + +def _propose_moves(rows, locations, contributions, node_world_size, budget): + """Deterministic interleaving avoids spending the whole budget on one rank.""" + ranks, capacity = rows.shape + rank_order = contributions.sum(dim=(0, 1)).argsort(descending=True, stable=True).tolist() + expert_order = contributions.sum(dim=(0, 2)).argsort(descending=True, stable=True).tolist() + occupancy = locations.tolist() + per_copy_load = contributions.sum(dim=0).tolist() + row_ids = rows.tolist() + per_rank = [] + for rank in rank_order: + node_start = rank // node_world_size * node_world_size + node_end = node_start + node_world_size + # Replacement candidates include cold redundant copies, freeing space + # for hot experts. Last copies on a node may only move through swaps. + + def replacements(rank=rank, node_start=node_start, node_end=node_end): + count = 0 + for new in expert_order: + if occupancy[new][rank]: + continue + if not any(occupancy[new][node_start:node_end]): + continue + for slot, old in enumerate(row_ids[rank]): + if sum(occupancy[old][node_start:node_end]) <= 1: + continue + yield rank, slot, old, new, -1, -1 + count += 1 + if count >= budget: + return + + # Swap hot rows of heavy ranks with cold rows of lighter ranks. In + # fixed-primary mode only the redundant rows are interchangeable. + def swaps(rank=rank, node_start=node_start, node_end=node_end): + count = 0 + hot_slots = sorted(range(capacity), key=lambda s: -per_copy_load[row_ids[rank][s]][rank]) + for other in reversed(rank_order): + if other == rank or other // node_world_size != rank // node_world_size: + continue + cold_slots = sorted(range(capacity), key=lambda s: per_copy_load[row_ids[other][s]][other]) + for slot in hot_slots: + old = row_ids[rank][slot] + if occupancy[old][other]: + continue + for other_slot in cold_slots: + new = row_ids[other][other_slot] + if occupancy[new][rank]: + continue + yield rank, slot, old, new, other, other_slot + count += 1 + if count >= budget: + return + + def interleaved(replacements=replacements, swaps=swaps): + replacement_iterator = replacements() + swap_iterator = swaps() + exhausted_replacements = exhausted_swaps = False + while not (exhausted_replacements and exhausted_swaps): + if not exhausted_replacements: + try: + yield next(replacement_iterator) + except StopIteration: + exhausted_replacements = True + if not exhausted_swaps: + try: + yield next(swap_iterator) + except StopIteration: + exhausted_swaps = True + + per_rank.append(interleaved()) + proposals = [] + for _ in range(budget): + for moves in per_rank: + try: + proposals.append(next(moves)) + except StopIteration: + continue + if len(proposals) == budget: + return proposals + return proposals + + +def plan_eplb_candidate( + expert_load, current, *, node_world_size, expert_alignment, stickiness, replica_hash_load=None, incremental=False +): + """Return a complete physical placement for the manager's CPU planner.""" + if incremental: + # Preserve the existing packing and bound steady-state migrations. + # A full initial pack handles large placement errors; subsequent + # windows need a few hot/cold swaps, not another global reshuffle. + return refine_placement_candidates( + expert_load, + current, + [], + expert_alignment=expert_alignment, + node_world_size=node_world_size, + stickiness=stickiness, + search_steps=4, + candidate_budget=128, + replica_hash_load=replica_hash_load, + ) + return plan_full_experts( + expert_load, + current, + expert_alignment=expert_alignment, + node_world_size=node_world_size, + stickiness=stickiness, + replica_hash_load=replica_hash_load, + ) + + +def plan_full_experts( + expert_load: torch.Tensor, + current: torch.Tensor, + *, + expert_alignment: int | None = None, + node_world_size: int | None = None, + stickiness: float = 0.0, + replica_hash_load: torch.Tensor | None = None, +) -> torch.Tensor: + """Allocate load-driven replica counts, pack within nodes, then refine. + + The capacity and expert set of each node stay fixed. A load-aware pack is + attempted first; a largest-remaining-capacity pack guarantees feasibility + for equal per-rank capacities and replica counts bounded by node size. + """ + validate_physical_placement(current, expert_load.shape[-1]) + node_world_size = _resolve_node_world_size(expert_load, current.shape[1], node_world_size) + load = expert_load.to(device="cpu", dtype=torch.float64) + locations = _expert_locations(current, load.shape[-1]) + # Sum unpadded traffic reaching a node under the current node-local policy. + traffic = _expert_rank_load_all( + load, locations, load.shape[2], node_world_size, None, replica_hash_load=replica_hash_load + ).sum(dim=0) + candidate = current.clone() + for layer in range(current.shape[0]): + for start in range(0, current.shape[1], node_world_size): + end = start + node_world_size + ids = sorted(set(current[layer, start:end].flatten().tolist())) + weights = traffic[layer, :, start:end].sum(dim=-1).tolist() + counts = {expert: 1 for expert in ids} + for _ in range(node_world_size * current.shape[2] - len(ids)): + legal = [expert for expert in ids if counts[expert] < node_world_size] + expert = max(legal, key=lambda e: (weights[e] / counts[e], -e)) + counts[expert] += 1 + packed = _pack_node(counts, weights, node_world_size, current.shape[2], prefer_load=True) + if packed is None: + packed = _pack_node(counts, weights, node_world_size, current.shape[2], prefer_load=False) + assert packed is not None + candidate[layer, start:end] = torch.tensor(packed, dtype=current.dtype) + return refine_placement_candidates( + load, + current, + [candidate], + expert_alignment=expert_alignment, + node_world_size=node_world_size, + stickiness=stickiness, + replica_hash_load=replica_hash_load, + ) + + +def _pack_node(counts, weights, ranks, capacity, *, prefer_load): + rows = [[] for _ in range(ranks)] + loads = [0.0] * ranks + order = sorted(counts, key=lambda e: (-weights[e] / counts[e], -counts[e], e)) + for expert in order: + available = [rank for rank in range(ranks) if len(rows[rank]) < capacity] + if len(available) < counts[expert]: + return None + available.sort(key=lambda r: (loads[r], len(rows[r]), r) if prefer_load else (len(rows[r]), loads[r], r)) + for rank in available[: counts[expert]]: + rows[rank].append(expert) + loads[rank] += weights[expert] / counts[expert] + return rows diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/expert_parallel_state.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/expert_parallel_state.py new file mode 100644 index 0000000000..4f3d5fe12f --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/expert_parallel_state.py @@ -0,0 +1,135 @@ +from dataclasses import dataclass +from typing import Callable, Optional + +import torch + + +@dataclass +class EPLBState: + num_redundant_experts_per_rank: int + logical_to_physical_map: torch.Tensor + logical_replica_count: torch.Tensor + route_counter: torch.Tensor + recording: bool = False + recorded_sample_count: int = 0 + physical_to_logical: Optional[torch.Tensor] = None + placement_generation: int = 0 + migration_poll: Optional[Callable[[], None]] = None + replica_rotation_rank: int = 0 + tp_output_sync_src: Optional[int] = None + + def next_sample_index(self) -> int: + if not self.recording: + return 0 + sample_index = self.recorded_sample_count % self.route_counter.shape[0] + self.recorded_sample_count += 1 + return sample_index + + +@dataclass(frozen=True) +class ExpertParallelState: + num_logical_experts: int + world_size: int + eplb: Optional[EPLBState] = None + + @property + def num_primary_experts_per_rank(self) -> int: + return self.num_logical_experts // self.world_size + + @property + def num_total_physical_experts(self) -> int: + num_redundant_experts_per_rank = 0 if self.eplb is None else self.eplb.num_redundant_experts_per_rank + return self.num_logical_experts + self.world_size * num_redundant_experts_per_rank + + def sync_output(self, output: torch.Tensor, is_prefill: Optional[bool], dist_group=None) -> torch.Tensor: + """Apply the output-consistency contract of this expert-parallel layout.""" + eplb = self.eplb + if is_prefill is True and eplb is not None and eplb.tp_output_sync_src is not None: + from lightllm.distributed import broadcast + + broadcast(output, src=eplb.tp_output_sync_src, group=dist_group) + return output + + +def create_expert_parallel_state( + num_logical_experts: int, + world_size: int, + rank: int, + *, + tp_rank: int, + tp_world_size: int, + is_mtp_draft_model: bool = False, +) -> ExpertParallelState: + from lightllm.common.eplb_utils import EPLB_STEADY_SAMPLE_STEPS, is_eplb_enabled + from lightllm.utils.dist_utils import get_node_world_size + from lightllm.utils.envs_utils import ( + get_env_start_args, + get_prefill_eplb_step_interval, + get_prefill_eplb_steady_sample_steps, + ) + from .eplb_placement import ( + build_initial_redundant_expert_ids, + expand_redundant_placement, + build_logical_to_physical_map, + get_max_reachable_replicas, + ) + + args = get_env_start_args() + eplb = None + if is_eplb_enabled(args): + node_world_size = get_node_world_size() + replica_rotation_rank = rank + tp_output_sync_src = None + if tp_world_size > 1 and not args.enable_tpsp_mix_mode: + # TP ranks hold the same tokens. Identical physical routes keep + # DeepEP's BF16 partial-sum groups consistent without a collective. + leader = rank - tp_rank + if leader // node_world_size == (leader + tp_world_size - 1) // node_world_size: + replica_rotation_rank = leader + else: + # Preserve node-local replica selection for cross-node TP and + # synchronize the complete routed output after MoE execution. + tp_output_sync_src = leader + redundant = args.eplb_num_redundant_experts_per_rank + if is_mtp_draft_model: + from lightllm.distributed import dist_group_manager + + physical_experts = num_logical_experts + world_size * redundant + if getattr(dist_group_manager, "ll_num_experts", physical_experts) != physical_experts: + raise ValueError("EPLB target and draft must share the same DeepEP expert space") + physical_to_logical = expand_redundant_placement( + build_initial_redundant_expert_ids(num_logical_experts, world_size, redundant), num_logical_experts + ) + logical_to_physical, replica_count = build_logical_to_physical_map( + physical_to_logical, + num_logical_experts, + source_rank=rank, + node_world_size=node_world_size, + replica_rotation_rank=replica_rotation_rank, + ) + max_hash_replicas = get_max_reachable_replicas(physical_to_logical, num_logical_experts, node_world_size) + # Cover both initial/steady windows and two microbatch dispatches per step. + # Draft routing stays static and unrecorded, but must use the same + # physical expert count as the target: DeepEP shares its signal buffer. + sample_steps = ( + 1 + if is_mtp_draft_model + else max( + get_prefill_eplb_step_interval(), min(get_prefill_eplb_steady_sample_steps(), EPLB_STEADY_SAMPLE_STEPS) + ) + ) + eplb = EPLBState( + num_redundant_experts_per_rank=redundant, + physical_to_logical=physical_to_logical, + logical_to_physical_map=logical_to_physical.cuda(), + logical_replica_count=replica_count.cuda(), + # Logical routes plus exact hash buckets for every replica count. + route_counter=torch.zeros( + (2 * sample_steps, num_logical_experts, 1 + max_hash_replicas * (max_hash_replicas + 1) // 2), + dtype=torch.int64, + device="cuda", + ), + replica_rotation_rank=replica_rotation_rank, + tp_output_sync_src=tp_output_sync_src, + ) + return ExpertParallelState(num_logical_experts, world_size, eplb) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index fd3de5c9b4..8f33002b9d 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -8,10 +8,13 @@ get_col_slice_mixin, SliceMixinTpl, ) -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import select_fuse_moe_impl +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import create_fuse_moe_impl +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( + create_expert_parallel_state, +) from lightllm.common.basemodel.moe_route_info_manager import get_moe_capture_callback from lightllm.common.quantization.quantize_method import QuantizationMethod -from lightllm.utils.envs_utils import get_redundancy_expert_ids, get_redundancy_expert_num, get_env_start_args +from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_global_world_size, get_global_rank from lightllm.utils.log_utils import init_logger @@ -56,17 +59,14 @@ def __init__( self.n_routed_experts = n_routed_experts self.num_fused_shared_experts = num_fused_shared_experts self._init_config(network_config) - self._init_redundancy_expert_params() - self._init_parallel_params() - self.fuse_moe_impl = select_fuse_moe_impl(self.quant_method, self.enable_ep_moe)( + self._init_expert_parallel_state(network_config) + self._init_weight_partition() + self.fuse_moe_impl = create_fuse_moe_impl( n_routed_experts=self.n_routed_experts, num_fused_shared_experts=self.num_fused_shared_experts, routed_scaling_factor=self.routed_scaling_factor, quant_method=self.quant_method, - redundancy_expert_num=self.redundancy_expert_num, - redundancy_expert_ids_tensor=self.redundancy_expert_ids_tensor, - routed_expert_counter_tensor=self.routed_expert_counter_tensor, - auto_update_redundancy_expert=self.auto_update_redundancy_expert, + expert_parallel_state=self.expert_parallel_state, ) self.lock = threading.Lock() self._create_weight() @@ -80,16 +80,21 @@ def _init_config(self, network_config: Dict[str, Any]): self.routed_scaling_factor = network_config.get("routed_scaling_factor", 1.0) self.scoring_func = network_config.get("scoring_func", "softmax") - def _init_redundancy_expert_params(self): - self.redundancy_expert_num = get_redundancy_expert_num() - self.redundancy_expert_ids = get_redundancy_expert_ids(self.layer_num_) - self.auto_update_redundancy_expert: bool = get_env_start_args().auto_update_redundancy_expert - self.redundancy_expert_ids_tensor = torch.tensor(self.redundancy_expert_ids, dtype=torch.int64, device="cuda") - self.routed_expert_counter_tensor = torch.zeros((self.n_routed_experts,), dtype=torch.int64, device="cuda") - # TODO: find out the reason of failure of deepep when redundancy_expert_num is 1. - assert self.redundancy_expert_num != 1, "redundancy_expert_num can not be 1 for some unknown hang of deepep." + def _init_expert_parallel_state(self, network_config: Dict[str, Any]): + self.expert_parallel_state = ( + create_expert_parallel_state( + self.n_routed_experts, + self.global_world_size, + self.global_rank_, + tp_rank=self.tp_rank_, + tp_world_size=self.tp_world_size_, + is_mtp_draft_model=network_config.get("is_mtp_draft_model", False), + ) + if self.enable_ep_moe + else None + ) - def _init_parallel_params(self): + def _init_weight_partition(self): if self.enable_ep_moe: self.tp_rank_ = 0 self.tp_world_size_ = 1 @@ -103,27 +108,25 @@ def _init_parallel_params(self): self.split_inter_size = self.moe_intermediate_size // self.tp_world_size_ if self.enable_ep_moe: assert self.num_fused_shared_experts == 0, "num_fused_shared_experts must be 0 when enable_ep_moe" - logger.debug( - f"global_rank {self.global_rank_} layerindex {self.layer_num_} " - f"redundancy_expertids: {self.redundancy_expert_ids}" - ) - self.local_n_routed_experts = self.n_routed_experts // self.global_world_size + self.redundancy_expert_num - n_experts_per_rank = self.n_routed_experts // self.global_world_size - start_expert_id = self.global_rank_ * n_experts_per_rank - self.local_expert_ids = ( - list(range(start_expert_id, start_expert_id + n_experts_per_rank)) + self.redundancy_expert_ids - ) - self.expert_idx_to_local_idx = { - expert_idx: expert_idx - start_expert_id for expert_idx in self.local_expert_ids[:n_experts_per_rank] - } - self.redundancy_expert_idx_to_local_idx = { - redundancy_expert_idx: n_experts_per_rank + i - for (i, redundancy_expert_idx) in enumerate(self.redundancy_expert_ids) - } + eplb = self.expert_parallel_state.eplb + num_redundant_experts_per_rank = 0 if eplb is None else eplb.num_redundant_experts_per_rank + num_primary_experts_per_rank = self.expert_parallel_state.num_primary_experts_per_rank + self.local_n_routed_experts = num_primary_experts_per_rank + num_redundant_experts_per_rank + if eplb is None: + start_expert_id = self.global_rank_ * num_primary_experts_per_rank + self.expert_idx_to_local_idx = { + expert_idx: expert_idx - start_expert_id + for expert_idx in range(start_expert_id, start_expert_id + num_primary_experts_per_rank) + } + else: + # Weights, scales and zero points share one absolute local-row map. + self.expert_idx_to_local_idx = { + expert: slot for slot, expert in enumerate(eplb.physical_to_logical[self.global_rank_].tolist()) + } + else: self.local_expert_ids = list(range(self.n_routed_experts + self.num_fused_shared_experts)) self.expert_idx_to_local_idx = {expert_idx: i for (i, expert_idx) in enumerate(self.local_expert_ids)} - self.rexpert_idx_to_local_idx = {} def experts( self, @@ -162,6 +165,7 @@ def experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + dist_group=None if infer_state is None else infer_state.dist_group, ) def low_latency_dispatch( @@ -287,8 +291,6 @@ def load_hf_weights(self, weights): self._load_e_score_correction_bias(weights) self._load_per_expert_scale(weights) self._load_weight(self.expert_idx_to_local_idx, weights) - if self.redundancy_expert_num > 0: - self._load_weight(self.redundancy_expert_idx_to_local_idx, weights) def verify_load(self): weight_load_ok = all(all(_weight_pack.load_ok) for _weight_pack in self.w1_list + self.w2_list + self.w3_list) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py index 67bb90e4ef..c00a35f600 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/__init__.py @@ -2,13 +2,36 @@ from .triton_impl import FuseMoeTriton from .marlin_impl import FuseMoeMarlin from .deepgemm_impl import FuseMoeDeepGEMM +from ..expert_parallel_state import ExpertParallelState -def select_fuse_moe_impl(quant_method: QuantizationMethod, enable_ep_moe: bool): - if enable_ep_moe: - return FuseMoeDeepGEMM +def create_fuse_moe_impl( + *, + n_routed_experts: int, + num_fused_shared_experts: int, + routed_scaling_factor: float, + quant_method: QuantizationMethod, + expert_parallel_state: ExpertParallelState | None = None, +): + if expert_parallel_state is not None: + return FuseMoeDeepGEMM( + n_routed_experts=n_routed_experts, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + quant_method=quant_method, + expert_parallel_state=expert_parallel_state, + ) if quant_method.method_name == "awq_marlin": - return FuseMoeMarlin - else: - return FuseMoeTriton + return FuseMoeMarlin( + n_routed_experts=n_routed_experts, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + quant_method=quant_method, + ) + return FuseMoeTriton( + n_routed_experts=n_routed_experts, + num_fused_shared_experts=num_fused_shared_experts, + routed_scaling_factor=routed_scaling_factor, + quant_method=quant_method, + ) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index e242870d69..817b2954b7 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -1,53 +1,25 @@ import torch -from abc import abstractmethod -from typing import Callable, Optional +from abc import ABC, abstractmethod +from typing import Callable, Optional, Tuple from lightllm.common.quantization.quantize_method import ( WeightPack, QuantizationMethod, ) -from lightllm.utils.dist_utils import ( - get_global_rank, - get_global_world_size, -) -class FuseMoeBaseImpl: +class FuseMoeBaseImpl(ABC): def __init__( self, n_routed_experts: int, num_fused_shared_experts: int, routed_scaling_factor: float, quant_method: QuantizationMethod, - redundancy_expert_num: int, - redundancy_expert_ids_tensor: torch.Tensor, - routed_expert_counter_tensor: torch.Tensor, - auto_update_redundancy_expert: bool, ): self.n_routed_experts = n_routed_experts self.num_fused_shared_experts = num_fused_shared_experts self.routed_scaling_factor = routed_scaling_factor self.quant_method = quant_method - self.global_rank_ = get_global_rank() - self.global_world_size_ = get_global_world_size() - self.ep_n_routed_experts = self.n_routed_experts // self.global_world_size_ - self.total_expert_num_contain_redundancy = ( - self.n_routed_experts + redundancy_expert_num * self.global_world_size_ - ) - - # redundancy expert related - self.redundancy_expert_num = redundancy_expert_num - self.redundancy_expert_ids_tensor = redundancy_expert_ids_tensor - self.routed_expert_counter_tensor = routed_expert_counter_tensor - self.auto_update_redundancy_expert = auto_update_redundancy_expert - - # workspace for kernel optimization - self.workspace = self.create_workspace() - - @abstractmethod - def create_workspace(self): - pass - @abstractmethod def __call__( self, input_tensor: torch.Tensor, @@ -70,5 +42,71 @@ def __call__( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + dist_group=None, + ) -> torch.Tensor: + topk_weights, topk_ids, origin_topk_ids = self._select_experts( + input_tensor=input_tensor, + router_logits=router_logits, + correction_bias=correction_bias, + top_k=top_k, + renormalize=renormalize, + use_grouped_topk=use_grouped_topk, + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func=scoring_func, + per_expert_scale=per_expert_scale, + shared_expert_gate=shared_expert_gate, + is_prefill=is_prefill, + preserve_logical_ids=moe_capture_callback is not None, + ) + if moe_capture_callback is not None: + moe_capture_callback(origin_topk_ids) + return self._fused_experts( + input_tensor=input_tensor, + w13=w13, + w2=w2, + topk_weights=topk_weights, + topk_ids=topk_ids, + router_logits=router_logits, + is_prefill=is_prefill, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + dist_group=dist_group, + ) + + @abstractmethod + def _select_experts( + self, + input_tensor: torch.Tensor, + router_logits: torch.Tensor, + correction_bias: Optional[torch.Tensor], + top_k: int, + renormalize: bool, + use_grouped_topk: bool, + topk_group: int, + num_expert_group: int, + scoring_func: str, + per_expert_scale: Optional[torch.Tensor] = None, + shared_expert_gate: Optional[torch.Tensor] = None, + is_prefill: Optional[bool] = None, + preserve_logical_ids: bool = False, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + pass + + @abstractmethod + def _fused_experts( + self, + input_tensor: torch.Tensor, + w13: WeightPack, + w2: WeightPack, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + router_logits: Optional[torch.Tensor] = None, + is_prefill: Optional[bool] = None, + alpha: Optional[float] = None, + limit: Optional[float] = None, + clamp_up_add_one: bool = True, + dist_group=None, ) -> torch.Tensor: pass diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 09b538a33d..c1e4df2a89 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -1,6 +1,7 @@ import torch from typing import Optional, Tuple, Any -from .triton_impl import FuseMoeTriton +from .base_impl import FuseMoeBaseImpl +from ..expert_parallel_state import ExpertParallelState from lightllm.distributed import dist_group_manager from lightllm.common.quantization.quantize_method import WeightPack from lightllm.utils.envs_utils import ( @@ -16,10 +17,15 @@ ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair -class FuseMoeDeepGEMM(FuseMoeTriton): +class FuseMoeDeepGEMM(FuseMoeBaseImpl): + def __init__(self, *args, expert_parallel_state: ExpertParallelState, **kwargs): + super().__init__(*args, **kwargs) + self.expert_parallel_state = expert_parallel_state + self.eplb = expert_parallel_state.eplb + self.ep_balance_counters = None + def _select_experts( self, input_tensor: torch.Tensor, @@ -33,38 +39,62 @@ def _select_experts( scoring_func: str, per_expert_scale: Optional[torch.Tensor] = None, shared_expert_gate: Optional[torch.Tensor] = None, + is_prefill: Optional[bool] = None, + preserve_logical_ids: bool = False, ): - """Select experts and return topk weights and ids.""" + """EPLB routes both phases through the published physical layout.""" assert shared_expert_gate is None, "fused shared expert as MoE is not supported by DeepGEMM fused MoE" - from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts + eplb = self.eplb + if eplb is not None: + if is_prefill is True and eplb.migration_poll is not None: + eplb.migration_poll() + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import ( + get_group_score_used_topk_num, + triton_grouped_topk_eplb, + ) - topk_weights, topk_ids = select_experts( - hidden_states=input_tensor, - router_logits=router_logits, - correction_bias=correction_bias, - use_grouped_topk=use_grouped_topk, - top_k=top_k, - renormalize=renormalize, - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func=scoring_func, - ) + topk_weights, topk_ids, logical_topk_ids = triton_grouped_topk_eplb( + hidden_states=input_tensor, + gating_output=router_logits, + correction_bias=correction_bias, + topk=top_k, + renormalize=renormalize, + num_expert_group=num_expert_group, + topk_group=topk_group, + scoring_func=scoring_func, + logical_to_physical_map=eplb.logical_to_physical_map, + logical_replica_count=eplb.logical_replica_count, + expert_counter=eplb.route_counter, + sample_index=eplb.next_sample_index() if is_prefill is True else 0, + record_load=is_prefill is True and eplb.recording, + use_grouped_topk=use_grouped_topk, + return_logical_ids=preserve_logical_ids or per_expert_scale is not None, + group_score_used_topk_num=get_group_score_used_topk_num(topk_group, num_expert_group, top_k), + replica_rotation_rank=eplb.replica_rotation_rank, + use_replicas=is_prefill is True, + ) + if per_expert_scale is not None: + topk_weights = topk_weights * per_expert_scale[logical_topk_ids.to(torch.long)].to(topk_weights.dtype) + origin_topk_ids = logical_topk_ids if logical_topk_ids is not None else topk_ids + else: + from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts + + topk_weights, topk_ids = select_experts( + hidden_states=input_tensor, + router_logits=router_logits, + correction_bias=correction_bias, + use_grouped_topk=use_grouped_topk, + top_k=top_k, + renormalize=renormalize, + topk_group=topk_group, + num_expert_group=num_expert_group, + scoring_func=scoring_func, + ) + if per_expert_scale is not None: + topk_weights = topk_weights * per_expert_scale[topk_ids.to(torch.long)].to(topk_weights.dtype) + origin_topk_ids = topk_ids if self.routed_scaling_factor != 1.0: topk_weights.mul_(self.routed_scaling_factor) - if per_expert_scale is not None: - topk_weights = topk_weights * per_expert_scale[topk_ids.to(torch.long)].to(topk_weights.dtype) - origin_topk_ids = topk_ids - if self.redundancy_expert_num > 0: - # 因为 redundancy_topk_ids_repair 会修改 topk_ids，所以需要先复制一份 - origin_topk_ids = topk_ids.clone() - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=self.redundancy_expert_ids_tensor, - ep_expert_num=self.ep_n_routed_experts, - global_rank=self.global_rank_, - expert_counter=self.routed_expert_counter_tensor, - enable_counter=self.auto_update_redundancy_expert, - ) return topk_weights, topk_ids, origin_topk_ids def _fused_experts( @@ -79,6 +109,7 @@ def _fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + dist_group=None, ): output = fused_experts( hidden_states=input_tensor, @@ -86,15 +117,16 @@ def _fused_experts( w2=w2, topk_weights=topk_weights, topk_idx=topk_ids.to(torch.long), - num_experts=self.total_expert_num_contain_redundancy, # number of all experts contain redundancy + num_experts=self.expert_parallel_state.num_total_physical_experts, quant_method=self.quant_method, is_prefill=is_prefill, previous_event=None, # for overlap alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + ep_balance_counters=self.ep_balance_counters, ) - return output + return self.expert_parallel_state.sync_output(output, is_prefill, dist_group) def low_latency_dispatch( self, @@ -118,6 +150,7 @@ def low_latency_dispatch( topk_group=topk_group, num_expert_group=n_group, scoring_func=scoring_func, + is_prefill=False, ) topk_idx = topk_idx.to(torch.long) @@ -127,7 +160,7 @@ def low_latency_dispatch( topk_idx=topk_idx, x=hidden_states, num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, - num_experts=self.total_expert_num_contain_redundancy, + num_experts=self.expert_parallel_state.num_total_physical_experts, use_fp8=use_fp8_w8a8, async_finish=False, return_recv_hook=True, @@ -157,6 +190,7 @@ def select_experts_and_quant_input( topk_group=topk_group, num_expert_group=n_group, scoring_func=scoring_func, + is_prefill=True, ) qinput_tensor = quantize_fused_experts_input(hidden_states, w13, self.quant_method) return topk_weights, topk_idx.to(torch.long), qinput_tensor @@ -174,7 +208,7 @@ def dispatch( qinput_tensor, topk_idx=topk_idx, topk_weights=topk_weights, - num_experts=self.total_expert_num_contain_redundancy, + num_experts=self.expert_parallel_state.num_total_physical_experts, num_max_tokens_per_rank=num_max_tokens_per_rank, expert_alignment=128, num_sms=get_ep_num_sms(), @@ -187,8 +221,23 @@ def dispatch( use_tma_aligned_col_major_sf=True, ) - def hook(): - event.current_stream_wait() + counters = self.ep_balance_counters + if counters is None: + + def hook(): + event.current_stream_wait() + + else: + # Sent routes are globally conserved by all-to-all; recv_x[0] is the 128-aligned expanded compute load. + route_load = topk_idx.numel() + compute_load = recv_x[0].shape[0] + + def hook(): + event.current_stream_wait() + counters.accumulate( + route_load=route_load, + compute_load=compute_load, + ) return recv_x, recv_topk_idx, recv_topk_weights, handle.num_recv_tokens_per_expert_list, handle, hook diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py index 5ccdbb4e9a..811e531d32 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py @@ -11,6 +11,10 @@ class FuseMoeMarlin(FuseMoeTriton): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.workspace = self.create_workspace() + def create_workspace(self): from lightllm.utils.vllm_utils import HAS_VLLM @@ -33,6 +37,7 @@ def _fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + dist_group=None, ): if alpha is not None or limit is not None: raise NotImplementedError("FuseMoeMarlin does not support clamped SwiGLU") diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index f748aea467..4043ee980e 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -1,13 +1,10 @@ import torch -from typing import Callable, Optional +from typing import Optional from lightllm.common.quantization.no_quant import WeightPack from .base_impl import FuseMoeBaseImpl class FuseMoeTriton(FuseMoeBaseImpl): - def create_workspace(self): - return None - def _select_experts( self, input_tensor: torch.Tensor, @@ -21,6 +18,8 @@ def _select_experts( scoring_func: str, per_expert_scale: Optional[torch.Tensor] = None, shared_expert_gate: Optional[torch.Tensor] = None, + is_prefill: Optional[bool] = None, + preserve_logical_ids: bool = False, ): """Select experts and return topk weights and ids.""" from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import select_experts @@ -67,6 +66,7 @@ def _fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + dist_group=None, ): w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale @@ -89,56 +89,3 @@ def _fused_experts( clamp_up_add_one=clamp_up_add_one, ) return input_tensor - - def __call__( - self, - input_tensor: torch.Tensor, - router_logits: torch.Tensor, - w13: WeightPack, - w2: WeightPack, - correction_bias: Optional[torch.Tensor], - scoring_func: str, - top_k: int, - renormalize: bool, - use_grouped_topk: bool, - topk_group: int, - num_expert_group: int, - is_prefill: Optional[bool] = None, - # Callback to capture MoE topk expert ids (routed experts metadata). - moe_capture_callback: Optional[Callable[[torch.Tensor], None]] = None, - per_expert_scale: Optional[torch.Tensor] = None, - shared_expert_gate: Optional[torch.Tensor] = None, - alpha: Optional[float] = None, - limit: Optional[float] = None, - clamp_up_add_one: bool = True, - ): - topk_weights, topk_ids, origin_topk_ids = self._select_experts( - input_tensor=input_tensor, - router_logits=router_logits, - correction_bias=correction_bias, - top_k=top_k, - renormalize=renormalize, - use_grouped_topk=use_grouped_topk, - topk_group=topk_group, - num_expert_group=num_expert_group, - scoring_func=scoring_func, - per_expert_scale=per_expert_scale, - shared_expert_gate=shared_expert_gate, - ) - - if moe_capture_callback is not None: - moe_capture_callback(origin_topk_ids) - - output = self._fused_experts( - input_tensor=input_tensor, - w13=w13, - w2=w2, - topk_weights=topk_weights, - topk_ids=topk_ids, - router_logits=router_logits, - is_prefill=is_prefill, - alpha=alpha, - limit=limit, - clamp_up_add_one=clamp_up_add_one, - ) - return output diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 712860bd8c..bf39b3e47e 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -19,6 +19,7 @@ ep_gather_chunk, ep_zero_padding, ) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_balance import PrefillEPBalanceCounters from lightllm.utils.envs_utils import ( get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, @@ -214,6 +215,7 @@ def fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + ep_balance_counters: Optional[PrefillEPBalanceCounters] = None, ): assert (limit is None and alpha is None) or (limit is not None and alpha is not None) check_ep_expert_dtype(quant_method) @@ -242,6 +244,7 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + ep_balance_counters=ep_balance_counters, ) @@ -264,6 +267,7 @@ def fused_experts_impl( limit: Optional[float] = None, clamp_up_add_one: bool = True, block_size_k: int = 128, + ep_balance_counters: Optional[PrefillEPBalanceCounters] = None, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -309,6 +313,12 @@ def fused_experts_impl( do_expand=True, use_tma_aligned_col_major_sf=True, ) + if ep_balance_counters is not None: + # Sent routes are globally conserved by all-to-all; recv_x[0] is the 128-aligned expanded compute load. + ep_balance_counters.accumulate( + route_load=topk_idx.numel(), + compute_load=recv_x[0].shape[0], + ) # Dispatch is synchronous in this path. Its FP8 source is no longer # needed once the received tensors have been produced. del qinput_tensor, input_scale diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py index 8c882b6c1d..469763e09e 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py @@ -1,92 +1,20 @@ -# adopt from https://github.com/triton-lang/triton/issues/3698#issuecomment-2067681396 import torch import triton import triton.language as tl -from triton.language.standard import _log2, sum, zeros_like -@triton.jit -def _compare_and_swap(x, x_1, ids, flip, i: tl.core.constexpr, n_dims: tl.core.constexpr): - n_outer: tl.core.constexpr = x.numel >> n_dims - shape: tl.core.constexpr = [n_outer * 2 ** i, 2, 2 ** (n_dims - i - 1)] - y = tl.core.reshape(x, shape) - y_1 = tl.core.reshape(x_1, shape) - # slice left/right with 'stride' 2**(n_dims - i - 1) - mask = tl.core.arange(0, 2)[None, :, None] - left = tl.core.broadcast_to(sum(y * (1 - mask), 1)[:, None, :], shape) - right = tl.core.broadcast_to(sum(y * mask, 1)[:, None, :], shape) - left = tl.core.reshape(left, x.shape) - right = tl.core.reshape(right, x.shape) - - left_1 = tl.core.broadcast_to(sum(y_1 * (1 - mask), 1)[:, None, :], shape) - right_1 = tl.core.broadcast_to(sum(y_1 * mask, 1)[:, None, :], shape) - left_1 = tl.core.reshape(left_1, x_1.shape) - right_1 = tl.core.reshape(right_1, x_1.shape) - # idx - y_idx = tl.core.reshape(ids, shape) - left_idx = tl.core.broadcast_to(sum(y_idx * (1 - mask), 1)[:, None, :], shape) - right_idx = tl.core.broadcast_to(sum(y_idx * mask, 1)[:, None, :], shape) - left_idx = tl.core.reshape(left_idx, x.shape) - right_idx = tl.core.reshape(right_idx, x.shape) - - # actual compare-and-swap - idtype = tl.core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) - ileft = left.to(idtype, bitcast=True) - iright = right.to(idtype, bitcast=True) - ix = x.to(idtype, bitcast=True) - - cond = ((left > right) ^ flip).to(tl.int1) - - ret = ix ^ tl.core.where(cond, ileft ^ iright, zeros_like(ix)) - - new_ids = ids ^ tl.core.where(cond, left_idx ^ right_idx, zeros_like(ids)) - - # swap x_1 - idtype_1 = tl.core.get_int_dtype(bitwidth=x_1.dtype.primitive_bitwidth, signed=True) - ileft_1 = left_1.to(idtype_1, bitcast=True) - iright_1 = right_1.to(idtype_1, bitcast=True) - ix_1 = x_1.to(idtype, bitcast=True) - ret_1 = ix_1 ^ tl.core.where(cond, ileft_1 ^ iright_1, zeros_like(ix_1)) - - return ret.to(x.dtype, bitcast=True), ret_1.to(x_1.dtype, bitcast=True), new_ids - - -@triton.jit -def _bitonic_merge(x, x_1, ids, stage: tl.core.constexpr, order: tl.core.constexpr, n_dims: tl.core.constexpr): - """ - order_type 0 == ascending - order_type 1 == descending - order_type 2 == alternating - """ - n_outer: tl.core.constexpr = x.numel >> n_dims - tl.core.static_assert(stage <= n_dims) - # flip denotes whether to re-arrange sub-sequences of elements in ascending or - # descending order. - # if flip = 00000000... then all elements will be re-arranged ascendingly at this stage - # if flip = 00110011... then all the elements will be re-arranged alternatingly (with - # a stride of 2) at this stage - if order == 2: - shape: tl.core.constexpr = [n_outer * 2 ** (n_dims - 1 - stage), 2, 2 ** stage] - flip = tl.core.reshape(tl.core.broadcast_to(tl.core.arange(0, 2)[None, :, None], shape), x.shape) - else: - flip = order - # perform `stage` rounds of `compare-and-swap` - for i in tl.core.static_range(stage): - x, x_1, ids = _compare_and_swap(x, x_1, ids, flip, i + (n_dims - stage), n_dims) - return x, x_1, ids +def get_group_score_used_topk_num(topk_group: int, num_expert_group: int, top_k: int) -> int: + return 2 if topk_group == 4 and num_expert_group == 8 and top_k == 8 else 1 @triton.jit -def argsort(x, x_1, ids, dim: tl.core.constexpr = None, descending: tl.core.constexpr = tl.core.CONSTEXPR_0): - # handle default dimension or check that it is the most minor dim - _dim: tl.core.constexpr = len(x.shape) - 1 if dim is None else dim - tl.core.static_assert(_dim == len(x.shape) - 1, "only minor dimension is currently supported") - # iteratively run bitonic merge-sort steps - n_dims: tl.core.constexpr = _log2(x.shape[_dim]) - - for i in tl.core.static_range(1, n_dims + 1): - x, x_1, ids = _bitonic_merge(x, x_1, ids, i, 2 if i < n_dims else descending, n_dims) - return x, x_1, ids +def _eplb_replica_index(token_index, logical_id, replica_count): + """Mix high and low bits before modulo, including periodic token indices.""" + value = token_index.to(tl.uint32) ^ ((logical_id.to(tl.uint32) + 1) * 0x9E3779B9) + value = (value ^ (value >> 16)) * 0x7FEB352D + value = (value ^ (value >> 15)) * 0x846CA68B + value = value ^ (value >> 16) + return value % replica_count.to(tl.uint32) @triton.jit @@ -95,111 +23,190 @@ def grouped_topk_kernel( gating_output_stride_m, gating_output_stride_n, correction_bias_ptr, - scores_buffer_ptr, # [token_num, total_expert_num] - scores_stride_m, - scores_stride_n, - scores_stride_token_m, - scores_stride_group, - scores_stride_group_v, out_topk_weights, out_topk_weights_stride_m, out_topk_weights_stride_n, out_topk_ids, out_topk_ids_stride_m, out_topk_ids_stride_n, + out_logical_ids, + out_logical_ids_stride_m, + out_logical_ids_stride_n, + logical_to_physical_ptr, + logical_replica_count_ptr, + expert_counter_ptr, + sample_index, group_num, group_expert_num, - total_expert_num, # group_num * group_expert_num == total_expert_num - topk_num, + total_expert_num, group_topk_num, IS_SIGMOID: tl.constexpr, + USE_GROUPED_TOPK: tl.constexpr, HAS_CORRECTION_BIAS: tl.constexpr, - EXPERT_BLOCK_SIZE: tl.constexpr, # tl.next_power_two_of(total_expert_num) - EXPERT_GROUP_NUM: tl.constexpr, # tl.next_power_two_of(group_num) - EXPERT_GROUP_SIZE: tl.constexpr, # tl.next_power_two_of(group_expert_num) + RETURN_LOGICAL_IDS: tl.constexpr, + EXPERT_GROUP_NUM: tl.constexpr, + EXPERT_GROUP_SIZE: tl.constexpr, + TOPK_NUM: tl.constexpr, + TOPK_BLOCK_SIZE: tl.constexpr, RENORMALIZE: tl.constexpr, GROUP_SCORE_USED_TOPK_NUM: tl.constexpr, + COUNTER_NUM_EXPERTS: tl.constexpr, + COUNTER_CHANNELS: tl.constexpr, + HASH_REPLICA_COUNTS: tl.constexpr, + REPLICA_ROTATION_RANK: tl.constexpr, + MAP_SLOTS: tl.constexpr, + APPLY_EPLB: tl.constexpr, + USE_REPLICAS: tl.constexpr, + RECORD_LOAD: tl.constexpr, + SINGLE_TOKEN: tl.constexpr, ): + """Shared logical top-k; ties choose the lowest expert ID, with optional EPLB mapping.""" token_index = tl.program_id(axis=0) - offs_n = tl.arange(0, EXPERT_BLOCK_SIZE) + offs_group = tl.arange(0, EXPERT_GROUP_NUM) + offs_group_v = tl.arange(0, EXPERT_GROUP_SIZE) + logical_ids = offs_group[:, None] * group_expert_num + offs_group_v[None, :] + valid_expert = ( + (offs_group < group_num)[:, None] + & (offs_group_v < group_expert_num)[None, :] + & (logical_ids < total_expert_num) + ) hidden_states = tl.load( - gating_output_ptr + token_index * gating_output_stride_m + offs_n, - mask=offs_n < total_expert_num, - other=-10000000.0, + gating_output_ptr + token_index * gating_output_stride_m + logical_ids * gating_output_stride_n, + mask=valid_expert, + other=-float("inf"), ).to(tl.float32) + if IS_SIGMOID: old_scores = tl.sigmoid(hidden_states) else: - old_scores = tl.softmax(hidden_states) + group_max = tl.max(hidden_states, axis=1) + global_max = tl.max(group_max, axis=0) + numerators = tl.where(valid_expert, tl.exp(hidden_states - global_max), 0.0) + denominator = tl.sum(tl.sum(numerators, axis=1), axis=0) + old_scores = numerators / denominator if HAS_CORRECTION_BIAS: - scores = old_scores + tl.load(correction_bias_ptr + offs_n, mask=offs_n < total_expert_num, other=-10000000.0) + correction_bias = tl.load(correction_bias_ptr + logical_ids, mask=valid_expert, other=0.0) + scores = tl.where(valid_expert, old_scores + correction_bias, -float("inf")) else: - scores = old_scores - - offs_group = tl.arange(0, EXPERT_GROUP_NUM) - offs_group_v = tl.arange(0, EXPERT_GROUP_SIZE) - tl.store(scores_buffer_ptr + scores_stride_m * token_index + offs_n, scores, mask=offs_n < total_expert_num) - tl.debug_barrier() - group_scores = tl.load( - scores_buffer_ptr - + scores_stride_token_m * token_index - + offs_group[:, None] * scores_stride_group - + offs_group_v[None, :] * scores_stride_group_v, - mask=(offs_group < group_num)[:, None] & (offs_group_v < group_expert_num)[None, :], - other=-10000000.0, - ) # [group, group_size] + scores = tl.where(valid_expert, old_scores, -float("inf")) - group_value = tl.sum( - tl.where( - (offs_group < group_num)[:, None] & (offs_group_v < GROUP_SCORE_USED_TOPK_NUM)[None, :], - tl.sort(group_scores, dim=1, descending=True), - 0.0, - ), - axis=1, - ) + if USE_GROUPED_TOPK: + if GROUP_SCORE_USED_TOPK_NUM == 1: + group_value = tl.max(scores, axis=1) + elif GROUP_SCORE_USED_TOPK_NUM == 2: + first_score, first_index = tl.max(scores, axis=1, return_indices=True) + second_score = tl.max( + tl.where(offs_group_v[None, :] == first_index[:, None], -float("inf"), scores), + axis=1, + ) + group_value = first_score + second_score + else: + sorted_group_scores = tl.sort(scores, dim=1, descending=True) + group_value = tl.sum( + tl.where(offs_group_v[None, :] < GROUP_SCORE_USED_TOPK_NUM, sorted_group_scores, 0.0), + axis=1, + ) - if EXPERT_GROUP_NUM > 1: - sorted_group_value = tl.sort(group_value, descending=True) + if EXPERT_GROUP_NUM > 1: + sorted_group_value = tl.sort(group_value, descending=True) + else: + sorted_group_value = group_value + group_topk_value = tl.sum(tl.where(offs_group == group_topk_num - 1, sorted_group_value, 0.0)) + candidate_scores = tl.where( + (group_value >= group_topk_value)[:, None] & valid_expert, + scores, + -float("inf"), + ) else: - sorted_group_value = group_value - group_topk_value = tl.sum(tl.where(offs_group == group_topk_num - 1, sorted_group_value, 0.0)) - mask_group_scores = tl.where( - ((group_value >= group_topk_value)[:, None]) & ((offs_group_v < group_expert_num)[None, :]), - group_scores, - -10000000.0, - ) + candidate_scores = tl.where(valid_expert, old_scores, -float("inf")) - tl.store( - scores_buffer_ptr - + scores_stride_token_m * token_index - + offs_group[:, None] * scores_stride_group - + offs_group_v[None, :] * scores_stride_group_v, - mask_group_scores, - mask=((offs_group < group_num)[:, None]) & ((offs_group_v < group_expert_num)[None, :]), - ) # [group, group_size] - tl.debug_barrier() - mask_scores = tl.load( - scores_buffer_ptr + scores_stride_m * token_index + offs_n, mask=offs_n < total_expert_num, other=-10000000.0 - ) - _, sorted_scores, sorted_indexes = argsort(mask_scores, old_scores, offs_n, descending=True) + sort_block_size: tl.constexpr = EXPERT_GROUP_NUM * EXPERT_GROUP_SIZE + flat_offsets = tl.arange(0, sort_block_size) + candidate_scores = tl.reshape(candidate_scores, (sort_block_size,)) + topk_offsets = tl.arange(0, TOPK_BLOCK_SIZE) + selected_weights = tl.zeros((TOPK_BLOCK_SIZE,), tl.float32) + selected_logical_ids = tl.zeros((TOPK_BLOCK_SIZE,), tl.int32) + sum_scores = 0.0 + for topk_index in range(TOPK_NUM): + selected_offset = tl.argmax(candidate_scores, axis=0) + selected_group = selected_offset // EXPERT_GROUP_SIZE + selected_group_offset = selected_offset % EXPERT_GROUP_SIZE + selected_logical_id = selected_group * group_expert_num + selected_group_offset + selected_hidden_state = tl.load( + gating_output_ptr + token_index * gating_output_stride_m + selected_logical_id * gating_output_stride_n + ).to(tl.float32) + if IS_SIGMOID: + selected_weight = tl.sigmoid(selected_hidden_state) + else: + selected_weight = tl.exp(selected_hidden_state - global_max) / denominator + sum_scores += selected_weight + topk_lane = topk_offsets == topk_index + selected_weights = tl.where(topk_lane, selected_weight, selected_weights) + selected_logical_ids = tl.where(topk_lane, selected_logical_id, selected_logical_ids) + candidate_scores = tl.where(flat_offsets == selected_offset, -float("inf"), candidate_scores) - if RENORMALIZE: - sum_scores = tl.sum(tl.where(offs_n < topk_num, sorted_scores, 0.0)) - renormlize_scores = sorted_scores / sum_scores - - tl.store( - out_topk_weights + token_index * out_topk_weights_stride_m + offs_n, - renormlize_scores, - mask=offs_n < topk_num, + topk_mask = topk_offsets < TOPK_NUM + if RECORD_LOAD: + tl.atomic_add( + expert_counter_ptr + (sample_index * COUNTER_NUM_EXPERTS + selected_logical_ids) * COUNTER_CHANNELS, + 1, + mask=topk_mask, + sem="relaxed", + ) + if COUNTER_CHANNELS > 1: + for copies in tl.static_range(1, HASH_REPLICA_COUNTS + 1): + counts = tl.full((TOPK_BLOCK_SIZE,), copies, tl.int32) + if SINGLE_TOKEN: + bucket = tl.full((TOPK_BLOCK_SIZE,), REPLICA_ROTATION_RANK % copies, tl.int32) + else: + bucket = ( + _eplb_replica_index(token_index, selected_logical_ids, counts) + REPLICA_ROTATION_RANK + ) % copies + channel = 1 + copies * (copies - 1) // 2 + bucket + tl.atomic_add( + expert_counter_ptr + + (sample_index * COUNTER_NUM_EXPERTS + selected_logical_ids) * COUNTER_CHANNELS + + channel, + 1, + mask=topk_mask, + sem="relaxed", + ) + if APPLY_EPLB: + if SINGLE_TOKEN or not USE_REPLICAS: + replica_indices = tl.zeros((TOPK_BLOCK_SIZE,), tl.int32) + else: + replica_counts = tl.load( + logical_replica_count_ptr + selected_logical_ids, + mask=topk_mask, + other=1, + ) + replica_indices = _eplb_replica_index(token_index, selected_logical_ids, replica_counts) + selected_physical_ids = tl.load( + logical_to_physical_ptr + selected_logical_ids * MAP_SLOTS + replica_indices, + mask=topk_mask, + other=-1, ) - tl.store(out_topk_ids + token_index * out_topk_ids_stride_m + offs_n, sorted_indexes, mask=offs_n < topk_num) else: + selected_physical_ids = selected_logical_ids + if RENORMALIZE: + selected_weights /= sum_scores + tl.store( + out_topk_weights + token_index * out_topk_weights_stride_m + topk_offsets * out_topk_weights_stride_n, + selected_weights, + mask=topk_mask, + ) + tl.store( + out_topk_ids + token_index * out_topk_ids_stride_m + topk_offsets * out_topk_ids_stride_n, + selected_physical_ids, + mask=topk_mask, + ) + if RETURN_LOGICAL_IDS: tl.store( - out_topk_weights + token_index * out_topk_weights_stride_m + offs_n, sorted_scores, mask=offs_n < topk_num + out_logical_ids + token_index * out_logical_ids_stride_m + topk_offsets * out_logical_ids_stride_n, + selected_logical_ids, + mask=topk_mask, ) - tl.store(out_topk_ids + token_index * out_topk_ids_stride_m + offs_n, sorted_indexes, mask=offs_n < topk_num) - return def triton_grouped_topk( @@ -213,53 +220,115 @@ def triton_grouped_topk( scoring_func: str = "softmax", group_score_used_topk_num=2, ): + weights, ids, _ = triton_grouped_topk_eplb( + hidden_states=hidden_states, + gating_output=gating_output, + correction_bias=correction_bias, + topk=topk, + renormalize=renormalize, + num_expert_group=num_expert_group, + topk_group=topk_group, + scoring_func=scoring_func, + logical_to_physical_map=None, + logical_replica_count=None, + expert_counter=None, + sample_index=0, + record_load=False, + use_grouped_topk=True, + group_score_used_topk_num=group_score_used_topk_num, + ) + return weights, ids - if correction_bias is not None: - has_correction_bias = True - else: - has_correction_bias = False - - token_num, total_expert_num = gating_output.shape - if gating_output.dtype == torch.float64: - dtype = torch.float64 - else: - dtype = torch.float32 - - scores_buffer = torch.empty((token_num, total_expert_num), dtype=dtype, device="cuda") - out_topk_weights = torch.empty((token_num, topk), dtype=torch.float32, device="cuda") - out_topk_ids = torch.empty((token_num, topk), dtype=torch.long, device="cuda") - assert total_expert_num % num_expert_group == 0 +def triton_grouped_topk_eplb( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + correction_bias: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: int, + topk_group: int, + scoring_func: str, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + expert_counter: torch.Tensor, + sample_index: int, + record_load: bool, + use_grouped_topk: bool, + return_logical_ids: bool = False, + group_score_used_topk_num: int = 2, + replica_rotation_rank: int = 0, + use_replicas: bool = True, +): + """One logical selector; absent EPLB maps compile out accounting and replica selection.""" + from math import isqrt - if token_num <= 256: - num_warps = 4 + channels = expert_counter.shape[2] if expert_counter is not None and expert_counter.ndim == 3 else 1 + max_hash_replicas = (isqrt(1 + 8 * (channels - 1)) - 1) // 2 + if expert_counter is not None and ( + expert_counter.ndim not in (2, 3) + or (channels > 1 and max_hash_replicas * (max_hash_replicas + 1) // 2 != channels - 1) + ): + raise ValueError("EPLB counters must contain logical counts and complete replica hash buckets") + token_num, total_expert_num = gating_output.shape + out_topk_weights = torch.empty((token_num, topk), dtype=torch.float32, device=gating_output.device) + out_topk_ids = torch.empty((token_num, topk), dtype=torch.long, device=gating_output.device) + out_logical_ids = ( + torch.empty((token_num, topk), dtype=torch.long, device=gating_output.device) if return_logical_ids else None + ) + if token_num == 0: + return out_topk_weights, out_topk_ids, out_logical_ids + if use_grouped_topk: + assert total_expert_num % num_expert_group == 0 + group_num = num_expert_group + group_expert_num = total_expert_num // num_expert_group + group_topk_num = topk_group else: - num_warps = 1 - + group_num = 1 + group_expert_num = total_expert_num + group_topk_num = 1 + expert_group_num = triton.next_power_of_2(group_num) + expert_group_size = triton.next_power_of_2(group_expert_num) + sort_block_size = expert_group_num * expert_group_size + num_warps = min(max(1, sort_block_size // 256), 8) grouped_topk_kernel[(token_num,)]( gating_output, *gating_output.stride(), correction_bias, - scores_buffer, - *scores_buffer.stride(), - *scores_buffer.view(token_num, num_expert_group, -1).stride(), out_topk_weights, *out_topk_weights.stride(), out_topk_ids, *out_topk_ids.stride(), - group_num=num_expert_group, - group_expert_num=total_expert_num // num_expert_group, + out_logical_ids if out_logical_ids is not None else out_topk_ids, + *(out_logical_ids.stride() if out_logical_ids is not None else out_topk_ids.stride()), + logical_to_physical_map, + logical_replica_count, + expert_counter, + sample_index, + group_num=group_num, + group_expert_num=group_expert_num, total_expert_num=total_expert_num, - topk_num=topk, - group_topk_num=topk_group, - IS_SIGMOID=scoring_func == "sigmoid", - HAS_CORRECTION_BIAS=has_correction_bias, - EXPERT_BLOCK_SIZE=triton.next_power_of_2(total_expert_num), - EXPERT_GROUP_NUM=triton.next_power_of_2(num_expert_group), - EXPERT_GROUP_SIZE=triton.next_power_of_2(total_expert_num // num_expert_group), + group_topk_num=group_topk_num, + IS_SIGMOID=use_grouped_topk and scoring_func == "sigmoid", + USE_GROUPED_TOPK=use_grouped_topk, + HAS_CORRECTION_BIAS=use_grouped_topk and correction_bias is not None, + RETURN_LOGICAL_IDS=return_logical_ids, + EXPERT_GROUP_NUM=expert_group_num, + EXPERT_GROUP_SIZE=expert_group_size, + TOPK_NUM=topk, + TOPK_BLOCK_SIZE=triton.next_power_of_2(topk), RENORMALIZE=renormalize, GROUP_SCORE_USED_TOPK_NUM=group_score_used_topk_num, + COUNTER_NUM_EXPERTS=expert_counter.shape[1] if expert_counter is not None else total_expert_num, + COUNTER_CHANNELS=channels, + HASH_REPLICA_COUNTS=max_hash_replicas, + REPLICA_ROTATION_RANK=replica_rotation_rank, + MAP_SLOTS=logical_to_physical_map.shape[1] if logical_to_physical_map is not None else 1, + APPLY_EPLB=logical_to_physical_map is not None, + USE_REPLICAS=use_replicas, + RECORD_LOAD=record_load, + SINGLE_TOKEN=token_num == 1, num_warps=num_warps, num_stages=1, ) - return out_topk_weights, out_topk_ids + return out_topk_weights, out_topk_ids, out_logical_ids diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py index 87cda6ce16..ad2aff269b 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py @@ -136,17 +136,15 @@ def select_experts( custom_routing_function: Optional[Callable] = None, ): from lightllm.common.basemodel.triton_kernel.fused_moe.topk_select import fused_topk - from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import ( + get_group_score_used_topk_num, + triton_grouped_topk, + ) # DeekSeekv2 uses grouped_top_k if use_grouped_topk: assert topk_group is not None assert num_expert_group is not None - group_score_topk_num = 1 - # for deepseek v3 - if topk_group == 4 and num_expert_group == 8 and top_k == 8: - group_score_topk_num = 2 - topk_weights, topk_ids = triton_grouped_topk( hidden_states=hidden_states, gating_output=router_logits, @@ -156,7 +154,7 @@ def select_experts( num_expert_group=num_expert_group, topk_group=topk_group, scoring_func=scoring_func, - group_score_used_topk_num=group_score_topk_num, + group_score_used_topk_num=get_group_score_used_topk_num(topk_group, num_expert_group, top_k), ) elif custom_routing_function is None: @@ -168,8 +166,7 @@ def select_experts( hidden_states=hidden_states, gating_output=router_logits, topk=top_k, renormalize=renormalize ) - ######################################## warning ################################################## - # here is used to match autotune feature, make topk_ids more random + # Preserve the ordinary MoE autotune workload; EPLB uses its own selector. if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): rand_gen = torch.Generator(device="cuda") rand_gen.manual_seed(router_logits.shape[0]) diff --git a/lightllm/common/basemodel/triton_kernel/redundancy_topk_ids_repair.py b/lightllm/common/basemodel/triton_kernel/redundancy_topk_ids_repair.py deleted file mode 100644 index ba48f414db..0000000000 --- a/lightllm/common/basemodel/triton_kernel/redundancy_topk_ids_repair.py +++ /dev/null @@ -1,111 +0,0 @@ -import torch -import triton -import triton.language as tl - - -@triton.jit -def _redundancy_topk_ids_repair_kernel( - topk_ids_ptr, - topk_total_num, - ep_expert_num, - redundancy_expert_num, - global_rank, - redundancy_expert_ids_ptr, - expert_counter_ptr, - BLOCK_SIZE: tl.constexpr, - ENABLE_COUNTER: tl.constexpr, -): - block_index = tl.program_id(0) - offs_d = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offs_d < topk_total_num - current_topk_ids = tl.load(topk_ids_ptr + offs_d, mask=mask, other=0) - - if ENABLE_COUNTER: - tl.atomic_add(expert_counter_ptr + current_topk_ids, 1, mask=mask) - - # Remap original expert IDs to a new space that accounts for redundant expert slots. - new_current_topk_ids = (current_topk_ids // ep_expert_num) * redundancy_expert_num + current_topk_ids - - for i in tl.range(0, redundancy_expert_num, step=1, num_stages=3): - cur_redundancy_expert_id = tl.load(redundancy_expert_ids_ptr + i) - cur_redundancy_expert_id = ( - cur_redundancy_expert_id // ep_expert_num - ) * redundancy_expert_num + cur_redundancy_expert_id - new_current_topk_ids = tl.where( - new_current_topk_ids == cur_redundancy_expert_id, - (ep_expert_num + redundancy_expert_num) * (global_rank) + ep_expert_num + i, - new_current_topk_ids, - ) - - tl.store(topk_ids_ptr + offs_d, new_current_topk_ids, mask=mask) - return - - -@torch.no_grad() -def redundancy_topk_ids_repair( - topk_ids: torch.Tensor, - redundancy_expert_ids: torch.Tensor, - ep_expert_num: int, - global_rank: int, - expert_counter: torch.Tensor = None, - enable_counter: bool = False, -): - assert topk_ids.is_contiguous() - assert len(topk_ids.shape) == 2 - assert redundancy_expert_ids is not None - redundancy_expert_num = redundancy_expert_ids.shape[0] - BLOCK_SIZE = 512 - grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) - num_warps = 4 - - _redundancy_topk_ids_repair_kernel[grid]( - topk_ids_ptr=topk_ids, - topk_total_num=topk_ids.numel(), - ep_expert_num=ep_expert_num, - redundancy_expert_num=redundancy_expert_num, - global_rank=global_rank, - redundancy_expert_ids_ptr=redundancy_expert_ids, - expert_counter_ptr=expert_counter, - BLOCK_SIZE=BLOCK_SIZE, - ENABLE_COUNTER=enable_counter, - num_warps=num_warps, - num_stages=3, - ) - return - - -@triton.jit -def _expert_id_counter_kernel( - topk_ids_ptr, - topk_total_num, - expert_counter_ptr, - BLOCK_SIZE: tl.constexpr, -): - block_index = tl.program_id(0) - offs_d = block_index * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offs_d < topk_total_num - current_topk_ids = tl.load(topk_ids_ptr + offs_d, mask=mask, other=0) - tl.atomic_add(expert_counter_ptr + current_topk_ids, 1, mask=mask) - return - - -@torch.no_grad() -def expert_id_counter( - topk_ids: torch.Tensor, - expert_counter: torch.Tensor, -): - assert topk_ids.is_contiguous() - assert len(topk_ids.shape) == 2 - BLOCK_SIZE = 512 - grid = (triton.cdiv(topk_ids.numel(), BLOCK_SIZE),) - num_warps = 4 - - _expert_id_counter_kernel[grid]( - topk_ids_ptr=topk_ids, - topk_total_num=topk_ids.numel(), - expert_counter_ptr=expert_counter, - BLOCK_SIZE=BLOCK_SIZE, - num_warps=num_warps, - num_stages=1, - ) - return diff --git a/lightllm/common/eplb_planner_worker.py b/lightllm/common/eplb_planner_worker.py new file mode 100644 index 0000000000..41de36b04c --- /dev/null +++ b/lightllm/common/eplb_planner_worker.py @@ -0,0 +1,249 @@ +"""CPU-only persistent EPLB candidate planner actor.""" +import atexit +import multiprocessing as mp +import os +import subprocess +import sys +import time +import threading + + +class _PlannerProcess: + """Small Process-compatible owner for the independently executed actor.""" + + def __init__(self, child, parent_pid): + self._popen = subprocess.Popen( + [sys.executable, os.path.abspath(__file__), str(child.fileno()), str(parent_pid)], + pass_fds=(child.fileno(),), + env={**os.environ, "CUDA_VISIBLE_DEVICES": ""}, + ) + + def is_alive(self): + return self._popen.poll() is None + + def join(self, timeout): + try: + self._popen.wait(timeout=timeout) + except subprocess.TimeoutExpired: + pass + + def terminate(self): + self._popen.terminate() + + def kill(self): + self._popen.kill() + + def close(self): + self._popen.wait() + + +class EPLBPlannerWorker: + def __init__(self, timeout_s=120.0): + self._timeout_s = timeout_s + self._job_id = 0 + self._rpc_lock = threading.Lock() + self._parent, child = mp.get_context("spawn").Pipe(duplex=True) + self._process = None + try: + # spawn reimports the serving main before its target can hide CUDA. + # exec this file directly with a CPU-only environment instead. + self._process = _PlannerProcess(child, os.getpid()) + child.close() + # Torch's CPU tensor reducers share descriptors through resource_sharer. + # Synchronize its authentication key privately, before importing Torch. + self._parent.send_bytes(bytes(mp.current_process().authkey)) + hello = self._receive(timeout_s) + if hello.get("kind") != "ready" or hello.get("cuda_initialized"): + raise RuntimeError(f"EPLB planner worker startup failed: {hello}") + self.info = hello + except BaseException: + try: + child.close() + except OSError: + pass + self.close() + self._parent.close() + raise + atexit.register(self.close) + + def _receive(self, timeout_s): + if not self._parent.poll(timeout_s): + raise TimeoutError("EPLB planner worker RPC timed out") + try: + reply = self._parent.recv() + except EOFError as exc: + raise RuntimeError("EPLB planner worker exited unexpectedly") from exc + if not isinstance(reply, dict): + raise RuntimeError("EPLB planner worker returned non-dict response") + if reply.get("kind") == "error": + raise RuntimeError("EPLB planner worker failed: " + reply.get("message", "unknown")) + return reply + + def plan(self, global_load, current_placement, **settings): + if not self._rpc_lock.acquire(blocking=False): + raise RuntimeError("EPLB planner worker permits only one in-flight RPC") + try: + if self._process is None or not self._process.is_alive(): + raise RuntimeError("EPLB planner worker is not alive") + if ( + not hasattr(global_load, "device") + or not hasattr(current_placement, "device") + or global_load.device.type != "cpu" + or current_placement.device.type != "cpu" + or global_load.ndim != 4 + or current_placement.ndim != 3 + ): + raise ValueError( + "EPLB planner worker accepts [samples,layers,nodes,experts] and " + "[layers,ranks,slots] CPU tensors only" + ) + self._job_id += 1 + started = time.perf_counter() + self._parent.send( + { + "kind": "plan", + "job_id": self._job_id, + "global_load": global_load, + "current_placement": current_placement, + "settings": settings, + } + ) + reply = self._receive(self._timeout_s) + if reply.get("kind") != "result" or reply.get("job_id") != self._job_id: + raise RuntimeError("EPLB planner worker returned mismatched response/job id") + candidate = reply.get("candidate") + if ( + not hasattr(candidate, "device") + or candidate.device.type != "cpu" + or candidate.dtype != current_placement.dtype + or candidate.shape != current_placement.shape + ): + raise RuntimeError("EPLB planner worker returned invalid candidate") + return candidate, { + "compute_ms": reply["compute_ms"], + "roundtrip_ms": (time.perf_counter() - started) * 1000.0, + } + except BaseException: + self.close() + raise + finally: + self._rpc_lock.release() + + def close(self): + process = getattr(self, "_process", None) + if process is None: + return + try: + if process.is_alive(): + try: + self._parent.send({"kind": "stop"}) + except (BrokenPipeError, EOFError, OSError): + pass + process.join(5) + if process.is_alive(): + process.terminate() + process.join(5) + if process.is_alive(): + process.kill() + process.join(5) + if process.is_alive(): + raise RuntimeError("EPLB planner worker did not exit during close") + finally: + try: + self._parent.close() + except OSError: + pass + if not process.is_alive(): + process.close() + self._process = None + + +def _watch_parent(expected_parent_pid): + while True: + if os.getppid() != expected_parent_pid: + os._exit(0) + time.sleep(0.25) + + +def _load_cpu_planner(): + """Load the CPU modules without executing the GPU model package initializers.""" + import importlib.util + import types + from pathlib import Path + + name = "_lightllm_eplb_cpu" + directory = Path(__file__).resolve().parent / "basemodel/layer_weights/meta_weights/fused_moe" + package = types.ModuleType(name) + package.__path__ = [str(directory)] + sys.modules[name] = package + for module_name in ("eplb_placement", "eplb_planner"): + qualified = f"{name}.{module_name}" + spec = importlib.util.spec_from_file_location(qualified, directory / f"{module_name}.py") + module = importlib.util.module_from_spec(spec) + sys.modules[qualified] = module + spec.loader.exec_module(module) + return module + + +def _serve(conn, expected_parent_pid): + os.environ["CUDA_VISIBLE_DEVICES"] = "" + threading.Thread(target=_watch_parent, args=(expected_parent_pid,), daemon=True).start() + try: + import torch + + torch.set_num_threads(1) + torch.set_num_interop_threads(1) + eplb_planner = _load_cpu_planner() + + # Startup prewarm is intentionally small and CPU-only. + eplb_planner.plan_eplb_candidate( + torch.ones((1, 1, 1, 8), dtype=torch.int64), + torch.arange(8).view(1, 8, 1), + node_world_size=8, + expert_alignment=128, + stickiness=0.1, + ) + conn.send( + { + "kind": "ready", + "pid": os.getpid(), + "module": eplb_planner.__file__, + "torch_threads": torch.get_num_threads(), + "cuda_initialized": torch.cuda.is_initialized(), + } + ) + while True: + if os.getppid() != expected_parent_pid: + os._exit(0) + message = conn.recv() + if message.get("kind") == "stop": + return + if message.get("kind") != "plan": + raise RuntimeError("unknown worker request") + started = time.perf_counter() + candidate = eplb_planner.plan_eplb_candidate( + message["global_load"], message["current_placement"], **message["settings"] + ) + conn.send( + { + "kind": "result", + "job_id": message["job_id"], + "candidate": candidate, + "compute_ms": (time.perf_counter() - started) * 1000.0, + } + ) + except BaseException as exc: + try: + conn.send({"kind": "error", "message": f"{type(exc).__name__}: {exc}"}) + except BaseException: + pass + finally: + conn.close() + + +if __name__ == "__main__": + from multiprocessing.connection import Connection + + connection = Connection(int(sys.argv[1])) + mp.current_process().authkey = connection.recv_bytes() + _serve(connection, int(sys.argv[2])) diff --git a/lightllm/common/eplb_utils.py b/lightllm/common/eplb_utils.py new file mode 100644 index 0000000000..b3fd0b2095 --- /dev/null +++ b/lightllm/common/eplb_utils.py @@ -0,0 +1,99 @@ +"""EPLB initialization, transfer and model profiling helpers.""" + +import math + + +EPLB_FULL_STAGING_PARTS = 1 +EPLB_STEADY_SAMPLE_STEPS = 4 + + +def is_eplb_enabled(args) -> bool: + """Resolve and validate EPLB at initialization; PD decode keeps its own layout.""" + if not args.enable_prefill_eplb or args.run_mode == "decode": + return False + if args.run_mode not in ("prefill", "normal"): + raise ValueError("EPLB requires --run_mode prefill or single-node normal") + if args.run_mode == "normal" and getattr(args, "nnodes", 1) != 1: + raise ValueError("EPLB normal mode requires a single node") + assert args.enable_ep_moe, "--enable_prefill_eplb requires --enable_ep_moe" + assert not args.enable_prefill_cudagraph, "--enable_prefill_eplb does not support --enable_prefill_cudagraph" + # Online updates use the initial expert-row map and do not invalidate staged migrations. + assert not args.enable_rl, "--enable_prefill_eplb does not support --enable_rl" + + # Keep disabled/PD decode resolution and planner imports independent of CUDA. + from lightllm.utils.device_utils import is_sm100_gpu + + # SM100 Mega-MoE caches transformed weights by data_ptr; EPLB updates them in place. + assert not is_sm100_gpu(), "--enable_prefill_eplb does not support SM100" + assert args.eplb_num_redundant_experts_per_rank >= 0, "redundant expert count must be nonnegative" + return True + + +def find_ep_moe_weights(model): + # Deferred import keeps planner/transfer helpers independent of model imports. + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight + + weights = {} + for layer in model.trans_layers_weight: + for weight in getattr(layer, "__dict__", {}).values(): + if isinstance(weight, FusedMoeWeight) and weight.enable_ep_moe: + weights[id(weight)] = weight + return sorted(weights.values(), key=lambda weight: weight.layer_num_) + + +def get_eplb_staging_shape(weights): + """Stage one complete layer before modifying any live expert rows.""" + if not weights: + return 0, 0 + parallel = weights[0].expert_parallel_state + return 1, parallel.num_primary_experts_per_rank + parallel.eplb.num_redundant_experts_per_rank + + +def get_eplb_staging_tensor_shape(tensor, staged_rows): + """Allocation shape shared by EPLB transfer and memory profiling.""" + row_numel = tensor[0].numel() + return (staged_rows, (row_numel + EPLB_FULL_STAGING_PARTS - 1) // EPLB_FULL_STAGING_PARTS) + + +def get_eplb_staging_nbytes(weights) -> int: + """Owned staging bytes, using the transfer's allocation shapes including scales.""" + if not weights: + return 0 + depth, staged_rows = get_eplb_staging_shape(weights) + return depth * sum( + math.prod(get_eplb_staging_tensor_shape(tensor, staged_rows)) * tensor.element_size() + for _, tensor in extract_eplb_expert_tensors(weights[0]) + ) + + +def get_eplb_sampling_peak_nbytes(weights) -> int: + """Peak stack/index-select temporaries; persistent route counters are already allocated.""" + counters = [] + seen = set() + for weight in weights: + state = getattr(weight, "expert_parallel_state", None) + counter = getattr(getattr(state, "eplb", None), "route_counter", None) + if counter is not None and id(counter) not in seen: + seen.add(id(counter)) + counters.append(counter) + if not counters: + return 0 + first = counters[0] + if any(tuple(counter.shape) != tuple(first.shape) or counter.dtype != first.dtype for counter in counters): + raise ValueError("EPLB route counters must have identical shape and dtype") + align = lambda value: (value + 511) // 512 * 512 + stack_bytes = align(len(counters) * first.numel() * first.element_size()) + index_temp_bytes = align(first.shape[0] * 8) + return 2 * stack_bytes + index_temp_bytes + + +def extract_eplb_expert_tensors(weight): + result = [] + for pack_name in ("w13", "w2"): + pack = getattr(weight, pack_name) + for value_name in ("weight", "weight_scale", "weight_zero_point"): + tensor = getattr(pack, value_name, None) + if tensor is not None: + assert tensor.ndim >= 1 and tensor.is_contiguous(), f"{pack_name}.{value_name} must be contiguous" + result.append((f"{pack_name}.{value_name}", tensor)) + return result diff --git a/lightllm/common/quantization/awq.py b/lightllm/common/quantization/awq.py index f3c7623975..a2949ccad8 100644 --- a/lightllm/common/quantization/awq.py +++ b/lightllm/common/quantization/awq.py @@ -109,11 +109,11 @@ def apply( return out def _create_weight( - self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: Optional[int] = None ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) group_size = self.hf_quantization_config["group_size"] - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (in_dim, out_dim // self.pack_factor), dtype=torch.int32).cuda(device_id) weight_scale = torch.empty(expert_prefix + (in_dim // group_size, out_dim), dtype=dtype).cuda(device_id) weight_zero_point = torch.empty( @@ -206,13 +206,13 @@ def apply( return out def _create_weight( - self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: Optional[int] = None ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) self.n = out_dim self.k = in_dim group_size = self.hf_quantization_config["group_size"] - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty( expert_prefix + (in_dim // self.tile_size, out_dim * self.tile_size // self.pack_factor), dtype=torch.int32 ).cuda(device_id) diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index c953afd27a..3839fdef5f 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -119,7 +119,12 @@ def apply( return out def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: self.use_ue8m0_scales = ( self.hf_quantization_config is not None and self.hf_quantization_config.get("scale_fmt") == "ue8m0" @@ -135,7 +140,7 @@ def _create_weight( ) weight_scale_out_dim = sum(weight_scale_out_dims) weight_scale_in_dim = (in_dim + self.block_size - 1) // self.block_size - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.float8_e4m3fn).cuda(device_id) if self.use_packed_ue8m0: aligned_out_dim = (out_dim + 3) // 4 * 4 @@ -211,12 +216,17 @@ def apply( raise NotImplementedError("fp4fp8-b32-deepgemm is only implemented for fused MoE expert weights") def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims assert in_dim % 2 == 0, "FP4 packed weight requires even input dimension" assert in_dim % self.block_size == 0, "FP4 scale dimension must be divisible by block_size" - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (out_dim, in_dim // 2), dtype=torch.int8).cuda(device_id) weight_scale = torch.empty(expert_prefix + (out_dim, in_dim // self.block_size), dtype=torch.int32).cuda( device_id diff --git a/lightllm/common/quantization/no_quant.py b/lightllm/common/quantization/no_quant.py index fa926ad6f0..9fc782f83a 100644 --- a/lightllm/common/quantization/no_quant.py +++ b/lightllm/common/quantization/no_quant.py @@ -35,10 +35,15 @@ def apply( return torch.addmm(bias, input_tensor, weight, out=out) def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=dtype).cuda(device_id) mm_param = WeightPack(weight=weight, weight_scale=None, weight_zero_point=None) # weight layout is (out_dim, in_dim), so the split dimension is -2. diff --git a/lightllm/common/quantization/quantize_method.py b/lightllm/common/quantization/quantize_method.py index 95d8d806f9..641c72d67e 100644 --- a/lightllm/common/quantization/quantize_method.py +++ b/lightllm/common/quantization/quantize_method.py @@ -112,8 +112,9 @@ def _check_weight_need_quanted(self, weight: torch.Tensor) -> bool: return weight.dtype in [torch.bfloat16, torch.float16, torch.float32, torch.float64] def _create_weight( - self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: Optional[int] = None ) -> Tuple[WeightPack, List[WeightPack]]: + """Create dense weights for None; an expert count always retains the expert axis.""" pass def _split_weight_pack( diff --git a/lightllm/common/quantization/w8a8.py b/lightllm/common/quantization/w8a8.py index 472db0d543..fbc34b9598 100644 --- a/lightllm/common/quantization/w8a8.py +++ b/lightllm/common/quantization/w8a8.py @@ -53,7 +53,12 @@ def method_name(self): return "w8a8-base" def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: raise NotImplementedError("Not implemented") @@ -103,10 +108,15 @@ def method_name(self): return "w8a8-vllm" def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.int8).cuda(device_id) weight_scale = torch.empty(expert_prefix + (out_dim,), dtype=torch.float32).cuda(device_id) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) @@ -166,10 +176,15 @@ def method_name(self): return "fp8w8a8-vllm" def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.float8_e4m3fn).cuda(device_id) weight_scale = torch.empty(expert_prefix + (out_dim,), dtype=torch.float32).cuda(device_id) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) @@ -285,9 +300,14 @@ def _dynamic_quant_input( return qweight, weight_scale, x_q, x_scale, m, n def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: - if num_experts != 1: + if num_experts is not None: raise NotImplementedError("FP8 per-tensor quantization currently supports dense weights only") if isinstance(out_dims, int): out_dims = [out_dims] @@ -454,10 +474,15 @@ def _dynamic_quant_input( return qinput_tensor, qweight, input_scale, weight_scale, out def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.float8_e4m3fn).cuda(device_id) weight_scale = torch.empty( expert_prefix + (out_dim // self.block_size, in_dim // self.block_size), dtype=torch.float32 diff --git a/lightllm/common/quantization/w8a8gx.py b/lightllm/common/quantization/w8a8gx.py index 3ffb5afc13..9eb2619fe1 100644 --- a/lightllm/common/quantization/w8a8gx.py +++ b/lightllm/common/quantization/w8a8gx.py @@ -32,7 +32,12 @@ def method_name(self): return "w8a8gx-base" def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: raise NotImplementedError("Not implemented") @@ -119,10 +124,15 @@ def method_name(self): return f"fp8w8a8g{self.act_quant_group_size}-triton" def _create_weight( - self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1 + self, + out_dims: Union[int, List[int]], + in_dim: int, + dtype: torch.dtype, + device_id: int, + num_experts: Optional[int] = None, ) -> Tuple[WeightPack, List[WeightPack]]: out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims - expert_prefix = (num_experts,) if num_experts > 1 else () + expert_prefix = () if num_experts is None else (num_experts,) weight = torch.empty(expert_prefix + (out_dim, in_dim), dtype=torch.float8_e4m3fn).cuda(device_id) weight_scale = torch.empty(expert_prefix + (out_dim,), dtype=torch.float32).cuda(device_id) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 93c603212d..c8ee8eb5eb 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -24,17 +24,18 @@ import torch.distributed as dist from torch.distributed import ReduceOp, ProcessGroup from typing import List, Dict, Optional, Set, Union +from lightllm.common.eplb_utils import is_eplb_enabled from lightllm.utils.log_utils import init_logger from lightllm.utils.device_utils import has_nvlink from lightllm.utils.envs_utils import ( get_env_start_args, get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, - get_redundancy_expert_num, ) from lightllm.utils.dist_utils import ( get_global_world_size, get_dp_world_size, + get_node_world_size, create_new_group_for_current_dp, create_dp_special_inter_group, ) @@ -67,7 +68,14 @@ def __init__(self): self.autotune_group = dist.new_group([i for i in range(get_global_world_size())], backend="gloo") def _support_custom_allreduce(self) -> bool: - return has_nvlink() and self.dp_world_size in [2, 4, 6, 8] + # CUDA IPC and symmetric-memory handles are local to one host. + node_world_size = get_node_world_size() + return ( + self.dp_world_size <= node_world_size + and node_world_size % self.dp_world_size == 0 + and self.dp_world_size in [2, 4, 6, 8] + and has_nvlink() + ) def init_symm_mem_reduce(self) -> None: if not self._support_custom_allreduce(): @@ -193,7 +201,13 @@ def new_deepep_group( self.ll_num_tokens = prefill_num_max_dispatch_tokens_per_rank self.ll_decode_num_tokens = decode_num_max_dispatch_tokens_per_rank self.ll_hidden = hidden_size - self.ll_num_experts = n_routed_experts + get_redundancy_expert_num() * global_world_size + total_redundant_experts = ( + get_env_start_args().eplb_num_redundant_experts_per_rank * global_world_size + if is_eplb_enabled(get_env_start_args()) + else 0 + ) + # Both phases address the same published physical expert rows. + self.ll_num_experts = n_routed_experts + total_redundant_experts self.ep_buffer = deep_ep.ElasticBuffer( deepep_group, num_max_tokens_per_rank=self.ll_num_tokens, @@ -233,7 +247,10 @@ def new_deepep_group( # FP8 MoE 的 decode 使用 legacy low-latency buffer；prefill 阶段还会将其 # 空闲的本地 RDMA storage 复用为分块 grouped GEMM 的临时 workspace。 decode_size_hint = deep_ep.Buffer.get_low_latency_rdma_size_hint( - self.ll_decode_num_tokens, self.ll_hidden, global_world_size, self.ll_num_experts + self.ll_decode_num_tokens, + self.ll_hidden, + global_world_size, + self.ll_num_experts, ) microbatch_count = len(self.groups) min_prefill_reuse_buffer_bytes = _calculate_min_chunked_expanded_moe_reuse_buffer_bytes( @@ -274,9 +291,11 @@ def new_deepep_group( moe_intermediate_size, ) logger.info( - "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, expert_quant_method_names=%s", - enable_low_latency_buffer, + "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, " + "ll_num_experts=%s, expert_quant_method_names=%s", + self.ep_low_latency_buffer is not None, enable_mega_moe_buffer, + self.ll_num_experts, sorted(expert_quant_method_names), ) theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) diff --git a/lightllm/models/glm5_next_mtp/model.py b/lightllm/models/glm5_next_mtp/model.py index 19cd9857df..5e2dae7a4a 100644 --- a/lightllm/models/glm5_next_mtp/model.py +++ b/lightllm/models/glm5_next_mtp/model.py @@ -25,6 +25,7 @@ def __init__(self, kvargs): def _init_config(self): # Reuse the validated target configuration without changing its mHC layout. self.config = deepcopy(self.main_model.config) + self.config["is_mtp_draft_model"] = self.is_mtp_draft_model assert self.config.get("num_nextn_predict_layers") == 1, "GLM NextN requires one native MTP block" self.config["mhc"] = False diff --git a/lightllm/models/qwen3_5/model.py b/lightllm/models/qwen3_5/model.py index 3315831f53..f6b3332526 100644 --- a/lightllm/models/qwen3_5/model.py +++ b/lightllm/models/qwen3_5/model.py @@ -64,6 +64,7 @@ def _init_config(self): all_config = json.load(json_file) self.config = all_config["text_config"] + self.config["is_mtp_draft_model"] = self.is_mtp_draft_model self.vision_config = all_config.get("vision_config", None) if self.vision_config is None: diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index c45a76f9df..70274c8508 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -802,15 +802,24 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""Whether to enable ep moe for deepseekv3 model.""", ) parser.add_argument( - "--ep_redundancy_expert_config_path", - type=str, - default=None, - help="""Path of the redundant expert config. It can be used for deepseekv3 model.""", + "--disable_ep_balance_monitor", + action="store_true", + help="""Disable the prefill expert balance monitor enabled by default for EP-MoE.""", ) parser.add_argument( - "--auto_update_redundancy_expert", + "--enable_prefill_eplb", action="store_true", - help="""Whether to update the redundant expert for deepseekv3 model by online expert used counter.""", + help="""Balance experts from prefill samples. PD decode automatically disables EPLB. + Single-node normal mode shares the published layout between prefill and decode; + decode does not sample loads, plan, or commit migrations.""", + ) + parser.add_argument( + "--eplb_num_redundant_experts_per_rank", + type=int, + default=2, + help="""Number of redundant physical experts per EP rank for each MoE layer used by EPLB. + Shared by both phases in normal mode; ignored on PD decode nodes. + The value must be nonnegative; zero enables rearrangement without extra replicas.""", ) parser.add_argument( "--enable_fused_shared_experts", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index c67ab9582d..5456722973 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -3,6 +3,7 @@ import uuid import subprocess import math +from lightllm.common.eplb_utils import is_eplb_enabled from lightllm.utils.start_utils import process_manager from .metrics.manager import start_metric_manager from .httpserver.req_id_generator_server import start_req_id_generator_server @@ -173,6 +174,8 @@ def _launch_subprocesses(args: StartArgs): if args.enable_dp_prefill_balance: assert args.enable_tpsp_mix_mode and args.dp > 1, "need set --enable_tpsp_mix_mode firstly and --dp > 1" + args.enable_prefill_eplb = is_eplb_enabled(args) + if args.enable_ep_moe: allowed_ep_prefill_att_backends = {"auto", "fa3", "triton", "flashqla"} for backend in args.llm_prefill_att_backend: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 3f9747831b..ecf0d9818c 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -193,8 +193,9 @@ class StartArgs: default="gpu_counter", metadata={"choices": ["cpu_counter", "pin_mem_counter", "gpu_counter"]} ) enable_ep_moe: bool = field(default=False) - ep_redundancy_expert_config_path: Optional[str] = field(default=None) - auto_update_redundancy_expert: bool = field(default=False) + disable_ep_balance_monitor: bool = field(default=False) + enable_prefill_eplb: bool = field(default=False) + eplb_num_redundant_experts_per_rank: int = field(default=2) enable_fused_shared_experts: bool = field(default=False) mtp_mode: Optional[str] = field( default=None, diff --git a/lightllm/server/metrics/metrics.py b/lightllm/server/metrics/metrics.py index 0d42462c3f..c19d756c17 100644 --- a/lightllm/server/metrics/metrics.py +++ b/lightllm/server/metrics/metrics.py @@ -32,6 +32,15 @@ "lightllm_cache_hit_rate": "Prefix cache hit rate of latest completed request", "lightllm_gen_throughput": "Generation throughput of latest completed request (tokens/s)", "lightllm_num_running_reqs": "Number of running requests", + "lightllm_prefill_ep_critical_overhead_gflops_per_routed_token": ( + "Estimated critical-path excess compute per logical routed source token, GFLOPs/token" + ), + "lightllm_prefill_ep_compute_critical_overhead_ratio": ( + "Estimated excess critical compute divided by balanced compute; 0.3 means +30%" + ), + "lightllm_prefill_ep_placement_pressure_drift": ( + "Normalized temporal drift of overloaded-rank pressure from the latest complete prefill report" + ), } @@ -111,6 +120,9 @@ def init_metrics(self, args): self.create_gauge("lightllm_cache_hit_rate") self.create_gauge("lightllm_gen_throughput") self.create_gauge("lightllm_num_running_reqs") + self.create_gauge("lightllm_prefill_ep_critical_overhead_gflops_per_routed_token") + self.create_gauge("lightllm_prefill_ep_compute_critical_overhead_ratio") + self.create_gauge("lightllm_prefill_ep_placement_pressure_drift") def create_histogram(self, name, buckets, labelnames=None): all_labels = ["model_name"] + (labelnames or []) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 6c621b6adb..2585a42957 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -51,6 +51,7 @@ from .prefill_queue_strategy import create_prefill_queue_strategy from .multi_level_kv_cache import MultiLevelKvCacheModule from lightllm.utils.profiler import ProcessProfiler, ProfilerCmd +from .ep_runtime import install_ep_runtime class ModeBackend: @@ -261,13 +262,17 @@ def init_model(self, kvargs): # 策略可能读取模型、缓存管理器和 rank 等全局调度状态，因此在 backend 完成初始化后再创建。 self.prefill_queue_strategy = create_prefill_queue_strategy(self) + install_ep_runtime(self) + self.start_infer_loops() + return + + def start_infer_loops(self): # 启动infer_loop_thread, 启动两个线程进行推理，对于具备双batch推理折叠得场景 # 可以降低 cpu overhead，大幅提升gpu得使用率。 self.infer_loop_thread = threading.Thread(target=self.infer_loop, daemon=True) self.infer_loop_thread.start() self.infer_loop_thread1 = threading.Thread(target=self.infer_loop, daemon=True) self.infer_loop_thread1.start() - return def init_custom(self): pass @@ -346,7 +351,8 @@ def init_mtp_draft_model(self, main_kvargs: dict): model_cfg=draft_model_cfg, spec_mode=spec_mode, ) - self.draft_models.append(draft_model_class(draft_model_kvargs)) + draft_model = draft_model_class(draft_model_kvargs) + self.draft_models.append(draft_model) self.logger.info(f"loaded speculative draft model class {self.draft_models[i].__class__}") return diff --git a/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py b/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py new file mode 100644 index 0000000000..5d460ed597 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py @@ -0,0 +1,359 @@ +import threading +from array import array +from typing import Optional, Tuple + +import torch +import torch.distributed as dist + +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.eplb_utils import find_ep_moe_weights +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_balance import PrefillEPBalanceCounters +from lightllm.server.metrics.manager import MetricClient +from lightllm.utils.device_utils import is_sm100_gpu +from lightllm.utils.dist_utils import ( + get_global_rank, + get_global_world_size, +) +from lightllm.utils.log_utils import init_logger +from lightllm.utils.shm_port_args import get_shm_port_args + + +logger = init_logger(__name__) + +EP_BALANCE_PREFILL_ROUNDS_PER_REPORT = 100 +EP_BALANCE_ROUND_BUFFER_CAPACITY = 4096 +EP_BALANCE_PRESSURE_DRIFT_BUCKET_ROUNDS = 20 +EP_BALANCE_PRESSURE_DRIFT_STABLE_THRESHOLD = 0.10 +ROUTE_LOAD = 0 +COMPUTE_LOAD = 1 +GFLOP = 1_000_000_000 + + +def should_enable_ep_balance_monitor(args) -> bool: + if args.enable_prefill_cudagraph or is_sm100_gpu(): + return False + return args.enable_ep_moe and not args.disable_ep_balance_monitor and args.run_mode != "decode" + + +def calculate_prefill_balance_stats( + round_stats: torch.Tensor, # [num_rounds, num_layers, world_size, 2] (route/compute) + layer_routed_experts: torch.Tensor, # [num_layers] + layer_flops_per_expert_token: torch.Tensor, # [num_layers] + layer_topks: torch.Tensor, # [num_layers] + source_token_replication: int, + report_min_route_samples_per_expert: int = 100, +) -> Optional[dict]: + """Summarize complete-prefill samples from [round, layer, rank, route/compute]. + + MoE layers execute sequentially, and every layer waits for its slowest EP + rank. Preserve the layer dimension until after taking the cross-rank max so + that different slow ranks in different layers cannot cancel each other. + """ + assert source_token_replication > 0 + + layer_route_load = round_stats[:, :, :, ROUTE_LOAD].sum(dim=(0, 2)) + minimum_route_samples = layer_routed_experts * report_min_route_samples_per_expert + if torch.any(layer_route_load < minimum_route_samples): + return None + total_route_load = layer_route_load.sum() + + compute_rank_load = round_stats[:, :, :, COMPUTE_LOAD] + compute_rank_load_float = compute_rank_load.to(torch.float64) + excess_compute_load = compute_rank_load_float.max(dim=2).values - compute_rank_load_float.mean(dim=2) + + # Every MoE expert token executes the two projections packed in w13 plus + # the w2 projection. Weight each padded compute token by the layer's + # actual matrix sizes so the metric remains comparable across models. + excess_compute_flops = (excess_compute_load * layer_flops_per_expert_token.to(torch.float64)).sum() + balanced_compute_flops = ( + compute_rank_load_float.mean(dim=2) * layer_flops_per_expert_token.to(torch.float64) + ).sum() + if balanced_compute_flops == 0: + return None + + # Non-TPSP prefill gathers one route-load copy per TP rank. Divide the + # replica count out so GFLOP/token uses logical source tokens. + source_tokens = total_route_load.to(torch.float64) / ( + layer_topks.to(torch.float64).sum() * source_token_replication + ) + if source_tokens == 0: + return None + + return { + "prefill_rounds": int(compute_rank_load.shape[0]), + "critical_overhead_gflops_per_routed_token": float((excess_compute_flops / source_tokens / GFLOP).item()), + "prefill_ep_compute_critical_overhead_ratio": float((excess_compute_flops / balanced_compute_flops).item()), + } + + +def calculate_prefill_placement_pressure_drift( + round_stats: torch.Tensor, + previous_pressure_signature: Optional[torch.Tensor] = None, + bucket_rounds: int = EP_BALANCE_PRESSURE_DRIFT_BUCKET_ROUNDS, +) -> Tuple[float, torch.Tensor]: + """Measure how prefill rank-pressure placement changes across time buckets. + + This is rank-0 CPU-only report analysis. The returned final bucket is a + compact signature that allows the next report to include the boundary pair. + """ + if bucket_rounds <= 0: + raise ValueError(f"bucket_rounds must be positive, got {bucket_rounds}") + if round_stats.ndim != 4 or round_stats.shape[-1] != 2: + raise ValueError( + "round_stats must have shape [num_rounds, num_layers, world_size, 2], " f"got {tuple(round_stats.shape)}" + ) + num_rounds, num_layers, world_size, _ = round_stats.shape + if num_rounds <= 0: + raise ValueError("round_stats must contain at least one round") + if num_layers <= 0 or world_size <= 0: + raise ValueError("round_stats must contain at least one layer and rank") + if num_rounds % bucket_rounds != 0: + raise ValueError(f"num_rounds ({num_rounds}) must be divisible by bucket_rounds ({bucket_rounds})") + + num_buckets = num_rounds // bucket_rounds + bucket_rank_load = ( + round_stats[:, :, :, COMPUTE_LOAD] + .to(torch.float64) + .reshape(num_buckets, bucket_rounds, num_layers, world_size) + .sum(dim=1) + ) + mean_rank_load = bucket_rank_load.mean(dim=2, keepdim=True).clamp_min(1) + pressure = torch.relu(bucket_rank_load / mean_rank_load - 1) + + if previous_pressure_signature is not None: + expected_shape = (num_layers, world_size) + if tuple(previous_pressure_signature.shape) != expected_shape: + raise ValueError( + "previous_pressure_signature must have shape " + f"{expected_shape}, got {tuple(previous_pressure_signature.shape)}" + ) + left = torch.cat((previous_pressure_signature.to(torch.float64).unsqueeze(0), pressure[:-1]), dim=0) + right = pressure + else: + left = pressure[:-1] + right = pressure[1:] + + total_pressure = (left + right).sum() + if total_pressure == 0: + drift = 0.0 + else: + drift = float((left - right).abs().sum().div(total_pressure).item()) + return drift, pressure[-1].clone() + + +def classify_prefill_placement_pressure_drift(drift: float) -> str: + if drift < EP_BALANCE_PRESSURE_DRIFT_STABLE_THRESHOLD: + return "stable" + return "dynamic" + + +class EPBalanceMonitor: + """Report cross-rank imbalance for non-overlapping blocks of complete prefill rounds.""" + + def __init__(self, model: TpPartBaseModel): + self.global_rank = get_global_rank() + self.world_size = get_global_world_size() + self.weights = find_ep_moe_weights(model) + self.enabled = bool(self.weights) + + if not self.enabled: + return + + self.source_token_replication = 1 if model.args.enable_tpsp_mix_mode else model.tp_world_size_ + self.counters: list[PrefillEPBalanceCounters] = [PrefillEPBalanceCounters() for _ in self.weights] + for weight, counter in zip(self.weights, self.counters): + weight.fuse_moe_impl.ep_balance_counters = counter + self.layer_routed_experts = torch.tensor( + [weight.n_routed_experts for weight in self.weights], dtype=torch.int64 + ) + self.layer_flops_per_expert_token = torch.tensor( + [ + # Each expert-token performs gate, up, and down projections; each MAC counts as 2 FLOPs. + 2 * 3 * weight.hidden_size * weight.moe_intermediate_size + for weight in self.weights + ], + dtype=torch.float64, + ) + self.layer_topks = torch.tensor( + [weight.num_experts_per_tok for weight in self.weights], + dtype=torch.float64, + ) + self._round_buffer_storage = array("q", [0]) * (EP_BALANCE_ROUND_BUFFER_CAPACITY * len(self.weights) * 2) + self._round_buffer = torch.frombuffer(self._round_buffer_storage, dtype=torch.int64).view( + EP_BALANCE_ROUND_BUFFER_CAPACITY, len(self.weights), 2 + ) + self._round_ready = threading.Event() + self._written_round_count = 0 # Prefill rounds fully written to the ring buffer. + self._processed_round_count = 0 # Prefill rounds consumed by the monitor thread. + self._overflowed = False + self._previous_pressure_signature: Optional[torch.Tensor] = None + self._common_round_end = torch.zeros((), dtype=torch.int64) + self._monitor_error = torch.zeros((), dtype=torch.int32) + + self.gloo_group = dist.new_group(ranks=list(range(self.world_size)), backend="gloo") + self.metric_client = MetricClient(get_shm_port_args().metric_port) if self.global_rank == 0 else None + threading.Thread(target=self._monitor_loop, daemon=True, name="ep-balance-monitor").start() + + def record_prefill_round(self): + """Publish one complete all-layer prefill sample to the SPSC ring.""" + if not self.enabled: + return + + written_round_count = self._written_round_count + if written_round_count - self._processed_round_count >= EP_BALANCE_ROUND_BUFFER_CAPACITY: + if not self._overflowed: + self._overflowed = True + self._round_ready.set() + return + + storage_index = (written_round_count % EP_BALANCE_ROUND_BUFFER_CAPACITY) * len(self.counters) * 2 + for counter in self.counters: + self._round_buffer_storage[storage_index] = counter.route_load + self._round_buffer_storage[storage_index + 1] = counter.compute_load + counter.route_load = 0 + counter.compute_load = 0 + storage_index += 2 + + # Publish only after the entire slot is written. The SPSC producer and + # monitor thread run under the CPython GIL, so this count is the release + # point for the corresponding ring slot. + self._written_round_count = written_round_count + 1 + if self._written_round_count - self._processed_round_count >= EP_BALANCE_PREFILL_ROUNDS_PER_REPORT: + self._round_ready.set() + + def _get_common_round_end(self) -> int: + """Return the exclusive round boundary completed by every rank.""" + self._common_round_end.fill_(self._written_round_count) + dist.all_reduce(self._common_round_end, op=dist.ReduceOp.MIN, group=self.gloo_group) + return int(self._common_round_end.item()) + + def _raise_buffer_overflow(self, phase: str, common_round_end: Optional[int] = None): + message = ( + "EP balance prefill-round buffer overflowed " + f"phase={phase} written={self._written_round_count} " + f"processed={self._processed_round_count} capacity={EP_BALANCE_ROUND_BUFFER_CAPACITY}" + ) + if common_round_end is not None: + message += f" common_round_end={common_round_end}" + raise RuntimeError(message) + + def _copy_local_rounds(self, start: int, end: int) -> torch.Tensor: + """Copy local prefill-round loads in the half-open range [start, end).""" + num_rounds = end - start + if num_rounds > EP_BALANCE_ROUND_BUFFER_CAPACITY: + raise ValueError("requested EP balance round range exceeds ring capacity") + start_index = start % EP_BALANCE_ROUND_BUFFER_CAPACITY + if start_index + num_rounds <= EP_BALANCE_ROUND_BUFFER_CAPACITY: + return self._round_buffer[start_index : start_index + num_rounds].clone() + end_index = (start_index + num_rounds) % EP_BALANCE_ROUND_BUFFER_CAPACITY + return torch.cat((self._round_buffer[start_index:], self._round_buffer[:end_index]), dim=0) + + def _gather_round_stats(self, local_round_stats: torch.Tensor) -> Optional[torch.Tensor]: + """Gather rank-local stats as [round, layer, rank, route/compute].""" + gathered = ( + [torch.empty_like(local_round_stats) for _ in range(self.world_size)] if self.global_rank == 0 else None + ) + dist.gather(local_round_stats, gather_list=gathered, dst=0, group=self.gloo_group) + if self.global_rank != 0: + return None + # [rank, round, layer, route/compute] + # -> [round, layer, rank, route/compute] + return torch.stack(gathered).permute(1, 2, 0, 3) + + def _log_stats(self, round_stats: torch.Tensor): + """Compute and log balance statistics for one complete global window.""" + compute = calculate_prefill_balance_stats( + round_stats, + self.layer_routed_experts, + self.layer_flops_per_expert_token, + self.layer_topks, + self.source_token_replication, + ) + if compute is None: + return + + drift, self._previous_pressure_signature = calculate_prefill_placement_pressure_drift( + round_stats, + previous_pressure_signature=self._previous_pressure_signature, + ) + drift_state = classify_prefill_placement_pressure_drift(drift) + + logger.info( + "ep_balance " + f"phase=prefill prefill_rounds={compute['prefill_rounds']} " + "prefill_ep_critical_overhead_gflops_per_routed_token=" + f"{compute['critical_overhead_gflops_per_routed_token']:.4f} " + "prefill_ep_compute_critical_overhead_ratio=" + f"{compute['prefill_ep_compute_critical_overhead_ratio']:.4f} " + f"prefill_ep_placement_pressure_drift={drift:.4f} " + f"prefill_ep_placement_pressure_state={drift_state}" + ) + self.metric_client.gauge_set( + "lightllm_prefill_ep_critical_overhead_gflops_per_routed_token", + compute["critical_overhead_gflops_per_routed_token"], + ) + self.metric_client.gauge_set( + "lightllm_prefill_ep_compute_critical_overhead_ratio", + compute["prefill_ep_compute_critical_overhead_ratio"], + ) + self.metric_client.gauge_set("lightllm_prefill_ep_placement_pressure_drift", drift) + + def _sync_monitor_error(self, error: Optional[Exception], phase: str): + """Stop every rank at the same boundary if one rank cannot continue.""" + self._monitor_error.fill_(int(error is not None)) + dist.all_reduce(self._monitor_error, op=dist.ReduceOp.MAX, group=self.gloo_group) + if self._monitor_error.item(): + if error is not None: + raise RuntimeError(f"EP balance monitor failed phase={phase}: {error}") from error + raise RuntimeError(f"EP balance monitor failed on another rank phase={phase}") + + def _monitor_loop(self): + """Consume commonly completed rounds in background report-sized windows.""" + try: + while True: + self._round_ready.wait() + self._round_ready.clear() + error = None + try: + if self._overflowed: + self._raise_buffer_overflow("before_sync") + except Exception as exc: + error = exc + self._sync_monitor_error(error, "before_sync") + common_round_end = self._get_common_round_end() + + while common_round_end - self._processed_round_count >= EP_BALANCE_PREFILL_ROUNDS_PER_REPORT: + round_start = self._processed_round_count + round_end = round_start + EP_BALANCE_PREFILL_ROUNDS_PER_REPORT + error = None + try: + if common_round_end - round_start > EP_BALANCE_ROUND_BUFFER_CAPACITY: + self._raise_buffer_overflow("common_round_lag", common_round_end=common_round_end) + local_round_stats = self._copy_local_rounds(round_start, round_end) + if self._overflowed: + self._raise_buffer_overflow("after_copy") + except Exception as exc: + error = exc + self._sync_monitor_error(error, "copy") + error = None + try: + round_stats = self._gather_round_stats(local_round_stats) + self._processed_round_count = round_end + if self.global_rank == 0: + self._log_stats(round_stats) + except Exception as exc: + error = exc + self._sync_monitor_error(error, "report") + + if self._written_round_count - self._processed_round_count >= EP_BALANCE_PREFILL_ROUNDS_PER_REPORT: + self._round_ready.set() + except Exception as exc: + logger.exception(f"EP balance monitor stopped unexpectedly: {exc}") + self._disable() + return + + def _disable(self): + """Detach counters from MoE weights and disable monitoring.""" + for weight in self.weights: + weight.fuse_moe_impl.ep_balance_counters = None + self.enabled = False diff --git a/lightllm/server/router/model_infer/mode_backend/ep_runtime.py b/lightllm/server/router/model_infer/mode_backend/ep_runtime.py new file mode 100644 index 0000000000..bfdcb1384a --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/ep_runtime.py @@ -0,0 +1,78 @@ +"""Bind optional EP services at initialization, outside the inference loop.""" + +from functools import wraps + +import torch.distributed as dist +from lightllm.common.eplb_utils import is_eplb_enabled + +from .ep_balance_monitor import EPBalanceMonitor, should_enable_ep_balance_monitor + + +def install_ep_runtime(backend): + """Install on the target model only, before either inference thread starts.""" + manager = None + if is_eplb_enabled(backend.args): + from .eplb_manager import EPLBManager + + manager = EPLBManager(backend.model) + if not backend.args.enable_prefill_microbatch_overlap and not backend.args.enable_tpsp_mix_mode: + # These MoE calls are globally ordered under the forward event. + # Commit ready layers here, but finish the generation only at + # the outer forward boundary. + for state in manager._eplb_states: + state.migration_poll = manager.poll_migration_at_layer_boundary + + monitor = EPBalanceMonitor(backend.model) if should_enable_ep_balance_monitor(backend.args) else None + if monitor is not None and not monitor.enabled: + monitor = None + + if manager is None and monitor is None: + return + + def record_prefill(prefill): + @wraps(prefill) + def run(*args, **kwargs): + if manager is not None: + if backend.args.run_mode == "normal": + # Run control collectives and publish weights only while + # owning a prefill forward, never on decode/idle rounds. + manager.poll() + manager.before_prefill() + output = prefill(*args, **kwargs) + if monitor is not None: + monitor.record_prefill_round() + if manager is not None: + manager.step() + return output + + return run + + # These are the two completed-prefill boundaries; decode and draft models + # retain their original methods. Failed forwards never publish a sample. + model = backend.model + model._prefill = record_prefill(model._prefill) + model._microbatch_overlap_prefill_cuda = record_prefill(model._microbatch_overlap_prefill_cuda) + + if manager is not None: + dist.barrier() + if backend.args.run_mode == "normal": + return + get_event_pack = backend.overlap_event_manager.get_overlap_event_pack + + @wraps(get_event_pack) + def get_overlap_event_pack(): + event_pack = get_event_pack() + # Capture the event rather than a bound pack method: storing the + # latter back on the pack would create a cycle on every batch. + wait_to_forward = event_pack.wait_forward_event.wait + + def wait_and_poll(): + wait_to_forward() + # The existing event serializes both inference threads. Poll + # here even on idle rounds, before request/DP collectives. + manager.poll() + + event_pack.wait_to_forward = wait_and_poll + return event_pack + + backend.overlap_event_manager.get_overlap_event_pack = get_overlap_event_pack diff --git a/lightllm/server/router/model_infer/mode_backend/eplb_manager.py b/lightllm/server/router/model_infer/mode_backend/eplb_manager.py new file mode 100644 index 0000000000..8012c4a99f --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb_manager.py @@ -0,0 +1,658 @@ +import os +import threading +import time +from pathlib import Path +from typing import Dict +from collections import deque + +import torch +import torch.distributed as dist + +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.eplb_utils import find_ep_moe_weights +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import ( + build_logical_to_physical_maps_for_layers, + select_improving_placements, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_planner import ( + plan_eplb_candidate, +) +from lightllm.server.router.model_infer.mode_backend.eplb_transfer import ( + NixlEPLBTransfer, + align_target_placement, + build_transfer_plan, +) +from lightllm.utils.dist_utils import get_global_rank, get_global_world_size, get_node_world_size +from lightllm.common.eplb_planner_worker import EPLBPlannerWorker +from lightllm.common.eplb_utils import EPLB_STEADY_SAMPLE_STEPS +from lightllm.utils.envs_utils import ( + get_eplb_placement_stickiness, + get_eplb_planner_process, + get_eplb_rebalance_gain_threshold, + get_prefill_eplb_step_interval, + get_prefill_eplb_steady_sample_steps, + get_prefill_eplb_min_rebalance_interval, +) +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) +EPLB_MIN_AVG_TOKENS_PER_EXPERT = 100 +EPLB_EXPERT_ALIGNMENT = 128 +EPLB_CONTROL_ERROR = -1 + + +class EPLBManager: + """Online EPLB with asynchronous GPU expert migration.""" + + def __init__(self, model: TpPartBaseModel): + self.weights = find_ep_moe_weights(model) + assert self.weights, "EPLB requires at least one EP MoE layer" + self.global_rank = get_global_rank() + self.world_size = get_global_world_size() + self.node_world_size = get_node_world_size() + self._eplb_states = [weight.expert_parallel_state.eplb for weight in self.weights] + self.step_interval = get_prefill_eplb_step_interval() + self.rebalance_gain_threshold = get_eplb_rebalance_gain_threshold() + self.placement_stickiness = get_eplb_placement_stickiness() + self.steady_sample_steps = get_prefill_eplb_steady_sample_steps() + self.min_rebalance_interval = get_prefill_eplb_min_rebalance_interval() + self.sampling_interval = self.step_interval + self.prefill_steps = 0 + self._initial_plan_complete = False + self._use_full_window = True + self._next_evaluation_step = self.step_interval + self._last_evaluation_step = 0 + self._sample_step_counts = deque([0], maxlen=max(self.step_interval, EPLB_STEADY_SAMPLE_STEPS) + 1) + self._evaluation_snapshot = None + routed = {weight.expert_parallel_state.num_logical_experts for weight in self.weights} + redundant = {state.num_redundant_experts_per_rank for state in self._eplb_states} + assert len(routed) == len(redundant) == 1 + self.num_logical_experts = routed.pop() + self.num_redundant_experts_per_rank = redundant.pop() + self.current_placement = torch.stack([state.physical_to_logical for state in self._eplb_states]) + self.in_flight = False + self.target_placement = None + self.target_metadata = None + self.in_flight_started_at = None + self.evaluation_in_flight = False + self._evaluation_lock = threading.Lock() + self._evaluation_result = None + self._evaluation_error = None + self._evaluation_thread = None + self._diagnostics_dir = os.getenv("LIGHTLLM_EPLB_DIAGNOSTICS_DIR") or None + self._diagnostic_dump_sequence = 0 + self._planner_worker = EPLBPlannerWorker() if self.global_rank == 0 and get_eplb_planner_process() else None + self._reset_recorded_samples() + self._set_recording(True) + # Keep background evaluation collectives separate from the main-thread + # control/poll collectives: their ordering is intentionally independent. + self.evaluation_group = dist.new_group(list(range(self.world_size)), backend="gloo") + self.control_group = dist.new_group(list(range(self.world_size)), backend="gloo") + self.transfer_group = dist.new_group(list(range(self.world_size)), backend="gloo") + # This control-group scalar is only touched from the main inference + # thread, never by the background evaluation thread. + self._control_ready_count = torch.empty(1, dtype=torch.int32) + self.transfer = NixlEPLBTransfer(self.weights, self.transfer_group, self.global_rank, self.world_size) + if self.global_rank == 0: + logger.info( + "eplb enabled " + f"layers={len(self.weights)} num_logical_experts={self.num_logical_experts} " + f"num_redundant_experts_per_rank={self.num_redundant_experts_per_rank} " + f"step_interval={self.step_interval} " + f"rebalance_gain_threshold={self.rebalance_gain_threshold:.4f} " + f"placement_stickiness={self.placement_stickiness:.4f} " + f"steady_sample_steps={self.steady_sample_steps} min_rebalance_interval={self.min_rebalance_interval} " + f"planner_process={self._planner_worker is not None} " + f"planner_worker={getattr(self._planner_worker, 'info', None)}" + ) + + def before_prefill(self): + """Clear the next dispatch slots on the forward's own CUDA stream.""" + slots = [] + for state in self._eplb_states: + capacity = state.route_counter.shape[0] + index = state.recorded_sample_count % capacity + # A forward dispatches once, or twice with microbatch overlap. + # Capacity covers twice the largest forward window, so clearing + # both next slots also preserves the latest window for one call. + slots.extend((state.route_counter[index], state.route_counter[(index + 1) % capacity])) + torch._foreach_zero_(slots) + + def poll(self): + """Poll only from a globally ordered pre-forward boundary.""" + if self.in_flight: + self._poll_in_flight() + return + if self.evaluation_in_flight and self._evaluation_ready_on_all_ranks(): + self._poll_evaluation() + + def step(self): + self.prefill_steps += 1 + count = self._eplb_states[0].recorded_sample_count + previous = self._sample_step_counts[-1] + if count - previous not in (1, 2) or any(state.recorded_sample_count != count for state in self._eplb_states): + raise RuntimeError("EPLB forward must record one or two matching dispatches per layer") + self._sample_step_counts.append(count) + if self.in_flight or self.evaluation_in_flight or self.prefill_steps < self._next_evaluation_step: + return + self._last_evaluation_step = self.prefill_steps + self._start_evaluation() + self._next_evaluation_step = (self.prefill_steps // self.sampling_interval + 1) * self.sampling_interval + + def _set_recording(self, enabled: bool): + for state in self._eplb_states: + state.recording = enabled + + def _reset_recorded_samples(self): + counters = [state.route_counter for state in self._eplb_states] + if counters: + torch._foreach_zero_(counters) + for state in self._eplb_states: + state.recorded_sample_count = 0 + + def _control_count(self, value: int) -> torch.Tensor: + """Return the main-thread-only reusable control collective scalar.""" + return self._control_ready_count.fill_(value) + + def _steady_sample_window_steps(self) -> int: + # Preserve the original four-step sampling window when the base interval is smaller. + return min( + getattr(self, "steady_sample_steps", EPLB_STEADY_SAMPLE_STEPS), + max(self.step_interval, EPLB_STEADY_SAMPLE_STEPS), + self.sampling_interval, + ) + + def _prepare_next_sampling_window(self): + """Schedule the next boundary without resetting the live ring.""" + self._set_recording(True) + self._next_evaluation_step = (self._last_evaluation_step // self.sampling_interval + 1) * self.sampling_interval + + def _collect_local_samples(self) -> torch.Tensor: + snapshot = self._evaluation_snapshot + values = ( + snapshot["count"], + snapshot["capacity"], + snapshot["total_count"], + snapshot["steps"], + snapshot["completed_step"], + ) + metadata = torch.tensor([value for item in values for value in (item, -item)], dtype=torch.int64) + dist.all_reduce(metadata, op=dist.ReduceOp.MIN, group=self.evaluation_group) + if any(metadata[i] != -metadata[i + 1] for i in range(0, len(metadata), 2)): + raise RuntimeError("EPLB frozen sample counts differ between ranks") + boundaries = torch.tensor( + [value for item in snapshot["forward_end_counts"] for value in (item, -item)], dtype=torch.int64 + ) + dist.all_reduce(boundaries, op=dist.ReduceOp.MIN, group=self.evaluation_group) + if any(boundaries[i] != -boundaries[i + 1] for i in range(0, len(boundaries), 2)): + raise RuntimeError("EPLB forward-to-dispatch boundaries differ between ranks") + samples = snapshot["samples"].cpu() + # Planning only needs CPU tensors; release the frozen GPU allocation. + snapshot["samples"] = None + return samples + + def _commit_layer_metadata(self, layer_index: int): + eplb_state = self._eplb_states[layer_index] + logical_to_physical, replica_count = self.target_metadata[layer_index] + eplb_state.logical_to_physical_map.copy_(logical_to_physical, non_blocking=True) + eplb_state.logical_replica_count.copy_(replica_count, non_blocking=True) + placement = self.target_placement[layer_index] + eplb_state.physical_to_logical = placement.clone() + eplb_state.placement_generation += 1 + # Keep the CPU state truthful even while only a prefix of layers has + # been committed; subsequent planning waits for the whole generation. + self.current_placement[layer_index].copy_(placement) + + def _finish_rebalance(self): + self.current_placement = self.target_placement + self.target_placement = None + self.target_metadata = None + self.in_flight = False + self._prepare_next_sampling_window() + if self.global_rank == 0: + timing = getattr(self, "_inflight_timing", {}) + logger.info( + "eplb completed wall_time=%.2fs control_cpu_ms=%.3f control_poll_calls=%s " + "commit_cpu_ms=%.3f commit_max_ms=%.3f commit_calls=%s " + "(CPU wall; commit includes existing stream waits, not net stall)", + time.time() - self.in_flight_started_at, + timing.get("control_cpu_ms", 0.0), + timing.get("control_poll_calls", 0), + timing.get("commit_cpu_ms", 0.0), + timing.get("commit_max_ms", 0.0), + timing.get("commit_calls", 0), + ) + + def poll_migration_at_layer_boundary(self): + """Commit ready layers without publishing a completed generation.""" + if self.in_flight and self.in_flight_layers: + self._poll_in_flight(finish_generation=False) + + def _poll_in_flight(self, finish_generation=True): + if not self.in_flight_layers: + if finish_generation: + self.transfer.finish() + self._finish_rebalance() + return + local_error = None + try: + pending = self.transfer.pending_layers() + except BaseException as exc: + pending = [] + local_error = exc + ready_count = self._control_count(EPLB_CONTROL_ERROR if local_error is not None else len(pending)) + control_started = time.perf_counter() + dist.all_reduce(ready_count, op=dist.ReduceOp.MIN, group=self.control_group) + timing = getattr(self, "_inflight_timing", None) + if timing is not None: + timing["control_poll_calls"] += 1 + timing["control_cpu_ms"] += (time.perf_counter() - control_started) * 1000.0 + ready_count = int(ready_count.item()) + if ready_count < 0: + if local_error is not None: + raise RuntimeError("EPLB transfer worker failed on this rank") from local_error + raise RuntimeError("EPLB transfer worker failed on another rank") + if ready_count == 0: + return + if ready_count > len(pending) or ready_count > len(self.in_flight_layers): + raise RuntimeError("EPLB global ready count exceeds the local ordered prefix") + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + # Previous forward is queued on the shared overlap stream; order the + # live-weight commit after it. The subsequent wait orders the next forward. + torch.cuda.current_stream().wait_stream(g_infer_context.get_overlap_stream()) + for layer_index, buffer_index in pending[:ready_count]: + if layer_index != self.in_flight_layers[0]: + raise RuntimeError( + f"EPLB pending layer {layer_index} does not match expected {self.in_flight_layers[0]}" + ) + commit_started = time.perf_counter() + self.transfer.commit(layer_index, buffer_index, lambda: self._commit_layer_metadata(layer_index)) + commit_ms = (time.perf_counter() - commit_started) * 1000.0 + timing = getattr(self, "_inflight_timing", None) + if timing is not None: + timing["commit_calls"] += 1 + timing["commit_cpu_ms"] += commit_ms + timing["commit_max_ms"] = max(timing["commit_max_ms"], commit_ms) + self.in_flight_layers.pop(0) + if not self.in_flight_layers and finish_generation: + self.transfer.finish() + self._finish_rebalance() + + def _plan_and_broadcast(self, global_load: torch.Tensor): + """Plan on rank zero and share the serializable result on the evaluation group.""" + replica_hash_load = global_load[..., 1:] if global_load.ndim == 5 else None + if replica_hash_load is not None: + global_load = global_load[..., 0] + result = None + local_error = None + current_placement = self._evaluation_snapshot["placement"] + if self.global_rank == 0: + try: + minimum = self.num_logical_experts * EPLB_MIN_AVG_TOKENS_PER_EXPERT + layer_samples = global_load.sum(dim=(0, 2, 3)) + if torch.any(layer_samples < minimum): + result = { + "kind": "insufficient", + "minimum_layer_samples": int(layer_samples.min().item()), + "minimum": minimum, + } + else: + settings = { + "node_world_size": self.node_world_size, + "expert_alignment": EPLB_EXPERT_ALIGNMENT, + "stickiness": self.placement_stickiness, + } + if replica_hash_load is not None: + settings["replica_hash_load"] = replica_hash_load + if self._initial_plan_complete: + settings["incremental"] = True + worker = getattr(self, "_planner_worker", None) + if worker is None: + candidate = plan_eplb_candidate(global_load, current_placement, **settings) + candidate_timing = {"compute_ms": None, "roundtrip_ms": None} + else: + candidate, candidate_timing = worker.plan( + global_load.cpu(), current_placement.cpu(), **settings + ) + select_align_started = time.perf_counter() + placement, improved, metrics, before_load, after_load = select_improving_placements( + global_load, + current_placement, + candidate, + expert_alignment=EPLB_EXPERT_ALIGNMENT, + node_world_size=self.node_world_size, + rebalance_gain_threshold=self.rebalance_gain_threshold, + replica_hash_load=replica_hash_load, + ) + if bool(torch.any(improved)): + # A planner placement identifies experts by rank, not + # by redundant slot. Canonicalize every selected row + # before broadcasting so transfer, metadata, and the + # next current_placement all describe the same live + # physical expert rows. + placement = placement.clone() + for layer_index in torch.nonzero(improved, as_tuple=False).flatten().tolist(): + placement[layer_index] = align_target_placement( + current_placement[layer_index], placement[layer_index] + ) + select_align_ms = (time.perf_counter() - select_align_started) * 1000.0 + result = { + "candidate_timing_ms": { + **candidate_timing, + "select_align_ms": select_align_ms, + }, + "kind": "planned" if bool(torch.any(improved)) else "no_improvement", + "placement": placement, + "improved": improved, + "before": _imbalance_summary(before_load), + "after": _imbalance_summary(after_load), + **metrics, + } + except BaseException as exc: + local_error = exc + result = {"kind": "error", "message": f"{type(exc).__name__}: {exc}"} + broadcast_started = time.perf_counter() + if self.world_size > 1: + result_list = [result] + dist.broadcast_object_list(result_list, src=0, group=self.evaluation_group) + result = result_list[0] + broadcast_ms = (time.perf_counter() - broadcast_started) * 1000.0 + if result["kind"] == "error": + if local_error is not None: + raise RuntimeError("EPLB planner failed on rank zero") from local_error + raise RuntimeError(f"EPLB planner failed on rank zero: {result['message']}") + result["broadcast_ms"] = broadcast_ms + return result + + def _dump_evaluation_samples(self, global_load, recorded_sample_count, sample_window_steps): + diagnostics_dir = getattr(self, "_diagnostics_dir", None) + if not diagnostics_dir or self.global_rank != 0: + return 0.0 + started = time.perf_counter() + try: + directory = Path(diagnostics_dir) + directory.mkdir(parents=True, exist_ok=True) + self._diagnostic_dump_sequence = getattr(self, "_diagnostic_dump_sequence", 0) + 1 + payload = { + "global_load": global_load[..., 0] if global_load.ndim == 5 else global_load, + "current_placement": self._evaluation_snapshot["placement"], + "prefill_steps": self._evaluation_snapshot["completed_step"], + "recorded_sample_count": recorded_sample_count, + "sample_window_steps": sample_window_steps, + } + torch.save(payload, directory / f"evaluation-{self._diagnostic_dump_sequence:06d}.pt") + except Exception as exc: + logger.warning("EPLB diagnostics dump disabled after write failure: %s", exc) + self._diagnostics_dir = None + return (time.perf_counter() - started) * 1000.0 + + def _evaluate_after_event(self, event: torch.cuda.Event): + """Run the CPU/Gloo planning phase after the frozen CUDA counters are ready.""" + try: + timing_started = time.perf_counter() + thread_cpu_started = time.thread_time() + torch.cuda.set_device(self._eplb_states[0].route_counter.device) + event.synchronize() + event_wait_ms = (time.perf_counter() - timing_started) * 1000.0 + collect_started = time.perf_counter() + local_load = self._collect_local_samples() + recorded_sample_count = int(local_load.shape[0]) + sample_window_steps = self._evaluation_snapshot["steps"] + num_nodes = self.world_size // self.node_world_size + # Preserve source nodes until physical-replica loads are combined; + # DeepEP applies expert alignment after traffic from all sources + # reaches each destination expert. + global_load = torch.zeros((*local_load.shape[:2], num_nodes, *local_load.shape[2:]), dtype=local_load.dtype) + global_load[:, :, self.global_rank // self.node_world_size] = local_load + dist.all_reduce(global_load, op=dist.ReduceOp.SUM, group=self.evaluation_group) + collect_reduce_ms = (time.perf_counter() - collect_started) * 1000.0 + diagnostic_dump_ms = self._dump_evaluation_samples(global_load, recorded_sample_count, sample_window_steps) + plan_started = time.perf_counter() + result = self._plan_and_broadcast(global_load) + plan_broadcast_ms = (time.perf_counter() - plan_started) * 1000.0 + prepare_started = time.perf_counter() + result["recorded_sample_count"] = recorded_sample_count + result["sample_window_steps"] = sample_window_steps + if result["kind"] == "planned": + metadata = [None] * len(self.weights) + layer_plans = [] + improved_layer_indices = torch.nonzero(result["improved"], as_tuple=False).flatten() + if improved_layer_indices.numel(): + maps_for_improved_layers, counts_for_improved_layers = build_logical_to_physical_maps_for_layers( + result["placement"][improved_layer_indices], + self.num_logical_experts, + source_rank=self.global_rank, + node_world_size=self.node_world_size, + replica_rotation_rank=self._eplb_states[0].replica_rotation_rank, + ) + for improved_layer_offset, layer_index in enumerate(improved_layer_indices.tolist()): + placement = result["placement"][layer_index] + metadata[layer_index] = ( + maps_for_improved_layers[improved_layer_offset], + counts_for_improved_layers[improved_layer_offset], + ) + layer_plans.append( + ( + layer_index, + build_transfer_plan( + self.current_placement[layer_index], + placement, + self.num_logical_experts, + self.world_size, + self.node_world_size, + ), + ) + ) + result["metadata"] = metadata + result["layer_plans"] = layer_plans + result["prepared_batches"] = self.transfer.prepare_transfer(layer_plans) + prepare_ms = (time.perf_counter() - prepare_started) * 1000.0 + candidate_timing = result.get("candidate_timing_ms", {}) + result["evaluation_timing_ms"] = { + "event_wait": event_wait_ms, + "collect_reduce": collect_reduce_ms, + "plan_broadcast": plan_broadcast_ms, + "candidate_compute": candidate_timing.get("compute_ms"), + "candidate_roundtrip": candidate_timing.get("roundtrip_ms"), + "select_align": candidate_timing.get("select_align_ms"), + "broadcast": result.get("broadcast_ms"), + "prepare": prepare_ms, + "diagnostic_dump": diagnostic_dump_ms, + "total": (time.perf_counter() - timing_started) * 1000.0, + "parent_thread_cpu": (time.thread_time() - thread_cpu_started) * 1000.0, + } + with self._evaluation_lock: + self._evaluation_result = result + except BaseException as exc: + with self._evaluation_lock: + self._evaluation_error = exc + + def _start_evaluation(self): + with self._evaluation_lock: + self._evaluation_result = None + self._evaluation_error = None + self._freeze_evaluation_samples() + event = torch.cuda.Event() + event.record(torch.cuda.current_stream()) + self.evaluation_in_flight = True + self._evaluation_thread = threading.Thread(target=self._evaluate_after_event, args=(event,), daemon=True) + self._evaluation_thread.start() + + def _freeze_evaluation_samples(self): + """Freeze a chronological forward window before live-ring writes resume.""" + requested = self.step_interval if self._use_full_window else self._steady_sample_window_steps() + steps = min(requested, len(self._sample_step_counts) - 1) + total = self._sample_step_counts[-1] + count = total - self._sample_step_counts[-steps - 1] + counters = [state.route_counter for state in self._eplb_states] + capacity = counters[0].shape[0] + if ( + steps == 0 + or count > capacity + or any( + state.recorded_sample_count != total or state.route_counter.shape != counters[0].shape + for state in self._eplb_states + ) + ): + raise RuntimeError("EPLB continuous window exceeds its ring capacity or differs between layers") + indices = torch.arange(total - count, total, dtype=torch.int64, device=counters[0].device) % capacity + samples = torch.stack(counters, dim=1).index_select(0, indices) + self._evaluation_snapshot = dict( + samples=samples, + count=count, + total_count=total, + capacity=capacity, + steps=steps, + placement=self.current_placement.clone(), + completed_step=self.prefill_steps, + forward_end_counts=tuple(self._sample_step_counts)[-steps - 1 :], + ) + + def _poll_evaluation(self): + if not self.evaluation_in_flight: + return False + with self._evaluation_lock: + error = self._evaluation_error + result = self._evaluation_result + if error is not None or result is not None: + self._evaluation_result = None + self._evaluation_error = None + if error is not None: + self._evaluation_thread.join() + self.evaluation_in_flight = False + self._evaluation_thread = None + raise error + if result is None: + return True + self._evaluation_thread.join() + self.evaluation_in_flight = False + self._evaluation_thread = None + if self.global_rank == 0 and result.get("evaluation_timing_ms"): + logger.info( + "eplb evaluation CPU wall ms event_wait=%.3f collect_reduce=%.3f " + "plan_broadcast=%.3f prepare=%.3f diagnostic_dump=%.3f total=%.3f " + "candidate_compute=%s candidate_roundtrip=%s select_align=%s broadcast=%s parent_thread_cpu=%s", + *[ + result["evaluation_timing_ms"][key] + for key in ( + "event_wait", + "collect_reduce", + "plan_broadcast", + "prepare", + "diagnostic_dump", + "total", + "candidate_compute", + "candidate_roundtrip", + "select_align", + "broadcast", + "parent_thread_cpu", + ) + ], + ) + if result["kind"] == "insufficient": + self._use_full_window = True + self._prepare_next_sampling_window() + return False + if result["kind"] == "no_improvement": + self._initial_plan_complete = True + self._use_full_window = False + self.sampling_interval = max(self.step_interval, self.min_rebalance_interval) + if self.global_rank == 0: + logger.info( + "eplb skip rearrangement: no model improvement model_imbalance_ratio=%.4f " + "candidate_model_imbalance_ratio=%.4f candidate_rebalance_gain=%.4f " + "candidate_changed_layer_count=%s actual_changed_layer_count=0 next_sampling_interval=%s " + "recorded_sample_count=%s sample_window_steps=%s", + result["model_imbalance_ratio"], + result["candidate_model_imbalance_ratio"], + result["candidate_rebalance_gain"], + result["candidate_changed_layer_count"], + self.sampling_interval, + result.get("recorded_sample_count"), + result.get("sample_window_steps"), + ) + self._prepare_next_sampling_window() + return False + self._start_rebalance(result) + return True + + def _evaluation_ready_on_all_ranks(self) -> bool: + with self._evaluation_lock: + local_error = self._evaluation_error + local_result = self._evaluation_result + local_status = EPLB_CONTROL_ERROR if local_error is not None else int(local_result is not None) + ready_count = self._control_count(local_status) + dist.all_reduce(ready_count, op=dist.ReduceOp.MIN, group=self.control_group) + ready_count = int(ready_count.item()) + if ready_count < 0: + if local_error is not None: + raise RuntimeError("EPLB evaluation failed on this rank") from local_error + raise RuntimeError("EPLB evaluation failed on another rank") + return bool(ready_count) + + def _start_rebalance(self, result): + placement = result["placement"] + layer_plans = result["layer_plans"] + self.sampling_interval = max(self.step_interval, getattr(self, "min_rebalance_interval", 0)) + self._initial_plan_complete = True + self._use_full_window = False + self.target_placement = placement + self.target_metadata = result["metadata"] + self.in_flight_layers = [layer_index for layer_index, _ in layer_plans] + self.in_flight = True + self.in_flight_started_at = time.time() + self._inflight_timing = { + "control_poll_calls": 0, + "control_cpu_ms": 0.0, + "commit_calls": 0, + "commit_cpu_ms": 0.0, + "commit_max_ms": 0.0, + } + self.transfer.start(layer_plans, result["prepared_batches"]) + if self.global_rank == 0: + actual_changed_slot_count = sum(len(plan) for _, plan in layer_plans) + cross_node_transfer_count = sum( + step.src_rank // self.node_world_size != step.dst_rank // self.node_world_size + for _, plan in layer_plans + for step in plan + ) + logger.info( + "eplb started prefill_steps=%s max_before=%.4f max_after=%.4f p95_before=%.4f p95_after=%.4f " + "model_imbalance_ratio=%.4f candidate_model_imbalance_ratio=%.4f " + "candidate_rebalance_gain=%.4f candidate_changed_layer_count=%s " + "actual_changed_layer_count=%s actual_changed_slot_count=%s cross_node_transfer_count=%s " + "recorded_sample_count=%s sample_window_steps=%s", + self.prefill_steps, + result["before"]["max"], + result["after"]["max"], + result["before"]["p95"], + result["after"]["p95"], + result["model_imbalance_ratio"], + result["candidate_model_imbalance_ratio"], + result["candidate_rebalance_gain"], + result["candidate_changed_layer_count"], + len(layer_plans), + actual_changed_slot_count, + cross_node_transfer_count, + result.get("recorded_sample_count"), + result.get("sample_window_steps"), + ) + + +def _imbalance_summary(rank_load: torch.Tensor) -> Dict[str, float]: + if rank_load.ndim == 2: + critical = rank_load.max(dim=1).values + mean = rank_load.mean(dim=1) + elif rank_load.ndim == 3: + critical = rank_load.max(dim=2).values.sum(dim=0) + mean = rank_load.mean(dim=2).sum(dim=0) + else: + raise ValueError("rank_load must be [layers, ranks] or [samples, layers, ranks]") + layer_imbalance = critical / mean.clamp_min(1.0) + sorted_imbalance = torch.sort(layer_imbalance).values + p95_index = max(0, (95 * layer_imbalance.numel() + 99) // 100 - 1) + return { + "max": float(layer_imbalance.max().item()), + "p95": float(sorted_imbalance[p95_index].item()), + } diff --git a/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py b/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py new file mode 100644 index 0000000000..406c5c29ff --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py @@ -0,0 +1,785 @@ +"""Asynchronous expert-row migration for EPLB.""" +import ctypes +import math +import os +import re +import socket +import threading +from collections import defaultdict, deque +from dataclasses import dataclass +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +import torch +import torch.distributed as dist + +from lightllm.common.eplb_utils import ( + EPLB_FULL_STAGING_PARTS, + extract_eplb_expert_tensors, + get_eplb_staging_shape, + get_eplb_staging_tensor_shape, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import validate_physical_placement + + +@dataclass(frozen=True) +class TransferStep: + """Source and destination rows are absolute positions in the physical layout.""" + + dst_rank: int + dst_slot: int + src_rank: int + src_local_row: int + + +def align_target_placement(current: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Canonicalize a target row layout without moving retained experts. + + EPLB placement is rank-based: movable slots on one rank are + interchangeable. Retained experts therefore keep their live physical + slot, while new experts fill freed slots in the planner's target-row + order. The returned placement is the single canonical layout that must + be used both for transfers and for published routing metadata. + """ + assert current.ndim == target.ndim == 2 + assert tuple(current.shape) == tuple(target.shape) + + current_rows = current.tolist() + target_rows = target.tolist() + aligned_target_rows = [] + for current_row, target_row in zip(current_rows, target_rows): + current_slots = {expert: slot for slot, expert in enumerate(current_row)} + target_experts = set(target_row) + aligned_row = list(current_row) + freed_slots = [slot for slot, expert in enumerate(current_row) if expert not in target_experts] + new_experts = [expert for expert in target_row if expert not in current_slots] + assert len(freed_slots) == len(new_experts) + for slot, expert in zip(freed_slots, new_experts): + aligned_row[slot] = expert + aligned_target_rows.append(aligned_row) + return target.new_tensor(aligned_target_rows) + + +def build_transfer_plan( + current: torch.Tensor, + target: torch.Tensor, + num_logical_experts: int, + world_size: int, + node_world_size: int, +) -> List[TransferStep]: + assert tuple(current.shape) == tuple(target.shape) == (world_size, current.shape[1]) + validate_physical_placement(current[None], num_logical_experts) + validate_physical_placement(target[None], num_logical_experts) + current_rows = current.tolist() + aligned_target_rows = align_target_placement(current, target).tolist() + # Every source comes from the current physical layout, with at most one copy per rank. + candidates_by_expert = [[] for _ in range(num_logical_experts)] + for rank, row in enumerate(current_rows): + for slot, expert in enumerate(row): + candidates_by_expert[expert].append((rank, slot)) + source_load = [0] * world_size + plan = [] + for dst_rank in range(world_size): + for dst_slot, expert in enumerate(aligned_target_rows[dst_rank]): + if expert == current_rows[dst_rank][dst_slot]: + continue + src_rank, src_row = min( + candidates_by_expert[expert], + key=lambda item: ( + item[0] // node_world_size != dst_rank // node_world_size, + source_load[item[0]], + item[0], + item[1], + ), + ) + source_load[src_rank] += 1 + plan.append(TransferStep(dst_rank, dst_slot, src_rank, src_row)) + return plan + + +class NixlEPLBTransfer: + """GPU-direct UCX/NIXL migration through one complete staging layer.""" + + backend = "nixl" + staging_depth = 1 + _DEFAULT_UCX_TLS = "self,sm,cuda_ipc,cuda_copy,rc_x" + + @dataclass + class _PreparedPart: + remote_entries: Dict[int, list] + push_batch: "_PreparedCudaMemcpyBatch | None" + part_index: int + + @dataclass + class _PreparedLayer: + layer_index: int + changed_dst_slots: tuple + parts: tuple + + def _init_staging(self, weights, transfer_group, global_rank, world_size): + self._eplb_states = [weight.expert_parallel_state.eplb for weight in weights] + self.transfer_group = transfer_group + self.global_rank = global_rank + self.world_size = world_size + self.device = weights[0].w13.weight.device + self.live = [extract_eplb_expert_tensors(weight) for weight in weights] + self._validate_live_layout(weights) + _, staged_rows = get_eplb_staging_shape(weights) + self.staging = [ + [ + ( + name, + torch.empty( + get_eplb_staging_tensor_shape(tensor, staged_rows), + dtype=tensor.dtype, + device=tensor.device, + ), + ) + for name, tensor in self.live[0] + ] + for _ in range(self.staging_depth) + ] + self._release = [threading.Event() for _ in range(self.staging_depth)] + for release in self._release: + release.set() + self._error = None + self._consumed_events = [torch.cuda.Event() for _ in range(self.staging_depth)] + self._consumed_recorded = [False] * self.staging_depth + self._pending = deque() + self._pending_lock = threading.Lock() + self._thread = None + self._needs_staging_reuse_barrier = False + + def _validate_live_layout(self, weights) -> None: + reference = [(name, tuple(tensor.shape[1:]), tensor.dtype, tensor.device) for name, tensor in self.live[0]] + num_redundant_slots_per_rank = self._eplb_states[0].num_redundant_experts_per_rank + for layer_index, (state, tensors) in enumerate(zip(self._eplb_states, self.live)): + layout = [(name, tuple(tensor.shape[1:]), tensor.dtype, tensor.device) for name, tensor in tensors] + assert layout == reference, f"EPLB layer {layer_index} has incompatible expert tensor layout" + assert ( + state.num_redundant_experts_per_rank == num_redundant_slots_per_rank + ), "EPLB redundant slot count must match" + + def start(self, layer_plans: Sequence[Tuple[int, Sequence[TransferStep]]], prepared_layers) -> None: + if self._thread is not None and self._thread.is_alive(): + raise RuntimeError("EPLB transfer is already in flight") + if len(prepared_layers) != len(layer_plans): + raise ValueError("EPLB prepared layer count does not match layer plans") + # 预构造批次与描述符一起传入，避免推理线程重建。 + self._start_transfer_generation() + self._error = None + with self._pending_lock: + self._pending.clear() + + def worker() -> None: + try: + torch.cuda.set_device(self.device) + for layer_offset, prepared in enumerate(prepared_layers): + release = self._release[0] + # A buffer cannot be reused until its prior committed rows are no longer read by CUDA. + release.wait() + if self._error is not None: + return + release.clear() + if self._consumed_recorded[0]: + self._consumed_events[0].synchronize() + if layer_offset > 0 and self._needs_staging_reuse_barrier: + # All destinations must finish consuming the prior IPC staging generation + # before a source can reuse the peer buffer for this batch. + dist.barrier(group=self.transfer_group) + self._prepared_layer = prepared + self._copy_part(prepared.parts[0]) + with self._pending_lock: + self._pending.append((prepared.layer_index, 0)) + except BaseException as exc: + self._error = exc + + self._thread = threading.Thread(target=worker, name=f"eplb-{self.backend}", daemon=True) + self._thread.start() + + def pending_layers(self): + if self._error is not None: + raise RuntimeError("EPLB migration worker failed") from self._error + with self._pending_lock: + return list(self._pending) + + def commit(self, layer_index: int, buffer_index: int, post_copy=None) -> None: + with self._pending_lock: + if not self._pending or self._pending[0] != (layer_index, buffer_index): + raise RuntimeError("EPLB commit does not match the pending FIFO") + self._pending.popleft() + changed_dst_slots = self._prepared_layer.changed_dst_slots + try: + self._commit_full_parts(layer_index, buffer_index, changed_dst_slots) + except BaseException as exc: + self._error = exc + for release in self._release: + release.set() + raise + if post_copy is not None: + post_copy() + self._consumed_events[buffer_index].record(torch.cuda.current_stream()) + self._consumed_recorded[buffer_index] = True + self._release[buffer_index].set() + + def _full_phase_gate(self, error): + flag = torch.tensor([int(error is not None)], dtype=torch.int32, device="cpu") + dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=self.transfer_group) + if flag.item(): + raise RuntimeError("full EPLB migration phase failed") from error + + def _commit_full_parts(self, layer_index, buffer_index, changed_dst_slots): + """Queue complete rows before routing publication on the inference stream.""" + error = None + try: + prepared = self._prepared_layer + if prepared.layer_index != layer_index or len(prepared.parts) != 1: + raise RuntimeError("invalid full EPLB prepared layer") + # All ranks finished staging the complete layer before the manager's + # ready-MIN gate. Keep copies and routing publication on the same + # stream; the worker waits on the consumed event before reuse. + self._commit_full_part(layer_index, buffer_index, changed_dst_slots, 0) + except BaseException as exc: + error = exc + self._full_phase_gate(error) # Propagate enqueue failures before any rank publishes. + self._prepared_layer = None + + def _commit_full_part(self, layer_index, buffer_index, changed_dst_slots, part_index): + for (_, live), (_, staging) in zip(self.live[layer_index], self.staging[buffer_index]): + live_flat = live.view(live.shape[0], -1) + width = staging.shape[1] + offset = part_index * width + valid = min(width, live_flat.shape[1] - offset) + if valid > 0: + _commit_staging_rows(live_flat[:, offset : offset + valid], staging[:, :valid], changed_dst_slots) + + def finish(self) -> None: + """Wait for the released migration worker to exit before another rebalance.""" + thread = self._thread + if thread is None: + return + thread.join() + self._thread = None + if self._error is not None: + raise RuntimeError("EPLB migration worker failed") from self._error + self._finish_transfer_generation() + + def __init__(self, weights, transfer_group, global_rank, world_size): + # Stage complete incoming rows before any overwrite, including cycles. + self._init_staging(weights, transfer_group, global_rank, world_size) + self._nixl_agent = None + self._registered_descs = None + self._remote_agents: Dict[int, str] = {} + self._remote_layouts = {} + self._xfer_cache = {} + self._used_xfer_cache_keys = set() + self._ipc_staging = {} + self._prepared_layer = None + self._same_node_ranks = set() + self._cross_node_ranks = set() + self._push_stream = torch.cuda.Stream(device=self.device) + self._batch_memcpy = _CudaBatchMemcpy() + try: + self._init_ipc_metadata() + self._init_push_layouts() + except Exception as exc: + self.shutdown() + if isinstance(exc, RuntimeError): + raise + raise RuntimeError("NIXL EPLB initialization failed") from exc + + def _local_layout(self): + return [ + [(name, tensor.data_ptr(), tensor.get_device(), tensor[0].nbytes) for name, tensor in layer] + for layer in self.live + ] + + def _init_ipc_metadata(self) -> None: + hostnames = [None] * self.world_size + dist.all_gather_object(hostnames, socket.gethostname(), group=self.transfer_group) + self._rank_hostnames = hostnames + local_hostname = hostnames[self.global_rank] + self._needs_staging_reuse_barrier = len(set(hostnames)) < len(hostnames) + self._same_node_ranks = {rank for rank, hostname in enumerate(hostnames) if hostname == local_hostname} + self._cross_node_ranks = set(range(self.world_size)) - self._same_node_ranks + from lightllm.server.router.model_infer.mode_backend.pd.p2p_fix import ( + p2p_fix_rebuild_cuda_tensor, + reduce_tensor, + ) + + exports = {} + for target_rank in self._same_node_ranks - {self.global_rank}: + exports[target_rank] = { + "staging": [ + [(name, tuple(tensor.shape), tensor.dtype, reduce_tensor(tensor)[1]) for name, tensor in staging] + for staging in self.staging + ], + } + all_exports = [None] * self.world_size + dist.all_gather_object(all_exports, exports, group=self.transfer_group) + + torch.cuda.set_device(self.device) + for dst_rank in self._same_node_ranks - {self.global_rank}: + metadata = all_exports[dst_rank].get(self.global_rank) + if metadata is None or len(metadata["staging"]) != self.staging_depth: + raise RuntimeError(f"NIXL IPC destination rank {dst_rank} has incompatible staging metadata") + rebuilt_staging = [] + for remote_staging, local_staging in zip(metadata["staging"], self.staging): + if len(remote_staging) != len(local_staging): + raise RuntimeError(f"NIXL IPC destination rank {dst_rank} staging tensor count mismatch") + rebuilt = [] + for (name, shape, dtype, args), (local_name, local_tensor) in zip(remote_staging, local_staging): + if name != local_name or shape != tuple(local_tensor.shape) or dtype != local_tensor.dtype: + raise RuntimeError(f"NIXL IPC destination rank {dst_rank} staging layout mismatch for {name}") + tensor = p2p_fix_rebuild_cuda_tensor(*args) + if tuple(tensor.shape) != shape or tensor.dtype != dtype or tensor.device != local_tensor.device: + raise RuntimeError( + f"NIXL IPC destination rank {dst_rank} staging rebuild validation failed for {name}" + ) + rebuilt.append((name, tensor)) + rebuilt_staging.append(rebuilt) + self._ipc_staging[dst_rank] = rebuilt_staging + + def _init_remote_metadata(self) -> None: + metadata = self._nixl_agent.get_agent_metadata() + all_metadata = [None] * self.world_size + all_layouts = [None] * self.world_size + dist.all_gather_object(all_metadata, metadata, group=self.transfer_group) + dist.all_gather_object(all_layouts, self._local_layout(), group=self.transfer_group) + for rank in self._cross_node_ranks: + layout = all_layouts[rank] + if len(layout) != len(self.live): + raise RuntimeError(f"NIXL remote rank {rank} has incompatible layer layout") + self._remote_agents[rank] = self._nixl_agent.add_remote_agent(all_metadata[rank]) + self._remote_layouts[rank] = layout + + def _wait_xfers(self, xfers) -> None: + pending = [] + for item in xfers: + state = self._nixl_agent.transfer(item[2]) + if state == "ERR": + raise RuntimeError("NIXL READ post failed") + if state == "PROC": + pending.append(item) + while pending: + remaining = [] + for item in pending: + state = self._nixl_agent.check_xfer_state(item[2]) + if state == "ERR": + raise RuntimeError("NIXL READ transfer failed") + if state != "DONE": + remaining.append(item) + pending = remaining + + def _release_xfers(self, xfers) -> None: + unreleased = [] + errors = [] + for local_dlist, remote_dlist, xfer in xfers: + remaining = [local_dlist, remote_dlist, xfer] + for remaining_index, handle, release in ( + (2, xfer, self._nixl_agent.release_xfer_handle), + (1, remote_dlist, self._nixl_agent.release_dlist_handle), + (0, local_dlist, self._nixl_agent.release_dlist_handle), + ): + if handle is not None: + try: + release(handle) + except Exception as exc: + errors.append(exc) + else: + remaining[remaining_index] = None + if any(handle is not None for handle in remaining): + unreleased.append(tuple(remaining)) + if errors: + error = RuntimeError("NIXL transfer handle release failed") + error.unreleased_xfers = unreleased + raise error from errors[0] + + def _init_push_layouts(self) -> None: + self._live_row_layout = [ + [(name, tensor.data_ptr(), tensor[0].nbytes) for name, tensor in layer] for layer in self.live + ] + self._push_staging_row_layout = {} + for dst_rank in self._same_node_ranks: + layouts = [] + for buffer_index in range(self.staging_depth): + staging = ( + self.staging[buffer_index] + if dst_rank == self.global_rank + else self._ipc_staging[dst_rank][buffer_index] + ) + layout = [(name, tensor.data_ptr(), tensor[0].nbytes) for name, tensor in staging] + expected = [ + ( + name, + math.prod(get_eplb_staging_tensor_shape(tensor, 1)) * tensor.element_size(), + ) + for name, tensor in self.live[0] + ] + if [(name, row_nbytes) for name, _, row_nbytes in layout] != expected: + raise RuntimeError("NIXL source-push staging row layout mismatch") + layouts.append(layout) + self._push_staging_row_layout[dst_rank] = layouts + + def _prepare_part(self, layer_index, plan, part_index): + if part_index not in range(EPLB_FULL_STAGING_PARTS): + raise ValueError("full EPLB requires a valid staging part index") + staging = self.staging[0] + remote_entries = defaultdict(list) + push_descriptors = [] + steps_by_source = defaultdict(list) + by_destination = defaultdict(list) + for step in plan: + if step.dst_rank == self.global_rank and step.src_rank not in self._same_node_ranks: + steps_by_source[step.src_rank].append(step) + if step.src_rank == self.global_rank and step.dst_rank in self._same_node_ranks: + by_destination[step.dst_rank].append(step) + for src_rank, steps in steps_by_source.items(): + runs = ([step] for step in steps) + remote_entries[src_rank].extend((layer_index, run, staging) for run in runs) + source_layout = self._live_row_layout[layer_index] + for dst_rank, steps in by_destination.items(): + destination_layout = self._push_staging_row_layout[dst_rank][0] + for first in steps: + for (_, source_ptr, row_nbytes), (_, destination_ptr, staged_row_nbytes) in zip( + source_layout, destination_layout + ): + offset = part_index * staged_row_nbytes + size = min(staged_row_nbytes, row_nbytes - offset) + if size <= 0: + continue + push_descriptors.append( + ( + source_ptr + first.src_local_row * row_nbytes + offset, + destination_ptr + first.dst_slot * staged_row_nbytes, + size, + ) + ) + push_batch = self._batch_memcpy.prepare(push_descriptors) if push_descriptors else None + return self._PreparedPart(dict(remote_entries), push_batch, part_index) + + def _ensure_cross_node_transfer(self, layer_plans): + # Plans contain global steps on every rank. Make the same decision even + # on nonparticipants, since remote metadata exchange is collective. + if not any( + self._rank_hostnames[step.src_rank] != self._rank_hostnames[step.dst_rank] + for _, plan in layer_plans + for step in plan + ): + return + if self._nixl_agent is not None: + return + try: + os.environ.setdefault("UCX_TLS", self._DEFAULT_UCX_TLS) + try: + import nixl + except Exception as exc: + raise RuntimeError("NIXL EPLB backend requires the nixl package for cross-node transfer") from exc + agent_name = f"lightllm-eplb-{socket.gethostname()}-{os.getpid()}-rank-{self.global_rank}" + config = nixl.nixl_agent_config(enable_prog_thread=True, enable_listen_thread=False, backends=["UCX"]) + self._nixl_agent = nixl.nixl_agent(agent_name, config) + reg_tensors = [tensor for layer in self.live for _, tensor in layer] + [ + tensor for staging in self.staging for _, tensor in staging + ] + self._registered_descs = self._nixl_agent.get_reg_descs(reg_tensors) + self._nixl_agent.register_memory(self._registered_descs, backends=["UCX"]) + self._init_remote_metadata() + except Exception as exc: + self.shutdown() + if isinstance(exc, RuntimeError): + raise + raise RuntimeError("NIXL EPLB initialization failed") from exc + + def prepare_transfer(self, layer_plans: Sequence[Tuple[int, Sequence[TransferStep]]]): + self._ensure_cross_node_transfer(layer_plans) + return [ + self._PreparedLayer( + layer_index, + tuple(step.dst_slot for step in plan if step.dst_rank == self.global_rank), + tuple(self._prepare_part(layer_index, plan, part) for part in range(EPLB_FULL_STAGING_PARTS)), + ) + for layer_index, plan in layer_plans + ] + + def _get_remote_read(self, src_rank: int, entries, part_index: int): + if part_index not in range(EPLB_FULL_STAGING_PARTS): + raise ValueError("invalid EPLB staging part index") + local_descs = [] + remote_descs = [] + local_dlist = remote_dlist = xfer = None + try: + for layer_index, run, staging in entries: + remote_layer = self._remote_layouts[src_rank][layer_index] + if len(remote_layer) != len(staging): + raise RuntimeError(f"NIXL remote rank {src_rank} has incompatible layer layout") + if len(run) != 1: + raise RuntimeError("EPLB part requires one transfer step per run") + first = run[0] + for tensor_index, (_, staging_tensor) in enumerate(staging): + name, remote_ptr, remote_device, remote_nbytes = remote_layer[tensor_index] + live_tensor = self.live[layer_index][tensor_index][1] + if name != self.live[layer_index][tensor_index][0] or remote_nbytes != live_tensor[0].nbytes: + raise RuntimeError(f"NIXL remote rank {src_rank} descriptor range mismatch") + staging_nbytes = staging_tensor[0].nbytes + offset = part_index * staging_nbytes + size = min(staging_nbytes, remote_nbytes - offset) + if size <= 0: + continue + local_descs.append((staging_tensor[first.dst_slot].data_ptr(), size, staging_tensor.get_device())) + remote_descs.append( + (remote_ptr + first.src_local_row * remote_nbytes + offset, size, remote_device) + ) + if not local_descs: + return None + cache_key = (src_rank, part_index, tuple(local_descs), tuple(remote_descs)) + cached = self._xfer_cache.get(cache_key) + if cached is not None: + self._used_xfer_cache_keys.add(cache_key) + return cached + local_dlist = self._nixl_agent.prep_xfer_dlist( + "NIXL_INIT_AGENT", self._nixl_agent.get_xfer_descs(local_descs, "VRAM"), backends=["UCX"] + ) + remote_dlist = self._nixl_agent.prep_xfer_dlist( + self._remote_agents[src_rank], self._nixl_agent.get_xfer_descs(remote_descs, "VRAM"), backends=["UCX"] + ) + xfer = self._nixl_agent.make_prepped_xfer( + "READ", + local_dlist, + list(range(len(local_descs))), + remote_dlist, + list(range(len(remote_descs))), + backends=["UCX"], + ) + if self._nixl_agent.query_xfer_backend(xfer) != "UCX": + raise RuntimeError("NIXL EPLB READ did not select UCX") + self._xfer_cache[cache_key] = (local_dlist, remote_dlist, xfer) + self._used_xfer_cache_keys.add(cache_key) + return self._xfer_cache[cache_key] + except Exception: + self._release_xfers([(local_dlist, remote_dlist, xfer)]) + raise + + def _copy_part(self, prepared_part) -> None: + if prepared_part.push_batch is not None: + self._batch_memcpy.enqueue(prepared_part.push_batch, self._push_stream.cuda_stream) + xfers = [ + self._get_remote_read(src_rank, entries, prepared_part.part_index) + for src_rank, entries in prepared_part.remote_entries.items() + ] + xfers = [xfer for xfer in xfers if xfer is not None] + self._wait_xfers(xfers) + self._push_stream.synchronize() + # Before a rank publishes this batch it has completed its outgoing source-pushes and + # incoming UCX READs. The manager's global MIN-ready gate therefore means all transfers + # are complete before any rank commits, without a destination-side GPU wait. + + def _start_transfer_generation(self) -> None: + self._used_xfer_cache_keys.clear() + + def _finish_transfer_generation(self) -> None: + errors = [] + for cache_key in set(self._xfer_cache) - self._used_xfer_cache_keys: + xfer = self._xfer_cache[cache_key] + try: + self._release_xfers([xfer]) + except Exception as exc: + unreleased = getattr(exc, "unreleased_xfers", None) + if unreleased: + self._xfer_cache[cache_key] = unreleased[0] + errors.append(exc) + else: + del self._xfer_cache[cache_key] + if errors: + raise RuntimeError("NIXL EPLB cache eviction failed") from errors[0] + + def shutdown(self) -> None: + agent = self._nixl_agent + errors = [] + getattr(self, "_used_xfer_cache_keys", set()).clear() + if agent is not None: + for cache_key, xfer in list(self._xfer_cache.items()): + try: + self._release_xfers([xfer]) + except Exception as exc: + unreleased = getattr(exc, "unreleased_xfers", None) + if unreleased: + self._xfer_cache[cache_key] = unreleased[0] + errors.append(exc) + else: + del self._xfer_cache[cache_key] + if errors: + raise RuntimeError("NIXL EPLB shutdown failed") from errors[0] + for remote_name in list(self._remote_agents.values()): + if agent is not None: + try: + agent.remove_remote_agent(remote_name) + except Exception as exc: + errors.append(exc) + self._remote_agents.clear() + self._remote_layouts.clear() + if agent is not None and self._registered_descs is not None: + try: + agent.deregister_memory(self._registered_descs, backends=["UCX"]) + except Exception as exc: + errors.append(exc) + self._registered_descs = None + self._nixl_agent = None + getattr(self, "_ipc_staging", {}).clear() + if errors: + raise RuntimeError("NIXL EPLB shutdown failed") from errors[0] + + def __del__(self): + try: + self.shutdown() + except Exception: + pass + + +def _commit_staging_rows( + live: torch.Tensor, + staging: torch.Tensor, + changed_dst_slots: Sequence[int], +) -> None: + slots = sorted(set(changed_dst_slots)) + if not slots: + return + run_start = previous = slots[0] + for dst_slot in (*slots[1:], None): + if dst_slot is not None and dst_slot == previous + 1: + previous = dst_slot + continue + run_length = previous - run_start + 1 + live.narrow(0, run_start, run_length).copy_(staging.narrow(0, run_start, run_length), non_blocking=True) + if dst_slot is not None: + run_start = previous = dst_slot + + +class _CudaMemLocation(ctypes.Structure): + _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)] + + +class _CudaMemcpyAttributes(ctypes.Structure): + _fields_ = [ + ("srcAccessOrder", ctypes.c_int), + ("srcLocHint", _CudaMemLocation), + ("dstLocHint", _CudaMemLocation), + ("flags", ctypes.c_uint), + ] + + +@dataclass +class _PreparedCudaMemcpyBatch: + """Host-side arrays retained for one cudaMemcpyBatchAsync submission.""" + + dsts: object + srcs: object + sizes: object + attrs: _CudaMemcpyAttributes + attrs_idxs: object + count: int + + +class _CudaBatchMemcpy: + """CUDA 13.x ``cudaMemcpyBatchAsync`` binding for EPLB source-push.""" + + _SRC_ACCESS_ORDER_STREAM = 1 + _PREFER_OVERLAP_WITH_COMPUTE = 1 + _CUDA_13_0 = 13000 + _CUDA_14_0 = 14000 + + def __init__(self, library=None): + if library is None: + path = self._find_loaded_cudart() + if path is None: + raise RuntimeError( + "NIXL same-node source-push requires CUDA Runtime 13.x cudaMemcpyBatchAsync; " + "libcudart.so.13 is not loaded" + ) + try: + library = ctypes.CDLL(path) + except OSError as exc: + raise RuntimeError(f"cannot load libcudart: {exc}") from exc + + try: + runtime_get_version = library.cudaRuntimeGetVersion + self._batch_async = library.cudaMemcpyBatchAsync + self._get_error_string = library.cudaGetErrorString + except AttributeError as exc: + raise RuntimeError("cudaMemcpyBatchAsync is unavailable") from exc + + runtime_get_version.restype = ctypes.c_int + runtime_get_version.argtypes = [ctypes.POINTER(ctypes.c_int)] + self._get_error_string.restype = ctypes.c_char_p + self._get_error_string.argtypes = [ctypes.c_int] + runtime_version = ctypes.c_int() + result = runtime_get_version(ctypes.byref(runtime_version)) + if result != 0: + raise RuntimeError(f"cudaRuntimeGetVersion failed with CUDA error {result}") + if not self._CUDA_13_0 <= runtime_version.value < self._CUDA_14_0: + raise RuntimeError( + f"cudaMemcpyBatchAsync requires CUDA Runtime 13.x (13.0 ABI), found {runtime_version.value}" + ) + + pointer_array = ctypes.POINTER(ctypes.c_void_p) + self._batch_async.restype = ctypes.c_int + self._batch_async.argtypes = [ + pointer_array, + pointer_array, + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ctypes.POINTER(_CudaMemcpyAttributes), + ctypes.POINTER(ctypes.c_size_t), + ctypes.c_size_t, + ctypes.c_void_p, + ] + + @staticmethod + def prepare(copies: Iterable[Tuple[int, int, int]]) -> _PreparedCudaMemcpyBatch: + copies = tuple(copies) + if not copies: + raise ValueError("cudaMemcpyBatchAsync requires at least one copy") + for src, dst, size in copies: + if not src or not dst or size <= 0: + raise ValueError("cudaMemcpyBatchAsync requires non-null pointers and positive sizes") + count = len(copies) + dsts = (ctypes.c_void_p * count)(*(dst for _, dst, _ in copies)) + srcs = (ctypes.c_void_p * count)(*(src for src, _, _ in copies)) + sizes = (ctypes.c_size_t * count)(*(size for _, _, size in copies)) + attrs = _CudaMemcpyAttributes() + attrs.srcAccessOrder = _CudaBatchMemcpy._SRC_ACCESS_ORDER_STREAM + attrs.flags = _CudaBatchMemcpy._PREFER_OVERLAP_WITH_COMPUTE + attrs_idxs = (ctypes.c_size_t * 1)(0) + return _PreparedCudaMemcpyBatch(dsts, srcs, sizes, attrs, attrs_idxs, count) + + def enqueue(self, prepared: _PreparedCudaMemcpyBatch, stream: int) -> None: + result = self._batch_async( + prepared.dsts, + prepared.srcs, + prepared.sizes, + prepared.count, + ctypes.byref(prepared.attrs), + prepared.attrs_idxs, + 1, + ctypes.c_void_p(stream), + ) + if result != 0: + message = self._get_error_string(result) + error = message.decode("utf-8") if message else f"CUDA error {result}" + raise RuntimeError(f"cudaMemcpyBatchAsync failed: {error}") + + @staticmethod + def _find_loaded_cudart() -> Optional[str]: + """Return a mapped CUDA 13 runtime without loading CUDA as a side effect.""" + try: + with open("/proc/self/maps") as maps: + for line in maps: + if "libcudart" not in line: + continue + path_start = line.find("/") + if path_start < 0: + continue + path = line[path_start:].strip().removesuffix(" (deleted)") + if re.search(r"libcudart[^/]*\.so\.13(?:\D|$)", path): + return path + except OSError: + pass + return None diff --git a/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py b/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py deleted file mode 100644 index 596eca4f24..0000000000 --- a/lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py +++ /dev/null @@ -1,158 +0,0 @@ -# 对于 deepseekv3 模型在 ep 运行模式下，自动分析统计各个专家的出现频率，然后 -# 自动更新当前的冗余专家为新的冗余专家。 -import torch -import time -import enum -import lightllm.utils.petrel_helper as utils -import threading -import json -from typing import List -from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_redundancy import ( - FusedMoeWeightEPAutoRedundancy, -) -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight -from lightllm.utils.envs_utils import get_env_start_args, get_redundancy_expert_update_interval -from lightllm.utils.envs_utils import get_redundancy_expert_update_max_load_count -from lightllm.utils.envs_utils import get_redundancy_expert_num -from lightllm.utils.dist_utils import get_global_rank -from lightllm.common.basemodel.layer_weights.hf_load_utils import load_func -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -class RedundancyExpertManager: - def __init__(self, model: TpPartBaseModel): - self.args = get_env_start_args() - self.model = model - self.ep_fused_moeweights: List[FusedMoeWeightEPAutoRedundancy] = [] - for layer in self.model.trans_layers_weight: - ep_weights = self._find_members_of_class(layer, FusedMoeWeight) - assert len(ep_weights) <= 1 - self.ep_fused_moeweights.extend([FusedMoeWeightEPAutoRedundancy(e) for e in ep_weights]) - - # save load params - self.use_safetensors = True - files = utils.PetrelHelper.list(self.args.model_dir, extension="all") - candidate_files = list(filter(lambda x: x.endswith(".safetensors"), files)) - if len(candidate_files) == 0: - self.use_safetensors = False - candidate_files = list(filter(lambda x: x.endswith(".bin"), files)) - assert len(candidate_files) != 0, "can only support pytorch tensor and safetensors format for weights." - self.candidate_files = candidate_files - - # state 1. check_to_update 2. prepare_update 3. start_load_hf_weights 4. wait_load_ready, 5. commit - self.state: _STATE = _STATE.CHECK_TO_UPDATE - self.update_time = time.time() - self.update_interval = get_redundancy_expert_update_interval() - self.load_thread: threading.Thread = None - self.global_rank = get_global_rank() - # 冗余专家的最大加载次数 - self.load_count = 0 - self.max_load_count = get_redundancy_expert_update_max_load_count() - - # 清理counter - self._clear_all_counter() - - self.rank0_redundancy_expert_config = { - "redundancy_expert_num": get_redundancy_expert_num(), - "default": list(range(get_redundancy_expert_num())), - } - - def step(self): - if self.load_count >= self.max_load_count: - return - - if self.state == _STATE.CHECK_TO_UPDATE: - cur_time = time.time() - if cur_time - self.update_time > self.update_interval: - self.update_time = cur_time - self.state = _STATE.PREPARE_UPDATE - logger.info(f"global_rank {self.global_rank} state to prepare update") - elif self.state == _STATE.PREPARE_UPDATE: - self._prepare_load_new_redundancy_expert() - self.state = _STATE.START_LOAD_HF_WEIGHTS - logger.info(f"global_rank {self.global_rank} state to start load hf weights") - - elif self.state == _STATE.START_LOAD_HF_WEIGHTS: - self.load_thread = threading.Thread(target=self._load_hf_weights, daemon=True) - self.load_thread.start() - self.state = _STATE.WAIT_LOAD_READY - logger.info(f"global_rank {self.global_rank} state to wait load ready") - - elif self.state == _STATE.WAIT_LOAD_READY: - if not self.load_thread.is_alive(): - self.load_thread = None - self.state = _STATE.COMMIT - logger.info(f"global_rank {self.global_rank} state to commit") - - elif self.state == _STATE.COMMIT: - self._commit() - self.state = _STATE.CHECK_TO_UPDATE - self.load_count += 1 - logger.info(f"global_rank {self.global_rank} state to check to update") - return - - def _prepare_load_new_redundancy_expert(self): - for w in self.ep_fused_moeweights: - topk_redundancy_expert_ids = w.prepare_redundancy_experts() - if self.global_rank == 0: - self.rank0_redundancy_expert_config[str(w._ep_w.layer_num)] = topk_redundancy_expert_ids - - if self.global_rank == 0: - try: - with open("./redundancy_expert_config.json", "w") as f: - json.dump(self.rank0_redundancy_expert_config, f, indent=4) - logger.info( - f"rank {self.global_rank} save redundancy_expert_config.json to ./redundancy_expert_config.json" - ) - except BaseException as e: - logger.exception(str(e)) - logger.error(f"global rank {self.global_rank} save redundancy_expert_config.json failed") - - return - - def _load_hf_weights(self): - start = time.time() - try: - for file in self.candidate_files: - load_func( - file, - use_safetensors=self.use_safetensors, - pre_post_layer=None, - transformer_layer_list=self.ep_fused_moeweights, - weight_dir=self.args.model_dir, - ) - except BaseException as e: - logger.exception(str(e)) - raise e - cost_time = time.time() - start - logger.info(f"global rank {self.global_rank} load redundancy_expert cost time: {cost_time} s") - return - - def _commit(self): - for w in self.ep_fused_moeweights: - w.commit() - return - - def _find_members_of_class(self, obj, cls): - members = [] - for attr in dir(obj): - value = getattr(obj, attr) - if isinstance(value, cls): - members.append(value) - return members - - def _clear_all_counter(self): - for w in self.ep_fused_moeweights: - w.clear_counter() - return - - -class _STATE(enum.Enum): - CHECK_TO_UPDATE = 0 - PREPARE_UPDATE = 1 - START_LOAD_HF_WEIGHTS = 2 - WAIT_LOAD_READY = 3 - COMMIT = 4 diff --git a/lightllm/server/router/model_infer/model_rpc.py b/lightllm/server/router/model_infer/model_rpc.py index c66ea0ec5b..d8cf2bf022 100644 --- a/lightllm/server/router/model_infer/model_rpc.py +++ b/lightllm/server/router/model_infer/model_rpc.py @@ -23,7 +23,6 @@ PDDecodeNode, PDDPForDecodeNode, ) -from lightllm.server.router.model_infer.mode_backend.redundancy_expert_manager import RedundancyExpertManager from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.log_utils import init_logger @@ -86,12 +85,6 @@ def exposed_init_model(self, kvargs): self.backend.init_model(kvargs) self.rl_backend_ops = RlBackendOps(self.backend) if self.args.enable_rl else None - # only deepseekv3 can support auto_update_redundancy_expert - if self.args.auto_update_redundancy_expert: - self.redundancy_expert_manager = RedundancyExpertManager(self.backend.model) - logger.info("init redundancy_expert_manager") - else: - self.redundancy_expert_manager = None return def exposed_get_max_total_token_num(self): diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 454abdbea2..eb44a8c449 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -107,64 +107,63 @@ def get_lightllm_websocket_max_message_size(): return int(os.getenv("LIGHTLLM_WEBSOCKET_MAX_SIZE", 128 * 1024 * 1024)) -# get_redundancy_expert_ids and get_redundancy_expert_num are primarily -# used to obtain the IDs and number of redundant experts during inference. -# They depend on a configuration file specified by ep_redundancy_expert_config_path, -# which is a JSON formatted text file. -# The content format is as follows: -# { -# "redundancy_expert_num": 1, # Number of redundant experts per rank -# "0": [0], # Key: layer_index (string), -# # Value: list of original expert IDs that are redundant for this layer -# "1": [0], -# "default": [0] # Default list of redundant expert IDs if layer-specific entry is not found -# } - - -@lru_cache(maxsize=None) -def get_redundancy_expert_ids(layer_index: int): - """ - Get the redundancy expert ids from the environment variable. - :return: List of redundancy expert ids. - """ - args = get_env_start_args() - if args.ep_redundancy_expert_config_path is None: - return [] +@lru_cache(maxsize=None) +def get_eplb_planner_process() -> bool: + raw = os.getenv("LIGHTLLM_EPLB_PLANNER_PROCESS", "0") + if raw not in ("0", "1"): + raise ValueError("LIGHTLLM_EPLB_PLANNER_PROCESS must be 0 or 1") + return raw == "1" - with open(args.ep_redundancy_expert_config_path, "r") as f: - config = json.load(f) - if str(layer_index) in config: - return config[str(layer_index)] - else: - return config.get("default", []) + +@lru_cache(maxsize=None) +def get_prefill_eplb_step_interval(): + """Return the number of prefill forwards between EPLB attempts.""" + interval = int(os.getenv("LIGHTLLM_PREFILL_EPLB_STEP_INTERVAL", 20)) + if interval <= 0: + raise ValueError("LIGHTLLM_PREFILL_EPLB_STEP_INTERVAL must be greater than 0") + return interval @lru_cache(maxsize=None) -def get_redundancy_expert_num(): - """ - Get the number of redundancy experts from the environment variable. - :return: Number of redundancy experts. - """ - args = get_env_start_args() - if args.ep_redundancy_expert_config_path is None: - return 0 +def get_prefill_eplb_steady_sample_steps() -> int: + value = int(os.getenv("LIGHTLLM_PREFILL_EPLB_STEADY_SAMPLE_STEPS", "4")) + if value <= 0: + raise ValueError("LIGHTLLM_PREFILL_EPLB_STEADY_SAMPLE_STEPS must be greater than 0") + return value - with open(args.ep_redundancy_expert_config_path, "r") as f: - config = json.load(f) - if "redundancy_expert_num" in config: - return config["redundancy_expert_num"] - else: - return 0 + +@lru_cache(maxsize=None) +def get_prefill_eplb_min_rebalance_interval() -> int: + """Return the fixed-boundary steady sampling period, not a strict commit-to-commit distance.""" + value = int(os.getenv("LIGHTLLM_PREFILL_EPLB_MIN_REBALANCE_INTERVAL", "0")) + if value < 0: + raise ValueError("LIGHTLLM_PREFILL_EPLB_MIN_REBALANCE_INTERVAL must be non-negative") + step_interval = get_prefill_eplb_step_interval() + if value and value < step_interval: + raise ValueError("LIGHTLLM_PREFILL_EPLB_MIN_REBALANCE_INTERVAL must be 0 or at least step interval") + return value @lru_cache(maxsize=None) -def get_redundancy_expert_update_interval(): - return int(os.getenv("LIGHTLLM_REDUNDANCY_EXPERT_UPDATE_INTERVAL", 30 * 60)) +def get_eplb_rebalance_gain_threshold() -> float: + """Return the EPLB gain threshold: estimated critical-load reduction ratio; 0.05 means 5%.""" + env_name = "LIGHTLLM_EPLB_REBALANCE_GAIN_THRESHOLD" + raw_value = os.getenv(env_name, "0.05") + value = float(raw_value) + if not 0.0 <= value <= 1.0: + raise ValueError(f"{env_name} must be a ratio between 0.0 and 1.0, got {raw_value!r}") + return value @lru_cache(maxsize=None) -def get_redundancy_expert_update_max_load_count(): - return int(os.getenv("LIGHTLLM_REDUNDANCY_EXPERT_UPDATE_MAX_LOAD_COUNT", 1)) +def get_eplb_placement_stickiness() -> float: + """Return the EPLB placement stickiness: a keep-bonus, as a fraction of the mean per-layer expert load.""" + env_name = "LIGHTLLM_EPLB_PLACEMENT_STICKINESS" + raw_value = os.getenv(env_name, "0.1") + value = float(raw_value) + if not 0.0 <= value <= 1.0: + raise ValueError(f"{env_name} must be a ratio between 0.0 and 1.0, got {raw_value!r}") + return value @lru_cache(maxsize=None) diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index 5ec2a9145e..cb60377b5e 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -6,6 +6,7 @@ from transformers import AutoModelForCausalLM import argparse from lightllm.common.build_utils import repair_config +from lightllm.common.eplb_utils import is_eplb_enabled from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_mtp_weight_layer_num @@ -61,25 +62,53 @@ def get_mtp_adjusted_mem_fraction( return adjusted_mem_fraction -@contextmanager -def profile_mtp_weight_memory(model): - """Reserve auto-profiled KV budget for MTP weights allocated after the target model.""" - should_adjust = ( - model.max_total_token_num is None and not model.is_mtp_draft_model and model.args.mtp_mode is not None +def get_eplb_adjusted_mem_fraction(model) -> float: + """Reserve auto-profiled KV budget for EPLB buffers allocated later.""" + if model.max_total_token_num is not None or model.is_mtp_draft_model or not is_eplb_enabled(model.args): + return model.mem_fraction + from lightllm.common.eplb_utils import ( + find_ep_moe_weights, + get_eplb_staging_nbytes, + get_eplb_sampling_peak_nbytes, ) - if not should_adjust: + + weights = [weight for weight in find_ep_moe_weights(model) if weight.expert_parallel_state.eplb is not None] + reserved_bytes = get_eplb_staging_nbytes(weights) + get_eplb_sampling_peak_nbytes(weights) + if not reserved_bytes: + return model.mem_fraction + device_id = get_current_device_id() + total_gpu_bytes = torch.cuda.get_device_properties(device_id).total_memory + adjusted_mem_fraction = model.mem_fraction - reserved_bytes / total_gpu_bytes + if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1: + adjusted_mem_fraction_tensor = torch.tensor( + adjusted_mem_fraction, dtype=torch.float32, device=f"cuda:{device_id}" + ) + torch.distributed.all_reduce(adjusted_mem_fraction_tensor, op=torch.distributed.ReduceOp.MIN) + adjusted_mem_fraction = adjusted_mem_fraction_tensor.item() + + assert adjusted_mem_fraction > 0, "EPLB buffers leave no memory for the target model KV cache" + return adjusted_mem_fraction + + +@contextmanager +def profile_model_memory(model): + """Profile weights and reserve KV budget for later MTP/EPLB allocations.""" + if model.max_total_token_num is not None or model.is_mtp_draft_model: yield return - weight_memory_before = torch.cuda.memory_allocated() + should_adjust_mtp = model.args.mtp_mode is not None + weight_memory_before = torch.cuda.memory_allocated() if should_adjust_mtp else 0 yield - target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before - model.mem_fraction = get_mtp_adjusted_mem_fraction( - mem_fraction=model.mem_fraction, - target_weight_bytes=target_weight_bytes, - target_layer_num=model.config["n_layer"], - mtp_layer_num=get_mtp_weight_layer_num(), - ) + if should_adjust_mtp: + target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before + model.mem_fraction = get_mtp_adjusted_mem_fraction( + mem_fraction=model.mem_fraction, + target_weight_bytes=target_weight_bytes, + target_layer_num=model.config["n_layer"], + mtp_layer_num=get_mtp_weight_layer_num(), + ) + model.mem_fraction = get_eplb_adjusted_mem_fraction(model) def load_config(weight_dir_): diff --git a/test/advanced_config/redundancy_expert/test_redundancy_expert_config.json b/test/advanced_config/redundancy_expert/test_redundancy_expert_config.json deleted file mode 100644 index 241ab25ea3..0000000000 --- a/test/advanced_config/redundancy_expert/test_redundancy_expert_config.json +++ /dev/null @@ -1,180 +0,0 @@ -{ - "redundancy_expert_num": 1, - "default": [ - 0 - ], - "3": [ - 226 - ], - "4": [ - 123 - ], - "5": [ - 187 - ], - "6": [ - 138 - ], - "7": [ - 132 - ], - "8": [ - 240 - ], - "9": [ - 4 - ], - "10": [ - 88 - ], - "11": [ - 60 - ], - "12": [ - 161 - ], - "13": [ - 178 - ], - "14": [ - 80 - ], - "15": [ - 144 - ], - "16": [ - 195 - ], - "17": [ - 251 - ], - "18": [ - 226 - ], - "19": [ - 87 - ], - "20": [ - 149 - ], - "21": [ - 45 - ], - "22": [ - 214 - ], - "23": [ - 41 - ], - "24": [ - 46 - ], - "25": [ - 156 - ], - "26": [ - 112 - ], - "27": [ - 185 - ], - "28": [ - 58 - ], - "29": [ - 156 - ], - "30": [ - 147 - ], - "31": [ - 199 - ], - "32": [ - 16 - ], - "33": [ - 188 - ], - "34": [ - 227 - ], - "35": [ - 136 - ], - "36": [ - 84 - ], - "37": [ - 15 - ], - "38": [ - 204 - ], - "39": [ - 96 - ], - "40": [ - 226 - ], - "41": [ - 25 - ], - "42": [ - 69 - ], - "43": [ - 122 - ], - "44": [ - 152 - ], - "45": [ - 113 - ], - "46": [ - 98 - ], - "47": [ - 68 - ], - "48": [ - 13 - ], - "49": [ - 102 - ], - "50": [ - 214 - ], - "51": [ - 201 - ], - "52": [ - 182 - ], - "53": [ - 235 - ], - "54": [ - 162 - ], - "55": [ - 125 - ], - "56": [ - 62 - ], - "57": [ - 121 - ], - "58": [ - 105 - ], - "59": [ - 236 - ], - "60": [ - 117 - ] -} \ No newline at end of file diff --git a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py b/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py deleted file mode 100644 index 16131ef935..0000000000 --- a/unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py +++ /dev/null @@ -1,151 +0,0 @@ -import torch -import pytest -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair -from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import expert_id_counter -from lightllm.utils.log_utils import init_logger - -logger = init_logger(__name__) - - -def test_redundancy_topk_ids_repair(): - ep_expert_num = 4 - global_rank = 0 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 0, - ], - dtype=torch.int64, - device="cuda", - ) - - expert_id_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - expert_counter=expert_id_counter, - enable_counter=True, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - assert torch.equal( - expert_id_counter, torch.tensor([1, 2, 1, 2, 0, 1, 0, 2, 0, 1, 1, 1], dtype=torch.int64, device="cuda") - ) - - ep_expert_num = 4 - global_rank = 1 - redundancy_expert_num = 1 - topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - - redundancy_expert_ids = torch.tensor( - [ - 5, - ], - dtype=torch.int64, - device="cuda", - ) - redundancy_topk_ids_repair( - topk_ids=topk_ids, - redundancy_expert_ids=redundancy_expert_ids, - ep_expert_num=ep_expert_num, - global_rank=global_rank, - ) - - ans_topk_ids = torch.tensor( - [ - [0, 1, 2, 3], - [7, 9, 10, 11], - [1, 3, 5, 7], - ], - dtype=torch.int64, - device="cuda", - ) - ans_topk_ids = (ans_topk_ids // ep_expert_num) * redundancy_expert_num + ans_topk_ids - new_redundancy_expert_ids = (redundancy_expert_ids // ep_expert_num) * redundancy_expert_num + redundancy_expert_ids - ans_topk_ids[ans_topk_ids == new_redundancy_expert_ids[0]] = ( - (ep_expert_num + redundancy_expert_num) * global_rank + ep_expert_num + 0 - ) - - assert torch.equal(topk_ids, ans_topk_ids) - - -def test_expert_id_counter(): - token_num = 256 - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - - ans_expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - ids, counts = torch.unique(tok_ids.view(-1), return_counts=True) - ans_expert_counter[ids] = counts - - assert torch.equal(expert_counter, ans_expert_counter) - - # test speed - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - for _ in range(100): - tok_ids = torch.randint( - low=0, - high=12, - size=(token_num, 8), - dtype=torch.int64, - device="cuda", - ) - expert_counter = torch.zeros(12, dtype=torch.int64, device="cuda") - expert_id_counter(topk_ids=tok_ids, expert_counter=expert_counter) - graph.replay() - - start_event = torch.cuda.Event(enable_timing=True) - start_event.record() - graph.replay() - end_event = torch.cuda.Event(enable_timing=True) - end_event.record() - torch.cuda.synchronize() - logger.info(f"expert_id_counter time cost: {start_event.elapsed_time(end_event)} ms") - - -if __name__ == "__main__": - pytest.main() diff --git a/unit_tests/common/fused_moe/test_activation_config.py b/unit_tests/common/fused_moe/test_activation_config.py index 311b1fc652..4649a92821 100644 --- a/unit_tests/common/fused_moe/test_activation_config.py +++ b/unit_tests/common/fused_moe/test_activation_config.py @@ -103,10 +103,6 @@ def test_marlin_rejects_clamp_at_call(monkeypatch): num_fused_shared_experts=0, routed_scaling_factor=1.0, quant_method=None, - redundancy_expert_num=0, - redundancy_expert_ids_tensor=None, - routed_expert_counter_tensor=None, - auto_update_redundancy_expert=False, ) with pytest.raises(NotImplementedError, match="does not support clamped SwiGLU"): impl( diff --git a/unit_tests/common/fused_moe/test_eplb.py b/unit_tests/common/fused_moe/test_eplb.py new file mode 100644 index 0000000000..39a37891cb --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb.py @@ -0,0 +1,3279 @@ +import builtins +import io +import threading +import time +from collections import deque +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + + +def _replica_index_reference(token, expert, counts): + value = token ^ (((expert.to(torch.int64) + 1) * 0x9E3779B9) & 0xFFFFFFFF) + value = ((value ^ (value >> 16)) * 0x7FEB352D) & 0xFFFFFFFF + value = ((value ^ (value >> 15)) * 0x846CA68B) & 0xFFFFFFFF + return (value ^ (value >> 16)) % counts + + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import ( + build_initial_redundant_expert_ids, + expand_redundant_placement, + build_logical_to_physical_map, + build_logical_to_physical_maps_for_layers, + _estimate_rank_load, + select_improving_placements, +) +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.server.router.model_infer.infer_batch import g_infer_context +from lightllm.server.router.model_infer.mode_backend import ( + eplb_manager as manager_module, +) +from lightllm.server.router.model_infer.mode_backend import ( + eplb_transfer as transfer_module, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import ( + deepgemm_impl as deepgemm_module, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( + EPLBState, + ExpertParallelState, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import ( + create_fuse_moe_impl, + FuseMoeMarlin, + FuseMoeTriton, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.base_impl import ( + FuseMoeBaseImpl, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import ( + fused_moe_weight as fused_weight_module, +) +from lightllm.utils import device_utils, envs_utils, dist_utils +from lightllm.common.eplb_utils import extract_eplb_expert_tensors +from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_topk +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import eplb_planner as eplb_planner_module +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import eplb_placement as placement_module +from lightllm.server.router.model_infer.mode_backend.eplb_transfer import ( + TransferStep, + _CudaBatchMemcpy, + _commit_staging_rows, + align_target_placement, + build_transfer_plan, +) + + +def _test_parallel_state( + *, + eplb=False, + num_logical_experts=128, + world_size=16, + num_redundant_experts_per_rank=1, + route_counter=None, + recording=False, + recorded_sample_count=0, +): + eplb_state = None + if eplb: + if route_counter is None: + route_counter = torch.zeros((2, num_logical_experts), dtype=torch.int64) + initial_layout_world_size = max(world_size, 2) + eplb_state = EPLBState( + num_redundant_experts_per_rank=num_redundant_experts_per_rank, + logical_to_physical_map=torch.zeros((num_logical_experts, 2), dtype=torch.int32), + logical_replica_count=torch.ones(num_logical_experts, dtype=torch.int32), + route_counter=route_counter, + physical_to_logical=expand_redundant_placement( + build_initial_redundant_expert_ids( + num_logical_experts, initial_layout_world_size, num_redundant_experts_per_rank + ), + num_logical_experts, + ), + recording=recording, + recorded_sample_count=recorded_sample_count, + ) + return ExpertParallelState( + num_logical_experts=num_logical_experts, + world_size=world_size, + eplb=eplb_state, + ) + + +def _validated_expert_parallel_state( + *, + eplb=True, + n_routed_experts=4, + world_size=2, + num_redundant_experts_per_rank=1, + device="cpu", +): + runtime = None + if eplb: + runtime = EPLBState( + num_redundant_experts_per_rank=num_redundant_experts_per_rank, + logical_to_physical_map=torch.empty((n_routed_experts, 2), dtype=torch.int32, device=device), + logical_replica_count=torch.ones(n_routed_experts, dtype=torch.int32, device=device), + route_counter=torch.zeros((2, n_routed_experts), dtype=torch.int64, device=device), + physical_to_logical=expand_redundant_placement( + build_initial_redundant_expert_ids(n_routed_experts, world_size, num_redundant_experts_per_rank), + n_routed_experts, + ), + ) + return ExpertParallelState( + num_logical_experts=n_routed_experts, + world_size=world_size, + eplb=runtime, + ) + + +def _set_expert_parallel_state(impl, state): + impl.expert_parallel_state = state + impl.eplb = state.eplb + + +def _manual_runtime_rank_load(source_load, placement, node_world_size, alignment): + """Reference the committed runtime logical-to-physical maps on CPU.""" + samples, layers, nodes, num_logical_experts = source_load.shape + ranks, redundant = placement.shape[1:] + num_physical_experts_per_rank = redundant + raw = torch.zeros((samples, layers, ranks, num_logical_experts), dtype=torch.float64) + for layer in range(layers): + for source_node in range(nodes): + logical_to_physical, replica_count = build_logical_to_physical_map( + placement[layer], + num_logical_experts, + source_rank=source_node * node_world_size, + node_world_size=node_world_size, + ) + for expert in range(num_logical_experts): + count = int(replica_count[expert].item()) + for physical_id in logical_to_physical[expert, :count].tolist(): + rank = physical_id // num_physical_experts_per_rank + raw[:, layer, rank, expert] += source_load[:, layer, source_node, expert] / count + return (torch.ceil(raw / alignment) * alignment).sum(dim=3) + + +@pytest.mark.parametrize("alpha,limit,clamp_up_add_one", [(None, None, True), (1.0, 10.0, False), (1.702, 7.0, True)]) +def test_base_call_template_forwards_selection_and_capture_callback(alpha, limit, clamp_up_add_one): + class Impl(FuseMoeBaseImpl): + def _select_experts( + self, + input_tensor, + router_logits, + correction_bias, + top_k, + renormalize, + use_grouped_topk, + topk_group, + num_expert_group, + scoring_func, + per_expert_scale=None, + shared_expert_gate=None, + is_prefill=None, + preserve_logical_ids=False, + ): + seen["select"] = {"preserve_logical_ids": preserve_logical_ids} + return "weights", "physical_ids", "logical_ids" + + def _fused_experts( + self, + input_tensor, + w13, + w2, + topk_weights, + topk_ids, + router_logits=None, + is_prefill=None, + alpha=None, + limit=None, + clamp_up_add_one=True, + dist_group=None, + ): + seen["fused"] = { + "topk_ids": topk_ids, + "activation": (alpha, limit, clamp_up_add_one), + "dist_group": dist_group, + } + return "output" + + seen, captured = {}, [] + dist_group = object() + impl = Impl(4, 0, 1.0, SimpleNamespace()) + result = impl( + "input", + "logits", + "w13", + "w2", + None, + "softmax", + 2, + False, + False, + 0, + 0, + moe_capture_callback=captured.append, + alpha=alpha, + limit=limit, + clamp_up_add_one=clamp_up_add_one, + dist_group=dist_group, + ) + assert result == "output" + assert captured == ["logical_ids"] + assert seen["select"]["preserve_logical_ids"] + assert seen["fused"]["topk_ids"] == "physical_ids" + assert seen["fused"]["activation"] == (alpha, limit, clamp_up_add_one) + assert seen["fused"]["dist_group"] is dist_group + + +def test_parallel_state_derives_expert_layout(): + state = _validated_expert_parallel_state(eplb=True) + assert state.num_primary_experts_per_rank == 2 + assert state.num_total_physical_experts == 6 + + +def test_factory_selects_all_paths_and_requires_ep_state(): + plain_quant = SimpleNamespace(method_name="none") + marlin_quant = SimpleNamespace(method_name="awq_marlin") + state = _validated_expert_parallel_state(eplb=False) + ep_impl = create_fuse_moe_impl( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=plain_quant, + expert_parallel_state=state, + ) + assert isinstance(ep_impl, deepgemm_module.FuseMoeDeepGEMM) + assert ep_impl.expert_parallel_state is state + assert state.eplb is None + assert isinstance( + create_fuse_moe_impl( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=plain_quant, + ), + FuseMoeTriton, + ) + assert isinstance( + create_fuse_moe_impl( + n_routed_experts=4, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=marlin_quant, + ), + FuseMoeMarlin, + ) + + +def test_find_fused_moe_weights_discovers_direct_layer_attributes(monkeypatch): + class FakeFusedMoeWeight: + def __init__(self, layer_num, enable_ep_moe=True): + self.layer_num_ = layer_num + self.enable_ep_moe = enable_ep_moe + + monkeypatch.setattr(fused_weight_module, "FusedMoeWeight", FakeFusedMoeWeight) + first = FakeFusedMoeWeight(3) + alternate = FakeFusedMoeWeight(1) + aliased = FakeFusedMoeWeight(2) + disabled = FakeFusedMoeWeight(0, enable_ep_moe=False) + model = SimpleNamespace( + trans_layers_weight=[ + SimpleNamespace(moe_weight=first), + SimpleNamespace(alternate_moe_weight=alternate), + SimpleNamespace(moe_weight=aliased, alternate_moe_weight=aliased), + SimpleNamespace(moe_weight=disabled), + ] + ) + + assert manager_module.find_ep_moe_weights(model) == [alternate, aliased, first] + + +def test_eplb_redundant_experts_defaults_per_ep_rank(): + parser = make_argument_parser() + + assert parser.parse_args([]).eplb_num_redundant_experts_per_rank == 2 + assert parser.parse_args(["--eplb_num_redundant_experts_per_rank", "3"]).eplb_num_redundant_experts_per_rank == 3 + assert StartArgs().eplb_num_redundant_experts_per_rank == 2 + assert parser.parse_args(["--eplb_num_redundant_experts_per_rank", "0"]).eplb_num_redundant_experts_per_rank == 0 + assert not hasattr(StartArgs(), "eplb_placement_mode") + with pytest.raises(SystemExit): + parser.parse_args(["--eplb_placement_mode", "full"]) + + +@pytest.mark.parametrize( + ("num_logical_experts", "num_ranks", "num_redundant_experts_per_rank", "expected"), + [ + (8, 4, 2, [[2, 3], [4, 5], [6, 7], [0, 1]]), + (6, 3, 4, [[2, 3, 4, 5], [4, 5, 0, 1], [0, 1, 2, 3]]), + ], +) +def test_build_initial_redundant_expert_ids( + num_logical_experts, + num_ranks, + num_redundant_experts_per_rank, + expected, +): + actual = build_initial_redundant_expert_ids( + num_logical_experts, + num_ranks, + num_redundant_experts_per_rank, + ) + + assert actual.dtype == torch.int64 + assert actual.shape == (num_ranks, num_redundant_experts_per_rank) + assert torch.equal(actual, torch.tensor(expected, dtype=torch.int64)) + + +def test_select_improving_placements_rejects_regressing_layer(): + expert_load = torch.tensor([[8649, 5740, 5002, 3441]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (expert_load).shape[-1]) + regressing_candidate = expand_redundant_placement(torch.tensor([[[1], [0]]]), (expert_load).shape[-1]) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, current, regressing_candidate, rebalance_gain_threshold=0.05 + ) + + current_ratio = _estimate_rank_load(expert_load, current).max() / _estimate_rank_load(expert_load, current).mean() + candidate_ratio = ( + _estimate_rank_load(expert_load, regressing_candidate).max() + / _estimate_rank_load(expert_load, regressing_candidate).mean() + ) + assert current_ratio.item() == pytest.approx(1.1007, abs=1e-4) + assert candidate_ratio.item() == pytest.approx(1.1184, abs=1e-4) + assert not improved.item() + assert torch.equal(selected, current) + + +def test_select_improving_placements_rejects_near_balance_when_gain_is_below_threshold(): + expert_load = torch.tensor([[1, 2, 1, 17]]) + current = expand_redundant_placement(torch.tensor([[[3], [0]]]), (expert_load).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[3], [1]]]), (expert_load).shape[-1]) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, current, candidate, rebalance_gain_threshold=0.05 + ) + + current_ratio = _estimate_rank_load(expert_load, current).max() / _estimate_rank_load(expert_load, current).mean() + candidate_ratio = ( + _estimate_rank_load(expert_load, candidate).max() / _estimate_rank_load(expert_load, candidate).mean() + ) + assert current_ratio.item() == pytest.approx(1.047619, abs=1e-6) + assert candidate_ratio.item() == pytest.approx(1.0) + assert not improved.item() + assert torch.equal(selected, current) + + +def test_select_improving_placements_accepts_alignment_aware_gain_even_when_current_ranks_are_balanced(): + expert_load = torch.tensor([[100, 129, 100, 129]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (expert_load).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[3], [1]]]), (expert_load).shape[-1]) + + selected, improved, metrics, _before_load, _after_load = select_improving_placements( + expert_load, + current, + candidate, + rebalance_gain_threshold=0.05, + expert_alignment=128, + ) + + assert metrics["model_imbalance_ratio"] == pytest.approx(1.0) + assert metrics["candidate_rebalance_gain"] == pytest.approx(0.25) + assert improved.item() + assert torch.equal(selected, candidate) + + +def test_select_improving_placements_rejects_insufficient_rebalance_gain(): + expert_load = torch.tensor([[1, 1, 6, 7]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (expert_load).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[3], [0]]]), (expert_load).shape[-1]) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, current, candidate, rebalance_gain_threshold=0.05 + ) + + current_ratio = _estimate_rank_load(expert_load, current).max() / _estimate_rank_load(expert_load, current).mean() + candidate_ratio = ( + _estimate_rank_load(expert_load, candidate).max() / _estimate_rank_load(expert_load, candidate).mean() + ) + relative_improvement = (current_ratio - candidate_ratio) / current_ratio + assert current_ratio.item() == pytest.approx(1.4) + assert candidate_ratio.item() == pytest.approx(1.333333, abs=1e-6) + assert relative_improvement.item() == pytest.approx(0.047619, abs=1e-6) + assert not improved.item() + assert torch.equal(selected, current) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, + current, + candidate, + rebalance_gain_threshold=0.04, + ) + + assert improved.item() + assert torch.equal(selected, candidate) + + +@pytest.mark.parametrize("rebalance_gain_threshold", [-0.01, 1.01, float("nan"), float("inf")]) +def test_select_improving_placements_rejects_invalid_rebalance_gain_threshold( + rebalance_gain_threshold, +): + expert_load = torch.tensor([[1, 1, 1, 2]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (expert_load).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[3], [0]]]), (expert_load).shape[-1]) + + with pytest.raises(ValueError, match="rebalance_gain_threshold"): + select_improving_placements( + expert_load, + current, + candidate, + rebalance_gain_threshold=rebalance_gain_threshold, + ) + + +def test_select_improving_placements_accepts_sufficient_rebalance_gain(): + expert_load = torch.tensor([[1, 1, 1, 2]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (expert_load).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[3], [0]]]), (expert_load).shape[-1]) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, current, candidate, rebalance_gain_threshold=0.05 + ) + + current_ratio = _estimate_rank_load(expert_load, current).max() / _estimate_rank_load(expert_load, current).mean() + candidate_ratio = ( + _estimate_rank_load(expert_load, candidate).max() / _estimate_rank_load(expert_load, candidate).mean() + ) + relative_improvement = (current_ratio - candidate_ratio) / current_ratio + assert current_ratio.item() == pytest.approx(1.2) + assert candidate_ratio.item() == pytest.approx(1.0) + assert relative_improvement.item() == pytest.approx(1 / 6) + assert improved.item() + assert torch.equal(selected, candidate) + + +def test_select_improving_placements_rejects_raw_improvement_that_does_not_improve_aligned_compute(): + expert_load = torch.tensor([[1, 1, 1, 8]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (expert_load).shape[-1]) + raw_improving_candidate = expand_redundant_placement(torch.tensor([[[3], [0]]]), (expert_load).shape[-1]) + + _, raw_improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, current, raw_improving_candidate, rebalance_gain_threshold=0.05 + ) + selected, aligned_improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, + current, + raw_improving_candidate, + rebalance_gain_threshold=0.05, + expert_alignment=128, + ) + + assert raw_improved.item() + assert not aligned_improved.item() + assert torch.equal(selected, current) + + +def test_estimate_rank_load_aligns_each_sample_before_accumulation(): + samples = torch.tensor([[[20, 0, 0, 0]], [[20, 0, 0, 0]]]) + placement = expand_redundant_placement(torch.tensor([[[2], [0]]]), (samples.sum(dim=0)).shape[-1]) + + per_sample = _estimate_rank_load(samples, placement, expert_alignment=128) + accumulated = _estimate_rank_load(samples.sum(dim=0), placement, expert_alignment=128) + + assert torch.equal(per_sample[:, 0], torch.tensor([[128.0, 128.0], [128.0, 128.0]])) + assert torch.equal(per_sample.sum(dim=0)[0], torch.tensor([256.0, 256.0])) + assert torch.equal(accumulated[0], torch.tensor([128.0, 128.0])) + + +def test_select_improving_placements_rejects_lower_ratio_when_critical_is_unchanged(): + samples = torch.tensor( + [ + [[255, 220, 226, 254]], + [[172, 278, 51, 238]], + [[249, 291, 284, 183]], + ] + ) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (samples).shape[-1]) + mean_inflating_candidate = expand_redundant_placement(torch.tensor([[[2], [1]]]), (samples).shape[-1]) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + samples, + current, + mean_inflating_candidate, + rebalance_gain_threshold=0.05, + expert_alignment=128, + ) + current_load = _estimate_rank_load(samples, current, expert_alignment=128) + candidate_load = _estimate_rank_load(samples, mean_inflating_candidate, expert_alignment=128) + current_critical = current_load.max(dim=2).values.sum() + candidate_critical = candidate_load.max(dim=2).values.sum() + + assert candidate_load.mean(dim=2).sum() > current_load.mean(dim=2).sum() + assert current_critical == candidate_critical + assert not improved.item() + assert torch.equal(selected, current) + + +def test_select_improving_placements_accepts_five_percent_critical_reduction(): + samples = torch.tensor( + [ + [[13, 352, 348, 141]], + [[287, 175, 236, 179]], + [[316, 99, 266, 353]], + ] + ) + current = expand_redundant_placement(torch.tensor([[[2], [0]]]), (samples).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[2], [1]]]), (samples).shape[-1]) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + samples, current, candidate, rebalance_gain_threshold=0.05, expert_alignment=128 + ) + + assert improved.item() + assert torch.equal(selected, candidate) + + +def test_select_improving_placements_rejects_single_layer_gain_below_model_threshold(): + # Layer 0 becomes better, but layer 1 dominates model critical load. The + # aggregate estimated critical-load reduction gain is below 5%, so neither layer may be changed. + expert_load = torch.tensor([[1, 1, 1, 2], [0, 0, 0, 10]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]], [[2], [0]]]), (expert_load).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[3], [0]], [[2], [0]]]), (expert_load).shape[-1]) + + selected, improved, metrics, _before_load, _after_load = select_improving_placements( + expert_load, current, candidate, rebalance_gain_threshold=0.05 + ) + + assert not torch.any(improved) + assert torch.equal(selected, current) + assert metrics["candidate_rebalance_gain"] == pytest.approx(0.5 / 13) + assert metrics["candidate_changed_layer_count"] == 1 + + +def test_select_improving_placements_accepts_only_when_model_gain_reaches_threshold(): + expert_load = torch.tensor([[1, 1, 1, 2], [0, 0, 0, 5]]) + current = expand_redundant_placement(torch.tensor([[[2], [0]], [[2], [0]]]), (expert_load).shape[-1]) + candidate = expand_redundant_placement(torch.tensor([[[3], [0]], [[2], [0]]]), (expert_load).shape[-1]) + + selected, improved, _metrics, _before_load, _after_load = select_improving_placements( + expert_load, current, candidate, rebalance_gain_threshold=0.05 + ) + + assert torch.equal(improved, torch.tensor([True, False])) + assert torch.equal(selected, candidate) + + +def test_logical_to_physical_map_has_at_most_one_slot_per_rank(): + redundant_expert_ids = expand_redundant_placement(torch.tensor([[2, 3], [0, 1]]), 4) + logical_to_physical, replica_count = build_logical_to_physical_map(redundant_expert_ids, num_logical_experts=4) + + assert logical_to_physical.shape == (4, 2) + assert torch.equal(replica_count, torch.full((4,), 2, dtype=torch.int64)) + + +def test_logical_to_physical_map_requires_expert_count_divisible_by_rank_count_without_source_rank(): + redundant_expert_ids = torch.tensor([[0, 1], [2, 3]]) + + with pytest.raises(AssertionError): + build_logical_to_physical_map(redundant_expert_ids, num_logical_experts=5) + + +def test_logical_to_physical_map_prefers_source_node_replicas(): + # Four ranks, two ranks per node, two primary experts/rank and one + # redundant slot/rank. Expert 0 is primary on rank 0 and replicated on + # rank 2 (the other node). + redundant = expand_redundant_placement(torch.tensor([[4], [5], [0], [1]], dtype=torch.int64), 8) + rank0_map, rank0_count = build_logical_to_physical_map( + redundant, num_logical_experts=8, source_rank=0, node_world_size=2 + ) + rank1_map, rank1_count = build_logical_to_physical_map( + redundant, num_logical_experts=8, source_rank=1, node_world_size=2 + ) + fallback_redundant = expand_redundant_placement(torch.tensor([[4], [5], [0], [1], [2], [3]], dtype=torch.int64), 12) + rank4_map, rank4_count = build_logical_to_physical_map(fallback_redundant, 12, source_rank=4, node_world_size=2) + + assert rank0_count[0].item() == rank1_count[0].item() == 1 + assert torch.equal(rank0_map[0, :1], torch.tensor([0])) + assert torch.equal(rank1_map[0, :1], torch.tensor([0])) + assert rank4_count[0].item() == 2 + assert set(rank4_map[0, :2].tolist()) == {0, 8} + + +def test_source_node_local_maps_fall_back_to_global_replicas(): + redundant = expand_redundant_placement(torch.tensor([[4], [5], [0], [1]], dtype=torch.int64), 8) + maps = [build_logical_to_physical_map(redundant, 8, source_rank=rank, node_world_size=2) for rank in range(4)] + assert maps[0][1][0].item() == maps[1][1][0].item() == 1 + assert maps[2][1][0].item() == maps[3][1][0].item() == 1 + assert maps[0][0][0, 0].item() == maps[1][0][0, 0].item() == 0 + assert maps[2][0][0, 0].item() == maps[3][0][0, 0].item() == 8 + + +def test_source_rank_rotates_selected_replica_order_without_changing_copies(): + # Expert 0 is primary on rank 0 and redundant on rank 1, so both ranks + # on node 0 have the same two local copies. Their source-rank phases + # must differ while their selected set/count remain identical. + redundant = expand_redundant_placement(torch.tensor([[1], [0], [3], [2]], dtype=torch.int64), 4) + rank0_map, rank0_count = build_logical_to_physical_map(redundant, 4, source_rank=0, node_world_size=2) + rank1_map, rank1_count = build_logical_to_physical_map(redundant, 4, source_rank=1, node_world_size=2) + + assert rank0_count[0].item() == rank1_count[0].item() == 2 + assert set(rank0_map[0, :2].tolist()) == set(rank1_map[0, :2].tolist()) == {0, 3} + assert torch.equal(rank1_map[0, :2], torch.tensor([3, 0], dtype=torch.int32)) + + +@pytest.mark.parametrize( + "source_rank,node_world_size", + [(None, None), (0, 2), (1, 2), (2, 2), (3, 2)], +) +def test_logical_to_physical_maps_for_layers_match_single_layer_api(source_rank, node_world_size): + placements_by_layer = expand_redundant_placement( + torch.tensor( + [ + [[4, 5], [0, 1], [0, 1], [2, 3]], + [[6, 7], [0, 1], [0, 1], [2, 3]], + [[4, 5], [0, 1], [0, 1], [2, 3]], + ], + dtype=torch.int64, + ), + 8, + ) + + maps_by_layer, counts_by_layer = build_logical_to_physical_maps_for_layers( + placements_by_layer, + num_logical_experts=8, + source_rank=source_rank, + node_world_size=node_world_size, + ) + expected_by_layer = [ + build_logical_to_physical_map( + placement, + num_logical_experts=8, + source_rank=source_rank, + node_world_size=node_world_size, + ) + for placement in placements_by_layer + ] + + assert maps_by_layer.shape == (3, 8, 4) + assert counts_by_layer.shape == (3, 8) + assert maps_by_layer.dtype == counts_by_layer.dtype == torch.int32 + assert torch.equal(maps_by_layer, torch.stack([item[0] for item in expected_by_layer])) + assert torch.equal(counts_by_layer, torch.stack([item[1] for item in expected_by_layer])) + if source_rank is None: + # expert 0 的主副本在 rank 0，且 rank 1、2 都有一个冗余副本；第 1、2 + # 个冗余副本必须分别写入映射表的第 1、2 列，而不能互相覆盖。 + assert torch.equal(counts_by_layer[:, 0], torch.tensor([3, 3, 3], dtype=torch.int32)) + assert torch.equal( + maps_by_layer[:, 0, :3], + torch.tensor([[0, 6, 10], [0, 6, 10], [0, 6, 10]], dtype=torch.int32), + ) + positions = torch.arange(maps_by_layer.shape[-1]).view(1, 1, -1) + valid = positions < counts_by_layer.unsqueeze(-1) + assert torch.all(maps_by_layer[valid] >= 0) + assert torch.all(maps_by_layer[~valid] == -1) + + +def test_source_node_estimate_matches_local_first_runtime_replica_sharing(): + # Expert 0 has a copy on each node. Node 0 and node 1 issue unequal + # traffic, so collapsing them before planning produces the wrong result. + placement = expand_redundant_placement(torch.tensor([[[4], [5], [0], [1]]], dtype=torch.int64), 8) + source_load = torch.zeros((1, 1, 2, 8), dtype=torch.int64) + source_load[0, 0, 0, 0] = 256 + source_load[0, 0, 1, 0] = 128 + + predicted = _estimate_rank_load(source_load, placement, expert_alignment=128, node_world_size=2) + runtime = _manual_runtime_rank_load(source_load, placement, node_world_size=2, alignment=128) + collapsed_global = _estimate_rank_load(source_load.sum(dim=2), placement, expert_alignment=128) + + assert torch.equal(predicted, runtime) + assert torch.equal(predicted[0, 0], torch.tensor([256.0, 0.0, 128.0, 0.0])) + assert not torch.equal(predicted, collapsed_global) + + +def _count_moved_slots(current: torch.Tensor, target: torch.Tensor) -> int: + """Count rank rows gaining an expert: one migrated row per new expert id.""" + moved = 0 + for layer in range(current.shape[0]): + for rank in range(current.shape[1]): + moved += len(set(target[layer, rank].tolist()) - set(current[layer, rank].tolist())) + return moved + + +def test_transfer_plan_slot_permutation_is_free(): + current = torch.tensor([[4, 5], [6, 7], [0, 1], [2, 3]]) + target = torch.tensor([[5, 4], [7, 6], [1, 0], [3, 2]]) + + assert torch.equal(align_target_placement(current, target), current) + assert build_transfer_plan(current, target, num_logical_experts=8, world_size=4, node_world_size=2) == [] + + +def test_align_target_placement_keeps_retained_experts_in_live_slots(): + current = torch.tensor([[4, 5], [6, 7], [0, 1], [2, 3]]) + target = torch.tensor([[5, 6], [7, 0], [1, 0], [3, 2]]) + + canonical = align_target_placement(current, target) + + assert torch.equal(canonical, torch.tensor([[6, 5], [0, 7], [0, 1], [2, 3]])) + + +def test_canonical_placement_keeps_transfer_rows_and_published_map_consistent(): + current = torch.tensor([[0, 1, 2], [2, 3, 4], [4, 5, 0]]) + target = current.roll(1, dims=0).flip(dims=(1,)) + canonical = align_target_placement(current, target) + plan = build_transfer_plan(current, canonical, 6, 3, 3) + live = current.clone() + for step in plan: + live[step.dst_rank, step.dst_slot] = current[step.src_rank, step.src_local_row] + assert torch.equal(live, canonical) + maps, counts = build_logical_to_physical_map(canonical, 6) + for expert, count in enumerate(counts.tolist()): + for physical_id in maps[expert, :count].tolist(): + rank, row = divmod(physical_id, current.shape[1]) + assert live[rank, row] == expert + + +def test_plan_and_broadcast_publishes_canonical_placement(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.global_rank = 0 + manager.world_size = 4 + manager.node_world_size = 2 + manager.num_logical_experts = 8 + manager.num_redundant_experts_per_rank = 2 + manager.current_placement = torch.tensor([[[4, 5], [6, 7], [0, 1], [2, 3]]]) + manager._initial_plan_complete = False + manager._evaluation_snapshot = {"placement": manager.current_placement} + manager.placement_stickiness = 0.1 + manager.rebalance_gain_threshold = 0.05 + manager.evaluation_group = object() + candidate = torch.tensor([[[5, 4], [7, 6], [1, 0], [3, 2]]]) + broadcasts = [] + + def fixed_selector(*_args, **_kwargs): + rank_load = torch.full((1, 4), 100.0) + return candidate.clone(), torch.tensor([True]), {}, rank_load, rank_load + + def record_broadcast(result_list, **_kwargs): + broadcasts.append(result_list[0]) + + monkeypatch.setattr(manager_module, "plan_eplb_candidate", lambda *_args, **_kwargs: candidate.clone()) + monkeypatch.setattr(manager_module, "select_improving_placements", fixed_selector) + monkeypatch.setattr(manager_module.dist, "broadcast_object_list", record_broadcast) + + result = manager._plan_and_broadcast(torch.full((1, 1, 2, 8), 100, dtype=torch.int64)) + + assert torch.equal(result["placement"], manager.current_placement) + assert broadcasts and torch.equal(broadcasts[0]["placement"], manager.current_placement) + + +def test_plan_and_broadcast_propagates_rank_zero_error_after_existing_broadcast(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.global_rank = 0 + manager.world_size = 2 + manager.node_world_size = 2 + manager.num_logical_experts = 8 + manager.num_redundant_experts_per_rank = 2 + manager.current_placement = torch.tensor([[[4, 5], [6, 7], [0, 1], [2, 3]]]) + manager._initial_plan_complete = False + manager._evaluation_snapshot = {"placement": manager.current_placement} + manager.placement_stickiness = 0.1 + manager.rebalance_gain_threshold = 0.05 + manager.evaluation_group = object() + broadcasted = [] + + def broken_planner(*_args, **_kwargs): + raise RuntimeError("planner boom") + + def record_broadcast(result_list, **_kwargs): + broadcasted.append(result_list[0]) + + monkeypatch.setattr(manager_module, "plan_eplb_candidate", broken_planner) + monkeypatch.setattr(manager_module.dist, "broadcast_object_list", record_broadcast) + + with pytest.raises(RuntimeError, match="EPLB planner failed on rank zero") as exc_info: + manager._plan_and_broadcast(torch.full((1, 1, 1, 8), 100, dtype=torch.int64)) + + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert str(exc_info.value.__cause__) == "planner boom" + assert broadcasted == [{"kind": "error", "message": "RuntimeError: planner boom"}] + + +def test_evaluation_no_improvement_logs_model_fields_without_reopening_interval_window( + monkeypatch, +): + class DoneThread: + def join(self): + pass + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.evaluation_in_flight = True + manager._evaluation_lock = threading.Lock() + manager._evaluation_error = None + manager._evaluation_thread = DoneThread() + manager._evaluation_result = { + "kind": "no_improvement", + "model_imbalance_ratio": 1.2, + "candidate_model_imbalance_ratio": 1.1, + "candidate_rebalance_gain": 0.01, + "candidate_changed_layer_count": 2, + } + manager.global_rank = 0 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.min_rebalance_interval = 0 + manager._last_evaluation_step = 0 + manager.weights = [] + manager._eplb_states = [] + recordings, logs = [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + monkeypatch.setattr(manager_module.logger, "info", lambda *args: logs.append(args)) + + assert not manager._poll_evaluation() + assert recordings == [True] + assert "model_imbalance_ratio" in logs[0][0] + assert "candidate_rebalance_gain" in logs[0][0] + assert "candidate_changed_layer_count" in logs[0][0] + assert "actual_changed_layer_count" in logs[0][0] + assert "next_sampling_interval" in logs[0][0] + assert manager.sampling_interval == 20 + + +def test_evaluation_worker_error_is_raised_by_main_thread(): + class DoneThread: + def join(self): + pass + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.evaluation_in_flight = True + manager._evaluation_lock = threading.Lock() + manager._evaluation_error = RuntimeError("planner failed") + manager._evaluation_result = None + manager._evaluation_thread = DoneThread() + + with pytest.raises(RuntimeError, match="planner failed"): + manager._poll_evaluation() + + +def test_evaluation_state_is_cleared_before_second_round(monkeypatch): + class DoneThread: + def join(self): + pass + + class PendingThread: + def __init__(self, **_kwargs): + self.started = False + + def start(self): + self.started = True + + def join(self): + pytest.fail("pending worker must not be joined") + + class Event: + def record(self, _stream): + pass + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.evaluation_in_flight = True + manager._evaluation_lock = threading.Lock() + manager._evaluation_error = None + manager._evaluation_thread = DoneThread() + manager._evaluation_result = { + "kind": "no_improvement", + "model_imbalance_ratio": 1.2, + "candidate_model_imbalance_ratio": 1.1, + "candidate_rebalance_gain": 0.01, + "candidate_changed_layer_count": 1, + } + manager.global_rank = 1 + manager.min_rebalance_interval = 0 + manager._last_evaluation_step = 0 + manager._freeze_evaluation_samples = lambda: None + manager.prefill_steps = 0 + manager.step_interval = 1 + manager.sampling_interval = 1 + manager.weights = [] + manager._eplb_states = [] + manager._set_recording = lambda _enabled: None + monkeypatch.setattr(manager_module.threading, "Thread", PendingThread) + monkeypatch.setattr(manager_module.torch.cuda, "Event", Event) + monkeypatch.setattr(manager_module.torch.cuda, "current_stream", lambda: object()) + + assert not manager._poll_evaluation() + assert manager._evaluation_result is None + assert manager._evaluation_error is None + + manager._start_evaluation() + assert manager.evaluation_in_flight + assert manager._evaluation_result is None + assert manager._poll_evaluation() # New worker has not produced a result. + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + "step_interval,steady_steps,capacity", [(1, 4, 8), (2, 4, 8), (3, 4, 8), (1, 1, 2), (20, 64, 40)] +) +def test_eplb_counter_preserves_each_dispatch_in_maximum_window(monkeypatch, step_interval, steady_steps, capacity): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk_eplb + + args = StartArgs( + enable_ep_moe=True, + enable_prefill_eplb=True, + eplb_num_redundant_experts_per_rank=2, + run_mode="prefill", + ) + weight = object.__new__(fused_weight_module.FusedMoeWeight) + weight.n_routed_experts = 4 + weight.global_world_size = 2 + weight.global_rank_ = 0 + weight.enable_ep_moe = True + weight.tp_rank_, weight.tp_world_size_ = 0, 1 + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr(envs_utils, "get_prefill_eplb_step_interval", lambda: step_interval) + monkeypatch.setattr(envs_utils, "get_prefill_eplb_steady_sample_steps", lambda: steady_steps) + monkeypatch.setattr(dist_utils, "get_node_world_size", lambda: 2) + monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor: tensor) + original_zeros = torch.zeros + + def cpu_zeros(*shape, **kwargs): + kwargs.pop("device", None) + return original_zeros(*shape, **kwargs) + + monkeypatch.setattr(fused_weight_module.torch, "zeros", cpu_zeros) + + weight._init_expert_parallel_state({}) + + state = weight.expert_parallel_state.eplb + assert state.route_counter.shape == (capacity, 4, 4) + state.route_counter = state.route_counter.to("cuda") + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager._eplb_states = [state] + manager.step_interval, manager.steady_sample_steps = step_interval, steady_steps + manager.sampling_interval = 16 * step_interval + steps = max(step_interval, manager._steady_sample_window_steps()) + manager._use_full_window = False + manager.current_placement = state.physical_to_logical.unsqueeze(0) + manager.prefill_steps = 0 + manager._sample_step_counts = deque([0], maxlen=max(step_interval, 4) + 1) + manager.in_flight = True + state.recording = True + expected = [] + for dispatch in range(2 * steps): + if dispatch % 2 == 0: + manager.before_prefill() + # Move the hotspot between rank-local expert pairs across dispatches. + expert = 2 * (dispatch % 2) + scores = torch.full((256, 4), -10.0, device="cuda") + scores[:128, expert] = 10 + scores[128:, expert + 1] = 10 + triton_grouped_topk_eplb( + hidden_states=scores, + gating_output=scores, + correction_bias=None, + topk=1, + renormalize=True, + num_expert_group=1, + topk_group=1, + scoring_func="softmax", + logical_to_physical_map=state.logical_to_physical_map.to("cuda"), + logical_replica_count=state.logical_replica_count.to("cuda"), + expert_counter=state.route_counter, + sample_index=state.next_sample_index(), + record_load=state.recording, + use_grouped_topk=False, + ) + expected.append([128, 128, 0, 0] if expert == 0 else [0, 0, 128, 128]) + if dispatch % 2 == 1: + manager.step() + manager._freeze_evaluation_samples() + samples = manager._evaluation_snapshot["samples"][:, 0].cpu() + assert torch.equal(samples[..., 0], torch.tensor(expected, dtype=torch.int64)) + for copies in (1, 2): + start = 1 + copies * (copies - 1) // 2 + assert torch.equal(samples[..., start : start + copies].sum(-1), samples[..., 0]) + + +def test_steady_sampling_resets_fixed_ring_without_retained_history(): + counter = torch.ones((8, 4), dtype=torch.int64) + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + state = _test_parallel_state( + eplb=True, + route_counter=counter, + recording=False, + recorded_sample_count=99, + num_logical_experts=4, + world_size=1, + ).eplb + manager._eplb_states = [state] + + manager._reset_recorded_samples() + manager._reset_recorded_samples() + + assert state.route_counter.shape == (8, 4) + assert torch.count_nonzero(state.route_counter) == 0 + assert state.recorded_sample_count == 0 + assert not hasattr(manager, "_retained_local_samples") + assert not hasattr(manager, "_sample_history") + + +@pytest.mark.parametrize( + "enabled,run_mode,nnodes", [(False, "prefill", 1), (False, "normal", 1), (True, "decode", 1), (True, "decode", 2)] +) +@pytest.mark.parametrize("draft", [False, True]) +def test_inactive_eplb_creates_logical_layout_without_runtime_state(monkeypatch, enabled, run_mode, nnodes, draft): + args = type( + "Args", + (), + { + "enable_prefill_eplb": enabled, + "eplb_num_redundant_experts_per_rank": 2, + "run_mode": run_mode, + "nnodes": nnodes, + }, + )() + weight = object.__new__(fused_weight_module.FusedMoeWeight) + weight.n_routed_experts = 4 + weight.global_world_size = 2 + weight.global_rank_ = 0 + weight.enable_ep_moe = True + weight.tp_rank_, weight.tp_world_size_ = 0, 1 + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr( + placement_module, + "build_initial_redundant_expert_ids", + lambda *args: pytest.fail("inactive EPLB must not initialize physical replicas"), + ) + + weight._init_expert_parallel_state({"is_mtp_draft_model": draft}) + + assert weight.expert_parallel_state is not None + assert weight.expert_parallel_state.eplb is None + assert weight.expert_parallel_state.num_total_physical_experts == weight.n_routed_experts + assert not hasattr(weight, "_initial_redundant_expert_ids") + assert not hasattr(weight, "route_counter") + assert not hasattr(weight, "routed_expert_counter_tensor") + + +@pytest.mark.parametrize("run_mode", ["prefill", "normal"]) +@pytest.mark.parametrize("transport_experts", [4, 8]) +def test_draft_keeps_static_unrecorded_layout_in_shared_expert_space(monkeypatch, run_mode, transport_experts): + # This checks initialization contracts; actual CUDA storage is covered by + # the checkpoint-loading and graph routing tests. + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: False) + real_zeros = torch.zeros + monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor, *args, **kwargs: tensor) + monkeypatch.setattr(torch, "zeros", lambda *args, **kwargs: real_zeros(*args, **{**kwargs, "device": "cpu"})) + args = StartArgs( + enable_ep_moe=True, + enable_prefill_eplb=True, + eplb_num_redundant_experts_per_rank=2, + run_mode=run_mode, + ) + weight = object.__new__(fused_weight_module.FusedMoeWeight) + weight.n_routed_experts = 4 + weight.global_world_size = 2 + weight.global_rank_ = 0 + weight.enable_ep_moe = True + weight.tp_rank_, weight.tp_world_size_ = 0, 1 + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr(dist_utils, "get_node_world_size", lambda: 2) + monkeypatch.setattr(deepgemm_module.dist_group_manager, "ll_num_experts", transport_experts, raising=False) + if transport_experts != 8: + monkeypatch.setattr( + placement_module, + "build_initial_redundant_expert_ids", + lambda *args: pytest.fail("reject incompatible transport before creating draft weights"), + ) + with pytest.raises(ValueError, match="same DeepEP expert space"): + weight._init_expert_parallel_state({"is_mtp_draft_model": True}) + return + + weight._init_expert_parallel_state({"is_mtp_draft_model": True}) + + state = weight.expert_parallel_state.eplb + assert weight.expert_parallel_state.num_total_physical_experts == 8 + assert not state.recording + assert state.next_sample_index() == state.recorded_sample_count == 0 + assert state.route_counter.shape[0] == 2 + assert state.migration_poll is None + assert state.placement_generation == 0 + assert not hasattr(weight, "_initial_redundant_expert_ids") + + +def test_manager_evaluation_collective_preserves_source_node_axis(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.weights = [ + type( + "Weight", + (), + { + "expert_parallel_state": _test_parallel_state( + eplb=True, + route_counter=torch.zeros((1, 4), dtype=torch.int64), + num_logical_experts=4, + world_size=1, + ) + }, + )() + ] + manager._eplb_states = [manager.weights[0].expert_parallel_state.eplb] + manager.global_rank = 2 + manager.world_size = 4 + manager.node_world_size = 2 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.num_logical_experts = 4 + manager.num_redundant_experts_per_rank = 1 + manager.current_placement = build_initial_redundant_expert_ids(4, 4, 1).unsqueeze(0) + manager.evaluation_group = object() + manager._evaluation_lock = threading.Lock() + manager._evaluation_result = None + manager._evaluation_error = None + manager._evaluation_snapshot = {"steps": 4} + local = torch.full((1, 1, 4), 100, dtype=torch.int64) + manager._collect_local_samples = lambda: local + seen = {} + + def all_reduce(tensor, **kwargs): + seen["before"] = tensor.clone() + seen["group"] = kwargs["group"] + # Simulate source node 0's contribution from the other ranks. + tensor[:, :, 0].fill_(100) + + monkeypatch.setattr(manager_module.dist, "all_reduce", all_reduce) + monkeypatch.setattr(manager_module.torch.cuda, "set_device", lambda _device: None) + + def plan_and_broadcast(global_load): + seen["global_load"] = global_load.clone() + return {"kind": "insufficient"} + + manager._plan_and_broadcast = plan_and_broadcast + + manager._evaluate_after_event(type("Event", (), {"synchronize": lambda self: None})()) + + assert seen["group"] is manager.evaluation_group + expected_local = local + assert seen["before"].shape == (1, 1, 2, 4) + assert torch.equal(seen["before"][:, :, 0], torch.zeros_like(expected_local)) + assert torch.equal(seen["before"][:, :, 1], expected_local) + assert torch.equal(seen["global_load"][:, :, 0], torch.full_like(expected_local, 100)) + assert torch.equal(seen["global_load"][:, :, 1], expected_local) + assert manager._evaluation_error is None + assert manager._evaluation_result["recorded_sample_count"] == 1 + assert manager._evaluation_result["sample_window_steps"] == 4 + + +@pytest.mark.parametrize("phase", [0, 1]) +def test_manager_planned_evaluation_builds_improved_metadata_in_one_multilayer_call(monkeypatch, phase): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.weights = [ + type( + "Weight", + (), + { + "expert_parallel_state": _test_parallel_state( + eplb=True, + route_counter=torch.zeros((1, 4), dtype=torch.int64), + num_logical_experts=4, + world_size=4, + ) + }, + )() + for _ in range(3) + ] + manager._eplb_states = [weight.expert_parallel_state.eplb for weight in manager.weights] + for state in manager._eplb_states: + state.replica_rotation_rank = phase + manager.global_rank = 1 + manager.world_size = 4 + manager.node_world_size = 2 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.num_logical_experts = 4 + manager.num_redundant_experts_per_rank = 1 + manager.current_placement = ( + expand_redundant_placement(build_initial_redundant_expert_ids(4, 4, 1), 4) + .unsqueeze(0) + .expand(3, -1, -1) + .clone() + ) + manager.evaluation_group = object() + manager._evaluation_lock = threading.Lock() + manager._evaluation_result = None + manager._evaluation_error = None + manager._evaluation_snapshot = {"steps": 4} + prepare_calls = [] + prepared_batches = object() + + def prepare_transfer(layer_plans): + prepare_calls.append(layer_plans) + return prepared_batches + + manager.transfer = SimpleNamespace(prepare_transfer=prepare_transfer) + manager._collect_local_samples = lambda: torch.full((1, 3, 4), 100, dtype=torch.int64) + planned_placement = torch.tensor( + [ + [[3, 0], [0, 1], [1, 2], [2, 3]], + [[2, 3], [3, 0], [0, 1], [1, 2]], + [[1, 2], [2, 3], [3, 0], [0, 1]], + ], + dtype=torch.int64, + ) + manager._plan_and_broadcast = lambda _global_load: { + "kind": "planned", + "placement": planned_placement, + "improved": torch.tensor([True, False, True]), + } + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda _tensor, **_kwargs: None) + monkeypatch.setattr(manager_module.torch.cuda, "set_device", lambda _device: None) + calls = [] + original_build_maps_for_layers = manager_module.build_logical_to_physical_maps_for_layers + + def build_maps_for_layers(*args, **kwargs): + calls.append(args[0].shape) + return original_build_maps_for_layers(*args, **kwargs) + + monkeypatch.setattr( + manager_module, + "build_logical_to_physical_maps_for_layers", + build_maps_for_layers, + ) + + manager._evaluate_after_event(type("Event", (), {"synchronize": lambda self: None})()) + + assert manager._evaluation_error is None + assert calls == [torch.Size([2, 4, 2])] + metadata = manager._evaluation_result["metadata"] + assert metadata[1] is None + assert [layer_index for layer_index, _plan in manager._evaluation_result["layer_plans"]] == [0, 2] + assert prepare_calls == [manager._evaluation_result["layer_plans"]] + assert manager._evaluation_result["prepared_batches"] is prepared_batches + for layer_index in (0, 2): + item = metadata[layer_index] + expected = build_logical_to_physical_map( + planned_placement[layer_index], + 4, + source_rank=manager.global_rank, + node_world_size=manager.node_world_size, + replica_rotation_rank=phase, + ) + assert torch.equal(item[0], expected[0]) + assert torch.equal(item[1], expected[1]) + + +def test_manager_preparation_error_is_saved_as_evaluation_error(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.weights = [ + type( + "Weight", + (), + { + "expert_parallel_state": _test_parallel_state( + eplb=True, + route_counter=torch.zeros((1, 2), dtype=torch.int64), + num_logical_experts=2, + world_size=1, + ) + }, + )() + ] + manager._eplb_states = [manager.weights[0].expert_parallel_state.eplb] + manager.global_rank = 0 + manager.world_size = 1 + manager.node_world_size = 1 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.num_logical_experts = 2 + manager.current_placement = torch.tensor([[[0], [1]]], dtype=torch.int64) + manager.evaluation_group = object() + manager._evaluation_lock = threading.Lock() + manager._evaluation_result = None + manager._evaluation_error = None + manager._evaluation_snapshot = {"steps": 4} + manager._collect_local_samples = lambda: torch.ones((1, 1, 2), dtype=torch.int64) + manager._plan_and_broadcast = lambda _global_load: { + "kind": "planned", + "placement": torch.tensor([[[0], [1]]], dtype=torch.int64), + "improved": torch.tensor([True]), + } + + def fail_prepare(_layer_plans): + raise RuntimeError("prepare failed") + + manager.transfer = SimpleNamespace(prepare_transfer=fail_prepare) + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda _tensor, **_kwargs: None) + monkeypatch.setattr(manager_module.torch.cuda, "set_device", lambda _device: None) + monkeypatch.setattr(manager_module, "build_transfer_plan", lambda *_args, **_kwargs: []) + + manager._evaluate_after_event(type("Event", (), {"synchronize": lambda self: None})()) + + assert manager._evaluation_result is None + assert isinstance(manager._evaluation_error, RuntimeError) + assert str(manager._evaluation_error) == "prepare failed" + + +def test_decode_dispatch_uses_published_physical_ids_and_expert_count(monkeypatch): + class Buffer: + def low_latency_dispatch(self, **kwargs): + calls.append(kwargs) + return "recv", "masked", "handle", "event", "hook" + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.quant_method = type("Quant", (), {"method_name": "fp8"})() + impl.n_routed_experts = 128 + _set_expert_parallel_state(impl, _test_parallel_state(eplb=True)) + impl._select_experts = lambda **_kwargs: ( + torch.ones((1, 2)), + torch.tensor([[0, 143]], dtype=torch.int32), + torch.tensor([[0, 127]], dtype=torch.int32), + ) + calls = [] + monkeypatch.setattr( + deepgemm_module, + "get_deepep_num_max_dispatch_tokens_per_rank_decode", + lambda: 16, + ) + monkeypatch.setattr(deepgemm_module.dist_group_manager, "ep_low_latency_buffer", Buffer()) + + result = impl.low_latency_dispatch( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + False, + 2, + False, + 0, + 0, + "softmax", + ) + + assert result[2].tolist() == [[0, 143]] + assert calls[0]["num_experts"] == 144 + + +def test_decode_select_maps_experts_without_sampling_or_polling(monkeypatch): + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 1.0 + _set_expert_parallel_state(impl, _test_parallel_state(eplb=True)) + logical_ids = torch.tensor([[3, 127]], dtype=torch.int64) + physical_ids = torch.tensor([[13, 143]], dtype=torch.int64) + counter = impl.eplb.route_counter.clone() + impl.eplb.recording = True + + def forbidden(): + pytest.fail("decode must not sample or poll EPLB") + + impl.eplb.next_sample_index = forbidden + impl.eplb.migration_poll = forbidden + + def select(**kwargs): + assert kwargs["record_load"] is False + assert kwargs["use_replicas"] is False + assert kwargs["sample_index"] == 0 + assert kwargs["logical_to_physical_map"] is impl.eplb.logical_to_physical_map + return torch.ones((1, 2)), physical_ids, logical_ids + + monkeypatch.setattr(grouped_topk, "triton_grouped_topk_eplb", select) + _, selected, origin = impl._select_experts( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + 2, + False, + False, + 0, + 0, + "softmax", + is_prefill=False, + ) + + assert selected is physical_ids + assert origin is logical_ids + assert torch.equal(impl.eplb.route_counter, counter) + assert impl.eplb.recorded_sample_count == 0 + + +def test_eplb_prefill_uses_single_fused_path_for_global_topk(monkeypatch): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_topk + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 1.0 + impl.quant_method = object() + _set_expert_parallel_state(impl, _test_parallel_state(eplb=True)) + physical_ids = torch.tensor([[130, 131]], dtype=torch.long) + calls = [] + + def fused_topk(**kwargs): + calls.append(kwargs) + return torch.ones((1, 2)), physical_ids, None + + monkeypatch.setattr(grouped_topk, "triton_grouped_topk_eplb", fused_topk) + monkeypatch.setattr(deepgemm_module, "quantize_fused_experts_input", lambda *_args: "qinput") + + weights, topk_idx, qinput = impl.select_experts_and_quant_input( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + object(), + False, + 2, + False, + 0, + 0, + "softmax", + ) + + assert weights.tolist() == [[1.0, 1.0]] + assert topk_idx is physical_ids + assert topk_idx.dtype is torch.long + assert qinput == "qinput" + assert not calls[0]["use_grouped_topk"] + assert not calls[0]["return_logical_ids"] + + +def test_eplb_prefill_dispatch_consumes_physical_ids_and_event(monkeypatch): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_topk + + class Buffer: + def dispatch(self, _qinput, **kwargs): + calls.append(kwargs) + return ( + (torch.empty((4, 2)),), + "recv_idx", + "recv_weight", + SimpleNamespace(num_recv_tokens_per_expert_list=[4]), + SimpleNamespace(current_stream_wait=lambda: None), + ) + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 1.0 + impl.quant_method = object() + state = _test_parallel_state( + eplb=True, + route_counter=torch.zeros((3, 128), dtype=torch.int64), + recording=True, + ) + _set_expert_parallel_state(impl, state) + impl.ep_balance_counters = None + calls, fused_calls = [], [] + physical_ids = torch.tensor([[130, 131]], dtype=torch.long) + + def fused_topk(**kwargs): + fused_calls.append(kwargs) + return torch.ones((1, 2)), physical_ids, None + + monkeypatch.setattr(grouped_topk, "triton_grouped_topk_eplb", fused_topk) + monkeypatch.setattr(deepgemm_module, "quantize_fused_experts_input", lambda *_args: "qinput") + monkeypatch.setattr(deepgemm_module.dist_group_manager, "ep_buffer", Buffer()) + monkeypatch.setattr( + deepgemm_module, + "get_deepep_num_max_dispatch_tokens_per_rank_prefill", + lambda: 16, + ) + monkeypatch.setattr(deepgemm_module, "get_ep_num_sms", lambda: 8) + + weights, topk_idx, qinput = impl.select_experts_and_quant_input( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + object(), + True, + 2, + False, + 1, + 8, + "sigmoid", + ) + caller_event = object() + impl.dispatch( + qinput, + topk_idx, + weights, + overlap_event=caller_event, + ) + + assert topk_idx is physical_ids + assert len(fused_calls) == 1 + assert fused_calls[0]["sample_index"] == 0 + assert fused_calls[0]["record_load"] + assert state.eplb.recorded_sample_count == 1 + assert calls[0]["topk_idx"] is physical_ids + assert calls[0]["topk_idx"].dtype is torch.long + assert calls[0]["previous_event"] is caller_event + + +def test_prefill_dispatch_preserves_event(monkeypatch): + class Buffer: + def dispatch(self, _qinput, **kwargs): + calls.append(kwargs) + return ( + (torch.empty((4, 2)),), + "recv_idx", + "recv_weight", + SimpleNamespace(num_recv_tokens_per_expert_list=[4]), + SimpleNamespace(current_stream_wait=lambda: None), + ) + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + _set_expert_parallel_state(impl, _test_parallel_state(eplb=True)) + impl.ep_balance_counters = None + calls = [] + caller_event = object() + monkeypatch.setattr(deepgemm_module.dist_group_manager, "ep_buffer", Buffer()) + monkeypatch.setattr( + deepgemm_module, + "get_deepep_num_max_dispatch_tokens_per_rank_prefill", + lambda: 16, + ) + monkeypatch.setattr(deepgemm_module, "get_ep_num_sms", lambda: 8) + + impl.dispatch( + "qinput", + torch.tensor([[1, 2]], dtype=torch.long), + torch.ones((1, 2)), + caller_event, + ) + + assert calls[0]["previous_event"] is caller_event + assert calls[0]["topk_idx"].dtype is torch.long + + +def test_deepgemm_constructor_configures_eplb(): + state = _validated_expert_parallel_state() + impl = deepgemm_module.FuseMoeDeepGEMM(4, 0, 1.0, SimpleNamespace(), expert_parallel_state=state) + assert impl.expert_parallel_state is state + + +def test_prefill_eplb_returns_requested_logical_ids(monkeypatch): + from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_topk + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 1.0 + _set_expert_parallel_state(impl, _test_parallel_state(eplb=True)) + + def fused_topk(**kwargs): + assert kwargs["return_logical_ids"] + return ( + torch.ones((1, 2)), + torch.tensor([[13, 14]], dtype=torch.int32), + torch.tensor([[3, 4]], dtype=torch.int32), + ) + + monkeypatch.setattr(grouped_topk, "triton_grouped_topk_eplb", fused_topk) + _, physical_ids, logical_ids = impl._select_experts( + torch.empty((1, 4)), + torch.empty((1, 128)), + None, + 2, + False, + False, + 0, + 0, + "softmax", + is_prefill=True, + preserve_logical_ids=True, + ) + + assert physical_ids.tolist() == [[13, 14]] + assert logical_ids.tolist() == [[3, 4]] + + +def test_decode_masked_group_gemm_uses_all_published_physical_rows( + monkeypatch, +): + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + _set_expert_parallel_state(impl, _test_parallel_state(eplb=True, num_logical_experts=8, world_size=1)) + captured = {} + + def masked(*args, **kwargs): + captured["w13"] = args[3] + captured["w13_scale"] = args[4] + captured["w2"] = args[5] + captured["w2_scale"] = args[6] + return "out" + + monkeypatch.setattr(deepgemm_module, "masked_group_gemm", masked) + pack = lambda: type( + "Pack", + (), + {"weight": torch.empty((10, 4)), "weight_scale": torch.empty((10, 1))}, + )() + + assert impl.masked_group_gemm((torch.empty((1, 4)),), pack(), pack(), torch.empty(8), torch.float16, 1) == "out" + assert captured["w13"].shape[0] == captured["w2"].shape[0] == 10 + assert captured["w13_scale"].shape[0] == captured["w2_scale"].shape[0] == 10 + + +def test_decode_fused_experts_uses_original_weight_packs_and_physical_experts( + monkeypatch, +): + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.n_routed_experts = 128 + _set_expert_parallel_state( + impl, + _test_parallel_state(eplb=True, num_logical_experts=128, world_size=16, num_redundant_experts_per_rank=2), + ) + impl.quant_method = object() + impl.ep_balance_counters = None + captured = [] + + def fused(**kwargs): + captured.append(kwargs) + return "out" + + monkeypatch.setattr(deepgemm_module, "fused_experts", fused) + pack = lambda: type( + "Pack", + (), + { + "weight": torch.empty((10, 4)), + "weight_scale": torch.empty((10, 1)), + "weight_zero_point": None, + }, + )() + w13, w2 = pack(), pack() + + for _ in range(2): + assert ( + impl._fused_experts( + torch.empty((1, 4)), + w13, + w2, + torch.ones((1, 2)), + torch.zeros((1, 2), dtype=torch.int64), + is_prefill=False, + ) + == "out" + ) + + assert [call["num_experts"] for call in captured] == [160, 160] + assert all(call["w13"].weight.shape[0] == call["w2"].weight.shape[0] == 10 for call in captured) + assert captured[0]["w13"] is captured[1]["w13"] + assert captured[0]["w2"] is captured[1]["w2"] + + +def test_extract_expert_tensors_includes_weight_scale_and_zero_point_in_order(): + class Pack: + def __init__(self, offset, scale=True, zero=True): + self.weight = torch.full((3, 2), offset) + self.weight_scale = torch.full((3, 1), offset + 1) if scale else None + self.weight_zero_point = torch.full((3, 1), offset + 2) if zero else None + + weight = type("Weight", (), {"w13": Pack(1), "w2": Pack(10, scale=False, zero=False)})() + tensors = extract_eplb_expert_tensors(weight) + assert [name for name, _ in tensors] == [ + "w13.weight", + "w13.weight_scale", + "w13.weight_zero_point", + "w2.weight", + ] + + +def test_commit_staging_rows_only_overwrites_changed_absolute_rows(): + live = torch.arange(20).reshape(5, 4) + staging = torch.full((5, 4), -1) + _commit_staging_rows( + live, + staging, + changed_dst_slots=(3, 4), + ) + assert torch.equal(live[:3], torch.arange(12).reshape(3, 4)) + assert torch.equal(live[3:], staging[3:]) + + +def test_commit_staging_rows_preserves_unchanged_destination_slots(): + live = torch.arange(28).reshape(7, 4) + staging = -torch.arange(28).reshape(7, 4) + original = live.clone() + + _commit_staging_rows( + live, + staging, + changed_dst_slots=(6, 4), + ) + + assert torch.equal(live[:4], original[:4]) + assert torch.equal(live[4], staging[4]) + assert torch.equal(live[5], original[5]) + assert torch.equal(live[6], staging[6]) + + +def test_commit_staging_rows_merges_contiguous_changed_slots(): + copies = [] + + class View: + def __init__(self, owner, start, length): + self.owner = owner + self.start = start + self.length = length + + def copy_(self, source, **_kwargs): + copies.append( + ( + self.owner, + self.start, + self.length, + source.owner, + source.start, + source.length, + ) + ) + + class Tensor: + def __init__(self, owner, rows): + self.owner = owner + self.shape = (rows,) + + def narrow(self, _dim, start, length): + return View(self.owner, start, length) + + _commit_staging_rows( + Tensor("live", 20), + Tensor("staging", 20), + changed_dst_slots=(13, 11, 12), + ) + + assert copies == [("live", 11, 3, "staging", 11, 3)] + + +def test_manager_inflight_ready_gate_commits_ordered_prefix_and_propagates_worker_error( + monkeypatch, +): + class Transfer: + def __init__(self): + self.pending = [(0, 0), (1, 1), (2, 2)] + self.commits = [] + self.finished = 0 + + def pending_layers(self): + return self.pending + + def commit(self, layer, buffer_index, post_copy=None): + assert self.pending.pop(0) == (layer, buffer_index) + self.commits.append((layer, buffer_index)) + if post_copy is not None: + post_copy() + + def finish(self): + self.finished += 1 + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.transfer = Transfer() + manager.in_flight = True + manager.world_size = 2 + manager.control_group = object() + manager._control_ready_count = torch.empty(1, dtype=torch.int32) + manager.in_flight_layers = [0, 1, 2] + manager._commit_layer_metadata = lambda layer: committed.append(layer) + manager._finish_rebalance = lambda: finished.append(True) + committed, finished = [], [] + operations = [] + current_stream_calls = [] + + class CurrentStream: + def wait_stream(self, stream): + operations.append(("wait", stream)) + + overlap_stream = object() + + def current_stream(): + current_stream_calls.append(True) + return CurrentStream() + + monkeypatch.setattr(manager_module.torch.cuda, "current_stream", current_stream) + monkeypatch.setattr(g_infer_context, "get_overlap_stream", lambda: overlap_stream) + + original_commit = manager.transfer.commit + + def record_commit(*args, **kwargs): + operations.append(("commit", args[0])) + return original_commit(*args, **kwargs) + + manager.transfer.commit = record_commit + + def set_global_ready(count): + return lambda tensor, **kwargs: tensor.fill_(count) + + monkeypatch.setattr(manager_module.dist, "all_reduce", set_global_ready(0)) + manager._poll_in_flight() + assert manager.transfer.commits == [] + assert operations == [] + assert current_stream_calls == [] + + # Local rank has three prefetched layers, but global MIN-ready only permits two. + monkeypatch.setattr(manager_module.dist, "all_reduce", set_global_ready(2)) + manager._poll_in_flight() + assert manager.transfer.commits == [(0, 0), (1, 1)] + assert committed == [0, 1] + assert not finished + assert manager.transfer.finished == 0 + assert operations == [("wait", overlap_stream), ("commit", 0), ("commit", 1)] + assert current_stream_calls == [True] + + monkeypatch.setattr(manager_module.dist, "all_reduce", set_global_ready(1)) + manager._poll_in_flight() + assert finished == [True] + assert manager.transfer.finished == 1 + assert operations == [ + ("wait", overlap_stream), + ("commit", 0), + ("commit", 1), + ("wait", overlap_stream), + ("commit", 2), + ] + assert current_stream_calls == [True, True] + + manager.in_flight_layers = [3] + manager.transfer.pending = [(9, 0)] + monkeypatch.setattr(manager_module.dist, "all_reduce", set_global_ready(1)) + with pytest.raises(RuntimeError, match="does not match expected"): + manager._poll_in_flight() + + class BrokenTransfer: + def pending_layers(self): + raise RuntimeError("boom") + + manager.transfer = BrokenTransfer() + encoded_statuses = [] + + def retain_local_error(tensor, **_kwargs): + encoded_statuses.append(int(tensor.item())) + + monkeypatch.setattr(manager_module.dist, "all_reduce", retain_local_error) + with pytest.raises(RuntimeError, match="EPLB transfer worker failed on this rank") as exc_info: + manager._poll_in_flight() + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert str(exc_info.value.__cause__) == "boom" + assert encoded_statuses == [manager_module.EPLB_CONTROL_ERROR] + + +def test_manager_inflight_remote_worker_error_does_not_commit(monkeypatch): + class Transfer: + def __init__(self): + self.commits = [] + + def pending_layers(self): + return [(0, 0)] + + def commit(self, *args): + self.commits.append(args) + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.transfer = Transfer() + manager.control_group = object() + manager._control_ready_count = torch.empty(1, dtype=torch.int32) + manager.in_flight_layers = [0] + manager._commit_layer_metadata = lambda _layer: None + manager._finish_rebalance = lambda: None + statuses = [] + + def remote_error(tensor, **_kwargs): + statuses.append(int(tensor.item())) + tensor.fill_(manager_module.EPLB_CONTROL_ERROR) + + monkeypatch.setattr(manager_module.dist, "all_reduce", remote_error) + + with pytest.raises(RuntimeError, match="EPLB transfer worker failed on another rank"): + manager._poll_in_flight() + assert statuses == [1] + assert manager.transfer.commits == [] + + +def test_evaluation_ready_gate_propagates_local_and_remote_errors(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager._evaluation_lock = threading.Lock() + manager.control_group = object() + manager._control_ready_count = torch.empty(1, dtype=torch.int32) + manager._evaluation_error = RuntimeError("evaluation boom") + manager._evaluation_result = None + statuses = [] + + def retain_local_error(tensor, **_kwargs): + statuses.append(int(tensor.item())) + + monkeypatch.setattr(manager_module.dist, "all_reduce", retain_local_error) + with pytest.raises(RuntimeError, match="EPLB evaluation failed on this rank") as exc_info: + manager._evaluation_ready_on_all_ranks() + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert str(exc_info.value.__cause__) == "evaluation boom" + assert statuses == [manager_module.EPLB_CONTROL_ERROR] + + manager._evaluation_error = None + manager._evaluation_result = {"kind": "no_improvement"} + statuses.clear() + + def remote_error(tensor, **_kwargs): + statuses.append(int(tensor.item())) + tensor.fill_(manager_module.EPLB_CONTROL_ERROR) + + monkeypatch.setattr(manager_module.dist, "all_reduce", remote_error) + with pytest.raises(RuntimeError, match="EPLB evaluation failed on another rank"): + manager._evaluation_ready_on_all_ranks() + assert statuses == [1] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_manager_inflight_commit_orders_live_weights_between_overlap_forwards( + monkeypatch, +): + class Transfer: + def __init__(self, live, staging): + self.live = live + self.staging = staging + self.pending = [(0, 0)] + + def pending_layers(self): + return self.pending + + def commit(self, layer, buffer_index, post_copy=None): + assert self.pending.pop(0) == (layer, buffer_index) + self.live.copy_(self.staging, non_blocking=True) + if post_copy is not None: + post_copy() + + def finish(self): + pass + + live = torch.tensor([1.0], device="cuda") + staging = torch.tensor([2.0], device="cuda") + previous_read = torch.empty_like(live) + next_read = torch.empty_like(live) + source_stream = torch.cuda.Stream(device=live.device) + destination_stream = torch.cuda.Stream(device=live.device) + initial_stream = torch.cuda.current_stream(device=live.device) + original_overlap_stream = g_infer_context.overlap_stream + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.transfer = Transfer(live, staging) + manager.control_group = object() + manager._control_ready_count = torch.empty(1, dtype=torch.int32) + manager.in_flight_layers = [0] + manager._commit_layer_metadata = lambda _layer: None + manager._finish_rebalance = lambda: None + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda tensor, **_kwargs: tensor.fill_(1)) + + try: + g_infer_context.overlap_stream = source_stream + with torch.cuda.stream(source_stream): + source_stream.wait_stream(initial_stream) + torch.cuda._sleep(20_000_000) + previous_read.copy_(live, non_blocking=True) + with torch.cuda.stream(destination_stream): + manager._poll_in_flight() + with torch.cuda.stream(source_stream): + source_stream.wait_stream(destination_stream) + next_read.copy_(live, non_blocking=True) + source_stream.synchronize() + + assert previous_read.item() == 1.0 + assert next_read.item() == 2.0 + finally: + g_infer_context.overlap_stream = original_overlap_stream + + +def test_transfer_ring_reuses_a_buffer_only_after_commit_and_consumption(monkeypatch): + operations = [] + + class Event: + def record(self, stream): + operations.append("consumed record") + + def synchronize(self): + operations.append("consumed synchronize") + + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer.backend, transfer.device, transfer.global_rank = "test", "cuda:0", 0 + transfer.staging = [[]] + transfer._release = [threading.Event()] + transfer._release[0].set() + transfer._consumed_events, transfer._consumed_recorded = [Event()], [False] + transfer._pending, transfer._pending_lock = deque(), threading.Lock() + transfer._error, transfer._thread = None, None + transfer._needs_staging_reuse_barrier, transfer.transfer_group = True, object() + transfer._copy_part = lambda part: operations.append(("copy", part[0])) + transfer._prepare_part = lambda layer, plan, part: (layer, part) + transfer._start_transfer_generation = lambda: None + transfer._finish_transfer_generation = lambda: None + transfer._commit_full_parts = lambda *args: None + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(transfer_module.torch.cuda, "current_stream", lambda: object()) + monkeypatch.setattr(transfer_module.dist, "barrier", lambda **kwargs: operations.append("barrier")) + plans = [(0, []), (1, []), (2, [])] + prepared = transfer.prepare_transfer(plans) + monkeypatch.setattr(transfer, "prepare_transfer", lambda plans: pytest.fail("must reuse prepared layers")) + transfer.start(plans, prepared) + for layer in range(3): + deadline = time.monotonic() + 2 + while not transfer.pending_layers() and time.monotonic() < deadline: + time.sleep(0.001) + assert transfer.pending_layers() == [(layer, 0)] + assert operations[-1] == ("copy", layer) + transfer.commit(layer, 0) + transfer.finish() + assert operations == [ + ("copy", 0), + "consumed record", + "consumed synchronize", + "barrier", + ("copy", 1), + "consumed record", + "consumed synchronize", + "barrier", + ("copy", 2), + "consumed record", + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_complete_layer_commit_queues_weights_and_mapping_without_waiting_on_host(monkeypatch): + live = torch.ones((2, 8), device="cuda") + staging = torch.full_like(live, 2) + mapping = torch.ones(1, device="cuda") + old_read = torch.empty_like(live) + new_read = torch.empty_like(live) + stream = torch.cuda.Stream() + current = torch.cuda.current_stream() + original_overlap = g_infer_context.overlap_stream + consumed = torch.cuda.Event() + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer.live, transfer.staging = [[("weight", live)]], [[("weight", staging)]] + transfer._prepared_layer = SimpleNamespace(layer_index=0, parts=(None,), changed_dst_slots=(0,)) + transfer._pending, transfer._pending_lock = deque([(0, 0)]), threading.Lock() + transfer._release = [threading.Event()] + transfer._consumed_events, transfer._consumed_recorded = [consumed], [False] + transfer._error = None + manager = object.__new__(manager_module.EPLBManager) + manager.transfer, manager.in_flight_layers = transfer, [0] + manager.control_group = object() + manager._control_ready_count = torch.empty(1, dtype=torch.int32) + manager._commit_layer_metadata = lambda _layer: mapping.fill_(2) + transfer.transfer_group = object() + + def reduce_ready(tensor, **kwargs): + if kwargs["group"] is manager.control_group: + tensor.fill_(1) + + monkeypatch.setattr(manager_module.dist, "all_reduce", reduce_ready) + try: + g_infer_context.overlap_stream = stream + with torch.cuda.stream(stream): + stream.wait_stream(current) + torch.cuda._sleep(100_000_000) + old_read.copy_(live) + manager._poll_in_flight(finish_generation=False) + assert transfer._consumed_recorded == [True] + assert not consumed.query(), "commit waited for the preceding GPU work on the host" + with torch.cuda.stream(stream): + stream.wait_stream(current) + new_read.copy_(live * mapping) + stream.synchronize() + torch.testing.assert_close(old_read, torch.ones_like(live)) + torch.testing.assert_close(new_read[0], torch.full_like(new_read[0], 4)) + torch.testing.assert_close(new_read[1], torch.full_like(new_read[1], 2)) + finally: + g_infer_context.overlap_stream = original_overlap + + +def test_transfer_finalizes_after_join_and_propagates_finalization_failure(monkeypatch): + def make_transfer(finalize): + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer.backend = "test" + transfer.device = torch.device("cuda", 0) + transfer.global_rank = 0 + transfer.staging_depth = 1 + transfer.staging = [[]] + transfer._release = [threading.Event()] + transfer._release[0].set() + transfer._consumed_events = [object()] + transfer._consumed_recorded = [False] + transfer._pending = deque() + transfer._pending_lock = threading.Lock() + transfer._error = None + transfer._thread = None + transfer._needs_staging_reuse_barrier = False + transfer._copy_part = lambda part: None + transfer._prepare_part = lambda layer, plan, part: None + transfer._start_transfer_generation = lambda: None + transfer._finish_transfer_generation = finalize + return transfer + + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) + + finalized_before_publish = [] + success = make_transfer(lambda: finalized_before_publish.append(len(success._pending))) + success.start([(0, [])], success.prepare_transfer([(0, [])])) + success.finish() + assert finalized_before_publish == [1] + assert success.pending_layers() == [(0, 0)] + + failed = make_transfer(lambda: (_ for _ in ()).throw(RuntimeError("cache boom"))) + failed.start([(0, [])], failed.prepare_transfer([(0, [])])) + failed._thread.join() + assert list(failed._pending) == [(0, 0)] + with pytest.raises(RuntimeError, match="cache boom"): + failed.finish() + assert failed._thread is None + + +def test_manager_rearms_after_rebalance_for_interval_one(): + recording_calls = [] + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.step_interval = 1 + manager.sampling_interval = 1 + manager._last_evaluation_step = 1 + manager.weights = [] + manager._eplb_states = [] + manager.target_placement = torch.zeros(1) + manager.in_flight_started_at = 0 + manager.global_rank = 1 + manager._set_recording = lambda enabled: recording_calls.append(enabled) + manager._finish_rebalance() + assert manager.in_flight is False + assert recording_calls == [True] + assert manager._next_evaluation_step == 2 + + +def test_planned_rebalance_resets_sampling_interval_to_base(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.current_placement = torch.zeros((1, 1, 1), dtype=torch.int64) + manager.num_logical_experts = 1 + manager.world_size = 1 + manager.node_world_size = 1 + manager.step_interval = 20 + manager.sampling_interval = 320 + manager.global_rank = 1 + manager.transfer = type( + "Transfer", + (), + {"start": lambda self, plans, prepared_batches: setattr(self, "started", (plans, prepared_batches))}, + )() + manager._reset_recorded_samples = lambda: None + + prepared_batches = [object()] + manager._start_rebalance( + { + "placement": torch.zeros((1, 1, 1), dtype=torch.int64), + "improved": torch.tensor([True]), + "metadata": [None], + "layer_plans": [(0, object())], + "prepared_batches": prepared_batches, + "before": {"max": 1.0, "p95": 1.0}, + "after": {"max": 1.0, "p95": 1.0}, + "model_imbalance_ratio": 1.0, + "candidate_model_imbalance_ratio": 1.0, + "candidate_rebalance_gain": 0.1, + "candidate_changed_layer_count": 1, + } + ) + + assert manager.sampling_interval == 20 + assert manager.in_flight + assert manager._initial_plan_complete + assert not manager._use_full_window + assert len(manager.transfer.started[0]) == 1 + assert manager.transfer.started[1] is prepared_batches + + +def test_transfer_start_rejects_prepared_layers_with_wrong_count(): + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer._thread = None + transfer.staging_depth = 2 + transfer.staging = [object(), object()] + + with pytest.raises(ValueError, match="prepared layer count"): + transfer.start([(0, []), (1, []), (2, [])], prepared_layers=[object()]) + + +def test_manager_poll_advances_inflight_before_evaluation(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = True + manager.evaluation_in_flight = True + calls = [] + manager._poll_in_flight = lambda: calls.append("inflight") + manager._poll_evaluation = lambda: calls.append("evaluation") + + manager.poll() + + assert calls == ["inflight"] + + +def test_manager_poll_waits_for_all_evaluation_results(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager.evaluation_in_flight = True + manager._evaluation_lock = threading.Lock() + manager._evaluation_result = None + manager._evaluation_error = None + manager.control_group = object() + manager._control_ready_count = torch.empty(1, dtype=torch.int32) + calls = [] + manager._poll_evaluation = lambda: calls.append("evaluation") + + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda tensor, **_kwargs: tensor.fill_(0)) + manager.poll() + assert calls == [] + + manager._evaluation_result = {"kind": "no_improvement"} + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda tensor, **_kwargs: tensor.fill_(1)) + manager.poll() + assert calls == ["evaluation"] + + +def test_nixl_prepare_part_compiles_hot_path_without_tensor_views(monkeypatch): + class BatchMemcpy: + def __init__(self): + self.prepared = [] + self.enqueued = [] + + def prepare(self, descriptors): + descriptor = tuple(descriptors) + self.prepared.append(descriptor) + return descriptor + + def enqueue(self, descriptor, stream): + self.enqueued.append((descriptor, stream)) + + stream = SimpleNamespace(cuda_stream=123, synchronize=lambda: None) + batch_memcpy = BatchMemcpy() + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer._push_stream = stream + transfer._batch_memcpy = batch_memcpy + transfer.global_rank = 0 + transfer._same_node_ranks = {0, 1, 2} + transfer._live_row_layout = [ + [ + ("w13.weight", 1000, 32), + ("w13.weight_scale", 2000, 32), + ("w13.weight_zero_point", 3000, 32), + ] + ] + transfer._push_staging_row_layout = { + 1: [[("w13.weight", 4000, 32), ("w13.weight_scale", 5000, 32), ("w13.weight_zero_point", 6000, 32)]], + 2: [[("w13.weight", 7000, 32), ("w13.weight_scale", 8000, 32), ("w13.weight_zero_point", 9000, 32)]], + } + transfer._get_remote_read = lambda *_args: None + transfer._wait_xfers = lambda _xfers: None + + run_a = [TransferStep(1, 3, 0, 5), TransferStep(1, 4, 0, 6)] + run_b = [TransferStep(2, 1, 0, 2)] + local_inbound = TransferStep(0, 0, 1, 0) + remote_inbound = TransferStep(0, 1, 3, 2) + staging = object() + transfer.staging = [staging] + prepared = transfer._prepare_part(0, run_a + run_b + [local_inbound, remote_inbound], part_index=0) + monkeypatch.setattr(transfer, "_prepare_part", lambda *args: pytest.fail("hot path must not prepare descriptors")) + monkeypatch.setattr(transfer_module.torch.cuda, "stream", lambda _stream: pytest.fail("must not switch streams")) + transfer._copy_part(prepared) + + expected = ( + (1160, 4096, 32), + (2160, 5096, 32), + (3160, 6096, 32), + (1192, 4128, 32), + (2192, 5128, 32), + (3192, 6128, 32), + (1064, 7032, 32), + (2064, 8032, 32), + (3064, 9032, 32), + ) + assert batch_memcpy.prepared == [expected] + assert batch_memcpy.enqueued == [(expected, 123)] + assert prepared.remote_entries == {3: [(0, [remote_inbound], staging)]} + + +def test_nixl_prepare_transfer_builds_complete_layers_before_transfer(): + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer.staging_depth, transfer.staging = 1, ["staging"] + seen = [] + + def prepare_part(layer_index, plan, part_index): + seen.append((layer_index, part_index)) + return f"prepared-{layer_index}-{part_index}" + + transfer.global_rank = 0 + transfer._prepare_part = prepare_part + prepared = transfer.prepare_transfer([(3, []), (4, []), (5, [])]) + assert seen == [(3, 0), (4, 0), (5, 0)] + assert [layer.layer_index for layer in prepared] == [3, 4, 5] + assert [layer.parts for layer in prepared] == [ + ("prepared-3-0",), + ("prepared-4-0",), + ("prepared-5-0",), + ] + + +def test_cuda_batch_memcpy_cuda13_abi_and_descriptor_layout(): + class Function: + def __init__(self, callback): + self.callback = callback + self.restype = None + self.argtypes = None + + def __call__(self, *args): + return self.callback(*args) + + class Library: + def __init__(self): + def get_version(pointer): + ctypes.cast(pointer, ctypes.POINTER(ctypes.c_int))[0] = 13000 + return 0 + + self.cudaRuntimeGetVersion = Function(get_version) + self.cudaMemcpyBatchAsync = Function(lambda *args: self.calls.append(args) or 0) + self.cudaGetErrorString = Function(lambda _result: b"fake cuda error") + self.calls = [] + + import ctypes + + library = Library() + batch_memcpy = _CudaBatchMemcpy(library) + prepared = batch_memcpy.prepare(((101, 201, 64), (102, 202, 128))) + batch_memcpy.enqueue(prepared, 777) + + assert len(library.cudaMemcpyBatchAsync.argtypes) == 8 + assert library.calls[0][3] == 2 + assert library.calls[0][6] == 1 + assert library.calls[0][7].value == 777 + assert [pointer for pointer in library.calls[0][0]] == [201, 202] + assert [pointer for pointer in library.calls[0][1]] == [101, 102] + assert list(library.calls[0][2]) == [64, 128] + attrs = library.calls[0][4]._obj + assert attrs.srcAccessOrder == 1 + assert attrs.srcLocHint.type == attrs.srcLocHint.id == 0 + assert attrs.dstLocHint.type == attrs.dstLocHint.id == 0 + assert attrs.flags == 1 + + +def test_cuda_batch_memcpy_rejects_unsupported_runtime_and_invalid_descriptors(): + class Function: + def __init__(self, callback): + self.callback = callback + self.restype = None + self.argtypes = None + + def __call__(self, *args): + return self.callback(*args) + + class OldRuntimeLibrary: + def __init__(self): + def get_version(pointer): + ctypes.cast(pointer, ctypes.POINTER(ctypes.c_int))[0] = 12080 + return 0 + + self.cudaRuntimeGetVersion = Function(get_version) + self.cudaMemcpyBatchAsync = Function(lambda *_args: 0) + self.cudaGetErrorString = Function(lambda _result: b"fake cuda error") + + import ctypes + + with pytest.raises(RuntimeError, match="13.0"): + _CudaBatchMemcpy(OldRuntimeLibrary()) + + class FutureRuntimeLibrary: + def __init__(self): + def get_version(pointer): + ctypes.cast(pointer, ctypes.POINTER(ctypes.c_int))[0] = 14000 + return 0 + + self.cudaRuntimeGetVersion = Function(get_version) + self.cudaMemcpyBatchAsync = Function(lambda *_args: 0) + self.cudaGetErrorString = Function(lambda _result: b"fake cuda error") + + with pytest.raises(RuntimeError, match="13.x"): + _CudaBatchMemcpy(FutureRuntimeLibrary()) + + class MissingBatchSymbolLibrary: + def __init__(self): + def get_version(pointer): + ctypes.cast(pointer, ctypes.POINTER(ctypes.c_int))[0] = 13000 + return 0 + + self.cudaRuntimeGetVersion = Function(get_version) + self.cudaGetErrorString = Function(lambda _result: b"fake cuda error") + + with pytest.raises(RuntimeError, match="cudaMemcpyBatchAsync"): + _CudaBatchMemcpy(MissingBatchSymbolLibrary()) + with pytest.raises(ValueError, match="at least one"): + _CudaBatchMemcpy.prepare(()) + with pytest.raises(ValueError, match="positive"): + _CudaBatchMemcpy.prepare(((1, 2, 0),)) + + +def test_nixl_transfer_fails_fast_without_cuda13_batch_memcpy(monkeypatch): + failure = RuntimeError("missing cudaMemcpyBatchAsync") + + def unavailable(): + raise failure + + monkeypatch.setattr( + transfer_module.NixlEPLBTransfer, "_init_staging", lambda self, *_args: setattr(self, "device", "mock") + ) + monkeypatch.setattr(transfer_module.torch.cuda, "Stream", lambda device: ("stream", device)) + monkeypatch.setattr(transfer_module, "_CudaBatchMemcpy", unavailable) + weight = SimpleNamespace(expert_parallel_state=_test_parallel_state(eplb=True)) + with pytest.raises(RuntimeError, match="missing cudaMemcpyBatchAsync") as exc_info: + transfer_module.NixlEPLBTransfer([weight], object(), 0, 1) + assert exc_info.value is failure + + +@pytest.mark.parametrize( + ("maps", "open_raises", "expected"), + [ + ( + "7f /cuda-12/libcudart.so.12.8\n" + "7f /first/libcudart.so.13 (deleted)\n" + "7f /second/libcudart.so.13\n" + "7f /cuda-14/libcudart.so.14\n" + "7f /cuda-130/libcudart.so.130\n", + False, + "/first/libcudart.so.13", + ), + ("7f /cuda-12/libcudart.so.12\n7f /cuda-130/libcudart.so.130\n", False, None), + ("", True, None), + ], +) +def test_cuda_batch_memcpy_finds_first_loaded_cuda13_runtime(monkeypatch, maps, open_raises, expected): + def open_maps(_path): + if open_raises: + raise OSError("maps unavailable") + return io.StringIO(maps) + + monkeypatch.setattr(builtins, "open", open_maps) + assert _CudaBatchMemcpy._find_loaded_cudart() == expected + + +@pytest.mark.parametrize(("layers", "expected_depth"), [(1, 1), (8, 1), (9, 1), (43, 1)]) +def test_nixl_transfer_bounds_staging_depth(monkeypatch, layers, expected_depth): + def base_init(self, *_args): + self.device = "mock-device" + + monkeypatch.setattr(transfer_module.NixlEPLBTransfer, "_init_staging", base_init) + monkeypatch.setattr(transfer_module.torch.cuda, "Stream", lambda device: ("stream", device)) + monkeypatch.setattr(transfer_module, "_CudaBatchMemcpy", lambda: object()) + monkeypatch.setattr(transfer_module.NixlEPLBTransfer, "_init_ipc_metadata", lambda self: None) + monkeypatch.setattr(transfer_module.NixlEPLBTransfer, "_init_push_layouts", lambda self: None) + + weight = SimpleNamespace(expert_parallel_state=_test_parallel_state(eplb=True)) + transfer = transfer_module.NixlEPLBTransfer([weight] * layers, object(), 0, 1) + + assert transfer.staging_depth == expected_depth + + +def test_nixl_remote_read_cache_reuses_exact_batch_key_and_releases_on_shutdown(): + class Tensor: + nbytes = 8 + + def __init__(self, pointer, nbytes=8): + self.pointer = pointer + self.nbytes = nbytes + + def data_ptr(self): + return self.pointer + + def get_device(self): + return 0 + + def __getitem__(self, index): + return Tensor(self.pointer + index * self.nbytes, self.nbytes) + + class Agent: + def __init__(self): + self.prepared = 0 + self.made = 0 + self.released_xfers = 0 + self.released_dlists = 0 + self.removed_agents = [] + + def get_xfer_descs(self, descriptors, _): + return descriptors + + def prep_xfer_dlist(self, *_args, **_kwargs): + self.prepared += 1 + return f"dlist-{self.prepared}" + + def make_prepped_xfer(self, *_args, **_kwargs): + self.made += 1 + return f"xfer-{self.made}" + + def query_xfer_backend(self, _): + return "UCX" + + def release_xfer_handle(self, _): + self.released_xfers += 1 + + def release_dlist_handle(self, _): + self.released_dlists += 1 + + def remove_remote_agent(self, remote_name): + self.removed_agents.append(remote_name) + + agent = Agent() + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer._nixl_agent = agent + transfer._xfer_cache = {} + transfer._used_xfer_cache_keys = set() + transfer._remote_agents = {1: "remote-1"} + transfer._remote_layouts = {1: [[("weight", 1000, 0, 8)]]} + transfer._registered_descs = None + transfer.live = [[("weight", Tensor(100))]] + staging = [("weight", Tensor(200, 4))] + first_entries = [(0, [TransferStep(0, 0, 1, 0)], staging)] + + first = transfer._get_remote_read(1, first_entries, 0) + assert transfer._get_remote_read(1, first_entries, 0) is first + assert agent.made == 1 + + changed_entries = [(0, [TransferStep(0, 1, 1, 0)], staging)] + transfer._get_remote_read(1, changed_entries, 0) + assert agent.made == 2 + + # Cache identity must include both exact address ranges. + transfer._remote_layouts[1][0][0] = ("weight", 1008, 0, 8) + transfer._get_remote_read(1, first_entries, 0) + assert agent.made == 3 + staging[0][1].pointer += 16 + transfer._get_remote_read(1, first_entries, 0) + assert agent.made == 4 + + transfer.shutdown() + assert agent.released_xfers == 4 + assert agent.released_dlists == 8 + assert agent.removed_agents == ["remote-1"] + + +def test_nixl_remote_read_cache_is_bounded_to_the_current_transfer_generation( + monkeypatch, +): + class Tensor: + nbytes = 8 + + def __init__(self, pointer): + self.pointer = pointer + + def data_ptr(self): + return self.pointer + + def get_device(self): + return 0 + + def __getitem__(self, index): + return Tensor(self.pointer + index * self.nbytes) + + class Agent: + def __init__(self): + self.made = 0 + self.released_xfers = 0 + self.released_dlists = 0 + + def get_xfer_descs(self, descriptors, _): + return descriptors + + def prep_xfer_dlist(self, *_args, **_kwargs): + return object() + + def make_prepped_xfer(self, *_args, **_kwargs): + self.made += 1 + return object() + + def query_xfer_backend(self, _): + return "UCX" + + def release_xfer_handle(self, _): + self.released_xfers += 1 + + def release_dlist_handle(self, _): + self.released_dlists += 1 + + def remove_remote_agent(self, _): + pass + + agent = Agent() + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer._nixl_agent = agent + transfer._xfer_cache = {} + transfer._used_xfer_cache_keys = set() + transfer._remote_agents = {1: "remote-1"} + transfer._remote_layouts = {1: [[("weight", 1000, 0, 8)]]} + transfer._registered_descs = None + transfer._ipc_staging = {} + transfer.live = [[("weight", Tensor(100))]] + transfer.device = torch.device("cuda", 0) + transfer.transfer_group = object() + transfer.global_rank = 0 + transfer.world_size = 2 + transfer.staging_depth = 1 + transfer.staging = [[]] + transfer.num_experts_per_rank = 0 + transfer._release = [threading.Event()] + transfer._release[0].set() + transfer._consumed_events = [object()] + transfer._consumed_recorded = [False] + transfer._pending = deque() + transfer._pending_lock = threading.Lock() + transfer._error = None + transfer._thread = None + transfer._needs_staging_reuse_barrier = False + + staging = [("weight", Tensor(200))] + entries_a = [(0, [TransferStep(0, 0, 1, 0)], staging)] + entries_b = [(0, [TransferStep(0, 1, 1, 0)], staging)] + generation = [entries_a] + + def copy_part(_prepared_part): + transfer._get_remote_read(1, generation[0], 0) + + transfer._copy_part = copy_part + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) + + prepared_batches = [transfer._PreparedLayer(0, (), (None,))] + transfer.start([(0, [])], prepared_batches) + transfer.finish() + assert agent.made == 1 + assert agent.released_xfers == agent.released_dlists == 0 + assert len(transfer._xfer_cache) == 1 + + # The real manager releases each staging buffer through commit(). This + # focused cache test has no commits, so model that hand-off before the + # next generation reuses buffer zero. + transfer._release[0].set() + transfer.start([(0, [])], prepared_batches) + transfer.finish() + assert agent.made == 1 + assert agent.released_xfers == agent.released_dlists == 0 + assert len(transfer._xfer_cache) == 1 + + generation[:] = [entries_b] + transfer._release[0].set() + transfer.start([(0, [])], prepared_batches) + transfer.finish() + assert agent.made == 2 + assert agent.released_xfers == 1 + assert agent.released_dlists == 2 + assert len(transfer._xfer_cache) == 1 + transfer.shutdown() + + +def test_nixl_ipc_metadata_exports_staging_per_local_target(monkeypatch): + from lightllm.server.router.model_infer.mode_backend.pd import p2p_fix + + class Tensor: + shape = (4, 2) + dtype = torch.float16 + device = torch.device("cuda", 0) + nbytes = 16 + + def __init__(self, label): + self.label = label + + def numel(self): + return 3 + + def __getitem__(self, _index): + return self + + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer.transfer_group = object() + transfer.global_rank = 0 + transfer.world_size = 3 + transfer.device = torch.device("cuda", 0) + transfer.live = [[("w13.weight", Tensor("local-w13")), ("w2.weight", Tensor("local-w2"))]] + transfer.staging_depth = 8 + transfer.staging = [[("w13.weight", Tensor(f"local-staging-{index}"))] for index in range(8)] + transfer._ipc_staging = {} + + reduce_calls, rebuild_calls, gathers = [], [], [] + + def reduce_tensor(tensor): + reduce_calls.append(tensor.label) + return None, (f"export-{tensor.label}",) + + def rebuild_tensor(export): + rebuild_calls.append(export) + return Tensor(export) + + source_one = [[("w13.weight", (4, 2), torch.float16, (f"rank1-staging-{index}",))] for index in range(8)] + + def all_gather(output, value, **_kwargs): + gathers.append(value) + if len(gathers) == 1: + output[:] = ["node-a", "node-a", "node-b"] + else: + output[:] = [value, {0: {"staging": source_one}}, {}] + + monkeypatch.setattr(p2p_fix, "reduce_tensor", reduce_tensor) + monkeypatch.setattr(p2p_fix, "p2p_fix_rebuild_cuda_tensor", rebuild_tensor) + monkeypatch.setattr(transfer_module.dist, "all_gather_object", all_gather) + monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) + + transfer._init_ipc_metadata() + + assert reduce_calls == [*(f"local-staging-{index}" for index in range(8))] + assert rebuild_calls == [*(f"rank1-staging-{index}" for index in range(8))] + assert transfer._same_node_ranks == {0, 1} + assert transfer._cross_node_ranks == {2} + assert transfer._needs_staging_reuse_barrier + assert [name for name, _ in transfer._ipc_staging[1][0]] == ["w13.weight"] + + +def test_nixl_copy_part_source_pushes_local_rows_and_keeps_remote_ucx_reads( + monkeypatch, +): + class Stream: + def __init__(self): + self.synchronized = 0 + self.cuda_stream = 123 + + def synchronize(self): + self.synchronized += 1 + + stream = Stream() + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer._push_stream = stream + transfer._batch_memcpy = SimpleNamespace(enqueue=lambda descriptor, stream: enqueued.append((descriptor, stream))) + remote_reads, waited_xfers, enqueued = [], [], [] + transfer._get_remote_read = lambda src_rank, entries, part_index=None: remote_reads.append((src_rank, entries)) or ( + None, + None, + "xfer", + ) + transfer._wait_xfers = waited_xfers.extend + + monkeypatch.setattr(transfer_module.torch.cuda, "stream", lambda _stream: pytest.fail("must not switch streams")) + prepared_push = transfer_module.NixlEPLBTransfer._PreparedPart({}, "push", 0) + transfer._copy_part(prepared_push) + assert enqueued == [("push", stream.cuda_stream)] + + remote_step = TransferStep(0, 1, 2, 0) + prepared_remote = transfer_module.NixlEPLBTransfer._PreparedPart({2: [(0, [remote_step], [])]}, None, 0) + transfer._copy_part(prepared_remote) + assert [rank for rank, _ in remote_reads] == [2] + assert waited_xfers == [(None, None, "xfer")] + + +def test_nixl_copy_part_self_only_rank_pushes_and_synchronizes(monkeypatch): + class Stream: + def __init__(self): + self.synchronized = 0 + self.cuda_stream = 456 + + def synchronize(self): + self.synchronized += 1 + + transfer = object.__new__(transfer_module.NixlEPLBTransfer) + transfer._push_stream = Stream() + enqueued = [] + transfer._batch_memcpy = SimpleNamespace(enqueue=lambda descriptor, stream: enqueued.append((descriptor, stream))) + transfer._wait_xfers = lambda _xfers: None + monkeypatch.setattr(transfer_module.torch.cuda, "stream", lambda _stream: pytest.fail("must not switch streams")) + + prepared = transfer_module.NixlEPLBTransfer._PreparedPart({}, "push", 0) + transfer._copy_part(prepared) + + assert enqueued == [("push", transfer._push_stream.cuda_stream)] + assert transfer._push_stream.synchronized == 1 + + +def test_manager_constructs_nixl_transfer(monkeypatch): + weight = type( + "Weight", + (), + { + "n_routed_experts": 4, + "expert_parallel_state": _test_parallel_state( + eplb=True, + num_logical_experts=4, + world_size=2, + num_redundant_experts_per_rank=2, + route_counter=torch.zeros((2, 4), dtype=torch.int64), + ), + }, + )() + transfer = object() + groups = [object(), object(), object()] + new_group_calls = [] + monkeypatch.setattr(manager_module, "find_ep_moe_weights", lambda model: [weight]) + monkeypatch.setattr(manager_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(manager_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(manager_module, "get_node_world_size", lambda: 2) + monkeypatch.setattr(manager_module, "get_prefill_eplb_step_interval", lambda: 20) + monkeypatch.setattr(manager_module, "get_eplb_rebalance_gain_threshold", lambda: 0.07) + + def new_group(*args, **kwargs): + new_group_calls.append((args, kwargs)) + return groups[len(new_group_calls) - 1] + + monkeypatch.setattr(manager_module.dist, "new_group", new_group) + transfer_calls = [] + monkeypatch.setattr( + manager_module, + "NixlEPLBTransfer", + lambda weights, group, rank, world_size: ( + transfer_calls.append((weights, group, rank, world_size)) or transfer + ), + ) + logs = [] + monkeypatch.setattr(manager_module.logger, "info", lambda message: logs.append(message)) + manager = manager_module.EPLBManager(type("Model", (), {})()) + assert manager.transfer is transfer + assert ( + manager.evaluation_group, + manager.control_group, + manager.transfer_group, + ) == tuple(groups) + assert new_group_calls == [(([0, 1],), {"backend": "gloo"})] * 3 + assert transfer_calls == [([weight], groups[2], 0, 2)] + assert manager.rebalance_gain_threshold == 0.07 + assert "rebalance_gain_threshold=0.0700" in logs[0] + assert manager._next_evaluation_step == manager.step_interval + assert tuple(manager._sample_step_counts) == (0,) + assert weight.expert_parallel_state.eplb.recording + assert manager._eplb_states[0] is weight.expert_parallel_state.eplb + assert not hasattr(weight.expert_parallel_state.eplb, "record_load") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for physical decode routing") +@pytest.mark.parametrize("tokens", [1, 4, 16]) +@pytest.mark.parametrize("capture", [False, True]) +def test_decode_reads_republished_layout_without_writing_counters(tokens, capture): + torch.manual_seed(57) + experts, topk = 8, 2 + logits = torch.randn((tokens, experts), device="cuda") + hidden = torch.empty((tokens, 1), device="cuda") + placement = expand_redundant_placement(build_initial_redundant_expert_ids(experts, 2, 1), experts) + mapping, replicas = build_logical_to_physical_map(placement, experts) + mapping, replicas = mapping.cuda(), replicas.cuda() + counter = torch.full((2, experts, 4), 17, dtype=torch.int64, device="cuda") + draft_mapping, draft_replicas, draft_counter = mapping.clone(), replicas.clone(), counter.clone() + + def run(route_map=mapping, counts=replicas, load_counter=counter): + return grouped_topk.triton_grouped_topk_eplb( + hidden, + logits, + None, + topk, + True, + 1, + 1, + "sigmoid", + route_map, + counts, + load_counter, + sample_index=0, + record_load=False, + use_grouped_topk=True, + use_replicas=False, + ) + + run() + run(draft_mapping, draft_replicas, draft_counter) + if capture: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + weights, physical_ids, _ = run() + draft_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(draft_graph): + draft_weights, draft_ids, _ = run(draft_mapping, draft_replicas, draft_counter) + for layout in (placement, placement.roll(1, dims=0), placement.roll(2, dims=1)): + next_mapping, next_replicas = build_logical_to_physical_map(layout, experts) + mapping.copy_(next_mapping) + replicas.copy_(next_replicas) + if capture: + graph.replay() + else: + weights, physical_ids, _ = run() + logical_ids = layout.flatten().cuda()[physical_ids] + reference_scores, reference_ids = torch.sigmoid(logits).topk(topk, dim=-1) + reference_weights = reference_scores / reference_scores.sum(dim=-1, keepdim=True) + torch.testing.assert_close(physical_ids, next_mapping[:, 0].cuda().long()[reference_ids]) + torch.testing.assert_close(logical_ids, reference_ids) + torch.testing.assert_close(weights, reference_weights) + # An expert-specific output catches stale layouts even if all IDs stay in bounds. + actual = (weights * (logical_ids + 1).square()).sum(dim=-1) + expected = (reference_weights * (reference_ids + 1).square()).sum(dim=-1) + torch.testing.assert_close(actual, expected) + if capture: + draft_graph.replay() + else: + draft_weights, draft_ids, _ = run(draft_mapping, draft_replicas, draft_counter) + draft_logical_ids = placement.flatten().cuda()[draft_ids] + torch.testing.assert_close(draft_logical_ids, reference_ids) + torch.testing.assert_close(draft_weights, reference_weights) + if capture: + graph.replay() + else: + weights, physical_ids, _ = run() + torch.testing.assert_close(layout.flatten().cuda()[physical_ids], reference_ids) + assert torch.all(counter == 17) + assert torch.all(draft_counter == 17) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +@pytest.mark.parametrize("record_load", [False, True]) +@pytest.mark.parametrize("tokens", [1, 32]) +@pytest.mark.parametrize("scoring_func", ["sigmoid", "softmax"]) +@pytest.mark.parametrize("renormalize", [False, True]) +def test_grouped_topk_eplb_matches_topk_mapping_and_counting(record_load, tokens, scoring_func, renormalize): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import ( + triton_grouped_topk, + triton_grouped_topk_eplb, + ) + + torch.manual_seed(1234) + topk = 8 + experts = 256 + num_expert_group = 8 + gating_output = torch.randn((tokens, experts), dtype=torch.bfloat16, device="cuda") + correction_bias = torch.randn((experts,), dtype=torch.float32, device="cuda") + hidden_states = torch.empty((tokens, 1), dtype=torch.bfloat16, device="cuda") + logical_to_physical = torch.stack( + ( + torch.arange(experts, dtype=torch.int32, device="cuda"), + torch.arange(experts, dtype=torch.int32, device="cuda") + experts, + ), + dim=1, + ) + logical_replica_count = torch.where( + torch.arange(experts, device="cuda") % 3 == 0, + torch.full((experts,), 2, dtype=torch.int32, device="cuda"), + torch.ones((experts,), dtype=torch.int32, device="cuda"), + ) + expected_counter = torch.zeros((2, experts), dtype=torch.int64, device="cuda") + fused_counter = torch.zeros_like(expected_counter) + + expected_weights, logical_ids = triton_grouped_topk( + hidden_states, + gating_output, + correction_bias, + topk, + renormalize, + num_expert_group, + 4, + scoring_func, + 2, + ) + if tokens == 1: + replica_indices = torch.zeros_like(logical_ids) + else: + token_indices = torch.arange(tokens, device="cuda", dtype=torch.int64).unsqueeze(1) + replica_indices = _replica_index_reference( + token_indices, logical_ids, logical_replica_count[logical_ids.to(torch.long)].to(torch.int64) + ) + expected_ids = logical_to_physical[logical_ids.to(torch.long), replica_indices.to(torch.long)].to(torch.long) + if record_load: + expected_counter[1].scatter_add_( + 0, + logical_ids.reshape(-1).to(torch.long), + torch.ones(logical_ids.numel(), dtype=torch.int64, device="cuda"), + ) + fused_weights, fused_ids, fused_logical_ids = triton_grouped_topk_eplb( + hidden_states, + gating_output, + correction_bias, + topk, + renormalize, + num_expert_group, + 4, + scoring_func, + logical_to_physical, + logical_replica_count, + fused_counter, + sample_index=1, + record_load=record_load, + use_grouped_topk=True, + group_score_used_topk_num=2, + ) + torch.cuda.synchronize() + + torch.testing.assert_close(fused_weights, expected_weights, rtol=1e-5, atol=1e-6) + assert torch.equal(fused_ids, expected_ids) + assert fused_logical_ids is None + assert torch.equal(fused_counter, expected_counter) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +@pytest.mark.parametrize("experts,groups,group_topk", [(288, 1, 1), (256, 8, 4), (160, 8, 3)]) +@pytest.mark.parametrize("tokens", [1, 32]) +@pytest.mark.parametrize("strided", [False, True]) +@pytest.mark.parametrize("scores", ["equal", "quantized", "saturated"]) +def test_grouped_topk_eplb_matches_logical_selection_at_ties(experts, groups, group_topk, tokens, scores, strided): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import ( + get_group_score_used_topk_num, + triton_grouped_topk, + triton_grouped_topk_eplb, + ) + + generator = torch.Generator(device="cuda").manual_seed(20261022) + logits = torch.randn((tokens, experts), device="cuda", generator=generator).round() + bias = (torch.arange(experts, device="cuda") % 11 / 64).to(torch.bfloat16) + if scores == "equal": + logits.zero_() + bias.zero_() + elif scores == "quantized": + # Exactly representable sigmoid scores isolate tie policy from libm rounding. + logits = (logits > 0).float() * 32 + elif scores == "saturated": + logits.add_(32) + if strided: + logits = torch.stack((logits, torch.empty_like(logits)), dim=-1)[..., 0] + common = dict( + hidden_states=torch.empty((tokens, 1), device="cuda", dtype=torch.bfloat16), + gating_output=logits, + correction_bias=bias, + topk=8, + renormalize=True, + num_expert_group=groups, + topk_group=group_topk, + scoring_func="sigmoid", + group_score_used_topk_num=get_group_score_used_topk_num(group_topk, groups, 8), + ) + expected_weights, expected_ids = triton_grouped_topk(**common) + # Both paths must match the routing formula, including deterministic ties. + reference_scores = logits.sigmoid() + choice_scores = reference_scores + bias + group_values = ( + choice_scores.reshape(tokens, groups, -1).topk(common["group_score_used_topk_num"], dim=-1).values.sum(-1) + ) + threshold = group_values.topk(group_topk, dim=-1).values[:, -1:] + allowed = (group_values >= threshold).repeat_interleave(experts // groups, dim=-1) + reference_ids = choice_scores.masked_fill(~allowed, -float("inf")).argsort(descending=True, stable=True)[:, :8] + torch.testing.assert_close(expected_ids, reference_ids, rtol=0, atol=0) + reference_weights = reference_scores.gather(1, reference_ids) + reference_weights /= reference_weights.sum(-1, keepdim=True) + torch.testing.assert_close(expected_weights, reference_weights, rtol=1e-6, atol=1e-7) + mapping = torch.arange(experts - 1, -1, -1, dtype=torch.int32, device="cuda")[:, None] + counter = torch.zeros((1, experts), dtype=torch.int64, device="cuda") + weights, physical_ids, logical_ids = triton_grouped_topk_eplb( + **common, + logical_to_physical_map=mapping, + logical_replica_count=torch.ones(experts, dtype=torch.int32, device="cuda"), + expert_counter=counter, + sample_index=0, + record_load=True, + use_grouped_topk=True, + return_logical_ids=True, + ) + assert torch.equal(logical_ids, expected_ids) + assert torch.equal(physical_ids, mapping[expected_ids, 0].long()) + torch.testing.assert_close(weights, expected_weights, rtol=1e-6, atol=1e-7) + expected_counts = torch.bincount(expected_ids.flatten(), minlength=experts) + assert torch.equal(counter[0], expected_counts) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +@pytest.mark.parametrize("record_load", [False, True]) +@pytest.mark.parametrize("tokens", [1, 32]) +def test_global_topk_eplb_supports_logical_ids_and_counting(record_load, tokens): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk_eplb + + torch.manual_seed(1234) + topk = 4 + experts = 64 + gating_output = torch.randn((tokens, experts), dtype=torch.float32, device="cuda") + logical_to_physical = torch.stack( + ( + torch.arange(experts, dtype=torch.int32, device="cuda"), + torch.arange(experts, dtype=torch.int32, device="cuda") + experts, + ), + dim=1, + ) + logical_replica_count = torch.where( + torch.arange(experts, device="cuda") % 3 == 0, + torch.full((experts,), 2, dtype=torch.int32, device="cuda"), + torch.ones((experts,), dtype=torch.int32, device="cuda"), + ) + expected_counter = torch.zeros((2, experts), dtype=torch.int64, device="cuda") + fused_counter = torch.zeros_like(expected_counter) + expected_weights, expected_logical_ids = torch.softmax(gating_output, dim=-1).topk(topk, dim=-1) + if tokens == 1: + replica_indices = torch.zeros_like(expected_logical_ids) + else: + token_indices = torch.arange(tokens, device="cuda", dtype=torch.int64).unsqueeze(1) + replica_indices = _replica_index_reference( + token_indices, expected_logical_ids, logical_replica_count[expected_logical_ids].to(torch.int64) + ) + expected_ids = logical_to_physical[expected_logical_ids, replica_indices.to(torch.long)].to(torch.long) + if record_load: + expected_counter[1].scatter_add_( + 0, + expected_logical_ids.reshape(-1), + torch.ones(expected_logical_ids.numel(), dtype=torch.int64, device="cuda"), + ) + expected_weights = expected_weights / expected_weights.sum(dim=-1, keepdim=True) + + weights, physical_ids, logical_ids = triton_grouped_topk_eplb( + hidden_states=torch.empty((tokens, 1), dtype=torch.float32, device="cuda"), + gating_output=gating_output, + correction_bias=torch.randn((experts,), dtype=torch.float32, device="cuda"), + topk=topk, + renormalize=True, + num_expert_group=8, + topk_group=4, + scoring_func="sigmoid", + logical_to_physical_map=logical_to_physical, + logical_replica_count=logical_replica_count, + expert_counter=fused_counter, + sample_index=1, + record_load=record_load, + use_grouped_topk=False, + return_logical_ids=True, + ) + torch.cuda.synchronize() + + torch.testing.assert_close(weights, expected_weights, rtol=1e-5, atol=1e-6) + assert torch.equal(physical_ids, expected_ids) + assert torch.equal(logical_ids, expected_logical_ids) + assert torch.equal(fused_counter, expected_counter) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required for the Triton EPLB kernel") +def test_triton_grouped_topk_eplb_empty_tokens_skips_kernel(): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk_eplb + + experts = 64 + counter = torch.zeros((1, experts), dtype=torch.int64, device="cuda") + weights, physical_ids, logical_ids = triton_grouped_topk_eplb( + hidden_states=torch.empty((0, 1), device="cuda"), + gating_output=torch.empty((0, experts), device="cuda"), + correction_bias=None, + topk=4, + renormalize=False, + num_expert_group=8, + topk_group=4, + scoring_func="softmax", + logical_to_physical_map=torch.zeros((experts, 1), dtype=torch.int32, device="cuda"), + logical_replica_count=torch.ones((experts,), dtype=torch.int32, device="cuda"), + expert_counter=counter, + sample_index=0, + record_load=True, + use_grouped_topk=False, + return_logical_ids=True, + ) + + assert weights.shape == physical_ids.shape == logical_ids.shape == (0, 4) + assert physical_ids.dtype is logical_ids.dtype is torch.long + assert torch.equal(counter, torch.zeros_like(counter)) + + +def test_last_layer_commit_finishes_at_outer_boundary_and_preserves_samples(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + state = SimpleNamespace(recording=True, recorded_sample_count=7) + operations = [] + + class Transfer: + pending = [(0, 0)] + + def pending_layers(self): + return self.pending[:] + + def commit(self, layer, buffer, post_copy): + assert self.pending.pop(0) == (layer, buffer) + post_copy() + + def finish(self): + operations.append("finish") + + manager.transfer = Transfer() + manager.in_flight = True + manager.in_flight_layers = [0] + manager.control_group = object() + manager._control_ready_count = torch.empty(1, dtype=torch.int32) + manager._inflight_timing = dict( + control_poll_calls=0, control_cpu_ms=0.0, commit_calls=0, commit_cpu_ms=0.0, commit_max_ms=0.0 + ) + manager._commit_layer_metadata = lambda layer: operations.append(("metadata", layer)) + + manager._eplb_states = [state] + manager.target_placement = torch.zeros(1) + manager.target_metadata = None + manager.global_rank = 1 + manager.sampling_interval = 4 + manager._last_evaluation_step = 4 + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda *args, **kwargs: None) + monkeypatch.setattr( + manager_module.torch.cuda, "current_stream", lambda: SimpleNamespace(wait_stream=lambda s: None) + ) + monkeypatch.setattr(g_infer_context, "get_overlap_stream", lambda: object()) + manager.poll_migration_at_layer_boundary() + assert operations == [("metadata", 0)] + assert manager.in_flight and manager.in_flight_layers == [] + assert state.recording and state.recorded_sample_count == 7 + manager.poll_migration_at_layer_boundary() + assert operations == [("metadata", 0)] + manager.poll() + assert operations == [("metadata", 0), "finish"] + assert not manager.in_flight and state.recording and state.recorded_sample_count == 7 + assert manager._next_evaluation_step == 8 + + +def test_layer_hook_does_not_start_an_evaluation(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager.evaluation_in_flight = True + manager.poll_migration_at_layer_boundary() + + +def test_migration_metadata_is_committed_before_selecting_experts(monkeypatch): + events = [] + state = EPLBState( + num_redundant_experts_per_rank=1, + logical_to_physical_map=torch.zeros((4, 2), dtype=torch.int32), + logical_replica_count=torch.ones(4, dtype=torch.int32), + route_counter=torch.zeros((4, 4), dtype=torch.int64), + recording=True, + ) + state.migration_poll = lambda: events.append(("migration", state.recorded_sample_count)) + impl = deepgemm_module.FuseMoeDeepGEMM.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.eplb, impl.routed_scaling_factor = state, 1.0 + + def topk(**kwargs): + events.append(("topk", kwargs["sample_index"])) + return torch.ones((2, 2)), torch.zeros((2, 2), dtype=torch.int64), None + + monkeypatch.setattr(grouped_topk, "triton_grouped_topk_eplb", topk) + impl._select_experts( + input_tensor=torch.zeros((2, 4)), + router_logits=torch.zeros((2, 4)), + correction_bias=None, + top_k=2, + renormalize=True, + use_grouped_topk=True, + topk_group=1, + num_expert_group=1, + scoring_func="sigmoid", + is_prefill=True, + ) + assert events == [("migration", 0), ("topk", 0)] + assert state.recorded_sample_count == 1 + + +def test_eplb_expert_scaling_uses_logical_ids_instead_of_physical_rows(monkeypatch): + state = EPLBState( + num_redundant_experts_per_rank=1, + logical_to_physical_map=torch.zeros((4, 2), dtype=torch.int32), + logical_replica_count=torch.ones(4, dtype=torch.int32), + route_counter=torch.zeros((4, 4), dtype=torch.int64), + ) + impl = deepgemm_module.FuseMoeDeepGEMM.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.eplb, impl.routed_scaling_factor = state, 2.5 + + def topk(**kwargs): + assert kwargs["return_logical_ids"] + return torch.ones((1, 2)), torch.tensor([[7, 8]]), torch.tensor([[1, 3]]) + + monkeypatch.setattr(grouped_topk, "triton_grouped_topk_eplb", topk) + weights, physical, logical = impl._select_experts( + input_tensor=torch.zeros((1, 4)), + router_logits=torch.zeros((1, 4)), + correction_bias=None, + top_k=2, + renormalize=True, + use_grouped_topk=True, + topk_group=1, + num_expert_group=1, + scoring_func="sigmoid", + is_prefill=True, + per_expert_scale=torch.tensor([10.0, 20.0, 30.0, 40.0]), + ) + assert weights.tolist() == [[50.0, 100.0]] + assert physical.tolist() == [[7, 8]] and logical.tolist() == [[1, 3]] diff --git a/unit_tests/common/fused_moe/test_eplb_full_layout.py b/unit_tests/common/fused_moe/test_eplb_full_layout.py new file mode 100644 index 0000000000..af7b30cf61 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_full_layout.py @@ -0,0 +1,358 @@ +"""Offline EPLB planning and publication regressions; no model/service startup.""" +import itertools +import math +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import ( + _estimate_rank_load, + _expert_locations, + build_initial_redundant_expert_ids, + build_logical_to_physical_maps_for_layers, + expand_redundant_placement, + select_improving_placements, + validate_physical_placement, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_planner import ( + _propose_moves, + plan_full_experts, + refine_placement_candidates, +) +from lightllm.common.eplb_utils import get_eplb_staging_shape +from lightllm.server.router.model_infer.mode_backend.eplb_transfer import ( + _commit_staging_rows, + align_target_placement, + build_transfer_plan, +) + + +@pytest.fixture(scope="module", autouse=True) +def single_cpu_thread(): + previous = torch.get_num_threads() + torch.set_num_threads(1) + yield + torch.set_num_threads(previous) + + +def _initial(redundant=1, layers=1): + extras = build_initial_redundant_expert_ids(8, 4, redundant)[None].expand(layers, -1, -1).clone() + return expand_redundant_placement(extras, 8) + + +def _independent_rank_load(source_load, physical, alignment, node_size): + # Slow scalar oracle: merge sources at each physical expert BEFORE padding. + rows = physical.tolist() + result = [0.0] * len(rows) + for expert in range(len(source_load[0])): + owners = [rank for rank, row in enumerate(rows) if expert in row] + received = [0.0] * len(rows) + for node, loads in enumerate(source_load): + local = [rank for rank in owners if rank // node_size == node] + selected = local or owners + for rank in selected: + received[rank] += loads[expert] / len(selected) + for rank, count in enumerate(received): + result[rank] += math.ceil(count / alignment) * alignment + return result + + +@pytest.mark.parametrize("redundant", [0, 1, 2]) +@pytest.mark.parametrize("node_size", [1, 2, 4]) +def test_full_maps_and_estimator_match_independent_physical_layout(redundant, node_size): + layout = _initial(redundant, layers=2) + # Move original experts: the map must use absolute rows, never E/W owners. + layout[0] = layout[0].roll(1, dims=0) + layout[1] = layout[1].flip(dims=(0, 1)) + capacity = layout.shape[-1] + for source_rank in range(4): + maps, counts = build_logical_to_physical_maps_for_layers( + layout, 8, source_rank=source_rank, node_world_size=node_size + ) + assert maps.shape == (2, 8, 4) + for layer, expert in itertools.product(range(2), range(8)): + candidates = [index for index, value in enumerate(layout[layer].flatten().tolist()) if value == expert] + local = [index for index in candidates if index // capacity // node_size == source_rank // node_size] + expected = local or candidates + chosen = maps[layer, expert, : counts[layer, expert]].tolist() + assert set(chosen) == set(expected) + assert chosen[0] == expected[source_rank % len(expected)] + assert maps[layer, expert, counts[layer, expert] :].eq(-1).all() + generator = torch.Generator().manual_seed(42) + loads = torch.randint(0, 1024, (3, 2, 4 // node_size, 8), generator=generator) + actual = _estimate_rank_load(loads, layout, 128, node_size) + for sample, layer in itertools.product(range(3), range(2)): + expected = _independent_rank_load(loads[sample, layer].tolist(), layout[layer], 128, node_size) + torch.testing.assert_close(actual[sample, layer], torch.tensor(expected, dtype=torch.float64)) + + +@pytest.mark.parametrize( + "layout,message", + [ + ([[[0, 0], [1, 2]]], "duplicate"), + ([[[0, 1], [1, 2]]], "cover"), + ([[[0, 1], [2, 4]]], "invalid"), + ([[[0, 1], [2, -1]]], "invalid"), + ], +) +def test_full_layout_rejects_unrepresentable_routes(layout, message): + with pytest.raises(ValueError, match=message): + build_logical_to_physical_maps_for_layers(torch.tensor(layout), 4) + + +@pytest.mark.parametrize("node_size", [1, 2, 4]) +def test_shared_replica_rotation_keeps_actual_source_node_selection(node_size): + layout = _initial(2, layers=2) + capacity = layout.shape[-1] + for source_rank, phase in itertools.product(range(4), range(4)): + maps, counts = build_logical_to_physical_maps_for_layers( + layout, 8, source_rank=source_rank, node_world_size=node_size, replica_rotation_rank=phase + ) + for layer, expert in itertools.product(range(2), range(8)): + candidates = [p for p, value in enumerate(layout[layer].flatten().tolist()) if value == expert] + local = [p for p in candidates if p // capacity // node_size == source_rank // node_size] + expected = local or candidates + chosen = maps[layer, expert, : counts[layer, expert]].tolist() + assert chosen == expected[phase % len(expected) :] + expected[: phase % len(expected)] + + +def test_replica_rotation_requires_a_source_and_valid_ep_rank(): + with pytest.raises(ValueError, match="requires source_rank"): + build_logical_to_physical_maps_for_layers(_initial(), 8, replica_rotation_rank=0) + with pytest.raises(ValueError, match="must be an EP rank"): + build_logical_to_physical_maps_for_layers( + _initial(), 8, source_rank=0, node_world_size=4, replica_rotation_rank=4 + ) + + +def test_refinement_matches_small_exhaustive_complete_layout_optimum(): + loads = torch.tensor([512, 128, 4096, 64]).reshape(1, 1, 1, 4) + current = torch.tensor([[[0, 2], [1, 3]]]) + refined = refine_placement_candidates(loads, current, [], expert_alignment=128) + optimum = min( + max(_independent_rank_load(loads[0, 0].tolist(), torch.tensor(rows).reshape(2, 2), 128, 2)) + for rows in itertools.permutations(range(4)) + ) + assert _estimate_rank_load(loads, refined, 128).max() == optimum + assert _estimate_rank_load(loads, current, 128).max() > optimum + assert torch.equal(refined, refine_placement_candidates(loads, current, [], expert_alignment=128)) + + +@pytest.mark.parametrize("redundant", [0, 1, 2]) +@pytest.mark.parametrize("node_size", [1, 2, 4]) +def test_full_planner_preserves_topology_and_never_worsens_samplewise_score(redundant, node_size): + current = _initial(redundant, layers=2) + generator = torch.Generator().manual_seed(1337 + node_size + redundant) + load = torch.randint(0, 10000, (4, 2, 4 // node_size, 8), generator=generator) + planned = plan_full_experts(load, current, expert_alignment=128, node_world_size=node_size) + validate_physical_placement(planned, 8) + assert planned.shape == current.shape + for layer, node in itertools.product(range(2), range(4 // node_size)): + ranks = slice(node * node_size, (node + 1) * node_size) + assert set(current[layer, ranks].flatten().tolist()) == set(planned[layer, ranks].flatten().tolist()) + old = _estimate_rank_load(load, current, 128, node_size).max(dim=-1).values.sum(dim=0) + new = _estimate_rank_load(load, planned, 128, node_size).max(dim=-1).values.sum(dim=0) + assert torch.all(new <= old) + assert torch.equal(current, _initial(redundant, layers=2)) + + +def test_uniform_workload_keeps_current_layout_and_zero_replicas_can_rebalance(): + current = _initial(0) + uniform = torch.full((2, 1, 1, 8), 12800) + assert torch.equal(plan_full_experts(uniform, current, expert_alignment=128), current) + skewed = torch.tensor([12800, 12800, 128, 128, 12800, 12800, 128, 128]).reshape(1, 1, 1, 8) + candidate = plan_full_experts(skewed, current, expert_alignment=128) + placement, improved, metrics, before, after = select_improving_placements( + skewed, current, candidate, rebalance_gain_threshold=0.05, expert_alignment=128 + ) + assert improved.all() and metrics["candidate_rebalance_gain"] > 0.4 + assert after.max() < before.max() + assert placement.numel() == 8 + + +def test_refinement_keeps_each_best_seed_layer_and_respects_zero_budget(): + current = _initial(1, layers=2) + load = torch.randint(1, 10000, (4, 2, 1, 8), generator=torch.Generator().manual_seed(12)) + greedy = plan_full_experts(load, current, expert_alignment=128) + result = refine_placement_candidates(load, current, [greedy], expert_alignment=128, search_steps=0) + cost = lambda placement: _estimate_rank_load(load, placement, 128).max(dim=-1).values.sum(dim=0) + torch.testing.assert_close(cost(result), torch.minimum(cost(current), cost(greedy))) + with pytest.raises(ValueError, match="budget"): + refine_placement_candidates(load, current, [greedy], candidate_budget=0) + + +def test_full_candidate_cannot_change_node_holdings(): + current = _initial(0) + with pytest.raises(ValueError, match="node's expert set"): + refine_placement_candidates(torch.ones((1, 1, 2, 8)), current, [current.roll(1, dims=1)], node_world_size=2) + + +def test_absolute_row_transfer_handles_overwrite_cycle_and_second_generation(): + current = torch.arange(6).reshape(3, 2) + live = [current[rank, :, None].to(torch.float32).repeat(1, 5) for rank in range(3)] + scales = [row[:, :1].clone() + 0.5 for row in live] + for shift in (1, -1): + target = align_target_placement(current, current.roll(shift, dims=0)) + plan = build_transfer_plan(current, target, 6, 3, 3) + assert len(plan) == 6 # all primary rows move, even with R=0 + staged_weights = [torch.empty_like(row) for row in live] + staged_scales = [torch.empty_like(row) for row in scales] + for step in plan: + staged_weights[step.dst_rank][step.dst_slot].copy_(live[step.src_rank][step.src_local_row]) + staged_scales[step.dst_rank][step.dst_slot].copy_(scales[step.src_rank][step.src_local_row]) + for rank in range(3): + slots = [step.dst_slot for step in plan if step.dst_rank == rank] + _commit_staging_rows(live[rank], staged_weights[rank], slots) + _commit_staging_rows(scales[rank], staged_scales[rank], slots) + torch.testing.assert_close(live[rank][:, 0], target[rank].float()) + torch.testing.assert_close(scales[rank][:, 0], target[rank].float() + 0.5) + current = target + + +def test_full_transfer_retains_slots_and_uses_sources_from_current_layout(): + current = torch.tensor([[2, 1, 4, 5], [3, 4, 0, 5]]) + target = torch.tensor([[4, 1, 0, 5], [2, 3, 4, 5]]) + aligned = align_target_placement(current, target) + plan = build_transfer_plan(current, target, 6, 2, 2) + assert len(plan) == 2 + for step in plan: + assert current[step.src_rank, step.src_local_row] == aligned[step.dst_rank, step.dst_slot] + + +@pytest.mark.parametrize("redundant,expected", [(0, (1, 32)), (2, (1, 34))]) +def test_staging_capacity_is_shared_with_memory_profiling(redundant, expected): + state = SimpleNamespace(num_redundant_experts_per_rank=redundant) + weight = SimpleNamespace(expert_parallel_state=SimpleNamespace(eplb=state, num_primary_experts_per_rank=32)) + assert get_eplb_staging_shape([weight] * 40) == expected + + +def test_manager_full_plan_is_canonical_and_publication_tracks_partial_generations(monkeypatch): + from lightllm.server.router.model_infer.mode_backend import eplb_manager as module + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import EPLBState + + manager = module.EPLBManager.__new__(module.EPLBManager) + manager.global_rank, manager.world_size, manager.node_world_size = 0, 4, 4 + manager.num_logical_experts, manager.num_redundant_experts_per_rank = 8, 0 + manager.current_placement = _initial(0, layers=2) + original = manager.current_placement.clone() + manager._evaluation_snapshot = {"placement": original} + manager._initial_plan_complete = False + manager.placement_stickiness, manager.rebalance_gain_threshold = 0.1, 0.05 + manager.evaluation_group = object() + broadcasts = [] + monkeypatch.setattr(module.dist, "broadcast_object_list", lambda result, **kwargs: broadcasts.append(result[0])) + loads = torch.tensor([12800, 12800, 128, 128, 12800, 12800, 128, 128]).reshape(1, 1, 1, 8).expand(1, 2, 1, 8) + result = manager._plan_and_broadcast(loads) + assert result["kind"] == "planned" and broadcasts[0] is result + assert torch.equal(manager.current_placement, original) + for current, target in zip(original, result["placement"]): + assert torch.equal(target, align_target_placement(current, target)) + maps, counts = build_logical_to_physical_maps_for_layers(original, 8) + manager._eplb_states = [ + EPLBState( + 0, + maps[layer].clone(), + counts[layer].clone(), + torch.zeros((1, 8)), + physical_to_logical=original[layer].clone(), + ) + for layer in range(2) + ] + targets, target_counts = build_logical_to_physical_maps_for_layers(result["placement"], 8) + manager.target_placement = result["placement"] + manager.target_metadata = list(zip(targets, target_counts)) + pointers = [state.logical_to_physical_map.data_ptr() for state in manager._eplb_states] + manager._commit_layer_metadata(1) + assert [state.placement_generation for state in manager._eplb_states] == [0, 1] + assert torch.equal(manager.current_placement[0], original[0]) + assert torch.equal(manager.current_placement[1], manager.target_placement[1]) + manager._commit_layer_metadata(0) + for layer, state in enumerate(manager._eplb_states): + assert state.placement_generation == 1 + assert state.logical_to_physical_map.data_ptr() == pointers[layer] + torch.testing.assert_close(state.logical_to_physical_map, targets[layer]) + torch.testing.assert_close(state.physical_to_logical, manager.target_placement[layer]) + + +def test_full_checkpoint_partition_uses_absolute_layout_rows(monkeypatch): + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight as module + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( + EPLBState, + ExpertParallelState, + ) + + weight = module.FusedMoeWeight.__new__(module.FusedMoeWeight) + weight.enable_ep_moe = True + weight.global_rank_, weight.layer_num_ = 1, 0 + weight.n_routed_experts, weight.num_fused_shared_experts, weight.moe_intermediate_size = 8, 0, 32 + weight.quant_method = SimpleNamespace(method_name="test") + physical = _initial(1)[0].roll(1, dims=0).flip(dims=(1,)) + state = EPLBState(1, None, None, None, physical_to_logical=physical) + weight.expert_parallel_state = ExpertParallelState(8, 4, state) + monkeypatch.setattr(module, "get_row_slice_mixin", lambda *a, **kw: None) + monkeypatch.setattr(module, "get_col_slice_mixin", lambda *a, **kw: None) + weight._init_weight_partition() + assert weight.local_n_routed_experts == 3 + assert weight.expert_idx_to_local_idx == {expert: slot for slot, expert in enumerate(physical[1].tolist())} + assert not hasattr(weight, "_initial_redundant_expert_idx_to_local_idx") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("redundant", [0, 1]) +def test_full_state_initialization_including_zero_replicas(monkeypatch, redundant): + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight as module + + from lightllm.server.core.objs.start_args_type import StartArgs + + args = StartArgs( + enable_ep_moe=True, + enable_prefill_eplb=True, + run_mode="prefill", + eplb_num_redundant_experts_per_rank=redundant, + ) + weight = module.FusedMoeWeight.__new__(module.FusedMoeWeight) + weight.n_routed_experts, weight.global_world_size, weight.global_rank_ = 8, 4, 1 + weight.enable_ep_moe = True + weight.tp_rank_, weight.tp_world_size_ = 0, 1 + from lightllm.utils import envs_utils, dist_utils + + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr(dist_utils, "get_node_world_size", lambda: 4) + monkeypatch.setattr(envs_utils, "get_prefill_eplb_step_interval", lambda: 4) + weight._init_expert_parallel_state({}) + state = weight.expert_parallel_state.eplb + assert state.placement_generation == 0 + assert state.route_counter.shape == (8, 8, 2 if redundant == 0 else 11) + assert state.physical_to_logical.shape == (4, 2 + redundant) + assert weight.expert_parallel_state.num_total_physical_experts == 8 + 4 * redundant + for expert, count in enumerate(state.logical_replica_count.tolist()): + for index in state.logical_to_physical_map[expert, :count].tolist(): + assert state.physical_to_logical.flatten()[index] == expert + + +def test_propose_moves_preserves_interleaving_and_budget_truncation(): + rows = torch.tensor([[0, 1], [2, 3]]) + locations = _expert_locations(rows[None], 4)[0] + contributions = torch.tensor([[[9.0, 1.0], [8.0, 1.0], [1.0, 6.0], [1.0, 5.0]]]) + all_moves = _propose_moves(rows, locations, contributions, 2, 256) + assert len(all_moves) == 8 and all(move[4] >= 0 for move in all_moves) + assert _propose_moves(rows, locations, contributions, 2, 1) == all_moves[:1] + assert _propose_moves(rows, locations, contributions, 2, 7) == all_moves[:7] + assert all_moves[0][0] != all_moves[1][0] + + +def test_propose_moves_full_layout_uses_swaps_when_replacements_are_illegal(): + rows = torch.tensor([[0, 1], [2, 3]]) + locations = _expert_locations(rows[None], 4)[0] + contributions = torch.tensor([[[9.0, 1.0], [8.0, 1.0], [1.0, 6.0], [1.0, 5.0]]]) + moves = _propose_moves(rows, locations, contributions, 2, 256) + assert moves + assert all(move[4] >= 0 for move in moves) + + +def test_propose_moves_returns_empty_when_every_copy_is_already_local(): + rows = torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]]) + locations = _expert_locations(rows[None], 4)[0] + contributions = torch.ones((1, 4, 2)) + assert _propose_moves(rows, locations, contributions, 2, 256) == [] diff --git a/unit_tests/common/fused_moe/test_eplb_hash.py b/unit_tests/common/fused_moe/test_eplb_hash.py new file mode 100644 index 0000000000..9531f575e9 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_hash.py @@ -0,0 +1,35 @@ +import pytest +import torch +import triton +import triton.language as tl + +from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import _eplb_replica_index + + +@triton.jit +def _hash_periodic_tokens(output, STRIDE: tl.constexpr, COPIES: tl.constexpr, COUNT: tl.constexpr): + index = tl.program_id(0) * 256 + tl.arange(0, 256) + replica = _eplb_replica_index(index * STRIDE, tl.full((256,), 7, tl.int32), tl.full((256,), COPIES, tl.int32)) + tl.store(output + index, replica, index < COUNT) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("copies", [2, 3, 4, 8]) +@pytest.mark.parametrize("stride", [1, 2, 4, 8, 256]) +def test_replica_hash_breaks_periodic_token_aliasing(copies, stride): + count = 4096 + result = torch.empty(count, device="cuda", dtype=torch.int32) + _hash_periodic_tokens[(triton.cdiv(count, 256),)](result, stride, copies, count) + histogram = torch.bincount(result.long(), minlength=copies).cpu() + # The old linear hash sends all stride-4 indices to one of four replicas. + # This tests distribution, independently of an implementation-shaped oracle. + expected = count / copies + assert histogram.min() > 0.80 * expected + assert histogram.max() < 1.20 * expected + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_original_periodic_hash_regression(): + result = torch.empty(1024, device="cuda", dtype=torch.int32) + _hash_periodic_tokens[(4,)](result, 4, 4, 1024) + assert torch.bincount(result.long(), minlength=4).tolist() == [266, 258, 244, 256] diff --git a/unit_tests/common/fused_moe/test_eplb_memory_profile.py b/unit_tests/common/fused_moe/test_eplb_memory_profile.py new file mode 100644 index 0000000000..815bccef63 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_memory_profile.py @@ -0,0 +1,421 @@ +"""EPLB buffers reduce the existing mem_fraction before KV profiling.""" +import json +from types import SimpleNamespace +from contextlib import nullcontext + +import pytest +import torch + +from lightllm.common import eplb_utils +from lightllm.common.basemodel import basemodel as base_module +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( + EPLBState, + ExpertParallelState, +) +from lightllm.common.kv_cache_mem_manager import mem_manager as memory_module +from lightllm.common.kv_cache_mem_manager.deepseek2_mem_manager import Deepseek2MemoryManager +from lightllm.common.kv_cache_mem_manager.glm5_next_mem_manager import Glm5NextMemManager +from lightllm.common.kv_cache_mem_manager.deepseek3_2mem_manager import Deepseek3_2MemoryManager +from lightllm.common.kv_cache_mem_manager.fp8_per_token_group_quant_deepseek3_2mem_manager import ( + FP8PerTokenGroupQuantDeepseek3_2MemoryManager, +) +from lightllm.common.kv_cache_mem_manager.fp8_static_per_head_quant_mem_manager import FP8StaticPerHeadQuantMemManager +from lightllm.common.kv_cache_mem_manager.fp8_static_per_tensor_quant_mem_manager import ( + FP8StaticPerTensorQuantMemManager, +) +from lightllm.common.kv_cache_mem_manager.ppl_int4kv_mem_manager import PPLINT4KVMemoryManager +from lightllm.common.kv_cache_mem_manager.ppl_int8kv_mem_manager import PPLINT8KVMemoryManager +from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import ( + Qwen3NextMemManager, + FP8StaticPerHeadQuantQwen3NextMemManager, + FP8StaticPerTensorQuantQwen3NextMemManager, +) +from lightllm.models.deepseek2 import model as deepseek_module +from lightllm.models.deepseek_mtp.model import Deepseek3MTPModel +from lightllm.models.glm5_next import model as glm_module +from lightllm.models.glm5_next_mtp.model import Glm5NextMTPModel +from lightllm.models.llama import model as llama_module +from lightllm.models.mixtral import model as mixtral_module +from lightllm.models.qwen2 import model as qwen_module +from lightllm.models.qwen3_moe.model import Qwen3MOEModel +from lightllm.models.qwen3_moe_mtp.model import Qwen3MOEMTPModel +from lightllm.models.qwen3_5.model import Qwen3_5TpPartModel +from lightllm.models.qwen3_5_mtp.model import Qwen3_5MTPModel +from lightllm.models.qwen3next import model as next_module +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import device_utils, profile_max_tokens + + +@pytest.fixture(autouse=True) +def supported_eplb_device(monkeypatch): + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: False) + + +class KVAllocationReached(Exception): + """Stop after profiling, before allocator, CUDA buffers or KV scale loading.""" + + +@pytest.fixture +def profile_without_allocating(monkeypatch): + budget = SimpleNamespace(free=40 * 1024 ** 3, total=80 * 1024 ** 3) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch.cuda, "get_device_properties", lambda _device: SimpleNamespace(total_memory=budget.total)) + monkeypatch.setattr(profile_max_tokens, "get_current_device_id", lambda: 0) + monkeypatch.setattr(memory_module.dist, "is_initialized", lambda: False) + monkeypatch.setattr(memory_module.dist, "get_world_size", lambda: 1) + monkeypatch.setattr(memory_module, "get_available_gpu_memory", lambda _world: budget.free / 1024 ** 3) + monkeypatch.setattr(memory_module, "get_total_gpu_memory", lambda: budget.total / 1024 ** 3) + monkeypatch.setattr(memory_module, "get_env_start_args", lambda: SimpleNamespace(page_size=1)) + monkeypatch.setattr(memory_module.MemoryManager, "get_pd_kv_move_buffer_size", lambda self: 0) + + def stop_allocation(size): + raise KVAllocationReached(size) + + monkeypatch.setattr(memory_module, "KvCacheAllocator", stop_allocation) + return budget + + +def make_expert(layer=0, redundant=0): + rows = 64 + redundant + expert = FusedMoeWeight.__new__(FusedMoeWeight) + expert.enable_ep_moe = True + expert.layer_num_ = layer + expert.expert_parallel_state = ExpertParallelState( + 512, + 8, + EPLBState( + redundant, + None, + None, + torch.empty((40, 512), dtype=torch.int64, device="meta"), + ), + ) + # Eight layers fit in an 80 GiB device; no actual large allocation is made. + expert.w13 = SimpleNamespace(weight=torch.empty((rows, 6144, 4096), dtype=torch.bfloat16, device="meta")) + expert.w2 = SimpleNamespace(weight=torch.empty((rows, 4096, 3072), dtype=torch.bfloat16, device="meta")) + return expert + + +def make_model(model_class=Qwen3MOEModel, redundant=0): + model = model_class.__new__(model_class) + model.args = StartArgs(enable_ep_moe=True, enable_prefill_eplb=True, run_mode="prefill", mtp_mode=None) + model.tp_world_size_ = 1 # DP8: the attention TP size is 1, while expert world size is 8. + model.data_type = torch.bfloat16 + model.max_total_token_num = None + model.mem_fraction = 0.98 + model.num_kv_heads = 8 + model.config = dict( + hidden_size=4096, + n_embed=4096, + num_attention_heads=32, + num_key_value_heads=8, + num_hidden_layers=8, + n_layer=8, + head_dim=128, + kv_lora_rank=512, + qk_rope_head_dim=64, + full_attention_interval=4, + linear_num_key_heads=8, + linear_num_value_heads=8, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_conv_kernel_dim=4, + ) + model.trans_layers_weight = [SimpleNamespace(experts=make_expert(layer, redundant)) for layer in range(8)] + return model + + +@pytest.mark.parametrize("redundant,staging", [(0, 4608), (2, 4752)]) +@pytest.mark.parametrize( + "model_class,module,manager_class", + [ + (base_module.TpPartBaseModel, base_module, memory_module.MemoryManager), + (llama_module.LlamaTpPartModel, llama_module, memory_module.MemoryManager), + (Qwen3MOEModel, qwen_module, memory_module.MemoryManager), + (deepseek_module.Deepseek2TpPartModel, deepseek_module, Deepseek2MemoryManager), + (deepseek_module.Deepseek2TpPartModel, deepseek_module, Deepseek3_2MemoryManager), + (mixtral_module.MixtralTpPartModel, mixtral_module, memory_module.MemoryManager), + (next_module.Qwen3NextTpPartModel, next_module, Qwen3NextMemManager), + (glm_module.Glm5NextTpPartModel, glm_module, Glm5NextMemManager), + ], + ids=["base", "llama", "qwen3", "deepseek2", "deepseek3_2", "mixtral", "qwen3next", "glm5next"], +) +def test_model_profile_reserves_eplb_before_kv_alloc( + monkeypatch, profile_without_allocating, model_class, module, manager_class, redundant, staging +): + model = make_model(model_class, redundant) + manager = manager_class.__new__(manager_class) + + def construct(*args, **kwargs): + manager_class.__init__(manager, *args, **kwargs) + + if module is mixtral_module: + monkeypatch.setattr(module, "MemoryManager", construct) + elif module is glm_module: + model.linear_config = SimpleNamespace( + full_att_head_dim=644, get_full_att_kv_layer_num_with_draft_model=lambda: 8 + ) + monkeypatch.setattr(module, "Glm5NextMemManager", construct) + elif module is next_module: + monkeypatch.setattr(module, "Qwen3NextMemManager", construct) + monkeypatch.setattr( + module, + "get_env_start_args", + lambda: SimpleNamespace(llm_kv_type="None", linear_att_ssm_data_type="bfloat16"), + ) + else: + monkeypatch.setattr(module, "select_mem_manager_class", lambda: construct) + monkeypatch.setattr(module, "get_added_mtp_kv_layer_num", lambda: 0, raising=False) + + budget = profile_without_allocating + original_fraction = model.mem_fraction + with profile_max_tokens.profile_model_memory(model): + pass + with pytest.raises(KVAllocationReached): + model._init_mem_manager() + + reserved_bytes = staging * 1024 ** 2 + 2 * 8 * 40 * 512 * 8 + 512 + assert model.mem_fraction == pytest.approx(original_fraction - reserved_bytes / budget.total) + available = budget.free - budget.total * (1 - model.mem_fraction) + assert manager.size == int(available / manager.get_cell_size()) + # The KV allocation must leave space for the complete staging layer. + assert budget.free - manager.size * manager.get_cell_size() >= reserved_bytes + + +@pytest.mark.parametrize( + "enabled,draft,size,run_mode", + [ + (False, False, None, "prefill"), + (True, True, None, "prefill"), + (True, False, 1024, "prefill"), + (True, False, None, "decode"), + ], +) +def test_inactive_eplb_draft_and_explicit_kv_size_keep_mem_fraction(monkeypatch, enabled, draft, size, run_mode): + model = make_model() + model.args.enable_prefill_eplb = enabled + model.args.run_mode = run_mode + model.is_mtp_draft_model = draft + model.max_total_token_num = size + monkeypatch.setattr(eplb_utils, "find_ep_moe_weights", lambda _: pytest.fail("EPLB budgeting must be skipped")) + with profile_max_tokens.profile_model_memory(model): + pass + assert model.mem_fraction == 0.98 + + +def test_weight_aliases_are_deduplicated_and_non_eplb_layers_are_excluded(profile_without_allocating): + model = make_model() + first, last = model.trans_layers_weight[0].experts, model.trans_layers_weight[-1].experts + ordinary = make_expert() + ordinary.expert_parallel_state = ExpertParallelState(512, 8) + model.trans_layers_weight = [ + SimpleNamespace(a=last, b=first, alias=first, ordinary=ordinary, dense=object()), + SimpleNamespace(experts=last), + ] + assert [ + weight for weight in eplb_utils.find_ep_moe_weights(model) if weight.expert_parallel_state.eplb is not None + ] == [first, last] + expected_sampling_bytes = 2 * 2 * 40 * 512 * 8 + 512 + assert eplb_utils.get_eplb_sampling_peak_nbytes([first, last]) == expected_sampling_bytes + reserved_bytes = eplb_utils.get_eplb_staging_nbytes([first, last]) + expected_sampling_bytes + assert profile_max_tokens.get_eplb_adjusted_mem_fraction(model) == pytest.approx( + model.mem_fraction - reserved_bytes / profile_without_allocating.total + ) + + +@pytest.mark.parametrize("redundant,expected", [(0, 36), (2, 60)]) +@pytest.mark.parametrize("zero_point", [False, True]) +def test_full_staging_includes_odd_rows_and_scalar_quantization_scales(redundant, expected, zero_point): + expert = make_expert(redundant=redundant) + expert.expert_parallel_state = ExpertParallelState(24, 8, expert.expert_parallel_state.eplb) + rows = 3 + redundant + expert.w13 = SimpleNamespace( + weight=torch.empty((rows, 5), dtype=torch.uint8), + weight_scale=torch.empty((rows, 1), dtype=torch.float32), + weight_zero_point=torch.empty((rows, 1), dtype=torch.float32) if zero_point else None, + ) + expert.w2 = SimpleNamespace(weight=torch.empty((rows, 3), dtype=torch.uint8)) + assert eplb_utils.get_eplb_staging_nbytes([expert] * 12) == expected + (4 * rows if zero_point else 0) + + +@pytest.mark.parametrize( + "manager_class", + [ + memory_module.MemoryManager, + Deepseek2MemoryManager, + Deepseek3_2MemoryManager, + FP8PerTokenGroupQuantDeepseek3_2MemoryManager, + PPLINT4KVMemoryManager, + PPLINT8KVMemoryManager, + FP8StaticPerHeadQuantMemManager, + FP8StaticPerTensorQuantMemManager, + Qwen3NextMemManager, + Glm5NextMemManager, + FP8StaticPerHeadQuantQwen3NextMemManager, + FP8StaticPerTensorQuantQwen3NextMemManager, + ], +) +def test_kv_backend_constructors_use_adjusted_mem_fraction_once(profile_without_allocating, manager_class): + budget = profile_without_allocating + model = make_model() + with profile_max_tokens.profile_model_memory(model): + pass + manager = manager_class.__new__(manager_class) + kwargs = dict( + size=None, dtype=torch.bfloat16, head_num=1, head_dim=576, layer_num=8, mem_fraction=model.mem_fraction + ) + if issubclass(manager_class, Qwen3NextMemManager): + kwargs["num_kv_heads"] = kwargs.pop("head_num") + kwargs["full_att_layer_num"] = kwargs.pop("layer_num") + kwargs["linear_config"] = SimpleNamespace() + with pytest.raises(KVAllocationReached): + manager_class.__init__(manager, **kwargs) + available = budget.free - budget.total * (1 - model.mem_fraction) + assert manager.size == int(available / manager.get_cell_size()) + + +@pytest.mark.parametrize("eplb", [False, True]) +def test_weight_profiling_reserves_mtp_and_optional_eplb(monkeypatch, profile_without_allocating, eplb): + model = make_model() + model.args.enable_prefill_eplb = eplb + model.args.mtp_mode = "vanilla" + original_fraction = model.mem_fraction + allocations = iter([2 * 1024 ** 3, 10 * 1024 ** 3]) + monkeypatch.setattr(torch.cuda, "memory_allocated", lambda: next(allocations)) + monkeypatch.setattr(profile_max_tokens, "get_mtp_weight_layer_num", lambda: 1) + weights = [layer.experts for layer in model.trans_layers_weight] + eplb_bytes = ( + eplb_utils.get_eplb_staging_nbytes(weights) + eplb_utils.get_eplb_sampling_peak_nbytes(weights) if eplb else 0 + ) + with profile_max_tokens.profile_model_memory(model): + assert model.mem_fraction == original_fraction + assert model.mem_fraction == pytest.approx( + original_fraction - (1024 ** 3 + eplb_bytes) / profile_without_allocating.total + ) + + +def test_distributed_budget_uses_minimum_mem_fraction(monkeypatch, profile_without_allocating): + model = make_model() + weights = [layer.experts for layer in model.trans_layers_weight] + reserved_bytes = eplb_utils.get_eplb_staging_nbytes(weights) + eplb_utils.get_eplb_sampling_peak_nbytes(weights) + expected_local = model.mem_fraction - reserved_bytes / profile_without_allocating.total + monkeypatch.setattr(profile_max_tokens, "get_current_device_id", lambda: 3) + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 8) + original_tensor = torch.tensor + devices = [] + + def cpu_tensor(value, *, dtype, device): + devices.append(device) + return original_tensor(value, dtype=dtype) + + def all_reduce(fraction, op): + assert fraction.item() == pytest.approx(expected_local) + assert op == torch.distributed.ReduceOp.MIN + fraction.fill_(0.8) # Another rank has less remaining memory. + + monkeypatch.setattr(torch, "tensor", cpu_tensor) + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + assert profile_max_tokens.get_eplb_adjusted_mem_fraction(model) == pytest.approx(0.8) + assert devices == ["cuda:3"] + + +def test_eplb_without_experts_keeps_mem_fraction(): + model = make_model() + model.trans_layers_weight = [] + assert profile_max_tokens.get_eplb_adjusted_mem_fraction(model) == model.mem_fraction + + +def test_insufficient_budget_is_rejected_before_kv_allocation(profile_without_allocating): + model = make_model() + model.mem_fraction = 0.001 + with pytest.raises(AssertionError, match="EPLB buffers leave no memory"): + profile_max_tokens.get_eplb_adjusted_mem_fraction(model) + + +@pytest.mark.parametrize( + "model_class,draft", + [ + (base_module.TpPartBaseModel, False), + (deepseek_module.Deepseek2TpPartModel, False), + (Deepseek3MTPModel, True), + (Qwen3MOEModel, False), + (Qwen3MOEMTPModel, True), + (Qwen3_5TpPartModel, False), + (Qwen3_5MTPModel, True), + (Glm5NextMTPModel, True), + ], +) +def test_config_initialization_sets_model_role_before_weights(tmp_path, model_class, draft): + config = dict( + num_attention_heads=2, + num_key_value_heads=2, + hidden_size=16, + num_hidden_layers=4, + num_nextn_predict_layers=1, + mhc=True, + is_mtp_draft_model=not draft, + ) + (tmp_path / "config.json").write_text(json.dumps(dict(config, text_config=config))) + model = model_class.__new__(model_class) + model.weight_dir_ = str(tmp_path) + model.finetune_config = None + model.tp_world_size_ = 1 + model.max_total_token_num = None + model.main_model = SimpleNamespace(config=config) + + model._init_config() + + assert model.config["is_mtp_draft_model"] is draft + assert config["is_mtp_draft_model"] is not draft + assert config["mhc"] is True + + +def test_base_model_passes_adjusted_fraction_to_overridden_kv_constructor(monkeypatch, profile_without_allocating): + template = make_model() + + class Model(base_module.TpPartBaseModel): + def _init_config(self): + self.config = template.config + + def _init_weights(self): + self.trans_layers_weight = template.trans_layers_weight + + def _init_req_manager(self): + # Weight profiling has already reserved EPLB memory; request-state + # allocation must still be visible to KV capacity profiling. + assert self.mem_fraction < 0.9 + profile_without_allocating.free -= 1024 + + def _init_mem_manager(self): + self.mem_manager = memory_module.MemoryManager.__new__(memory_module.MemoryManager) + memory_module.MemoryManager.__init__( + self.mem_manager, None, torch.bfloat16, 8, 128, 8, mem_fraction=self.mem_fraction + ) + + for name in ("_verify_must", "_verify_params", "_init_quant"): + monkeypatch.setattr(Model, name, lambda self: None) + monkeypatch.setattr( + base_module, "get_env_start_args", lambda: StartArgs(enable_ep_moe=True, enable_prefill_eplb=True) + ) + monkeypatch.setattr(base_module, "get_llm_data_type", lambda: torch.bfloat16) + monkeypatch.setattr(base_module, "get_dp_world_size", lambda: 1) + monkeypatch.setattr( + base_module, "TorchMemorySaverWrapper", lambda enabled: SimpleNamespace(region=lambda **kwargs: nullcontext()) + ) + monkeypatch.setattr( + base_module.MtpManager, + "get_instance", + lambda: SimpleNamespace( + get_decode_tokens_per_request=lambda draft: 1, get_decode_batch_alignment=lambda draft: 1 + ), + ) + model = Model.__new__(Model) + with pytest.raises(KVAllocationReached): + Model.__init__(model, dict(run_mode="prefill", weight_dir="unused", max_total_token_num=None)) + weights = [layer.experts for layer in model.trans_layers_weight] + reserved_bytes = eplb_utils.get_eplb_staging_nbytes(weights) + eplb_utils.get_eplb_sampling_peak_nbytes(weights) + assert model.mem_fraction == pytest.approx(0.9 - reserved_bytes / profile_without_allocating.total) + available = profile_without_allocating.free - profile_without_allocating.total * (1 - model.mem_fraction) + assert model.mem_manager.size == int(available / model.mem_manager.get_cell_size()) diff --git a/unit_tests/common/fused_moe/test_eplb_output_precision.py b/unit_tests/common/fused_moe/test_eplb_output_precision.py new file mode 100644 index 0000000000..c02ac64947 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_output_precision.py @@ -0,0 +1,129 @@ +"""Check exact EPLB expert selection and normal FP32/BF16 rounding bounds.""" + +import pytest +import torch +import torch.nn.functional as F + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _routed_expert_outputs(hidden, physical_ids, w13, w2): + """Keep the same BF16 expert MLP when physical expert rows move.""" + selected_w13 = w13[physical_ids.long()].float() + gate_up = torch.einsum("tkih,th->tki", selected_w13, hidden.float()).to(torch.bfloat16) + gate, up = gate_up.chunk(2, dim=-1) + intermediate = (F.silu(gate.float()) * up.float()).to(torch.bfloat16) + return torch.einsum("tkhi,tki->tkh", w2[physical_ids.long()].float(), intermediate.float()).to(torch.bfloat16) + + +def _gather_by_placement(expert_outputs, weights, physical_ids, capacity): + from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_expanded_layout_kernels import ep_gather_chunk + + tokens, topk, hidden = expert_outputs.shape + partials = [] + for rank in range(2): + routes = torch.nonzero((physical_ids // capacity).flatten() == rank).flatten() + routes = routes[physical_ids.flatten()[routes].argsort(stable=True)] + metadata = torch.full((tokens, topk + 2), -1, dtype=torch.int32, device="cuda") + metadata[routes // topk, routes % topk + 2] = torch.arange(routes.numel(), dtype=torch.int32, device="cuda") + partial = torch.zeros((tokens, hidden), dtype=torch.bfloat16, device="cuda") + if routes.numel(): + ep_gather_chunk( + expert_outputs.reshape(-1, hidden)[routes].contiguous(), + 0, + weights.flatten()[routes].contiguous(), + metadata, + partial, + ) + partials.append(partial) + # Isolate placement-dependent local gather; transport preserves these values. + return torch.stack(partials).float().sum(0).to(torch.bfloat16) + + +@pytest.mark.parametrize( + "experts,groups,group_topk,grouped", [(288, 1, 1, True), (256, 8, 4, True), (128, 1, 1, False)] +) +@pytest.mark.parametrize("tokens", [1, 8, 32]) +@pytest.mark.parametrize("tied", [False, True]) +@pytest.mark.parametrize("migrated", [False, True]) +@pytest.mark.parametrize("is_prefill", [False, True]) +def test_eplb_preserves_routing_and_moe_output( + experts, groups, group_topk, grouped, tokens, tied, migrated, is_prefill +): + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import ( + build_initial_redundant_expert_ids, + build_logical_to_physical_map, + expand_redundant_placement, + ) + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import EPLBState + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import FuseMoeDeepGEMM + + generator = torch.Generator(device="cuda").manual_seed(20261022) + hidden = torch.randn((tokens, 32), device="cuda", generator=generator).to(torch.bfloat16) + logits = torch.randn((tokens, experts), device="cuda", generator=generator) + bias = (torch.arange(experts, device="cuda") % 11 / 64).to(torch.bfloat16) if grouped else None + if tied: + logits.zero_() + if bias is not None: + bias.zero_() + w13 = (torch.randn((experts, 64, 32), device="cuda", generator=generator) / 8).to(torch.bfloat16) + w2 = (torch.randn((experts, 32, 32), device="cuda", generator=generator) / 8).to(torch.bfloat16) + moe = FuseMoeDeepGEMM.__new__(FuseMoeDeepGEMM) + moe.eplb = None + moe.routed_scaling_factor = 2.5 + arguments = dict( + input_tensor=hidden, + router_logits=logits, + correction_bias=bias, + top_k=8, + renormalize=True, + use_grouped_topk=grouped, + topk_group=group_topk, + num_expert_group=groups, + scoring_func="sigmoid" if grouped else "softmax", + is_prefill=is_prefill, + preserve_logical_ids=True, + ) + off_weights, off_physical, off_logical = moe._select_experts(**arguments) + off_experts = _routed_expert_outputs(hidden, off_physical, w13, w2) + off_output = _gather_by_placement(off_experts, off_weights, off_physical, experts // 2) + + placement = expand_redundant_placement(build_initial_redundant_expert_ids(experts, 2, 2), experts) + if migrated: + placement = placement.flip(0).roll(placement.shape[1] // 2, dims=1) + mapping, counts = build_logical_to_physical_map(placement, experts, source_rank=0, node_world_size=2) + moe.eplb = EPLBState( + num_redundant_experts_per_rank=2, + logical_to_physical_map=mapping.cuda(), + logical_replica_count=counts.cuda(), + route_counter=torch.zeros((1, experts), dtype=torch.int64, device="cuda"), + ) + on_weights, on_physical, on_logical = moe._select_experts(**arguments) + torch.testing.assert_close(on_logical.long(), off_logical.long(), rtol=0, atol=0) + # Different FP32 reduction orders are valid; verify the actual weight formula + # independently instead of requiring bitwise-identical normalized weights. + scores = torch.sigmoid(logits.double()) if grouped else logits.double().softmax(dim=-1) + reference_weights = scores.gather(1, off_logical.long()) + reference_weights = reference_weights / reference_weights.sum(dim=-1, keepdim=True) * moe.routed_scaling_factor + unit_fp32 = torch.finfo(torch.float32).eps / 2 + rounding_steps = arguments["top_k"] + 4 + weight_rtol = rounding_steps * unit_fp32 / (1 - rounding_steps * unit_fp32) + for weights in (off_weights, on_weights): + torch.testing.assert_close(weights.double(), reference_weights, rtol=weight_rtol, atol=0) + rows = placement.flatten().cuda() + torch.testing.assert_close(rows[on_physical.long()], off_logical.long(), rtol=0, atol=0) + on_experts = _routed_expert_outputs(hidden, on_physical, w13[rows], w2[rows]) + torch.testing.assert_close(on_experts, off_experts, rtol=0, atol=0) + on_output = _gather_by_placement(on_experts, on_weights, on_physical, placement.shape[1]) + + contributions = off_experts.double() * reference_weights.unsqueeze(-1) + reference = contributions.sum(dim=1) + magnitude = contributions.abs().sum(dim=1) + # One BF16 local cast and one final cast, plus FP32 routing/product/sum rounding. + unit_bf16 = torch.finfo(torch.bfloat16).eps / 2 + rounding_bound = 2 * unit_bf16 / (1 - 2 * unit_bf16) + 10 * unit_fp32 / (1 - 10 * unit_fp32) + bound = (rounding_bound + (1 + rounding_bound) * weight_rtol) * magnitude + for output in (off_output, on_output): + assert torch.all((output.double() - reference).abs() <= bound), "MoE output exceeds BF16 rounding bound" + assert torch.all((on_output.double() - off_output.double()).abs() <= 2 * bound) diff --git a/unit_tests/common/fused_moe/test_eplb_parallel.py b/unit_tests/common/fused_moe/test_eplb_parallel.py new file mode 100644 index 0000000000..9f934125de --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_parallel.py @@ -0,0 +1,234 @@ +"""EPLB routing and output consistency follow the TP/SP layout for every model.""" + +from datetime import timedelta +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight as weight_module +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import build_logical_to_physical_map +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import deepgemm_impl as deepgemm_module +from lightllm.common.quantization.quantize_method import WeightPack +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import device_utils, dist_utils, envs_utils + + +def _cpu_weight_factory(monkeypatch): + # Exercise the real common initialization without allocating CUDA buffers. + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: False) + monkeypatch.setattr(torch.Tensor, "cuda", lambda tensor: tensor) + original_zeros = torch.zeros + + def cpu_zeros(*args, **kwargs): + if kwargs.get("device") == "cuda": + kwargs["device"] = "cpu" + return original_zeros(*args, **kwargs) + + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(weight_module, "get_moe_capture_callback", lambda *args: None) + + def make_weight(rank, tp, node_size, *, world=8, tpsp=False, eplb=True, draft=False): + args = StartArgs( + run_mode="prefill", + tp=world, + dp=world // tp, + enable_ep_moe=True, + enable_prefill_eplb=eplb, + enable_tpsp_mix_mode=tpsp, + eplb_num_redundant_experts_per_rank=2, + ) + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr(dist_utils, "get_node_world_size", lambda: node_size) + weight = weight_module.FusedMoeWeight.__new__(weight_module.FusedMoeWeight) + weight.n_routed_experts = 2 * world + weight.global_rank_, weight.global_world_size = rank, world + weight.tp_rank_, weight.tp_world_size_ = rank % tp, tp + weight.enable_ep_moe = True + weight._init_expert_parallel_state({"is_mtp_draft_model": draft}) + weight.layer_num_ = 0 + num_local_experts = weight.expert_parallel_state.num_total_physical_experts // world + weight.w13 = weight.w2 = WeightPack(weight=torch.empty((num_local_experts, 1, 1))) + weight.e_score_correction_bias = weight.per_expert_scale = None + weight.scoring_func = "softmax" + weight.fuse_moe_impl = weight_module.create_fuse_moe_impl( + n_routed_experts=weight.n_routed_experts, + num_fused_shared_experts=0, + routed_scaling_factor=1.0, + quant_method=None, + expert_parallel_state=weight.expert_parallel_state, + ) + # Keep the real implementation and communication boundary; replace GPU kernels. + weight.fuse_moe_impl._select_experts = lambda input_tensor, **kwargs: ( + None, + torch.zeros((input_tensor.shape[0], 1), dtype=torch.long), + None, + ) + monkeypatch.setattr( + deepgemm_module, "fused_experts", lambda hidden_states, **kwargs: hidden_states.clone() + 100 * rank + ) + return weight + + return make_weight + + +@pytest.fixture +def make_weight(monkeypatch): + return _cpu_weight_factory(monkeypatch) + + +@pytest.mark.parametrize("tpsp", [False, True]) +@pytest.mark.parametrize("tp,node_size", [(1, 4), (2, 4), (4, 4), (4, 2), (8, 4), (2, 1)]) +def test_common_initialization_preserves_node_local_replicas_and_tp_consistency(make_weight, tp, node_size, tpsp): + states = [] + for rank in range(8): + state = make_weight(rank, tp, node_size, tpsp=tpsp).expert_parallel_state.eplb + leader = rank // tp * tp + same_node = leader // node_size == (leader + tp - 1) // node_size + replicated = tp > 1 and not tpsp + expected_rotation = leader if replicated and same_node else rank + assert state.replica_rotation_rank == expected_rotation + assert state.tp_output_sync_src == (leader if replicated and not same_node else None) + original_maps, original_counts = build_logical_to_physical_map( + state.physical_to_logical, 16, source_rank=rank, node_world_size=node_size + ) + # TP consistency changes ordering, preserving every eligible local replica. + torch.testing.assert_close(state.logical_to_physical_map.sort(-1).values, original_maps.sort(-1).values) + torch.testing.assert_close(state.logical_replica_count, original_counts) + states.append(state) + if not tpsp and 1 < tp <= node_size: + for rank, state in enumerate(states): + torch.testing.assert_close(state.logical_to_physical_map, states[rank // tp * tp].logical_to_physical_map) + + +@pytest.mark.parametrize("batch_size", [0, 3]) +@pytest.mark.parametrize( + "prefill,eplb,tp,node_size,tpsp,draft,synchronize", + [ + (True, True, 4, 2, False, False, True), + (True, True, 4, 8, False, False, False), + (False, True, 4, 2, False, False, False), + (None, True, 4, 2, False, False, False), + (True, False, 4, 2, False, False, False), + (True, True, 1, 2, False, False, False), + (True, True, 4, 2, True, False, False), + # Static draft layouts need the same cross-node TP output consistency. + (True, True, 4, 2, False, True, True), + ], +) +def test_ep_impl_synchronizes_only_replicated_eplb_prefill( + monkeypatch, make_weight, batch_size, prefill, eplb, tp, node_size, tpsp, draft, synchronize +): + weight = make_weight(6, tp, node_size, tpsp=tpsp, eplb=eplb, draft=draft) + group = object() + broadcast = Mock(side_effect=lambda output, **kwargs: output.fill_(4)) + monkeypatch.setattr(dist, "get_world_size", lambda group: tp) + monkeypatch.setattr(dist, "broadcast", broadcast) + inputs = torch.ones((batch_size, 4)) + output = weight.experts( + inputs, + router_logits=None, + top_k=2, + renormalize=True, + use_grouped_topk=False, + topk_group=0, + num_expert_group=0, + is_prefill=prefill, + infer_state=SimpleNamespace(dist_group=group), + ) + if synchronize: + broadcast.assert_called_once_with(output, src=4, group=group, async_op=False) + else: + broadcast.assert_not_called() + torch.testing.assert_close(output, torch.full_like(inputs, 4 if synchronize else 601)) + + +@pytest.mark.parametrize("prefill", [True, False, None]) +def test_eplb_sync_preserves_output_and_order(monkeypatch, make_weight, prefill): + weight = make_weight(6, 4, 2) + inputs = torch.ones((3, 4)) + output = inputs.clone() + group = object() + events = [] + + def compute(**kwargs): + events.append("compute") + return output + + def broadcast(tensor, **kwargs): + events.append("broadcast") + assert tensor is output + assert kwargs == {"src": 4, "group": group, "async_op": False} + tensor.add_(10) + + monkeypatch.setattr(deepgemm_module, "fused_experts", compute) + monkeypatch.setattr(dist, "get_world_size", lambda group: 4) + monkeypatch.setattr(dist, "broadcast", broadcast) + result = weight.fuse_moe_impl._fused_experts( + inputs, + weight.w13, + weight.w2, + torch.ones((3, 2)), + torch.zeros((3, 2), dtype=torch.int32), + is_prefill=prefill, + dist_group=group, + ) + assert result is output + assert events == (["compute", "broadcast"] if prefill is True else ["compute"]) + torch.testing.assert_close(output, inputs + (10 if prefill is True else 0)) + + +def test_failed_eplb_compute_does_not_broadcast(monkeypatch, make_weight): + weight = make_weight(6, 4, 2) + error = RuntimeError("expert computation failed") + + def fail(**kwargs): + raise error + + broadcast = Mock() + monkeypatch.setattr(deepgemm_module, "fused_experts", fail) + monkeypatch.setattr(dist, "get_world_size", lambda group: 4) + monkeypatch.setattr(dist, "broadcast", broadcast) + with pytest.raises(RuntimeError) as caught: + weight.fuse_moe_impl._fused_experts( + torch.ones((3, 4)), weight.w13, weight.w2, torch.ones((3, 2)), torch.zeros((3, 2)), is_prefill=True + ) + assert caught.value is error + broadcast.assert_not_called() + + +def _tp_output_worker(rank, rendezvous): + torch.set_num_threads(1) + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=4, timeout=timedelta(seconds=30)) + try: + groups = [dist.new_group(ranks) for ranks in ([0, 1], [2, 3])] + group = groups[rank // 2] + with pytest.MonkeyPatch.context() as monkeypatch: + make_weight = _cpu_weight_factory(monkeypatch) + for tpsp in (False, True): + # Each two-rank TP group spans simulated single-rank nodes. + weight = make_weight(rank, 2, 1, world=4, tpsp=tpsp) + input_value = rank + 1 if tpsp else 1 + inputs = torch.full((3, 4), input_value, dtype=torch.float32) + output = weight.experts( + inputs, + router_logits=None, + top_k=2, + renormalize=True, + use_grouped_topk=False, + topk_group=0, + num_expert_group=0, + is_prefill=True, + infer_state=SimpleNamespace(dist_group=group), + ) + expected_rank = rank if tpsp else rank // 2 * 2 + torch.testing.assert_close(output, torch.full_like(inputs, input_value + 100 * expected_rank)) + finally: + dist.destroy_process_group() + + +def test_real_tp_groups_synchronize_independently_and_preserve_sp_tokens(tmp_path): + mp.spawn(_tp_output_worker, args=((tmp_path / "gloo-init").as_uri(),), nprocs=4, join=True) diff --git a/unit_tests/common/fused_moe/test_eplb_planner_worker.py b/unit_tests/common/fused_moe/test_eplb_planner_worker.py new file mode 100644 index 0000000000..345dc0afce --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_planner_worker.py @@ -0,0 +1,130 @@ +import os +import signal +import subprocess +import sys +import time + +import pytest +import torch + +from lightllm.common.eplb_planner_worker import EPLBPlannerWorker +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_planner import plan_eplb_candidate + +FULL_LOAD = torch.ones((1, 1, 1, 8), dtype=torch.int64) +FULL_CURRENT = torch.arange(8, dtype=torch.int64).view(1, 8, 1) +FULL = dict( + node_world_size=8, + expert_alignment=128, + stickiness=0.1, +) +RED_LOAD = torch.ones((1, 1, 1, 8), dtype=torch.int64) +RED_CURRENT = torch.stack((torch.arange(8), (torch.arange(8) + 1) % 8), dim=-1)[None] +RED = dict( + node_world_size=8, + expert_alignment=128, + stickiness=0.1, +) + + +def dead(pid): + try: + return open(f"/proc/{pid}/stat").read().split()[2] == "Z" + except (FileNotFoundError, ProcessLookupError): + return True + + +def test_worker_preserves_noncontiguous_exact_hash_samples(): + from unit_tests.common.fused_moe.test_eplb_runtime_hash import _samples + + rng = torch.Generator().manual_seed(17) + samples = [[torch.rand((96, 8), generator=rng).argsort(-1)[:, :2] for _ in range(4)]] + logical, hist = _samples(samples, 8, 4, 4) + packed = torch.cat((logical[..., None], hist), -1) + current = torch.tensor([[[0, 1, 6, 7], [0, 2, 3, 4], [0, 3, 4, 5], [1, 2, 5, 6]]]) + settings = dict(node_world_size=4, expert_alignment=128, stickiness=0.1, replica_hash_load=packed[..., 1:]) + expected = plan_eplb_candidate(packed[..., 0], current, **settings) + worker = EPLBPlannerWorker(timeout_s=30) + try: + actual, timing = worker.plan(packed[..., 0], current, **settings) + assert torch.equal(actual, expected) + assert not worker.info["cuda_initialized"] + finally: + worker.close() + + +def test_worker_persistent_full_redundant_and_cpu(): + worker = EPLBPlannerWorker(timeout_s=30) + pid = worker.info["pid"] + try: + for load, current, settings in ((FULL_LOAD, FULL_CURRENT, FULL), (RED_LOAD, RED_CURRENT, RED)): + direct = plan_eplb_candidate(load, current, **settings) + actual, _ = worker.plan(load, current, **settings) + assert torch.equal(direct, actual) + assert worker.info["pid"] == pid and not worker.info["cuda_initialized"] + finally: + worker.close() + assert dead(pid) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires a CUDA-initialized serving parent") +def test_worker_exec_does_not_reimport_cuda_initialized_main(tmp_path): + script = tmp_path / "serving_main.py" + script.write_text( + """import torch +torch.cuda.init() +from lightllm.common.eplb_planner_worker import EPLBPlannerWorker +if __name__ == "__main__": + assert torch.cuda.is_initialized() + worker = EPLBPlannerWorker(timeout_s=30) + try: + assert not worker.info["cuda_initialized"] + load = torch.ones((1, 1, 1, 8), dtype=torch.int64) + current = torch.arange(8).view(1, 8, 1) + result, _ = worker.plan(load, current, node_world_size=8, expert_alignment=128, stickiness=0.1) + assert torch.equal(result, current) + finally: + worker.close() +""" + ) + subprocess.run([sys.executable, str(script)], capture_output=True, text=True, timeout=45, check=True) + + +def test_worker_rejects_concurrent_and_timeout_cleans_child(): + worker = EPLBPlannerWorker(timeout_s=30) + pid = worker.info["pid"] + try: + assert worker._rpc_lock.acquire(False) + with pytest.raises(RuntimeError, match="one in-flight"): + worker.plan(FULL_LOAD, FULL_CURRENT, **FULL) + worker._rpc_lock.release() + assert worker._process.is_alive() + os.kill(pid, signal.SIGSTOP) + worker._timeout_s = 0.05 + with pytest.raises(TimeoutError): + worker.plan(FULL_LOAD, FULL_CURRENT, **FULL) + finally: + worker.close() + assert dead(pid) + + +def test_worker_error_and_parent_watchdog_exit(): + worker = EPLBPlannerWorker(timeout_s=30) + pid = worker.info["pid"] + try: + with pytest.raises(RuntimeError): + worker.plan(FULL_LOAD, FULL_CURRENT.repeat(1, 1, 2), **FULL) + finally: + worker.close() + assert dead(pid) + code = """import os +from lightllm.common.eplb_planner_worker import EPLBPlannerWorker +worker = EPLBPlannerWorker() +print(worker.info["pid"], flush=True) +os._exit(0) +""" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=45, check=True) + child = next(int(line) for line in result.stdout.splitlines() if line.strip().isdigit()) + until = time.monotonic() + 5 + while not dead(child) and time.monotonic() < until: + time.sleep(0.1) + assert dead(child) diff --git a/unit_tests/common/fused_moe/test_eplb_runtime_hash.py b/unit_tests/common/fused_moe/test_eplb_runtime_hash.py new file mode 100644 index 0000000000..b9f5f6016c --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_runtime_hash.py @@ -0,0 +1,150 @@ +"""Route replay oracles independent of the planner's uniform-load approximation.""" +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import ( + _estimate_rank_load, + build_logical_to_physical_map, + build_initial_redundant_expert_ids, + expand_redundant_placement, + get_max_reachable_replicas, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_planner import plan_full_experts + + +def _hash(token, expert, copies): + value = token ^ (((expert + 1) * 0x9E3779B9) & 0xFFFFFFFF) + value = ((value ^ (value >> 16)) * 0x7FEB352D) & 0xFFFFFFFF + value = ((value ^ (value >> 15)) * 0x846CA68B) & 0xFFFFFFFF + return (value ^ (value >> 16)) % copies + + +def _samples(ids_by_sample, experts, world, node_size, tp_size=1): + load = torch.zeros((len(ids_by_sample), 1, world // node_size, experts), dtype=torch.int64) + hist = torch.zeros((*load.shape, world * (world + 1) // 2), dtype=torch.int64) + for sample, sources in enumerate(ids_by_sample): + for rank, ids in enumerate(sources): + token = torch.arange(len(ids))[:, None].expand_as(ids).flatten() + flat = ids.flatten() + node = rank // node_size + load[sample, 0, node] += torch.bincount(flat, minlength=experts) + for copies in range(1, world + 1): + bucket = torch.zeros_like(flat) if len(ids) == 1 else _hash(token, flat, copies) + bucket = (bucket + rank // tp_size * tp_size) % copies + start = copies * (copies - 1) // 2 + hist[sample, 0, node, :, start : start + copies] += torch.bincount( + flat * copies + bucket, minlength=experts * copies + ).reshape(experts, copies) + return load, hist + + +def _route_oracle(ids_by_sample, layout, node_size, tp_size=1): + world, capacity = layout.shape + experts = int(layout.max()) + 1 + results = [] + rows = layout.tolist() + for sources in ids_by_sample: + counts = torch.zeros((world, experts), dtype=torch.int64) + for source, ids in enumerate(sources): + for token, selected in enumerate(ids.tolist()): + for expert in selected: + owners = [rank for rank in range(world) if expert in rows[rank]] + local = [rank for rank in owners if rank // node_size == source // node_size] + owners = local or owners + bucket = 0 if len(ids) == 1 else _hash(token, expert, len(owners)) + destination = owners[(bucket + source // tp_size * tp_size) % len(owners)] + counts[destination, expert] += 1 + results.append(((counts + 127) // 128 * 128).sum(-1)) + return torch.stack(results)[:, None] + + +@pytest.mark.parametrize("node_size", [1, 2, 4]) +@pytest.mark.parametrize("tp_size", [1, 2, 4]) +def test_hash_estimator_matches_exact_source_local_route_replay(node_size, tp_size): + layout = torch.tensor([[0, 1, 6, 7], [0, 2, 3, 4], [0, 3, 4, 5], [1, 2, 5, 6]]) + rng = torch.Generator().manual_seed(791) + samples = [[torch.rand((tokens, 8), generator=rng).argsort(-1)[:, :2] for _ in range(4)] for tokens in [96, 1, 257]] + load, hist = _samples(samples, 8, 4, node_size, tp_size) + actual = _estimate_rank_load(load, layout[None], 128, node_size, replica_hash_load=hist) + torch.testing.assert_close(actual, _route_oracle(samples, layout, node_size, tp_size).double(), rtol=0, atol=0) + candidate = plan_full_experts( + load, layout[None], expert_alignment=128, node_world_size=node_size, replica_hash_load=hist + ) + candidate_load = _estimate_rank_load(load, candidate, 128, node_size, replica_hash_load=hist) + torch.testing.assert_close( + candidate_load, _route_oracle(samples, candidate[0], node_size, tp_size).double(), rtol=0, atol=0 + ) + assert candidate_load.max(-1).values.sum() <= actual.max(-1).values.sum() + + +def test_three_replicas_crossing_128_row_boundary_are_not_underestimated(): + layout = torch.cat((torch.arange(288).reshape(4, 72), torch.tensor([[72, 144], [0, 216], [0, 1], [2, 3]])), 1) + samples = [[torch.zeros((96, 1), dtype=torch.int64) for _ in range(4)]] + load, hist = _samples(samples, 288, 4, 4) + actual = _estimate_rank_load(load, layout[None], 128, 4, replica_hash_load=hist) + assert actual[0, 0].tolist() == [128, 256, 256, 0] + + +def test_hash_sampling_bound_uses_reachable_node_holdings_not_ep_world_size(): + for redundant, expected in [(0, 1), (2, 16)]: + placement = expand_redundant_placement(build_initial_redundant_expert_ids(256, 64, redundant), 256) + assert get_max_reachable_replicas(placement, 256, 8) == expected + + +def test_exact_estimator_accepts_only_the_reachable_hash_counts(): + layout = torch.arange(8).reshape(4, 2) + samples = [[torch.zeros((96, 1), dtype=torch.int64) for _ in range(4)]] + load, hist = _samples(samples, 8, 4, 4) + actual = _estimate_rank_load(load, layout[None], 128, 4, replica_hash_load=hist[..., :1]) + assert actual[0, 0].tolist() == [384, 0, 0, 0] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tokens", [1, 96]) +@pytest.mark.parametrize("source", range(4)) +@pytest.mark.parametrize("tp_size", [1, 4]) +def test_fused_topk_records_exact_hash_buckets_for_all_replica_counts(tokens, source, tp_size): + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk_eplb + + layout = torch.tensor([[0, 1, 6, 7], [0, 2, 3, 4], [0, 3, 4, 5], [1, 2, 5, 6]]) + phase = source // tp_size * tp_size + maps, copies = build_logical_to_physical_map( + layout, 8, source_rank=source, node_world_size=4, replica_rotation_rank=phase + ) + counter = torch.zeros((1, 8, 11), dtype=torch.int64, device="cuda") + _, physical, logical = triton_grouped_topk_eplb( + hidden_states=torch.zeros((tokens, 32), device="cuda"), + gating_output=torch.randn((tokens, 8), device="cuda"), + correction_bias=None, + topk=2, + renormalize=True, + num_expert_group=1, + topk_group=1, + scoring_func="sigmoid", + logical_to_physical_map=maps.cuda(), + logical_replica_count=copies.cuda(), + expert_counter=counter, + sample_index=0, + record_load=True, + use_grouped_topk=True, + return_logical_ids=True, + group_score_used_topk_num=1, + replica_rotation_rank=phase, + ) + ids = logical.cpu() + empty = torch.empty((0, 2), dtype=torch.int64) + sources = [empty] * 4 + sources[source] = ids + load, hist = _samples([sources], 8, 4, 4, tp_size) + expected = torch.cat((load[0, 0, 0, :, None], hist[0, 0, 0]), -1) + torch.testing.assert_close(counter[0].cpu(), expected, rtol=0, atol=0) + expected_physical = [] + flat = layout.flatten().tolist() + for token, selected in enumerate(ids.tolist()): + row = [] + for expert in selected: + owners = [p for p, value in enumerate(flat) if value == expert] + bucket = 0 if tokens == 1 else _hash(token, expert, len(owners)) + row.append(owners[(bucket + phase) % len(owners)]) + expected_physical.append(row) + assert physical.cpu().tolist() == expected_physical diff --git a/unit_tests/common/fused_moe/test_eplb_sampling_policy.py b/unit_tests/common/fused_moe/test_eplb_sampling_policy.py new file mode 100644 index 0000000000..1b343a96dd --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_sampling_policy.py @@ -0,0 +1,333 @@ +import threading +from collections import deque +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend import eplb_manager as manager_module +from lightllm.utils import envs_utils +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import eplb_planner, eplb_placement + + +@pytest.fixture(autouse=True) +def clear_eplb_env_caches(monkeypatch): + for name in ( + "LIGHTLLM_PREFILL_EPLB_STEP_INTERVAL", + "LIGHTLLM_PREFILL_EPLB_STEADY_SAMPLE_STEPS", + "LIGHTLLM_PREFILL_EPLB_MIN_REBALANCE_INTERVAL", + ): + monkeypatch.delenv(name, raising=False) + for getter in ( + envs_utils.get_prefill_eplb_step_interval, + envs_utils.get_prefill_eplb_steady_sample_steps, + envs_utils.get_prefill_eplb_min_rebalance_interval, + ): + getter.cache_clear() + yield + for getter in ( + envs_utils.get_prefill_eplb_step_interval, + envs_utils.get_prefill_eplb_steady_sample_steps, + envs_utils.get_prefill_eplb_min_rebalance_interval, + ): + getter.cache_clear() + + +def test_eplb_sampling_getters_defaults_and_valid_values(monkeypatch): + assert envs_utils.get_prefill_eplb_step_interval() == 20 + assert envs_utils.get_prefill_eplb_steady_sample_steps() == 4 + assert envs_utils.get_prefill_eplb_min_rebalance_interval() == 0 + monkeypatch.setenv("LIGHTLLM_PREFILL_EPLB_STEP_INTERVAL", "20") + monkeypatch.setenv("LIGHTLLM_PREFILL_EPLB_STEADY_SAMPLE_STEPS", "20") + monkeypatch.setenv("LIGHTLLM_PREFILL_EPLB_MIN_REBALANCE_INTERVAL", "80") + for getter in ( + envs_utils.get_prefill_eplb_step_interval, + envs_utils.get_prefill_eplb_steady_sample_steps, + envs_utils.get_prefill_eplb_min_rebalance_interval, + ): + getter.cache_clear() + assert envs_utils.get_prefill_eplb_steady_sample_steps() == 20 + assert envs_utils.get_prefill_eplb_min_rebalance_interval() == 80 + + +@pytest.mark.parametrize( + "name,value,match", + [ + ("LIGHTLLM_PREFILL_EPLB_STEADY_SAMPLE_STEPS", "0", "greater than 0"), + ("LIGHTLLM_PREFILL_EPLB_STEADY_SAMPLE_STEPS", "nope", "invalid literal"), + ("LIGHTLLM_PREFILL_EPLB_MIN_REBALANCE_INTERVAL", "-1", "non-negative"), + ("LIGHTLLM_PREFILL_EPLB_MIN_REBALANCE_INTERVAL", "19", "at least step interval"), + ], +) +def test_eplb_sampling_getters_reject_invalid_values(monkeypatch, name, value, match): + monkeypatch.setenv(name, value) + envs_utils.get_prefill_eplb_steady_sample_steps.cache_clear() + envs_utils.get_prefill_eplb_min_rebalance_interval.cache_clear() + getter = ( + envs_utils.get_prefill_eplb_steady_sample_steps + if "STEADY" in name + else envs_utils.get_prefill_eplb_min_rebalance_interval + ) + with pytest.raises(ValueError, match=match): + getter() + + +def test_steady_window_clamps_to_configured_intervals(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.step_interval = 20 + manager.sampling_interval = 80 + manager.steady_sample_steps = 64 + assert manager._steady_sample_window_steps() == 20 + manager.sampling_interval = 3 + assert manager._steady_sample_window_steps() == 3 + + +def _no_improvement_manager(interval, min_interval): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.evaluation_in_flight = True + manager._evaluation_lock = threading.Lock() + manager._evaluation_error = None + manager._evaluation_thread = type("Done", (), {"join": lambda self: None})() + manager._evaluation_result = { + "kind": "no_improvement", + "model_imbalance_ratio": 1.2, + "candidate_model_imbalance_ratio": 1.1, + "candidate_rebalance_gain": 0.01, + "candidate_changed_layer_count": 1, + } + manager.global_rank = 1 + manager.step_interval = 20 + manager.sampling_interval = interval + manager.min_rebalance_interval = min_interval + manager.steady_sample_steps = 4 + manager.prefill_steps = 0 + manager.weights = [] + manager._eplb_states = [] + manager._last_evaluation_step = 0 + manager._set_recording = lambda _enabled: None + manager._reset_recorded_samples = lambda: None + return manager + + +def test_success_rebalance_uses_min80_interval(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.current_placement = torch.zeros((1, 1, 1), dtype=torch.int64) + manager.num_logical_experts = manager.world_size = manager.node_world_size = 1 + manager.step_interval = 20 + manager.min_rebalance_interval = 80 + manager.sampling_interval = 320 + manager.global_rank = 1 + manager._reset_recorded_samples = lambda: None + manager.transfer = type("Transfer", (), {"start": lambda *_args: None})() + manager._start_rebalance( + { + "placement": torch.zeros((1, 1, 1), dtype=torch.int64), + "metadata": [None], + "layer_plans": [], + "prepared_batches": [], + "before": {"max": 1, "p95": 1}, + "after": {"max": 1, "p95": 1}, + "model_imbalance_ratio": 1, + "candidate_model_imbalance_ratio": 1, + "candidate_rebalance_gain": 1, + "candidate_changed_layer_count": 0, + } + ) + assert manager.sampling_interval == 80 + + +def test_eplb_diagnostics_dump_is_rank_zero_cpu_only(tmp_path): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.global_rank = 0 + manager._diagnostics_dir = str(tmp_path) + manager._diagnostic_dump_sequence = 0 + manager.current_placement = torch.tensor([[1, 2]], dtype=torch.int64) + manager._evaluation_snapshot = {"placement": manager.current_placement.clone(), "completed_step": 7} + manager.prefill_steps = 9 + manager._dump_evaluation_samples(torch.ones((1, 1, 1, 1)), 4, 4) + payload = torch.load(tmp_path / "evaluation-000001.pt", weights_only=True) + assert set(payload) == { + "global_load", + "current_placement", + "prefill_steps", + "recorded_sample_count", + "sample_window_steps", + } + assert payload["global_load"].device.type == "cpu" + assert payload["prefill_steps"] == 7 + manager.global_rank = 1 + manager._dump_evaluation_samples(torch.ones((1, 1, 1, 1)), 4, 4) + assert len(list(tmp_path.glob("*.pt"))) == 1 + + +def test_eplb_diagnostics_dump_failure_disables_future_writes(monkeypatch, tmp_path): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.global_rank = 0 + manager._diagnostics_dir = str(tmp_path) + manager._diagnostic_dump_sequence = 0 + manager.current_placement = torch.zeros(1) + manager.prefill_steps = 1 + manager._evaluation_snapshot = {"placement": manager.current_placement, "completed_step": 1} + monkeypatch.setattr(torch, "save", lambda *_args, **_kwargs: (_ for _ in ()).throw(OSError("nope"))) + assert manager._dump_evaluation_samples(torch.ones(1), 1, 1) >= 0 + assert manager._diagnostics_dir is None + + +def test_base_one_keeps_default_four_sample_window(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.step_interval = 1 + manager.sampling_interval = 4 + manager.steady_sample_steps = 4 + assert manager._steady_sample_window_steps() == 4 + + +def manager_for_test(step=4, steady=4): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager._use_full_window = True + manager._initial_plan_complete = False + manager.in_flight = manager.evaluation_in_flight = False + manager.step_interval = manager.sampling_interval = step + manager.steady_sample_steps = steady + manager.min_rebalance_interval = 0 + manager.prefill_steps = 0 + manager._next_evaluation_step = step + manager._last_evaluation_step = 0 + manager._sample_step_counts = deque([0], maxlen=max(step, 4) + 1) + manager._eplb_states = [ + SimpleNamespace( + route_counter=torch.zeros((2 * max(step, 4), 2, 4), dtype=torch.int64), + recorded_sample_count=0, + recording=True, + ) + for _ in range(2) + ] + manager.current_placement = torch.tensor([[[0], [1]], [[0], [1]]]) + manager.starts = [] + manager._start_evaluation = lambda: manager.starts.append(manager.prefill_steps) + return manager + + +def forward(manager, dispatches=1): + manager.before_prefill() + for _ in range(dispatches): + for layer, state in enumerate(manager._eplb_states): + index = state.recorded_sample_count % state.route_counter.shape[0] + state.route_counter[index].add_(state.recorded_sample_count + 1 + layer * 1000) + state.recorded_sample_count += 1 + manager.step() + + +def test_busy_work_counts_real_forwards_and_due_evaluation_uses_first_idle_forward(): + manager = manager_for_test() + for _ in range(4): + forward(manager) + assert manager.starts == [4] + manager.evaluation_in_flight = True + for _ in range(3): + forward(manager) + manager.evaluation_in_flight = False + manager.in_flight = True + for _ in range(2): + forward(manager) + assert manager.prefill_steps == 9 + assert manager.starts == [4] + manager.in_flight = False + forward(manager) + assert manager.starts == [4, 10] + assert manager._next_evaluation_step == 12 + assert all(state.recording and state.recorded_sample_count == 10 for state in manager._eplb_states) + + +@pytest.mark.parametrize("pattern", [(1,) * 12, (2,) * 12, (1, 2, 1, 2, 2, 1, 2, 1, 1, 2, 1, 1)]) +def test_wrapping_live_ring_preserves_exact_recent_forward_window(pattern): + manager = manager_for_test() + for dispatches in pattern: + forward(manager, dispatches) + manager._freeze_evaluation_samples() + frozen = manager._evaluation_snapshot + total = sum(pattern) + count = sum(pattern[-4:]) + assert frozen["count"] == count + assert frozen["forward_end_counts"][-1] == total + assert frozen["completed_step"] == 12 + expected = torch.arange(total - count + 1, total + 1) + assert torch.equal(frozen["samples"][:, 0, 0, 0], expected) + assert torch.equal(frozen["samples"][:, 1, 0, 0], expected + 1000) + # A new ring overwrite and a partial migration cannot change the frozen + # logical routes or the complete base layout used by the CPU planner. + saved = frozen["samples"].clone() + forward(manager, 2) + manager.current_placement.fill_(7) + assert torch.equal(frozen["samples"], saved) + assert not (frozen["placement"] == 7).any() + + +def test_background_copy_reads_frozen_samples_and_releases_gpu_ownership(monkeypatch): + manager = manager_for_test() + for _ in range(4): + forward(manager) + manager._freeze_evaluation_samples() + expected = manager._evaluation_snapshot["samples"].clone() + manager.evaluation_group = object() + metadata = [] + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda value, **kwargs: metadata.append(value.clone())) + for state in manager._eplb_states: + state.route_counter.fill_(-1) + state.recorded_sample_count += 2 + copied = manager._collect_local_samples() + assert torch.equal(copied, expected) + assert manager._evaluation_snapshot["samples"] is None + assert metadata[0].tolist() == [4, -4, 8, -8, 4, -4, 4, -4, 4, -4] + assert metadata[1].tolist() == [0, 0, 1, -1, 2, -2, 3, -3, 4, -4] + + +def test_sampling_resume_keeps_history_and_honors_changed_min_interval(): + manager = manager_for_test() + for _ in range(6): + forward(manager) + manager.sampling_interval = 16 + manager._prepare_next_sampling_window() + assert manager._next_evaluation_step == 16 + assert manager._sample_step_counts[-1] == 6 + assert all(state.recorded_sample_count == 6 for state in manager._eplb_states) + + +def test_forward_count_mismatch_is_rejected_before_freezing(): + manager = manager_for_test() + manager._eplb_states[0].recorded_sample_count = 1 + with pytest.raises(RuntimeError, match="matching dispatches"): + manager.step() + + +def test_incremental_planner_bounds_migration_and_preserves_expert_coverage(): + current = eplb_placement.expand_redundant_placement( + eplb_placement.build_initial_redundant_expert_ids(16, 4, 1), 16 + ).unsqueeze(0) + load = torch.full((4, 1, 1, 16), 10, dtype=torch.int64) + load[..., 0] = 2000 + load[..., 1] = 1500 + target = eplb_planner.plan_eplb_candidate( + load, current, node_world_size=4, expert_alignment=128, stickiness=0.1, incremental=True + ) + eplb_placement.validate_physical_placement(target, 16) + old = eplb_placement._expert_locations(current, 16) + new = eplb_placement._expert_locations(target, 16) + assert 0 < int((new & ~old).sum()) <= 8 + before = eplb_placement._estimate_rank_load(load, current, 128, 4).max(-1).values.sum(0) + after = eplb_placement._estimate_rank_load(load, target, 128, 4).max(-1).values.sum(0) + assert torch.all(after < before) + + +@pytest.mark.parametrize("kind", ["insufficient", "no_improvement"]) +@pytest.mark.parametrize("minimum_interval", [0, 80, 640]) +def test_evaluation_keeps_sampling_and_uses_full_window_only_when_needed(kind, minimum_interval): + manager = _no_improvement_manager(20, minimum_interval) + manager._initial_plan_complete = False + manager._use_full_window = False + manager._evaluation_result["kind"] = kind + assert not manager._poll_evaluation() + assert manager._initial_plan_complete == (kind == "no_improvement") + assert manager._use_full_window == (kind == "insufficient") + expected = 20 if kind == "insufficient" else max(20, minimum_interval) + assert manager.sampling_interval == expected + assert manager._next_evaluation_step == expected diff --git a/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py b/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py new file mode 100644 index 0000000000..a386ff2207 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py @@ -0,0 +1,265 @@ +"""Complete-layout EPLB migration and stream-ordered publication across CUDA GPUs.""" +import math +import os +import socket +import time +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from lightllm.common.eplb_utils import get_eplb_staging_nbytes, get_eplb_staging_tensor_shape +from lightllm.server.router.model_infer.mode_backend.eplb_transfer import ( + NixlEPLBTransfer, + align_target_placement, + build_transfer_plan, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.eplb_placement import ( + build_initial_redundant_expert_ids, + expand_redundant_placement, + build_logical_to_physical_maps_for_layers, +) +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( + EPLBState, + ExpertParallelState, +) + +pytest.importorskip("nixl", reason="NIXL package is required") + + +class _Pack: + def __init__(self, weight, weight_scale): + self.weight = weight + self.weight_scale = weight_scale + self.weight_zero_point = None + + +def _free_port(): + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + port = sock.getsockname()[1] + sock.close() + return port + + +def _wait_for_ready_prefix(transfer, control_group): + deadline = time.monotonic() + 30 + while True: + pending = transfer.pending_layers() + ready_count = torch.tensor([len(pending)], dtype=torch.int32) + dist.all_reduce(ready_count, op=dist.ReduceOp.MIN, group=control_group) + if int(ready_count.item()) > 0: + return pending[: int(ready_count.item())] + if time.monotonic() >= deadline: + raise TimeoutError("EPLB transfer worker did not publish a globally ready layer") + time.sleep(0.001) + + +def _run_layers( + transfer, + control_group, + layer_plans, + callback=lambda layer_index: None, + before_commit_callback=lambda layer_index: None, +): + transfer.start(layer_plans, transfer.prepare_transfer(layer_plans)) + committed = 0 + while committed < len(layer_plans): + pending = _wait_for_ready_prefix(transfer, control_group) + assert len(pending) <= len(layer_plans) - committed + for layer_index, buffer_index in pending: + assert layer_index == layer_plans[committed][0] + before_commit_callback(layer_index) + transfer.commit( + layer_index, + buffer_index, + lambda layer_index=layer_index: callback(layer_index), + ) + committed += 1 + transfer.finish() + + +def _full_layout_worker(rank, port, world_size, redundant=0, inject_failure=False): + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + torch.set_num_threads(1) + dist.init_process_group("gloo", rank=rank, world_size=world_size) + control_group = dist.new_group(list(range(world_size)), backend="gloo") + transfer_group = dist.new_group(list(range(world_size)), backend="gloo") + experts = world_size * 2 + initial_extras = build_initial_redundant_expert_ids(experts, world_size, redundant) + current = expand_redundant_placement(initial_extras, experts) + capacity = current.shape[1] + layers = 9 if world_size == 8 else 3 + weights = [] + for layer in range(layers): + state = EPLBState( + num_redundant_experts_per_rank=redundant, + logical_to_physical_map=torch.empty((experts, world_size), device="cuda", dtype=torch.int32), + logical_replica_count=torch.ones(experts, device="cuda", dtype=torch.int32), + route_counter=torch.zeros((1, experts), device="cuda", dtype=torch.int64), + physical_to_logical=current.clone(), + ) + values = (current[rank].cuda() + 10 * layer).view(-1, 1) + cols13 = torch.arange(17, device="cuda").view(1, -1) + cols2 = torch.arange(9, device="cuda").view(1, -1) + weights.append( + SimpleNamespace( + expert_parallel_state=ExpertParallelState(experts, world_size, state), + w13=_Pack((values * 17 + cols13).to(torch.uint8), (values + 0.5).float()), + w2=_Pack((values * 9 + cols2).to(torch.float16), (values + 0.25 + cols2[:, :3]).float()), + ) + ) + for layer, weight in enumerate(weights): + values = (current[rank].cuda() + 10 * layer).view(-1, 1) + weight.w13.weight_zero_point = (values + torch.arange(5, device="cuda")).to(torch.int32) + transfer = NixlEPLBTransfer(weights, transfer_group, rank, world_size) + original_copy_part = transfer._copy_part + delayed = [False] + + def delayed_copy_part(prepared): + if rank == 0 and not delayed[0] and prepared.part_index == 0: + delayed[0] = True + time.sleep(0.05) + return original_copy_part(prepared) + + transfer._copy_part = delayed_copy_part + assert transfer.staging_depth == 1 + assert all(tensor.shape[0] == capacity for _, tensor in transfer.staging[0]) + actual = sum(tensor.nbytes for staging in transfer.staging for _, tensor in staging) + expected = sum(tensor.nbytes for _, tensor in transfer.live[0]) + shared_expected = sum( + math.prod(get_eplb_staging_tensor_shape(tensor, capacity)) * tensor.element_size() + for _, tensor in transfer.live[0] + ) + assert actual == expected == shared_expected == get_eplb_staging_nbytes(weights) + + def assert_values(layer, placement): + values = (placement[rank].cuda() + 10 * layer).view(-1, 1) + weight = weights[layer] + torch.testing.assert_close( + weight.w13.weight, + (values * 17 + torch.arange(17, device="cuda")).to(torch.uint8), + ) + torch.testing.assert_close(weight.w13.weight_scale, (values + 0.5).float()) + torch.testing.assert_close( + weight.w13.weight_zero_point, (values + torch.arange(5, device="cuda")).to(torch.int32) + ) + torch.testing.assert_close( + weight.w2.weight, + (values * 9 + torch.arange(9, device="cuda")).to(torch.float16), + ) + torch.testing.assert_close( + weight.w2.weight_scale, + (values + 0.25 + torch.arange(3, device="cuda")).float(), + ) + + targets = [current.roll(1, dims=0)] + second = targets[0].clone() + second[[0, 1]] = second[[1, 0]] + targets.append(second) + for generation, target in enumerate(targets, start=1): + target = align_target_placement(current, target) + plan = build_transfer_plan(current, target, experts, world_size, world_size) + if world_size == 3 and generation == 2: + assert all(step.dst_rank != 2 for step in plan) + maps, counts = build_logical_to_physical_maps_for_layers( + target[None], experts, source_rank=rank, node_world_size=world_size + ) + maps, counts = maps.cuda(), counts.cuda() + + snapshots = {} + published = {} + + def assert_old(layer): + assert_values(layer, current) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + torch.cuda._sleep(1000000) + snapshots[layer] = tuple( + t.clone() + for pack in (weights[layer].w13, weights[layer].w2) + for t in (pack.weight, pack.weight_scale, pack.weight_zero_point) + if t is not None + ) + torch.cuda.current_stream().wait_stream(stream) + + def publish(layer): + state = weights[layer].expert_parallel_state.eplb + state.logical_to_physical_map.copy_(maps[0], non_blocking=True) + state.logical_replica_count.copy_(counts[0], non_blocking=True) + # Observe the published generation on the same stream without a + # host assertion that would mask missing asynchronous ordering. + published[layer] = tuple( + t.clone() + for pack in (weights[layer].w13, weights[layer].w2) + for t in (pack.weight, pack.weight_scale, pack.weight_zero_point) + if t is not None + ) + state.physical_to_logical = target.clone() + state.placement_generation += 1 + + _run_layers(transfer, control_group, [(layer, plan) for layer in reversed(range(layers))], publish, assert_old) + torch.cuda.synchronize() + for layer, weight in enumerate(weights): + assert_values(layer, target) + for tensors, placement in [(snapshots[layer], current), (published[layer], target)]: + values = (placement[rank].cuda() + 10 * layer).view(-1, 1) + expected = [ + (values * 17 + torch.arange(17, device="cuda")).to(torch.uint8), + (values + 0.5).float(), + (values + torch.arange(5, device="cuda")).to(torch.int32), + (values * 9 + torch.arange(9, device="cuda")).to(torch.float16), + (values + 0.25 + torch.arange(3, device="cuda")).float(), + ] + for observed, reference in zip(tensors, expected): + torch.testing.assert_close(observed, reference) + state = weight.expert_parallel_state.eplb + assert state.placement_generation == generation + for expert in range(experts): + physical = state.logical_to_physical_map[expert, 0].item() + dst, row = divmod(physical, capacity) + assert target[dst, row] == expert + current = target + dist.barrier(group=control_group) + if inject_failure: + target = current.roll(1, dims=0) + plan = build_transfer_plan(current, target, experts, world_size, world_size) + original_commit = transfer._commit_full_part + + def fail_commit(*args): + if rank == 0: + raise RuntimeError("injected commit enqueue failure") + return original_commit(*args) + + transfer._commit_full_part = fail_commit + plans = [(0, plan)] + transfer.start(plans, transfer.prepare_transfer(plans)) + layer, buffer_index = _wait_for_ready_prefix(transfer, control_group)[0] + with pytest.raises(RuntimeError, match="full EPLB migration phase failed"): + transfer.commit(layer, buffer_index, lambda: pytest.fail("must not publish")) + assert all(release.is_set() for release in transfer._release) + transfer._thread.join(timeout=5) + assert not transfer._thread.is_alive() + with pytest.raises(RuntimeError, match="migration worker failed"): + transfer.finish() + transfer.shutdown() + dist.barrier(group=control_group) + dist.destroy_process_group() + + +@pytest.mark.parametrize("world_size", [2, 3, 4, 8]) +@pytest.mark.parametrize("redundant", [0, 1, 2]) +def test_full_layout_gpu_overwrite_cycles_and_staging_reuse(world_size, redundant): + if not torch.cuda.is_available() or torch.cuda.device_count() < world_size: + pytest.skip(f"requires {world_size} CUDA GPUs") + mp.spawn(_full_layout_worker, args=(_free_port(), world_size, redundant), nprocs=world_size, join=True) + + +def test_full_layout_gpu_commit_enqueue_failure_propagates(): + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("requires two CUDA GPUs") + mp.spawn(_full_layout_worker, args=(_free_port(), 2, 0, True), nprocs=2, join=True) diff --git a/unit_tests/common/fused_moe/test_eplb_transfer_init.py b/unit_tests/common/fused_moe/test_eplb_transfer_init.py new file mode 100644 index 0000000000..07d77ba5bd --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_transfer_init.py @@ -0,0 +1,100 @@ +"""Remote registration is driven by global migration plans, not cluster size.""" + +import builtins +import sys +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from lightllm.server.router.model_infer.mode_backend import eplb_transfer as module + + +@pytest.fixture +def make_transfer(monkeypatch): + def staging(self, weights, group, rank, world): + self.device = "mock-device" + self.global_rank, self.world_size = rank, world + self.live = [[("weight", "live")]] + self.staging = [[("weight", "staging")]] + + def ipc(self): + self._rank_hostnames = ["a", "a", "b", "b", "c", "c"] + self._same_node_ranks = { + rank for rank, host in enumerate(self._rank_hostnames) if host == self._rank_hostnames[self.global_rank] + } + self._cross_node_ranks = set(range(self.world_size)) - self._same_node_ranks + + monkeypatch.setattr(module.NixlEPLBTransfer, "_init_staging", staging) + monkeypatch.setattr(module.NixlEPLBTransfer, "_init_ipc_metadata", ipc) + monkeypatch.setattr(module.NixlEPLBTransfer, "_init_push_layouts", lambda self: None) + monkeypatch.setattr(module.NixlEPLBTransfer, "_prepare_part", lambda self, layer, plan, part: part) + monkeypatch.setattr(module.torch.cuda, "Stream", lambda **kwargs: object()) + monkeypatch.setattr(module, "_CudaBatchMemcpy", lambda: object()) + monkeypatch.setenv("UCX_TLS", "test-existing-value") + return lambda rank: module.NixlEPLBTransfer([], object(), rank, 6) + + +@pytest.mark.parametrize("rank", [0, 2, 4]) +def test_multinode_initialization_and_node_local_plans_do_not_import_nixl(monkeypatch, make_transfer, rank): + original_import = builtins.__import__ + + def no_nixl(name, *args, **kwargs): + if name == "nixl": + pytest.fail("node-local migration must not require NIXL") + return original_import(name, *args, **kwargs) + + monkeypatch.setattr(builtins, "__import__", no_nixl) + transfer = make_transfer(rank) + assert transfer._cross_node_ranks + assert transfer._nixl_agent is None + assert transfer.prepare_transfer([]) == [] + plans = [(0, [module.TransferStep(1, 0, 0, 0), module.TransferStep(3, 0, 2, 0)])] + assert len(transfer.prepare_transfer(plans)) == 1 + assert transfer._nixl_agent is None + assert transfer._registered_descs is None + transfer.shutdown() + + +@pytest.mark.parametrize("rank", [0, 2, 4]) +def test_first_remote_plan_initializes_all_ranks_once_including_nonparticipants(monkeypatch, make_transfer, rank): + agent = Mock() + agent.get_reg_descs.return_value = "descriptors" + nixl = SimpleNamespace(nixl_agent=Mock(return_value=agent), nixl_agent_config=Mock(return_value="config")) + monkeypatch.setitem(sys.modules, "nixl", nixl) + exchange = Mock() + monkeypatch.setattr(module.NixlEPLBTransfer, "_init_remote_metadata", exchange) + transfer = make_transfer(rank) + nixl.nixl_agent.assert_not_called() + plans = [(0, [module.TransferStep(2, 0, 0, 0)])] + transfer.prepare_transfer(plans) + transfer.prepare_transfer(plans) + nixl.nixl_agent.assert_called_once() + exchange.assert_called_once() + agent.get_reg_descs.assert_called_once_with(["live", "staging"]) + agent.register_memory.assert_called_once_with("descriptors", backends=["UCX"]) + transfer.shutdown() + transfer.shutdown() + agent.deregister_memory.assert_called_once_with("descriptors", backends=["UCX"]) + + +@pytest.mark.parametrize("failure", ["import", "register", "metadata"]) +def test_remote_initialization_failure_cleans_partial_resources(monkeypatch, make_transfer, failure): + agent = Mock() + agent.get_reg_descs.return_value = "descriptors" + nixl = SimpleNamespace(nixl_agent=Mock(return_value=agent), nixl_agent_config=Mock(return_value="config")) + monkeypatch.setitem(sys.modules, "nixl", None if failure == "import" else nixl) + if failure == "register": + agent.register_memory.side_effect = RuntimeError("registration failed") + exchange = Mock(side_effect=ValueError("metadata failed") if failure == "metadata" else None) + monkeypatch.setattr(module.NixlEPLBTransfer, "_init_remote_metadata", exchange) + transfer = make_transfer(4) + plans = [(0, [module.TransferStep(2, 0, 0, 0)])] + with pytest.raises(RuntimeError): + transfer.prepare_transfer(plans) + assert transfer._nixl_agent is None + assert transfer._registered_descs is None + assert transfer._remote_agents == {} + if failure != "import": + agent.deregister_memory.assert_called_once_with("descriptors", backends=["UCX"]) + transfer.shutdown() diff --git a/unit_tests/common/fused_moe/test_moe_weight_creation.py b/unit_tests/common/fused_moe/test_moe_weight_creation.py new file mode 100644 index 0000000000..5f5f5de1e4 --- /dev/null +++ b/unit_tests/common/fused_moe/test_moe_weight_creation.py @@ -0,0 +1,161 @@ +"""Real CUDA weight creation and checkpoint loading for singleton MoE experts.""" +import pytest +import torch + +from lightllm.common.basemodel.layer_weights.meta_weights import base_weight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight +from lightllm.common.quantization import quantize_method +from lightllm.common.quantization.awq import AWQW4A16QuantizationMethod, HAS_VLLM +from lightllm.common.quantization.deepgemm import DeepGEMMFP8w8a8B128QuantizationMethod, HAS_DEEPGEMM +from lightllm.common.quantization.no_quant import NoQuantization +from lightllm.common.quantization.w8a8gx import FP8w8a8g128QuantizationMethod +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import envs_utils, dist_utils + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.fixture +def device_id(monkeypatch): + device_id = torch.cuda.current_device() + monkeypatch.setattr(quantize_method, "get_current_device_id", lambda: device_id) + return device_id + + +@pytest.mark.parametrize("num_experts", [None, 1, 3], ids=["dense", "single", "multi"]) +@pytest.mark.parametrize( + "method_class,shapes,split_dims", + [ + (NoQuantization, ((256, 128), None, None), (-2, None, None)), + (DeepGEMMFP8w8a8B128QuantizationMethod, ((256, 128), (2, 1), None), (-2, -2, None)), + (FP8w8a8g128QuantizationMethod, ((256, 128), (256,), None), (-2, -1, None)), + (AWQW4A16QuantizationMethod, ((128, 32), (1, 256), (1, 32)), (-1, -1, -1)), + ], + ids=["bf16", "fp8-block", "fp8-channel", "awq"], +) +def test_weight_creation_preserves_expert_axis_and_loading_views( + device_id, method_class, shapes, split_dims, num_experts +): + if method_class is DeepGEMMFP8w8a8B128QuantizationMethod and not HAS_DEEPGEMM: + pytest.skip("DeepGEMM is required") + if method_class is AWQW4A16QuantizationMethod and not HAS_VLLM: + pytest.skip("vLLM is required for AWQ") + method = method_class() + method.hf_quantization_config = {"group_size": 128} + kwargs = dict(out_dims=[128, 128], in_dim=128, dtype=torch.bfloat16, device_id=device_id) + if num_experts is None: + parent, children = method.create_weight(**kwargs) + prefix = () + else: + parent, children = method.create_moe_weight(**kwargs, num_experts=num_experts) + prefix = (num_experts,) + assert len(children) == 2 + fields = ("weight", "weight_scale", "weight_zero_point") + loaders = (method.load_weight, method.load_weight_scale, method.load_weight_zero_point) + for name, shape, split_dim in zip(fields, shapes, split_dims): + tensor = getattr(parent, name) + if shape is None: + assert tensor is None + assert all(getattr(child, name) is None for child in children) + continue + assert tensor.shape == prefix + shape + assert tensor.device == torch.device("cuda", device_id) + child_shape = list(shape) + child_shape[split_dim] //= 2 + for child in children: + view = getattr(child, name) + assert view.shape == prefix + tuple(child_shape) + assert view.untyped_storage().data_ptr() == tensor.untyped_storage().data_ptr() + + # Loading each projection/expert must update the original fused buffers. + for expert_idx in range(1 if num_experts is None else num_experts): + parent_expert = parent if num_experts is None else parent.get_expert(expert_idx) + for child_idx, child in enumerate(children): + child_expert = child if num_experts is None else child.get_expert(expert_idx) + for name, loader in zip(fields, loaders): + tensor = getattr(child_expert, name) + if tensor is not None: + loader(torch.full_like(tensor, expert_idx * 2 + child_idx + 1), child_expert) + assert all(child_expert.load_ok) + for name, shape, split_dim in zip(fields, shapes, split_dims): + if shape is None: + continue + child_shape = list(shape) + child_shape[split_dim] //= 2 + expected = torch.cat( + [ + torch.full(child_shape, expert_idx * 2 + child_idx + 1, device=f"cuda:{device_id}") + for child_idx in range(2) + ], + dim=split_dim, + ) + torch.testing.assert_close(getattr(parent_expert, name).float(), expected.float()) + + +@pytest.mark.parametrize("redundant", [0, 2]) +@pytest.mark.parametrize("draft", [False, True]) +def test_eplb_one_primary_expert_fp8_constructor_and_checkpoint_load(monkeypatch, device_id, redundant, draft): + if not HAS_DEEPGEMM: + pytest.skip("DeepGEMM is required") + args = StartArgs( + run_mode="prefill", + tp=8, + dp=8, + nnodes=1, + enable_ep_moe=True, + enable_prefill_eplb=True, + eplb_num_redundant_experts_per_rank=redundant, + expert_dtype="fp8", + data_type="bfloat16", + ) + monkeypatch.setattr(fused_moe_weight, "get_env_start_args", lambda: args) + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr(fused_moe_weight, "get_global_rank", lambda: 0) + monkeypatch.setattr(fused_moe_weight, "get_global_world_size", lambda: 8) + monkeypatch.setattr(dist_utils, "get_node_world_size", lambda: 8) + monkeypatch.setattr(base_weight, "get_dp_world_size", lambda: 1) + monkeypatch.setattr(base_weight, "get_current_rank_in_dp", lambda: 0) + monkeypatch.setattr(base_weight, "get_current_device_id", lambda: device_id) + weight = fused_moe_weight.FusedMoeWeight( + gate_proj_name="w1", + down_proj_name="w2", + up_proj_name="w3", + e_score_correction_bias_name="", + weight_prefix="model.experts", + n_routed_experts=8, + hidden_size=128, + moe_intermediate_size=128, + data_type=torch.bfloat16, + quant_method=DeepGEMMFP8w8a8B128QuantizationMethod(), + network_config={ + "num_experts_per_tok": 2, + "norm_topk_prob": True, + "scoring_func": "softmax", + "is_mtp_draft_model": draft, + }, + ) + assert weight.expert_parallel_state.eplb is not None + assert not weight.expert_parallel_state.eplb.recording + assert weight.expert_parallel_state.num_total_physical_experts == 8 * (1 + redundant) + assert weight.local_n_routed_experts == 1 + redundant + assert weight.w13.weight.shape == (1 + redundant, 256, 128) + assert weight.w2.weight.shape == (1 + redundant, 128, 128) + assert weight.w13.weight_scale.shape == (1 + redundant, 2, 1) + assert weight.w2.weight_scale.shape == (1 + redundant, 1, 1) + checkpoint = {} + for expert_idx in range(1 + redundant): + for projection_idx, name in enumerate(("w1", "w2", "w3")): + prefix = f"model.experts.{expert_idx}.{name}" + checkpoint[prefix + ".weight"] = torch.full( + (128, 128), expert_idx * 3 + projection_idx + 1, dtype=torch.float8_e4m3fn + ) + checkpoint[prefix + ".weight_scale_inv"] = torch.full((1, 1), expert_idx * 3 + projection_idx + 1.0) + weight.load_hf_weights(checkpoint) + for projection_idx, packs in enumerate((weight.w1_list, weight.w2_list, weight.w3_list)): + for expert_idx, pack in enumerate(packs): + value = expert_idx * 3 + projection_idx + 1 + assert all(pack.load_ok) + torch.testing.assert_close( + pack.weight.float(), torch.full((128, 128), value, device=f"cuda:{device_id}").float() + ) + torch.testing.assert_close(pack.weight_scale, torch.full((1, 1), float(value), device=f"cuda:{device_id}")) diff --git a/unit_tests/common/fused_moe/test_topk_select.py b/unit_tests/common/fused_moe/test_topk_select.py new file mode 100644 index 0000000000..95b2bada47 --- /dev/null +++ b/unit_tests/common/fused_moe/test_topk_select.py @@ -0,0 +1,47 @@ +from contextlib import nullcontext + +import pytest +import torch +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType +from lightllm.common.basemodel.triton_kernel.fused_moe import topk_select + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +@pytest.mark.parametrize("route", ["global", "grouped", "custom"]) +@pytest.mark.parametrize("phase", [None, AutotuneKernelType.GENERAL, AutotuneKernelType.DECODE_ATTENTION]) +def test_plain_moe_preserves_warmup_routing(monkeypatch, route, phase): + monkeypatch.setattr(Autotuner, "_autotune_warmup_kernel_type", None) + batch, experts, top_k = 7, 8, 2 + logits = torch.arange(experts, device="cuda", dtype=torch.float32).expand(batch, -1).contiguous() + hidden = torch.ones((batch, 16), device="cuda", dtype=torch.bfloat16) + + def custom_routing(**kwargs): + scores = kwargs["gating_output"].softmax(-1) + return torch.topk(scores, k=kwargs["topk"], dim=1) + + kwargs = dict( + hidden_states=hidden, + router_logits=logits, + correction_bias=None, + top_k=top_k, + use_grouped_topk=route == "grouped", + renormalize=False, + topk_group=1, + num_expert_group=2, + custom_routing_function=custom_routing if route == "custom" else None, + ) + with Autotuner.autotune_warmup(AutotuneKernelType.DECODE_ATTENTION): + expected_weights, original_ids = topk_select.select_experts(**kwargs) + context = Autotuner.autotune_warmup(phase) if phase is not None else nullcontext() + with context: + weights, ids = topk_select.select_experts(**kwargs) + torch.testing.assert_close(weights, expected_weights) + if phase == AutotuneKernelType.GENERAL: + generator = torch.Generator(device="cuda").manual_seed(batch) + reference_logits = torch.randn((batch, experts), generator=generator, device="cuda", dtype=torch.float32) + expected_ids = reference_logits.topk(top_k, dim=1).indices + assert not torch.equal(expected_ids, original_ids.long()) + else: + expected_ids = original_ids.long() + torch.testing.assert_close(ids.long(), expected_ids) diff --git a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py index 1fe14a96bc..837013f788 100644 --- a/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py +++ b/unit_tests/common/quantization/test_ue8m0_fp8_quantization.py @@ -378,27 +378,15 @@ def check_activation_scales(*args, **kwargs): @pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") @pytest.mark.parametrize("prepacked", [False, True]) @pytest.mark.parametrize("sm100", [False, True]) -def test_redundancy_expert_update_packs_scales(monkeypatch, packed_weight_layout, prepacked, sm100): - from types import SimpleNamespace - from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_redundancy import ( - FusedMoeWeightEPAutoRedundancy, - ) - +def test_moe_expert_row_load_packs_scales(monkeypatch, packed_weight_layout, prepacked, sm100): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") monkeypatch.setenv("LIGHTLLM_USE_UE8M0_SCALES", "1") monkeypatch.setattr(deepgemm, "is_sm100_gpu", lambda: sm100) method = Quantcfg({"n_layer": 1}, quant_type="fp8w8a8-b128-deepgemm").get_quant_method(0, "q_proj") w13, _ = method.create_moe_weight([128, 128], 512, torch.bfloat16, 0, 4) w2, _ = method.create_moe_weight([512], 128, torch.bfloat16, 0, 4) - update = object.__new__(FusedMoeWeightEPAutoRedundancy) - update.redundancy_expert_num = 2 - update.redundancy_expert_ids = [6, 7] - update._ep_w = SimpleNamespace( - w13=w13, - w2=w2, - quant_method=method, - redundancy_expert_ids_tensor=torch.zeros(2, device="cuda", dtype=torch.int64), - ) expected = [] + before = len(packed_weight_layout) for dest in (w13, w2): dest.weight.zero_() dest.weight_scale.zero_() @@ -409,11 +397,13 @@ def test_redundancy_expert_update_packs_scales(monkeypatch, packed_weight_layout source, _ = method.create_moe_weight([shape[-2]], shape[-1], torch.bfloat16, 0, 2) method.load_weight(raw, source) sf = source.weight_scale + for row in range(2): + expert = dest.get_expert(row + 2) + method.load_weight(q[row], expert) + method.load_weight_scale(sf[row], expert) + assert all(expert.load_ok) expected.append([q, sf]) - update.w13, update.w2 = expected - before = len(packed_weight_layout) - update.commit() - assert len(packed_weight_layout) == before + (2 if sm100 and not prepacked else 0) + assert len(packed_weight_layout) == before + (2 if sm100 and prepacked else 4 if sm100 else 0) for dest, (q, sf) in zip((w13, w2), expected): assert torch.count_nonzero(dest.weight[:2].float()) == 0 assert torch.count_nonzero(dest.weight_scale[:2]) == 0 @@ -429,7 +419,6 @@ def test_redundancy_expert_update_packs_scales(monkeypatch, packed_weight_layout else: actual, expected_sf = dest.weight_scale[2:], sf torch.testing.assert_close(actual, expected_sf, rtol=0, atol=0) - assert update._ep_w.redundancy_expert_ids_tensor.tolist() == [6, 7] @pytest.mark.skipif(not deepgemm.HAS_DEEPGEMM, reason="requires DeepGEMM") diff --git a/unit_tests/common/test_ep_communication.py b/unit_tests/common/test_ep_communication.py new file mode 100644 index 0000000000..0d9edacbba --- /dev/null +++ b/unit_tests/common/test_ep_communication.py @@ -0,0 +1,119 @@ +"""Transport selection and scratch storage must match actual execution paths.""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from lightllm.distributed import communication_op as comm +from lightllm.utils import device_utils + + +@pytest.mark.parametrize( + "tp,node,nvlink,expected", + [ + (8, 4, True, False), + (6, 8, True, False), + (4, 8, True, True), + (2, 4, True, True), + (4, 4, False, False), + (1, 4, True, False), + ], +) +def test_custom_allreduce_requires_whole_node_local_groups(monkeypatch, tp, node, nvlink, expected): + group = comm.CustomProcessGroup.__new__(comm.CustomProcessGroup) + group.dp_world_size = tp + monkeypatch.setattr(comm, "get_node_world_size", lambda: node) + monkeypatch.setattr(comm, "has_nvlink", lambda: nvlink) + assert group._support_custom_allreduce() is expected + + +@pytest.fixture +def buffers(monkeypatch): + args = SimpleNamespace( + enable_ep_moe=True, + enable_prefill_eplb=True, + enable_prefill_cudagraph=False, + enable_rl=False, + eplb_num_redundant_experts_per_rank=2, + data_type="bfloat16", + run_mode="prefill", + disable_cudagraph=True, + ) + monkeypatch.setattr(comm, "get_env_start_args", lambda: args) + monkeypatch.setattr(comm, "get_global_world_size", lambda: 8) + monkeypatch.setattr(comm, "get_deepep_num_max_dispatch_tokens_per_rank_prefill", lambda: 1024) + monkeypatch.setattr(comm, "get_deepep_num_max_dispatch_tokens_per_rank_decode", lambda: 128) + monkeypatch.setattr(comm, "is_sm100_gpu", lambda: False) + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: False) + monkeypatch.setattr(comm, "HAS_DEEPEP", True) + monkeypatch.setattr(comm.dist, "new_group", Mock(return_value=object())) + monkeypatch.setattr(comm.dist, "barrier", Mock()) + monkeypatch.setattr(comm.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(comm, "_calculate_min_chunked_expanded_moe_reuse_buffer_bytes", lambda **_: 769) + legacy_storage = torch.empty(1536, dtype=torch.uint8) + legacy = Mock() + legacy.get_local_buffer_tensor.return_value = legacy_storage + legacy_class = Mock(return_value=legacy) + legacy_class.get_low_latency_rdma_size_hint.return_value = 1281 + elastic = Mock() + elastic.get_theoretical_num_sms.return_value = 24 + monkeypatch.setattr(comm, "deep_ep", SimpleNamespace(Buffer=legacy_class, ElasticBuffer=Mock(return_value=elastic))) + monkeypatch.setattr(comm.DistributeGroupManager, "_set_num_sms_for_deep_gemm", Mock()) + manager = comm.DistributeGroupManager() + manager.groups = [object(), object()] + return args, manager, legacy_class, legacy, legacy_storage + + +def initialize(manager): + manager.new_deepep_group(256, 7168, {"fp8w8a8-b128-deepgemm"}, 8, 2048) + + +@pytest.mark.parametrize( + "mode,disable_graph,eplb,experts", + [ + ("prefill", True, True, 272), + ("prefill", False, True, 272), + ("normal", True, True, 272), + ("normal", False, True, 272), + ("decode", True, True, 256), + ("decode", False, True, 256), + ("prefill", True, False, 256), + ("prefill", False, False, 256), + ("normal", True, False, 256), + ("normal", False, False, 256), + ("decode", True, False, 256), + ("decode", False, False, 256), + ], +) +def test_eplb_preserves_transport_and_uses_physical_expert_count(buffers, mode, disable_graph, eplb, experts): + args, manager, constructor, legacy, storage = buffers + args.run_mode, args.disable_cudagraph = mode, disable_graph + args.enable_prefill_eplb = eplb + initialize(manager) + constructor.assert_called_once() + assert constructor.call_args.kwargs["num_rdma_bytes"] == 1536 + assert constructor.call_args.kwargs["num_qps_per_rank"] == experts // 8 + assert manager.ll_num_experts == experts + constructor.get_low_latency_rdma_size_hint.assert_called_once_with(128, 7168, 8, experts) + manager.ep_buffer.get_theoretical_num_sms.assert_called_once_with(experts, 8) + manager._set_num_sms_for_deep_gemm.assert_called_once_with(24) + first, second = [manager.get_deep_ep_prefill_moe_workspace(i) for i in range(2)] + assert first.numel() == second.numel() == 768 + assert first.data_ptr() == storage.data_ptr() + assert second.data_ptr() == storage.data_ptr() + 768 + with pytest.raises(AssertionError): + manager.get_deep_ep_prefill_moe_workspace(2) + manager.clear_deepep_buffer() + legacy.clean_low_latency_buffer.assert_called_once_with(128, 7168, experts) + + +def test_disabling_ep_clears_buffers(buffers): + args, manager, constructor, _, _ = buffers + initialize(manager) + args.enable_ep_moe = False + initialize(manager) + assert manager.ep_low_latency_buffer is None + assert manager.ep_buffer is None + constructor.assert_called_once() diff --git a/unit_tests/server/router/model_infer/test_ep_balance_monitor.py b/unit_tests/server/router/model_infer/test_ep_balance_monitor.py new file mode 100644 index 0000000000..e411844730 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_ep_balance_monitor.py @@ -0,0 +1,624 @@ +import threading +from array import array +from datetime import timedelta +from types import SimpleNamespace + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from prometheus_client import generate_latest + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_balance import PrefillEPBalanceCounters +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight as fused_weight_module +from lightllm.distributed import communication_op as communication_op_module +from lightllm.server.metrics.metrics import Monitor +from lightllm.server.router.model_infer.mode_backend import ep_balance_monitor as monitor_module +from lightllm.server.router.model_infer.mode_backend.ep_balance_monitor import ( + EP_BALANCE_PRESSURE_DRIFT_BUCKET_ROUNDS, + calculate_prefill_placement_pressure_drift, + calculate_prefill_balance_stats, + classify_prefill_placement_pressure_drift, + should_enable_ep_balance_monitor, +) + + +@pytest.fixture(autouse=True) +def _mock_non_sm100(monkeypatch): + monkeypatch.setattr(monitor_module, "is_sm100_gpu", lambda: False) + + +def _stats(source_token_replication: int): + return calculate_prefill_balance_stats( + torch.tensor([[[[800, 40], [800, 20]], [[800, 20], [800, 40]]]], dtype=torch.int64), + layer_routed_experts=torch.tensor([1, 1]), + layer_flops_per_expert_token=torch.tensor([2.0, 4.0]), + layer_topks=torch.tensor([2.0, 2.0]), + source_token_replication=source_token_replication, + ) + + +def _monitor_args(**overrides): + args = { + "enable_ep_moe": True, + "disable_ep_balance_monitor": False, + "run_mode": "normal", + "enable_prefill_cudagraph": False, + } + args.update(overrides) + return SimpleNamespace(**args) + + +def _pressure_round_stats(bucket_rank_loads): + """Build [round, layer=1, rank, route/compute] CPU samples for drift tests.""" + return torch.tensor([[[[0, load] for load in rank_loads]] for rank_loads in bucket_rank_loads], dtype=torch.int64) + + +def test_pressure_drift_is_zero_for_identical_pressure(): + drift, _ = calculate_prefill_placement_pressure_drift(_pressure_round_stats([[2, 1], [2, 1]]), bucket_rounds=1) + assert drift == 0.0 + + +def test_pressure_drift_is_one_for_complete_hot_rank_migration(): + drift, _ = calculate_prefill_placement_pressure_drift(_pressure_round_stats([[2, 0], [0, 2]]), bucket_rounds=1) + assert drift == 1.0 + + +def test_pressure_drift_tracks_same_rank_magnitude_change(): + drift, _ = calculate_prefill_placement_pressure_drift(_pressure_round_stats([[2, 1], [3, 1]]), bucket_rounds=1) + assert drift == pytest.approx(0.2) + + +def test_pressure_drift_is_invariant_to_uniform_load_scale(): + drift, _ = calculate_prefill_placement_pressure_drift(_pressure_round_stats([[2, 1], [4, 2]]), bucket_rounds=1) + assert drift == 0.0 + + +def test_pressure_drift_is_zero_for_balanced_inputs(): + drift, _ = calculate_prefill_placement_pressure_drift(_pressure_round_stats([[8, 8], [16, 16]]), bucket_rounds=1) + assert drift == 0.0 + + +def test_pressure_drift_previous_signature_bridges_report_boundary(): + _, signature = calculate_prefill_placement_pressure_drift(_pressure_round_stats([[2, 0]]), bucket_rounds=1) + drift, next_signature = calculate_prefill_placement_pressure_drift( + _pressure_round_stats([[0, 2]]), previous_pressure_signature=signature, bucket_rounds=1 + ) + assert drift == 1.0 + assert torch.equal(next_signature, torch.tensor([[0.0, 1.0]], dtype=torch.float64)) + + +def test_pressure_drift_default_bucket_handles_normal_report_window(): + round_stats = _pressure_round_stats([[4, 2]] * 100) + drift, signature = calculate_prefill_placement_pressure_drift(round_stats) + assert EP_BALANCE_PRESSURE_DRIFT_BUCKET_ROUNDS == 20 + assert drift == 0.0 + assert signature.shape == (1, 2) + + +@pytest.mark.parametrize( + ("drift", "expected"), + [ + (0.0, "stable"), + (0.0999, "stable"), + (0.10, "dynamic"), + (0.2999, "dynamic"), + (0.30, "dynamic"), + (0.75, "dynamic"), + (1.0, "dynamic"), + ], +) +def test_pressure_drift_classification_boundaries(drift, expected): + assert classify_prefill_placement_pressure_drift(drift) == expected + + +def test_monitor_log_stats_reports_pressure_drift_and_bridges_reports(monkeypatch): + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor.layer_routed_experts = torch.tensor([1], dtype=torch.int64) + monitor.layer_flops_per_expert_token = torch.tensor([2.0], dtype=torch.float64) + monitor.layer_topks = torch.tensor([2.0], dtype=torch.float64) + monitor.source_token_replication = 1 + monitor._previous_pressure_signature = None + metric_calls = [] + monitor.metric_client = SimpleNamespace(gauge_set=lambda name, value: metric_calls.append((name, value))) + logs = [] + monkeypatch.setattr(monitor_module.logger, "info", logs.append) + + def report_for_hot_rank(hot_rank): + round_stats = torch.zeros((100, 1, 2, 2), dtype=torch.int64) + round_stats[:, :, :, monitor_module.ROUTE_LOAD] = 100 + round_stats[:, :, hot_rank, monitor_module.COMPUTE_LOAD] = 2 + monitor._log_stats(round_stats) + + report_for_hot_rank(0) + first_signature = monitor._previous_pressure_signature.clone() + report_for_hot_rank(1) + + assert "prefill_ep_placement_pressure_drift=0.0000" in logs[0] + assert "prefill_ep_placement_pressure_state=stable" in logs[0] + assert "prefill_ep_placement_pressure_drift=0.2000" in logs[1] + assert "prefill_ep_placement_pressure_state=dynamic" in logs[1] + assert torch.equal(first_signature, torch.tensor([[1.0, 0.0]], dtype=torch.float64)) + assert torch.equal(monitor._previous_pressure_signature, torch.tensor([[0.0, 1.0]], dtype=torch.float64)) + assert metric_calls == [ + ("lightllm_prefill_ep_critical_overhead_gflops_per_routed_token", pytest.approx(2e-11)), + ("lightllm_prefill_ep_compute_critical_overhead_ratio", pytest.approx(1.0)), + ("lightllm_prefill_ep_placement_pressure_drift", 0.0), + ("lightllm_prefill_ep_critical_overhead_gflops_per_routed_token", pytest.approx(2e-11)), + ("lightllm_prefill_ep_compute_critical_overhead_ratio", pytest.approx(1.0)), + ("lightllm_prefill_ep_placement_pressure_drift", pytest.approx(0.2)), + ] + + +def test_critical_overhead_preserves_per_layer_slowest_rank(): + stats = _stats(source_token_replication=1) + assert stats is not None + assert stats["critical_overhead_gflops_per_routed_token"] == pytest.approx(7.5e-11) + assert stats["prefill_ep_compute_critical_overhead_ratio"] == pytest.approx(1 / 3) + + +def test_non_tpsp_tp_replication_scales_gflops_per_token_but_not_ratio(): + tpsp_stats = _stats(source_token_replication=1) + non_tpsp_tp8_stats = _stats(source_token_replication=8) + assert tpsp_stats is not None and non_tpsp_tp8_stats is not None + assert non_tpsp_tp8_stats["critical_overhead_gflops_per_routed_token"] == pytest.approx( + tpsp_stats["critical_overhead_gflops_per_routed_token"] * 8 + ) + assert non_tpsp_tp8_stats["prefill_ep_compute_critical_overhead_ratio"] == pytest.approx( + tpsp_stats["prefill_ep_compute_critical_overhead_ratio"] + ) + + +def test_critical_overhead_is_zero_when_ranks_are_balanced(): + stats = calculate_prefill_balance_stats( + torch.tensor([[[[100, 32], [100, 32]], [[100, 64], [100, 64]]]], dtype=torch.int64), + layer_routed_experts=torch.tensor([1, 1]), + layer_flops_per_expert_token=torch.tensor([2.0, 4.0]), + layer_topks=torch.tensor([2.0, 2.0]), + source_token_replication=1, + ) + assert stats is not None + assert stats["critical_overhead_gflops_per_routed_token"] == 0.0 + assert stats["prefill_ep_compute_critical_overhead_ratio"] == 0.0 + + +@pytest.mark.parametrize( + "round_stats", + [ + torch.tensor([[[[1, 1], [1, 1]]]], dtype=torch.int64), + torch.tensor([[[[100, 0], [100, 0]]]], dtype=torch.int64), + ], +) +def test_critical_overhead_rejects_insufficient_or_zero_compute_samples(round_stats): + assert ( + calculate_prefill_balance_stats( + round_stats, + layer_routed_experts=torch.tensor([1]), + layer_flops_per_expert_token=torch.tensor([2.0]), + layer_topks=torch.tensor([2.0]), + source_token_replication=1, + ) + is None + ) + + +def test_cpu_counter_accumulates_multiple_prefill_dispatches(): + counters = PrefillEPBalanceCounters() + counters.accumulate(route_load=3, compute_load=128) + counters.accumulate(route_load=4, compute_load=256) + assert (counters.route_load, counters.compute_load) == (7, 384) + + +def test_monitor_creates_dedicated_gloo_group_before_starting_thread(monkeypatch): + sentinel_group = object() + impl = SimpleNamespace(ep_balance_counters=None) + weight = SimpleNamespace( + fuse_moe_impl=impl, + n_routed_experts=8, + hidden_size=16, + moe_intermediate_size=32, + num_experts_per_tok=2, + ) + model = SimpleNamespace(args=SimpleNamespace(enable_tpsp_mix_mode=True), tp_world_size_=1) + + monkeypatch.setattr(monitor_module, "find_ep_moe_weights", lambda _: [weight]) + monkeypatch.setattr(monitor_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(monitor_module, "get_global_world_size", lambda: 2) + metric_ports = [] + monkeypatch.setattr(monitor_module, "get_shm_port_args", lambda: SimpleNamespace(metric_port=4321)) + monkeypatch.setattr(monitor_module, "MetricClient", lambda port: metric_ports.append(port) or object()) + events = [] + + def create_group(*, ranks, backend): + assert ranks == [0, 1] + assert backend == "gloo" + events.append("group") + return sentinel_group + + monkeypatch.setattr(monitor_module.dist, "new_group", create_group) + monkeypatch.setattr( + monitor_module.threading, + "Thread", + lambda *args, **kwargs: SimpleNamespace(start=lambda: events.append("thread")), + ) + + monitor = monitor_module.EPBalanceMonitor(model) + + assert monitor.gloo_group is sentinel_group + assert impl.ep_balance_counters is monitor.counters[0] + assert metric_ports == [4321] + assert events == ["group", "thread"] + + +def test_monitor_without_experts_does_not_allocate_resources(monkeypatch): + monkeypatch.setattr(monitor_module, "get_global_rank", lambda: 0) + monkeypatch.setattr(monitor_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(monitor_module, "find_ep_moe_weights", lambda _: []) + monkeypatch.setattr(monitor_module.dist, "new_group", lambda **kwargs: pytest.fail("unexpected group")) + monkeypatch.setattr(monitor_module, "MetricClient", lambda port: pytest.fail("unexpected metric client")) + monkeypatch.setattr(monitor_module.threading, "Thread", lambda **kwargs: pytest.fail("unexpected thread")) + + monitor = monitor_module.EPBalanceMonitor(SimpleNamespace()) + + assert not monitor.enabled + + +def test_nonzero_rank_monitor_does_not_create_metric_client(monkeypatch): + sentinel_group = object() + impl = SimpleNamespace(ep_balance_counters=None) + weight = SimpleNamespace( + fuse_moe_impl=impl, + n_routed_experts=8, + hidden_size=16, + moe_intermediate_size=32, + num_experts_per_tok=2, + ) + model = SimpleNamespace(args=SimpleNamespace(enable_tpsp_mix_mode=True), tp_world_size_=1) + metric_client_calls = [] + monkeypatch.setattr(monitor_module, "find_ep_moe_weights", lambda _: [weight]) + monkeypatch.setattr(monitor_module, "get_global_rank", lambda: 1) + monkeypatch.setattr(monitor_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(monitor_module.dist, "new_group", lambda **kwargs: sentinel_group) + monkeypatch.setattr(monitor_module, "get_shm_port_args", lambda: pytest.fail("unexpected port lookup")) + monkeypatch.setattr( + monitor_module, + "MetricClient", + lambda port: metric_client_calls.append(port) or pytest.fail("unexpected metric client"), + ) + monkeypatch.setattr( + monitor_module.threading, + "Thread", + lambda *args, **kwargs: SimpleNamespace(start=lambda: None), + ) + + monitor = monitor_module.EPBalanceMonitor(model) + + assert monitor.metric_client is None + assert metric_client_calls == [] + + +@pytest.mark.parametrize("disable_monitor", [False, True]) +def test_common_group_creation_does_not_allocate_monitor_resources(monkeypatch, disable_monitor): + custom_groups = [] + + class FakeCustomProcessGroup: + def init_symm_mem_reduce(self): + pass + + def init_flashinfer_reduce(self): + pass + + args = SimpleNamespace( + enable_ep_moe=True, + disable_ep_balance_monitor=disable_monitor, + run_mode="normal", + enable_prefill_cudagraph=False, + disable_symm_mem_allreduce=True, + disable_flashinfer_allreduce=True, + ) + monkeypatch.setattr(communication_op_module, "get_env_start_args", lambda: args) + monkeypatch.setattr( + communication_op_module, + "CustomProcessGroup", + lambda: custom_groups.append(FakeCustomProcessGroup()) or custom_groups[-1], + ) + monkeypatch.setattr(communication_op_module, "get_global_world_size", lambda: 2) + monkeypatch.setattr(communication_op_module, "is_sm100_gpu", lambda: False) + calls = [] + monkeypatch.setattr( + communication_op_module.dist, + "new_group", + lambda *args, **kwargs: calls.append((args, kwargs)), + ) + + manager = communication_op_module.DistributeGroupManager() + manager.create_groups(group_size=2) + + assert len(manager.groups) == 2 + assert calls == [] + + +def test_monitor_registers_prefill_ep_gauges_with_model_label(): + monitor = Monitor( + SimpleNamespace( + metric_gateway=None, + job_name="test", + grouping_key=[], + enable_monitor_auth=False, + model_name="monitor-test-model", + max_req_total_len=128, + mtp_step=0, + ) + ) + values = { + "lightllm_prefill_ep_critical_overhead_gflops_per_routed_token": 1.25, + "lightllm_prefill_ep_compute_critical_overhead_ratio": 0.3, + "lightllm_prefill_ep_placement_pressure_drift": 0.125, + } + assert set(values).issubset(monitor.monitor_registry) + for name, value in values.items(): + monitor.gauge_set(name, value) + + exposition = generate_latest(monitor.registry).decode() + for name, value in values.items(): + assert f'{name}{{model_name="monitor-test-model"}} {value}' in exposition + + +def test_record_prefill_round_stores_cumulative_counter_deltas_in_ring_buffer(): + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor.enabled = True + monitor.counters = [PrefillEPBalanceCounters()] + monitor._round_buffer_storage = array("q", [0]) * (monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY * 2) + monitor._round_buffer = torch.frombuffer(monitor._round_buffer_storage, dtype=torch.int64).view( + monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY, 1, 2 + ) + monitor._round_ready = threading.Event() + monitor._written_round_count = 0 + monitor._processed_round_count = 0 + monitor._overflowed = False + + monitor.counters[0].accumulate(route_load=3, compute_load=128) + monitor.record_prefill_round() + assert (monitor.counters[0].route_load, monitor.counters[0].compute_load) == (0, 0) + monitor.counters[0].accumulate(route_load=2, compute_load=256) + monitor.record_prefill_round() + + assert torch.equal( + monitor._copy_local_rounds(0, 2), + torch.tensor([[[3, 128]], [[2, 256]]], dtype=torch.int64), + ) + assert not hasattr(monitor, "_round_lock") + + +def test_spsc_ring_copy_wraps_without_a_lock(): + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor.enabled = True + monitor.counters = [PrefillEPBalanceCounters()] + monitor._round_buffer_storage = array("q", [0]) * (monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY * 2) + monitor._round_buffer = torch.frombuffer(monitor._round_buffer_storage, dtype=torch.int64).view( + monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY, 1, 2 + ) + monitor._round_ready = threading.Event() + monitor._written_round_count = monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY - 2 + monitor._processed_round_count = monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY - 2 + monitor._overflowed = False + + for value in (11, 12, 13): + monitor.counters[0].accumulate(route_load=value, compute_load=value * 10) + monitor.record_prefill_round() + + assert torch.equal( + monitor._copy_local_rounds(monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY - 2, monitor._written_round_count), + torch.tensor([[[11, 110]], [[12, 120]], [[13, 130]]], dtype=torch.int64), + ) + assert not hasattr(monitor, "_round_lock") + + +def test_spsc_ring_overflow_is_deferred_to_the_monitor_thread(): + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor.enabled = True + monitor.counters = [PrefillEPBalanceCounters(route_load=7, compute_load=70)] + monitor._round_buffer_storage = array("q", [0]) * (monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY * 2) + monitor._round_buffer = torch.frombuffer(monitor._round_buffer_storage, dtype=torch.int64).view( + monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY, 1, 2 + ) + monitor._round_ready = threading.Event() + monitor._written_round_count = monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY + monitor._processed_round_count = 0 + monitor._overflowed = False + + monitor.record_prefill_round() + + assert monitor._overflowed + assert monitor._round_ready.is_set() + assert monitor._written_round_count == monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY + assert (monitor.counters[0].route_load, monitor.counters[0].compute_load) == (7, 70) + + +def test_raise_buffer_overflow_always_reports_phase_and_ring_counts(): + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor._written_round_count = 23 + monitor._processed_round_count = 7 + + with pytest.raises(RuntimeError) as exc_info: + monitor._raise_buffer_overflow("before_sync") + + assert str(exc_info.value) == ( + "EP balance prefill-round buffer overflowed " + f"phase=before_sync written=23 processed=7 capacity={monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY}" + ) + + +def test_raise_buffer_overflow_optionally_reports_common_round_end(): + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor._written_round_count = 23 + monitor._processed_round_count = 7 + + with pytest.raises(RuntimeError) as exc_info: + monitor._raise_buffer_overflow("common_round_lag", common_round_end=19) + + assert str(exc_info.value) == ( + "EP balance prefill-round buffer overflowed " + f"phase=common_round_lag written=23 processed=7 capacity={monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY} " + "common_round_end=19" + ) + + +def test_gather_round_stats_only_allocates_receive_buffers_on_rank_zero(monkeypatch): + local_round_stats = torch.tensor([[[3, 128]]], dtype=torch.int64) + sentinel_group = object() + + rank_zero_monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + rank_zero_monitor.global_rank = 0 + rank_zero_monitor.world_size = 2 + rank_zero_monitor.gloo_group = sentinel_group + + def root_gather(input_tensor, gather_list, dst, group): + assert dst == 0 and group is sentinel_group + assert len(gather_list) == 2 + gather_list[0].copy_(input_tensor) + gather_list[1].copy_(input_tensor + 1) + + monkeypatch.setattr(monitor_module.dist, "gather", root_gather) + result = rank_zero_monitor._gather_round_stats(local_round_stats) + assert torch.equal(result, torch.tensor([[[[3, 128], [4, 129]]]], dtype=torch.int64)) + + nonzero_monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + nonzero_monitor.global_rank = 1 + nonzero_monitor.world_size = 2 + nonzero_monitor.gloo_group = sentinel_group + + def nonroot_gather(input_tensor, gather_list, dst, group): + assert input_tensor is local_round_stats + assert gather_list is None + assert dst == 0 and group is sentinel_group + + monkeypatch.setattr(monitor_module.dist, "gather", nonroot_gather) + assert nonzero_monitor._gather_round_stats(local_round_stats) is None + + +def test_find_fused_moe_weights_discovers_any_layer_member_once_and_sorts(monkeypatch): + class FakeFusedMoeWeight: + def __init__(self, layer_num, enabled=True): + self.layer_num_ = layer_num + self.enable_ep_moe = enabled + + monkeypatch.setattr(fused_weight_module, "FusedMoeWeight", FakeFusedMoeWeight) + first = FakeFusedMoeWeight(3) + second = FakeFusedMoeWeight(1) + disabled = FakeFusedMoeWeight(0, enabled=False) + model = SimpleNamespace( + trans_layers_weight=[ + SimpleNamespace(experts_=first, alias=first, ignored=disabled), + SimpleNamespace(any_direct_member=second), + ] + ) + + assert monitor_module.find_ep_moe_weights(model) == [second, first] + + +def test_monitor_disable_detaches_counters_from_all_impls(): + impl = SimpleNamespace(ep_balance_counters="unset") + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor.weights = [SimpleNamespace(fuse_moe_impl=impl)] + monitor.enabled = True + monitor._disable() + assert impl.ep_balance_counters is None + assert not monitor.enabled + + +def test_critical_overhead_requires_minimum_samples_for_every_layer(): + round_stats = torch.tensor([[[[200, 32], [200, 32]], [[1, 32], [1, 32]]]], dtype=torch.int64) + assert ( + calculate_prefill_balance_stats( + round_stats, + layer_routed_experts=torch.tensor([1, 1]), + layer_flops_per_expert_token=torch.tensor([2.0, 4.0]), + layer_topks=torch.tensor([2.0, 2.0]), + source_token_replication=1, + ) + is None + ) + + +def test_ep_moe_normal_and_prefill_enable_monitor_by_default(): + assert should_enable_ep_balance_monitor(_monitor_args()) + assert should_enable_ep_balance_monitor(_monitor_args(run_mode="prefill")) + + +def test_disable_ep_balance_monitor_turns_monitor_off(): + assert not should_enable_ep_balance_monitor(_monitor_args(disable_ep_balance_monitor=True)) + + +def test_non_ep_moe_and_decode_mode_do_not_enable_monitor(): + assert not should_enable_ep_balance_monitor(_monitor_args(enable_ep_moe=False)) + assert not should_enable_ep_balance_monitor(_monitor_args(run_mode="decode")) + + +def test_prefill_cudagraph_silently_disables_monitor(): + assert not should_enable_ep_balance_monitor(_monitor_args(run_mode="prefill", enable_prefill_cudagraph=True)) + + +def test_sm100_silently_disables_monitor(monkeypatch): + monkeypatch.setattr(monitor_module, "is_sm100_gpu", lambda: True) + assert not should_enable_ep_balance_monitor(_monitor_args()) + + +def _monitor_failure_worker(rank, rendezvous, failure): + torch.set_num_threads(1) + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=2, timeout=timedelta(seconds=15)) + control_group = dist.new_group(backend="gloo", timeout=timedelta(seconds=15)) + monitor = monitor_module.EPBalanceMonitor.__new__(monitor_module.EPBalanceMonitor) + monitor.global_rank = rank + monitor.world_size = 2 + monitor.gloo_group = dist.group.WORLD + monitor._common_round_end = torch.zeros((), dtype=torch.int64) + monitor._monitor_error = torch.zeros((), dtype=torch.int32) + monitor._round_buffer = torch.ones((monitor_module.EP_BALANCE_ROUND_BUFFER_CAPACITY, 1, 2), dtype=torch.int64) + monitor._written_round_count = 200 + monitor._processed_round_count = 0 + monitor._overflowed = failure == "overflow" and rank == 1 + monitor._round_ready = threading.Event() + monitor._round_ready.set() + monitor.enabled = True + impl = SimpleNamespace(ep_balance_counters=object()) + monitor.weights = [SimpleNamespace(fuse_moe_impl=impl)] + errors = [] + monitor_module.logger.exception = lambda message: errors.append(str(message)) + + def fail(*args): + raise RuntimeError(f"injected {failure} failure") + + if failure == "copy" and rank == 1: + monitor._copy_local_rounds = fail + if failure == "after_copy" and rank == 1: + copy_rounds = monitor._copy_local_rounds + + def overflow_after_copy(start, end): + result = copy_rounds(start, end) + monitor._overflowed = True + return result + + monitor._copy_local_rounds = overflow_after_copy + monitor._log_stats = fail if failure == "log" else lambda stats: None + thread = threading.Thread(target=monitor._monitor_loop, daemon=True) + thread.start() + thread.join(timeout=5) + stuck = thread.is_alive() + # Keep the group alive until both ranks have measured the worker's exit. + dist.barrier(group=control_group) + dist.destroy_process_group() + assert not stuck, f"rank {rank}: monitor was left waiting in a collective ({errors})" + assert not monitor.enabled and impl.ep_balance_counters is None + if rank == (0 if failure == "log" else 1): + assert any( + ("overflow" if failure in ("overflow", "after_copy") else f"injected {failure}") in error + for error in errors + ) + else: + assert any("another rank" in error for error in errors), errors + + +@pytest.mark.parametrize("failure", ["overflow", "copy", "after_copy", "log"]) +def test_monitor_failure_stops_all_ranks(tmp_path, failure): + mp.spawn(_monitor_failure_worker, args=((tmp_path / "gloo-init").as_uri(), failure), nprocs=2, join=True) diff --git a/unit_tests/server/router/model_infer/test_ep_runtime.py b/unit_tests/server/router/model_infer/test_ep_runtime.py new file mode 100644 index 0000000000..6aefaf8144 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_ep_runtime.py @@ -0,0 +1,331 @@ +import threading +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.server.router.model_infer.mode_backend import ep_runtime, eplb_manager +from lightllm.server.router.model_infer.mode_backend.chunked_prefill.impl import ChunkedPrefillBackend +from lightllm.server.router.model_infer.mode_backend.dp_backend.impl import DPChunkedPrefillBackend +from lightllm.server.router.model_infer.mode_backend.overlap_events import OverlapEventManager +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import device_utils + + +@pytest.fixture +def runtime(monkeypatch): + monkeypatch.setenv("LIGHTLLM_CURRENT_DEVICE_ID", "0") + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: False) + events = [] + model = TpPartBaseModel.__new__(TpPartBaseModel) + output = object() + model._prefill = lambda model_input: events.append("prefill") or output + model._decode = lambda model_input: events.append("decode") or output + model._microbatch_overlap_prefill_cuda = lambda *inputs: events.append("overlap") or output + model.microbatch_overlap_decode = lambda *inputs: output + monitor = SimpleNamespace(enabled=True, record_prefill_round=lambda: events.append("record")) + manager = SimpleNamespace( + before_prefill=lambda: events.append("prepare"), + step=lambda: events.append("step"), + poll=lambda: events.append("poll"), + _eplb_states=[SimpleNamespace(migration_poll=None)], + poll_migration_at_layer_boundary=lambda: events.append("layer_poll"), + ) + backend = SimpleNamespace( + args=StartArgs( + enable_prefill_eplb=True, + enable_ep_moe=True, + run_mode="prefill", + disable_ep_balance_monitor=False, + enable_prefill_microbatch_overlap=False, + enable_tpsp_mix_mode=False, + ), + model=model, + overlap_event_manager=OverlapEventManager(), + ) + monkeypatch.setattr(ep_runtime, "EPBalanceMonitor", lambda model: monitor) + monkeypatch.setattr( + ep_runtime, "should_enable_ep_balance_monitor", lambda args: not args.disable_ep_balance_monitor + ) + monkeypatch.setattr(eplb_manager, "EPLBManager", lambda model: manager) + monkeypatch.setattr(ep_runtime.dist, "barrier", lambda: events.append("barrier")) + return backend, events, output, monitor, manager + + +@pytest.mark.parametrize("eplb,monitor_enabled", [(False, False), (False, True), (True, False), (True, True)]) +@pytest.mark.parametrize("batch_size", [0, 3]) +@pytest.mark.parametrize("run_mode", ["prefill", "normal"]) +def test_completed_rounds_and_decode_preserve_behavior(runtime, eplb, monitor_enabled, batch_size, run_mode): + backend, events, output, monitor, _ = runtime + backend.args.enable_prefill_eplb = eplb + backend.args.run_mode = run_mode + backend.args.disable_ep_balance_monitor = not monitor_enabled + model = backend.model + original_methods = vars(model).copy() + original_event_methods = vars(backend.overlap_event_manager).copy() + ep_runtime.install_ep_runtime(backend) + assert events == (["barrier"] if eplb else []) + events.clear() + + model_input = SimpleNamespace(is_prefill=True, batch_size=batch_size, to_cuda=lambda: None) + expected = (["record"] if monitor_enabled else []) + (["step"] if eplb else []) + assert model.forward(model_input) is output + preparation = (["poll"] if run_mode == "normal" else []) + ["prepare"] if eplb else [] + assert events == [*preparation, "prefill", *expected] + events.clear() + # Empty DP ranks and both empty microbatches still participate in the + # globally scheduled round. Count one full round, not one per microbatch. + assert model._microbatch_overlap_prefill_cuda(model_input, model_input) is output + assert events == [*preparation, "overlap", *expected] + events.clear() + model_input.is_prefill = False + assert model.forward(model_input) is output + assert events == ["decode"] + assert model._decode is original_methods["_decode"] + assert model.microbatch_overlap_decode is original_methods["microbatch_overlap_decode"] + if not eplb or run_mode == "normal": + assert vars(backend.overlap_event_manager) == original_event_methods + if not eplb and not monitor_enabled: + assert vars(model) == original_methods + + +def test_disabled_monitor_and_draft_methods_are_untouched(runtime): + backend, events, _, monitor, _ = runtime + backend.args.enable_prefill_eplb = False + monitor.enabled = False + original_methods = vars(backend.model).copy() + ep_runtime.install_ep_runtime(backend) + assert vars(backend.model) == original_methods + assert not events + # The installer binds only the supplied target model, never its class. + draft = TpPartBaseModel.__new__(TpPartBaseModel) + backend.args.enable_prefill_eplb = True + ep_runtime.install_ep_runtime(backend) + assert "_prefill" not in vars(draft) + assert draft._prefill.__func__ is TpPartBaseModel._prefill + + +def test_normal_decode_and_idle_never_poll_or_wait_for_eplb(runtime): + backend, events, output, _, manager = runtime + backend.args.run_mode = "normal" + + def forbidden(): + pytest.fail("decode/idle must not poll, commit, or wait for EPLB") + + manager.poll = forbidden + manager.before_prefill = forbidden + manager.step = forbidden + manager.poll_migration_at_layer_boundary = forbidden + ep_runtime.install_ep_runtime(backend) + events.clear() + event_pack = backend.overlap_event_manager.get_overlap_event_pack() + event_pack._close_overlap() + event_pack.wait_to_forward() + assert not events + assert backend.model.forward(SimpleNamespace(is_prefill=False, to_cuda=lambda: None)) is output + assert events == ["decode"] + + +def test_pd_decode_never_installs_eplb_even_when_requested(runtime, monkeypatch): + backend, events, output, _, _ = runtime + backend.args.run_mode = "decode" + backend.args.disable_ep_balance_monitor = True + + def forbidden(*args): + pytest.fail("PD decode must not construct or synchronize EPLB") + + monkeypatch.setattr(eplb_manager, "EPLBManager", forbidden) + monkeypatch.setattr(ep_runtime.dist, "barrier", forbidden) + original_methods = vars(backend.model).copy() + original_event_methods = vars(backend.overlap_event_manager).copy() + ep_runtime.install_ep_runtime(backend) + assert vars(backend.model) == original_methods + assert vars(backend.overlap_event_manager) == original_event_methods + assert not events + assert backend.model.forward(SimpleNamespace(is_prefill=False, to_cuda=lambda: None)) is output + assert events == ["decode"] + + +@pytest.mark.parametrize("method", ["_prefill", "_microbatch_overlap_prefill_cuda"]) +def test_failed_prefill_is_not_published(runtime, method): + backend, events, _, _, _ = runtime + + def fail(*args, **kwargs): + raise ValueError("forward failed") + + setattr(backend.model, method, fail) + ep_runtime.install_ep_runtime(backend) + events.clear() + with pytest.raises(ValueError, match="forward failed"): + getattr(backend.model, method)(object()) + assert events == ["prepare"] + + +def test_poll_waits_for_execution_ownership_and_propagates_failures(runtime): + backend, events, _, _, manager = runtime + ep_runtime.install_ep_runtime(backend) + first = backend.overlap_event_manager.get_overlap_event_pack() + second = backend.overlap_event_manager.get_overlap_event_pack() + first._close_overlap() + second._close_overlap() + events.clear() + started = threading.Event() + finished = threading.Event() + + def run_second(): + started.set() + second.wait_to_forward() + second.notify_post_handle_event.set() + finished.set() + + thread = threading.Thread(target=run_second, daemon=True) + thread.start() + assert started.wait(2) + assert not finished.wait(0.02) + assert not events + first.wait_to_forward() + assert events == ["poll"] + first.notify_pre_post_handle() + assert finished.wait(2) + thread.join(2) + assert not thread.is_alive() + assert events == ["poll", "poll"] + + def fail(): + raise RuntimeError("poll failed") + + manager.poll = fail + with pytest.raises(RuntimeError, match="poll failed"): + second.wait_to_forward() + + +def test_poll_binding_does_not_retain_completed_event_packs(runtime): + backend, _, _, _, _ = runtime + ep_runtime.install_ep_runtime(backend) + event_pack = backend.overlap_event_manager.get_overlap_event_pack() + reference = weakref.ref(event_pack) + del event_pack + assert reference() is None + + +@pytest.mark.parametrize("backend_class", [ChunkedPrefillBackend, DPChunkedPrefillBackend]) +@pytest.mark.parametrize("run_mode", ["prefill", "normal"]) +def test_real_infer_loops_poll_only_at_role_specific_boundaries(runtime, monkeypatch, backend_class, run_mode): + template, events, _, _, _ = runtime + backend = backend_class.__new__(backend_class) + vars(backend).update(vars(template)) + backend.args.run_mode = run_mode + backend.support_overlap = False + backend.classed_req_no_decode = False + backend.classed_req_strict_prefill = False + backend.processed_prefill_tokens = 0 + backend.logger = SimpleNamespace(exception=lambda message: None) + actions = iter(("prefill", "decode", "pass")) + backend.control_state_machine = SimpleNamespace( + try_recover_paused_reqs=lambda: False, + select_run_way=lambda **kwargs: _run_way(next(actions)), + ) + backend._get_classed_reqs = lambda **kwargs: ([], [], 0) + get_event_pack = backend.overlap_event_manager.get_overlap_event_pack + + def get_ready_event_pack(): + event_pack = get_event_pack() + # Simulate the other thread's completed post-handle stage; this test + # drives one real infer_loop. The test above checks actual handoff. + event_pack.wait_post_handle_event.set() + return event_pack + + backend.overlap_event_manager.get_overlap_event_pack = get_ready_event_pack + + class StopLoop(BaseException): + pass + + def read_requests(): + events.append("requests") + if events.count("requests") == 4: + raise StopLoop() + + def gather(**kwargs): + events.append("gather") + return [], [] + + def prefill(event_pack, prefill_reqs): + backend.model.forward(SimpleNamespace(is_prefill=True, to_cuda=lambda: None)) + event_pack.notify_pre_post_handle() + + def decode(event_pack, decode_reqs): + backend.model.forward(SimpleNamespace(is_prefill=False, to_cuda=lambda: None)) + event_pack.notify_pre_post_handle() + + backend._try_read_new_reqs = read_requests + backend._dp_all_gather_prefill_and_decode_req_num = gather + backend.prefill = prefill + backend.decode = decode + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: None) + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + monkeypatch.setattr(g_infer_context, "get_overlap_stream", lambda: SimpleNamespace(wait_stream=lambda stream: None)) + ep_runtime.install_ep_runtime(backend) + events.clear() + with pytest.raises(StopLoop): + backend.infer_loop() + gathered = ["gather"] if backend_class is DPChunkedPrefillBackend else [] + if run_mode == "normal": + assert events == [ + "requests", + *gathered, + "poll", + "prepare", + "prefill", + "record", + "step", + "requests", + *gathered, + "decode", + "requests", + *gathered, + "requests", + ] + return + assert events == [ + "poll", + "requests", + *gathered, + "prepare", + "prefill", + "record", + "step", + "poll", + "requests", + *gathered, + "decode", + "poll", + "requests", + *gathered, + "poll", + "requests", + ] + + +def _run_way(action): + return SimpleNamespace( + is_prefill=lambda: action == "prefill", is_decode=lambda: action == "decode", is_pass=lambda: action == "pass" + ) + + +@pytest.mark.parametrize( + "microbatch,tpsp,installed", [(False, False, True), (True, False, False), (False, True, False)] +) +def test_layer_migration_hooks_only_use_ordered_prefill_paths(runtime, microbatch, tpsp, installed): + backend, events, _, _, manager = runtime + backend.args.enable_prefill_microbatch_overlap = microbatch + backend.args.enable_tpsp_mix_mode = tpsp + ep_runtime.install_ep_runtime(backend) + hook = manager._eplb_states[0].migration_poll + assert (hook is not None) == installed + if installed: + hook() + assert events[-1] == "layer_poll" diff --git a/unit_tests/server/test_api_start_eplb.py b/unit_tests/server/test_api_start_eplb.py new file mode 100644 index 0000000000..5a03f798f7 --- /dev/null +++ b/unit_tests/server/test_api_start_eplb.py @@ -0,0 +1,213 @@ +import pytest + +from lightllm.common.eplb_utils import is_eplb_enabled +from lightllm.server import api_start +from lightllm.server.core.objs.start_args_type import StartArgs +from lightllm.utils import device_utils + + +@pytest.fixture(autouse=True) +def supported_eplb_device(monkeypatch): + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: False) + + +@pytest.mark.parametrize( + "redundant,run_mode,nnodes,error,message", + [ + (0, "normal", 2, ValueError, "requires a single node"), + (-1, "prefill", 1, AssertionError, "nonnegative"), + ], +) +def test_full_eplb_mode_rejected_before_any_process_starts(monkeypatch, redundant, run_mode, nnodes, error, message): + args = StartArgs( + enable_ep_moe=True, + enable_prefill_eplb=True, + eplb_num_redundant_experts_per_rank=redundant, + run_mode=run_mode, + nnodes=nnodes, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + for name in ( + "_set_envs_and_config", + "auto_set_max_req_total_len", + "auto_set_fused_shared_experts", + "set_unique_server_name", + ): + monkeypatch.setattr(api_start, name, lambda args: None) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *a, **kw: pytest.fail("must reject before process startup"), + ) + with pytest.raises(error, match=message): + api_start._launch_subprocesses(args) + + +@pytest.mark.parametrize( + "enabled,run_mode,nnodes,expected", + [ + (False, "prefill", 1, False), + (False, "normal", 2, False), + (False, "decode", 1, False), + (True, "prefill", 1, True), + (True, "prefill", 2, True), + (True, "normal", 1, True), + (True, "decode", 1, False), + (True, "decode", 2, False), + ], +) +def test_eplb_role_policy(enabled, run_mode, nnodes, expected): + args = StartArgs(enable_ep_moe=True, enable_prefill_eplb=enabled, run_mode=run_mode, nnodes=nnodes) + assert is_eplb_enabled(args) is expected + + +def test_eplb_rejects_unsupported_inference_role(): + with pytest.raises(ValueError, match="requires --run_mode prefill or single-node normal"): + is_eplb_enabled(StartArgs(enable_prefill_eplb=True, run_mode="visual_only")) + + +@pytest.mark.parametrize("run_mode", ["prefill", "normal"]) +@pytest.mark.parametrize( + "overrides,message", + [ + ({"enable_ep_moe": False}, "requires --enable_ep_moe"), + ({"enable_prefill_cudagraph": True}, "does not support --enable_prefill_cudagraph"), + ({"enable_rl": True}, "does not support --enable_rl"), + ({"eplb_num_redundant_experts_per_rank": -1}, "nonnegative"), + ], +) +def test_eplb_policy_owns_all_configuration_rejections(run_mode, overrides, message): + args = StartArgs(enable_ep_moe=True, enable_prefill_eplb=True, run_mode=run_mode) + for name, value in overrides.items(): + setattr(args, name, value) + with pytest.raises(AssertionError, match=message): + is_eplb_enabled(args) + + +@pytest.mark.parametrize("run_mode", ["prefill", "normal"]) +def test_eplb_policy_rejects_sm100(monkeypatch, run_mode): + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: True) + args = StartArgs(enable_ep_moe=True, enable_prefill_eplb=True, run_mode=run_mode) + with pytest.raises(AssertionError, match="does not support SM100"): + is_eplb_enabled(args) + + +@pytest.mark.parametrize("enabled,run_mode", [(False, "normal"), (False, "visual_only"), (True, "decode")]) +def test_inactive_eplb_skips_support_checks_and_device_queries(monkeypatch, enabled, run_mode): + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: pytest.fail("inactive EPLB must not query CUDA")) + args = StartArgs( + enable_prefill_eplb=enabled, + run_mode=run_mode, + nnodes=2, + enable_ep_moe=False, + enable_prefill_cudagraph=True, + enable_rl=True, + eplb_num_redundant_experts_per_rank=-1, + ) + assert is_eplb_enabled(args) is False + + +def test_eplb_rl_is_rejected_before_starting_subprocesses(monkeypatch): + args = StartArgs( + enable_ep_moe=True, + enable_prefill_eplb=True, + enable_rl=True, + run_mode="normal", + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *args, **kwargs: pytest.fail("subprocess startup must not be reached"), + ) + with pytest.raises(AssertionError, match="does not support --enable_rl"): + api_start._launch_subprocesses(args) + + +def test_eplb_prefill_cudagraph_is_rejected_before_starting_subprocesses(monkeypatch): + args = StartArgs( + enable_ep_moe=True, + enable_prefill_eplb=True, + enable_prefill_cudagraph=True, + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *args, **kwargs: pytest.fail("subprocess startup must not be reached"), + ) + + with pytest.raises(AssertionError, match="--enable_prefill_eplb does not support --enable_prefill_cudagraph"): + api_start._launch_subprocesses(args) + + +@pytest.mark.parametrize("redundant", [0, 2]) +@pytest.mark.parametrize( + "run_mode,nnodes,ep,sm100,expected", + [ + ("prefill", 1, True, False, True), + ("prefill", 2, True, False, True), + ("normal", 1, True, False, True), + ("decode", 1, True, False, False), + ("decode", 2, True, True, False), + ("decode", 1, False, True, False), + ], +) +def test_eplb_role_resolution_reaches_subprocesses_with_mtp( + monkeypatch, redundant, run_mode, nnodes, ep, sm100, expected +): + monkeypatch.setattr(api_start, "is_hybrid_att_model", lambda _: False) + args = StartArgs( + model_dir="test-model", + enable_ep_moe=ep, + enable_prefill_eplb=True, + eplb_num_redundant_experts_per_rank=redundant, + run_mode=run_mode, + nnodes=nnodes, + mtp_mode="vanilla_no_att", + mtp_step=1, + eos_id=0, + data_type="float16", + disable_vision=True, + disable_audio=True, + disable_shm_warning=True, + ) + + monkeypatch.setattr(api_start, "_set_envs_and_config", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_max_req_total_len", lambda args: None) + monkeypatch.setattr(api_start, "auto_set_fused_shared_experts", lambda args: None) + monkeypatch.setattr(api_start, "set_unique_server_name", lambda args: None) + monkeypatch.setattr(device_utils, "is_sm100_gpu", lambda: sm100) + monkeypatch.setattr(api_start, "auto_set_response_parsers", lambda args: None) + monkeypatch.setattr(api_start, "auto_configure_allreduce_flags_from_args", lambda args: None) + monkeypatch.setattr(api_start, "validate_ports", lambda ports: None) + exported = [] + monkeypatch.setattr(api_start, "set_env_start_args", lambda args: exported.append(args.enable_prefill_eplb)) + monkeypatch.setattr(api_start, "get_shm_port_args", lambda create=False: None) + monkeypatch.setattr(api_start, "send_and_receive_node_ip", lambda args: None) + monkeypatch.setattr(api_start.process_manager, "setup_exit_controller", lambda: None) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *args, **kwargs: (object(), None), + ) + monkeypatch.setattr(api_start.process_manager, "register_process_tree", lambda process: None) + + api_start._launch_subprocesses(args) + assert args.enable_prefill_eplb is expected + assert exported == [expected, expected]