From d257574f2f2310ad8465a6daf462181d2d4a4c1f Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Wed, 22 Jul 2026 15:49:26 +0800 Subject: [PATCH 1/4] feat: imbalance statistics --- lightllm/common/basemodel/basemodel.py | 13 +- .../meta_weights/fused_moe/ep_balance.py | 14 + .../fused_moe/impl/deepgemm_impl.py | 24 +- .../fused_moe/grouped_fused_moe_ep.py | 10 + lightllm/distributed/communication_op.py | 9 + lightllm/server/api_cli.py | 5 + lightllm/server/core/objs/start_args_type.py | 1 + lightllm/server/metrics/metrics.py | 12 + .../mode_backend/ep_balance_monitor.py | 346 +++++++++++ .../server/router/model_infer/model_rpc.py | 9 + .../model_infer/test_ep_balance_monitor.py | 546 ++++++++++++++++++ 11 files changed, 984 insertions(+), 5 deletions(-) create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_balance.py create mode 100644 lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py create mode 100644 unit_tests/server/router/model_infer/test_ep_balance_monitor.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 34a1c9c71c..3fad220894 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -73,6 +73,7 @@ class TpPartBaseModel: def __init__(self, kvargs): self.args = get_env_start_args() + self.ep_balance_monitor = None self.run_mode = kvargs["run_mode"] self.weight_dir_ = kvargs["weight_dir"] self.max_total_token_num = kvargs["max_total_token_num"] @@ -325,9 +326,14 @@ def forward(self, model_input: ModelInput): model_input.to_cuda() if model_input.is_prefill: - return self._prefill(model_input=model_input) - else: - return self._decode(model_input) + model_output = self._prefill(model_input=model_input) + self._record_prefill_ep_balance() + return model_output + return self._decode(model_input) + + def _record_prefill_ep_balance(self): + if self.ep_balance_monitor is not None: + self.ep_balance_monitor.record_prefill_round() def _select_mem_indexes(self, model_input: ModelInput): if model_input.is_prefill: @@ -790,6 +796,7 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input dist_group_manager.clear_deepep_buffer() model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event + self._record_prefill_ep_balance() return model_output0, model_output1 @torch.no_grad() 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/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 09b538a33d..30b6aa1322 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 @@ -20,6 +20,10 @@ class FuseMoeDeepGEMM(FuseMoeTriton): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.ep_balance_counters = None + def _select_experts( self, input_tensor: torch.Tensor, @@ -93,6 +97,7 @@ def _fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + ep_balance_counters=self.ep_balance_counters, ) return output @@ -187,8 +192,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/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/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 93c603212d..76b202f4d5 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -108,6 +108,7 @@ def all_gather_into_tensor(self, output_: torch.Tensor, input_: torch.Tensor, as class DistributeGroupManager: def __init__(self): self.groups = [] + self.ep_balance_monitor_group = None self.ep_buffer = None self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None @@ -125,6 +126,14 @@ def create_groups(self, group_size: int): if not args.disable_flashinfer_allreduce: group.init_flashinfer_reduce() self.groups.append(group) + if ( + getattr(args, "enable_ep_moe", False) + and not getattr(args, "disable_ep_balance_monitor", False) + and getattr(args, "run_mode", "normal") != "decode" + and not getattr(args, "enable_prefill_cudagraph", False) + and not is_sm100_gpu() + ): + self.ep_balance_monitor_group = dist.new_group(ranks=list(range(get_global_world_size())), backend="gloo") return def get_default_group(self) -> CustomProcessGroup: diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index c45a76f9df..92cce28a82 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -801,6 +801,11 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: action="store_true", help="""Whether to enable ep moe for deepseekv3 model.""", ) + parser.add_argument( + "--disable_ep_balance_monitor", + action="store_true", + help="""Disable the prefill expert balance monitor enabled by default for EP-MoE.""", + ) parser.add_argument( "--ep_redundancy_expert_config_path", type=str, diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 3f9747831b..9ebae524a9 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -193,6 +193,7 @@ class StartArgs: default="gpu_counter", metadata={"choices": ["cpu_counter", "pin_mem_counter", "gpu_counter"]} ) enable_ep_moe: bool = field(default=False) + disable_ep_balance_monitor: bool = field(default=False) ep_redundancy_expert_config_path: Optional[str] = field(default=None) auto_update_redundancy_expert: bool = field(default=False) enable_fused_shared_experts: bool = field(default=False) 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/ep_balance_monitor.py b/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py new file mode 100644 index 0000000000..107f4ba504 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py @@ -0,0 +1,346 @@ +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.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_balance import PrefillEPBalanceCounters +from lightllm.distributed.communication_op import dist_group_manager +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" + + +def _find_fused_moe_weights(model): + weights_by_id = {} + for layer in model.trans_layers_weight: + for value in getattr(layer, "__dict__", {}).values(): + if isinstance(value, FusedMoeWeight) and value.enable_ep_moe: + weights_by_id[id(value)] = value + return sorted(weights_by_id.values(), key=lambda weight: weight.layer_num_) + + +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_fused_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.gloo_group = dist_group_manager.ep_balance_monitor_group + if self.gloo_group is None: + raise RuntimeError("EP balance monitor requires a pre-created dedicated Gloo process group") + 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 _monitor_loop(self): + """Consume commonly completed rounds in background report-sized windows.""" + try: + while True: + self._round_ready.wait() + self._round_ready.clear() + if self._overflowed: + self._raise_buffer_overflow("before_sync") + common_round_end = self._get_common_round_end() + if common_round_end - self._processed_round_count > EP_BALANCE_ROUND_BUFFER_CAPACITY: + self._raise_buffer_overflow("common_round_lag", common_round_end=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 + local_round_stats = self._copy_local_rounds(round_start, round_end) + if self._overflowed: + self._raise_buffer_overflow("after_copy") + 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) + + 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/model_rpc.py b/lightllm/server/router/model_infer/model_rpc.py index c66ea0ec5b..2eee76a03d 100644 --- a/lightllm/server/router/model_infer/model_rpc.py +++ b/lightllm/server/router/model_infer/model_rpc.py @@ -25,6 +25,10 @@ ) 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.router.model_infer.mode_backend.ep_balance_monitor import ( + EPBalanceMonitor, + should_enable_ep_balance_monitor, +) from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.log_utils import init_logger from lightllm.utils.graceful_utils import graceful_registry @@ -92,6 +96,11 @@ def exposed_init_model(self, kvargs): logger.info("init redundancy_expert_manager") else: self.redundancy_expert_manager = None + + if should_enable_ep_balance_monitor(self.args): + monitor = EPBalanceMonitor(self.backend.model) + if monitor.enabled: + self.backend.model.ep_balance_monitor = monitor return def exposed_get_max_total_token_num(self): 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..2ddcaafddf --- /dev/null +++ b/unit_tests/server/router/model_infer/test_ep_balance_monitor.py @@ -0,0 +1,546 @@ +import threading +from array import array +from types import SimpleNamespace + +import pytest +import torch +from prometheus_client import generate_latest + +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.ep_balance import PrefillEPBalanceCounters +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_reuses_manager_precreated_dedicated_gloo_group(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_fused_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()) + monkeypatch.setattr(monitor_module, "dist_group_manager", SimpleNamespace(ep_balance_monitor_group=sentinel_group)) + monkeypatch.setattr(monitor_module.dist, "new_group", lambda *args, **kwargs: pytest.fail("unexpected new_group")) + monkeypatch.setattr( + monitor_module.threading, + "Thread", + lambda *args, **kwargs: SimpleNamespace(start=lambda: None), + ) + + 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] + + +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_fused_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_group_manager", SimpleNamespace(ep_balance_monitor_group=sentinel_group)) + monkeypatch.setattr(monitor_module.dist, "new_group", lambda *args, **kwargs: pytest.fail("unexpected new_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_group_manager_creates_monitor_gloo_group_only_when_enabled(monkeypatch, disable_monitor): + monitor_group = object() + 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)) or monitor_group, + ) + + manager = communication_op_module.DistributeGroupManager() + manager.create_groups(group_size=2) + + assert len(manager.groups) == 2 + if disable_monitor: + assert calls == [] + assert manager.ep_balance_monitor_group is None + else: + assert calls == [((), {"ranks": [0, 1], "backend": "gloo"})] + assert manager.ep_balance_monitor_group is monitor_group + + +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(monitor_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_fused_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()) From 0a3334dc54673a61d1162fb1825968bda63cfebd Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Wed, 5 Aug 2026 17:53:26 +0800 Subject: [PATCH 2/4] feat: remove auto_update_redundancy_expert function --- docs/CN/source/tutorial/api_server_args.rst | 11 - docs/EN/source/tutorial/api_server_args.rst | 11 - .../meta_weights/fused_moe/ep_redundancy.py | 199 ------------------ .../fused_moe/fused_moe_weight.py | 11 +- .../meta_weights/fused_moe/impl/base_impl.py | 4 - .../fused_moe/impl/deepgemm_impl.py | 12 -- .../redundancy_topk_ids_repair.py | 111 ---------- lightllm/distributed/communication_op.py | 3 +- lightllm/server/api_cli.py | 11 - lightllm/server/core/objs/start_args_type.py | 2 - .../mode_backend/redundancy_expert_manager.py | 158 -------------- .../server/router/model_infer/model_rpc.py | 8 - lightllm/utils/envs_utils.py | 60 ------ .../test_redundancy_expert_config.json | 180 ---------------- .../test_redundancy_topk_ids_repair.py | 151 ------------- 15 files changed, 3 insertions(+), 929 deletions(-) delete mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/ep_redundancy.py delete mode 100644 lightllm/common/basemodel/triton_kernel/redundancy_topk_ids_repair.py delete mode 100644 lightllm/server/router/model_infer/mode_backend/redundancy_expert_manager.py delete mode 100644 test/advanced_config/redundancy_expert/test_redundancy_expert_config.json delete mode 100644 unit_tests/common/basemodel/triton_kernel/test_redundancy_topk_ids_repair.py 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/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/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index fd3de5c9b4..03cb59c169 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 @@ -11,7 +11,6 @@ from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import select_fuse_moe_impl 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.dist_utils import get_global_world_size, get_global_rank from lightllm.utils.log_utils import init_logger @@ -64,9 +63,7 @@ def __init__( 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, ) self.lock = threading.Lock() self._create_weight() @@ -81,13 +78,9 @@ def _init_config(self, network_config: Dict[str, Any]): 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.redundancy_expert_num = 0 + self.redundancy_expert_ids = [] 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_parallel_params(self): if self.enable_ep_moe: 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..fc6433015e 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 @@ -19,9 +19,7 @@ def __init__( 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 @@ -36,9 +34,7 @@ def __init__( # 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() 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 30b6aa1322..bd9555d8a1 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 @@ -16,7 +16,6 @@ ) 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): @@ -58,17 +57,6 @@ def _select_experts( 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( 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/distributed/communication_op.py b/lightllm/distributed/communication_op.py index 76b202f4d5..b721585596 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -30,7 +30,6 @@ 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, @@ -202,7 +201,7 @@ 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 + self.ll_num_experts = n_routed_experts self.ep_buffer = deep_ep.ElasticBuffer( deepep_group, num_max_tokens_per_rank=self.ll_num_tokens, diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 92cce28a82..fec6c94b5d 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -806,17 +806,6 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: action="store_true", help="""Disable the prefill expert balance monitor enabled by default for EP-MoE.""", ) - 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.""", - ) - parser.add_argument( - "--auto_update_redundancy_expert", - action="store_true", - help="""Whether to update the redundant expert for deepseekv3 model by online expert used counter.""", - ) parser.add_argument( "--enable_fused_shared_experts", action="store_true", diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9ebae524a9..c23cf019eb 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -194,8 +194,6 @@ class StartArgs: ) enable_ep_moe: bool = field(default=False) disable_ep_balance_monitor: bool = field(default=False) - ep_redundancy_expert_config_path: Optional[str] = field(default=None) - auto_update_redundancy_expert: bool = field(default=False) enable_fused_shared_experts: bool = field(default=False) mtp_mode: Optional[str] = field( default=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 2eee76a03d..97fadd3db2 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.router.model_infer.mode_backend.ep_balance_monitor import ( EPBalanceMonitor, @@ -90,13 +89,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 - if should_enable_ep_balance_monitor(self.args): monitor = EPBalanceMonitor(self.backend.model) if monitor.enabled: diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 454abdbea2..f954f5850c 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -107,66 +107,6 @@ 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 [] - - 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_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 - - 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_redundancy_expert_update_interval(): - return int(os.getenv("LIGHTLLM_REDUNDANCY_EXPERT_UPDATE_INTERVAL", 30 * 60)) - - -@lru_cache(maxsize=None) -def get_redundancy_expert_update_max_load_count(): - return int(os.getenv("LIGHTLLM_REDUNDANCY_EXPERT_UPDATE_MAX_LOAD_COUNT", 1)) - - @lru_cache(maxsize=None) def get_triton_autotune_level(): return int(os.getenv("LIGHTLLM_TRITON_AUTOTUNE_LEVEL", 0)) 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() From 2b98bad421ad66c66a165a33e8315cf131fb4030 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Wed, 22 Jul 2026 15:49:26 +0800 Subject: [PATCH 3/4] feat: add EPLB --- lightllm/common/basemodel/basemodel.py | 19 +- .../meta_weights/fused_moe/eplb_placement.py | 385 ++ .../meta_weights/fused_moe/eplb_planner.py | 271 ++ .../fused_moe/expert_parallel_state.py | 97 + .../fused_moe/fused_moe_weight.py | 63 +- .../meta_weights/fused_moe/impl/__init__.py | 35 +- .../meta_weights/fused_moe/impl/base_impl.py | 93 +- .../fused_moe/impl/deepgemm_impl.py | 114 +- .../fused_moe/impl/marlin_impl.py | 4 + .../fused_moe/impl/triton_impl.py | 60 +- .../triton_kernel/fused_moe/grouped_topk.py | 255 +- .../triton_kernel/fused_moe/topk_select.py | 9 - lightllm/common/eplb_planner_worker.py | 184 + lightllm/common/eplb_utils.py | 92 + .../kv_cache_mem_manager/mem_manager.py | 35 +- .../common/quantization/quantize_method.py | 13 +- lightllm/distributed/communication_op.py | 38 +- lightllm/server/api_cli.py | 12 + lightllm/server/api_start.py | 9 + lightllm/server/core/objs/start_args_type.py | 2 + .../model_infer/mode_backend/base_backend.py | 14 +- .../mode_backend/ep_balance_monitor.py | 60 +- .../model_infer/mode_backend/ep_runtime.py | 63 + .../model_infer/mode_backend/eplb_manager.py | 685 ++++ .../model_infer/mode_backend/eplb_transfer.py | 778 ++++ .../server/router/model_infer/model_rpc.py | 8 - lightllm/utils/envs_utils.py | 58 + .../fused_moe/test_activation_config.py | 4 - unit_tests/common/fused_moe/test_eplb.py | 3470 +++++++++++++++++ .../common/fused_moe/test_eplb_full_layout.py | 327 ++ unit_tests/common/fused_moe/test_eplb_hash.py | 35 + .../fused_moe/test_eplb_memory_profile.py | 307 ++ .../fused_moe/test_eplb_planner_worker.py | 88 + .../fused_moe/test_eplb_sampling_policy.py | 229 ++ .../fused_moe/test_eplb_transfer_gpu.py | 258 ++ .../fused_moe/test_moe_weight_creation.py | 152 + .../test_ue8m0_fp8_quantization.py | 29 +- .../model_infer/test_ep_balance_monitor.py | 72 +- .../router/model_infer/test_ep_runtime.py | 234 ++ unit_tests/server/test_api_start_eplb.py | 102 + 40 files changed, 8524 insertions(+), 239 deletions(-) create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_placement.py create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_planner.py create mode 100644 lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/expert_parallel_state.py create mode 100644 lightllm/common/eplb_planner_worker.py create mode 100644 lightllm/common/eplb_utils.py create mode 100644 lightllm/server/router/model_infer/mode_backend/ep_runtime.py create mode 100644 lightllm/server/router/model_infer/mode_backend/eplb_manager.py create mode 100644 lightllm/server/router/model_infer/mode_backend/eplb_transfer.py create mode 100644 unit_tests/common/fused_moe/test_eplb.py create mode 100644 unit_tests/common/fused_moe/test_eplb_full_layout.py create mode 100644 unit_tests/common/fused_moe/test_eplb_hash.py create mode 100644 unit_tests/common/fused_moe/test_eplb_memory_profile.py create mode 100644 unit_tests/common/fused_moe/test_eplb_planner_worker.py create mode 100644 unit_tests/common/fused_moe/test_eplb_sampling_policy.py create mode 100644 unit_tests/common/fused_moe/test_eplb_transfer_gpu.py create mode 100644 unit_tests/common/fused_moe/test_moe_weight_creation.py create mode 100644 unit_tests/server/router/model_infer/test_ep_runtime.py create mode 100644 unit_tests/server/test_api_start_eplb.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 3fad220894..930d874ebd 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -15,6 +15,8 @@ from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class +from lightllm.common.eplb_utils import get_eplb_memory_reservations +from lightllm.common.kv_cache_mem_manager.mem_manager import memory_reservation_scope from lightllm.common.req_manager import ReqManager from lightllm.common.build_utils import repair_config from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( @@ -73,7 +75,6 @@ class TpPartBaseModel: def __init__(self, kvargs): self.args = get_env_start_args() - self.ep_balance_monitor = None self.run_mode = kvargs["run_mode"] self.weight_dir_ = kvargs["weight_dir"] self.max_total_token_num = kvargs["max_total_token_num"] @@ -115,7 +116,8 @@ def __init__(self, kvargs): self._init_weights() with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): self._init_req_manager() - self._init_mem_manager() + with memory_reservation_scope(get_eplb_memory_reservations(self)): + self._init_mem_manager() # Qwen3.5 等 linear attention 模型会在 req_manager 中保存大量运行时 state。先初始化 # req_manager、再初始化 mem_manager，可以让 KV cache 显存评估包含这些 state 的实际占用； @@ -326,14 +328,9 @@ def forward(self, model_input: ModelInput): model_input.to_cuda() if model_input.is_prefill: - model_output = self._prefill(model_input=model_input) - self._record_prefill_ep_balance() - return model_output - return self._decode(model_input) - - def _record_prefill_ep_balance(self): - if self.ep_balance_monitor is not None: - self.ep_balance_monitor.record_prefill_round() + return self._prefill(model_input=model_input) + else: + return self._decode(model_input) def _select_mem_indexes(self, model_input: ModelInput): if model_input.is_prefill: @@ -637,7 +634,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" @@ -796,7 +792,6 @@ def _microbatch_overlap_prefill_cuda(self, model_input0: ModelInput, model_input dist_group_manager.clear_deepep_buffer() model_output0.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event model_output1.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event - self._record_prefill_ep_balance() return model_output0, model_output1 @torch.no_grad() 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..a980abe7e8 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_placement.py @@ -0,0 +1,385 @@ +from typing import Dict, Tuple +import torch + + +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 非空时必填。 +) -> 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, + ) + 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 非空时必填。 +) -> 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 后，先筛选源节点内副本，再按 source_rank 轮转候选前缀。 + """ + return _build_layer_maps( + physical_expert_ids_by_layer, + num_logical_experts, + source_rank=source_rank, + node_world_size=node_world_size, + ) + + +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, +) -> 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) + candidate_rank_load = _estimate_rank_load(expert_load, candidate_placement, expert_alignment, node_world_size) + 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 非空时必填。 +) -> 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: + 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 + 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=source_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, +) -> 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, + ).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, +) -> torch.Tensor: + """Return aligned ``[samples, layers, expert, rank]`` contributions.""" + route = _source_route(locations, num_nodes, node_world_size) + physical_load = torch.einsum("slne,lner->sler", source_load.to(torch.float64), route) + 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..71a3ae862d --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/eplb_planner.py @@ -0,0 +1,271 @@ +"""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): + return _estimate_rank_load(load, placement, alignment, node_world_size).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, +) -> 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() + score = _critical_load(load, best, expert_alignment, node_world_size) + 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), + ): + proposal_score = _critical_load(load, proposal, expert_alignment, node_world_size) + 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): + result = placement.clone() + experts = load.shape[-1] + for layer in range(placement.shape[0]): + source = load[:, layer : layer + 1] + margin = stickiness * float(source.sum()) / experts + initial_score = _critical_load(source, placement[layer : layer + 1], alignment, node_world_size) + 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)[:, 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) + next_contributions = _expert_rank_load_all( + move_load, new_locations, source.shape[2], node_world_size, alignment + ).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) + 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): + """Return a complete physical placement for the manager's CPU planner.""" + return plan_full_experts( + expert_load, + current, + expert_alignment=expert_alignment, + node_world_size=node_world_size, + stickiness=stickiness, + ) + + +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, +) -> 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).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, + ) + + +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..b353348574 --- /dev/null +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/expert_parallel_state.py @@ -0,0 +1,97 @@ +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass +from typing import Iterator, Optional + +import torch + + +_eplb_model_init_disabled: ContextVar[bool] = ContextVar("eplb_model_init_disabled", default=False) + + +def is_eplb_model_init_disabled() -> bool: + return _eplb_model_init_disabled.get() + + +@contextmanager +def disable_eplb_model_init() -> Iterator[None]: + token = _eplb_model_init_disabled.set(True) + try: + yield + finally: + _eplb_model_init_disabled.reset(token) + + +@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 + + 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 create_expert_parallel_state(num_logical_experts: int, world_size: int, rank: int) -> ExpertParallelState: + from lightllm.common.eplb_utils import EPLB_STEADY_SAMPLE_STEPS + 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, + ) + + args = get_env_start_args() + eplb = None + if args.enable_prefill_eplb and not is_eplb_model_init_disabled(): + if args.run_mode != "prefill": + raise ValueError("EPLB requires --run_mode prefill") + redundant = args.eplb_num_redundant_experts_per_rank + 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=get_node_world_size() + ) + # Cover both initial/steady windows and two microbatch dispatches per step. + sample_steps = 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(), + route_counter=torch.zeros((2 * sample_steps, num_logical_experts), dtype=torch.int64, device="cuda"), + ) + 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 03cb59c169..67653a7359 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,9 +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_env_start_args from lightllm.utils.dist_utils import get_global_world_size, get_global_rank from lightllm.utils.log_utils import init_logger @@ -55,15 +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() + 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, - routed_expert_counter_tensor=self.routed_expert_counter_tensor, + expert_parallel_state=self.expert_parallel_state, ) self.lock = threading.Lock() self._create_weight() @@ -77,12 +80,14 @@ 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 = 0 - self.redundancy_expert_ids = [] - self.routed_expert_counter_tensor = torch.zeros((self.n_routed_experts,), dtype=torch.int64, device="cuda") + def _init_expert_parallel_state(self): + self.expert_parallel_state = ( + create_expert_parallel_state(self.n_routed_experts, self.global_world_size, self.global_rank_) + 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 @@ -96,27 +101,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, @@ -280,8 +283,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 fc6433015e..750aa0fa42 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,49 +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, - routed_expert_counter_tensor: torch.Tensor, ): 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.routed_expert_counter_tensor = routed_expert_counter_tensor - - # workspace for kernel optimization - self.workspace = self.create_workspace() - - @abstractmethod - def create_workspace(self): - pass - @abstractmethod def __call__( self, input_tensor: torch.Tensor, @@ -66,5 +42,68 @@ def __call__( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + ) -> 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, + ) + + @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, ) -> 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 bd9555d8a1..aa090c0067 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 ( @@ -18,10 +19,13 @@ from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType -class FuseMoeDeepGEMM(FuseMoeTriton): - def __init__(self, *args, **kwargs): +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 + self._primary_weight_pack_cache = {} def _select_experts( self, @@ -36,27 +40,55 @@ 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.""" + """选择 expert；EPLB prefill 统一由融合路径返回 physical ID。""" 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 + eplb_active = eplb is not None + if is_prefill is True and eplb_active: + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import 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, - ) + group_score_topk_num = 2 if topk_group == 4 and num_expert_group == 8 and top_k == 8 else 1 + 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(), + record_load=eplb.recording, + use_grouped_topk=use_grouped_topk, + return_logical_ids=preserve_logical_ids, + group_score_used_topk_num=group_score_topk_num, + ) + 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 return topk_weights, topk_ids, origin_topk_ids def _fused_experts( @@ -72,13 +104,20 @@ def _fused_experts( limit: Optional[float] = None, clamp_up_add_one: bool = True, ): + if is_prefill is False: + w13 = self._primary_weight_pack(w13) + w2 = self._primary_weight_pack(w2) + num_experts = self.n_routed_experts + else: + num_experts = self.expert_parallel_state.num_total_physical_experts + output = fused_experts( hidden_states=input_tensor, w13=w13, 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=num_experts, quant_method=self.quant_method, is_prefill=is_prefill, previous_event=None, # for overlap @@ -111,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) @@ -120,7 +160,8 @@ 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, + # decode 与 EPLB 的物理冗余行刻意隔离：DeepEP 使用原始 logical expert ID。 + num_experts=self.n_routed_experts, use_fp8=use_fp8_w8a8, async_finish=False, return_recv_hook=True, @@ -150,6 +191,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 @@ -167,7 +209,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(), @@ -209,6 +251,7 @@ def masked_group_gemm( dtype: torch.dtype, expected_m: int, ): + w13, w2 = self._primary_weight_pack(w13), self._primary_weight_pack(w2) w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale return masked_group_gemm( @@ -306,3 +349,30 @@ def hook(): event.current_stream_wait() return combined_x, hook + + def _primary_weight_pack(self, weight_pack: WeightPack) -> WeightPack: + """返回所有 decode 路径使用的缓存本地主副本视图。""" + if self.eplb is None: + return weight_pack + cache = getattr(self, "_primary_weight_pack_cache", None) + if cache is None: + cache = self._primary_weight_pack_cache = {} + cache_key = id(weight_pack) + primary = cache.get(cache_key) + if primary is None: + num_primary_experts_per_rank = self.expert_parallel_state.num_primary_experts_per_rank + primary = WeightPack( + weight=weight_pack.weight[:num_primary_experts_per_rank], + weight_scale=( + weight_pack.weight_scale[:num_primary_experts_per_rank] + if weight_pack.weight_scale is not None + else None + ), + weight_zero_point=( + getattr(weight_pack, "weight_zero_point", None)[:num_primary_experts_per_rank] + if getattr(weight_pack, "weight_zero_point", None) is not None + else None + ), + ) + cache[cache_key] = primary + return primary 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..d09336a102 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 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..abaec5fc22 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 @@ -89,56 +88,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_topk.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py index 8c882b6c1d..196e1d432f 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_topk.py @@ -5,6 +5,16 @@ from triton.language.standard import _log2, sum, zeros_like +@triton.jit +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 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 @@ -202,6 +212,172 @@ def grouped_topk_kernel( return +@triton.jit +def grouped_topk_eplb_kernel( + gating_output_ptr, + gating_output_stride_m, + gating_output_stride_n, + correction_bias_ptr, + 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_topk_num, + IS_SIGMOID: tl.constexpr, + USE_GROUPED_TOPK: tl.constexpr, + HAS_CORRECTION_BIAS: tl.constexpr, + 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, + MAP_SLOTS: tl.constexpr, + RECORD_LOAD: tl.constexpr, + SINGLE_TOKEN: tl.constexpr, +): + """Grouped top-k, EPLB accounting, and replica mapping without a global score workspace.""" + token_index = tl.program_id(axis=0) + 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 + 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: + 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: + 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 = tl.where(valid_expert, old_scores, -float("inf")) + + 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) + 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: + candidate_scores = tl.where(valid_expert, old_scores, -float("inf")) + + 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) + + topk_mask = topk_offsets < TOPK_NUM + if RECORD_LOAD: + tl.atomic_add( + expert_counter_ptr + sample_index * COUNTER_NUM_EXPERTS + selected_logical_ids, + 1, + mask=topk_mask, + sem="relaxed", + ) + if SINGLE_TOKEN: + 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, + ) + 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_logical_ids + token_index * out_logical_ids_stride_m + topk_offsets * out_logical_ids_stride_n, + selected_logical_ids, + mask=topk_mask, + ) + + def triton_grouped_topk( hidden_states: torch.Tensor, gating_output: torch.Tensor, @@ -213,7 +389,6 @@ def triton_grouped_topk( scoring_func: str = "softmax", group_score_used_topk_num=2, ): - if correction_bias is not None: has_correction_bias = True else: @@ -263,3 +438,81 @@ def triton_grouped_topk( num_stages=1, ) return out_topk_weights, out_topk_ids + + +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, +): + """Fused EPLB prefill top-k returning physical IDs and optional logical IDs.""" + 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: + 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_eplb_kernel[(token_num,)]( + gating_output, + *gating_output.stride(), + correction_bias, + out_topk_weights, + *out_topk_weights.stride(), + out_topk_ids, + *out_topk_ids.stride(), + 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, + 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], + MAP_SLOTS=logical_to_physical_map.shape[1], + RECORD_LOAD=record_load, + SINGLE_TOKEN=token_num == 1, + num_warps=num_warps, + num_stages=1, + ) + 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..d2f59de480 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py @@ -21,7 +21,6 @@ from lightllm.utils.sgl_utils import sgl_ops from typing import Callable, List, Optional, Tuple from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk -from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType def fused_topk( @@ -168,12 +167,4 @@ 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 - if Autotuner.is_kernel_autotune_warmup(AutotuneKernelType.GENERAL): - rand_gen = torch.Generator(device="cuda") - rand_gen.manual_seed(router_logits.shape[0]) - router_logits = torch.randn(size=router_logits.shape, generator=rand_gen, dtype=torch.float32, device="cuda") - _, topk_ids = torch.topk(router_logits, k=top_k, dim=1) - return topk_weights, topk_ids diff --git a/lightllm/common/eplb_planner_worker.py b/lightllm/common/eplb_planner_worker.py new file mode 100644 index 0000000000..3e64901395 --- /dev/null +++ b/lightllm/common/eplb_planner_worker.py @@ -0,0 +1,184 @@ +"""CPU-only persistent EPLB candidate planner actor.""" +import atexit +import multiprocessing as mp +import os +import time +import threading + + +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 = mp.get_context("spawn").Process(target=_serve, args=(child, os.getpid()), daemon=True) + try: + self._process.start() + child.close() + 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() + 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 _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) + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import eplb_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() diff --git a/lightllm/common/eplb_utils.py b/lightllm/common/eplb_utils.py new file mode 100644 index 0000000000..bdf5a9d9e7 --- /dev/null +++ b/lightllm/common/eplb_utils.py @@ -0,0 +1,92 @@ +"""Small, dependency-free EPLB helpers shared by transfer and model profiling.""" + +import math + + +EPLB_FULL_STAGING_PARTS = 2 +EPLB_STEADY_SAMPLE_STEPS = 4 + + +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_memory_reservations(model): + """Budget buffers created after KV profiling; persistent counters already exist.""" + if model.is_mtp_draft_model or not model.args.enable_prefill_eplb: + return {} + weights = [weight for weight in find_ep_moe_weights(model) if weight.expert_parallel_state.eplb is not None] + return { + name: value + for name, value in ( + ("eplb_staging", get_eplb_staging_nbytes(weights)), + ("eplb_sampling", get_eplb_sampling_peak_nbytes(weights)), + ) + if value + } + + +def get_eplb_staging_shape(weights): + """Stage two row slices of one complete layer, independent of layer count.""" + 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/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 38c870ce4e..9905438930 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -1,6 +1,8 @@ import math import os import re +from contextlib import contextmanager +from contextvars import ContextVar import torch import torch.distributed as dist import torch.multiprocessing as mp @@ -25,19 +27,37 @@ from .operator import BaseMemManagerOperator, NormalMemOperator logger = init_logger(__name__) +_memory_reservations = ContextVar("kv_memory_reservations", default=None) -class MemoryManager: +@contextmanager +def memory_reservation_scope(reservations): + """Supply post-profile budgets only to the synchronous KV construction stack.""" + token = _memory_reservations.set(dict(reservations)) + try: + yield + finally: + _memory_reservations.reset(token) + +class MemoryManager: operator_class = NormalMemOperator - def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): + def __init__( + self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9, memory_reservations=None + ): self.size = size self.head_num = head_num self.head_dim = head_dim self.layer_num = layer_num self.always_copy = always_copy self.dtype = dtype + # These buffers are allocated after KV profiling and deducted from its budget once. + if memory_reservations is None: + memory_reservations = _memory_reservations.get() + self.memory_reservations = dict(memory_reservations or {}) + if any(value < 0 for value in self.memory_reservations.values()): + raise ValueError(f"memory reservations must be non-negative: {self.memory_reservations}") # profile the max total token num if the size is None self.profile_size(mem_fraction) @@ -97,8 +117,16 @@ def profile_size(self, mem_fraction): world_size = dist.get_world_size() available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) cell_size = self.get_cell_size() + reservations = getattr(self, "memory_reservations", {}) + reserved_memory_size = sum(reservations.values()) pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size + available_memory_bytes = available_memory * 1024 ** 3 - reserved_memory_size - pd_kv_move_buffer_size + if available_memory_bytes <= 0: + raise RuntimeError( + f"{type(self).__name__} reservations require {reserved_memory_size / 1024**3:.2f} GB, " + f"plus {pd_kv_move_buffer_size / 1024**3:.2f} GB for the PD KV transfer buffer, " + f"but only {available_memory:.2f} GB is available" + ) self.size = int(available_memory_bytes / cell_size) if world_size > 1: tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") @@ -106,6 +134,7 @@ def profile_size(self, mem_fraction): self.size = tensor.item() logger.info( f"{str(available_memory)} GB space is available after load the model weight\n" + f"{reservations} bytes are reserved for post-profile model buffers\n" f"{str(pd_kv_move_buffer_size / 1024 ** 2)} MB is reserved for PD KV transfer buffer\n" f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n" f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n" diff --git a/lightllm/common/quantization/quantize_method.py b/lightllm/common/quantization/quantize_method.py index 95d8d806f9..e535bae0b4 100644 --- a/lightllm/common/quantization/quantize_method.py +++ b/lightllm/common/quantization/quantize_method.py @@ -76,13 +76,24 @@ def create_weight( def create_moe_weight( self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int ) -> Tuple[WeightPack, List[WeightPack]]: - return self._create_weight( + weight_pack, weight_pack_list = self._create_weight( out_dims=out_dims, in_dim=in_dim, dtype=dtype, device_id=device_id, num_experts=num_experts, ) + if num_experts == 1: + # Dense factories omit the singleton expert axis. Restore it on + # the fused pack and its loading views without copying storage. + for pack in [weight_pack, *weight_pack_list]: + if pack.weight.ndim == 2: + pack.weight = pack.weight.unsqueeze(0) + if pack.weight_scale is not None: + pack.weight_scale = pack.weight_scale.unsqueeze(0) + if pack.weight_zero_point is not None: + pack.weight_zero_point = pack.weight_zero_point.unsqueeze(0) + return weight_pack, weight_pack_list def load_weight(self, weight: torch.Tensor, weight_pack: WeightPack) -> None: if self._check_weight_need_quanted(weight): diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index b721585596..d18c4a780f 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -201,7 +201,15 @@ 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 + total_redundant_experts = ( + get_env_start_args().eplb_num_redundant_experts_per_rank * global_world_size + if get_env_start_args().enable_prefill_eplb + else 0 + ) + self.ll_prefill_num_experts = n_routed_experts + total_redundant_experts + # EPLB's redundant rows are a prefill-only physical layout; decode + # always routes the logical expert space. + self.ll_decode_num_experts = n_routed_experts self.ep_buffer = deep_ep.ElasticBuffer( deepep_group, num_max_tokens_per_rank=self.ll_num_tokens, @@ -241,7 +249,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_decode_num_experts, ) microbatch_count = len(self.groups) min_prefill_reuse_buffer_bytes = _calculate_min_chunked_expanded_moe_reuse_buffer_bytes( @@ -262,7 +273,7 @@ def new_deepep_group( deepep_group, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, - num_qps_per_rank=(self.ll_num_experts // global_world_size), + num_qps_per_rank=(self.ll_decode_num_experts // global_world_size), ) if enable_mega_moe_buffer: @@ -275,22 +286,26 @@ def new_deepep_group( self.ep_mega_moe_buffer = deep_gemm.get_symm_buffer_for_mega_moe( deepep_group, - self.ll_num_experts, + self.ll_decode_num_experts, self.ll_num_tokens, num_experts_per_tok, self.ll_hidden, moe_intermediate_size, ) logger.info( - "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, expert_quant_method_names=%s", + "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, " + "ll_prefill_num_experts=%s, ll_decode_num_experts=%s, expert_quant_method_names=%s", enable_low_latency_buffer, enable_mega_moe_buffer, + self.ll_prefill_num_experts, + self.ll_decode_num_experts, sorted(expert_quant_method_names), ) - theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) - self._set_num_sms_for_deep_gemm(theoretical_sms) + theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_prefill_num_experts, num_experts_per_tok) + low_latency_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_decode_num_experts, num_experts_per_tok) + self._set_num_sms_for_deep_gemm(theoretical_sms, low_latency_sms) - def _set_num_sms_for_deep_gemm(self, deepep_sms: int): + def _set_num_sms_for_deep_gemm(self, deepep_sms: int, low_latency_sms: int): try: try: from deep_gemm.jit_kernels.utils import set_num_sms @@ -299,9 +314,12 @@ def _set_num_sms_for_deep_gemm(self, deepep_sms: int): device_sms = get_device_sm_count() deepep_sms = max(0, min(deepep_sms, max(device_sms - 2, 0))) + low_latency_sms = max(0, min(low_latency_sms, max(device_sms - 2, 0))) self.ep_num_sms = deepep_sms if self.ep_low_latency_buffer is not None: - deep_ep.Buffer.set_num_sms(deepep_sms - deepep_sms % 2) + # This setting controls the legacy low-latency buffer; keep + # its SM reservation based on decode's logical expert count. + deep_ep.Buffer.set_num_sms(low_latency_sms - low_latency_sms % 2) set_num_sms(max(device_sms - deepep_sms, 2)) except BaseException as e: logger.warning(f"set num sms for deep_gemm failed: {e}") @@ -343,7 +361,7 @@ def clear_deepep_buffer(self): """ if self.ep_low_latency_buffer is not None: self.ep_low_latency_buffer.clean_low_latency_buffer( - self.ll_decode_num_tokens, self.ll_hidden, self.ll_num_experts + self.ll_decode_num_tokens, self.ll_hidden, self.ll_decode_num_experts ) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index fec6c94b5d..cee9b4f5d2 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -806,6 +806,18 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: action="store_true", help="""Disable the prefill expert balance monitor enabled by default for EP-MoE.""", ) + parser.add_argument( + "--enable_prefill_eplb", + action="store_true", + help="""Enable online expert load balancing with complete physical layouts; requires --run_mode prefill.""", + ) + 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 prefill EPLB. + The value must be nonnegative; zero enables rearrangement without extra replicas.""", + ) parser.add_argument( "--enable_fused_shared_experts", action="store_true", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index c67ab9582d..7bb29c400f 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -33,6 +33,7 @@ get_default_disable_chunked_mm_input, ) from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args +from lightllm.utils.device_utils import is_sm100_gpu logger = init_logger(__name__) @@ -173,6 +174,14 @@ 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" + if args.enable_prefill_eplb: + 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" + # EPLB updates expert weights in place, but SM100 Mega-MoE caches transformed weights by tensor data_ptr. + assert not is_sm100_gpu(), "--enable_prefill_eplb does not support SM100" + assert args.run_mode == "prefill", "--enable_prefill_eplb requires --run_mode prefill" + assert args.eplb_num_redundant_experts_per_rank >= 0, "redundant expert count must be nonnegative" + 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 c23cf019eb..ecf0d9818c 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -194,6 +194,8 @@ class StartArgs: ) enable_ep_moe: 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/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 6c621b6adb..9bf7b143d0 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,10 @@ 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 +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( + disable_eplb_model_init, +) class ModeBackend: @@ -261,13 +265,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 +354,9 @@ 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)) + with disable_eplb_model_init(): + 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 index 107f4ba504..56051f8e44 100644 --- a/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py +++ b/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py @@ -6,7 +6,7 @@ import torch.distributed as dist from lightllm.common.basemodel.basemodel import TpPartBaseModel -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +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.distributed.communication_op import dist_group_manager from lightllm.server.metrics.manager import MetricClient @@ -148,22 +148,13 @@ def classify_prefill_placement_pressure_drift(drift: float) -> str: return "dynamic" -def _find_fused_moe_weights(model): - weights_by_id = {} - for layer in model.trans_layers_weight: - for value in getattr(layer, "__dict__", {}).values(): - if isinstance(value, FusedMoeWeight) and value.enable_ep_moe: - weights_by_id[id(value)] = value - return sorted(weights_by_id.values(), key=lambda weight: weight.layer_num_) - - 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_fused_moe_weights(model) + self.weights = find_ep_moe_weights(model) self.enabled = bool(self.weights) if not self.enabled: @@ -198,6 +189,7 @@ def __init__(self, model: TpPartBaseModel): 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_group_manager.ep_balance_monitor_group if self.gloo_group is None: @@ -309,28 +301,52 @@ def _log_stats(self, round_stats: torch.Tensor): ) 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() - if self._overflowed: - self._raise_buffer_overflow("before_sync") + 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() - if common_round_end - self._processed_round_count > EP_BALANCE_ROUND_BUFFER_CAPACITY: - self._raise_buffer_overflow("common_round_lag", common_round_end=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 - local_round_stats = self._copy_local_rounds(round_start, round_end) - if self._overflowed: - self._raise_buffer_overflow("after_copy") - 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) + 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() 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..061231b2f4 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/ep_runtime.py @@ -0,0 +1,63 @@ +"""Bind optional EP services at initialization, outside the inference loop.""" + +from functools import wraps + +import torch.distributed as dist + +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 backend.args.enable_prefill_eplb: + from .eplb_manager import EPLBManager + + manager = EPLBManager(backend.model) + + 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): + 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: + 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 + dist.barrier() 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..5cd0a12a2b --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb_manager.py @@ -0,0 +1,685 @@ +import os +import threading +import time +from pathlib import Path +from typing import Dict, Optional + +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 + 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 + # A fresh manager starts with one continuous base window. After a + # sufficient evaluation, steady state returns to the cheap sparse + # probe. An insufficient sparse probe schedules one fresh continuous + # base window before the next fixed sampling boundary. + self._continuous_collection_start_step: Optional[int] = None + self._continuous_collection_end_step: Optional[int] = self.step_interval + self._steady_collection_end_step: Optional[int] = 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 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): + if self.in_flight or self.evaluation_in_flight: + return + self.prefill_steps += 1 + continuous_start = self._continuous_collection_start_step + continuous_end = self._continuous_collection_end_step + if continuous_end is not None: + # 启动或稀疏样本不足时连续采样，保证负载统计可靠。 + if continuous_start is not None and self.prefill_steps == continuous_start: + self._set_recording(True) + if self.prefill_steps >= continuous_end: + self._start_evaluation() + return + sampling_interval = self.sampling_interval + phase = self.prefill_steps % sampling_interval + steady_collection_end_step = self._steady_collection_end_step + if steady_collection_end_step is not None: + if self.prefill_steps >= steady_collection_end_step: + self._clear_steady_collection() + self._start_evaluation() + return + if sampling_interval == 1: + self._start_evaluation() + return + if phase == sampling_interval - self._steady_sample_window_steps(): + # 稳态仅在周期末采样少量 step，降低路由计数和评估开销。 + self._start_steady_sampling_window(self.prefill_steps + self._steady_sample_window_steps()) + + 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 _clear_continuous_collection(self): + self._continuous_collection_start_step = None + self._continuous_collection_end_step = None + + def _clear_steady_collection(self): + self._steady_collection_end_step = None + + 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 _sampling_interval_cap(self) -> int: + return max(self.step_interval * 16, getattr(self, "min_rebalance_interval", 0)) + + def _start_steady_sampling_window(self, collection_end_step: int): + """Start the fixed sparse window without moving its evaluation boundary.""" + self._reset_recorded_samples() + self._steady_collection_end_step = collection_end_step + self._set_recording(True) + + def _begin_continuous_collection(self): + minimum_end = self.prefill_steps + self.step_interval + collection_end = -(-minimum_end // self.sampling_interval) * self.sampling_interval + self._reset_recorded_samples() + self._clear_steady_collection() + self._continuous_collection_start_step = collection_end - self.step_interval + self._continuous_collection_end_step = collection_end + self._set_recording(self._continuous_collection_start_step == self.prefill_steps) + + def _prepare_next_sampling_window(self): + """Clear the current window and arm the next sparse sampling window.""" + self._clear_continuous_collection() + self._clear_steady_collection() + if self.sampling_interval == 1: + self._reset_recorded_samples() + self._set_recording(True) + elif self.sampling_interval <= self._steady_sample_window_steps(): + # There is no later pre-boundary manager step at which to arm a + # full clamped window, so arm immediately but keep the same next + # fixed boundary. + self._start_steady_sampling_window(self.prefill_steps + self.sampling_interval) + else: + self._reset_recorded_samples() + self._set_recording(False) + + @staticmethod + def _recent_ring_samples(counter: torch.Tensor, recorded_sample_count: int) -> torch.Tensor: + """Return the newest ring rows in chronological order.""" + capacity = counter.shape[0] + available = min(recorded_sample_count, capacity) + if available == 0: + return counter[:0] + start = (recorded_sample_count - available) % capacity + indices = (torch.arange(available, dtype=torch.int64, device=counter.device) + start) % capacity + return counter.index_select(0, indices) + + def _collect_local_samples(self) -> torch.Tensor: + counters = [state.route_counter for state in self._eplb_states] + capacities = [counter.shape[0] for counter in counters] + if len(set(capacities)) != 1 or any(counter.ndim != 2 for counter in counters): + raise RuntimeError("EPLB sample capacities differ between layers") + counts = [state.recorded_sample_count for state in self._eplb_states] + if len(set(counts)) != 1: + raise RuntimeError("EPLB recorded sample counts differ between layers") + sample_count = counts[0] + # Validate the metadata before copying the newest rows to the CPU. + metadata = torch.tensor([sample_count, -sample_count, capacities[0], -capacities[0]], dtype=torch.int64) + dist.all_reduce(metadata, op=dist.ReduceOp.MIN, group=self.evaluation_group) + if metadata[0] != -metadata[1] or metadata[2] != -metadata[3]: + raise RuntimeError("EPLB recorded sample count or capacity differs between ranks") + # Stack the fixed-size ring buffers in one GPU launch. Slicing each + # layer before stacking turns a single launch into one index_select per + # MoE layer and is measurably slower in the normal sparse path. + counter_samples = torch.stack(counters, dim=1) + return self._recent_ring_samples(counter_samples, sample_count).cpu() + + 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_in_flight(self): + 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: + 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.""" + result = None + local_error = None + 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, + } + worker = getattr(self, "_planner_worker", None) + if worker is None: + candidate = plan_eplb_candidate(global_load, self.current_placement, **settings) + candidate_timing = {"compute_ms": None, "roundtrip_ms": None} + else: + candidate, candidate_timing = worker.plan( + global_load.cpu(), self.current_placement.cpu(), **settings + ) + select_align_started = time.perf_counter() + placement, improved, metrics, before_load, after_load = select_improving_placements( + global_load, + self.current_placement, + candidate, + expert_alignment=EPLB_EXPERT_ALIGNMENT, + node_world_size=self.node_world_size, + rebalance_gain_threshold=self.rebalance_gain_threshold, + ) + 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( + self.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, + "current_placement": self.current_placement.clone(), + "prefill_steps": self.prefill_steps, + "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.step_interval + if self._continuous_collection_end_step is not None + else self._steady_sample_window_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, + ) + 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._set_recording(False) + 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 _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": + from_continuous_window = self._continuous_collection_end_step is not None + if from_continuous_window: + self.sampling_interval = min(self.sampling_interval * 4, self._sampling_interval_cap()) + self._prepare_next_sampling_window() + else: + self._begin_continuous_collection() + if self.global_rank == 0: + if from_continuous_window: + logger.info( + "eplb insufficient samples: prefill_steps=%s minimum_layer_samples=%s required=%s " + "next_sampling_interval=%s recorded_sample_count=%s sample_window_steps=%s", + self.prefill_steps, + result["minimum_layer_samples"], + result["minimum"], + self.sampling_interval, + result.get("recorded_sample_count"), + result.get("sample_window_steps"), + ) + else: + logger.info( + "eplb insufficient samples: prefill_steps=%s minimum_layer_samples=%s required=%s " + "scheduled_fresh_window_start=%s scheduled_fresh_window_end=%s " + "recorded_sample_count=%s sample_window_steps=%s", + self.prefill_steps, + result["minimum_layer_samples"], + result["minimum"], + self._continuous_collection_start_step, + self._continuous_collection_end_step, + result.get("recorded_sample_count"), + result.get("sample_window_steps"), + ) + return False + if result["kind"] == "no_improvement": + self.sampling_interval = min(self.sampling_interval * 4, self._sampling_interval_cap()) + 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._clear_continuous_collection() + self._reset_recorded_samples() + 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..285a04a90a --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py @@ -0,0 +1,778 @@ +"""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 layer's two staging parts.""" + + 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): + """Commit both slices before publishing full-layout routing metadata.""" + error = None + try: + prepared = self._prepared_layer + if prepared.layer_index != layer_index or len(prepared.parts) != EPLB_FULL_STAGING_PARTS: + raise RuntimeError("invalid full EPLB prepared layer") + self._commit_full_part(layer_index, buffer_index, changed_dst_slots, 0) + torch.cuda.current_stream().synchronize() + except BaseException as exc: + error = exc + self._full_phase_gate(error) # All ranks consumed part 0 before staging reuse. + error = None + try: + self._copy_part(prepared.parts[1]) + except BaseException as exc: + error = exc + self._full_phase_gate(error) # All incoming part-1 transfers are complete. + error = None + try: + self._commit_full_part(layer_index, buffer_index, changed_dst_slots, 1) + torch.cuda.current_stream().synchronize() + self._prepared_layer = None + except BaseException as exc: + error = exc + self._full_phase_gate(error) # All live weight slices are complete before publish. + + 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): + # Full layouts stage every changed row of one layer before publishing; + # this makes arbitrary overwrite cycles safe without eight full layers. + 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() + if self._cross_node_ranks: + 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-{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 _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) + 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 prepare_transfer(self, layer_plans: Sequence[Tuple[int, Sequence[TransferStep]]]): + 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/model_rpc.py b/lightllm/server/router/model_infer/model_rpc.py index 97fadd3db2..d8cf2bf022 100644 --- a/lightllm/server/router/model_infer/model_rpc.py +++ b/lightllm/server/router/model_infer/model_rpc.py @@ -24,10 +24,6 @@ PDDPForDecodeNode, ) from lightllm.server.router.model_infer.mode_backend.rl_backend_ops import RlBackendOps -from lightllm.server.router.model_infer.mode_backend.ep_balance_monitor import ( - EPBalanceMonitor, - should_enable_ep_balance_monitor, -) from lightllm.server.core.objs.start_args_type import StartArgs from lightllm.utils.log_utils import init_logger from lightllm.utils.graceful_utils import graceful_registry @@ -89,10 +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 - if should_enable_ep_balance_monitor(self.args): - monitor = EPBalanceMonitor(self.backend.model) - if monitor.enabled: - self.backend.model.ep_balance_monitor = monitor return def exposed_get_max_total_token_num(self): diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index f954f5850c..6d4544182c 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -108,6 +108,64 @@ def get_lightllm_websocket_max_message_size(): @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" + + +@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_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 + + +@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_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_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 + + def get_triton_autotune_level(): return int(os.getenv("LIGHTLLM_TRITON_AUTOTUNE_LEVEL", 0)) 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..8cbca5bdb2 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb.py @@ -0,0 +1,3470 @@ +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 envs_utils, dist_utils +from lightllm.common.eplb_utils import extract_eplb_expert_tensors +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( + disable_eplb_model_init, + is_eplb_model_init_disabled, +) +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, + ): + seen["fused"] = { + "topk_ids": topk_ids, + "activation": (alpha, limit, clamp_up_add_one), + } + return "output" + + seen, captured = {}, [] + 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, + ) + 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) + + +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.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.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_steady_state_sparse_sampling_records_steps_sixteen_to_nineteen_and_evaluates_step_twenty( + monkeypatch, +): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + counter = torch.tensor([[10, 20, 30, 40], [0, 0, 0, 0]], dtype=torch.int64) + manager.weights = [ + type( + "Weight", + (), + { + "expert_parallel_state": _test_parallel_state( + eplb=True, + route_counter=counter, + recording=False, + recorded_sample_count=1, + num_logical_experts=4, + world_size=1, + ), + }, + )() + ] + manager.in_flight = False + manager.prefill_steps = 15 + manager.step_interval = 20 + manager.sampling_interval = manager.step_interval + manager.evaluation_group = object() + manager.num_logical_experts = 4 + manager.global_rank = 1 + manager.evaluation_in_flight = False + manager._steady_collection_end_step = None + manager._continuous_collection_start_step = None + manager._continuous_collection_end_step = None + + recordings, resets, started = [], [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._reset_recorded_samples = lambda: resets.append(True) + monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(True)) + monkeypatch.setattr( + manager_module.torch.cuda, + "synchronize", + lambda: pytest.fail("step must not synchronize CUDA"), + ) + + manager.step() + assert manager.prefill_steps == 16 + assert recordings == [True] + assert resets == [True] + assert manager._steady_collection_end_step is not None + assert manager._steady_collection_end_step == 20 + + for _ in range(3): + manager.step() + assert manager.prefill_steps == 19 + assert recordings == [True] + assert started == [] + + manager.step() + assert manager.prefill_steps == 20 + assert manager._steady_collection_end_step is None + assert started == [True] + + +def test_steady_sampling_window_clamps_to_short_interval_without_moving_boundary(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager.evaluation_in_flight = False + manager.prefill_steps = 0 + manager.step_interval = 20 + manager.sampling_interval = 3 + manager._steady_collection_end_step = None + manager._continuous_collection_start_step = None + manager._continuous_collection_end_step = None + recordings, resets, started = [], [], [] + manager._set_recording = lambda enabled: recordings.append((manager.prefill_steps, enabled)) + manager._reset_recorded_samples = lambda: resets.append(manager.prefill_steps) + monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(manager.prefill_steps)) + + manager._prepare_next_sampling_window() + assert resets == [0] + assert recordings == [(0, True)] + assert manager._steady_collection_end_step == 3 + assert manager._steady_collection_end_step is not None + + manager.step() + manager.step() + assert started == [] + manager.step() + assert started == [3] + + +def test_eplb_step_does_not_start_a_second_evaluation_while_one_is_pending(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + counter = torch.tensor([[10, 20, 30, 40], [0, 0, 0, 0]], dtype=torch.int64) + manager.weights = [ + type( + "Weight", + (), + { + "expert_parallel_state": _test_parallel_state( + eplb=True, + route_counter=counter, + recording=False, + recorded_sample_count=1, + num_logical_experts=4, + world_size=1, + ), + }, + )() + ] + manager.in_flight = False + manager.prefill_steps = 1 + manager.step_interval = 2 + manager.sampling_interval = manager.step_interval + manager.evaluation_group = object() + manager.num_logical_experts = 4 + manager.global_rank = 1 + manager.evaluation_in_flight = True + + started = [] + monkeypatch.setattr(manager, "_poll_evaluation", lambda: True) + monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(True)) + + manager.step() + + assert manager.prefill_steps == 1 + assert started == [] + + +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.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 == [False] + 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 == 80 + + +def test_interval_one_rearms_after_evaluation_but_never_evaluates_empty_counter( + monkeypatch, +): + class DoneThread: + def join(self): + pass + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager.prefill_steps = 1 + manager.step_interval = 1 + manager.sampling_interval = 1 + manager._steady_collection_end_step = None + 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 = 1 + manager.weights = [] + manager._eplb_states = [] + recordings, starts = [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._start_evaluation = lambda: starts.append(True) + manager._evaluation_ready_on_all_ranks = lambda: True + + manager.poll() + + # The no-improvement backoff changes interval 1 to 4. The clamped + # steady window arms immediately but still waits for boundary step 5. + assert recordings == [True] + assert manager._steady_collection_end_step is not None + assert manager._steady_collection_end_step == 5 + assert starts == [] + assert manager.prefill_steps == 1 + + for _ in range(3): + manager.step() + assert starts == [] + manager.step() + assert starts == [True] + + +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.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. + + +def test_manager_collects_recent_ring_samples_in_chronological_order(monkeypatch): + counters = [ + torch.tensor([[10, 11], [20, 21], [30, 31]], dtype=torch.int64), + torch.tensor([[40, 41], [50, 51], [60, 61]], dtype=torch.int64), + ] + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.weights = [ + type( + "Weight", + (), + { + "expert_parallel_state": _test_parallel_state( + eplb=True, + route_counter=counter, + recording=False, + recorded_sample_count=5, + num_logical_experts=2, + world_size=1, + ), + }, + )() + for counter in counters + ] + manager._eplb_states = [weight.expert_parallel_state.eplb for weight in manager.weights] + manager.evaluation_group = object() + metadata_sizes = [] + + def all_reduce(metadata, **_kwargs): + metadata_sizes.append(metadata.numel()) + + monkeypatch.setattr(manager_module.dist, "all_reduce", all_reduce) + + samples = manager._collect_local_samples() + + assert metadata_sizes == [4] + assert torch.equal( + samples, + torch.tensor( + [ + [[30, 31], [60, 61]], + [[10, 11], [40, 41]], + [[20, 21], [50, 51]], + ], + dtype=torch.int64, + ), + ) + + +def test_manager_collects_only_two_recent_sparse_samples(monkeypatch): + counters = [ + torch.tensor([[10, 11], [20, 21], [30, 31], [40, 41]], dtype=torch.int64), + torch.tensor([[50, 51], [60, 61], [70, 71], [80, 81]], dtype=torch.int64), + ] + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.weights = [ + type( + "Weight", + (), + { + "expert_parallel_state": _test_parallel_state( + eplb=True, + route_counter=counter, + recording=False, + recorded_sample_count=2, + num_logical_experts=2, + world_size=1, + ), + }, + )() + for counter in counters + ] + manager._eplb_states = [weight.expert_parallel_state.eplb for weight in manager.weights] + manager.evaluation_group = object() + monkeypatch.setattr(manager_module.dist, "all_reduce", lambda tensor, **kwargs: None) + + samples = manager._collect_local_samples() + + assert torch.equal( + samples, + torch.tensor([[[10, 11], [50, 51]], [[20, 21], [60, 61]]], dtype=torch.int64), + ) + + +@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 = type( + "Args", + (), + {"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 + 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) + 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._start_steady_sampling_window(steps) + expected = [] + for dispatch in range(2 * steps): + # 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]) + samples = manager._recent_ring_samples(state.route_counter, state.recorded_sample_count).cpu() + assert torch.equal(samples, torch.tensor(expected, dtype=torch.int64)) + + +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") + + +def test_ep_without_eplb_creates_layout_without_eplb_runtime_state(monkeypatch): + args = type( + "Args", + (), + {"enable_prefill_eplb": False, "eplb_num_redundant_experts_per_rank": 2}, + )() + 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 + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + + weight._init_expert_parallel_state() + + 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") + + +def test_disable_eplb_model_init_skips_eplb_state(monkeypatch): + args = type( + "Args", + (), + {"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 + monkeypatch.setattr(envs_utils, "get_env_start_args", lambda: args) + monkeypatch.setattr( + placement_module, + "build_initial_redundant_expert_ids", + lambda *args, **kwargs: pytest.fail("disabled scope must not initialize EPLB"), + ) + + with disable_eplb_model_init(): + weight._init_expert_parallel_state() + + 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") + + +def test_disable_eplb_model_init_scope_restores_after_exception(): + assert not is_eplb_model_init_disabled() + with disable_eplb_model_init(): + assert is_eplb_model_init_disabled() + with disable_eplb_model_init(): + assert is_eplb_model_init_disabled() + assert is_eplb_model_init_disabled() + + with pytest.raises(RuntimeError): + with disable_eplb_model_init(): + raise RuntimeError + assert not is_eplb_model_init_disabled() + + +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._continuous_collection_end_step = None + 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 + + +def test_manager_planned_evaluation_builds_improved_metadata_in_one_multilayer_call(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=4, + ) + }, + )() + for _ in range(3) + ] + manager._eplb_states = [weight.expert_parallel_state.eplb for weight in manager.weights] + 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 = build_initial_redundant_expert_ids(4, 4, 1).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._continuous_collection_end_step = None + 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], [1], [2]], + [[2], [3], [0], [1]], + [[1], [2], [3], [0]], + ], + 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, 1])] + 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, + ) + 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._continuous_collection_end_step = None + 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_keeps_logical_ids_and_uses_logical_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, 127]], 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, 127]] + assert calls[0]["num_experts"] == 128 + + +def test_decode_select_does_not_clone_or_map_eplb_topk_ids(monkeypatch): + from lightllm.common.basemodel.triton_kernel.fused_moe import topk_select + + impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) + impl.routed_scaling_factor = 1.0 + _set_expert_parallel_state(impl, _test_parallel_state(eplb=True)) + topk_ids = torch.tensor([[3, 127]], dtype=torch.int32) + monkeypatch.setattr(topk_select, "select_experts", lambda **_kwargs: (torch.ones((1, 2)), topk_ids)) + _, 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.data_ptr() == origin.data_ptr() + + +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_primary_rows_only_when_eplb_is_enabled( + 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] == 8 + assert captured["w13_scale"].shape[0] == captured["w2_scale"].shape[0] == 8 + + +def test_decode_fused_experts_uses_cached_primary_weight_packs_and_logical_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] == [128, 128] + assert all(call["w13"].weight.shape[0] == call["w2"].weight.shape[0] == 8 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", + ] + + +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._steady_collection_end_step = None + manager._continuous_collection_start_step = 0 + 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._steady_collection_end_step is None + assert manager._continuous_collection_start_step is None + + +def test_manager_sparse_insufficient_schedules_bounded_fresh_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": "insufficient", + "minimum_layer_samples": 1, + "minimum": 2, + } + manager.global_rank = 0 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.prefill_steps = 37 + manager.weights = [] + manager._eplb_states = [] + manager._continuous_collection_start_step = None + manager._continuous_collection_end_step = None + recordings, resets, logs = [], [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._reset_recorded_samples = lambda: resets.append(True) + monkeypatch.setattr(manager_module.logger, "info", lambda *args: logs.append(args)) + + assert not manager._poll_evaluation() + assert not hasattr(manager, "_retained_local_samples") + assert manager._continuous_collection_start_step == 40 + assert manager._continuous_collection_end_step == 60 + assert recordings == [False] + assert resets == [True] + assert manager.sampling_interval == 20 + assert "insufficient samples" in logs[0][0] + assert "scheduled_fresh_window" in logs[0][0] + + manager.in_flight = False + manager.evaluation_in_flight = False + starts = [] + monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(manager.prefill_steps)) + manager.step() + manager.step() + assert manager.prefill_steps == 39 + assert starts == [] + manager.step() + assert manager.prefill_steps == 40 + assert starts == [] + for _ in range(19): + manager.step() + assert manager.prefill_steps == 59 + assert starts == [] + manager.step() + assert starts == [60] + + +def test_manager_full_window_insufficient_clears_and_backs_off(): + 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": "insufficient", + "minimum_layer_samples": 1, + "minimum": 2, + } + manager.global_rank = 1 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.prefill_steps = 60 + manager.weights = [] + manager._eplb_states = [] + manager._continuous_collection_start_step = 40 + manager._continuous_collection_end_step = 60 + recordings, resets = [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._reset_recorded_samples = lambda: resets.append(True) + + assert not manager._poll_evaluation() + assert not hasattr(manager, "_retained_local_samples") + assert manager._continuous_collection_start_step is None + assert manager._continuous_collection_end_step is None + assert manager.sampling_interval == 80 + assert recordings == [False] + assert resets == [True] + + +def test_begin_continuous_collection_uses_full_window_at_fixed_boundary(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager.evaluation_in_flight = False + manager.prefill_steps = 36 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager._steady_collection_end_step = 40 + recordings, resets, starts = [], [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._reset_recorded_samples = lambda: resets.append(True) + monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(True)) + + # The window is never truncated to the next boundary: it waits until 40, + # records a full 20 fresh steps, then evaluates at the boundary at 60. + manager._begin_continuous_collection() + assert manager._continuous_collection_start_step == 40 + assert manager._continuous_collection_end_step == 60 + assert recordings == [False] + assert resets == [True] + assert manager._steady_collection_end_step is None + + for _ in range(4): + manager.step() + assert manager.prefill_steps == 40 + assert recordings == [False, True] + assert starts == [] + for _ in range(19): + manager.step() + assert starts == [] + manager.step() # 60: the full window ends and triggers the evaluation. + assert starts == [True] + + +def test_begin_continuous_collection_preserves_full_window_at_sparse_boundary(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.prefill_steps = 80 + manager.step_interval = 20 + manager.sampling_interval = 80 + manager._steady_collection_end_step = None + recordings = [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._reset_recorded_samples = lambda: None + + manager._begin_continuous_collection() + assert manager._continuous_collection_start_step == 140 + assert manager._continuous_collection_end_step == 160 + assert recordings == [False] + + +def test_first_no_improvement_switches_to_sparse_sampling_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 = 1 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager._continuous_collection_start_step = 0 + manager.weights = [] + manager._eplb_states = [] + recordings = [] + manager._set_recording = lambda enabled: recordings.append(enabled) + + assert not manager._poll_evaluation() + assert manager._continuous_collection_start_step is None + assert recordings == [False] + assert manager.sampling_interval == 80 + + +def test_continuous_collection_evaluates_only_after_one_full_base_window(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager._continuous_collection_start_step = 0 + manager._continuous_collection_end_step = 20 + manager.prefill_steps = 0 + manager.step_interval = 20 + manager.sampling_interval = 320 + manager._steady_collection_end_step = None + manager.evaluation_in_flight = False + started = [] + monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(True)) + + for _ in range(19): + manager.step() + assert manager.prefill_steps == 19 + assert started == [] + + manager.step() + assert manager.prefill_steps == 20 + assert started == [True] + + +def test_no_improvement_exponentially_backs_off_sampling_interval_at_cap(): + class DoneThread: + def join(self): + pass + + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager._evaluation_lock = threading.Lock() + manager._set_recording = lambda _enabled: None + manager.global_rank = 1 + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.weights = [] + manager._eplb_states = [] + + for expected_interval in (80, 320, 320): + manager.evaluation_in_flight = True + 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, + } + assert not manager._poll_evaluation() + assert manager.sampling_interval == expected_interval + + +def test_sparse_backoff_arms_and_evaluates_only_at_new_interval_boundary(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager.prefill_steps = 18 + manager.step_interval = 20 + manager.sampling_interval = 80 + manager._steady_collection_end_step = None + manager._continuous_collection_start_step = None + manager._continuous_collection_end_step = None + manager.evaluation_in_flight = False + recordings, resets, starts = [], [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._reset_recorded_samples = lambda: resets.append(True) + monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(True)) + + manager.step() + manager.step() + assert manager.prefill_steps == 20 + assert recordings == [] + assert starts == [] + + for _ in range(55): + manager.step() + assert manager.prefill_steps == 75 + assert recordings == [] + assert starts == [] + + manager.step() + assert manager.prefill_steps == 76 + assert recordings == [True] + assert resets == [True] + assert manager._steady_collection_end_step is not None + + for _ in range(4): + manager.step() + assert manager.prefill_steps == 80 + assert starts == [True] + assert manager._steady_collection_end_step is None + + +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._continuous_collection_start_step = 0 + 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._continuous_collection_start_step is None + 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_first_rebalance_completion_switches_to_four_step_sparse_window(monkeypatch): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.step_interval = 20 + manager.sampling_interval = manager.step_interval + manager._steady_collection_end_step = None + manager.weights = [] + manager._eplb_states = [] + manager.target_placement = torch.zeros(1) + manager.in_flight_started_at = 0 + manager.global_rank = 1 + recordings, starts = [], [] + manager._set_recording = lambda enabled: recordings.append(enabled) + manager._finish_rebalance() + assert recordings == [False] + + manager.in_flight = False + manager.prefill_steps = 38 + manager.evaluation_in_flight = False + monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(True)) + manager.prefill_steps = 35 + manager.step() + assert recordings == [False, True] + for _ in range(4): + manager.step() + assert starts == [True] + + +def test_manager_inflight_step_does_not_poll(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = True + calls = [] + manager._poll_in_flight = lambda: calls.append("poll") + manager.step() + assert calls == [] + + +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_both_parts_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), (3, 1), (4, 0), (4, 1), (5, 0), (5, 1)] + assert [layer.layer_index for layer in prepared] == [3, 4, 5] + assert [layer.parts for layer in prepared] == [ + ("prepared-3-0", "prepared-3-1"), + ("prepared-4-0", "prepared-4-1"), + ("prepared-5-0", "prepared-5-1"), + ] + + +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 the part and both exact address ranges. + transfer._get_remote_read(1, first_entries, 1) + assert agent.made == 3 + transfer._remote_layouts[1][0][0] = ("weight", 1008, 0, 8) + transfer._get_remote_read(1, first_entries, 0) + assert agent.made == 4 + staging[0][1].pointer += 16 + transfer._get_remote_read(1, first_entries, 0) + assert agent.made == 5 + + transfer.shutdown() + assert agent.released_xfers == 5 + assert agent.released_dlists == 10 + 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, 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._continuous_collection_start_step is None + assert manager._continuous_collection_end_step == manager.step_interval + 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 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("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_eplb_configured_window_equal_interval_arms_immediately(): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.step_interval = 20 + manager.sampling_interval = 20 + manager.steady_sample_steps = 20 + manager._continuous_collection_start_step = None + manager._continuous_collection_end_step = None + manager._steady_collection_end_step = None + manager.prefill_steps = 0 + calls = [] + manager._reset_recorded_samples = lambda: calls.append("reset") + manager._set_recording = lambda value: calls.append(value) + manager._prepare_next_sampling_window() + assert manager._steady_collection_end_step == 20 + assert calls == ["reset", True] 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..500631f530 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_full_layout.py @@ -0,0 +1,327 @@ +"""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) + + +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.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 + + args = SimpleNamespace( + 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 + 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) + 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..fa4c74d6e8 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_memory_profile.py @@ -0,0 +1,307 @@ +"""EPLB reservations reach real KV profiling before any KV allocation.""" +from types import SimpleNamespace +from contextlib import nullcontext +from concurrent.futures import ThreadPoolExecutor + +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.glm5_next import model as glm_module +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.qwen3next import model as next_module +from lightllm.server.core.objs.start_args_type import StartArgs + + +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(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 = SimpleNamespace(enable_prefill_eplb=True) + 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, 2304), (2, 2376)]) +@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) + + with memory_module.memory_reservation_scope(eplb_utils.get_eplb_memory_reservations(model)): + with pytest.raises(KVAllocationReached): + model._init_mem_manager() + + expected = {"eplb_staging": staging * 1024 ** 2, "eplb_sampling": 2 * 8 * 40 * 512 * 8 + 512} + assert manager.memory_reservations == expected + budget = profile_without_allocating + available = budget.free - budget.total * (1 - model.mem_fraction) + assert manager.size == int((available - sum(expected.values())) / manager.get_cell_size()) + # The original full R0 bug leaves ~1.6 GiB for a 2.25 GiB staging allocation. + assert budget.free - manager.size * manager.get_cell_size() >= sum(expected.values()) + + +@pytest.mark.parametrize("enabled,draft", [(False, False), (True, True)]) +def test_disabled_eplb_and_draft_do_not_reserve_target_buffers(enabled, draft): + model = make_model() + model.args.enable_prefill_eplb = enabled + model.is_mtp_draft_model = draft + assert eplb_utils.get_eplb_memory_reservations(model) == {} + + +def test_weight_aliases_are_deduplicated_and_non_eplb_layers_are_excluded(): + 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] + assert eplb_utils.get_eplb_memory_reservations(model)["eplb_sampling"] == 2 * 2 * 40 * 512 * 8 + 512 + + +@pytest.mark.parametrize("redundant,expected", [(0, 27), (2, 45)]) +@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_deduct_reservations_once(profile_without_allocating, manager_class): + budget = profile_without_allocating + budget.free, budget.total = 64 * 1024 ** 2, 0 + reservations = {"eplb_staging": 1024, "eplb_sampling": 512} + manager = manager_class.__new__(manager_class) + kwargs = dict(size=None, dtype=torch.bfloat16, head_num=1, head_dim=576, layer_num=8) + 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 memory_module.memory_reservation_scope(reservations): + with pytest.raises(KVAllocationReached): + manager_class.__init__(manager, **kwargs) + assert manager.memory_reservations == reservations + assert manager.size == (budget.free - 1536) // manager.get_cell_size() + + +def test_reservation_scope_restores_after_nested_failure_and_is_thread_local(profile_without_allocating): + def construct(reservations=None): + manager = memory_module.MemoryManager.__new__(memory_module.MemoryManager) + with pytest.raises(KVAllocationReached): + memory_module.MemoryManager.__init__( + manager, None, torch.bfloat16, 1, 128, 8, memory_reservations=reservations + ) + return manager + + assert construct().memory_reservations == {} + with memory_module.memory_reservation_scope({"outer": 1024}): + outer = construct() + with pytest.raises(RuntimeError, match="nested failure"): + with memory_module.memory_reservation_scope({"inner": 512}): + assert construct().memory_reservations == {"inner": 512} + raise RuntimeError("nested failure") + assert construct().memory_reservations == {"outer": 1024} + assert construct({}).memory_reservations == {} + with ThreadPoolExecutor(max_workers=1) as executor: + assert executor.submit(construct).result().memory_reservations == {} + assert outer.memory_reservations == {"outer": 1024} + assert construct().memory_reservations == {} + + +def test_base_model_scopes_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): + # EPLB budgeting must happen after request-state allocation. + 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) + + 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_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, "profile_mtp_weight_memory", lambda model: nullcontext()) + 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)) + expected = eplb_utils.get_eplb_memory_reservations(model) + assert model.mem_manager.memory_reservations == expected + available = profile_without_allocating.free - profile_without_allocating.total * (1 - model.mem_fraction) + assert model.mem_manager.size == int((available - sum(expected.values())) / model.mem_manager.get_cell_size()) + assert memory_module._memory_reservations.get() is None 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..03d4e79a70 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_planner_worker.py @@ -0,0 +1,88 @@ +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: + return True + + +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) + + +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_sampling_policy.py b/unit_tests/common/fused_moe/test_eplb_sampling_policy.py new file mode 100644 index 0000000000..4d1304c816 --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_sampling_policy.py @@ -0,0 +1,229 @@ +import threading + +import pytest +import torch + +from lightllm.server.router.model_infer.mode_backend import eplb_manager as manager_module +from lightllm.utils import envs_utils + + +@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 _window_manager(step=20, interval=80, steady=20): + manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) + manager.in_flight = False + manager.evaluation_in_flight = False + manager.prefill_steps = 0 + manager.step_interval = step + manager.sampling_interval = interval + manager.steady_sample_steps = steady + manager.min_rebalance_interval = interval + manager._steady_collection_end_step = None + manager._continuous_collection_start_step = 0 + manager._continuous_collection_end_step = step + manager._steady_collection_end_step = None + manager._set_recording = lambda _enabled: None + manager._reset_recorded_samples = lambda: None + return manager + + +def test_initial_window_then_min80_sparse_window_arms_at_60_and_evaluates_at_80(monkeypatch): + manager = _window_manager() + starts = [] + monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(manager.prefill_steps)) + for _ in range(20): + manager.step() + assert starts == [20] # Initial continuous 20-step window is unchanged. + manager._continuous_collection_start_step = None + manager._continuous_collection_end_step = None + manager.prefill_steps = 40 + for _ in range(20): + manager.step() + assert manager.prefill_steps == 60 + assert manager._steady_collection_end_step == 80 + for _ in range(20): + manager.step() + assert starts == [20, 80] + + +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._continuous_collection_start_step = None + manager._continuous_collection_end_step = None + manager._set_recording = lambda _enabled: None + manager._reset_recorded_samples = lambda: None + return manager + + +def test_skip_backoff_honors_min_interval_above_default_cap(): + manager = _no_improvement_manager(80, 640) + for expected in (320, 640, 640): + assert not manager._poll_evaluation() + assert manager.sampling_interval == expected + manager.evaluation_in_flight = True + 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, + } + + +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.prefill_steps = 7 + 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" + 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 + 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 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..c6ab9ae4db --- /dev/null +++ b/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py @@ -0,0 +1,258 @@ +"""Complete-layout EPLB migration across two, three and eight 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(capacity * ((tensor[0].numel() + 1) // 2) * tensor.element_size() 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 + ) + + snapshots = {} + + 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): + assert_values(layer, target) + expected_old = [] + values = (current[rank].cuda() + 10 * layer).view(-1, 1) + expected_old.extend( + [ + (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 snapshot, expected in zip(snapshots[layer], expected_old): + torch.testing.assert_close(snapshot, expected) + state = weights[layer].expert_parallel_state.eplb + state.logical_to_physical_map.copy_(maps[0]) + state.logical_replica_count.copy_(counts[0]) + 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) + 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_copy = transfer._copy_part + + def fail_second(prepared): + if rank == 0 and prepared.part_index == 1: + raise RuntimeError("injected second part failure") + return original_copy(prepared) + + transfer._copy_part = fail_second + plans = [(layer, plan) for layer in (0, 1, 2)] + 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, 8]) +@pytest.mark.parametrize("redundant", [0, 1]) +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_second_part_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_moe_weight_creation.py b/unit_tests/common/fused_moe/test_moe_weight_creation.py new file mode 100644 index 0000000000..1c256975b2 --- /dev/null +++ b/unit_tests/common/fused_moe/test_moe_weight_creation.py @@ -0,0 +1,152 @@ +"""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]) +def test_eplb_one_primary_expert_fp8_constructor_and_checkpoint_load(monkeypatch, device_id, redundant): + 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"}, + ) + 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/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/server/router/model_infer/test_ep_balance_monitor.py b/unit_tests/server/router/model_infer/test_ep_balance_monitor.py index 2ddcaafddf..293d3d6f3f 100644 --- a/unit_tests/server/router/model_infer/test_ep_balance_monitor.py +++ b/unit_tests/server/router/model_infer/test_ep_balance_monitor.py @@ -1,12 +1,16 @@ 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 @@ -217,7 +221,7 @@ def test_monitor_reuses_manager_precreated_dedicated_gloo_group(monkeypatch): ) model = SimpleNamespace(args=SimpleNamespace(enable_tpsp_mix_mode=True), tp_world_size_=1) - monkeypatch.setattr(monitor_module, "_find_fused_moe_weights", lambda _: [weight]) + 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 = [] @@ -250,7 +254,7 @@ def test_nonzero_rank_monitor_does_not_create_metric_client(monkeypatch): ) model = SimpleNamespace(args=SimpleNamespace(enable_tpsp_mix_mode=True), tp_world_size_=1) metric_client_calls = [] - monkeypatch.setattr(monitor_module, "_find_fused_moe_weights", lambda _: [weight]) + 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_group_manager", SimpleNamespace(ep_balance_monitor_group=sentinel_group)) @@ -485,7 +489,7 @@ def __init__(self, layer_num, enabled=True): self.layer_num_ = layer_num self.enable_ep_moe = enabled - monkeypatch.setattr(monitor_module, "FusedMoeWeight", FakeFusedMoeWeight) + monkeypatch.setattr(fused_weight_module, "FusedMoeWeight", FakeFusedMoeWeight) first = FakeFusedMoeWeight(3) second = FakeFusedMoeWeight(1) disabled = FakeFusedMoeWeight(0, enabled=False) @@ -496,7 +500,7 @@ def __init__(self, layer_num, enabled=True): ] ) - assert monitor_module._find_fused_moe_weights(model) == [second, first] + assert monitor_module.find_ep_moe_weights(model) == [second, first] def test_monitor_disable_detaches_counters_from_all_impls(): @@ -544,3 +548,63 @@ def test_prefill_cudagraph_silently_disables_monitor(): 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..d17e33d950 --- /dev/null +++ b/unit_tests/server/router/model_infer/test_ep_runtime.py @@ -0,0 +1,234 @@ +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 + + +@pytest.fixture +def runtime(monkeypatch): + 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(step=lambda: events.append("step"), poll=lambda: events.append("poll")) + backend = SimpleNamespace( + args=SimpleNamespace(enable_prefill_eplb=True, disable_ep_balance_monitor=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]) +def test_completed_rounds_and_decode_preserve_behavior(runtime, eplb, monitor_enabled, batch_size): + backend, events, output, monitor, _ = runtime + backend.args.enable_prefill_eplb = eplb + 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 + assert events == ["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 == ["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: + 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 + + +@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 not events + + +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]) +def test_real_infer_loops_poll_before_requests_even_when_idle(runtime, monkeypatch, backend_class): + template, events, _, _, _ = runtime + backend = backend_class.__new__(backend_class) + vars(backend).update(vars(template)) + backend.support_overlap = False + backend.classed_req_no_decode = False + backend.classed_req_strict_prefill = False + 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: ([], []) + 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 [] + assert events == [ + "poll", + "requests", + *gathered, + "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" + ) 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..b1df345834 --- /dev/null +++ b/unit_tests/server/test_api_start_eplb.py @@ -0,0 +1,102 @@ +import pytest + +from lightllm.server import api_start +from lightllm.server.core.objs.start_args_type import StartArgs + + +@pytest.mark.parametrize( + "redundant,run_mode,message", + [ + (0, "normal", "requires --run_mode prefill"), + (0, "decode", "requires --run_mode prefill"), + (-1, "prefill", "nonnegative"), + ], +) +def test_full_eplb_mode_rejected_before_any_process_starts(monkeypatch, redundant, run_mode, message): + args = StartArgs( + enable_ep_moe=True, + enable_prefill_eplb=True, + eplb_num_redundant_experts_per_rank=redundant, + run_mode=run_mode, + 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, "is_sm100_gpu", lambda: False) + monkeypatch.setattr( + api_start.process_manager, + "start_submodule_processes", + lambda *a, **kw: pytest.fail("must reject before process startup"), + ) + with pytest.raises(AssertionError, match=message): + 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]) +def test_eplb_mtp_combination_is_not_rejected_before_starting_subprocesses(monkeypatch, redundant): + args = StartArgs( + model_dir="test-model", + enable_ep_moe=True, + enable_prefill_eplb=True, + eplb_num_redundant_experts_per_rank=redundant, + run_mode="prefill", + 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(api_start, "is_sm100_gpu", lambda: False) + 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) + monkeypatch.setattr(api_start, "set_env_start_args", lambda args: None) + 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) From 3b40d72db790fe7ccb6b23e34a302e579d69b509 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Fri, 16 Oct 2026 21:30:34 +0800 Subject: [PATCH 4/4] fix(eplb): correct routing costs and streamline rebalancing --- lightllm/common/basemodel/basemodel.py | 10 +- .../meta_weights/fused_moe/eplb_placement.py | 82 +- .../meta_weights/fused_moe/eplb_planner.py | 66 +- .../fused_moe/expert_parallel_state.py | 96 +- .../fused_moe/fused_moe_weight.py | 14 +- .../meta_weights/fused_moe/impl/base_impl.py | 3 + .../fused_moe/impl/deepgemm_impl.py | 69 +- .../fused_moe/impl/marlin_impl.py | 1 + .../fused_moe/impl/triton_impl.py | 1 + .../triton_kernel/fused_moe/grouped_topk.py | 346 ++--- .../triton_kernel/fused_moe/topk_select.py | 20 +- lightllm/common/eplb_planner_worker.py | 71 +- lightllm/common/eplb_utils.py | 43 +- .../kv_cache_mem_manager/mem_manager.py | 35 +- lightllm/common/quantization/awq.py | 8 +- lightllm/common/quantization/deepgemm.py | 18 +- lightllm/common/quantization/no_quant.py | 9 +- .../common/quantization/quantize_method.py | 16 +- lightllm/common/quantization/w8a8.py | 43 +- lightllm/common/quantization/w8a8gx.py | 16 +- lightllm/distributed/communication_op.py | 55 +- lightllm/models/glm5_next_mtp/model.py | 1 + lightllm/models/qwen3_5/model.py | 1 + lightllm/server/api_cli.py | 7 +- lightllm/server/api_start.py | 10 +- .../model_infer/mode_backend/base_backend.py | 6 +- .../mode_backend/ep_balance_monitor.py | 5 +- .../model_infer/mode_backend/ep_runtime.py | 19 +- .../model_infer/mode_backend/eplb_manager.py | 271 ++-- .../model_infer/mode_backend/eplb_transfer.py | 79 +- lightllm/utils/envs_utils.py | 1 + lightllm/utils/profile_max_tokens.py | 57 +- unit_tests/common/fused_moe/test_eplb.py | 1153 +++++++---------- .../common/fused_moe/test_eplb_full_layout.py | 37 +- .../fused_moe/test_eplb_memory_profile.py | 236 +++- .../fused_moe/test_eplb_output_precision.py | 129 ++ .../common/fused_moe/test_eplb_parallel.py | 234 ++++ .../fused_moe/test_eplb_planner_worker.py | 44 +- .../fused_moe/test_eplb_runtime_hash.py | 150 +++ .../fused_moe/test_eplb_sampling_policy.py | 216 ++- .../fused_moe/test_eplb_transfer_gpu.py | 63 +- .../fused_moe/test_eplb_transfer_init.py | 100 ++ .../fused_moe/test_moe_weight_creation.py | 13 +- .../common/fused_moe/test_topk_select.py | 47 + unit_tests/common/test_ep_communication.py | 119 ++ .../model_infer/test_ep_balance_monitor.py | 44 +- .../router/model_infer/test_ep_runtime.py | 115 +- unit_tests/server/test_api_start_eplb.py | 133 +- 48 files changed, 2738 insertions(+), 1574 deletions(-) create mode 100644 unit_tests/common/fused_moe/test_eplb_output_precision.py create mode 100644 unit_tests/common/fused_moe/test_eplb_parallel.py create mode 100644 unit_tests/common/fused_moe/test_eplb_runtime_hash.py create mode 100644 unit_tests/common/fused_moe/test_eplb_transfer_init.py create mode 100644 unit_tests/common/fused_moe/test_topk_select.py create mode 100644 unit_tests/common/test_ep_communication.py diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 930d874ebd..18b5d1c171 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -15,8 +15,6 @@ from lightllm.common.basemodel.infer_struct import InferStateInfo from lightllm.common.kv_cache_mem_manager import MemoryManager from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class -from lightllm.common.eplb_utils import get_eplb_memory_reservations -from lightllm.common.kv_cache_mem_manager.mem_manager import memory_reservation_scope from lightllm.common.req_manager import ReqManager from lightllm.common.build_utils import repair_config from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import ( @@ -30,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, @@ -111,13 +109,12 @@ 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): self._init_req_manager() - with memory_reservation_scope(get_eplb_memory_reservations(self)): - self._init_mem_manager() + self._init_mem_manager() # Qwen3.5 等 linear attention 模型会在 req_manager 中保存大量运行时 state。先初始化 # req_manager、再初始化 mem_manager，可以让 KV cache 显存评估包含这些 state 的实际占用； @@ -151,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"]) 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 index a980abe7e8..a62a5c3d08 100644 --- 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 @@ -1,7 +1,33 @@ 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, @@ -50,6 +76,7 @@ def build_logical_to_physical_map( 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] @@ -60,6 +87,7 @@ def build_logical_to_physical_map( 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) @@ -69,6 +97,7 @@ def build_logical_to_physical_maps_for_layers( 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] @@ -77,13 +106,14 @@ def build_logical_to_physical_maps_for_layers( 第一维是层，不是请求或 token 的 batch，也不会自动处理模型中的其他层。 输入包含全部物理行，副本按 rank-major、slot-major 的稳定顺序排列。 - 指定 source_rank 后，先筛选源节点内副本，再按 source_rank 轮转候选前缀。 + 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, ) @@ -95,13 +125,18 @@ def select_improving_placements( 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) - candidate_rank_load = _estimate_rank_load(expert_load, candidate_placement, expert_alignment, node_world_size) + 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 @@ -224,6 +259,7 @@ def _build_layer_maps( 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: @@ -244,10 +280,15 @@ def _build_layer_maps( ) 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, @@ -261,7 +302,7 @@ def _build_layer_maps( _rotate_selected_replicas( compact_maps_by_layer, selected_counts_by_layer, - source_rank=source_rank, + source_rank=rotation_rank, replica_positions=replica_positions, ), selected_counts_by_layer, @@ -273,6 +314,8 @@ def _estimate_rank_load( 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. @@ -298,6 +341,7 @@ def _estimate_rank_load( 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 @@ -376,10 +420,36 @@ def _expert_rank_load_all( 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.""" - route = _source_route(locations, num_nodes, node_world_size) - physical_load = torch.einsum("slne,lner->sler", source_load.to(torch.float64), route) + 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 index 71a3ae862d..2f18e80965 100644 --- 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 @@ -16,8 +16,12 @@ ) -def _critical_load(load, placement, alignment, node_world_size): - return _estimate_rank_load(load, placement, alignment, node_world_size).max(dim=-1).values.sum(dim=0) +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( @@ -30,6 +34,7 @@ def refine_placement_candidates( 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. @@ -44,7 +49,9 @@ def refine_placement_candidates( 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() - score = _critical_load(load, best, expert_alignment, node_world_size) + 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") @@ -59,26 +66,38 @@ def refine_placement_candidates( 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), + _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) + 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): +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) + 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)[:, 0] + 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: @@ -100,8 +119,9 @@ def _local_search(load, placement, alignment, node_world_size, stickiness, steps 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 + 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) @@ -112,7 +132,7 @@ def _local_search(load, placement, alignment, node_world_size, stickiness, steps 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) + 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 @@ -199,14 +219,32 @@ def interleaved(replacements=replacements, swaps=swaps): return proposals -def plan_eplb_candidate(expert_load, current, *, node_world_size, expert_alignment, stickiness): +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, ) @@ -217,6 +255,7 @@ def plan_full_experts( 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. @@ -229,7 +268,9 @@ def plan_full_experts( 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).sum(dim=0) + 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): @@ -253,6 +294,7 @@ def plan_full_experts( expert_alignment=expert_alignment, node_world_size=node_world_size, stickiness=stickiness, + replica_hash_load=replica_hash_load, ) 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 index b353348574..4f3d5fe12f 100644 --- 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 @@ -1,27 +1,9 @@ -from contextlib import contextmanager -from contextvars import ContextVar from dataclasses import dataclass -from typing import Iterator, Optional +from typing import Callable, Optional import torch -_eplb_model_init_disabled: ContextVar[bool] = ContextVar("eplb_model_init_disabled", default=False) - - -def is_eplb_model_init_disabled() -> bool: - return _eplb_model_init_disabled.get() - - -@contextmanager -def disable_eplb_model_init() -> Iterator[None]: - token = _eplb_model_init_disabled.set(True) - try: - yield - finally: - _eplb_model_init_disabled.reset(token) - - @dataclass class EPLBState: num_redundant_experts_per_rank: int @@ -32,6 +14,9 @@ class EPLBState: 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: @@ -56,9 +41,26 @@ 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 create_expert_parallel_state(num_logical_experts: int, world_size: int, rank: int) -> ExpertParallelState: - from lightllm.common.eplb_utils import EPLB_STEADY_SAMPLE_STEPS + 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, @@ -69,29 +71,65 @@ def create_expert_parallel_state(num_logical_experts: int, world_size: int, rank 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 args.enable_prefill_eplb and not is_eplb_model_init_disabled(): - if args.run_mode != "prefill": - raise ValueError("EPLB requires --run_mode prefill") + 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=get_node_world_size() + 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. - sample_steps = max( - get_prefill_eplb_step_interval(), min(get_prefill_eplb_steady_sample_steps(), EPLB_STEADY_SAMPLE_STEPS) + # 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(), - route_counter=torch.zeros((2 * sample_steps, num_logical_experts), dtype=torch.int64, device="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 67653a7359..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 @@ -59,7 +59,7 @@ 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_expert_parallel_state() + 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, @@ -80,9 +80,16 @@ 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_expert_parallel_state(self): + 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_) + 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 ) @@ -158,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( 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 750aa0fa42..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 @@ -42,6 +42,7 @@ 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, @@ -71,6 +72,7 @@ def __call__( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + dist_group=dist_group, ) @abstractmethod @@ -105,5 +107,6 @@ def _fused_experts( 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 aa090c0067..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 @@ -25,7 +25,6 @@ def __init__(self, *args, expert_parallel_state: ExpertParallelState, **kwargs): self.expert_parallel_state = expert_parallel_state self.eplb = expert_parallel_state.eplb self.ep_balance_counters = None - self._primary_weight_pack_cache = {} def _select_experts( self, @@ -43,14 +42,17 @@ def _select_experts( is_prefill: Optional[bool] = None, preserve_logical_ids: bool = False, ): - """选择 expert；EPLB prefill 统一由融合路径返回 physical ID。""" + """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" eplb = self.eplb - eplb_active = eplb is not None - if is_prefill is True and eplb_active: - from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_topk import triton_grouped_topk_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, + ) - group_score_topk_num = 2 if topk_group == 4 and num_expert_group == 8 and top_k == 8 else 1 topk_weights, topk_ids, logical_topk_ids = triton_grouped_topk_eplb( hidden_states=input_tensor, gating_output=router_logits, @@ -63,12 +65,16 @@ def _select_experts( 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(), - record_load=eplb.recording, + 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, - group_score_used_topk_num=group_score_topk_num, + 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 @@ -103,21 +109,15 @@ def _fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + dist_group=None, ): - if is_prefill is False: - w13 = self._primary_weight_pack(w13) - w2 = self._primary_weight_pack(w2) - num_experts = self.n_routed_experts - else: - num_experts = self.expert_parallel_state.num_total_physical_experts - output = fused_experts( hidden_states=input_tensor, w13=w13, w2=w2, topk_weights=topk_weights, topk_idx=topk_ids.to(torch.long), - num_experts=num_experts, + num_experts=self.expert_parallel_state.num_total_physical_experts, quant_method=self.quant_method, is_prefill=is_prefill, previous_event=None, # for overlap @@ -126,7 +126,7 @@ def _fused_experts( 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, @@ -160,8 +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, - # decode 与 EPLB 的物理冗余行刻意隔离：DeepEP 使用原始 logical expert ID。 - num_experts=self.n_routed_experts, + num_experts=self.expert_parallel_state.num_total_physical_experts, use_fp8=use_fp8_w8a8, async_finish=False, return_recv_hook=True, @@ -251,7 +250,6 @@ def masked_group_gemm( dtype: torch.dtype, expected_m: int, ): - w13, w2 = self._primary_weight_pack(w13), self._primary_weight_pack(w2) w13_weight, w13_scale = w13.weight, w13.weight_scale w2_weight, w2_scale = w2.weight, w2.weight_scale return masked_group_gemm( @@ -349,30 +347,3 @@ def hook(): event.current_stream_wait() return combined_x, hook - - def _primary_weight_pack(self, weight_pack: WeightPack) -> WeightPack: - """返回所有 decode 路径使用的缓存本地主副本视图。""" - if self.eplb is None: - return weight_pack - cache = getattr(self, "_primary_weight_pack_cache", None) - if cache is None: - cache = self._primary_weight_pack_cache = {} - cache_key = id(weight_pack) - primary = cache.get(cache_key) - if primary is None: - num_primary_experts_per_rank = self.expert_parallel_state.num_primary_experts_per_rank - primary = WeightPack( - weight=weight_pack.weight[:num_primary_experts_per_rank], - weight_scale=( - weight_pack.weight_scale[:num_primary_experts_per_rank] - if weight_pack.weight_scale is not None - else None - ), - weight_zero_point=( - getattr(weight_pack, "weight_zero_point", None)[:num_primary_experts_per_rank] - if getattr(weight_pack, "weight_zero_point", None) is not None - else None - ), - ) - cache[cache_key] = primary - return primary 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 d09336a102..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 @@ -37,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 abaec5fc22..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 @@ -66,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 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 196e1d432f..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,8 +1,10 @@ -# 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 + + +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 @@ -15,205 +17,8 @@ def _eplb_replica_index(token_index, logical_id, replica_count): return value % replica_count.to(tl.uint32) -@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 - - -@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 - - @triton.jit def grouped_topk_kernel( - gating_output_ptr, - 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, - group_num, - group_expert_num, - total_expert_num, # group_num * group_expert_num == total_expert_num - topk_num, - group_topk_num, - IS_SIGMOID: 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) - RENORMALIZE: tl.constexpr, - GROUP_SCORE_USED_TOPK_NUM: tl.constexpr, -): - token_index = tl.program_id(axis=0) - offs_n = tl.arange(0, EXPERT_BLOCK_SIZE) - hidden_states = tl.load( - gating_output_ptr + token_index * gating_output_stride_m + offs_n, - mask=offs_n < total_expert_num, - other=-10000000.0, - ).to(tl.float32) - if IS_SIGMOID: - old_scores = tl.sigmoid(hidden_states) - else: - old_scores = tl.softmax(hidden_states) - - if HAS_CORRECTION_BIAS: - scores = old_scores + tl.load(correction_bias_ptr + offs_n, mask=offs_n < total_expert_num, other=-10000000.0) - 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] - - 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 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)) - mask_group_scores = tl.where( - ((group_value >= group_topk_value)[:, None]) & ((offs_group_v < group_expert_num)[None, :]), - group_scores, - -10000000.0, - ) - - 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) - - 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, - ) - tl.store(out_topk_ids + token_index * out_topk_ids_stride_m + offs_n, sorted_indexes, mask=offs_n < topk_num) - else: - tl.store( - out_topk_weights + token_index * out_topk_weights_stride_m + offs_n, sorted_scores, mask=offs_n < topk_num - ) - tl.store(out_topk_ids + token_index * out_topk_ids_stride_m + offs_n, sorted_indexes, mask=offs_n < topk_num) - return - - -@triton.jit -def grouped_topk_eplb_kernel( gating_output_ptr, gating_output_stride_m, gating_output_stride_n, @@ -246,11 +51,16 @@ def grouped_topk_eplb_kernel( 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, ): - """Grouped top-k, EPLB accounting, and replica mapping without a global score workspace.""" + """Shared logical top-k; ties choose the lowest expert ID, with optional EPLB mapping.""" token_index = tl.program_id(axis=0) offs_group = tl.arange(0, EXPERT_GROUP_NUM) offs_group_v = tl.arange(0, EXPERT_GROUP_SIZE) @@ -339,25 +149,46 @@ def grouped_topk_eplb_kernel( topk_mask = topk_offsets < TOPK_NUM if RECORD_LOAD: tl.atomic_add( - expert_counter_ptr + sample_index * COUNTER_NUM_EXPERTS + selected_logical_ids, + expert_counter_ptr + (sample_index * COUNTER_NUM_EXPERTS + selected_logical_ids) * COUNTER_CHANNELS, 1, mask=topk_mask, sem="relaxed", ) - if SINGLE_TOKEN: - replica_indices = tl.zeros((TOPK_BLOCK_SIZE,), tl.int32) - else: - replica_counts = tl.load( - logical_replica_count_ptr + selected_logical_ids, + 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, + 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, - ) + else: + selected_physical_ids = selected_logical_ids if RENORMALIZE: selected_weights /= sum_scores tl.store( @@ -389,55 +220,24 @@ def triton_grouped_topk( scoring_func: str = "softmax", group_score_used_topk_num=2, ): - 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 - - if token_num <= 256: - num_warps = 4 - else: - num_warps = 1 - - 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, - 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), - RENORMALIZE=renormalize, - GROUP_SCORE_USED_TOPK_NUM=group_score_used_topk_num, - num_warps=num_warps, - num_stages=1, + 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 out_topk_weights, out_topk_ids + return weights, ids def triton_grouped_topk_eplb( @@ -457,8 +257,19 @@ def triton_grouped_topk_eplb( 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, ): - """Fused EPLB prefill top-k returning physical IDs and optional logical IDs.""" + """One logical selector; absent EPLB maps compile out accounting and replica selection.""" + from math import isqrt + + 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) @@ -480,7 +291,7 @@ def triton_grouped_topk_eplb( 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_eplb_kernel[(token_num,)]( + grouped_topk_kernel[(token_num,)]( gating_output, *gating_output.stride(), correction_bias, @@ -508,8 +319,13 @@ def triton_grouped_topk_eplb( 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], - MAP_SLOTS=logical_to_physical_map.shape[1], + 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, 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 d2f59de480..ad2aff269b 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/topk_select.py @@ -21,6 +21,7 @@ from lightllm.utils.sgl_utils import sgl_ops from typing import Callable, List, Optional, Tuple from lightllm.common.basemodel.triton_kernel.fused_moe.softmax_topk import softmax_topk +from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType def fused_topk( @@ -135,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, @@ -155,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: @@ -167,4 +166,11 @@ def select_experts( hidden_states=hidden_states, gating_output=router_logits, topk=top_k, renormalize=renormalize ) + # 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]) + router_logits = torch.randn(size=router_logits.shape, generator=rand_gen, dtype=torch.float32, device="cuda") + _, topk_ids = torch.topk(router_logits, k=top_k, dim=1) + return topk_weights, topk_ids diff --git a/lightllm/common/eplb_planner_worker.py b/lightllm/common/eplb_planner_worker.py index 3e64901395..41de36b04c 100644 --- a/lightllm/common/eplb_planner_worker.py +++ b/lightllm/common/eplb_planner_worker.py @@ -2,20 +2,56 @@ 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 = mp.get_context("spawn").Process(target=_serve, args=(child, os.getpid()), daemon=True) + self._process = None try: - self._process.start() + # 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}") @@ -26,6 +62,7 @@ def __init__(self, timeout_s=120.0): except OSError: pass self.close() + self._parent.close() raise atexit.register(self.close) @@ -128,6 +165,26 @@ def _watch_parent(expected_parent_pid): 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() @@ -136,7 +193,7 @@ def _serve(conn, expected_parent_pid): torch.set_num_threads(1) torch.set_num_interop_threads(1) - from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import eplb_planner + eplb_planner = _load_cpu_planner() # Startup prewarm is intentionally small and CPU-only. eplb_planner.plan_eplb_candidate( @@ -182,3 +239,11 @@ def _serve(conn, expected_parent_pid): 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 index bdf5a9d9e7..b3fd0b2095 100644 --- a/lightllm/common/eplb_utils.py +++ b/lightllm/common/eplb_utils.py @@ -1,12 +1,34 @@ -"""Small, dependency-free EPLB helpers shared by transfer and model profiling.""" +"""EPLB initialization, transfer and model profiling helpers.""" import math -EPLB_FULL_STAGING_PARTS = 2 +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 @@ -19,23 +41,8 @@ def find_ep_moe_weights(model): return sorted(weights.values(), key=lambda weight: weight.layer_num_) -def get_eplb_memory_reservations(model): - """Budget buffers created after KV profiling; persistent counters already exist.""" - if model.is_mtp_draft_model or not model.args.enable_prefill_eplb: - return {} - weights = [weight for weight in find_ep_moe_weights(model) if weight.expert_parallel_state.eplb is not None] - return { - name: value - for name, value in ( - ("eplb_staging", get_eplb_staging_nbytes(weights)), - ("eplb_sampling", get_eplb_sampling_peak_nbytes(weights)), - ) - if value - } - - def get_eplb_staging_shape(weights): - """Stage two row slices of one complete layer, independent of layer count.""" + """Stage one complete layer before modifying any live expert rows.""" if not weights: return 0, 0 parallel = weights[0].expert_parallel_state diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 9905438930..38c870ce4e 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -1,8 +1,6 @@ import math import os import re -from contextlib import contextmanager -from contextvars import ContextVar import torch import torch.distributed as dist import torch.multiprocessing as mp @@ -27,37 +25,19 @@ from .operator import BaseMemManagerOperator, NormalMemOperator logger = init_logger(__name__) -_memory_reservations = ContextVar("kv_memory_reservations", default=None) - - -@contextmanager -def memory_reservation_scope(reservations): - """Supply post-profile budgets only to the synchronous KV construction stack.""" - token = _memory_reservations.set(dict(reservations)) - try: - yield - finally: - _memory_reservations.reset(token) class MemoryManager: + operator_class = NormalMemOperator - def __init__( - self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9, memory_reservations=None - ): + def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False, mem_fraction=0.9): self.size = size self.head_num = head_num self.head_dim = head_dim self.layer_num = layer_num self.always_copy = always_copy self.dtype = dtype - # These buffers are allocated after KV profiling and deducted from its budget once. - if memory_reservations is None: - memory_reservations = _memory_reservations.get() - self.memory_reservations = dict(memory_reservations or {}) - if any(value < 0 for value in self.memory_reservations.values()): - raise ValueError(f"memory reservations must be non-negative: {self.memory_reservations}") # profile the max total token num if the size is None self.profile_size(mem_fraction) @@ -117,16 +97,8 @@ def profile_size(self, mem_fraction): world_size = dist.get_world_size() available_memory = get_available_gpu_memory(world_size) - get_total_gpu_memory() * (1 - mem_fraction) cell_size = self.get_cell_size() - reservations = getattr(self, "memory_reservations", {}) - reserved_memory_size = sum(reservations.values()) pd_kv_move_buffer_size = self.get_pd_kv_move_buffer_size() - available_memory_bytes = available_memory * 1024 ** 3 - reserved_memory_size - pd_kv_move_buffer_size - if available_memory_bytes <= 0: - raise RuntimeError( - f"{type(self).__name__} reservations require {reserved_memory_size / 1024**3:.2f} GB, " - f"plus {pd_kv_move_buffer_size / 1024**3:.2f} GB for the PD KV transfer buffer, " - f"but only {available_memory:.2f} GB is available" - ) + available_memory_bytes = available_memory * 1024 ** 3 - pd_kv_move_buffer_size self.size = int(available_memory_bytes / cell_size) if world_size > 1: tensor = torch.tensor(self.size, dtype=torch.int64, device=f"cuda:{get_current_device_id()}") @@ -134,7 +106,6 @@ def profile_size(self, mem_fraction): self.size = tensor.item() logger.info( f"{str(available_memory)} GB space is available after load the model weight\n" - f"{reservations} bytes are reserved for post-profile model buffers\n" f"{str(pd_kv_move_buffer_size / 1024 ** 2)} MB is reserved for PD KV transfer buffer\n" f"{str(cell_size / 1024 ** 2)} MB is the size of one token kv cache\n" f"{self.size} is the profiled max_total_token_num with the mem_fraction {mem_fraction}\n" 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 e535bae0b4..641c72d67e 100644 --- a/lightllm/common/quantization/quantize_method.py +++ b/lightllm/common/quantization/quantize_method.py @@ -76,24 +76,13 @@ def create_weight( def create_moe_weight( self, out_dims: List[int], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int ) -> Tuple[WeightPack, List[WeightPack]]: - weight_pack, weight_pack_list = self._create_weight( + return self._create_weight( out_dims=out_dims, in_dim=in_dim, dtype=dtype, device_id=device_id, num_experts=num_experts, ) - if num_experts == 1: - # Dense factories omit the singleton expert axis. Restore it on - # the fused pack and its loading views without copying storage. - for pack in [weight_pack, *weight_pack_list]: - if pack.weight.ndim == 2: - pack.weight = pack.weight.unsqueeze(0) - if pack.weight_scale is not None: - pack.weight_scale = pack.weight_scale.unsqueeze(0) - if pack.weight_zero_point is not None: - pack.weight_zero_point = pack.weight_zero_point.unsqueeze(0) - return weight_pack, weight_pack_list def load_weight(self, weight: torch.Tensor, weight_pack: WeightPack) -> None: if self._check_weight_need_quanted(weight): @@ -123,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 d18c4a780f..c8ee8eb5eb 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -24,6 +24,7 @@ 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 ( @@ -34,6 +35,7 @@ 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, ) @@ -66,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(): @@ -107,7 +116,6 @@ def all_gather_into_tensor(self, output_: torch.Tensor, input_: torch.Tensor, as class DistributeGroupManager: def __init__(self): self.groups = [] - self.ep_balance_monitor_group = None self.ep_buffer = None self.ep_low_latency_buffer = None self.ep_mega_moe_buffer = None @@ -125,14 +133,6 @@ def create_groups(self, group_size: int): if not args.disable_flashinfer_allreduce: group.init_flashinfer_reduce() self.groups.append(group) - if ( - getattr(args, "enable_ep_moe", False) - and not getattr(args, "disable_ep_balance_monitor", False) - and getattr(args, "run_mode", "normal") != "decode" - and not getattr(args, "enable_prefill_cudagraph", False) - and not is_sm100_gpu() - ): - self.ep_balance_monitor_group = dist.new_group(ranks=list(range(get_global_world_size())), backend="gloo") return def get_default_group(self) -> CustomProcessGroup: @@ -203,13 +203,11 @@ def new_deepep_group( self.ll_hidden = hidden_size total_redundant_experts = ( get_env_start_args().eplb_num_redundant_experts_per_rank * global_world_size - if get_env_start_args().enable_prefill_eplb + if is_eplb_enabled(get_env_start_args()) else 0 ) - self.ll_prefill_num_experts = n_routed_experts + total_redundant_experts - # EPLB's redundant rows are a prefill-only physical layout; decode - # always routes the logical expert space. - self.ll_decode_num_experts = n_routed_experts + # 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, @@ -252,7 +250,7 @@ def new_deepep_group( self.ll_decode_num_tokens, self.ll_hidden, global_world_size, - self.ll_decode_num_experts, + self.ll_num_experts, ) microbatch_count = len(self.groups) min_prefill_reuse_buffer_bytes = _calculate_min_chunked_expanded_moe_reuse_buffer_bytes( @@ -273,7 +271,7 @@ def new_deepep_group( deepep_group, num_rdma_bytes=num_rdma_bytes, low_latency_mode=True, - num_qps_per_rank=(self.ll_decode_num_experts // global_world_size), + num_qps_per_rank=(self.ll_num_experts // global_world_size), ) if enable_mega_moe_buffer: @@ -286,7 +284,7 @@ def new_deepep_group( self.ep_mega_moe_buffer = deep_gemm.get_symm_buffer_for_mega_moe( deepep_group, - self.ll_decode_num_experts, + self.ll_num_experts, self.ll_num_tokens, num_experts_per_tok, self.ll_hidden, @@ -294,18 +292,16 @@ def new_deepep_group( ) logger.info( "Initialize DeepEP MoE buffers: low_latency=%s, mega_moe=%s, " - "ll_prefill_num_experts=%s, ll_decode_num_experts=%s, expert_quant_method_names=%s", - enable_low_latency_buffer, + "ll_num_experts=%s, expert_quant_method_names=%s", + self.ep_low_latency_buffer is not None, enable_mega_moe_buffer, - self.ll_prefill_num_experts, - self.ll_decode_num_experts, + self.ll_num_experts, sorted(expert_quant_method_names), ) - theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_prefill_num_experts, num_experts_per_tok) - low_latency_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_decode_num_experts, num_experts_per_tok) - self._set_num_sms_for_deep_gemm(theoretical_sms, low_latency_sms) + theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok) + self._set_num_sms_for_deep_gemm(theoretical_sms) - def _set_num_sms_for_deep_gemm(self, deepep_sms: int, low_latency_sms: int): + def _set_num_sms_for_deep_gemm(self, deepep_sms: int): try: try: from deep_gemm.jit_kernels.utils import set_num_sms @@ -314,12 +310,9 @@ def _set_num_sms_for_deep_gemm(self, deepep_sms: int, low_latency_sms: int): device_sms = get_device_sm_count() deepep_sms = max(0, min(deepep_sms, max(device_sms - 2, 0))) - low_latency_sms = max(0, min(low_latency_sms, max(device_sms - 2, 0))) self.ep_num_sms = deepep_sms if self.ep_low_latency_buffer is not None: - # This setting controls the legacy low-latency buffer; keep - # its SM reservation based on decode's logical expert count. - deep_ep.Buffer.set_num_sms(low_latency_sms - low_latency_sms % 2) + deep_ep.Buffer.set_num_sms(deepep_sms - deepep_sms % 2) set_num_sms(max(device_sms - deepep_sms, 2)) except BaseException as e: logger.warning(f"set num sms for deep_gemm failed: {e}") @@ -361,7 +354,7 @@ def clear_deepep_buffer(self): """ if self.ep_low_latency_buffer is not None: self.ep_low_latency_buffer.clean_low_latency_buffer( - self.ll_decode_num_tokens, self.ll_hidden, self.ll_decode_num_experts + self.ll_decode_num_tokens, self.ll_hidden, self.ll_num_experts ) 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 cee9b4f5d2..70274c8508 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -809,13 +809,16 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--enable_prefill_eplb", action="store_true", - help="""Enable online expert load balancing with complete physical layouts; requires --run_mode prefill.""", + 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 prefill EPLB. + 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( diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 7bb29c400f..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 @@ -33,7 +34,6 @@ get_default_disable_chunked_mm_input, ) from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args -from lightllm.utils.device_utils import is_sm100_gpu logger = init_logger(__name__) @@ -174,13 +174,7 @@ 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" - if args.enable_prefill_eplb: - 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" - # EPLB updates expert weights in place, but SM100 Mega-MoE caches transformed weights by tensor data_ptr. - assert not is_sm100_gpu(), "--enable_prefill_eplb does not support SM100" - assert args.run_mode == "prefill", "--enable_prefill_eplb requires --run_mode prefill" - assert args.eplb_num_redundant_experts_per_rank >= 0, "redundant expert count must be nonnegative" + args.enable_prefill_eplb = is_eplb_enabled(args) if args.enable_ep_moe: allowed_ep_prefill_att_backends = {"auto", "fa3", "triton", "flashqla"} 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 9bf7b143d0..2585a42957 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -52,9 +52,6 @@ from .multi_level_kv_cache import MultiLevelKvCacheModule from lightllm.utils.profiler import ProcessProfiler, ProfilerCmd from .ep_runtime import install_ep_runtime -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( - disable_eplb_model_init, -) class ModeBackend: @@ -354,8 +351,7 @@ def init_mtp_draft_model(self, main_kvargs: dict): model_cfg=draft_model_cfg, spec_mode=spec_mode, ) - with disable_eplb_model_init(): - draft_model = 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__}") 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 index 56051f8e44..5d460ed597 100644 --- a/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py +++ b/lightllm/server/router/model_infer/mode_backend/ep_balance_monitor.py @@ -8,7 +8,6 @@ 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.distributed.communication_op import dist_group_manager from lightllm.server.metrics.manager import MetricClient from lightllm.utils.device_utils import is_sm100_gpu from lightllm.utils.dist_utils import ( @@ -191,9 +190,7 @@ def __init__(self, model: TpPartBaseModel): self._common_round_end = torch.zeros((), dtype=torch.int64) self._monitor_error = torch.zeros((), dtype=torch.int32) - self.gloo_group = dist_group_manager.ep_balance_monitor_group - if self.gloo_group is None: - raise RuntimeError("EP balance monitor requires a pre-created dedicated Gloo process group") + 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() diff --git a/lightllm/server/router/model_infer/mode_backend/ep_runtime.py b/lightllm/server/router/model_infer/mode_backend/ep_runtime.py index 061231b2f4..bfdcb1384a 100644 --- a/lightllm/server/router/model_infer/mode_backend/ep_runtime.py +++ b/lightllm/server/router/model_infer/mode_backend/ep_runtime.py @@ -3,6 +3,7 @@ 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 @@ -10,10 +11,16 @@ def install_ep_runtime(backend): """Install on the target model only, before either inference thread starts.""" manager = None - if backend.args.enable_prefill_eplb: + 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: @@ -25,6 +32,12 @@ def install_ep_runtime(backend): 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() @@ -41,6 +54,9 @@ def run(*args, **kwargs): 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) @@ -60,4 +76,3 @@ def wait_and_poll(): return event_pack backend.overlap_event_manager.get_overlap_event_pack = get_overlap_event_pack - dist.barrier() diff --git a/lightllm/server/router/model_infer/mode_backend/eplb_manager.py b/lightllm/server/router/model_infer/mode_backend/eplb_manager.py index 5cd0a12a2b..8012c4a99f 100644 --- a/lightllm/server/router/model_infer/mode_backend/eplb_manager.py +++ b/lightllm/server/router/model_infer/mode_backend/eplb_manager.py @@ -2,7 +2,8 @@ import threading import time from pathlib import Path -from typing import Dict, Optional +from typing import Dict +from collections import deque import torch import torch.distributed as dist @@ -57,6 +58,12 @@ def __init__(self, model: TpPartBaseModel): 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 @@ -75,13 +82,6 @@ def __init__(self, model: TpPartBaseModel): 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 - # A fresh manager starts with one continuous base window. After a - # sufficient evaluation, steady state returns to the cheap sparse - # probe. An insufficient sparse probe schedules one fresh continuous - # base window before the next fixed sampling boundary. - self._continuous_collection_start_step: Optional[int] = None - self._continuous_collection_end_step: Optional[int] = self.step_interval - self._steady_collection_end_step: Optional[int] = None self._reset_recorded_samples() self._set_recording(True) # Keep background evaluation collectives separate from the main-thread @@ -106,6 +106,18 @@ def __init__(self, model: TpPartBaseModel): 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: @@ -115,32 +127,17 @@ def poll(self): self._poll_evaluation() def step(self): - if self.in_flight or self.evaluation_in_flight: - return self.prefill_steps += 1 - continuous_start = self._continuous_collection_start_step - continuous_end = self._continuous_collection_end_step - if continuous_end is not None: - # 启动或稀疏样本不足时连续采样，保证负载统计可靠。 - if continuous_start is not None and self.prefill_steps == continuous_start: - self._set_recording(True) - if self.prefill_steps >= continuous_end: - self._start_evaluation() - return - sampling_interval = self.sampling_interval - phase = self.prefill_steps % sampling_interval - steady_collection_end_step = self._steady_collection_end_step - if steady_collection_end_step is not None: - if self.prefill_steps >= steady_collection_end_step: - self._clear_steady_collection() - self._start_evaluation() - return - if sampling_interval == 1: - self._start_evaluation() + 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 - if phase == sampling_interval - self._steady_sample_window_steps(): - # 稳态仅在周期末采样少量 step，降低路由计数和评估开销。 - self._start_steady_sampling_window(self.prefill_steps + self._steady_sample_window_steps()) + 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: @@ -157,13 +154,6 @@ 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 _clear_continuous_collection(self): - self._continuous_collection_start_step = None - self._continuous_collection_end_step = None - - def _clear_steady_collection(self): - self._steady_collection_end_step = None - def _steady_sample_window_steps(self) -> int: # Preserve the original four-step sampling window when the base interval is smaller. return min( @@ -172,70 +162,34 @@ def _steady_sample_window_steps(self) -> int: self.sampling_interval, ) - def _sampling_interval_cap(self) -> int: - return max(self.step_interval * 16, getattr(self, "min_rebalance_interval", 0)) - - def _start_steady_sampling_window(self, collection_end_step: int): - """Start the fixed sparse window without moving its evaluation boundary.""" - self._reset_recorded_samples() - self._steady_collection_end_step = collection_end_step - self._set_recording(True) - - def _begin_continuous_collection(self): - minimum_end = self.prefill_steps + self.step_interval - collection_end = -(-minimum_end // self.sampling_interval) * self.sampling_interval - self._reset_recorded_samples() - self._clear_steady_collection() - self._continuous_collection_start_step = collection_end - self.step_interval - self._continuous_collection_end_step = collection_end - self._set_recording(self._continuous_collection_start_step == self.prefill_steps) - def _prepare_next_sampling_window(self): - """Clear the current window and arm the next sparse sampling window.""" - self._clear_continuous_collection() - self._clear_steady_collection() - if self.sampling_interval == 1: - self._reset_recorded_samples() - self._set_recording(True) - elif self.sampling_interval <= self._steady_sample_window_steps(): - # There is no later pre-boundary manager step at which to arm a - # full clamped window, so arm immediately but keep the same next - # fixed boundary. - self._start_steady_sampling_window(self.prefill_steps + self.sampling_interval) - else: - self._reset_recorded_samples() - self._set_recording(False) - - @staticmethod - def _recent_ring_samples(counter: torch.Tensor, recorded_sample_count: int) -> torch.Tensor: - """Return the newest ring rows in chronological order.""" - capacity = counter.shape[0] - available = min(recorded_sample_count, capacity) - if available == 0: - return counter[:0] - start = (recorded_sample_count - available) % capacity - indices = (torch.arange(available, dtype=torch.int64, device=counter.device) + start) % capacity - return counter.index_select(0, indices) + """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: - counters = [state.route_counter for state in self._eplb_states] - capacities = [counter.shape[0] for counter in counters] - if len(set(capacities)) != 1 or any(counter.ndim != 2 for counter in counters): - raise RuntimeError("EPLB sample capacities differ between layers") - counts = [state.recorded_sample_count for state in self._eplb_states] - if len(set(counts)) != 1: - raise RuntimeError("EPLB recorded sample counts differ between layers") - sample_count = counts[0] - # Validate the metadata before copying the newest rows to the CPU. - metadata = torch.tensor([sample_count, -sample_count, capacities[0], -capacities[0]], dtype=torch.int64) + 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 metadata[0] != -metadata[1] or metadata[2] != -metadata[3]: - raise RuntimeError("EPLB recorded sample count or capacity differs between ranks") - # Stack the fixed-size ring buffers in one GPU launch. Slicing each - # layer before stacking turns a single launch into one index_select per - # MoE layer and is measurably slower in the normal sparse path. - counter_samples = torch.stack(counters, dim=1) - return self._recent_ring_samples(counter_samples, sample_count).cpu() + 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] @@ -269,7 +223,17 @@ def _finish_rebalance(self): timing.get("commit_calls", 0), ) - def _poll_in_flight(self): + 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() @@ -311,14 +275,18 @@ def _poll_in_flight(self): 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: + 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 @@ -335,22 +303,27 @@ def _plan_and_broadcast(self, global_load: torch.Tensor): "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, self.current_placement, **settings) + 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(), self.current_placement.cpu(), **settings + 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, - self.current_placement, + 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 @@ -361,7 +334,7 @@ def _plan_and_broadcast(self, global_load: torch.Tensor): placement = placement.clone() for layer_index in torch.nonzero(improved, as_tuple=False).flatten().tolist(): placement[layer_index] = align_target_placement( - self.current_placement[layer_index], placement[layer_index] + current_placement[layer_index], placement[layer_index] ) select_align_ms = (time.perf_counter() - select_align_started) * 1000.0 result = { @@ -402,9 +375,9 @@ def _dump_evaluation_samples(self, global_load, recorded_sample_count, sample_wi directory.mkdir(parents=True, exist_ok=True) self._diagnostic_dump_sequence = getattr(self, "_diagnostic_dump_sequence", 0) + 1 payload = { - "global_load": global_load, - "current_placement": self.current_placement.clone(), - "prefill_steps": self.prefill_steps, + "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, } @@ -425,16 +398,12 @@ def _evaluate_after_event(self, event: torch.cuda.Event): collect_started = time.perf_counter() local_load = self._collect_local_samples() recorded_sample_count = int(local_load.shape[0]) - sample_window_steps = ( - self.step_interval - if self._continuous_collection_end_step is not None - else self._steady_sample_window_steps() - ) + 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 = 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 @@ -455,6 +424,7 @@ def _evaluate_after_event(self, event: torch.cuda.Event): 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] @@ -502,13 +472,43 @@ def _start_evaluation(self): with self._evaluation_lock: self._evaluation_result = None self._evaluation_error = None - self._set_recording(False) + 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 @@ -551,40 +551,13 @@ def _poll_evaluation(self): ], ) if result["kind"] == "insufficient": - from_continuous_window = self._continuous_collection_end_step is not None - if from_continuous_window: - self.sampling_interval = min(self.sampling_interval * 4, self._sampling_interval_cap()) - self._prepare_next_sampling_window() - else: - self._begin_continuous_collection() - if self.global_rank == 0: - if from_continuous_window: - logger.info( - "eplb insufficient samples: prefill_steps=%s minimum_layer_samples=%s required=%s " - "next_sampling_interval=%s recorded_sample_count=%s sample_window_steps=%s", - self.prefill_steps, - result["minimum_layer_samples"], - result["minimum"], - self.sampling_interval, - result.get("recorded_sample_count"), - result.get("sample_window_steps"), - ) - else: - logger.info( - "eplb insufficient samples: prefill_steps=%s minimum_layer_samples=%s required=%s " - "scheduled_fresh_window_start=%s scheduled_fresh_window_end=%s " - "recorded_sample_count=%s sample_window_steps=%s", - self.prefill_steps, - result["minimum_layer_samples"], - result["minimum"], - self._continuous_collection_start_step, - self._continuous_collection_end_step, - result.get("recorded_sample_count"), - result.get("sample_window_steps"), - ) + self._use_full_window = True + self._prepare_next_sampling_window() return False if result["kind"] == "no_improvement": - self.sampling_interval = min(self.sampling_interval * 4, self._sampling_interval_cap()) + 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 " @@ -622,8 +595,8 @@ 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._clear_continuous_collection() - self._reset_recorded_samples() + 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] diff --git a/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py b/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py index 285a04a90a..406c5c29ff 100644 --- a/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py +++ b/lightllm/server/router/model_infer/mode_backend/eplb_transfer.py @@ -97,7 +97,7 @@ def build_transfer_plan( class NixlEPLBTransfer: - """GPU-direct UCX/NIXL migration through one layer's two staging parts.""" + """GPU-direct UCX/NIXL migration through one complete staging layer.""" backend = "nixl" staging_depth = 1 @@ -228,31 +228,20 @@ def _full_phase_gate(self, error): raise RuntimeError("full EPLB migration phase failed") from error def _commit_full_parts(self, layer_index, buffer_index, changed_dst_slots): - """Commit both slices before publishing full-layout routing metadata.""" + """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) != EPLB_FULL_STAGING_PARTS: + 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) - torch.cuda.current_stream().synchronize() except BaseException as exc: error = exc - self._full_phase_gate(error) # All ranks consumed part 0 before staging reuse. - error = None - try: - self._copy_part(prepared.parts[1]) - except BaseException as exc: - error = exc - self._full_phase_gate(error) # All incoming part-1 transfers are complete. - error = None - try: - self._commit_full_part(layer_index, buffer_index, changed_dst_slots, 1) - torch.cuda.current_stream().synchronize() - self._prepared_layer = None - except BaseException as exc: - error = exc - self._full_phase_gate(error) # All live weight slices are complete before publish. + 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]): @@ -275,8 +264,7 @@ def finish(self) -> None: self._finish_transfer_generation() def __init__(self, weights, transfer_group, global_rank, world_size): - # Full layouts stage every changed row of one layer before publishing; - # this makes arbitrary overwrite cycles safe without eight full layers. + # 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 @@ -293,21 +281,6 @@ def __init__(self, weights, transfer_group, global_rank, world_size): try: self._init_ipc_metadata() self._init_push_layouts() - if self._cross_node_ranks: - 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-{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): @@ -323,6 +296,7 @@ def _local_layout(self): 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} @@ -483,7 +457,40 @@ def _prepare_part(self, layer_index, plan, part_index): 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, diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 6d4544182c..eb44a8c449 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -166,6 +166,7 @@ def get_eplb_placement_stickiness() -> float: return value +@lru_cache(maxsize=None) def get_triton_autotune_level(): return int(os.getenv("LIGHTLLM_TRITON_AUTOTUNE_LEVEL", 0)) 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/unit_tests/common/fused_moe/test_eplb.py b/unit_tests/common/fused_moe/test_eplb.py index 8cbca5bdb2..39a37891cb 100644 --- a/unit_tests/common/fused_moe/test_eplb.py +++ b/unit_tests/common/fused_moe/test_eplb.py @@ -52,12 +52,9 @@ def _replica_index_reference(token, expert, counts): from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import ( fused_moe_weight as fused_weight_module, ) -from lightllm.utils import envs_utils, dist_utils +from lightllm.utils import device_utils, envs_utils, dist_utils from lightllm.common.eplb_utils import extract_eplb_expert_tensors -from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.expert_parallel_state import ( - disable_eplb_model_init, - is_eplb_model_init_disabled, -) +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 ( @@ -193,14 +190,17 @@ def _fused_experts( 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", @@ -218,12 +218,14 @@ def _fused_experts( 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(): @@ -740,6 +742,8 @@ def test_plan_and_broadcast_publishes_canonical_placement(monkeypatch): 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() @@ -771,6 +775,8 @@ def test_plan_and_broadcast_propagates_rank_zero_error_after_existing_broadcast( 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() @@ -793,134 +799,6 @@ def record_broadcast(result_list, **_kwargs): assert broadcasted == [{"kind": "error", "message": "RuntimeError: planner boom"}] -def test_steady_state_sparse_sampling_records_steps_sixteen_to_nineteen_and_evaluates_step_twenty( - monkeypatch, -): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - counter = torch.tensor([[10, 20, 30, 40], [0, 0, 0, 0]], dtype=torch.int64) - manager.weights = [ - type( - "Weight", - (), - { - "expert_parallel_state": _test_parallel_state( - eplb=True, - route_counter=counter, - recording=False, - recorded_sample_count=1, - num_logical_experts=4, - world_size=1, - ), - }, - )() - ] - manager.in_flight = False - manager.prefill_steps = 15 - manager.step_interval = 20 - manager.sampling_interval = manager.step_interval - manager.evaluation_group = object() - manager.num_logical_experts = 4 - manager.global_rank = 1 - manager.evaluation_in_flight = False - manager._steady_collection_end_step = None - manager._continuous_collection_start_step = None - manager._continuous_collection_end_step = None - - recordings, resets, started = [], [], [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._reset_recorded_samples = lambda: resets.append(True) - monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(True)) - monkeypatch.setattr( - manager_module.torch.cuda, - "synchronize", - lambda: pytest.fail("step must not synchronize CUDA"), - ) - - manager.step() - assert manager.prefill_steps == 16 - assert recordings == [True] - assert resets == [True] - assert manager._steady_collection_end_step is not None - assert manager._steady_collection_end_step == 20 - - for _ in range(3): - manager.step() - assert manager.prefill_steps == 19 - assert recordings == [True] - assert started == [] - - manager.step() - assert manager.prefill_steps == 20 - assert manager._steady_collection_end_step is None - assert started == [True] - - -def test_steady_sampling_window_clamps_to_short_interval_without_moving_boundary(monkeypatch): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.in_flight = False - manager.evaluation_in_flight = False - manager.prefill_steps = 0 - manager.step_interval = 20 - manager.sampling_interval = 3 - manager._steady_collection_end_step = None - manager._continuous_collection_start_step = None - manager._continuous_collection_end_step = None - recordings, resets, started = [], [], [] - manager._set_recording = lambda enabled: recordings.append((manager.prefill_steps, enabled)) - manager._reset_recorded_samples = lambda: resets.append(manager.prefill_steps) - monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(manager.prefill_steps)) - - manager._prepare_next_sampling_window() - assert resets == [0] - assert recordings == [(0, True)] - assert manager._steady_collection_end_step == 3 - assert manager._steady_collection_end_step is not None - - manager.step() - manager.step() - assert started == [] - manager.step() - assert started == [3] - - -def test_eplb_step_does_not_start_a_second_evaluation_while_one_is_pending(monkeypatch): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - counter = torch.tensor([[10, 20, 30, 40], [0, 0, 0, 0]], dtype=torch.int64) - manager.weights = [ - type( - "Weight", - (), - { - "expert_parallel_state": _test_parallel_state( - eplb=True, - route_counter=counter, - recording=False, - recorded_sample_count=1, - num_logical_experts=4, - world_size=1, - ), - }, - )() - ] - manager.in_flight = False - manager.prefill_steps = 1 - manager.step_interval = 2 - manager.sampling_interval = manager.step_interval - manager.evaluation_group = object() - manager.num_logical_experts = 4 - manager.global_rank = 1 - manager.evaluation_in_flight = True - - started = [] - monkeypatch.setattr(manager, "_poll_evaluation", lambda: True) - monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(True)) - - manager.step() - - assert manager.prefill_steps == 1 - assert started == [] - - def test_evaluation_no_improvement_logs_model_fields_without_reopening_interval_window( monkeypatch, ): @@ -943,6 +821,8 @@ def join(self): 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 = [], [] @@ -950,62 +830,13 @@ def join(self): monkeypatch.setattr(manager_module.logger, "info", lambda *args: logs.append(args)) assert not manager._poll_evaluation() - assert recordings == [False] + 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 == 80 - - -def test_interval_one_rearms_after_evaluation_but_never_evaluates_empty_counter( - monkeypatch, -): - class DoneThread: - def join(self): - pass - - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.in_flight = False - manager.prefill_steps = 1 - manager.step_interval = 1 - manager.sampling_interval = 1 - manager._steady_collection_end_step = None - 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 = 1 - manager.weights = [] - manager._eplb_states = [] - recordings, starts = [], [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._start_evaluation = lambda: starts.append(True) - manager._evaluation_ready_on_all_ranks = lambda: True - - manager.poll() - - # The no-improvement backoff changes interval 1 to 4. The clamped - # steady window arms immediately but still waits for boundary step 5. - assert recordings == [True] - assert manager._steady_collection_end_step is not None - assert manager._steady_collection_end_step == 5 - assert starts == [] - assert manager.prefill_steps == 1 - - for _ in range(3): - manager.step() - assert starts == [] - manager.step() - assert starts == [True] + assert manager.sampling_interval == 20 def test_evaluation_worker_error_is_raised_by_main_thread(): @@ -1056,6 +887,9 @@ def record(self, _stream): "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 @@ -1076,89 +910,6 @@ def record(self, _stream): assert manager._poll_evaluation() # New worker has not produced a result. -def test_manager_collects_recent_ring_samples_in_chronological_order(monkeypatch): - counters = [ - torch.tensor([[10, 11], [20, 21], [30, 31]], dtype=torch.int64), - torch.tensor([[40, 41], [50, 51], [60, 61]], dtype=torch.int64), - ] - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.weights = [ - type( - "Weight", - (), - { - "expert_parallel_state": _test_parallel_state( - eplb=True, - route_counter=counter, - recording=False, - recorded_sample_count=5, - num_logical_experts=2, - world_size=1, - ), - }, - )() - for counter in counters - ] - manager._eplb_states = [weight.expert_parallel_state.eplb for weight in manager.weights] - manager.evaluation_group = object() - metadata_sizes = [] - - def all_reduce(metadata, **_kwargs): - metadata_sizes.append(metadata.numel()) - - monkeypatch.setattr(manager_module.dist, "all_reduce", all_reduce) - - samples = manager._collect_local_samples() - - assert metadata_sizes == [4] - assert torch.equal( - samples, - torch.tensor( - [ - [[30, 31], [60, 61]], - [[10, 11], [40, 41]], - [[20, 21], [50, 51]], - ], - dtype=torch.int64, - ), - ) - - -def test_manager_collects_only_two_recent_sparse_samples(monkeypatch): - counters = [ - torch.tensor([[10, 11], [20, 21], [30, 31], [40, 41]], dtype=torch.int64), - torch.tensor([[50, 51], [60, 61], [70, 71], [80, 81]], dtype=torch.int64), - ] - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.weights = [ - type( - "Weight", - (), - { - "expert_parallel_state": _test_parallel_state( - eplb=True, - route_counter=counter, - recording=False, - recorded_sample_count=2, - num_logical_experts=2, - world_size=1, - ), - }, - )() - for counter in counters - ] - manager._eplb_states = [weight.expert_parallel_state.eplb for weight in manager.weights] - manager.evaluation_group = object() - monkeypatch.setattr(manager_module.dist, "all_reduce", lambda tensor, **kwargs: None) - - samples = manager._collect_local_samples() - - assert torch.equal( - samples, - torch.tensor([[[10, 11], [50, 51]], [[20, 21], [60, 61]]], dtype=torch.int64), - ) - - @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)] @@ -1166,16 +917,18 @@ def test_manager_collects_only_two_recent_sparse_samples(monkeypatch): 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 = type( - "Args", - (), - {"enable_prefill_eplb": True, "eplb_num_redundant_experts_per_rank": 2, "run_mode": "prefill"}, - )() + 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) @@ -1189,19 +942,26 @@ def cpu_zeros(*shape, **kwargs): monkeypatch.setattr(fused_weight_module.torch, "zeros", cpu_zeros) - weight._init_expert_parallel_state() + weight._init_expert_parallel_state({}) state = weight.expert_parallel_state.eplb - assert state.route_counter.shape == (capacity, 4) + 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._start_steady_sampling_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") @@ -1224,8 +984,14 @@ def cpu_zeros(*shape, **kwargs): use_grouped_topk=False, ) expected.append([128, 128, 0, 0] if expert == 0 else [0, 0, 128, 128]) - samples = manager._recent_ring_samples(state.route_counter, state.recorded_sample_count).cpu() - assert torch.equal(samples, torch.tensor(expected, dtype=torch.int64)) + 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(): @@ -1251,20 +1017,35 @@ def test_steady_sampling_resets_fixed_ring_without_retained_history(): assert not hasattr(manager, "_sample_history") -def test_ep_without_eplb_creates_layout_without_eplb_runtime_state(monkeypatch): +@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": False, "eplb_num_redundant_experts_per_rank": 2}, + { + "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() + 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 @@ -1274,46 +1055,52 @@ def test_ep_without_eplb_creates_layout_without_eplb_runtime_state(monkeypatch): assert not hasattr(weight, "routed_expert_counter_tensor") -def test_disable_eplb_model_init_skips_eplb_state(monkeypatch): - args = type( - "Args", - (), - {"enable_prefill_eplb": True, "eplb_num_redundant_experts_per_rank": 2, "run_mode": "prefill"}, - )() +@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( - placement_module, - "build_initial_redundant_expert_ids", - lambda *args, **kwargs: pytest.fail("disabled scope must not initialize EPLB"), - ) + 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 - with disable_eplb_model_init(): - weight._init_expert_parallel_state() + weight._init_expert_parallel_state({"is_mtp_draft_model": True}) - assert weight.expert_parallel_state.eplb is None - assert weight.expert_parallel_state.num_total_physical_experts == weight.n_routed_experts + 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_disable_eplb_model_init_scope_restores_after_exception(): - assert not is_eplb_model_init_disabled() - with disable_eplb_model_init(): - assert is_eplb_model_init_disabled() - with disable_eplb_model_init(): - assert is_eplb_model_init_disabled() - assert is_eplb_model_init_disabled() - - with pytest.raises(RuntimeError): - with disable_eplb_model_init(): - raise RuntimeError - assert not is_eplb_model_init_disabled() - - def test_manager_evaluation_collective_preserves_source_node_axis(monkeypatch): manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) manager.weights = [ @@ -1343,7 +1130,7 @@ def test_manager_evaluation_collective_preserves_source_node_axis(monkeypatch): manager._evaluation_lock = threading.Lock() manager._evaluation_result = None manager._evaluation_error = None - manager._continuous_collection_end_step = None + manager._evaluation_snapshot = {"steps": 4} local = torch.full((1, 1, 4), 100, dtype=torch.int64) manager._collect_local_samples = lambda: local seen = {} @@ -1377,7 +1164,8 @@ def plan_and_broadcast(global_load): assert manager._evaluation_result["sample_window_steps"] == 4 -def test_manager_planned_evaluation_builds_improved_metadata_in_one_multilayer_call(monkeypatch): +@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( @@ -1395,6 +1183,8 @@ def test_manager_planned_evaluation_builds_improved_metadata_in_one_multilayer_c 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 @@ -1402,12 +1192,17 @@ def test_manager_planned_evaluation_builds_improved_metadata_in_one_multilayer_c 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).expand(3, -1, -1).clone() + 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._continuous_collection_end_step = None + manager._evaluation_snapshot = {"steps": 4} prepare_calls = [] prepared_batches = object() @@ -1419,9 +1214,9 @@ def prepare_transfer(layer_plans): manager._collect_local_samples = lambda: torch.full((1, 3, 4), 100, dtype=torch.int64) planned_placement = torch.tensor( [ - [[3], [0], [1], [2]], - [[2], [3], [0], [1]], - [[1], [2], [3], [0]], + [[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, ) @@ -1448,7 +1243,7 @@ def build_maps_for_layers(*args, **kwargs): manager._evaluate_after_event(type("Event", (), {"synchronize": lambda self: None})()) assert manager._evaluation_error is None - assert calls == [torch.Size([2, 4, 1])] + 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] @@ -1461,6 +1256,7 @@ def build_maps_for_layers(*args, **kwargs): 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]) @@ -1494,7 +1290,7 @@ def test_manager_preparation_error_is_saved_as_evaluation_error(monkeypatch): manager._evaluation_lock = threading.Lock() manager._evaluation_result = None manager._evaluation_error = None - manager._continuous_collection_end_step = 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", @@ -1517,7 +1313,7 @@ def fail_prepare(_layer_plans): assert str(manager._evaluation_error) == "prepare failed" -def test_decode_dispatch_keeps_logical_ids_and_uses_logical_expert_count(monkeypatch): +def test_decode_dispatch_uses_published_physical_ids_and_expert_count(monkeypatch): class Buffer: def low_latency_dispatch(self, **kwargs): calls.append(kwargs) @@ -1529,7 +1325,7 @@ def low_latency_dispatch(self, **kwargs): _set_expert_parallel_state(impl, _test_parallel_state(eplb=True)) impl._select_experts = lambda **_kwargs: ( torch.ones((1, 2)), - torch.tensor([[0, 127]], dtype=torch.int32), + torch.tensor([[0, 143]], dtype=torch.int32), torch.tensor([[0, 127]], dtype=torch.int32), ) calls = [] @@ -1552,18 +1348,33 @@ def low_latency_dispatch(self, **kwargs): "softmax", ) - assert result[2].tolist() == [[0, 127]] - assert calls[0]["num_experts"] == 128 + assert result[2].tolist() == [[0, 143]] + assert calls[0]["num_experts"] == 144 -def test_decode_select_does_not_clone_or_map_eplb_topk_ids(monkeypatch): - from lightllm.common.basemodel.triton_kernel.fused_moe import topk_select - +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)) - topk_ids = torch.tensor([[3, 127]], dtype=torch.int32) - monkeypatch.setattr(topk_select, "select_experts", lambda **_kwargs: (torch.ones((1, 2)), topk_ids)) + 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)), @@ -1577,7 +1388,10 @@ def test_decode_select_does_not_clone_or_map_eplb_topk_ids(monkeypatch): is_prefill=False, ) - assert selected.data_ptr() == origin.data_ptr() + 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): @@ -1765,7 +1579,7 @@ def fused_topk(**kwargs): assert logical_ids.tolist() == [[3, 4]] -def test_decode_masked_group_gemm_uses_primary_rows_only_when_eplb_is_enabled( +def test_decode_masked_group_gemm_uses_all_published_physical_rows( monkeypatch, ): impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) @@ -1787,11 +1601,11 @@ def masked(*args, **kwargs): )() 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] == 8 - assert captured["w13_scale"].shape[0] == captured["w2_scale"].shape[0] == 8 + 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_cached_primary_weight_packs_and_logical_experts( +def test_decode_fused_experts_uses_original_weight_packs_and_physical_experts( monkeypatch, ): impl = object.__new__(deepgemm_module.FuseMoeDeepGEMM) @@ -1833,8 +1647,8 @@ def fused(**kwargs): == "out" ) - assert [call["num_experts"] for call in captured] == [128, 128] - assert all(call["w13"].weight.shape[0] == call["w2"].weight.shape[0] == 8 for call in captured) + 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"] @@ -2211,6 +2025,56 @@ def synchronize(self): ] +@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) @@ -2257,8 +2121,7 @@ def test_manager_rearms_after_rebalance_for_interval_one(): manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) manager.step_interval = 1 manager.sampling_interval = 1 - manager._steady_collection_end_step = None - manager._continuous_collection_start_step = 0 + manager._last_evaluation_step = 1 manager.weights = [] manager._eplb_states = [] manager.target_placement = torch.zeros(1) @@ -2268,275 +2131,7 @@ def test_manager_rearms_after_rebalance_for_interval_one(): manager._finish_rebalance() assert manager.in_flight is False assert recording_calls == [True] - assert manager._steady_collection_end_step is None - assert manager._continuous_collection_start_step is None - - -def test_manager_sparse_insufficient_schedules_bounded_fresh_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": "insufficient", - "minimum_layer_samples": 1, - "minimum": 2, - } - manager.global_rank = 0 - manager.step_interval = 20 - manager.sampling_interval = 20 - manager.prefill_steps = 37 - manager.weights = [] - manager._eplb_states = [] - manager._continuous_collection_start_step = None - manager._continuous_collection_end_step = None - recordings, resets, logs = [], [], [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._reset_recorded_samples = lambda: resets.append(True) - monkeypatch.setattr(manager_module.logger, "info", lambda *args: logs.append(args)) - - assert not manager._poll_evaluation() - assert not hasattr(manager, "_retained_local_samples") - assert manager._continuous_collection_start_step == 40 - assert manager._continuous_collection_end_step == 60 - assert recordings == [False] - assert resets == [True] - assert manager.sampling_interval == 20 - assert "insufficient samples" in logs[0][0] - assert "scheduled_fresh_window" in logs[0][0] - - manager.in_flight = False - manager.evaluation_in_flight = False - starts = [] - monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(manager.prefill_steps)) - manager.step() - manager.step() - assert manager.prefill_steps == 39 - assert starts == [] - manager.step() - assert manager.prefill_steps == 40 - assert starts == [] - for _ in range(19): - manager.step() - assert manager.prefill_steps == 59 - assert starts == [] - manager.step() - assert starts == [60] - - -def test_manager_full_window_insufficient_clears_and_backs_off(): - 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": "insufficient", - "minimum_layer_samples": 1, - "minimum": 2, - } - manager.global_rank = 1 - manager.step_interval = 20 - manager.sampling_interval = 20 - manager.prefill_steps = 60 - manager.weights = [] - manager._eplb_states = [] - manager._continuous_collection_start_step = 40 - manager._continuous_collection_end_step = 60 - recordings, resets = [], [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._reset_recorded_samples = lambda: resets.append(True) - - assert not manager._poll_evaluation() - assert not hasattr(manager, "_retained_local_samples") - assert manager._continuous_collection_start_step is None - assert manager._continuous_collection_end_step is None - assert manager.sampling_interval == 80 - assert recordings == [False] - assert resets == [True] - - -def test_begin_continuous_collection_uses_full_window_at_fixed_boundary(monkeypatch): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.in_flight = False - manager.evaluation_in_flight = False - manager.prefill_steps = 36 - manager.step_interval = 20 - manager.sampling_interval = 20 - manager._steady_collection_end_step = 40 - recordings, resets, starts = [], [], [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._reset_recorded_samples = lambda: resets.append(True) - monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(True)) - - # The window is never truncated to the next boundary: it waits until 40, - # records a full 20 fresh steps, then evaluates at the boundary at 60. - manager._begin_continuous_collection() - assert manager._continuous_collection_start_step == 40 - assert manager._continuous_collection_end_step == 60 - assert recordings == [False] - assert resets == [True] - assert manager._steady_collection_end_step is None - - for _ in range(4): - manager.step() - assert manager.prefill_steps == 40 - assert recordings == [False, True] - assert starts == [] - for _ in range(19): - manager.step() - assert starts == [] - manager.step() # 60: the full window ends and triggers the evaluation. - assert starts == [True] - - -def test_begin_continuous_collection_preserves_full_window_at_sparse_boundary(): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.prefill_steps = 80 - manager.step_interval = 20 - manager.sampling_interval = 80 - manager._steady_collection_end_step = None - recordings = [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._reset_recorded_samples = lambda: None - - manager._begin_continuous_collection() - assert manager._continuous_collection_start_step == 140 - assert manager._continuous_collection_end_step == 160 - assert recordings == [False] - - -def test_first_no_improvement_switches_to_sparse_sampling_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 = 1 - manager.step_interval = 20 - manager.sampling_interval = 20 - manager._continuous_collection_start_step = 0 - manager.weights = [] - manager._eplb_states = [] - recordings = [] - manager._set_recording = lambda enabled: recordings.append(enabled) - - assert not manager._poll_evaluation() - assert manager._continuous_collection_start_step is None - assert recordings == [False] - assert manager.sampling_interval == 80 - - -def test_continuous_collection_evaluates_only_after_one_full_base_window(monkeypatch): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.in_flight = False - manager._continuous_collection_start_step = 0 - manager._continuous_collection_end_step = 20 - manager.prefill_steps = 0 - manager.step_interval = 20 - manager.sampling_interval = 320 - manager._steady_collection_end_step = None - manager.evaluation_in_flight = False - started = [] - monkeypatch.setattr(manager, "_start_evaluation", lambda: started.append(True)) - - for _ in range(19): - manager.step() - assert manager.prefill_steps == 19 - assert started == [] - - manager.step() - assert manager.prefill_steps == 20 - assert started == [True] - - -def test_no_improvement_exponentially_backs_off_sampling_interval_at_cap(): - class DoneThread: - def join(self): - pass - - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager._evaluation_lock = threading.Lock() - manager._set_recording = lambda _enabled: None - manager.global_rank = 1 - manager.step_interval = 20 - manager.sampling_interval = 20 - manager.weights = [] - manager._eplb_states = [] - - for expected_interval in (80, 320, 320): - manager.evaluation_in_flight = True - 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, - } - assert not manager._poll_evaluation() - assert manager.sampling_interval == expected_interval - - -def test_sparse_backoff_arms_and_evaluates_only_at_new_interval_boundary(monkeypatch): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.in_flight = False - manager.prefill_steps = 18 - manager.step_interval = 20 - manager.sampling_interval = 80 - manager._steady_collection_end_step = None - manager._continuous_collection_start_step = None - manager._continuous_collection_end_step = None - manager.evaluation_in_flight = False - recordings, resets, starts = [], [], [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._reset_recorded_samples = lambda: resets.append(True) - monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(True)) - - manager.step() - manager.step() - assert manager.prefill_steps == 20 - assert recordings == [] - assert starts == [] - - for _ in range(55): - manager.step() - assert manager.prefill_steps == 75 - assert recordings == [] - assert starts == [] - - manager.step() - assert manager.prefill_steps == 76 - assert recordings == [True] - assert resets == [True] - assert manager._steady_collection_end_step is not None - - for _ in range(4): - manager.step() - assert manager.prefill_steps == 80 - assert starts == [True] - assert manager._steady_collection_end_step is None + assert manager._next_evaluation_step == 2 def test_planned_rebalance_resets_sampling_interval_to_base(monkeypatch): @@ -2547,7 +2142,6 @@ def test_planned_rebalance_resets_sampling_interval_to_base(monkeypatch): manager.node_world_size = 1 manager.step_interval = 20 manager.sampling_interval = 320 - manager._continuous_collection_start_step = 0 manager.global_rank = 1 manager.transfer = type( "Transfer", @@ -2575,7 +2169,8 @@ def test_planned_rebalance_resets_sampling_interval_to_base(monkeypatch): assert manager.sampling_interval == 20 assert manager.in_flight - assert manager._continuous_collection_start_step is None + 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 @@ -2590,42 +2185,6 @@ def test_transfer_start_rejects_prepared_layers_with_wrong_count(): transfer.start([(0, []), (1, []), (2, [])], prepared_layers=[object()]) -def test_first_rebalance_completion_switches_to_four_step_sparse_window(monkeypatch): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.step_interval = 20 - manager.sampling_interval = manager.step_interval - manager._steady_collection_end_step = None - manager.weights = [] - manager._eplb_states = [] - manager.target_placement = torch.zeros(1) - manager.in_flight_started_at = 0 - manager.global_rank = 1 - recordings, starts = [], [] - manager._set_recording = lambda enabled: recordings.append(enabled) - manager._finish_rebalance() - assert recordings == [False] - - manager.in_flight = False - manager.prefill_steps = 38 - manager.evaluation_in_flight = False - monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(True)) - manager.prefill_steps = 35 - manager.step() - assert recordings == [False, True] - for _ in range(4): - manager.step() - assert starts == [True] - - -def test_manager_inflight_step_does_not_poll(): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.in_flight = True - calls = [] - manager._poll_in_flight = lambda: calls.append("poll") - manager.step() - assert calls == [] - - def test_manager_poll_advances_inflight_before_evaluation(): manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) manager.in_flight = True @@ -2723,7 +2282,7 @@ def enqueue(self, descriptor, stream): assert prepared.remote_entries == {3: [(0, [remote_inbound], staging)]} -def test_nixl_prepare_transfer_builds_both_parts_before_transfer(): +def test_nixl_prepare_transfer_builds_complete_layers_before_transfer(): transfer = object.__new__(transfer_module.NixlEPLBTransfer) transfer.staging_depth, transfer.staging = 1, ["staging"] seen = [] @@ -2735,12 +2294,12 @@ def prepare_part(layer_index, plan, part_index): transfer.global_rank = 0 transfer._prepare_part = prepare_part prepared = transfer.prepare_transfer([(3, []), (4, []), (5, [])]) - assert seen == [(3, 0), (3, 1), (4, 0), (4, 1), (5, 0), (5, 1)] + 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-3-1"), - ("prepared-4-0", "prepared-4-1"), - ("prepared-5-0", "prepared-5-1"), + ("prepared-3-0",), + ("prepared-4-0",), + ("prepared-5-0",), ] @@ -2969,19 +2528,17 @@ def remove_remote_agent(self, remote_name): transfer._get_remote_read(1, changed_entries, 0) assert agent.made == 2 - # Cache identity must include the part and both exact address ranges. - transfer._get_remote_read(1, first_entries, 1) - assert agent.made == 3 + # 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 == 4 + assert agent.made == 3 staging[0][1].pointer += 16 transfer._get_remote_read(1, first_entries, 0) - assert agent.made == 5 + assert agent.made == 4 transfer.shutdown() - assert agent.released_xfers == 5 - assert agent.released_dlists == 10 + assert agent.released_xfers == 4 + assert agent.released_dlists == 8 assert agent.removed_agents == ["remote-1"] @@ -3069,7 +2626,7 @@ def copy_part(_prepared_part): transfer._copy_part = copy_part monkeypatch.setattr(transfer_module.torch.cuda, "set_device", lambda _device: None) - prepared_batches = [transfer._PreparedLayer(0, (), (None, None))] + prepared_batches = [transfer._PreparedLayer(0, (), (None,))] transfer.start([(0, [])], prepared_batches) transfer.finish() assert agent.made == 1 @@ -3268,13 +2825,89 @@ def new_group(*args, **kwargs): 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._continuous_collection_start_step is None - assert manager._continuous_collection_end_step == manager.step_interval + 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]) @@ -3358,6 +2991,75 @@ def test_grouped_topk_eplb_matches_topk_mapping_and_counting(record_load, tokens 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]) @@ -3453,18 +3155,125 @@ def test_triton_grouped_topk_eplb_empty_tokens_skips_kernel(): assert torch.equal(counter, torch.zeros_like(counter)) -def test_eplb_configured_window_equal_interval_arms_immediately(): +def test_last_layer_commit_finishes_at_outer_boundary_and_preserves_samples(monkeypatch): manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.step_interval = 20 - manager.sampling_interval = 20 - manager.steady_sample_steps = 20 - manager._continuous_collection_start_step = None - manager._continuous_collection_end_step = None - manager._steady_collection_end_step = None - manager.prefill_steps = 0 - calls = [] - manager._reset_recorded_samples = lambda: calls.append("reset") - manager._set_recording = lambda value: calls.append(value) - manager._prepare_next_sampling_window() - assert manager._steady_collection_end_step == 20 - assert calls == ["reset", True] + 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 index 500631f530..af7b30cf61 100644 --- a/unit_tests/common/fused_moe/test_eplb_full_layout.py +++ b/unit_tests/common/fused_moe/test_eplb_full_layout.py @@ -101,6 +101,31 @@ def test_full_layout_rejects_unrepresentable_routes(layout, 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]]]) @@ -211,6 +236,8 @@ def test_manager_full_plan_is_canonical_and_publication_tracks_partial_generatio 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 = [] @@ -276,7 +303,10 @@ def test_full_checkpoint_partition_uses_absolute_layout_rows(monkeypatch): 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 - args = SimpleNamespace( + 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, @@ -284,15 +314,16 @@ def test_full_state_initialization_including_zero_replicas(monkeypatch, redundan 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() + weight._init_expert_parallel_state({}) state = weight.expert_parallel_state.eplb assert state.placement_generation == 0 - assert state.route_counter.shape == (8, 8) + 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()): diff --git a/unit_tests/common/fused_moe/test_eplb_memory_profile.py b/unit_tests/common/fused_moe/test_eplb_memory_profile.py index fa4c74d6e8..815bccef63 100644 --- a/unit_tests/common/fused_moe/test_eplb_memory_profile.py +++ b/unit_tests/common/fused_moe/test_eplb_memory_profile.py @@ -1,7 +1,7 @@ -"""EPLB reservations reach real KV profiling before any KV allocation.""" +"""EPLB buffers reduce the existing mem_fraction before KV profiling.""" +import json from types import SimpleNamespace from contextlib import nullcontext -from concurrent.futures import ThreadPoolExecutor import pytest import torch @@ -32,13 +32,24 @@ 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): @@ -49,6 +60,9 @@ class KVAllocationReached(Exception): 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) @@ -85,7 +99,7 @@ def make_expert(layer=0, redundant=0): def make_model(model_class=Qwen3MOEModel, redundant=0): model = model_class.__new__(model_class) - model.args = SimpleNamespace(enable_prefill_eplb=True) + 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 @@ -112,7 +126,7 @@ def make_model(model_class=Qwen3MOEModel, redundant=0): return model -@pytest.mark.parametrize("redundant,staging", [(0, 2304), (2, 2376)]) +@pytest.mark.parametrize("redundant,staging", [(0, 4608), (2, 4752)]) @pytest.mark.parametrize( "model_class,module,manager_class", [ @@ -154,28 +168,43 @@ def construct(*args, **kwargs): monkeypatch.setattr(module, "select_mem_manager_class", lambda: construct) monkeypatch.setattr(module, "get_added_mtp_kv_layer_num", lambda: 0, raising=False) - with memory_module.memory_reservation_scope(eplb_utils.get_eplb_memory_reservations(model)): - with pytest.raises(KVAllocationReached): - model._init_mem_manager() - - expected = {"eplb_staging": staging * 1024 ** 2, "eplb_sampling": 2 * 8 * 40 * 512 * 8 + 512} - assert manager.memory_reservations == expected 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 - sum(expected.values())) / manager.get_cell_size()) - # The original full R0 bug leaves ~1.6 GiB for a 2.25 GiB staging allocation. - assert budget.free - manager.size * manager.get_cell_size() >= sum(expected.values()) + 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", [(False, False), (True, True)]) -def test_disabled_eplb_and_draft_do_not_reserve_target_buffers(enabled, draft): +@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 - assert eplb_utils.get_eplb_memory_reservations(model) == {} + 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(): +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() @@ -187,10 +216,15 @@ def test_weight_aliases_are_deduplicated_and_non_eplb_layers_are_excluded(): assert [ weight for weight in eplb_utils.find_ep_moe_weights(model) if weight.expert_parallel_state.eplb is not None ] == [first, last] - assert eplb_utils.get_eplb_memory_reservations(model)["eplb_sampling"] == 2 * 2 * 40 * 512 * 8 + 512 + 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, 27), (2, 45)]) +@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) @@ -222,48 +256,123 @@ def test_full_staging_includes_odd_rows_and_scalar_quantization_scales(redundant FP8StaticPerTensorQuantQwen3NextMemManager, ], ) -def test_kv_backend_constructors_deduct_reservations_once(profile_without_allocating, manager_class): +def test_kv_backend_constructors_use_adjusted_mem_fraction_once(profile_without_allocating, manager_class): budget = profile_without_allocating - budget.free, budget.total = 64 * 1024 ** 2, 0 - reservations = {"eplb_staging": 1024, "eplb_sampling": 512} + 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) + 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 memory_module.memory_reservation_scope(reservations): - with pytest.raises(KVAllocationReached): - manager_class.__init__(manager, **kwargs) - assert manager.memory_reservations == reservations - assert manager.size == (budget.free - 1536) // manager.get_cell_size() + 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()) -def test_reservation_scope_restores_after_nested_failure_and_is_thread_local(profile_without_allocating): - def construct(reservations=None): - manager = memory_module.MemoryManager.__new__(memory_module.MemoryManager) - with pytest.raises(KVAllocationReached): - memory_module.MemoryManager.__init__( - manager, None, torch.bfloat16, 1, 128, 8, memory_reservations=reservations - ) - return manager - - assert construct().memory_reservations == {} - with memory_module.memory_reservation_scope({"outer": 1024}): - outer = construct() - with pytest.raises(RuntimeError, match="nested failure"): - with memory_module.memory_reservation_scope({"inner": 512}): - assert construct().memory_reservations == {"inner": 512} - raise RuntimeError("nested failure") - assert construct().memory_reservations == {"outer": 1024} - assert construct({}).memory_reservations == {} - with ThreadPoolExecutor(max_workers=1) as executor: - assert executor.submit(construct).result().memory_reservations == {} - assert outer.memory_reservations == {"outer": 1024} - assert construct().memory_reservations == {} - - -def test_base_model_scopes_overridden_kv_constructor(monkeypatch, profile_without_allocating): +@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): @@ -274,19 +383,24 @@ def _init_weights(self): self.trans_layers_weight = template.trans_layers_weight def _init_req_manager(self): - # EPLB budgeting must happen after request-state allocation. + # 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) + 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_prefill_eplb=True)) + 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, "profile_mtp_weight_memory", lambda model: nullcontext()) monkeypatch.setattr( base_module, "TorchMemorySaverWrapper", lambda enabled: SimpleNamespace(region=lambda **kwargs: nullcontext()) ) @@ -300,8 +414,8 @@ def _init_mem_manager(self): model = Model.__new__(Model) with pytest.raises(KVAllocationReached): Model.__init__(model, dict(run_mode="prefill", weight_dir="unused", max_total_token_num=None)) - expected = eplb_utils.get_eplb_memory_reservations(model) - assert model.mem_manager.memory_reservations == expected + 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 - sum(expected.values())) / model.mem_manager.get_cell_size()) - assert memory_module._memory_reservations.get() is None + 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 index 03d4e79a70..345dc0afce 100644 --- a/unit_tests/common/fused_moe/test_eplb_planner_worker.py +++ b/unit_tests/common/fused_moe/test_eplb_planner_worker.py @@ -29,10 +29,29 @@ def dead(pid): try: return open(f"/proc/{pid}/stat").read().split()[2] == "Z" - except FileNotFoundError: + 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"] @@ -47,6 +66,29 @@ def test_worker_persistent_full_redundant_and_cpu(): 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"] 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 index 4d1304c816..1b343a96dd 100644 --- a/unit_tests/common/fused_moe/test_eplb_sampling_policy.py +++ b/unit_tests/common/fused_moe/test_eplb_sampling_policy.py @@ -1,10 +1,13 @@ 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) @@ -79,43 +82,6 @@ def test_steady_window_clamps_to_configured_intervals(): assert manager._steady_sample_window_steps() == 3 -def _window_manager(step=20, interval=80, steady=20): - manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) - manager.in_flight = False - manager.evaluation_in_flight = False - manager.prefill_steps = 0 - manager.step_interval = step - manager.sampling_interval = interval - manager.steady_sample_steps = steady - manager.min_rebalance_interval = interval - manager._steady_collection_end_step = None - manager._continuous_collection_start_step = 0 - manager._continuous_collection_end_step = step - manager._steady_collection_end_step = None - manager._set_recording = lambda _enabled: None - manager._reset_recorded_samples = lambda: None - return manager - - -def test_initial_window_then_min80_sparse_window_arms_at_60_and_evaluates_at_80(monkeypatch): - manager = _window_manager() - starts = [] - monkeypatch.setattr(manager, "_start_evaluation", lambda: starts.append(manager.prefill_steps)) - for _ in range(20): - manager.step() - assert starts == [20] # Initial continuous 20-step window is unchanged. - manager._continuous_collection_start_step = None - manager._continuous_collection_end_step = None - manager.prefill_steps = 40 - for _ in range(20): - manager.step() - assert manager.prefill_steps == 60 - assert manager._steady_collection_end_step == 80 - for _ in range(20): - manager.step() - assert starts == [20, 80] - - def _no_improvement_manager(interval, min_interval): manager = manager_module.EPLBManager.__new__(manager_module.EPLBManager) manager.evaluation_in_flight = True @@ -137,29 +103,12 @@ def _no_improvement_manager(interval, min_interval): manager.prefill_steps = 0 manager.weights = [] manager._eplb_states = [] - manager._continuous_collection_start_step = None - manager._continuous_collection_end_step = None + manager._last_evaluation_step = 0 manager._set_recording = lambda _enabled: None manager._reset_recorded_samples = lambda: None return manager -def test_skip_backoff_honors_min_interval_above_default_cap(): - manager = _no_improvement_manager(80, 640) - for expected in (320, 640, 640): - assert not manager._poll_evaluation() - assert manager.sampling_interval == expected - manager.evaluation_in_flight = True - 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, - } - - 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) @@ -193,7 +142,8 @@ def test_eplb_diagnostics_dump_is_rank_zero_cpu_only(tmp_path): manager._diagnostics_dir = str(tmp_path) manager._diagnostic_dump_sequence = 0 manager.current_placement = torch.tensor([[1, 2]], dtype=torch.int64) - manager.prefill_steps = 7 + 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) == { @@ -204,6 +154,7 @@ def test_eplb_diagnostics_dump_is_rank_zero_cpu_only(tmp_path): "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 @@ -216,6 +167,7 @@ def test_eplb_diagnostics_dump_failure_disables_future_writes(monkeypatch, tmp_p 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 @@ -227,3 +179,155 @@ def test_base_one_keeps_default_four_sample_window(): 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 index c6ab9ae4db..a386ff2207 100644 --- a/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py +++ b/unit_tests/common/fused_moe/test_eplb_transfer_gpu.py @@ -1,4 +1,4 @@ -"""Complete-layout EPLB migration across two, three and eight CUDA GPUs.""" +"""Complete-layout EPLB migration and stream-ordered publication across CUDA GPUs.""" import math import os import socket @@ -130,7 +130,7 @@ def delayed_copy_part(prepared): 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(capacity * ((tensor[0].numel() + 1) // 2) * tensor.element_size() for _, tensor in transfer.live[0]) + 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] @@ -169,8 +169,10 @@ def assert_values(layer, placement): 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) @@ -186,23 +188,17 @@ def assert_old(layer): torch.cuda.current_stream().wait_stream(stream) def publish(layer): - assert_values(layer, target) - expected_old = [] - values = (current[rank].cuda() + 10 * layer).view(-1, 1) - expected_old.extend( - [ - (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 snapshot, expected in zip(snapshots[layer], expected_old): - torch.testing.assert_close(snapshot, expected) state = weights[layer].expert_parallel_state.eplb - state.logical_to_physical_map.copy_(maps[0]) - state.logical_replica_count.copy_(counts[0]) + 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 @@ -210,6 +206,17 @@ def publish(layer): 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): @@ -221,15 +228,15 @@ def publish(layer): if inject_failure: target = current.roll(1, dims=0) plan = build_transfer_plan(current, target, experts, world_size, world_size) - original_copy = transfer._copy_part + original_commit = transfer._commit_full_part - def fail_second(prepared): - if rank == 0 and prepared.part_index == 1: - raise RuntimeError("injected second part failure") - return original_copy(prepared) + def fail_commit(*args): + if rank == 0: + raise RuntimeError("injected commit enqueue failure") + return original_commit(*args) - transfer._copy_part = fail_second - plans = [(layer, plan) for layer in (0, 1, 2)] + 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"): @@ -244,15 +251,15 @@ def fail_second(prepared): dist.destroy_process_group() -@pytest.mark.parametrize("world_size", [2, 3, 8]) -@pytest.mark.parametrize("redundant", [0, 1]) +@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_second_part_failure_propagates(): +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 index 1c256975b2..5f5f5de1e4 100644 --- a/unit_tests/common/fused_moe/test_moe_weight_creation.py +++ b/unit_tests/common/fused_moe/test_moe_weight_creation.py @@ -93,7 +93,8 @@ def test_weight_creation_preserves_expert_axis_and_loading_views( @pytest.mark.parametrize("redundant", [0, 2]) -def test_eplb_one_primary_expert_fp8_constructor_and_checkpoint_load(monkeypatch, device_id, redundant): +@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( @@ -126,8 +127,16 @@ def test_eplb_one_primary_expert_fp8_constructor_and_checkpoint_load(monkeypatch 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"}, + 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) 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/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 index 293d3d6f3f..e411844730 100644 --- a/unit_tests/server/router/model_infer/test_ep_balance_monitor.py +++ b/unit_tests/server/router/model_infer/test_ep_balance_monitor.py @@ -209,7 +209,7 @@ def test_cpu_counter_accumulates_multiple_prefill_dispatches(): assert (counters.route_load, counters.compute_load) == (7, 384) -def test_monitor_reuses_manager_precreated_dedicated_gloo_group(monkeypatch): +def test_monitor_creates_dedicated_gloo_group_before_starting_thread(monkeypatch): sentinel_group = object() impl = SimpleNamespace(ep_balance_counters=None) weight = SimpleNamespace( @@ -227,12 +227,19 @@ def test_monitor_reuses_manager_precreated_dedicated_gloo_group(monkeypatch): 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()) - monkeypatch.setattr(monitor_module, "dist_group_manager", SimpleNamespace(ep_balance_monitor_group=sentinel_group)) - monkeypatch.setattr(monitor_module.dist, "new_group", lambda *args, **kwargs: pytest.fail("unexpected new_group")) + 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: None), + lambda *args, **kwargs: SimpleNamespace(start=lambda: events.append("thread")), ) monitor = monitor_module.EPBalanceMonitor(model) @@ -240,6 +247,20 @@ def test_monitor_reuses_manager_precreated_dedicated_gloo_group(monkeypatch): 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): @@ -257,8 +278,7 @@ def test_nonzero_rank_monitor_does_not_create_metric_client(monkeypatch): 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_group_manager", SimpleNamespace(ep_balance_monitor_group=sentinel_group)) - monkeypatch.setattr(monitor_module.dist, "new_group", lambda *args, **kwargs: pytest.fail("unexpected new_group")) + 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, @@ -278,8 +298,7 @@ def test_nonzero_rank_monitor_does_not_create_metric_client(monkeypatch): @pytest.mark.parametrize("disable_monitor", [False, True]) -def test_group_manager_creates_monitor_gloo_group_only_when_enabled(monkeypatch, disable_monitor): - monitor_group = object() +def test_common_group_creation_does_not_allocate_monitor_resources(monkeypatch, disable_monitor): custom_groups = [] class FakeCustomProcessGroup: @@ -309,19 +328,14 @@ def init_flashinfer_reduce(self): monkeypatch.setattr( communication_op_module.dist, "new_group", - lambda *args, **kwargs: calls.append((args, kwargs)) or monitor_group, + lambda *args, **kwargs: calls.append((args, kwargs)), ) manager = communication_op_module.DistributeGroupManager() manager.create_groups(group_size=2) assert len(manager.groups) == 2 - if disable_monitor: - assert calls == [] - assert manager.ep_balance_monitor_group is None - else: - assert calls == [((), {"ranks": [0, 1], "backend": "gloo"})] - assert manager.ep_balance_monitor_group is monitor_group + assert calls == [] def test_monitor_registers_prefill_ep_gauges_with_model_label(): diff --git a/unit_tests/server/router/model_infer/test_ep_runtime.py b/unit_tests/server/router/model_infer/test_ep_runtime.py index d17e33d950..6aefaf8144 100644 --- a/unit_tests/server/router/model_infer/test_ep_runtime.py +++ b/unit_tests/server/router/model_infer/test_ep_runtime.py @@ -10,10 +10,14 @@ 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() @@ -22,9 +26,22 @@ def runtime(monkeypatch): 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(step=lambda: events.append("step"), poll=lambda: events.append("poll")) + 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=SimpleNamespace(enable_prefill_eplb=True, disable_ep_balance_monitor=False), + 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(), ) @@ -39,9 +56,11 @@ def runtime(monkeypatch): @pytest.mark.parametrize("eplb,monitor_enabled", [(False, False), (False, True), (True, False), (True, True)]) @pytest.mark.parametrize("batch_size", [0, 3]) -def test_completed_rounds_and_decode_preserve_behavior(runtime, eplb, monitor_enabled, batch_size): +@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() @@ -53,19 +72,20 @@ def test_completed_rounds_and_decode_preserve_behavior(runtime, eplb, monitor_en 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 - assert events == ["prefill", *expected] + 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 == ["overlap", *expected] + 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: + 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 @@ -87,6 +107,47 @@ def test_disabled_monitor_and_draft_methods_are_untouched(runtime): 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 @@ -99,7 +160,7 @@ def fail(*args, **kwargs): events.clear() with pytest.raises(ValueError, match="forward failed"): getattr(backend.model, method)(object()) - assert not events + assert events == ["prepare"] def test_poll_waits_for_execution_ownership_and_propagates_failures(runtime): @@ -150,20 +211,23 @@ def test_poll_binding_does_not_retain_completed_event_packs(runtime): @pytest.mark.parametrize("backend_class", [ChunkedPrefillBackend, DPChunkedPrefillBackend]) -def test_real_infer_loops_poll_before_requests_even_when_idle(runtime, monkeypatch, backend_class): +@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: ([], []) + backend._get_classed_reqs = lambda **kwargs: ([], [], 0) get_event_pack = backend.overlap_event_manager.get_overlap_event_pack def get_ready_event_pack(): @@ -209,10 +273,28 @@ def decode(event_pack, decode_reqs): 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", @@ -232,3 +314,18 @@ 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 index b1df345834..5a03f798f7 100644 --- a/unit_tests/server/test_api_start_eplb.py +++ b/unit_tests/server/test_api_start_eplb.py @@ -1,23 +1,30 @@ 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,message", + "redundant,run_mode,nnodes,error,message", [ - (0, "normal", "requires --run_mode prefill"), - (0, "decode", "requires --run_mode prefill"), - (-1, "prefill", "nonnegative"), + (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, message): +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, @@ -29,13 +36,99 @@ def test_full_eplb_mode_rejected_before_any_process_starts(monkeypatch, redundan "set_unique_server_name", ): monkeypatch.setattr(api_start, name, lambda args: None) - monkeypatch.setattr(api_start, "is_sm100_gpu", lambda: False) 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) @@ -64,13 +157,28 @@ def test_eplb_prefill_cudagraph_is_rejected_before_starting_subprocesses(monkeyp @pytest.mark.parametrize("redundant", [0, 2]) -def test_eplb_mtp_combination_is_not_rejected_before_starting_subprocesses(monkeypatch, redundant): +@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=True, + enable_ep_moe=ep, enable_prefill_eplb=True, eplb_num_redundant_experts_per_rank=redundant, - run_mode="prefill", + run_mode=run_mode, + nnodes=nnodes, mtp_mode="vanilla_no_att", mtp_step=1, eos_id=0, @@ -84,11 +192,12 @@ def test_eplb_mtp_combination_is_not_rejected_before_starting_subprocesses(monke 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, "is_sm100_gpu", lambda: False) + 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) - monkeypatch.setattr(api_start, "set_env_start_args", lambda args: 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) @@ -100,3 +209,5 @@ def test_eplb_mtp_combination_is_not_rejected_before_starting_subprocesses(monke 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]