# SPDX-License-Identifier: MIT AND Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation # SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project # # Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. import json import logging import random from dataclasses import dataclass from pathlib import Path import torch import torch.distributed import torch.nn.functional as F from tokenspeed.runtime.configs.model_config import ModelConfig from tokenspeed.runtime.model_loader import get_model_architecture from tokenspeed.runtime.moe import eplb_algorithms from tokenspeed.runtime.utils.server_args import ServerArgs logger = logging.getLogger(__name__) @dataclass class ExpertLocationMetadata: physical_to_logical_map: torch.Tensor # (layers, num_physical_experts) physical_to_logical_map_cpu: torch.Tensor logical_to_all_physical_map: torch.Tensor # (layers, num_logical_experts, X) logical_to_all_physical_map_num_valid: torch.Tensor # (layers, num_logical_experts) # (layers, num_logical_experts) logical_to_rank_dispatch_physical_map: torch.Tensor | None # -------------------------------- properties ------------------------------------ @property def num_layers(self) -> int: return self.physical_to_logical_map.shape[0] @property def num_physical_experts(self) -> int: return self.physical_to_logical_map.shape[1] @property def num_local_physical_experts(self) -> int: count, remainder = divmod(self.num_physical_experts, self.ep_size) if remainder != 0: raise ValueError( f"num_physical_experts={self.num_physical_experts} must be divisible by ep_size={self.ep_size}." ) return count @property def num_logical_experts(self) -> int: return self.logical_to_all_physical_map.shape[1] @property def ep_size(self): return torch.distributed.get_world_size() def __post_init__(self): num_layers_0, num_physical_experts_0 = self.physical_to_logical_map.shape num_layers_1, num_logical_experts_0, num_physical_experts_1 = ( self.logical_to_all_physical_map.shape ) num_layers_2, num_logical_experts_1 = ( self.logical_to_all_physical_map_num_valid.shape ) if not num_layers_0 == num_layers_1 == num_layers_2: raise ValueError( "Expert location maps disagree on layer count: " f"{num_layers_0}, {num_layers_1}, {num_layers_2}." ) if num_logical_experts_0 != num_logical_experts_1: raise ValueError( "Expert location maps disagree on logical expert count: " f"{num_logical_experts_0}, {num_logical_experts_1}." ) if num_physical_experts_0 != num_physical_experts_1: raise ValueError( "Expert location maps disagree on physical expert count: " f"{num_physical_experts_0}, {num_physical_experts_1}." ) # -------------------------------- construction ------------------------------------ @staticmethod def init_trivial(server_args: ServerArgs, model_config: ModelConfig): """Trivial location - logical expert i corresponds to physical expert i""" common = ExpertLocationMetadata._init_common(server_args, model_config) num_physical_experts = common["num_physical_experts"] model_config_for_expert_location = common["model_config_for_expert_location"] num_layers = model_config_for_expert_location.num_layers num_logical_experts = model_config_for_expert_location.num_logical_experts physical_to_logical_map = ( torch.arange(0, num_physical_experts).repeat(num_layers, 1) % num_logical_experts ) return ExpertLocationMetadata.init_by_mapping( server_args, model_config, physical_to_logical_map=physical_to_logical_map, ) @staticmethod def init_by_mapping( server_args: ServerArgs, model_config: ModelConfig, physical_to_logical_map, ): if not isinstance(physical_to_logical_map, torch.Tensor): physical_to_logical_map = torch.tensor(physical_to_logical_map) physical_to_logical_map = physical_to_logical_map.to(server_args.device) common = ExpertLocationMetadata._init_common(server_args, model_config) model_config_for_expert_location = common["model_config_for_expert_location"] logical_to_all_physical_map = _compute_logical_to_all_physical_map( physical_to_logical_map, num_logical_experts=model_config_for_expert_location.num_logical_experts, ) return ExpertLocationMetadata._init_raw( server_args=server_args, ep_size=common["ep_size"], physical_to_logical_map=physical_to_logical_map, logical_to_all_physical_map=logical_to_all_physical_map, ) @staticmethod def init_by_eplb( server_args: ServerArgs, model_config: ModelConfig, logical_count: torch.Tensor ): if not isinstance(logical_count, torch.Tensor): logical_count = torch.tensor(logical_count) if len(logical_count.shape) == 2: logical_count = logical_count.unsqueeze(0) logical_count = logical_count.to(server_args.device) common = ExpertLocationMetadata._init_common(server_args, model_config) model_config_for_expert_location = common["model_config_for_expert_location"] num_physical_experts = common["num_physical_experts"] num_groups = model_config_for_expert_location.num_groups num_nodes = server_args.mapping.nnodes physical_to_logical_map, logical_to_all_physical_map, expert_count = ( eplb_algorithms.rebalance_experts( tokens_per_expert=logical_count, num_physical_experts=num_physical_experts, num_local_physical_experts=num_physical_experts // common["ep_size"], num_groups=num_groups, num_nodes=num_nodes, algorithm=eplb_algorithms.compute_algorithm( raw_algorithm=server_args.eplb_algorithm, num_groups=num_groups, num_nodes=num_nodes, ), ) ) return ExpertLocationMetadata._init_raw( server_args=server_args, ep_size=common["ep_size"], physical_to_logical_map=physical_to_logical_map.to(server_args.device), logical_to_all_physical_map=logical_to_all_physical_map.to( server_args.device ), ) @staticmethod def _init_common(server_args: ServerArgs, model_config: ModelConfig): model_config_for_expert_location = ( ModelConfigForExpertLocation.from_model_config(model_config) ) num_physical_experts = ( model_config_for_expert_location.num_logical_experts + server_args.ep_num_redundant_experts ) ep_size = server_args.mapping.moe.ep_size if ep_size <= 0 or num_physical_experts % ep_size != 0: raise ValueError(f"{num_physical_experts=} {ep_size=}") num_local_physical_experts = num_physical_experts // ep_size return dict( model_config_for_expert_location=model_config_for_expert_location, num_physical_experts=num_physical_experts, num_local_physical_experts=num_local_physical_experts, ep_size=ep_size, ) @staticmethod def _init_raw( server_args: ServerArgs, ep_size: int, physical_to_logical_map: torch.Tensor, logical_to_all_physical_map: torch.Tensor, ): _, num_physical_experts = physical_to_logical_map.shape logical_to_all_physical_map_padded = F.pad( logical_to_all_physical_map, (0, num_physical_experts - logical_to_all_physical_map.shape[-1]), value=-1, ) logical_to_all_physical_map_num_valid = torch.count_nonzero( logical_to_all_physical_map != -1, dim=-1 ) return ExpertLocationMetadata( physical_to_logical_map=physical_to_logical_map, physical_to_logical_map_cpu=physical_to_logical_map.cpu(), logical_to_all_physical_map=logical_to_all_physical_map_padded, logical_to_all_physical_map_num_valid=logical_to_all_physical_map_num_valid, logical_to_rank_dispatch_physical_map=( compute_logical_to_rank_dispatch_physical_map( logical_to_all_physical_map=logical_to_all_physical_map, num_gpus=ep_size, num_physical_experts=num_physical_experts, ep_rank=torch.distributed.get_rank() % ep_size, ) if server_args.ep_dispatch_algorithm == "static" or server_args.ep_dispatch_algorithm == "static_with_zero_expert" else None ), ) # -------------------------------- mutation ------------------------------------ def update( self, other: "ExpertLocationMetadata", update_layer_ids: list[int], ): for field in [ "ep_size", ]: if getattr(self, field) != getattr(other, field): raise ValueError( f"Cannot update ExpertLocationMetadata with different {field}." ) for field in [ "physical_to_logical_map", "physical_to_logical_map_cpu", "logical_to_all_physical_map", "logical_to_all_physical_map_num_valid", "logical_to_rank_dispatch_physical_map", ]: other_field = getattr(other, field) self_field = getattr(self, field) if (other_field is not None) != (self_field is not None): raise ValueError( f"Cannot update ExpertLocationMetadata with incompatible {field}." ) if self_field is not None: mask_update = torch.tensor( [i in update_layer_ids for i in range(self.num_layers)] ) mask_update = mask_update.view(*([-1] + [1] * (self_field.dim() - 1))) mask_update = mask_update.to(self_field.device, non_blocking=True) self_field[...] = torch.where(mask_update, other_field, self_field) # -------------------------------- usage ------------------------------------ def logical_to_all_physical( self, layer_id: int, logical_expert_id: int ) -> list[int]: return [ physical_expert_id for physical_expert_id in self.logical_to_all_physical_map[ layer_id, logical_expert_id ].tolist() if physical_expert_id != -1 ] def _compute_logical_to_all_physical_map( physical_to_logical_map: torch.Tensor, num_logical_experts: int ): # This is rarely called, so we use for loops for maximum clarity num_layers, num_physical_experts = physical_to_logical_map.shape logical_to_all_physical_map = [ [[] for _ in range(num_logical_experts)] for _ in range(num_layers) ] for layer_id in range(num_layers): for physical_expert_id in range(num_physical_experts): logical_expert_id = physical_to_logical_map[ layer_id, physical_expert_id ].item() logical_to_all_physical_map[layer_id][logical_expert_id].append( physical_expert_id ) logical_to_all_physical_map = _pad_nested_array( logical_to_all_physical_map, pad_value=-1 ) return torch.tensor( logical_to_all_physical_map, device=physical_to_logical_map.device ) def _pad_nested_array(arr, pad_value): max_len = max(len(inner) for outer in arr for inner in outer) padded = [ [inner + [pad_value] * (max_len - len(inner)) for inner in outer] for outer in arr ] return padded def compute_logical_to_rank_dispatch_physical_map( logical_to_all_physical_map: torch.Tensor, num_gpus: int, num_physical_experts: int, ep_rank: int, seed: int = 42, ): r = random.Random(seed) num_local_physical_experts = num_physical_experts // num_gpus num_layers, num_logical_experts, _ = logical_to_all_physical_map.shape dtype = logical_to_all_physical_map.dtype logical_to_rank_dispatch_physical_map = torch.full( size=(num_gpus, num_layers, num_logical_experts), fill_value=-1, dtype=dtype, ) for layer_id in range(num_layers): for logical_expert_id in range(num_logical_experts): candidate_physical_expert_ids = _logical_to_all_physical_raw( logical_to_all_physical_map, layer_id, logical_expert_id ) output_partial = logical_to_rank_dispatch_physical_map[ :, layer_id, logical_expert_id ] for gpu_id in range(num_gpus): same_gpu_physical_expert_ids = [ physical_expert_id for physical_expert_id in candidate_physical_expert_ids if _compute_gpu_id_of_physical_expert( physical_expert_id, num_local_physical_experts ) == gpu_id ] if len(same_gpu_physical_expert_ids) > 0: output_partial[gpu_id] = same_gpu_physical_expert_ids[0] num_remain = torch.sum(output_partial == -1).item() output_partial[output_partial == -1] = torch.tensor( _fair_choices(candidate_physical_expert_ids, k=num_remain, r=r), dtype=dtype, ) if not torch.all(logical_to_rank_dispatch_physical_map != -1): raise RuntimeError( "logical_to_rank_dispatch_physical_map contains unassigned entries." ) device = logical_to_all_physical_map.device return logical_to_rank_dispatch_physical_map[ep_rank, :, :].to(device) def _logical_to_all_physical_raw( logical_to_all_physical_map, layer_id: int, logical_expert_id: int ) -> list[int]: return [ physical_expert_id for physical_expert_id in logical_to_all_physical_map[ layer_id, logical_expert_id ].tolist() if physical_expert_id != -1 ] def _compute_gpu_id_of_physical_expert( physical_expert_id: int, num_local_physical_experts: int ) -> int: return physical_expert_id // num_local_physical_experts def _fair_choices(arr: list, k: int, r: random.Random) -> list: quotient, remainder = divmod(k, len(arr)) choices = arr * quotient + r.sample(arr, k=remainder) r.shuffle(choices) return choices @dataclass class ModelConfigForExpertLocation: num_layers: int num_logical_experts: int num_groups: int | None = None @staticmethod def init_dummy(): return ModelConfigForExpertLocation(num_layers=1, num_logical_experts=1) @staticmethod def from_model_config(model_config: ModelConfig): model_class, _ = get_model_architecture(model_config) if hasattr(model_class, "get_model_config_for_expert_location"): return model_class.get_model_config_for_expert_location( model_config.hf_config ) else: return ModelConfigForExpertLocation.init_dummy() def compute_initial_expert_location_metadata( server_args: ServerArgs, model_config: ModelConfig ) -> ExpertLocationMetadata: data = server_args.init_expert_location if data == "trivial": return ExpertLocationMetadata.init_trivial(server_args, model_config) if data.endswith(".pt"): data_dict = torch.load(data, weights_only=True) elif data.endswith(".json"): data_dict = json.loads(Path(data).read_text()) else: data_dict = json.loads(data) if "physical_to_logical_map" in data_dict: logger.info( "init_expert_location from init_by_mapping using ServerArgs.init_expert_location" ) return ExpertLocationMetadata.init_by_mapping( server_args, model_config, **data_dict ) elif "logical_count" in data_dict: logger.info( "init_expert_location from init_by_eplb using ServerArgs.init_expert_location" ) return ExpertLocationMetadata.init_by_eplb( server_args, model_config, logical_count=data_dict["logical_count"] ) else: raise NotImplementedError( f"Unknown init_expert_location format ({list(data_dict.keys())=})" )