From 4100f8461dd79c97a7a978cf21ad58894bf6fe8e Mon Sep 17 00:00:00 2001 From: yingxudeng Date: Mon, 24 Aug 2026 15:13:13 +0800 Subject: [PATCH] feat: support EPLB for DeepSeek-V3.2 Python model executor. Add Expert Parallel Load Balancing (EPLB) support to the Python MoE path. The implementation uses slot-reuse mode where device_experts_num stays at num_local_experts to avoid exceeding the NPU fused GMM kernel groupList length limit. EPLB replaces cold expert weights in-place within fixed slots rather than allocating additional redundant slots. Key changes: - New xllm/python/layers/eplb.py with helper functions ported from C++ - PyCausalLM bridges prepare/start_transfer/update/last_ok to Python - DeepseekV3MoE: log2phy_map remap in forward, dynamic lifecycle - grouped_moe kernel: log2phy_map parameter for expert id remapping Verified EP=2 (redundant=0) and EP=4 (redundant=1) outputs match baseline (EPLB off) token-for-token on 6-layer w8a8 model. --- tests/python/test_deepseek_v32_parallel.py | 187 ++++++++++++++++++++- xllm/models/llm/py_causal_lm.cpp | 48 ++++++ xllm/models/llm/py_causal_lm.h | 8 +- xllm/python/kernels_cuda/moe.py | 1 + xllm/python/kernels_npu/moe.py | 8 +- xllm/python/layers/eplb.py | 143 ++++++++++++++++ xllm/python/models/deepseek_v32.py | 137 +++++++++++++++ 7 files changed, 528 insertions(+), 4 deletions(-) create mode 100644 xllm/python/layers/eplb.py diff --git a/tests/python/test_deepseek_v32_parallel.py b/tests/python/test_deepseek_v32_parallel.py index e649b9e8d5..4345b077e4 100644 --- a/tests/python/test_deepseek_v32_parallel.py +++ b/tests/python/test_deepseek_v32_parallel.py @@ -99,7 +99,10 @@ def _config(**overrides) -> DeepseekV3Config: "world_size": 1, } values.update(overrides) - return DeepseekV3Config.from_dict(values) + cfg = DeepseekV3Config.from_dict(values) + cfg.enable_eplb = values.get("enable_eplb", False) + cfg.redundant_experts_num = values.get("redundant_experts_num", 0) + return cfg # --------------------------------------------------------------------------- @@ -365,3 +368,185 @@ def test_dp2_eager_output_sliced_rank1(self): # dp_rank=1: offset=sum([3])=3, narrow(0, 3, 4) → [4, 64] assert result.shape[0] == 4 + + +# --------------------------------------------------------------------------- +# EPLB (Expert Parallel Load Balancing) tests +# --------------------------------------------------------------------------- + + +class TestEplbHelpers: + """Unit tests for xllm.python.layers.eplb helper functions.""" + + def test_build_initial_expert_ids_basic(self): + from xllm.python.layers.eplb import build_initial_expert_ids + + ids = build_initial_expert_ids(num_total_experts=16, ep_size=2, device_experts_num=9, redundant_experts_num=1) + assert len(ids) == 18 + assert ids[:8] == list(range(8)) + assert ids[8] == 7 + assert ids[9:17] == list(range(8, 16)) + assert ids[17] == 15 + + def test_slice_rank_expert_ids(self): + from xllm.python.layers.eplb import build_initial_expert_ids, slice_rank_expert_ids + + ids = build_initial_expert_ids(16, 2, 9, 1) + rank0 = slice_rank_expert_ids(ids, 0, 9) + rank1 = slice_rank_expert_ids(ids, 1, 9) + assert len(rank0) == 9 + assert rank0 == ids[:9] + assert rank1 == ids[9:] + + def test_build_log2phy_map_all_mapped(self): + from xllm.python.layers.eplb import build_initial_expert_ids, build_log2phy_map + + ids = build_initial_expert_ids(16, 2, 9, 1) + log2phy = build_log2phy_map(ids, 16, ep_rank=0) + assert len(log2phy) == 16 + assert all(p >= 0 for p in log2phy) + + def test_build_log2phy_map_rotation(self): + from xllm.python.layers.eplb import build_initial_expert_ids, build_log2phy_map + + ids = build_initial_expert_ids(16, 2, 9, 1) + map_r0 = build_log2phy_map(ids, 16, ep_rank=0, moe_tp_rank_in_group=0) + map_r0_tp1 = build_log2phy_map(ids, 16, ep_rank=0, moe_tp_rank_in_group=1) + assert map_r0[7] != map_r0_tp1[7] + + def test_remap_expert_ids_tensor(self): + from xllm.python.layers.eplb import remap_expert_ids + + log2phy = torch.tensor([5, 3, 1, 0, 4, 2], dtype=torch.int32) + topk_ids = torch.tensor([[0, 2], [4, 5]], dtype=torch.int32) + remapped = remap_expert_ids(topk_ids, log2phy) + expected = torch.tensor([[5, 1], [4, 2]], dtype=torch.int32) + assert torch.equal(remapped, expected) + + def test_expand_redundant_weight_storage(self): + from xllm.python.layers.eplb import expand_redundant_weight_storage + + tensor = torch.randn(10, 4, 4) + tensor[8] = torch.ones(4, 4) * 99.0 + expand_redundant_weight_storage(tensor, num_local_experts=9, device_experts_num=10) + assert torch.equal(tensor[9], tensor[8]) + + +class TestDeepseekV3MoEEplb: + """Test DeepseekV3MoE with EPLB enabled.""" + + def setup_method(self): + distributed.all_gather.reset_mock() + distributed.all_gather.side_effect = lambda x, **kw: x + distributed.all_reduce_.reset_mock() + kernels.grouped_moe.reset_mock() + + def test_eplb_moe_has_log2phy_map(self): + cfg = _config( + ep_size=2, + ep_rank=0, + moe_tp_size=1, + world_size=2, + enable_eplb=True, + redundant_experts_num=1, + ) + moe = DeepseekV3MoE(cfg, layer_id=0, dtype=torch.float32, device=torch.device("cpu")) + assert hasattr(moe, "log2phy_map") + assert moe.log2phy_map.shape[0] == 16 + # Slot-reuse: device_experts_num == num_local_experts (no extra slots) + assert moe.device_experts_num == 8 + + +# --------------------------------------------------------------------------- +# EPLB dynamic prepare/activate lifecycle tests +# --------------------------------------------------------------------------- + + +class TestEplbLifecycle: + """Test the prepare/activate dynamic EPLB bridge.""" + + def _make_causal_lm(self) -> DeepseekV3ForCausalLM: + from xllm.python.models.deepseek_v32 import DeepseekV3ForCausalLM + + config = { + "hidden_size": 64, + "n_layers": 2, + "n_heads": 4, + "head_dim": 16, + "intermediate_size": 128, + "vocab_size": 1024, + "q_lora_rank": 32, + "kv_lora_rank": 16, + "qk_nope_head_dim": 8, + "qk_rope_head_dim": 8, + "v_head_dim": 16, + "index_n_heads": 4, + "index_head_dim": 16, + "index_topk": 64, + "first_k_dense_replace": 0, + "moe_layer_freq": 1, + "n_routed_experts": 16, + "n_shared_experts": 1, + "num_experts_per_tok": 4, + "n_group": 4, + "topk_group": 2, + "routed_scaling_factor": 2.5, + "topk_method": "noaux_tc", + "norm_topk_prob": True, + "moe_intermediate_size": 32, + "tp_size": 1, + "tp_rank": 0, + "ep_size": 2, + "ep_rank": 0, + "dp_size": 1, + "dp_rank": 0, + "moe_tp_size": 1, + "moe_tp_rank": 0, + "world_size": 2, + "enable_eplb": True, + "redundant_experts_num": 1, + "device": "cpu", + } + return DeepseekV3ForCausalLM(config, build_model=True) + + def test_prepare_sets_pending_state(self): + from xllm.python.layers.eplb import build_initial_expert_ids + + model = self._make_causal_lm() + moe_layers = model._moe_layers() + assert len(moe_layers) >= 1 + + # Simulate C++ EplbManager expert_ids (length = ep_size * cpp_device_experts_num) + # cpp_device_experts_num = num_local + redundant = 8 + 1 = 9 + new_expert_ids = build_initial_expert_ids(16, 2, 9, 1) + new_expert_ids[8] = 5 # swap redundant slot to expert 5 + + model.prepare_expert_weight(0, new_expert_ids) + moe = moe_layers[0] + assert hasattr(moe, "_pending_log2phy") + assert moe._pending_log2phy.shape[0] == 16 + + def test_update_activates_new_map(self): + from xllm.python.layers.eplb import build_initial_expert_ids + + model = self._make_causal_lm() + moe = model._moe_layers()[0] + old_map = moe.log2phy_map.clone() + + # Simulate C++ moving expert 10 into rank 0's slot 7 (replacing expert 7) + # C++ expert_ids: rank 0 has 9 slots, rank 1 has 9 slots + new_expert_ids = build_initial_expert_ids(16, 2, 9, 1) + # Replace slot 7 on rank 0 with expert 10 (from rank 1) + new_expert_ids[7] = 10 + + model.prepare_expert_weight(0, new_expert_ids) + model.start_expert_weight_transfer(0) + model.update_expert_weight(0) + + # Expert 10 should now map to local slot 7 + assert moe.log2phy_map[10].item() == 7 + assert not hasattr(moe, "_pending_log2phy") + + def test_last_prepare_ok_returns_true(self): + model = self._make_causal_lm() + assert model.last_prepare_expert_weight_ok(0) is True diff --git a/xllm/models/llm/py_causal_lm.cpp b/xllm/models/llm/py_causal_lm.cpp index cedfbdc5da..5c31835089 100644 --- a/xllm/models/llm/py_causal_lm.cpp +++ b/xllm/models/llm/py_causal_lm.cpp @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "core/framework/config/eplb_config.h" #include "core/framework/config/execution_config.h" #include "core/framework/model/model_output.h" #include "core/framework/model_loader.h" @@ -87,6 +88,7 @@ PyCausalLM::PyCausalLM(const ModelContext& context) dp_size_ = (dp_group != nullptr) ? dp_group->world_size() : 1; dp_rank_ = (dp_group != nullptr) ? dp_group->rank() : 0; ep_size_ = parallel_args.ep_size(); + enable_eplb_ = ::xllm::EPLBConfig::get_instance().enable_eplb(); CHECK(parallel_args.moe_tp_group_ != nullptr); ProcessGroup* moe_tp_group = parallel_args.moe_tp_group_; @@ -208,6 +210,9 @@ py::dict PyCausalLM::build_config_dict( d["enable_graph"] = ExecutionConfig::get_instance().enable_graph(); d["python_graph_backend"] = ExecutionConfig::get_instance().python_graph_backend(); + d["enable_eplb"] = ::xllm::EPLBConfig::get_instance().enable_eplb(); + d["redundant_experts_num"] = + ::xllm::EPLBConfig::get_instance().redundant_experts_num(); return d; } @@ -258,4 +263,47 @@ bool PyCausalLM::share_weights_from(CausalLM& source) { return true; } +void PyCausalLM::prepare_expert_weight(int32_t layer_id, + const std::vector& expert_ids) { + if (!enable_eplb_) { + return; + } + py::gil_scoped_acquire gil; + if (py::hasattr(py_model_, "prepare_expert_weight")) { + py_model_.attr("prepare_expert_weight")(layer_id, expert_ids); + } +} + +void PyCausalLM::start_expert_weight_transfer(int32_t layer_id) { + if (!enable_eplb_) { + return; + } + py::gil_scoped_acquire gil; + if (py::hasattr(py_model_, "start_expert_weight_transfer")) { + py_model_.attr("start_expert_weight_transfer")(layer_id); + } +} + +void PyCausalLM::update_expert_weight(int32_t layer_id) { + if (!enable_eplb_) { + return; + } + py::gil_scoped_acquire gil; + if (py::hasattr(py_model_, "update_expert_weight")) { + py_model_.attr("update_expert_weight")(layer_id); + } +} + +bool PyCausalLM::last_prepare_expert_weight_ok(int32_t layer_id) const { + if (!enable_eplb_) { + return true; + } + py::gil_scoped_acquire gil; + if (py::hasattr(py_model_, "last_prepare_expert_weight_ok")) { + return py_model_.attr("last_prepare_expert_weight_ok")(layer_id) + .cast(); + } + return true; +} + } // namespace xllm diff --git a/xllm/models/llm/py_causal_lm.h b/xllm/models/llm/py_causal_lm.h index 87c6d23e47..f12c50ce89 100644 --- a/xllm/models/llm/py_causal_lm.h +++ b/xllm/models/llm/py_causal_lm.h @@ -73,8 +73,11 @@ class __attribute__((visibility("hidden"))) PyCausalLM : public CausalVLM { torch::Device device() const override { return device_; } const torch::TensorOptions& options() const override { return options_; } - void prepare_expert_weight(int32_t, const std::vector&) override {} - void update_expert_weight(int32_t) override {} + void prepare_expert_weight(int32_t layer_id, + const std::vector& expert_ids) override; + void start_expert_weight_transfer(int32_t layer_id) override; + void update_expert_weight(int32_t layer_id) override; + bool last_prepare_expert_weight_ok(int32_t layer_id) const override; bool share_weights_from(CausalLM& source) override; @@ -99,6 +102,7 @@ class __attribute__((visibility("hidden"))) PyCausalLM : public CausalVLM { int64_t ep_rank_ = 0; int64_t cp_size_ = 1; int64_t cp_rank_ = 0; + bool enable_eplb_ = false; ProcessGroup* tp_group_ = nullptr; pybind11::object py_model_; diff --git a/xllm/python/kernels_cuda/moe.py b/xllm/python/kernels_cuda/moe.py index 31a28bc5b6..1e30d29778 100644 --- a/xllm/python/kernels_cuda/moe.py +++ b/xllm/python/kernels_cuda/moe.py @@ -161,6 +161,7 @@ def grouped_moe( num_expert_groups: int, renormalize: bool, active_expert_range: list[int] | None = None, + log2phy_map: torch.Tensor | None = None, ) -> torch.Tensor: """Route and run grouped quantized experts as one fused operator. diff --git a/xllm/python/kernels_npu/moe.py b/xllm/python/kernels_npu/moe.py index b5e0b81f2f..7fa781c771 100644 --- a/xllm/python/kernels_npu/moe.py +++ b/xllm/python/kernels_npu/moe.py @@ -97,6 +97,7 @@ def grouped_moe( renormalize: bool, routed_scaling_factor: float, active_expert_range: list[int] | None = None, + log2phy_map: torch.Tensor | None = None, ) -> torch.Tensor: """Route and run grouped quantized experts as one fused operator. @@ -134,15 +135,18 @@ def grouped_moe( routed_scaling_factor=routed_scaling_factor, eps=1e-20, ) + if log2phy_map is not None: + topk_ids = log2phy_map[topk_ids.long()].to(torch.int32) num_tokens = hidden_states.shape[0] num_experts = gating_output.shape[1] expert_range = active_expert_range if active_expert_range is not None else [0, num_experts] + routing_num_experts = expert_range[1] if log2phy_map is not None else num_experts sorted_hidden_i8, expanded_row_idx, group_list, pertoken_scale = torch_npu.npu_moe_init_routing_v2( hidden_states, topk_ids.to(torch.int32), scale=None, active_num=num_tokens * topk, - expert_num=num_experts, + expert_num=routing_num_experts, # GMM v2 consumes cumulative expert-token offsets. expert_tokens_num_type=0, expert_tokens_num_flag=True, @@ -195,6 +199,7 @@ def _grouped_moe_fake( renormalize: bool, routed_scaling_factor: float, active_expert_range: list[int] | None = None, + log2phy_map: torch.Tensor | None = None, ) -> torch.Tensor: del ( gating_output, @@ -209,6 +214,7 @@ def _grouped_moe_fake( renormalize, routed_scaling_factor, active_expert_range, + log2phy_map, ) return torch.empty_like(hidden_states) diff --git a/xllm/python/layers/eplb.py b/xllm/python/layers/eplb.py new file mode 100644 index 0000000000..f68cbead34 --- /dev/null +++ b/xllm/python/layers/eplb.py @@ -0,0 +1,143 @@ +# Copyright 2026 The xLLM Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://github.com/xLLM-AI/xllm/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""EPLB (Expert Parallel Load Balancing) utilities. + +Ports the core logic from the C++ implementation in +``xllm/core/layers/npu_torch/deepseek_v4_eplb_utils.h``. +""" + +from __future__ import annotations + +import torch + + +def local_physical_experts_num( + num_total_experts: int, + ep_size: int, + redundant_experts_num: int, +) -> int: + """Number of physical expert slots per device (routed + redundant).""" + return num_total_experts // ep_size + redundant_experts_num + + +def build_initial_expert_ids( + num_total_experts: int, + ep_size: int, + device_experts_num: int, + redundant_experts_num: int, +) -> list[int]: + """Build the initial flat expert-id distribution across all EP ranks. + + Returns a list of length ``ep_size * device_experts_num``. Each segment + of ``device_experts_num`` corresponds to one rank. Routed slots hold + sequential expert ids; redundant slots duplicate the last routed expert + of that rank. + """ + if num_total_experts <= 0 or ep_size <= 0 or device_experts_num <= redundant_experts_num: + return [] + routed_per_rank = device_experts_num - redundant_experts_num + expert_ids: list[int] = [] + for rank in range(ep_size): + base = rank * routed_per_rank + for slot in range(routed_per_rank): + expert_ids.append(base + slot) + duplicate_expert = base + routed_per_rank - 1 + for _ in range(redundant_experts_num): + expert_ids.append(duplicate_expert) + return expert_ids + + +def slice_rank_expert_ids( + expert_ids: list[int], + ep_rank: int, + device_experts_num: int, +) -> list[int]: + """Extract the expert-id segment for a single EP rank.""" + if ep_rank < 0 or device_experts_num <= 0: + return [] + begin = ep_rank * device_experts_num + end = begin + device_experts_num + if end > len(expert_ids): + return [] + return expert_ids[begin:end] + + +def build_log2phy_map( + expert_ids: list[int], + num_total_experts: int, + ep_rank: int, + moe_tp_rank_in_group: int = 0, +) -> list[int]: + """Build logical-expert-id to physical-slot-id mapping. + + For each logical expert id (0..num_total_experts-1), resolves which + physical slot index (in the global flat expert_ids list) this consumer + should route to. When an expert has multiple duplicates, the pick is + rotated by ``(ep_rank + moe_tp_rank_in_group)`` so that ranks sharing + the same EP position but different moe_tp positions land on different + duplicates. + """ + log2phy: list[int] = [-1] * num_total_experts + if num_total_experts <= 0 or not expert_ids: + return log2phy + + rotation_base = max(ep_rank, 0) + max(moe_tp_rank_in_group, 0) + + duplicate_counts = [0] * num_total_experts + for eid in expert_ids: + if 0 <= eid < num_total_experts: + duplicate_counts[eid] += 1 + + selected_duplicates = [-1] * num_total_experts + for eid in range(num_total_experts): + cnt = duplicate_counts[eid] + if cnt > 0: + selected_duplicates[eid] = rotation_base % cnt + + duplicate_indices = [0] * num_total_experts + for physical_id, eid in enumerate(expert_ids): + if eid < 0 or eid >= num_total_experts: + continue + if duplicate_indices[eid] == selected_duplicates[eid]: + log2phy[eid] = physical_id + duplicate_indices[eid] += 1 + + return log2phy + + +def remap_expert_ids( + topk_ids: torch.Tensor, + log2phy_map: torch.Tensor, +) -> torch.Tensor: + """Remap logical expert ids to physical slot ids via gather.""" + flat = topk_ids.reshape(-1).long() + remapped = log2phy_map.index_select(0, flat).reshape(topk_ids.shape) + return remapped.to(torch.int32) + + +def expand_redundant_weight_storage( + tensor: torch.Tensor, + num_local_experts: int, + device_experts_num: int, +) -> None: + """Fill redundant slots by copying the last routed expert (in-place). + + After weight loading fills slots [0, num_local_experts), this copies the + last loaded expert into slots [num_local_experts, device_experts_num). + """ + if tensor is None or tensor.dim() == 0 or device_experts_num <= num_local_experts: + return + for slot in range(num_local_experts, device_experts_num): + tensor[slot].copy_(tensor[num_local_experts - 1]) diff --git a/xllm/python/models/deepseek_v32.py b/xllm/python/models/deepseek_v32.py index c17eee4990..1dd1e09942 100644 --- a/xllm/python/models/deepseek_v32.py +++ b/xllm/python/models/deepseek_v32.py @@ -1157,6 +1157,28 @@ def __init__( self.local_expert_start = self.ep_rank * num_local_experts self.local_expert_end = self.local_expert_start + num_local_experts + enable_eplb = getattr(cfg, "enable_eplb", False) + redundant_experts_num = getattr(cfg, "redundant_experts_num", 0) if enable_eplb else 0 + self.enable_eplb = enable_eplb + self.redundant_experts_num = redundant_experts_num + # Slot-reuse mode: device_experts_num stays at num_local_experts to + # avoid exceeding the NPU fused GMM kernel groupList length limit. + # EPLB replaces cold expert weights in-place within fixed slots. + self.device_experts_num = num_local_experts + + if enable_eplb: + from xllm.python.layers.eplb import build_initial_expert_ids, build_log2phy_map + + moe_tp_rank = getattr(cfg, "moe_tp_rank", 0) + # Slot-reuse: no redundant slots at init, each slot holds one unique expert. + initial_expert_ids = build_initial_expert_ids(self.num_experts, self.ep_size, self.device_experts_num, 0) + log2phy_list = build_log2phy_map(initial_expert_ids, self.num_experts, self.ep_rank, moe_tp_rank) + self.register_buffer( + "log2phy_map", + torch.tensor(log2phy_list, dtype=torch.int32, device=device), + persistent=False, + ) + # Match the ATB router's FP32 precision. self.gate = nn.Linear( cfg.hidden_size, @@ -1309,6 +1331,23 @@ def process_weights_after_loading(self, *, skip_expert_format: bool = False) -> def _run_routed_experts(self, hidden: torch.Tensor) -> torch.Tensor: logits = self.gate(hidden.to(torch.float32)) + if self.enable_eplb: + return kernels.grouped_moe( + hidden, + logits, + self.experts_w13, + self.experts_w2, + self.experts_w13_scale, + self.experts_w2_scale_compute, + self.e_score_correction_bias, + self.topk, + self.topk_group, + self.n_group, + self.cfg.norm_topk_prob, + self.routed_scaling, + [self.ep_rank * self.device_experts_num, (self.ep_rank + 1) * self.device_experts_num], + self.log2phy_map, + ) return kernels.grouped_moe( hidden, logits, @@ -1659,6 +1698,8 @@ def __init__(self, config: dict, build_model: bool = True) -> None: f"dp_size ({self.cfg.dp_size}) when ep_size > 1" ) self.cfg.moe_tp_size = self.cfg.moe_tp_size // self.cfg.dp_size + self.cfg.enable_eplb = bool(config.get("enable_eplb", False)) + self.cfg.redundant_experts_num = int(config.get("redundant_experts_num", 0)) if hasattr(self.cfg, "validate"): self.cfg.validate() dtype = self.resolve_dtype(config.get("dtype") or config.get("torch_dtype")) @@ -1823,3 +1864,99 @@ def load_weights( "lm_head.weight", loader.shard(loader.load_tensor("lm_head.weight"), dim=0), ) + + def _moe_layers(self) -> list[DeepseekV3MoE]: + """Return all MoE layers indexed by their position in the layer stack.""" + layers: list[DeepseekV3MoE] = [] + if self.model is None: + return layers + for layer in self.model.layers: + if isinstance(layer.mlp, DeepseekV3MoE): + layers.append(layer.mlp) + return layers + + def prepare_expert_weight(self, layer_id: int, expert_ids: list[int]) -> None: + """Prepare pending EPLB state for a single MoE layer. + + Called from C++ EplbExecutor worker thread (via GIL). + C++ passes expert_ids of length ep_size * cpp_device_experts_num where + cpp_device_experts_num = num_local_experts + redundant_experts_num. + Python uses slot-reuse mode (device_experts_num = num_local_experts), + so we slice using the C++ stride and take only the first num_local slots. + """ + from xllm.python.layers.eplb import ( + build_log2phy_map, + slice_rank_expert_ids, + ) + + moe_layers = self._moe_layers() + if layer_id < 0 or layer_id >= len(moe_layers): + return + moe = moe_layers[layer_id] + if not moe.enable_eplb: + return + + # C++ stride includes redundant slots + cpp_device_experts_num = moe.num_local_experts + moe.redundant_experts_num + pending_local_full = slice_rank_expert_ids(expert_ids, moe.ep_rank, cpp_device_experts_num) + # Take only the first num_local_experts slots (slot-reuse: no extra physical slots) + pending_local = pending_local_full[: moe.num_local_experts] + moe_tp_rank = getattr(moe.cfg, "moe_tp_rank", 0) + pending_log2phy = build_log2phy_map(expert_ids, moe.num_experts, moe.ep_rank, moe_tp_rank) + moe._pending_expert_ids = expert_ids + moe._pending_local_ids = pending_local + moe._pending_log2phy = torch.tensor(pending_log2phy, dtype=torch.int32, device=moe.log2phy_map.device) + moe._last_prepare_ok = True + + def start_expert_weight_transfer(self, layer_id: int) -> None: + """Begin weight copy for changed slots (local D2D for now).""" + moe_layers = self._moe_layers() + if layer_id < 0 or layer_id >= len(moe_layers): + return + moe = moe_layers[layer_id] + if not moe.enable_eplb or not hasattr(moe, "_pending_local_ids"): + return + if moe.experts_w13.numel() == 0: + return + active_local = getattr(moe, "_active_local_ids", None) + if active_local is None: + from xllm.python.layers.eplb import build_initial_expert_ids, slice_rank_expert_ids + + initial = build_initial_expert_ids(moe.num_experts, moe.ep_size, moe.device_experts_num, 0) + active_local = slice_rank_expert_ids(initial, moe.ep_rank, moe.device_experts_num) + pending_local = moe._pending_local_ids + for slot in range(moe.device_experts_num): + if slot < len(pending_local) and slot < len(active_local): + if pending_local[slot] != active_local[slot]: + src_expert = pending_local[slot] + src_slot = None + for s, eid in enumerate(active_local): + if eid == src_expert: + src_slot = s + break + if src_slot is not None and src_slot != slot: + moe.experts_w13.data[slot].copy_(moe.experts_w13.data[src_slot]) + moe.experts_w2.data[slot].copy_(moe.experts_w2.data[src_slot]) + moe.experts_w13_scale.data[slot].copy_(moe.experts_w13_scale.data[src_slot]) + moe.experts_w2_scale_compute.data[slot].copy_(moe.experts_w2_scale_compute.data[src_slot]) + + def update_expert_weight(self, layer_id: int) -> None: + """Atomically activate pending EPLB state.""" + moe_layers = self._moe_layers() + if layer_id < 0 or layer_id >= len(moe_layers): + return + moe = moe_layers[layer_id] + if not moe.enable_eplb or not hasattr(moe, "_pending_log2phy"): + return + moe.log2phy_map.copy_(moe._pending_log2phy) + moe._active_local_ids = moe._pending_local_ids + del moe._pending_expert_ids + del moe._pending_local_ids + del moe._pending_log2phy + + def last_prepare_expert_weight_ok(self, layer_id: int) -> bool: + moe_layers = self._moe_layers() + if layer_id < 0 or layer_id >= len(moe_layers): + return True + moe = moe_layers[layer_id] + return getattr(moe, "_last_prepare_ok", True)