FlashLB algorithm

[Mercykid-bash]

Purpose

This Pull Request enhances the EPLB (Expert Parallelism Load Balancing) system by introducing a novel balancing algorithm: FlashLB.

Motivation

1. The default algorithm adopts a two-stage greedy strategy:
   a. Replica allotment: Determine the number of expert replicas by minimizing the maximum load per replica (Min Max Replica, MMR).
   b. Replica placement: Distribute replicas across devices by repeatedly assigning the heaviest replica to the least loaded device (Longest Processing Time First, LPT).

   However, this sequential process lacks inter-stage collaborative optimization, often leading to suboptimal load balancing. For example, in the simple case shown in the figure below: given 8 logical experts with hotness values of 600, 560, 120, 120, 20, 10, 10, 10, and 2 replicas allocated per device across 8 devices, the EPLB algorithm yields a maximum per-device hotness of 232, while our proposed FlashLB algorithm can reduce this value to 205.

   

2. The default algorithm relies on the averaged expert hotness over a fixed time window for optimization. While this provides a coarse approximation of the hotness distribution, it fails to capture oscillatory deviations and temporal correlations of expert hotness observed across iterations in real-world scenarios, limiting optimization quality.

3. The default algorithm periodically regenerates the expert placement table. However, it generates the table for each individual layer, and the new table does not account for correlations with the previous one; these two factors collectively lead to nearly full-scale expert reassignment.

FlashLB Algorithm Principle

1. Joint Optimization
   FlashLB achieves joint optimization of replica allotment and placement through group-based decision-making. Each group gradually determines the replica count and placement for a subset of experts, ensuring that the expected inter-device load balance (considering both deployed and pending expert replicas) is holistically optimized. To attain superior load balancing, FlashLB employs tree search to expand the solution space while integrating pruning and precompilation techniques for acceleration, thereby delivering load balancing that is both high-quality and practically efficient.

2. Multi-Shot Enhancement
   FlashLB partitions each profiling interval (e.g., 1024 iterations) into consecutive smaller sub-intervals (e.g., 16 iterations), each capturing independent hotness measurements. It then performs multi-shot optimization to co-optimize these sub-intervals simultaneously—enabling adaptation to time-variant expert hotness while enhancing robustness.

3. Incremental Adjustment
   To reduce the overhead of frequent expert re-deployment, FlashLB introduces an incremental adjustment scheme operating at both inter-layer and intra-layer levels:
   a. Inter-Layer: Hotness variations are tracked at the layer level. Only layers with fluctuations exceeding a predefined threshold trigger re-computation of expert placement, avoiding unnecessary redeployment for stable layers；
   b. Intra-Layer (Optional): A lightweight incremental LPT algorithm (LPT-Incremental) is applied. Instead of recomputing full placement for all experts in a layer, it selectively adjusts only the hottest experts or those with replica count changes, further reducing migration overhead.

   This incremental strategy significantly reduces adjustment costs while maintaining balanced performance across layers and devices.

Co-author:

Co-authored-by: Skywalker-EP 173723846@qq.com

• vLLM version: v0.10.2
• vLLM main: vllm-project/vllm@9607d5e

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

[gemini-code-assist]

Code Review

This pull request introduces FlashLB, a new and sophisticated load balancing algorithm. The implementation leverages Numba for performance, which is a good choice for the complex numerical computations involved. However, I've identified several critical issues in the implementation that could lead to incorrect behavior, crashes, or instability. These include the use of shared class state, improper error handling that terminates the process, multiple potential division-by-zero errors, and unhandled edge cases that can cause crashes. I have provided detailed comments and suggestions to address these critical problems. Additionally, there's a high-severity issue with a module-level warmup function that causes a significant side effect on import. Addressing these points is crucial for the stability and correctness of the new policy.

[gemini-code-assist]

Calling warm_up() at the module level is a significant side effect that occurs upon import. This can make module loading slow and unpredictable. This function, which seems to be for JIT compilation warmup, should be called explicitly by the application that initializes the load balancing system, rather than being a side effect of importing this module.

Suggested change

	warm_up()
	# warm_up()

[Mercykid-bash]

The warm-up call at the module level is intentionally designed to trigger JIT compilation during import rather than at inference time. This approach ensures that numba's just-in-time compilation overhead is incurred upfront during module loading, preventing unexpected latency spikes during actual inference operations.

While module imports may take slightly longer, this trade-off is necessary to guarantee consistent performance during runtime, where low and predictable latency is critical for the load balancing system's effectiveness. Explicitly deferring this warm-up to application initialization could lead to missed compilation steps and subsequent runtime delays if the initialization sequence is not strictly followed.

The current implementation prioritizes inference-time performance predictability over import speed, which aligns with the requirements of latency-sensitive load balancing operations.

[wangxiyuan]

Thanks for the contribution. Only one question: warm_up should be called when import the module? is there any way to let it only runs when the policy is set?

[Mercykid-bash]

Another approach is to incorporate a check within policy_factory.py that triggers the warm_up operation exclusively when the RealTimeLB policy is selected.

This way, the warm_up process will only run if and when the RealTimeLB policy is actually instantiated (i.e., when this specific policy is set), avoiding unnecessary execution for other policy types.

Would this align with your intended implementation? Please let me know if you need further adjustments.

class PolicyFactory:

@staticmethod
def generate_policy(policy_type: int, config: DynamicConfig) -> EplbPolicy:
    policy = {
        # Constraint applying Dynamic EPLB policy V2:
        # If there exists redundant expert:
        # only one redundant expert can be placed in one NPU and its physical expert index must be 0

        # Applying greedy d2d expert weight update composing
        0:
        RandomLoadBalance,  # RandomLoadBalance: shuffle last physical expert on NPU 1 and 3
        1:
        DynamicEplb,  # Dynamic EPLB policy: overall expert replacement based on current moe load
        2:
        DynamicEplbV2,  # Dynamic EPLB policy V2:  expert replacement with constrained number of expert shuffle
        3:
        FlashLB,  # FlashLB EPLB policy: expert replacement based on Joint Optimization, Multi-Shot Enhancement and Incremental Adjustment
    }
    policy_class = policy.get(policy_type, RandomLoadBalance)
    policy_instance = policy_class(config)
    if policy_type == 3:
        policy_instance.warm_up()
    return policy_instance