xc-llm-ascend/vllm_ascend/eplb/core/eplb_utils.py

#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. 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
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.
# This file is a part of the vllm-ascend project.
#
# Todo: Once https://github.com/vllm-project/vllm/issues/22246 is merged in vllm. Remove eplb utils.
import json
import os.path
import sys
from collections import defaultdict

import numpy as np
import torch
from vllm.logger import logger

import vllm_ascend.envs as envs_ascend


def expert_file_to_tensor(expert_map_path, layer_id):
    with open(expert_map_path, "r") as f:
        data = json.load(f)
    physical_count = 0
    device_data = []
    if layer_id > data["moe_layer_count"]:
        raise ValueError("Invalid EPLB Table")
    if layer_id == data["moe_layer_count"]:
        logger.warning("Init expert map of mtp/eagle when using sample.")
        return None, None
    for device in data["layer_list"][layer_id]["device_list"]:
        physical_count += len(device["device_expert"])
        device_data.append(device["device_expert"])
    global_placement = torch.tensor(device_data, dtype=torch.int32)
    return global_placement, physical_count


def generate_global_placement(n_expert, ep_size, n_redundant):
    all_experts = np.arange(n_expert)
    groups = np.array_split(all_experts, ep_size)
    for i in range(n_redundant):
        j = i % ep_size + 1
        if len(groups[-j]) == 0:
            groups[-j] = np.append(groups[-j], j)
        else:
            groups[-j] = np.append(groups[-j], (groups[-j][-1] + 1) % n_expert)
    return torch.tensor(groups, dtype=torch.int32)


def init_eplb_config(ascend_config, layer_id, moe_config):
    expert_map_path = ascend_config.expert_map_path
    n_experts = moe_config.num_experts
    ep_size = moe_config.ep_size
    global_placement = None
    eplb_enable = ascend_config.dynamic_eplb or ascend_config.expert_map_record_path
    n_redundant = ascend_config.init_redundancy_expert if eplb_enable else 0
    if expert_map_path:
        if not (os.path.exists(expert_map_path)
                and os.access(expert_map_path, os.R_OK)):
            raise ValueError("Invalid EPLB path")
        eplb_enable = True
        global_placement, physical_count = expert_file_to_tensor(
            expert_map_path, layer_id)
        if physical_count is not None:
            n_redundant = physical_count - n_experts
            if not moe_config.supports_eplb:
                raise ValueError(
                    "Eplb supports only w8a8_dynamic quantization.")
        else:
            eplb_enable = False

    if global_placement is None:
        global_placement = generate_global_placement(n_experts, ep_size,
                                                     n_redundant)

    if ep_size == 1:
        return None, None, n_redundant
    global_expert_map = []
    for rankid in range(ep_size):
        expert_map = torch.full((n_experts, ), -1, dtype=torch.int32)
        local_placement = global_placement[rankid]
        expert_map[local_placement] = torch.arange(local_placement.shape[0],
                                                   dtype=torch.int32)
        global_expert_map.append(expert_map)
    local_expert_map = global_expert_map[moe_config.ep_rank].npu()
    log2phy = generate_log2phy_map(
        global_expert_map, moe_config.ep_rank).npu() if eplb_enable else None

    return local_expert_map, log2phy, n_redundant


def generate_log2phy_map(global_expert_map, ep_rank):
    log2phy_map = defaultdict(list)
    valid_count = torch.sum(global_expert_map[0] != -1)
    for rankid, map_per_rank in enumerate(global_expert_map):
        for idx, val in enumerate(map_per_rank):
            val = val.item()
            if val != -1:
                log2phy_map[idx].append(val + rankid * valid_count)

    for key in log2phy_map.keys():
        num_of_duplications = len(log2phy_map[key])
        log2phy_map[key] = log2phy_map[key][ep_rank % num_of_duplications]

    log2phy_map = torch.scatter(
        torch.zeros(len(log2phy_map.keys()), dtype=torch.int32), 0,
        torch.tensor(list(log2phy_map.keys()), dtype=torch.int64),
        torch.tensor(list(log2phy_map.values()), dtype=torch.int32))

    return log2phy_map


class EPLBParamUtils:

    @staticmethod
    def check_iterations(iterations):
        if not isinstance(iterations, int):
            raise TypeError(f"The {iterations} is not int.")
        if iterations <= 0:
            raise ValueError(
                f"The {iterations} can not less than or equal to 0.")
        if iterations > sys.maxsize:
            raise ValueError(
                f"The {iterations} can not large than {sys.maxsize}")

    @staticmethod
    def check_dynamic_eplb(dynamic_eplb):
        if dynamic_eplb is None:
            return
        if not isinstance(dynamic_eplb, bool):
            raise TypeError("The dynamic_eplb is not bool.")

        if dynamic_eplb and envs_ascend.DYNAMIC_EPLB not in ("true", "1"):
            raise ValueError(
                'Can not enable dynamic_eplb when DYNAMIC_EPLB is not set to "true" or "1".'
            )

    @staticmethod
    def check_expert_map_path(expert_map):
        if expert_map is None:
            return
        if not isinstance(expert_map, str):
            raise TypeError("The expert_map is not str.")
        if not expert_map.strip():
            raise ValueError("The expert_map is not empty.")
        _, ext = os.path.splitext(expert_map)
        if ext.lower() != ".json":
            raise TypeError("The expert_map is not json.")
        if not os.path.exists(expert_map):
            raise ValueError("The expert_map is not exist.")
        try:
            with open(expert_map, "w", encoding='utf-8') as f:
                f.read()
        except Exception as e:
            raise IOError(
                f"Fail read expert info from {expert_map}, please check the reading permission of {expert_map} : {e}"
            )

    @staticmethod
    def check_expert_map_record_path(expert_map_record_path):
        if expert_map_record_path is None:
            return
        if not isinstance(expert_map_record_path, str):
            raise TypeError("The expert_map_record_path is not str.")
        if not expert_map_record_path.strip():
            raise ValueError("The expert_map_record_path is empty.")
        _, ext = os.path.splitext(expert_map_record_path)
        if ext.lower() != ".json":
            raise TypeError("The expert_map_record_path is not json.")
        if os.getenv("EXPERT_MAP_RECORD", "false") != "true":
            raise ValueError(
                'Can not enable expert_map_record_path when not export EXPERT_MAP_RECORD="true".'
            )
        try:
            with open(expert_map_record_path, "w", encoding='utf-8') as f:
                f.write("")
        except Exception as e:
            raise IOError(
                f"Fail write expert info to {expert_map_record_path}, please check the writing permission of {expert_map_record_path} : {e}"
            )