From 60dd916e4f749b96eab278894e85f1c7576f3541 Mon Sep 17 00:00:00 2001 From: Yun Yang Date: Tue, 29 Sep 2026 15:31:15 +0800 Subject: [PATCH] [EPLB] Add the MoE Load Balancer policy: pluggable expert placement and replica routing `--eplb-config '{"policy": "mlb", ...}'` delegates EPLB's expert placement and, optionally, per-token replica routing to the MoE Load Balancer library (moe_load_balancer): DeepSeek EPLB hierarchical/global placement, LPLB and Waterfill routing, and UltraEP's online placement/quota refresh with its own expert-weight mover. vLLM keeps the mechanism, the library keeps the policies; nothing changes unless the policy is selected and the package is imported lazily. Adds: - vllm/distributed/eplb/policy/mlb.py: the `mlb` EPLB placement policy. - vllm/distributed/eplb/mlb_runtime.py: the routing runtime. Replica shares or resolved physical ids from a policy go through the existing fused EPLB map-and-record kernel, so load recording, the record gate and the padding mask are unchanged; UltraEP's refresh re-solves every N representative batches and moves replica weights pairwise over the EPLB communicator. - --eplb-config keys l2_algorithm, init_placement_path, save_placement_path. - VLLM_FUSE_SHARED_EXPERTS: dispatch the shared expert through the EP all-to-all as one slot per rank so a policy can place it (Waterfill). - A DP-synced uniform_decode_across_dp flag on the forward context, and the replay of a policy's per-step collective on the dummy-run path, so a policy that issues a collective inside the forward reaches the same verdict on every rank. - rank_imbalance next to balancedness in the EPLB log. - DeepEP: VLLM_DEEPEP_RDMA_BUFFER_SIZE_MB and multi-node NVLink knobs. - Docs section and unit tests (skip without the package). Co-Authored-By: Claude Fable 5.1 Signed-off-by: Yun Yang --- .buildkite/test_areas/expert_parallelism.yaml | 13 + docs/serving/expert_parallel_deployment.md | 49 +- tests/distributed/test_eplb_execute.py | 68 + tests/distributed/test_eplb_mlb_policy.py | 181 ++ tests/distributed/test_eplb_mlb_routing.py | 895 ++++++++ .../test_eplb_replica_shares_kernel.py | 204 ++ tests/kernels/moe/test_routing.py | 185 ++ vllm/config/parallel.py | 75 +- .../device_communicators/all2all.py | 35 +- vllm/distributed/eplb/eplb_state.py | 390 +++- vllm/distributed/eplb/mlb_runtime.py | 1993 +++++++++++++++++ vllm/distributed/eplb/policy/__init__.py | 7 +- vllm/distributed/eplb/policy/mlb.py | 134 ++ vllm/distributed/eplb/rebalance_execute.py | 25 + vllm/envs.py | 108 + vllm/forward_context.py | 16 +- .../model_executor/layers/fused_moe/config.py | 8 + vllm/model_executor/layers/fused_moe/layer.py | 82 +- .../layers/fused_moe/routed_experts.py | 10 +- .../layers/fused_moe/router/base_router.py | 239 +- .../fused_moe/router/fused_moe_router.py | 10 + .../layers/fused_moe/shared_expert_fusion.py | 262 +++ vllm/model_executor/models/deepseek_v2.py | 119 +- vllm/models/deepseek_v32/nvidia/model.py | 42 +- vllm/models/deepseek_v4/nvidia/model.py | 2 + vllm/models/kimi_k3/nvidia/model.py | 2 + vllm/v1/spec_decode/extract_hidden_states.py | 2 +- vllm/v1/spec_decode/llm_base_proposer.py | 2 +- vllm/v1/worker/dp_utils.py | 83 +- vllm/v1/worker/gpu_model_runner.py | 111 +- 30 files changed, 5230 insertions(+), 122 deletions(-) create mode 100644 tests/distributed/test_eplb_mlb_policy.py create mode 100644 tests/distributed/test_eplb_mlb_routing.py create mode 100644 tests/distributed/test_eplb_replica_shares_kernel.py create mode 100644 vllm/distributed/eplb/mlb_runtime.py create mode 100644 vllm/distributed/eplb/policy/mlb.py create mode 100644 vllm/model_executor/layers/fused_moe/shared_expert_fusion.py diff --git a/.buildkite/test_areas/expert_parallelism.yaml b/.buildkite/test_areas/expert_parallelism.yaml index 930a9a448a16..810cb87c1c63 100644 --- a/.buildkite/test_areas/expert_parallelism.yaml +++ b/.buildkite/test_areas/expert_parallelism.yaml @@ -9,11 +9,20 @@ steps: working_dir: "/vllm-workspace/tests" source_file_dependencies: - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/router/base_router.py - tests/distributed/test_eplb_algo.py - tests/distributed/test_eplb_utils.py + - tests/distributed/test_eplb_mlb_policy.py + - tests/distributed/test_eplb_mlb_routing.py + - tests/distributed/test_eplb_replica_shares_kernel.py commands: - pytest -v -s distributed/test_eplb_algo.py - pytest -v -s distributed/test_eplb_utils.py + # The MoE Load Balancer tests skip themselves when the optional + # moe_load_balancer package is not installed. + - pytest -v -s distributed/test_eplb_mlb_policy.py + - pytest -v -s distributed/test_eplb_mlb_routing.py + - pytest -v -s distributed/test_eplb_replica_shares_kernel.py mirror: amd: label: ":amd: (MI355 DPX) EPLB Algorithm" @@ -24,8 +33,12 @@ steps: - image-build-amd source_file_dependencies: - vllm/distributed/eplb + - vllm/model_executor/layers/fused_moe/router/base_router.py - tests/distributed/test_eplb_algo.py - tests/distributed/test_eplb_utils.py + - tests/distributed/test_eplb_mlb_policy.py + - tests/distributed/test_eplb_mlb_routing.py + - tests/distributed/test_eplb_replica_shares_kernel.py - vllm/platforms/rocm.py - label: ":nvidia: (L4) EPLB Execution" diff --git a/docs/serving/expert_parallel_deployment.md b/docs/serving/expert_parallel_deployment.md index ea602336455a..e7929cb3eefd 100644 --- a/docs/serving/expert_parallel_deployment.md +++ b/docs/serving/expert_parallel_deployment.md @@ -154,7 +154,10 @@ Configure EPLB with the `--eplb-config` argument, which accepts a JSON string. T | `log_balancedness` | Log balancedness metrics (avg tokens per expert ÷ max tokens per expert) | `false` | | `num_redundant_experts` | Additional global experts per EP rank beyond equal distribution | `0` | | `use_async` | Use non-blocking EPLB for reduced latency overhead | `true` | -| `policy` | The policy type for expert parallel load balancing | `"default"` | +| `policy` | The policy type for expert parallel load balancing: `"default"` (built-in) or `"mlb"` (MoE Load Balancer, see below) | `"default"` | +| `l2_algorithm` | MoE Load Balancer replica-routing expression (only with `policy: "mlb"`); empty keeps vLLM's built-in replica choice | `""` | +| `init_placement_path` | Load a saved expert placement at start-up instead of the trivial layout | `null` | +| `save_placement_path` | Save the expert placement to this path after every rearrangement | `null` | | `communicator` | Backend for expert weight transfers: `"torch_nccl"`, `"torch_gloo"`, `"pynccl"`, `"nixl"`, or `null` (auto) | `null` | For example: @@ -176,6 +179,50 @@ vllm serve Qwen/Qwen3-30B-A3B \ --eplb-config.log_balancedness true ``` +### MoE Load Balancer policy + +`policy: "mlb"` delegates expert placement, and optionally per-token replica routing, to the +[MoE Load Balancer](https://github.com/xutizhou/moe_load_balancer) (`moe_load_balancer`), a +framework-neutral library of expert-parallel load-balancing policies. vLLM imports it lazily, so +nothing changes unless the policy is selected; install it into the serving environment first. + +Two stages are configurable: + +- **L1, placement** (`VLLM_MLB_L1_ALGORITHM`): which logical expert lives in which physical slot, + decided at every rearrangement. `auto` (default) selects the balancer's default algorithm for the + model's expert-group layout; `ultraep` re-solves placement and per-replica quotas online every + `MLB_ULTRAEP_REFRESH_INTERVAL` representative batches and moves the affected expert weights itself. +- **L2, replica routing** (`--eplb-config.l2_algorithm` or `VLLM_MLB_L2_ALGORITHM`): which replica + of a logical expert serves each token. Expressions combine a routed-expert policy (`static`, + `lplb`, `ultraep`) with an optional shared-expert policy (`+waterfill`), e.g. `lplb+waterfill`. + `waterfill` dispatches the shared expert through the expert-parallel all-to-all so that it can be + assigned to the least-loaded rank; this widens the expert space by one slot per rank and can be + enabled on its own with `VLLM_FUSE_SHARED_EXPERTS=1`. L2 routing is skipped automatically when + there are no redundant experts, where it could only reproduce the built-in choice. + +```bash +# Placement fitted by the balancer, online UltraEP refresh, Waterfill for the shared expert +VLLM_MLB_L1_ALGORITHM=ultraep MLB_ULTRAEP_REFRESH_INTERVAL=8 \ +vllm serve deepseek-ai/DeepSeek-V3 \ + --tensor-parallel-size 1 --data-parallel-size 8 --enable-expert-parallel \ + --enable-eplb \ + --eplb-config '{"policy":"mlb","use_async":false,"num_redundant_experts":16,"l2_algorithm":"ultraep+waterfill","log_balancedness":true}' +``` + +`use_async` must be `false` with this policy: the balancer plans synchronously on the caller's +thread. To fit a placement once and serve with it held, run a short profiling serve with +`"save_placement_path": "/path/placement.pt"` and start production with +`"init_placement_path": "/path/placement.pt"` and a `step_interval` that is never reached. + +Tuning knobs for the online UltraEP refresh (all read through `vllm.envs`): + +| Variable | Default | Meaning | +| -------- | ------- | ------- | +| `MLB_ULTRAEP_REFRESH_INTERVAL` | `64` | Re-solve every N representative batches; `8` is a good production value, `1` maximizes balance at the cost of one refresh per step | +| `MLB_ULTRAEP_REFRESH_MIN_TOKENS` | `512` | A batch counts as representative when it routes at least this many tokens | +| `MLB_ULTRAEP_MOVER` | `direct` | How re-solved replica weights move between ranks: `direct` (send/recv over the EPLB communicator), `symm` (torch symmetric memory; one host or a multi-node NVLink fabric), `auto` (symm when every rank agrees it can, otherwise direct) | +| `MLB_ULTRAEP_LAGGED_APPLY` | `0` | `1` lands a plan on the layer's next forward, avoiding a per-layer device sync at the price of one step of staleness; worth about +3% at interval 1 on Blackwell-class step times, nothing at interval 8 | + ### Expert Distribution Formula - **Default**: Each EP rank has `NUM_TOTAL_EXPERTS ÷ NUM_EP_RANKS` experts diff --git a/tests/distributed/test_eplb_execute.py b/tests/distributed/test_eplb_execute.py index 2037535e3ddf..6c5f293f3d7b 100644 --- a/tests/distributed/test_eplb_execute.py +++ b/tests/distributed/test_eplb_execute.py @@ -3,6 +3,7 @@ import random +import numpy as np import pytest import torch import torch.distributed @@ -15,6 +16,8 @@ has_nixl, ) from vllm.distributed.eplb.rebalance_execute import ( + experts_being_acquired, + get_ep_ranks_with_experts_batch, move_from_buffer, rearrange_expert_weights_inplace, transfer_layer, @@ -849,3 +852,68 @@ def test_rearrange_expert_weights_profile_mode(world_size): _test_rearrange_expert_weights_profile_mode, world_size, ) + + +@pytest.mark.parametrize("seed", range(20)) +def test_narrowed_send_query_drops_only_experts_with_nowhere_to_send(seed: int): + """Experts move_to_buffer stops asking about could not have sent anything. + + move_to_buffer asks get_ep_ranks_with_experts_batch only about experts that + occupy a position where the new placement differs from the old, rather than + about every expert resident on the rank. That is sound only because the + send loop skips any expert whose ranks_to_recv is empty, and an expert + acquired by no rank has exactly that. The narrowing and the skip live in + different functions with nothing tying them together, so assert the + implication against the real mapping instead of trusting the argument for + it -- a change to how recv ranks are derived would otherwise turn this into + silently dropped transfers. + """ + rng = np.random.default_rng(seed) + ep_size, num_local = 8, 34 + num_physical = ep_size * num_local + num_logical = 256 + + old_indices = np.concatenate( + [ + np.arange(num_logical), + rng.choice(num_logical, num_physical - num_logical, replace=False), + ] + ) + rng.shuffle(old_indices) + new_indices = old_indices.copy() + moved = rng.choice(num_physical, int(rng.integers(1, 24)), replace=False) + new_indices[moved] = rng.choice(num_logical, moved.size, replace=False) + + arriving = experts_being_acquired(old_indices, new_indices) + + dropped_any = False + kept_with_work = False + for ep_rank in range(ep_size): + local = slice(ep_rank * num_local, (ep_rank + 1) * num_local) + resident = np.unique(old_indices[local]) + + dropped = resident[~np.isin(resident, arriving)] + if dropped.size: + dropped_any = True + _, recv_map = get_ep_ranks_with_experts_batch( + dropped, num_local, old_indices, new_indices + ) + for expert in dropped.tolist(): + assert not recv_map[int(expert)], ( + f"rank {ep_rank} expert {expert} was dropped by the " + f"narrowing but has receivers {recv_map[int(expert)]}" + ) + + kept = resident[np.isin(resident, arriving)] + if kept.size: + _, recv_map = get_ep_ranks_with_experts_batch( + kept, num_local, old_indices, new_indices + ) + kept_with_work |= any(recv_map[int(e)] for e in kept.tolist()) + + # Both guard against a vacuous pass: nothing dropped means the narrowing + # was never exercised, and nothing kept with receivers means the placement + # asked for no transfers at all, in which case an empty recv_map says + # nothing about whether the narrowing is selective. + assert dropped_any, "narrowing dropped nothing -- the case is not exercised" + assert kept_with_work, "no expert had receivers -- recv_map is empty for all" diff --git a/tests/distributed/test_eplb_mlb_policy.py b/tests/distributed/test_eplb_mlb_policy.py new file mode 100644 index 000000000000..e612c3c3f073 --- /dev/null +++ b/tests/distributed/test_eplb_mlb_policy.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the MLB-backed EPLB placement policy. + +These run on CPU and need no distributed group: the placement contract is a +pure function of the load statistics and the topology. + +They are skipped unless the ``moe-load-balancer`` package is importable. +""" + +import numpy as np +import pytest +import torch + +from vllm.distributed.eplb.policy import EPLB_POLICIES, DefaultEplbPolicy + +pytest.importorskip("moe_load_balancer") + +from vllm.distributed.eplb.policy.mlb import MlbEplbPolicy # noqa: E402 + +# DeepSeek-R1 geometry: 58 MoE layers, 256 routed experts, EP=8, 32 redundant. +NUM_LAYERS = 58 +NUM_LOGICAL = 256 +NUM_RANKS = 8 +NUM_REPLICAS = NUM_LOGICAL + 32 +NUM_GROUPS = 8 + + +def _skewed_load(seed: int = 0) -> torch.Tensor: + """A realistic per-layer expert load: a hot head plus a long tail.""" + rng = np.random.default_rng(seed) + base = rng.pareto(1.6, size=(NUM_LAYERS, NUM_LOGICAL)) + 1.0 + # Shuffle which experts are hot per layer, as real routing does. + for layer in range(NUM_LAYERS): + rng.shuffle(base[layer]) + return torch.from_numpy((base * 1000).astype(np.float32)) + + +def _rank_loads(phy2log: torch.Tensor, weight: torch.Tensor) -> np.ndarray: + """Per-layer, per-rank load if each replica takes an equal share.""" + phy = phy2log.numpy() + w = weight.numpy() + replicas = np.zeros_like(w) + np.add.at( + replicas, + (np.arange(phy.shape[0])[:, None].repeat(phy.shape[1], 1), phy), + 1, + ) + per_replica = w / np.maximum(replicas, 1) + slots_per_rank = phy.shape[1] // NUM_RANKS + loads = np.take_along_axis(per_replica, phy, axis=1) + return loads.reshape(phy.shape[0], NUM_RANKS, slots_per_rank).sum(axis=2) + + +def _imbalance(phy2log: torch.Tensor, weight: torch.Tensor) -> float: + loads = _rank_loads(phy2log, weight) + return float(np.mean(loads.max(axis=1) / np.maximum(loads.mean(axis=1), 1e-9))) + + +def _assert_valid_placement(phy2log: torch.Tensor) -> None: + assert phy2log.shape == (NUM_LAYERS, NUM_REPLICAS) + assert phy2log.dtype == torch.int64 + assert phy2log.device.type == "cpu" + arr = phy2log.numpy() + assert arr.min() >= 0 and arr.max() < NUM_LOGICAL + # Every logical expert must own at least one physical slot in every layer, + # otherwise its tokens have nowhere to go. + for layer in range(NUM_LAYERS): + assert len(np.unique(arr[layer])) == NUM_LOGICAL + + +def test_mlb_policy_is_registered(): + assert EPLB_POLICIES["mlb"] is MlbEplbPolicy + + +def test_mlb_placement_is_valid(): + weight = _skewed_load() + phy2log = MlbEplbPolicy.rebalance_experts( + weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS + ) + _assert_valid_placement(phy2log) + + +def test_mlb_placement_is_deterministic(): + weight = _skewed_load() + a = MlbEplbPolicy.rebalance_experts(weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS) + b = MlbEplbPolicy.rebalance_experts(weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS) + assert torch.equal(a, b), "placement must be reproducible across EP ranks" + + +def test_mlb_balances_at_least_as_well_as_default(): + """MLB's DeepSeek EPLB and vLLM's are ports of the same algorithm, so the + balance quality should be equivalent -- not merely 'not much worse'.""" + weight = _skewed_load() + mlb = MlbEplbPolicy.rebalance_experts( + weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS + ) + default = DefaultEplbPolicy.rebalance_experts( + weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS + ) + mlb_imb = _imbalance(mlb, weight) + default_imb = _imbalance(default, weight) + assert mlb_imb <= default_imb * 1.02, ( + f"MLB imbalance {mlb_imb:.4f} materially worse than " + f"vLLM default {default_imb:.4f}" + ) + + +def test_slot_preservation_reduces_weight_movement(): + """The built-in policy preserves intra-GPU slots to avoid weight copies. + MLB's L1 policies do not consume a previous placement, so the glue applies + the same pass; without it, a re-plan churns far more slots.""" + weight = _skewed_load(seed=0) + first = MlbEplbPolicy.rebalance_experts( + weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS + ) + + drifted = _skewed_load(seed=1) * 0.25 + weight * 0.75 + without = MlbEplbPolicy.rebalance_experts( + drifted, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS + ) + with_preserve = MlbEplbPolicy.rebalance_experts( + drifted, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS, old_global_expert_indices=first + ) + + _assert_valid_placement(with_preserve) + moved_without = int((without != first).sum()) + moved_with = int((with_preserve != first).sum()) + assert moved_with < moved_without, ( + f"slot preservation did not reduce churn: {moved_with} vs {moved_without}" + ) + + # Preservation must only permute slots inside a rank, never move an expert + # to a different rank than MLB placed it on. + slots = NUM_REPLICAS // NUM_RANKS + a = without.numpy().reshape(NUM_LAYERS, NUM_RANKS, slots) + b = with_preserve.numpy().reshape(NUM_LAYERS, NUM_RANKS, slots) + assert np.array_equal(np.sort(a, axis=2), np.sort(b, axis=2)), ( + "slot preservation changed which experts live on which rank" + ) + + +@pytest.mark.parametrize( + "algorithm", + ["auto", "deepseek", "deepseek_hierarchical", "deepseek_vec"], +) +def test_mlb_algorithms_selectable(monkeypatch, algorithm): + monkeypatch.setenv("VLLM_MLB_L1_ALGORITHM", algorithm) + weight = _skewed_load() + phy2log = MlbEplbPolicy.rebalance_experts( + weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS + ) + _assert_valid_placement(phy2log) + + +def test_no_expert_groups(): + """VLLM passes num_groups=0 for models without expert groups.""" + weight = _skewed_load() + phy2log = MlbEplbPolicy.rebalance_experts(weight, NUM_REPLICAS, 0, 1, NUM_RANKS) + _assert_valid_placement(phy2log) + + +def test_ep_rank_reaches_the_placement_request(monkeypatch): + """Ultraep's placement kernel solves once per EP rank and validates the + value is in range, so the caller's rank must reach the request -- every + other algorithm ignores it, but the plumbing is shared.""" + import vllm.distributed.eplb.mlb_runtime as mlb_runtime + + real_placement_request = mlb_runtime.placement_request + captured = {} + + def spy(*args, **kwargs): + captured["ep_rank"] = kwargs.get("ep_rank") + return real_placement_request(*args, **kwargs) + + monkeypatch.setattr(mlb_runtime, "placement_request", spy) + weight = _skewed_load() + MlbEplbPolicy.rebalance_experts( + weight, NUM_REPLICAS, NUM_GROUPS, 1, NUM_RANKS, ep_rank=5 + ) + assert captured["ep_rank"] == 5 diff --git a/tests/distributed/test_eplb_mlb_routing.py b/tests/distributed/test_eplb_mlb_routing.py new file mode 100644 index 000000000000..9be3c7e5d7ee --- /dev/null +++ b/tests/distributed/test_eplb_mlb_routing.py @@ -0,0 +1,895 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for MLB-backed L2 token routing. + +Run on a single device without a process group: ``static`` and ``dynamic`` +replica policies make a local decision, so no EP collective is issued. The +``lplb`` policy does issue one and is covered by the multi-rank CUDA-graph +probe instead. + +The important property under test is the accuracy-preserving invariant: MLB +may only change *which physical replica* serves a token, never which logical +expert the router chose. +""" + +import pytest +import torch + +pytest.importorskip("moe_load_balancer") + +from vllm.distributed.eplb.eplb_state import ( # noqa: E402 + EplbLayerState, + compute_logical_maps, +) +from vllm.distributed.eplb.mlb_runtime import ( # noqa: E402 + MlbRoutingRuntime, + reset_mlb_routing, +) + +NUM_LAYERS = 4 +NUM_LOGICAL = 32 +NUM_PHYSICAL = 40 # 8 redundant replicas +EP_SIZE = 4 +NUM_TOKENS = 64 +TOPK = 4 + + +@pytest.fixture(autouse=True) +def _reset(): + yield + reset_mlb_routing() + + +def _placement() -> torch.Tensor: + """A placement where the first 8 logical experts get a second replica.""" + rows = [] + for _ in range(NUM_LAYERS): + row = list(range(NUM_LOGICAL)) + list(range(NUM_PHYSICAL - NUM_LOGICAL)) + rows.append(row) + return torch.tensor(rows, dtype=torch.int64) + + +def _runtime(algorithm: str) -> MlbRoutingRuntime: + phy2log = _placement() + log2phy, replica_count = compute_logical_maps(phy2log, NUM_LOGICAL) + rt = MlbRoutingRuntime( + algorithm, + ep_size=EP_SIZE, + ep_rank=0, + num_logical_experts=NUM_LOGICAL, + num_physical_experts=NUM_PHYSICAL, + physical_to_logical_map=phy2log, + ) + rt.register_logical_maps(log2phy, replica_count) + return rt + + +def _layer_state(rt: MlbRoutingRuntime, layer_id: int = 0) -> EplbLayerState: + state = EplbLayerState() + state.set_layer_state( + layer_id, + torch.zeros(NUM_LAYERS, NUM_PHYSICAL, dtype=torch.int64), + rt._logical_to_physical_map, + rt._logical_replica_count, + ) + state.should_record_tensor = torch.tensor(1, dtype=torch.int32) + return state + + +def test_every_replica_policy_resolves_ids_rather_than_going_silent(): + """A policy that returned nothing here was not falling back gracefully -- + it never ran at all, and the caller silently kept its built-in choice.""" + rt = _runtime("dynamic") + state = _layer_state(rt) + torch.manual_seed(0) + logical = torch.randint(0, NUM_LOGICAL, (NUM_TOKENS, TOPK), dtype=torch.int64) + shares, ids = rt.resolve_routing(logical, torch.rand(NUM_TOKENS, TOPK), state, None) + # Every replica policy resolves its own ids now; none hands back a share + # table for this method to apply. + assert shares is None + assert ids is not None and ids.shape == logical.shape + + +def test_ultraep_fast_refresh_degrades_gracefully_without_a_mover(): + """The refresh makes memory match the placement it solves by moving expert + weight with ``EplbState``'s own staging buffer and P2P communicator. A + caller that does not supply them has given the policy nothing to move + weight with, so the fast-refresh path must stay off and leave placement and + routing working on the slow rearrange cadence alone -- not fail, and not + half-enable itself.""" + num_local_physical = NUM_PHYSICAL // EP_SIZE + expert_weights = [ + [torch.zeros(num_local_physical, 8, 8), torch.zeros(num_local_physical, 8, 8)] + for _ in range(NUM_LAYERS) + ] + + phy2log = _placement() + rt = MlbRoutingRuntime( + "ultraep", + ep_size=EP_SIZE, + ep_rank=0, + num_logical_experts=NUM_LOGICAL, + num_physical_experts=NUM_PHYSICAL, + physical_to_logical_map=phy2log, + expert_weights=expert_weights, + # expert_buffer / communicator deliberately omitted. + ) + assert rt._ultraep_expert_weights is None + assert rt._ultraep_communicator is None + assert rt._ultraep_refresh_gate is None + + # Routing must still work normally -- the fast-refresh path is purely + # additive, so its absence must be silent, not a degraded dispatch. + log2phy, replica_count = compute_logical_maps(phy2log, NUM_LOGICAL) + rt.register_logical_maps(log2phy, replica_count) + state = _layer_state(rt) + torch.manual_seed(0) + logical = torch.randint(0, NUM_LOGICAL, (NUM_TOKENS, TOPK), dtype=torch.int64) + shares, ids = rt.resolve_routing(logical, torch.rand(NUM_TOKENS, TOPK), state, None) + assert shares is None + assert ids is not None and ids.shape == logical.shape + + +def test_ultraep_fast_refresh_skips_decode_batches(monkeypatch): + """UltraEP is a prefill-time algorithm (the paper's own scoping; SGLang's + reference should_refresh gates the same way, rejecting only its provable + "decode" stage). A decode batch must return before touching the refresh + gate's counter or moving any weight.""" + import vllm.distributed.eplb.mlb_runtime as mlb_runtime + + num_local_physical = NUM_PHYSICAL // EP_SIZE + expert_weights = [ + [torch.zeros(num_local_physical, 8, 8), torch.zeros(num_local_physical, 8, 8)] + for _ in range(NUM_LAYERS) + ] + phy2log = _placement() + rt = MlbRoutingRuntime( + "ultraep", + ep_size=EP_SIZE, + ep_rank=0, + num_logical_experts=NUM_LOGICAL, + num_physical_experts=NUM_PHYSICAL, + physical_to_logical_map=phy2log, + expert_weights=expert_weights, + ) + + # Force the gate open -- pre-existing quota buffers, so a non-decode call + # would proceed into real collective and weight-move work. A decode-stage + # call must never reach it: the stub raises if it does. + def _explode(*args, **kwargs): + raise AssertionError("must not move weight for a decode-stage batch") + + monkeypatch.setattr(rt, "_commit_placement_weights", _explode) + rt._ultraep_rank_quota_prefix = torch.zeros(NUM_LAYERS, NUM_LOGICAL, EP_SIZE) + monkeypatch.setattr(mlb_runtime, "_current_stage", lambda: "decode") + + topk_ids = torch.randint(0, NUM_LOGICAL, (NUM_TOKENS, TOPK), dtype=torch.int64) + rt._ultraep_fast_refresh(0, topk_ids, num_unpadded_tokens=None) + + +def test_current_stage_under_dp_ignores_this_ranks_own_decode_shape(monkeypatch): + """Everything the refresh does past its stage check issues a collective, so + the check has to reach the same verdict on every EP rank. Under DP it does + not: a rank with no requests runs `_dummy_run(uniform_decode=True)` and + looks like a textbook decode to itself while its peers run real prefill. + Answering "decode" off this rank's own descriptor is what splits the group, + so the DP-reduced flag wins and a locally-decode-shaped batch still reports + "mixed" when its peers are not decoding.""" + import vllm.distributed.eplb.mlb_runtime as mlb_runtime + from vllm.forward_context import BatchDescriptor + + class _Ctx: + # Shaped exactly like a decode to this rank, and vLLM's cross-rank + # reduction says not every rank agrees. + batch_descriptor = BatchDescriptor(num_tokens=4, num_reqs=4, uniform=True) + uniform_decode_across_dp = False + dp_metadata = object() + + monkeypatch.setattr( + mlb_runtime, "is_forward_context_available", lambda: True, raising=False + ) + monkeypatch.setattr( + "vllm.forward_context.is_forward_context_available", lambda: True + ) + monkeypatch.setattr("vllm.forward_context.get_forward_context", lambda: _Ctx()) + assert mlb_runtime._current_stage() == "mixed" + + # And when every rank really is decoding, all of them say so together. + _Ctx.uniform_decode_across_dp = True + assert mlb_runtime._current_stage() == "decode" + + +def test_dp_token_counts_reads_the_tensor_that_actually_exists(): + """The cross-rank token count is the input that makes every rank agree on + whether a batch is worth refreshing for, so reading it has to actually + return it. + + It lives on `DPMetadata.num_tokens_across_dp_cpu`. `set_forward_context` + accepts an argument called `num_tokens_across_dp`, but does not keep a field + by that name -- it is consumed into DPMetadata. Asking the context for it + directly returned None on every rank, which silently routed this check to + its single-rank fallback everywhere, i.e. produced precisely the per-rank + disagreement it exists to prevent, while looking like a working fix. The + assertion that matters here is simply that it is not None.""" + from vllm.distributed.eplb.mlb_runtime import _dp_token_counts + from vllm.forward_context import DPMetadata + + counts = torch.tensor([4096, 0, 0, 0], dtype=torch.int32) + + class _Ctx: + dp_metadata = DPMetadata(num_tokens_across_dp_cpu=counts) + + import vllm.forward_context as fc + + saved_avail, saved_get = fc.is_forward_context_available, fc.get_forward_context + try: + fc.is_forward_context_available = lambda: True + fc.get_forward_context = lambda: _Ctx() + got = _dp_token_counts() + assert got is not None, ( + "_dp_token_counts returned None with DP metadata present -- the " + "cross-rank check has silently degraded to a per-rank one" + ) + assert torch.equal(got, counts) + assert int(got.sum()) == 4096 + + # No DP: nothing to read, and the caller's local fallback is correct + # there because no peer exists to disagree with it. + class _NoDp: + dp_metadata = None + + fc.get_forward_context = lambda: _NoDp() + assert _dp_token_counts() is None + finally: + fc.is_forward_context_available = saved_avail + fc.get_forward_context = saved_get + + +def test_ultraep_fast_refresh_arms_for_composed_expressions(): + """`VLLM_MLB_L2_ALGORITHM` carries a whole pipeline expression, not a bare + policy name. UltraEP owns placement in every expression it appears as the + placement policy of -- `ultraep` and `ultraep+waterfill` alike -- and the + refresh has to arm for all of them. + + It used to be selected by `algorithm == "ultraep"`, so any composition + silently disarmed it, and the failure was not graceful: `_snapshot` still + saw an allocated quota tensor and routed through placement tables no + refresh had ever written.""" + from vllm.distributed.eplb.mlb_runtime import _algorithm_places_with_ultraep + + assert _algorithm_places_with_ultraep("ultraep") + assert _algorithm_places_with_ultraep("ultraep+waterfill") + assert not _algorithm_places_with_ultraep("lplb") + assert not _algorithm_places_with_ultraep("static") + assert not _algorithm_places_with_ultraep("waterfill") + # An unparsable expression is not UltraEP's to claim, and must not raise here. + assert not _algorithm_places_with_ultraep("not-a-policy") + + +def test_snapshot_ignores_uncommitted_ultraep_tables(): + """Allocating the quota tensor and filling a layer's slice are different + events: the bootstrap solve does the first for the whole model, a refresh + does the second one layer at a time. Between them the MLB-side tables + describe a layout the framework never applied, so routing through them + sends tokens to slots holding other experts. Until a layer is committed, + the snapshot must hand back the framework's own candidates instead.""" + num_local_physical = NUM_PHYSICAL // EP_SIZE + expert_weights = [ + [torch.zeros(num_local_physical, 8, 8), torch.zeros(num_local_physical, 8, 8)] + for _ in range(NUM_LAYERS) + ] + phy2log = _placement() + rt = MlbRoutingRuntime( + "ultraep", + ep_size=EP_SIZE, + ep_rank=0, + num_logical_experts=NUM_LOGICAL, + num_physical_experts=NUM_PHYSICAL, + physical_to_logical_map=phy2log, + expert_weights=expert_weights, + ) + log2phy, replica_count = compute_logical_maps(phy2log, NUM_LOGICAL) + rt.register_logical_maps(log2phy, replica_count) + state = _layer_state(rt) + + # Quota allocated whole, as the bootstrap solve leaves it, and deliberately + # inconsistent with the committed placement -- exactly the state that made + # the real failure route 90% of tokens to the wrong expert. + rt._ultraep_rank_quota_prefix = torch.zeros(NUM_LAYERS, NUM_LOGICAL, EP_SIZE) + rt._ultraep_logical_to_physical = torch.zeros( + NUM_LAYERS, NUM_LOGICAL, EP_SIZE, dtype=torch.int64 + ) + rt._ultraep_replica_counts = torch.ones(NUM_LAYERS, NUM_LOGICAL, dtype=torch.int64) + assert rt._ultraep_committed_layers == set() + + snap = rt._snapshot(state, 0) + # Uncommitted: the framework's candidates, which always agree with + # physical_to_logical_map. + assert torch.equal( + snap.logical_to_physical_candidates, + state.logical_to_physical_map[:, : rt._max_replicas], + ) + + # Committed: MLB's own tables win, which is the whole point of the branch. + rt._ultraep_committed_layers.add(0) + snap = rt._snapshot(state, 0) + assert torch.equal( + snap.logical_to_physical_candidates, rt._ultraep_logical_to_physical[0] + ) + + +def test_synthesized_default_prefers_local_replicas(): + """VLLM keeps no per-rank dispatch table, so the adapter builds one. + It must prefer a replica this rank owns -- that is the locality SGLang + gets for free by collapsing its candidate list.""" + from moe_load_balancer.adapters.vllm import nearest_replica_table + + phy2log = _placement() + log2phy, counts = compute_logical_maps(phy2log, NUM_LOGICAL) + slots = NUM_PHYSICAL // EP_SIZE + + for ep_rank in range(EP_SIZE): + table = nearest_replica_table( + log2phy[0], counts[0], ep_rank=ep_rank, num_local_physical_experts=slots + ) + assert table.shape == (NUM_LOGICAL,) + # Chosen slot must really hold that logical expert. + assert torch.equal(phy2log[0][table], torch.arange(NUM_LOGICAL)) + low, high = ep_rank * slots, (ep_rank + 1) * slots + for logical in range(NUM_LOGICAL): + valid = log2phy[0, logical, : counts[0, logical]] + local = valid[(valid >= low) & (valid < high)] + if local.numel() > 0: + assert int(table[logical]) in local.tolist(), ( + f"expert {logical} has a replica on rank {ep_rank} " + "but the default table points elsewhere" + ) + + +def test_waterfill_is_rejected_unless_the_shared_expert_is_dispatched(monkeypatch): + """A shared-expert decision needs somewhere to be written back, and on this + runtime that place exists only when shared-expert fusion is on: without it + the shared expert is a per-rank replicated MLP that never enters dispatch, + so there is no destination rank to choose. + + The guard used to live in the id converter, which raised on any decision + carrying a shared_expert_rank. It now lives at configuration time instead, + which is the better place -- a pipeline that cannot work is refused before + the server starts rather than part-way through a forward -- but the + property being guarded is the same one: never accept the decision and then + drop it. + """ + from vllm.distributed.eplb.mlb_runtime import l2_inapplicable_reason + + monkeypatch.setenv("VLLM_FUSE_SHARED_EXPERTS", "0") + import importlib + + import vllm.envs + + importlib.reload(vllm.envs) + for expr in ("waterfill", "ultraep+waterfill"): + reason = l2_inapplicable_reason(expr, num_redundant_experts=16) + assert reason is not None, f"{expr} must be refused without fusion" + assert "shared expert" in reason + + # ultraep alone routes nothing shared and is unaffected either way. + assert l2_inapplicable_reason("ultraep", num_redundant_experts=16) is None + + monkeypatch.setenv("VLLM_FUSE_SHARED_EXPERTS", "1") + importlib.reload(vllm.envs) + for expr in ("waterfill", "ultraep+waterfill"): + assert l2_inapplicable_reason(expr, num_redundant_experts=16) is None + + +def test_l2_applicability_is_decided_on_the_config_the_run_uses(monkeypatch): + """Only ParallelConfig may disable L2, because only it is certainly used. + + EPLBConfig defaults l2_algorithm from the environment, so the instance + EngineArgs builds from the field default carries an algorithm while its + redundancy is still zero. Deciding applicability there announced a disable + that never applied to the run -- once per engine core, in the same log that + the workers were filling with "L2 routing enabled". The decision belongs on + the object that ends up in the running config. + """ + from vllm.config.parallel import EPLBConfig, ParallelConfig + + monkeypatch.setenv("VLLM_MLB_L2_ALGORITHM", "ultraep") + + # The default instance takes the algorithm from the environment and leaves + # it alone: its zero redundancy is a default, not a deployment. + assert EPLBConfig().l2_algorithm == "ultraep" + + # The assembled config of a run that enables EPLB decides. Redundancy + # stays zero here, so ultraep -- which only has replicas to choose between + # when there are redundant experts -- is switched off. + assembled = ParallelConfig( + data_parallel_size=2, + enable_expert_parallel=True, + enable_eplb=True, + eplb_config=EPLBConfig(l2_algorithm="ultraep"), + ) + assert assembled.eplb_config.l2_algorithm == "" + + # A config that does not enable EPLB has no routing boundary for L2 to sit + # in and is left alone, silently: this is also what a throwaway instance + # built from the field defaults looks like, and it must not announce a + # disable that never applied to any run. + idle = ParallelConfig(eplb_config=EPLBConfig(l2_algorithm="ultraep")) + assert idle.eplb_config.l2_algorithm == "ultraep" + + +def test_placement_commit_refreshes_policy_state(): + """A committed rearrangement must be observable without re-registration: + vLLM updates its maps in place, and the runtime holds those tensors.""" + rt = _runtime("dynamic") + rt.on_placement_committed() # must not raise + # Simulate vLLM committing a new placement in place. + rt.physical_to_logical_map[0, NUM_LOGICAL:] = torch.arange( + NUM_PHYSICAL - NUM_LOGICAL, dtype=torch.int64 + ).flip(0) + new_log2phy, new_counts = compute_logical_maps( + rt.physical_to_logical_map, NUM_LOGICAL + ) + rt._logical_to_physical_map.copy_(new_log2phy) + rt._logical_replica_count.copy_(new_counts) + rt.on_placement_committed() + + state = _layer_state(rt) + torch.manual_seed(0) + logical = torch.randint(0, NUM_LOGICAL, (NUM_TOKENS, TOPK), dtype=torch.int64) + # The snapshot handed to MLB must reflect the committed maps, not the ones + # captured at registration. + snapshot = rt._snapshot(state, 0) + assert torch.equal( + snapshot.logical_to_physical_candidates, rt._logical_to_physical_map[0] + ) + shares, ids = rt.resolve_routing(logical, torch.rand(NUM_TOKENS, TOPK), state, None) + assert (shares is None) != (ids is None) + + +def test_capabilities_gate_the_work_the_framework_does(): + """MLB declares what the framework must prepare; Glue must honour it + instead of doing every lifecycle step for every policy.""" + dyn = _runtime("dynamic") + # dynamic keeps no placement-derived state and needs no dispatch table. + assert dyn.requires_post_topk_routing + assert not dyn.requires_placement_state + assert not dyn.requires_rank_dispatch_map + assert dyn._default_replicas is None, ( + "the rank dispatch table was built for a policy that never reads it" + ) + + st = _runtime("static") + assert st.requires_rank_dispatch_map + assert st._default_replicas is not None + + +def test_replica_routing_sees_every_replica_not_just_the_local_one(): + """A replica-routing policy must be shown the whole candidate list. + + Narrowing each expert to its local replica is what a rank-dispatch policy + asks for; handing the same view to LPLB decides in advance the very thing + LPLB is there to decide, and on a rank that holds a copy of everything it + sees it removes the redundancy entirely, so the LP is skipped. It also + keeps the returned shares indexable by vLLM's own map, since + ``replica_shares`` passes on the probabilities without the candidates. + """ + rt = _runtime("lplb") + assert not rt.requires_rank_dispatch_map + snapshot = rt._snapshot(_layer_state(rt), 0) + assert torch.equal( + snapshot.logical_to_physical_candidates, rt._logical_to_physical_map[0] + ), "LPLB was handed a narrowed candidate list" + assert torch.equal(snapshot.logical_to_physical_count, rt._logical_replica_count[0]) + # The first 8 logical experts are replicated by _placement(); the choice + # between their copies is exactly what the policy has to solve. + assert int(snapshot.logical_to_physical_count[0]) == 2 + + # `static` is the policy the collapse exists for, and it still gets it. + st = _runtime("static") + assert st.requires_rank_dispatch_map + st_snapshot = st._snapshot(_layer_state(st), 0) + assert int(st_snapshot.logical_to_physical_count[0]) == 1 + + +def test_ultraep_quota_and_candidates_reach_the_snapshot_only_when_set(): + """The per-layer quota table and candidate table UltraEP's L1 solve + produces have to reach L2 routing through the snapshot *together*, or the + quota's replica columns get read against vLLM's own candidates -- rebuilt + independently from physical_to_logical_map, not guaranteed to share MLB's + width or column ordering -- rather than the table MLB solved it against. + """ + rt = _runtime("ultraep") + state = _layer_state(rt) + + # Before any L1 solve has run, there is nothing to thread through, and + # the snapshot falls back to vLLM's own candidates untouched. + baseline = rt._snapshot(state, 0) + assert baseline.metadata.get("rank_quota_prefix") is None + assert torch.equal( + baseline.logical_to_physical_candidates, rt._logical_to_physical_map[0] + ) + + # A refresh stashes the whole per-layer tensors and commits the layer it + # wrote; _snapshot() must slice each by the layer actually being routed, + # and must prefer MLB's own candidate table over vLLM's for those layers. + # + # Committing is a separate step from stashing on purpose: the bootstrap + # solve allocates these tensors for the whole model before any refresh has + # written a layer, and what it leaves in them describes a placement the + # framework never applied. Routing through that sends tokens to slots + # holding other experts, so an uncommitted layer falls back instead -- + # covered by test_snapshot_ignores_uncommitted_ultraep_tables. + quota = torch.arange(NUM_LAYERS * 3 * 5, dtype=torch.int32).reshape( + NUM_LAYERS, 3, 5 + ) + candidates = ( + torch.arange(NUM_LAYERS * 3 * 5, dtype=torch.int64).reshape(NUM_LAYERS, 3, 5) + + 1000 + ) + counts = torch.full((NUM_LAYERS, 3), 2, dtype=torch.int64) + rt._ultraep_rank_quota_prefix = quota + rt._ultraep_logical_to_physical = candidates + rt._ultraep_replica_counts = counts + rt._ultraep_committed_layers = set(range(NUM_LAYERS)) + for layer_id in range(NUM_LAYERS): + snapshot = rt._snapshot(state, layer_id) + assert torch.equal(snapshot.metadata["rank_quota_prefix"], quota[layer_id]) + assert torch.equal( + snapshot.logical_to_physical_candidates, candidates[layer_id] + ) + assert torch.equal(snapshot.logical_to_physical_count, counts[layer_id]) + + +def test_commit_is_skipped_for_policies_without_placement_state(monkeypatch): + rt = _runtime("dynamic") + calls = [] + monkeypatch.setattr(rt, "_commit_layers", lambda ids: calls.append(list(ids))) + rt.on_placement_committed([0, 1]) + rt.announce_initial_placement() + assert calls == [], "dynamic does not keep placement state; nothing to refresh" + + +def test_only_changed_layers_are_refreshed(monkeypatch): + """Refreshing an untouched layer is pure stall: a shape-changing rebuild + costs a JIT build plus warmup per layer.""" + rt = _runtime("static") # any policy; we intercept the MLB call + seen = [] + monkeypatch.setattr(rt, "requires_placement_state", True) + monkeypatch.setattr(rt, "_commit_layers", lambda ids: seen.append(list(ids))) + + rt.on_placement_committed([1, 3]) + assert seen == [[1, 3]] + + seen.clear() + rt.on_placement_committed([]) + assert seen == [], "an empty change set must not touch any layer" + + seen.clear() + rt.on_placement_committed(None) # unknown -> conservative full refresh + assert seen == [list(range(NUM_LAYERS))] + + +def test_graphs_plus_rearranging_placement_state_is_refused(monkeypatch): + """LPLB replaces its solver tensors when a committed placement changes + their layout, and a captured CUDA graph cannot follow that -- replay reads + freed memory, observed as Xid 43 on two devices. It is intermittent (it + needs a rearrangement that actually changes the layout), so it has to be + refused up front rather than left to surface in a long serving run.""" + import sys + import types + + from vllm.distributed.eplb import mlb_runtime + + class _Mode: + name = "FULL_AND_PIECEWISE" + + class _Compilation: + cudagraph_mode = _Mode() + + class _Cfg: + compilation_config = _Compilation() + + stub = types.ModuleType("vllm.config") + stub.get_current_vllm_config = lambda: _Cfg() # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "vllm.config", stub) + + with pytest.raises(ValueError, match="freed memory|Xid 43"): + mlb_runtime._reject_graphs_with_rearranging_placement_state(rearranges=True) + + # Pinned placement is the supported combination and must stay allowed -- + # that is how the reported LPLB numbers were measured. + mlb_runtime._reject_graphs_with_rearranging_placement_state(rearranges=False) + + _Mode.name = "NONE" # eager is the other escape + mlb_runtime._reject_graphs_with_rearranging_placement_state(rearranges=True) + + +def test_initial_placement_is_announced(): + """The contract is 'after initial weight loading AND after each commit'; + skipping the first half leaves the policy to build state lazily inside the + first forward, from whatever placement that forward happens to see.""" + rt = _runtime("static") + seen = [] + rt.requires_placement_state = True + rt._commit_layers = lambda ids: seen.append(list(ids)) + rt.announce_initial_placement() + assert seen == [list(range(NUM_LAYERS))] + + +def _runtime_with_redundancy(num_redundant: int) -> MlbRoutingRuntime: + """A runtime whose only interesting property is whether experts repeat.""" + num_physical = NUM_LOGICAL + num_redundant + phy2log = torch.stack( + [torch.arange(num_physical) % NUM_LOGICAL for _ in range(NUM_LAYERS)] + ) + return MlbRoutingRuntime( + "lplb", + ep_size=EP_SIZE, + ep_rank=0, + num_logical_experts=NUM_LOGICAL, + num_physical_experts=num_physical, + physical_to_logical_map=phy2log, + ) + + +@pytest.mark.parametrize("num_redundant, allocated", [(0, False), (8, True)]) +def test_count_buffers_exist_only_where_a_policy_consumes_them( + num_redundant, allocated, monkeypatch +): + """An LP with nothing to split must not cost a collective every step. + + Without redundant experts every logical expert holds exactly one replica, + so LPLB returns identity dispatch and never reads the count. These buffers + are what make vLLM run a per-layer count kernel and one EP all-reduce per + step to produce it -- work whose result is then discarded. + """ + monkeypatch.delenv("MLB_KEEP_ZERO_REDUNDANCY_COUNTS", raising=False) + rt = _runtime_with_redundancy(num_redundant) + assert (rt._logical_count_local is not None) is allocated + assert (rt._logical_count_global is not None) is allocated + + +def test_the_zero_redundancy_override_restores_the_buffers(monkeypatch): + """The override exists so the skipped work has a measured size rather than + only an argued one; a benchmark that cannot restore the old path cannot + report what removing it bought.""" + monkeypatch.setenv("MLB_KEEP_ZERO_REDUNDANCY_COUNTS", "1") + rt = _runtime_with_redundancy(0) + assert rt._logical_count_local is not None + assert rt._logical_count_global is not None + + +def test_finalize_step_counts_is_inert_without_buffers(monkeypatch): + """It must not reach for the EP group when there is nothing to reduce. + + These tests run with no distributed init, so an unguarded get_ep_group() + would raise -- which is exactly the assertion. + """ + monkeypatch.delenv("MLB_KEEP_ZERO_REDUNDANCY_COUNTS", raising=False) + _runtime_with_redundancy(0).finalize_step_counts() + + +def _sized_runtime(algorithm: str, num_redundant: int = 8) -> MlbRoutingRuntime: + num_physical = NUM_LOGICAL + num_redundant + phy2log = torch.stack( + [torch.arange(num_physical) % NUM_LOGICAL for _ in range(NUM_LAYERS)] + ) + return MlbRoutingRuntime( + algorithm, + ep_size=EP_SIZE, + ep_rank=0, + num_logical_experts=NUM_LOGICAL, + num_physical_experts=num_physical, + physical_to_logical_map=phy2log, + ) + + +@pytest.mark.parametrize( + "algorithm, consumes", + [("lplb", True), ("static", False), ("dynamic", False), ("ultraep", False)], +) +def test_only_policies_that_read_the_load_pay_for_gathering_it(algorithm, consumes): + """The count machinery follows what the policy declares, not "has a + post-TopK policy". + + A replica policy can route from placement alone. Keying the count on + requires_post_topk_routing charged every such policy a per-layer kernel and + one collective per step for a number it never reads. + """ + rt = _sized_runtime(algorithm) + assert rt.consumes_global_logical_count is consumes + assert (rt._logical_count_local is not None) is consumes + + +def test_the_runtime_reports_what_the_framework_must_respect(): + """Concurrency and graph stability are declared, not inferred from LPLB. + + Both guards used to key off properties LPLB happens to have, which barred + combinations that were never at risk of the fault being guarded against. + """ + st = _sized_runtime("static") + assert st.supports_concurrent_microbatches is True + assert st.graph_stability == "stable" + + lp = _sized_runtime("lplb") + assert lp.supports_concurrent_microbatches is False + assert lp.graph_stability == "realloc_on_placement_change" + + # ultraep caches a per-layer state dict on every route_tokens() call, and + # the quota tensor it caches is a fresh allocation each L1 replan rather + # than an in-place update -- the same hazard LPLB's solver state has. + ultraep_rt = _sized_runtime("ultraep") + assert ultraep_rt.supports_concurrent_microbatches is False + assert ultraep_rt.graph_stability == "realloc_on_placement_change" + + +def test_placement_and_routing_share_one_balancer(): + """An engine gets one balancer serving both layers. + + Two instances cannot see each other, which rules out any placement decision + that depends on what routing observed. + """ + from vllm.distributed.eplb.mlb_runtime import ( + get_mlb_integration, + reset_mlb_routing, + ) + + reset_mlb_routing() + try: + integration = get_mlb_integration() + integration.algorithm = "lplb" + num_physical = NUM_LOGICAL + 8 + phy2log = torch.stack( + [torch.arange(num_physical) % NUM_LOGICAL for _ in range(NUM_LAYERS)] + ) + routing = integration.bind_routing( + ep_size=EP_SIZE, + ep_rank=0, + num_logical_experts=NUM_LOGICAL, + num_physical_experts=num_physical, + physical_to_logical_map=phy2log, + ) + assert integration.balancer() is routing._mlb + assert get_mlb_integration() is integration + finally: + reset_mlb_routing() + + +def test_async_placement_commits_are_delivered_on_the_main_thread(): + """The async worker commits a layer's new map without telling the policy. + + Left undelivered, a policy holding placement-derived state keeps solving + against the layout the layer used to have. The dispatch kernel clamps to the + live replica count and reads the committed map, so the token still reaches a + valid replica of the right expert -- what degrades silently is the split, + which was computed to balance a placement that no longer exists. + + Delivery is queued rather than immediate because rebuilding that state does + GPU work and the commit runs on a worker thread. + """ + from vllm.distributed.eplb import eplb_state as st + + st._ASYNC_COMMITTED_LAYERS.clear() + st._ASYNC_COMMITTED_LAYERS.update({2, 0}) + + rt = _runtime("lplb") + seen: list[list[int]] = [] + rt.on_placement_committed = lambda ids: seen.append(list(ids)) + rt.finalize_step_counts = lambda: None + + # Mirror what prepare_forward does with the queue. + if st._ASYNC_COMMITTED_LAYERS and rt.requires_placement_state: + pending = sorted(st._ASYNC_COMMITTED_LAYERS) + st._ASYNC_COMMITTED_LAYERS.clear() + rt.on_placement_committed(pending) + + assert seen == [[0, 2]] + assert not st._ASYNC_COMMITTED_LAYERS + + +def test_a_policy_without_placement_state_needs_no_such_delivery(): + """Static routes from the committed map directly, so there is nothing + cached that a placement change could invalidate.""" + rt = _runtime("static") + assert rt.requires_placement_state is False + + +def test_the_nearest_replica_table_is_refreshed_on_every_commit(): + """This side's derived table has to follow the placement, whatever the + policy keeps. + + static reads this table and holds nothing else, so gating the refresh on + requires_placement_state left it routing by the placement the run started + with. Once experts move, most defaults are no longer in their expert's + candidate list, the share table falls back to column 0 for them, and the + traffic those replicas exist to spread piles onto one. + """ + for algorithm, keeps_state in (("static", False), ("lplb", True)): + rt = _runtime(algorithm) + assert rt.requires_placement_state is keeps_state + calls: list[int] = [] + rt._rebuild_default_replicas = lambda c=calls: c.append(1) + rt._commit_layers = lambda ids: None + rt.on_placement_committed([0]) + assert calls, ( + f"{algorithm}: the nearest-replica table was not refreshed, so it " + "still describes the placement before this commit" + ) + + +def test_bootstrap_solve_is_routed_against_not_just_stored(monkeypatch): + """A committed placement's own tables must be usable by the very next forward. + + `_snapshot` hands the router MLB's candidate table, replica counts and quota + only for layers in `_ultraep_committed_layers`. That set used to be filled in + one place -- inside `_ultraep_fast_refresh` -- while the commit path filled + the three tables for *every* layer and marked none of them. The router + therefore ignored the solve it had just been handed and fell back to vLLM's + own candidates until a fast refresh rewrote each layer individually, which at + interval=64 is most of a run: measured balance sat at the no-balancer + baseline for the first ~50 sampling windows of those cells. + + The tables come from the same plan whose physical_to_logical_map the caller + commits, so they are consistent for every layer at that moment; there is + nothing to wait for. + """ + from vllm.distributed.eplb import mlb_runtime + + num_layers, num_logical, max_replicas = 4, NUM_LOGICAL, 3 + + class _Plan: + physical_to_logical_map = torch.zeros( + (num_layers, NUM_PHYSICAL), dtype=torch.int32 + ) + logical_to_all_physical_map = torch.zeros( + (num_layers, num_logical, max_replicas), dtype=torch.int32 + ) + logical_to_physical_count = torch.ones( + (num_layers, num_logical), dtype=torch.int32 + ) + metadata = { + "rank_quota_prefix": torch.ones( + (num_layers, num_logical, max_replicas), dtype=torch.int32 + ) + } + + class _Routing: + _ultraep_rank_quota_prefix = None + _ultraep_logical_to_physical = None + _ultraep_replica_counts = None + _ultraep_committed_layers: set[int] = set() + + class _Balancer: + def plan_placement(self, request): + return _Plan() + + routing = _Routing() + + class _Integration: + def __init__(self): + self.routing = routing + + def balancer(self): + return _Balancer() + + integration = _Integration() + monkeypatch.setattr(mlb_runtime, "get_mlb_integration", lambda: integration) + # plan_placement imports this inside the function body, so the name has to + # be replaced where it is looked up rather than on mlb_runtime itself. + monkeypatch.setattr( + "moe_load_balancer.adapters.vllm.to_vllm_physical_to_logical", + lambda plan: plan.physical_to_logical_map, + ) + + mlb_runtime.plan_placement(object()) + + assert routing._ultraep_rank_quota_prefix is not None, ( + "the commit path must stash the quota" + ) + assert routing._ultraep_committed_layers == set(range(num_layers)), ( + "every layer the commit path wrote tables for must be marked usable; " + "leaving the set empty makes the router discard this solve until a fast " + "refresh rewrites each layer, and how long that takes scales with the " + "refresh interval" + ) diff --git a/tests/distributed/test_eplb_replica_shares_kernel.py b/tests/distributed/test_eplb_replica_shares_kernel.py new file mode 100644 index 000000000000..9b069f1fa873 --- /dev/null +++ b/tests/distributed/test_eplb_replica_shares_kernel.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The EPLB mapping kernel's replica-share path. + +A load balancer decides how much of a logical expert's traffic each physical +replica should take. Applying that split stays in this kernel, which also +records expert load -- so the split must be honoured *without* disturbing the +mapping's correctness or the recording. +""" + +import pytest +import torch + +from vllm.model_executor.layers.fused_moe.router.base_router import ( + _eplb_map_and_record_triton, +) +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="the mapping kernel is CUDA-only" +) + +NUM_LOGICAL = 8 +NUM_PHYSICAL = 12 # experts 0..3 get a second replica +NUM_TOKENS = 512 +TOPK = 4 + + +def _placement(device): + """physical_to_logical: 0..7 then a second replica of 0..3.""" + phy2log = torch.tensor( + list(range(NUM_LOGICAL)) + list(range(NUM_PHYSICAL - NUM_LOGICAL)), + dtype=torch.int64, + device=device, + ) + # logical -> [primary, secondary or -1] + log2phy = torch.full((NUM_LOGICAL, 2), -1, dtype=torch.int32, device=device) + counts = torch.ones(NUM_LOGICAL, dtype=torch.int32, device=device) + for logical in range(NUM_LOGICAL): + log2phy[logical, 0] = logical + for logical in range(NUM_PHYSICAL - NUM_LOGICAL): + log2phy[logical, 1] = NUM_LOGICAL + logical + counts[logical] = 2 + return phy2log, log2phy, counts + + +def _run(log2phy, counts, topk_ids, prob, device): + load = torch.zeros(NUM_PHYSICAL, dtype=torch.int64, device=device) + out = _eplb_map_and_record_triton( + topk_ids=topk_ids, + logical_to_physical_map=log2phy, + logical_replica_count=counts, + expert_load_view=load, + record_enabled=torch.tensor(1, dtype=torch.int32, device=device), + num_unpadded_tokens=None, + replica_prob=prob, + ) + return out, load + + +@pytest.mark.parametrize("target", [0, 1]) +def test_a_degenerate_share_sends_every_token_to_that_replica(target): + """The sharpest check that the table is actually consulted: put all of an + expert's share on one replica and every one of its tokens must land there.""" + device = torch.device("cuda") + phy2log, log2phy, counts = _placement(device) + # Only experts with two replicas can demonstrate a choice. + topk_ids = torch.randint( + 0, + NUM_PHYSICAL - NUM_LOGICAL, + (NUM_TOKENS, TOPK), + dtype=torch.int32, + device=device, + ) + prob = torch.zeros(NUM_LOGICAL, 2, dtype=torch.float32, device=device) + prob[:, target] = 1.0 + + out, _ = _run(log2phy, counts, topk_ids, prob, device) + + expected = log2phy[topk_ids.long(), target] + assert torch.equal(out.to(torch.int32), expected), ( + f"tokens did not follow a share of 1.0 on replica {target}" + ) + + +def test_split_shares_are_respected_in_proportion(): + device = torch.device("cuda") + phy2log, log2phy, counts = _placement(device) + # One expert, so the sample is the whole batch. + topk_ids = torch.zeros((NUM_TOKENS, TOPK), dtype=torch.int32, device=device) + prob = torch.zeros(NUM_LOGICAL, 2, dtype=torch.float32, device=device) + prob[0, 0] = 0.25 + prob[0, 1] = 0.75 + + out, _ = _run(log2phy, counts, topk_ids, prob, device) + + secondary = int((out == NUM_LOGICAL).sum()) + fraction = secondary / out.numel() + # The draw is a hash of the token index, so this is deterministic, not + # statistical; allow slack only for the finite sample. + assert 0.65 < fraction < 0.85, ( + f"expected ~0.75 on the second replica, got {fraction:.3f}" + ) + + +def test_shares_never_break_the_logical_choice(): + """The accuracy-preserving invariant: whatever replica is chosen, mapping + it back must return the logical expert the router asked for.""" + device = torch.device("cuda") + phy2log, log2phy, counts = _placement(device) + torch.manual_seed(0) + topk_ids = torch.randint( + 0, NUM_LOGICAL, (NUM_TOKENS, TOPK), dtype=torch.int32, device=device + ) + prob = torch.rand(NUM_LOGICAL, 2, dtype=torch.float32, device=device) + prob[counts == 1, 1] = 0.0 # padded column must stay unreachable + + out, _ = _run(log2phy, counts, topk_ids, prob, device) + + assert torch.equal(phy2log[out.long()].to(torch.int32), topk_ids) + + +def test_an_all_zero_row_falls_back_to_the_builtin_choice(): + """'No preference' must not collapse every token onto replica 0.""" + device = torch.device("cuda") + phy2log, log2phy, counts = _placement(device) + topk_ids = torch.zeros((NUM_TOKENS, TOPK), dtype=torch.int32, device=device) + prob = torch.zeros(NUM_LOGICAL, 2, dtype=torch.float32, device=device) + + with_table, _ = _run(log2phy, counts, topk_ids, prob, device) + without_table, _ = _run(log2phy, counts, topk_ids, None, device) + + assert torch.equal(with_table, without_table) + + +def test_recording_still_happens_on_the_share_path(): + """Recording lives in this kernel; supplying shares must not disturb it.""" + device = torch.device("cuda") + phy2log, log2phy, counts = _placement(device) + topk_ids = torch.randint( + 0, NUM_LOGICAL, (NUM_TOKENS, TOPK), dtype=torch.int32, device=device + ) + prob = torch.rand(NUM_LOGICAL, 2, dtype=torch.float32, device=device) + prob[counts == 1, 1] = 0.0 + + out, load = _run(log2phy, counts, topk_ids, prob, device) + + assert int(load.sum()) == topk_ids.numel() + assert torch.equal( + load, torch.bincount(out.reshape(-1).long(), minlength=NUM_PHYSICAL) + ) + + +def test_matching_the_collective_does_not_publish_zero_counts(): + """A dummy forward must not wipe the counts a real one just published. + + `finalize_step_counts` copies the per-layer local counts into the global + buffer the LP reads, then clears the local one. Running it a second time + within a step -- which is what a DP rank on a dummy batch does beside a peer + on a real batch -- copies the freshly cleared buffer over the published + counts. The solve then sees an all-zero load, produces an exactly uniform + split, and LPLB silently becomes `dynamic` while still paying for the solve. + + That is not hypothetical: it is what shipped, and every LPLB measurement + taken before this fix was of a policy solving against zeros. + + So the dummy path reduces a scratch tensor instead. This checks the counts + survive it. + """ + import torch + + from vllm.distributed.eplb.mlb_runtime import MlbRoutingRuntime + + rt = MlbRoutingRuntime.__new__(MlbRoutingRuntime) + rt._logical_count_local = torch.zeros(2, 8) + rt._logical_count_global = torch.zeros(2, 8) + rt._logical_count_ready = False + rt._count_scratch = None + + published = torch.tensor( + [[3.0, 1.0, 0, 0, 0, 0, 0, 0], [0, 2.0, 5.0, 0, 0, 0, 0, 0]] + ) + rt._logical_count_local.copy_(published) + + class _Group: + def all_reduce(self, t): # single rank: reduction is identity + return t + + import vllm.distributed as dist_mod + + orig = dist_mod.get_ep_group + dist_mod.get_ep_group = lambda: _Group() + try: + rt.finalize_step_counts() + assert torch.equal(rt._logical_count_global, published) + # The dummy path runs next, in the same step. + rt.match_step_counts_collective() + assert torch.equal(rt._logical_count_global, published), ( + "the dummy path overwrote the published counts; the solve would " + "see zeros and split uniformly" + ) + assert rt._count_scratch is not None, "no collective was issued" + finally: + dist_mod.get_ep_group = orig diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index b0c3e26db99a..897859c42668 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -10,6 +10,7 @@ from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.layers.fused_moe.config import RoutingMethodType from vllm.model_executor.layers.fused_moe.router.base_router import ( + _eplb_map_and_record_triton, eplb_map_to_physical_and_record, ) from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import ( @@ -819,6 +820,71 @@ def test_eplb_map_no_redundancy( assert load.sum().item() == 0 +def test_eplb_replica_shares_steer_routing_from_a_narrow_table(): + """A supplied share table decides the replica, and it may be narrower than + the candidate map. + + vLLM pads the map to ``MAX_EXPERT_REDUNDANCY + 1`` (1024) columns whatever + the redundancy, and the table is walked by a compile-time loop, so a table + as wide as the map cannot be compiled in reasonable time. Only the width + that can hold real replicas is needed, and the map is still indexed at its + own width. + """ + num_tokens, num_logical, R = 4096, 8, 2 + map_slots = 64 # padding, as vLLM does (it uses 1024) + num_physical = num_logical + R - 1 + + l2p = torch.full((num_logical, map_slots), -1, dtype=torch.int64, device="cuda") + l2p[0, :R] = torch.arange(R, dtype=torch.int64, device="cuda") + for i in range(1, num_logical): + l2p[i, 0] = R + i - 1 + rc = torch.tensor([R] + [1] * (num_logical - 1), dtype=torch.int64, device="cuda") + + torch.manual_seed(0) + topk_ids = torch.randint( + 1, num_logical, (num_tokens, 2), dtype=torch.int32, device="cuda" + ) + topk_ids[:, 0] = 0 # every token hits the replicated expert + rec = torch.tensor(True, dtype=torch.bool, device="cuda") + + for target in range(R): + prob = torch.zeros((num_logical, R), dtype=torch.float32, device="cuda") + prob[0, target] = 1.0 # send expert 0's traffic to exactly one replica + load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") + _eplb_map_and_record_triton( + topk_ids=topk_ids, + logical_to_physical_map=l2p, + logical_replica_count=rc, + expert_load_view=load, + record_enabled=rec, + num_unpadded_tokens=None, + replica_prob=prob, + ) + hot = load[:R] + assert int(hot[target]) == num_tokens, ( + f"share table asked for replica {target}, got {hot.tolist()}" + ) + assert int(hot.sum()) == num_tokens + + # An all-zero row means "no preference" and must fall back to the hash, + # not collapse every token onto replica 0. + prob = torch.zeros((num_logical, R), dtype=torch.float32, device="cuda") + load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") + _eplb_map_and_record_triton( + topk_ids=topk_ids, + logical_to_physical_map=l2p, + logical_replica_count=rc, + expert_load_view=load, + record_enabled=rec, + num_unpadded_tokens=None, + replica_prob=prob, + ) + hot = load[:R].float() + assert (hot.max() / hot.mean()).item() < 1.15, ( + f"all-zero row did not fall back to the hash: {hot.tolist()}" + ) + + @pytest.mark.parametrize("top_k,R", [(2, 2), (4, 2), (8, 4), (8, 8)]) def test_eplb_map_hot_expert_replica_balance(top_k, R): """Hot logical expert with R replicas must be balanced across replicas @@ -1006,3 +1072,122 @@ def test_eplb_map_num_unpadded_tokens( exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda") torch.testing.assert_close(load, exp_load) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_a_policy_may_resolve_replicas_itself_and_still_be_recorded(): + """The contract's primary return is per-token physical ids. + + A runtime that only accepts a share table forces every policy to express + its decision as a per-expert distribution, and silently ignores the ones + that cannot -- which is what made static and dynamic appear to do nothing + here. Taking the ids directly must still leave load recording in this + kernel, since that is the reason the boundary is inside it at all. + """ + from vllm.model_executor.layers.fused_moe.router.base_router import ( + _eplb_map_and_record_triton, + ) + + torch.manual_seed(0) + num_logical, num_physical = 8, 12 + tokens, topk = 32, 2 + logical_to_physical = torch.full( + (num_logical, 2), -1, dtype=torch.int32, device="cuda" + ) + logical_to_physical[:, 0] = torch.arange( + num_logical, dtype=torch.int32, device="cuda" + ) + logical_to_physical[:4, 1] = torch.arange( + num_logical, num_physical, dtype=torch.int32, device="cuda" + ) + counts = torch.ones(num_logical, dtype=torch.int32, device="cuda") + counts[:4] = 2 + + topk_ids = torch.randint( + 0, num_logical, (tokens, topk), dtype=torch.int32, device="cuda" + ) + # Resolve every token onto the *second* replica where one exists -- a choice + # the built-in hash would not make for every token, so the result is + # distinguishable from the fallback. + chosen = torch.where( + counts[topk_ids.long()] > 1, + logical_to_physical[topk_ids.long(), 1], + logical_to_physical[topk_ids.long(), 0], + ).to(torch.int32) + + load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") + out = _eplb_map_and_record_triton( + topk_ids=topk_ids, + logical_to_physical_map=logical_to_physical, + logical_replica_count=counts, + expert_load_view=load, + record_enabled=torch.tensor(1, dtype=torch.int32, device="cuda"), + num_unpadded_tokens=None, + physical_ids=chosen, + ) + assert torch.equal(out, chosen), "the policy's own ids must be used verbatim" + # Recording is the reason this path goes through the kernel rather than + # around it: every routed token must still land in the load view. + assert int(load.sum()) == tokens * topk + for pid in chosen.unique(): + assert int(load[int(pid)]) == int((chosen == pid).sum()) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_policy_resolved_ids_stay_valid_and_fully_recorded(): + """Whatever a policy hands back, two things must still hold. + + Every id has to name a live replica of the expert the router chose -- the + policy replaces the selection, not the routing decision -- and the load + view has to account for every routed slot, since recording is the reason + resolved ids come through this kernel rather than around it. + """ + torch.manual_seed(0) + NL, NP, EP = 64, 80, 8 + per_rank = NP // EP + l2p = torch.full((NL, 2), -1, dtype=torch.int32, device="cuda") + l2p[:, 0] = torch.arange(NL, dtype=torch.int32, device="cuda") + l2p[:16, 1] = torch.arange(NL, NP, dtype=torch.int32, device="cuda") + counts = torch.ones(NL, dtype=torch.int32, device="cuda") + counts[:16] = 2 + topk = torch.randint(0, NL, (2048, 6), dtype=torch.int32, device="cuda") + + # Pin every expert to the replica nearest rank 3 -- a decision the built-in + # hash would not make, so a fallback would be visible. + mlb_vllm = pytest.importorskip("moe_load_balancer.adapters.vllm") + nearest_replica_table = mlb_vllm.nearest_replica_table + + defaults = nearest_replica_table( + l2p.to(torch.int64), + counts.to(torch.int64), + ep_rank=3, + num_local_physical_experts=per_rank, + ).to(torch.int32) + resolved = defaults[topk.to(torch.int64)] + + load = torch.zeros(NP, dtype=torch.int32, device="cuda") + out = _eplb_map_and_record_triton( + topk_ids=topk, + logical_to_physical_map=l2p, + logical_replica_count=counts, + expert_load_view=load, + record_enabled=torch.tensor(1, dtype=torch.int32, device="cuda"), + num_unpadded_tokens=None, + physical_ids=resolved, + ) + assert torch.equal(out, resolved), "the policy's ids were not used verbatim" + + # Each id names a replica of the right logical expert. + logical_of = torch.full((NP,), -1, dtype=torch.int32, device="cuda") + for j in range(2): + col = l2p[:, j] + live = col >= 0 + logical_of[col[live].to(torch.int64)] = torch.arange( + NL, dtype=torch.int32, device="cuda" + )[live] + assert torch.equal(logical_of[out.to(torch.int64)], topk) + + # And every routed slot is accounted for, slot by slot. + assert int(load.sum()) == topk.numel() + expected = torch.bincount(out.flatten().to(torch.int64), minlength=NP) + assert torch.equal(load.to(torch.int64), expected) diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index b9136cfbb244..7c2338807cb0 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -37,7 +37,7 @@ ExpertPlacementStrategy = Literal["linear", "round_robin"] DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"] DataParallelBackend = Literal["ray", "mp"] -EPLBPolicyOption = Literal["default"] +EPLBPolicyOption = Literal["default", "mlb"] DCPCommBackend = Literal["ag_rs", "a2a"] EPLBCommunicatorBackend = Literal[ "torch_nccl", "torch_gloo", "torch_xccl", "nixl", "pynccl" @@ -93,6 +93,11 @@ class EPLBConfig: policy: EPLBPolicyOption = "default" """The policy type for expert parallel load balancing (EPLB).""" + l2_algorithm: str = "" + """MoE Load Balancer routing expression for the L2 stage, which decides + which *replica* of a logical expert each token goes to. Empty leaves vLLM's + built-in choice in place. Defaults to ``$VLLM_MLB_L2_ALGORITHM``.""" + communicator: EPLBCommunicatorBackend | None = None """ Backend for EPLB expert weight communication: @@ -105,6 +110,28 @@ class EPLBConfig: on CUDA, falls back to "torch_gloo") """ + init_placement_path: str | None = None + """ + Path to a placement checkpoint saved by a previous profiling run. + + When set, the initial physical_to_logical_map is loaded from this file + instead of using the trivial (uniform) assignment. This is the vLLM + equivalent of SGLang's ``--init-expert-location``: profile once, save the + MLB-optimised placement, then start production with that placement pre-loaded + so LPLB has hot experts already replicated. + + The file must have been written by ``save_placement_path`` on a model with + the same topology (num_moe_layers, num_physical_experts). + """ + + save_placement_path: str | None = None + """ + If set, write the current physical_to_logical_map to this path after every + rearrangement. The most recent rearrangement overwrites the previous file. + Use in a short profiling run to capture a good placement, then pass the + result to ``init_placement_path`` in the serving run. + """ + @model_validator(mode="after") def _validate_eplb_config(self) -> Self: if self.use_async and self.policy != "default": @@ -118,6 +145,10 @@ def _validate_eplb_config(self) -> Self: ) if self.log_balancedness and self.log_balancedness_interval <= 0: raise ValueError("log_balancedness_interval must be greater than 0.") + + if not self.l2_algorithm: + self.l2_algorithm = envs.VLLM_MLB_L2_ALGORITHM + return self @@ -562,6 +593,48 @@ def _validate_parallel_config(self) -> Self: "num_redundant_experts." ) + # An L2 policy chooses among a logical expert's replicas. With no + # redundant experts there is exactly one replica each, so there is + # nothing to choose: the policy can only reproduce the mapping vLLM + # would have used anyway. Leaving it on is not merely useless -- the + # routing boundary is crossed per layer per forward, and a policy that + # consumes the EP-wide expert load pays a collective to produce a count + # it then discards. + # + # Turning it off here rather than making it cheap deeper down is what + # keeps it honest: the decision is visible, made once, and taken before + # anything sizes buffers or declares capabilities against it. Output is + # unaffected, which is what makes doing it silently unacceptable but + # doing it at all safe. + # + # It belongs on this object rather than on EPLBConfig because only this + # one is certain to be the config the run uses. EPLBConfig defaults + # l2_algorithm from the environment, so a throwaway instance built + # from the field defaults carries an algorithm while its redundancy is + # still zero; gating on enable_eplb keeps such instances quiet, and a + # run without EPLB has no routing boundary for the policy to sit in. + if self.enable_eplb and self.eplb_config.l2_algorithm: + from vllm.distributed.eplb.mlb_runtime import l2_inapplicable_reason + + reason = l2_inapplicable_reason( + self.eplb_config.l2_algorithm, + self.eplb_config.num_redundant_experts, + ) + if reason is not None: + # The reason comes from the balancer: which deployments a + # policy can work on is a property of the policy, and a + # framework that decides it by assumption will switch off a + # policy that had work to do. + logger.warning( + "Disabling MoE Load Balancer L2 routing (%r): %s. Routing " + "is unchanged by this -- the policy could only have " + "reproduced it -- but the per-layer routing boundary and " + "any load collective it needed are now skipped.", + self.eplb_config.l2_algorithm, + reason, + ) + self.eplb_config.l2_algorithm = "" + tp = self.tensor_parallel_size pcp = self.prefill_context_parallel_size dcp = self.decode_context_parallel_size diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 21a1f00da4a6..56500c0a7ebe 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -226,7 +226,14 @@ def _make_all2all_kwargs(self) -> dict[Any, Any]: num_qps_per_rank = None if self.internode and not envs.VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: - num_rdma_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024 + # The RDMA buffer may be sized separately from the NVL one: DeepEP + # caps num_nvl_bytes at INT_MAX whenever num_rdma_bytes != 0, and + # on EP groups beyond NUM_MAX_NVL_PEERS (internode dispatch) it is + # the RDMA buffer that runs out. + rdma_mb = envs.VLLM_DEEPEP_RDMA_BUFFER_SIZE_MB + if rdma_mb is None: + rdma_mb = envs.VLLM_DEEPEP_BUFFER_SIZE_MB + num_rdma_bytes = rdma_mb * 1024 * 1024 num_qps_per_rank = self.num_sms // 2 else: num_rdma_bytes = 0 @@ -244,6 +251,23 @@ def _make_all2all_kwargs(self) -> dict[Any, Any]: num_qps_per_rank=num_qps_per_rank, explicitly_destroy=True, ) + # On a multi-node NVLink domain (e.g. Blackwell multi-node NVLink + # systems) an EP group spans trays, + # and the buffer's default IPC path is node-local: it dies at startup + # with cudaErrorInvalidResourceHandle. deep_ep 2.0 selects the fabric + # path through allow_mnnvl. Opt-in. + if envs.VLLM_DEEPEP_HT_USE_MNNVL: + import inspect + + import deep_ep # type: ignore[import-not-found] + + if "allow_mnnvl" in inspect.signature(deep_ep.Buffer.__init__).parameters: + kwargs["allow_mnnvl"] = True + else: + logger.warning( + "VLLM_DEEPEP_HT_USE_MNNVL=1 but this deep_ep build has no " + "allow_mnnvl parameter; leaving the buffer on the IPC path." + ) return kwargs def get_handle(self, kwargs): @@ -1092,9 +1116,18 @@ def get_handle(self, kwargs): buffer_kwargs, deep_ep.ElasticBuffer ) if self._num_sms is None: + # Link bandwidths for DeepEP v2's SM estimator. It only probes + # rdma_gbs when num_rdma_ranks > 1 but divides by it whenever + # num_scaleout_ranks > 1; on an MNNVL machine scale-out rides the + # NVLink fabric, so the probe returns 0 and the estimator raises + # ZeroDivisionError. These values only size the SM estimate. + rdma_gbs = envs.VLLM_DEEPEP_V2_RDMA_GBS + nvlink_gbs = envs.VLLM_DEEPEP_V2_NVLINK_GBS self._num_sms = handle.get_theoretical_num_sms( num_experts=num_experts, num_topk=kwargs["num_topk"], + rdma_gbs=rdma_gbs, + nvlink_gbs=nvlink_gbs, ) return handle diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 205226cde4a9..2461706b0f78 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -31,8 +31,9 @@ from dataclasses import dataclass import torch -from torch.distributed import ProcessGroup, all_reduce +from torch.distributed import ProcessGroup, all_gather_into_tensor, all_reduce +import vllm.envs as envs from vllm.config import ModelConfig, ParallelConfig from vllm.config.utils import compute_hash_cached from vllm.distributed.parallel_state import ( @@ -63,6 +64,11 @@ logger = init_logger(__name__) +# A step interval this large means the deployment never intends to re-plan; +# treat it as "placement is pinned" rather than as a very slow cadence. +_REARRANGEMENT_EFFECTIVELY_DISABLED = 1_000_000 + + def _compute_eplb_load_stats( num_tokens_per_rank: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -231,6 +237,12 @@ class EplbModelState: """ +# Layers committed by the async worker whose policy notification is still owed. +# Written from the worker thread, drained on the main thread; set operations are +# atomic under the GIL and a missed drain only defers by one forward. +_ASYNC_COMMITTED_LAYERS: set[int] = set() + + class EplbState: """EplbState of each expert parallel model. Key is the model config hash.""" @@ -361,13 +373,58 @@ def add_model( self.validate_ep_configuration(model) self.is_async = self.parallel_config.eplb_config.use_async + # Assuming 8 GPUs per node, this supports up to + # (1023 + 1) / 8 = 128 nodes for now. + # TODO(rui): make this configurable + MAX_EXPERT_REDUNDANCY = 1023 + assert model.num_redundant_experts <= MAX_EXPERT_REDUNDANCY, ( + f"num_redundant_experts {model.num_redundant_experts} " + f"must be less than or equal to {MAX_EXPERT_REDUNDANCY}" + ) + max_slots_per_logical_expert = MAX_EXPERT_REDUNDANCY + 1 + + init_path = self.parallel_config.eplb_config.init_placement_path + # A checkpoint says where the experts should *end up*, not where they + # are. The loader has just filled every slot according to the trivial + # layout, so the live maps must start trivial and the weights have to be + # physically moved -- see `_install_initial_placement` below. Installing + # the checkpoint's map here instead would leave every slot claiming an + # expert it does not hold: throughput is unaffected and every output is + # silently wrong, which is exactly the failure a throughput benchmark + # cannot see. + pending_init_p2l: torch.Tensor | None = None + if init_path: + # Saved as [num_moe_layers, num_physical_experts] -- each layer has + # its own optimised placement. + p2l_full = torch.load( + init_path, map_location="cpu", weights_only=True + ).long() + assert p2l_full.shape == ( + model.num_moe_layers, + model.num_physical_experts, + ), ( + f"Placement checkpoint shape {tuple(p2l_full.shape)} does not " + f"match model ({model.num_moe_layers}, {model.num_physical_experts})" + ) + # Left on CPU: this takes the same path as a periodic + # rearrangement, where the target comes straight from + # `policy.rebalance_experts` and `compute_logical_maps` asserts CPU. + pending_init_p2l = p2l_full + logger.info( + "EPLB: loaded per-layer placement from %s (layers=%d, " + "num_physical=%d); experts will be rearranged to match it", + init_path, + model.num_moe_layers, + model.num_physical_experts, + ) + physical_to_logical_map_list = ( EplbState.build_initial_global_physical_to_logical_map( model.num_routed_experts, model.num_redundant_experts, ) ) - physical_to_logical_map = torch.tensor( + p2l_1d = torch.tensor( physical_to_logical_map_list, device=self.device, ) @@ -379,20 +436,11 @@ def add_model( * get_ep_group().world_size // self.parallel_config.data_parallel_size ) - physical_to_logical_map = torch.nn.functional.pad( - physical_to_logical_map, + p2l_1d = torch.nn.functional.pad( + p2l_1d, (0, physical_expert_capacity - model.num_physical_experts), value=-1, ) - # Assuming 8 GPUs per node, this supports up to - # (1023 + 1) / 8 = 128 nodes for now. - # TODO(rui): make this configurable - MAX_EXPERT_REDUNDANCY = 1023 - assert model.num_redundant_experts <= MAX_EXPERT_REDUNDANCY, ( - f"num_redundant_experts {model.num_redundant_experts} " - f"must be less than or equal to {MAX_EXPERT_REDUNDANCY}" - ) - max_slots_per_logical_expert = MAX_EXPERT_REDUNDANCY + 1 logical_to_physical_map = torch.full( (model.num_logical_experts, max_slots_per_logical_expert), -1, @@ -403,20 +451,14 @@ def add_model( device=self.device, dtype=torch.long, ) - for i in range(model.num_physical_experts): - logical_idx = physical_to_logical_map[i] + logical_idx = p2l_1d[i] logical_to_physical_map[logical_idx, logical_replica_count[logical_idx]] = i logical_replica_count[logical_idx] += 1 # Duplicate initial mapping for all layers physical_to_logical_map = ( - physical_to_logical_map.unsqueeze(0) - .expand( - model.num_moe_layers, - -1, - ) - .contiguous() + p2l_1d.unsqueeze(0).expand(model.num_moe_layers, -1).contiguous() ) physical_to_logical_map_buffer = physical_to_logical_map physical_to_logical_map = physical_to_logical_map_buffer[ @@ -424,19 +466,12 @@ def add_model( ] logical_to_physical_map = ( logical_to_physical_map.unsqueeze(0) - .expand( - model.num_moe_layers, - -1, - -1, - ) + .expand(model.num_moe_layers, -1, -1) .contiguous() ) logical_replica_count = ( logical_replica_count.unsqueeze(0) - .expand( - model.num_moe_layers, - -1, - ) + .expand(model.num_moe_layers, -1) .contiguous() ) @@ -517,6 +552,109 @@ def add_model( ) self.model_states[model_config.compute_hash()] = model_state + # Move the experts to where the checkpoint says they belong. This has + # to happen before the routing policy is told about the placement, so + # that what it is told is what is actually on the device. + if pending_init_p2l is not None: + self._install_initial_placement(model_state, pending_init_p2l) + + # Optional: hand L2 replica choice to the MoE Load Balancer. A no-op + # unless an L2 algorithm is configured, in which case vLLM's fused + # mapping kernel is bypassed at the routing boundary. + # + # Read from the config rather than the environment: EPLBConfig resolves + # $VLLM_MLB_L2_ALGORITHM once and clears it for placements the policy + # cannot act on, so by here the answer already accounts for redundancy. + from vllm.distributed.eplb.mlb_runtime import ( + init_mlb_routing, + l2_pipeline_capabilities, + ) + + l2_algorithm = self.parallel_config.eplb_config.l2_algorithm + if l2_algorithm: + if self.parallel_config.num_ubatches > 1: + caps = l2_pipeline_capabilities(l2_algorithm) + if caps is not None and not caps.supports_concurrent_microbatches: + raise ValueError( + f"MoE Load Balancer L2 routing {l2_algorithm!r} is " + "incompatible with DBO: it keeps one solver state per " + "layer, and concurrent micro-batches would clobber each " + "other. Disable DBO, or select a policy that declares " + "supports_concurrent_microbatches. Clearing " + "eplb_config.l2_algorithm also removes the conflict." + ) + ep_group = get_ep_group() + routing = init_mlb_routing( + algorithm=l2_algorithm, + ep_size=ep_group.world_size, + ep_rank=ep_group.rank_in_group, + num_logical_experts=model.num_logical_experts, + num_physical_experts=model.num_physical_experts, + physical_to_logical_map=physical_to_logical_map, + logical_to_physical_map=logical_to_physical_map, + logical_replica_count=logical_replica_count, + expert_weights=model.expert_weights, + # The same staging buffer and P2P communicator this state uses + # for its own periodic rearrangement. A policy that re-plans + # placement mid-run has to move weight to match it, and that is + # this file's machinery, not something a policy brings its own + # copy of -- see MlbRoutingRuntime._commit_placement_weights. + expert_buffer=expert_buffer, + communicator=communicator, + # A step interval that the run can actually reach means the + # placement will be re-planned, which is the case a captured + # graph cannot follow when the policy keeps solver state. + rearranges=( + self.expert_rearrangement_step_interval + < _REARRANGEMENT_EFFECTIVELY_DISABLED + ), + ) + if routing is not None: + # The contract is "after initial weight loading and after each + # committed update"; weights are loaded by this point, so this + # is the first half. + routing.announce_initial_placement() + + def _install_initial_placement( + self, + model_state: EplbModelState, + target_physical_to_logical_map: torch.Tensor, + ) -> None: + """Rearrange the freshly loaded weights into a checkpointed placement. + + The weights sit in the trivial layout after loading, so a checkpoint can + only be honoured by physically moving them; writing its map into the + live state without the move makes every slot claim an expert it does not + hold. That is invisible to a throughput benchmark -- the run is exactly + as fast, only the text is wrong -- so it is done here rather than being + left to the periodic rearrangement, which `step_interval` may disable. + + Args: + model_state: The model whose experts should be moved. + target_physical_to_logical_map: ``[num_moe_layers, + num_physical_experts]`` placement the weights must end up in. + + """ + ep_group = get_ep_group().device_group + rearrange_expert_weights_inplace( + model_state.physical_to_logical_map, + target_physical_to_logical_map, + model_state.model.expert_weights, + model_state.expert_buffer, + ep_group, + model_state.communicator, + False, + None, + ) + _commit_eplb_maps( + model_state, + new_physical_to_logical_map=target_physical_to_logical_map, + ) + logger.info( + "EPLB: rearranged experts into the checkpointed placement (%d layers)", + target_physical_to_logical_map.shape[0], + ) + def prepare_forward( self, model_config: ModelConfig, @@ -535,6 +673,23 @@ def prepare_forward( token range. When ``None``, only ``tensors[0]`` is filled. """ + # Finalise LPLB count cache from the previous forward pass. + # count_logical_experts fills _lplb_local_count per-layer inside the + # graph; we all_reduce and move to _lplb_global_count here (outside + # the graph, one collective for all layers). This gives LP solve a + # stable, up-to-date input without any NCCL inside the capture stream. + from vllm.distributed.eplb.mlb_runtime import get_mlb_routing + + routing = get_mlb_routing() + if routing is not None: + # Deliver any placement changes the async worker committed since the + # last forward, on this thread, before anything routes against them. + if _ASYNC_COMMITTED_LAYERS: + pending = sorted(_ASYNC_COMMITTED_LAYERS) + _ASYNC_COMMITTED_LAYERS.clear() + routing.on_placement_committed(pending) + routing.finalize_step_counts() + model_state = self.model_states.get(compute_hash_cached(model_config)) if model_state is None or model_state.num_unpadded_tokens_tensors is None: return @@ -622,16 +777,31 @@ def step( avg_tokens, max_tokens = tokens_tensors balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 + # The figure above reduces over dim 0, which is layers, so it + # reports how much a rank's load varies between its own layers. + # A balancer acts on something else: how far the busiest rank + # is from the average *within a layer*, since a layer's MoE + # step ends when its slowest rank does. Reported alongside + # rather than instead, so a number that has been quoted before + # keeps meaning what it did. + per_layer_mean = num_tokens_per_rank.mean(dim=1) + per_layer_peak = num_tokens_per_rank.max(dim=1).values + rank_imbalance = float( + (per_layer_peak / per_layer_mean.clamp(min=1e-9)).mean() + ) + if ep_group.rank() == 0: logger.info( "EPLB step: %d for model %s: avg_tokens=%.2f, " "max_tokens=%d, balancedness=%.4f, " + "rank_imbalance=%.4f, " "steps until the next rearrangement: %d", self.expert_rearrangement_step, eplb_model_state.model_name, avg_tokens, max_tokens, balancedness, + rank_imbalance, self.expert_rearrangement_step_interval - self.expert_rearrangement_step, ) @@ -814,8 +984,54 @@ def rearrange( global_expert_load_window = logical_expert_load_window[..., :-1].sum(dim=0) global_expert_load_windows.append(global_expert_load_window) - # Perform all-reduce to get the expert load across all ranks for each model - global_expert_load_windows = self._allreduce_list(global_expert_load_windows) + # A policy that places redundant replicas by which rank is overloaded, + # not just which logical expert is hot, needs every rank's own count, + # which a sum-reducing all-reduce discards. Gather instead and derive + # the sum locally -- every other policy sees the identical value it + # always has, just computed one way. + wants_per_rank = getattr(self.policy, "wants_per_rank_weight", None) + if wants_per_rank is not None and wants_per_rank(): + per_rank_expert_load_windows = self._allgather_list( + global_expert_load_windows + ) + global_expert_load_windows = [ + window.sum(dim=0) for window in per_rank_expert_load_windows + ] + else: + per_rank_expert_load_windows = [None] * len(global_expert_load_windows) + global_expert_load_windows = self._allreduce_list( + global_expert_load_windows + ) + + # Env-gated, write-only observation hook (no-op unless the variable is + # set): dump the global *logical* expert load. rank_imbalance is a pure + # function of (logical load, placement, replica routing), so one recorded + # load lets every policy x redundancy combination be scored on CPU + # afterwards instead of re-serving the model once per arm. + _dump = envs.VLLM_EPLB_DUMP_LOAD_PATH + if _dump and not is_profile and ep_group.rank() == 0: + # Numbered per rearrangement: the step counter has already been + # reset when this runs, so it cannot tell two dumps apart. + self._dump_seq = getattr(self, "_dump_seq", 0) + 1 + _st = next(iter(self.model_states.values())) + torch.save( + { + "logical_expert_load": global_expert_load_windows[0].cpu(), + "physical_to_logical_map": _st.physical_to_logical_map.cpu(), + "num_logical_experts": _st.model.num_logical_experts, + "num_physical_experts": _st.model.num_physical_experts, + "num_moe_layers": _st.model.num_moe_layers, + "num_expert_groups": _st.model.num_expert_groups, + "ep_size": ep_group.size(), + "rearrangement": self._dump_seq, + }, + f"{_dump}.rearr{self._dump_seq:03d}.pt", + ) + logger.info( + "EPLB: dumped logical expert load to %s.rearr%03d.pt", + _dump, + self._dump_seq, + ) # TODO(bowen): Treat differently for prefill and decode nodes eplb_model_state = next(iter(self.model_states.values())) @@ -848,14 +1064,32 @@ def rearrange( ) # Get new expert mappings - for eplb_model_state, global_expert_load_window in zip( - self.model_states.values(), global_expert_load_windows - ): + rearrange_inputs = zip( + self.model_states.values(), + global_expert_load_windows, + per_rank_expert_load_windows, + ) + for ( + eplb_model_state, + global_expert_load_window, + per_rank_expert_load_window, + ) in rearrange_inputs: if not is_profile: eplb_model_state.last_expert_load = global_expert_load_window if not self.is_async or is_profile: # Get new expert mappings for the model. The policy runs on the # host, so the load window and current map have to come back. + # Unlike those, per_rank_weight stays on-device: it only ever + # feeds ultraep's CUDA placement kernel, never the CPU numpy + # solve every other algorithm here uses. + rebalance_kwargs = ( + {} + if per_rank_expert_load_window is None + else { + "per_rank_weight": per_rank_expert_load_window, + "ep_rank": ep_rank, + } + ) with gpu_sync_allowed(): new_physical_to_logical_map = self.policy.rebalance_experts( global_expert_load_window.cpu(), @@ -864,6 +1098,7 @@ def rearrange( num_nodes, num_gpus, eplb_model_state.physical_to_logical_map.cpu(), + **rebalance_kwargs, ) skip_rearrange = False @@ -923,6 +1158,27 @@ def rank_load_imbalance( ) if not skip_rearrange: + # Which layers actually move? A pluggable routing policy + # may have to rebuild per-layer state, which for `lplb` + # costs a JIT build plus warmup; refreshing untouched + # layers is pure stall. SGLang gets this list from its own + # updater (`update_layer_ids`); vLLM does not track it, so + # diff here -- before `_commit_eplb_maps` overwrites the + # live map in place, which would make every layer look + # unchanged. + changed_layer_ids = ( + ( + eplb_model_state.physical_to_logical_map + != new_physical_to_logical_map.to( + eplb_model_state.physical_to_logical_map.device + ) + ) + .any(dim=-1) + .nonzero() + .flatten() + .tolist() + ) + # Update expert weights rearrange_expert_weights_inplace( eplb_model_state.physical_to_logical_map, @@ -940,6 +1196,14 @@ def rank_load_imbalance( eplb_model_state, new_physical_to_logical_map=new_physical_to_logical_map, ) + # Weights have moved and the live maps are updated, so + # a pluggable routing policy may now rebuild any + # placement-derived state. + from vllm.distributed.eplb.mlb_runtime import get_mlb_routing + + routing = get_mlb_routing() + if routing is not None: + routing.on_placement_committed(changed_layer_ids) if is_main_rank: assert start_event is not None @@ -952,6 +1216,16 @@ def rank_load_imbalance( " (profile) " if is_profile else " ", gpu_elapsed, ) + save_path = self.parallel_config.eplb_config.save_placement_path + if save_path and not is_profile: + torch.save( + eplb_model_state.physical_to_logical_map.cpu(), + save_path, + ) + logger.info( + "EPLB: saved placement checkpoint to %s", + save_path, + ) else: eplb_model_state.eplb_stats = EplbStats( # We copy the tensor to snapshot the global_expert_load_window @@ -1067,6 +1341,26 @@ def _allreduce_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor] offset += shape[0] return all_reduce_list + def _allgather_list(self, tensor_list: list[torch.Tensor]) -> list[torch.Tensor]: + """All-gather a list of 2D tensors, one extra leading rank dimension each. + + Only a placement policy that reweighs *which rank* is overloaded, not + just which logical expert is hot, needs this -- everything _allreduce_list + already serves is a plain sum, recoverable from this by summing dim 0. + Kept separate from _allreduce_list rather than folded into it: that + helper is shared with _sync_load_pass, a stats path with no reason to + pay for a gather it would immediately reduce anyway. + """ + ep_group = get_ep_group().device_group + world_size = ep_group.size() + gathered = [] + for tensor in tensor_list: + assert tensor.dim() == 2, "All tensors must be 2D." + out = tensor.new_empty((world_size, *tensor.shape)) + all_gather_into_tensor(out, tensor.contiguous(), group=ep_group) + gathered.append(out) + return gathered + def _sync_load_pass(self) -> list[torch.Tensor]: """Sync the expert load pass across all ranks for log stats. Doesn't update the expert load pass in eplb_model_state. @@ -1203,6 +1497,20 @@ class EplbLayerState: Reference to the parent :class:`EplbModelState`'s tensor list so the router can read the correct per-[u]batch unpadded token count. """ + moe_layer_idx: int | None = None + """ + Index of this layer among the model's MoE layers. + + The built-in mapping kernel is stateless and does not need it, but a + pluggable routing policy (see ``mlb_runtime``) keys its per-layer state by + it. Recording it here is free: ``set_layer_state`` already receives it. + + Note that the layer state deliberately does *not* carry + ``physical_to_logical_map``: threading it down would change + ``MixtureOfExperts.set_eplb_state``, a public interface implemented by + every MoE model. A policy that needs the physical->logical direction takes + the model-level tensor from :class:`EplbState` and slices it by this index. + """ def set_layer_state( self, @@ -1211,6 +1519,7 @@ def set_layer_state( logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor, ) -> None: + self.moe_layer_idx = moe_layer_idx self.expert_load_view = expert_load_view[moe_layer_idx] self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx] self.logical_replica_count = logical_replica_count[moe_layer_idx] @@ -1435,6 +1744,17 @@ def _move_to_workspace( layer=result.layer_idx, ) + # A policy holding placement-derived state has to be told this layer moved, + # or it keeps solving against the layout the layer used to have: the + # decision stays a valid replica of the right expert, because the dispatch + # kernel clamps to the live count and reads the committed map, but the split + # it computes is the one that balanced the *previous* placement. + # + # The notification is queued rather than delivered here. This runs on the + # async worker thread, and rebuilding a policy's per-layer state does GPU + # work; the main thread drains the queue at the top of the next forward. + _ASYNC_COMMITTED_LAYERS.add(result.layer_idx) + if result.layer_idx == model_state.model.num_moe_layers - 1: model_state.rebalanced = False diff --git a/vllm/distributed/eplb/mlb_runtime.py b/vllm/distributed/eplb/mlb_runtime.py new file mode 100644 index 000000000000..fd84018134f0 --- /dev/null +++ b/vllm/distributed/eplb/mlb_runtime.py @@ -0,0 +1,1993 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM-owned glue for MLB's L2 token-routing layer. + +vLLM's built-in replica choice is a Knuth hash of the token index, fused into +the same Triton kernel that records expert load +(``fused_moe/router/base_router.py``). When an MLB routing algorithm is +selected the kernel is kept: MLB supplies the per-expert replica shares it +solved for, and the kernel samples from them instead of hashing. Recording, +the record gate and the padding mask therefore stay where vLLM put them, with +nothing reimplemented alongside. + +Enable with ``VLLM_MLB_L2_ALGORITHM``, e.g. ``lplb``, ``dynamic``, ``static``. +Unset (the default) leaves vLLM's fused kernel in charge, so this module costs +nothing when it is not used. + +Two things are deliberately *not* supported and fail loudly rather than +silently degrading: + +* **Waterfill / shared-expert routing.** On vLLM's CUDA path the shared expert + is a per-rank replicated MLP that never enters EP dispatch, so there is no + destination-rank decision to write back. +* **DBO (dual-batch overlap).** MLB's ``RoutingRequest`` carries a single + token count and its policies keep one solver state per layer, so two + concurrently-executing micro-batches would clobber each other. +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, Any + +import numpy as np +import torch + +import vllm.envs as envs +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.distributed.eplb.eplb_state import EplbLayerState + +logger = init_logger(__name__) + +ALGORITHM_ENV = "VLLM_MLB_L2_ALGORITHM" + +_runtime: MlbRoutingRuntime | None = None + + +def mlb_l2_algorithm() -> str: + """Configured MLB routing expression, or "" when MLB routing is off.""" + return envs.VLLM_MLB_L2_ALGORITHM + + +class VllmRoutingCollectives: + """Expose vLLM's EP process group through MLB's generic transport API. + + Validated to be CUDA-graph capturable: vLLM routes EP collectives through + pynccl, and capture happens inside vLLM's own ``graph_capture()`` context. + """ + + @property + def rank(self) -> int: + from vllm.distributed.parallel_state import get_ep_group + + return get_ep_group().rank_in_group + + @property + def world_size(self) -> int: + from vllm.distributed.parallel_state import get_ep_group + + return get_ep_group().world_size + + def all_reduce_sum(self, payload: torch.Tensor) -> torch.Tensor: + from vllm.distributed.parallel_state import get_ep_group + + return get_ep_group().all_reduce(payload) + + +def _algorithm_places_with_ultraep(algorithm: str) -> bool: + """Whether this L2 expression puts UltraEP in charge of placement. + + ``RoutingPipeline`` reports that as ``placement_policy``, which is exactly + the question the fast-refresh path needs answered: it owns the placement + solve and the weight movement that follows it, and it must run for every + expression UltraEP places for -- ``ultraep`` and ``ultraep+waterfill`` + alike. An unparsable expression is not UltraEP's to claim; the caller + reports that separately. + """ + from moe_load_balancer.core.routing_pipeline import RoutingPipeline + + try: + return RoutingPipeline.from_value(algorithm).placement_policy == "ultraep" + except Exception: + return False + + +def _current_stage() -> str | None: + """Map vLLM's current batch onto MLB's routing stage. + + SGLang hands MLB a discrete ForwardMode (EXTEND / DECODE / IDLE / ...). + vLLM has no equivalent: with continuous batching and chunked prefill a + single batch routinely mixes prefill and decode tokens, so a clean + "prefill" stage does not exist. Only decode is provable -- every request + contributing exactly one token -- and everything else is reported as mixed + rather than guessed at, so a stage-gated policy never runs on a stage it + did not ask for. + """ + from vllm.forward_context import ( + get_forward_context, + is_forward_context_available, + ) + + if not is_forward_context_available(): + return None + ctx = get_forward_context() + descriptor = getattr(ctx, "batch_descriptor", None) + if descriptor is None: + return None + # `uniform_decode_across_dp` is vLLM's own cross-rank reduction of exactly + # the predicate below (see `coordinate_batch_across_dp`): it is True only + # when *every* DP rank is decoding this step. The descriptor describes this + # rank alone, and under DP that is not the same question -- a rank with no + # requests runs a dummy decode batch while its peers prefill, so reading the + # descriptor makes a rank answer "decode" on a step its peers are treating + # as prefill. Any caller gating a collective on the result would then split + # the EP group. Prefer the reduced answer whenever the framework supplies + # one; fall back to the local predicate only where there is no DP peer that + # could disagree. + if getattr(ctx, "uniform_decode_across_dp", False): + return "decode" + if getattr(ctx, "dp_metadata", None) is not None: + # Under DP the reduced flag above is the only admissible answer. It is + # False here, so at least one rank is not decoding and every rank has + # to say so -- including this one, even if its own batch is a textbook + # decode. Falling through to the descriptor would be the split this + # function exists to avoid. Read off the forward context rather than + # the config so the "is this DP" test itself comes from the same + # per-step object every rank sees. + return "mixed" + if ( + descriptor.uniform + and descriptor.num_reqs is not None + and descriptor.num_reqs > 0 + and descriptor.num_tokens == descriptor.num_reqs + ): + return "decode" + return "mixed" + + +def _dp_token_counts() -> torch.Tensor | None: + """Cross-rank-synchronized per-DP-rank token count for this step, if any. + + vLLM computes this once per step -- real batch or dummy -- and publishes + it to every rank before any MoE layer's forward runs + (``set_forward_context(num_tokens_across_dp=...)``, both in + ``execute_model`` and in ``_dummy_run``). Every rank sees the identical + array regardless of whether it personally has real requests this step, + which is exactly the property ``_ultraep_fast_refresh`` needs from its + own "is this worth a refresh" check: reading this instead of a per-rank + local token count means every rank reaches the same verdict even when + some ranks are running a dummy batch and never see real ``topk_ids`` at + all. ``None`` outside DP (single rank, or TP/PP-only) -- the caller falls + back to a local check there, where no dummy-batch peer exists to + disagree with. + """ + from vllm.forward_context import get_forward_context, is_forward_context_available + + if not is_forward_context_available(): + return None + # On `dp_metadata`, not on the context itself. `set_forward_context` takes a + # `num_tokens_across_dp` argument but does not keep it under that name -- + # it is consumed into DPMetadata, and ForwardContext has no such field. A + # getattr for it therefore always answered None, which silently sent this + # check down its single-rank fallback on every rank, which is exactly the + # cross-rank disagreement it exists to prevent. Read the tensor where it + # actually lives. It is a CPU tensor, so summing it costs no GPU sync. + dp_metadata = getattr(get_forward_context(), "dp_metadata", None) + if dp_metadata is None: + return None + return getattr(dp_metadata, "num_tokens_across_dp_cpu", None) + + +class MlbRoutingRuntime: + """Holds the MoELoadBalancer instance and the placement geometry.""" + + def __init__( + self, + algorithm: str, + *, + ep_size: int, + ep_rank: int, + num_logical_experts: int, + num_physical_experts: int, + physical_to_logical_map: torch.Tensor, + balancer: Any | None = None, + expert_weights: Any | None = None, + expert_buffer: Any | None = None, + communicator: Any | None = None, + ) -> None: + from moe_load_balancer import MoELoadBalancer + + self.algorithm = algorithm + self.ep_size = ep_size + self.ep_rank = ep_rank + self.num_logical_experts = num_logical_experts + self.num_physical_experts = num_physical_experts + # Model-level [layers, num_physical] map. Kept here rather than on the + # layer state so that MixtureOfExperts.set_eplb_state -- a public model + # interface -- does not have to change. + self.physical_to_logical_map = physical_to_logical_map + # A logical expert can hold at most one slot plus every redundant one, + # so this bounds the replica count without reading the map off device. + self._max_replicas = num_physical_experts - num_logical_experts + 1 + self._default_replicas: torch.Tensor | None = None + # Until a placement commits, assume there is a choice to make; the + # first commit replaces this with what the placement actually offers. + self._placement_offers_replica_choice: bool = True + # The same table in the layout the fused kernel reads: one candidate + # column per logical expert, and a replica count of 1 so the kernel's + # `hash % count` lands on it. Built once per placement. + self._fixed_map: torch.Tensor | None = None + self._fixed_counts: torch.Tensor | None = None + # [num_physical_experts] — slot i belongs to rank (i // experts_per_rank). + # Passed to every PlacementSnapshot so LPLB can incorporate cross-GPU + # transfer cost; SGLang always provides this, we build it from topology. + from moe_load_balancer.adapters.vllm import build_physical_to_rank_map + + self._physical_to_rank_map = build_physical_to_rank_map( + num_physical_experts, + ep_size, + device=physical_to_logical_map.device, + ) + # Injected when an engine-scoped integration owns the balancer, so L1 + # and L2 are served by one instance rather than two that cannot see + # each other. + self._mlb: Any = ( + balancer + if balancer is not None + else MoELoadBalancer.from_algorithm( + algorithm, + ep_size=ep_size, + source_rank=ep_rank, + experts_per_rank=num_physical_experts // ep_size, + collectives=VllmRoutingCollectives(), + ) + ) + # MLB declares what the framework has to prepare for the selected + # policy, so none of the work below is done unconditionally: `lplb` + # needs placement state refreshed on commit, `static` needs a per-rank + # dispatch table, and only a pipeline with a post-TopK policy wants the + # routing boundary at all. + caps = self._mlb.routing_capabilities + self.requires_post_topk_routing = caps.requires_post_topk_routing + self.requires_placement_state = caps.requires_placement_state + self.requires_rank_dispatch_map = caps.requires_rank_dispatch_map + # Whether a policy in this pipeline picks the shared expert's rank. + # The router still has to have somewhere to put that answer -- it only + # asks when its layer was built with the EP shared-expert layout + # (VLLM_FUSE_SHARED_EXPERTS). getattr, because an older MLB has no + # such field. + self.routes_shared_expert = getattr(caps, "routes_shared_expert", False) + # A policy whose answer depends only on the committed placement does + # not need the per-forward call at all: bake the answer into a map the + # fused kernel already knows how to read, and the boundary disappears. + # getattr, because an older MLB has no such field. + self.dispatch_fixed_by_placement = getattr( + caps, "dispatch_fixed_by_placement", False + ) + # These two cannot both hold. Routing a shared expert means choosing a + # rank per batch from that batch's load; "fixed by placement" means the + # answer is a function of the committed placement alone and the + # per-forward call can be skipped. A pipeline claiming both gets its + # shared-expert decisions computed and then thrown away, and the stale + # ones from the previous layer read in their place -- with no symptom + # beyond a quality regression. Checked here rather than trusted: the + # default above protects against an MLB too old to have the field, but + # not against one old enough to have it and set it wrong. + if self.routes_shared_expert and self.dispatch_fixed_by_placement: + raise ValueError( + f"MLB routing pipeline {algorithm!r} reports both " + "routes_shared_expert and dispatch_fixed_by_placement. A " + "shared-expert rank is a per-batch decision, so it cannot be " + "baked in at placement-commit time; one of the two is wrong. " + "An MLB predating the fix to dispatch_fixed_by_placement " + "reports this for any pipeline containing waterfill." + ) + # Declared by the pipeline rather than inferred from "has a post-TopK + # policy": a replica policy can route from placement alone, and + # gathering the EP-wide load for it costs a collective per step that is + # then discarded. + self.consumes_global_logical_count = getattr( + caps, "consumes_global_logical_count", caps.requires_post_topk_routing + ) + self.max_input_staleness_steps = getattr(caps, "max_input_staleness_steps", 0) + self.supports_concurrent_microbatches = getattr( + caps, "supports_concurrent_microbatches", False + ) + self.graph_stability = getattr(caps, "graph_stability", "stable") + logger.info( + "MLB L2 routing enabled (algorithm=%s, ep_size=%d, logical=%d, " + "physical=%d, caps: post_topk=%s placement_state=%s " + "rank_dispatch_map=%s)", + algorithm, + ep_size, + num_logical_experts, + num_physical_experts, + self.requires_post_topk_routing, + self.requires_placement_state, + self.requires_rank_dispatch_map, + ) + + # Graph-safe global count cache for LPLB. + # + # LPLB's LP solve needs global_logical_count: how many tokens each + # logical expert received across all EP ranks. The naive path inside + # MLB does count_logical_experts(topk_ids) + EP all_reduce per layer + # per forward, which is 40 NCCL collectives per step -- and NCCL + # cannot be captured in a CUDA graph. + # + # The fix: split "compute local count" (CUDA kernel, graph-safe) from + # "all_reduce" (outside the graph, once per step). + # + # _logical_count_local[layer, expert]: local counts accumulated per-layer + # inside the graph by count_logical_experts; reset each step. + # _logical_count_global[layer, expert]: all-reduced result of the previous + # step, stable GPU address, read by LP solve inside the graph. + # + # The LP solve therefore uses counts stale by one step, which is an + # excellent approximation for decode (distribution barely changes) and + # acceptable for prefill (first step uses zeros → hash routing, then + # converges). + num_moe_layers = physical_to_logical_map.shape[0] + dev = physical_to_logical_map.device + # Gated on redundancy as well as on the policy. With no redundant + # experts every logical has exactly one replica, the LP is degenerate, + # and LPLB answers with identity dispatch without ever reading the + # count -- so keeping these buffers would buy a per-layer count kernel + # plus one collective per step for a result that is thrown away. + # Leaving them None also makes MLB decline to gather the count itself + # (needs_global_logical_count), so nothing downstream pays for it. + # MLB_KEEP_ZERO_REDUNDANCY_COUNTS=1 restores the ungated behaviour so the + # cost of that discarded work can be measured against this, rather than + # only argued for. + keep_degenerate = envs.MLB_KEEP_ZERO_REDUNDANCY_COUNTS + if self.consumes_global_logical_count and ( + self._max_replicas > 1 or keep_degenerate + ): + self._logical_count_local: torch.Tensor | None = torch.zeros( + num_moe_layers, num_logical_experts, dtype=torch.float32, device=dev + ) + self._logical_count_global: torch.Tensor | None = torch.zeros( + num_moe_layers, num_logical_experts, dtype=torch.float32, device=dev + ) + else: + self._logical_count_local = None + self._logical_count_global = None + # True once finalize_step_counts has run at least once (first step uses + # zeros → hash routing fallback via selection is None). + self._logical_count_ready = False + # Reduced by the dummy path so the collective stays symmetric. + self._count_scratch: torch.Tensor | None = None + + # MLB_FRESH_COUNTS=1 withholds the pre-computed count so MLB gathers it + # itself, per layer per forward. That is what the SGLang adapter does and + # what LPLB's published numbers were measured against; the count here is + # one step stale so the collective can live outside the CUDA graph. + # Staleness is a good approximation in decode, where the routing + # distribution moves slowly, and a worse one in chunked prefill, where + # consecutive chunks can differ sharply -- so a measured LPLB regression + # cannot be attributed to the algorithm without checking this. Read once: + # replica_shares runs per layer per forward. + self._fresh_counts = envs.MLB_FRESH_COUNTS + + # Diagnostic capture of the LP's inputs and output. Off unless a + # directory is named; bounded so a long run cannot fill the disk. + self._dump_lp_dir = envs.MLB_DUMP_LP + self._dump_lp_left = envs.MLB_DUMP_LP_N + # Skip the opening calls: the first steps legitimately carry zeros + # while the one-step-stale pipeline fills, and sampling only those + # would mistake a warm-up transient for steady state. + self._dump_lp_skip = envs.MLB_DUMP_LP_SKIP + + # MLB_TIME_L2= times n route_tokens calls with a device sync on each + # side. The sync makes the number meaningful and the measurement + # invasive, so it is off unless asked for -- diagnostic only. + # Set by replica_shares when a policy answers with ids; read by + # resolve_routing in the same call. + self._last_physical_ids: torch.Tensor | None = None + # Same one-call handoff for a shared-expert decision. Separate from the + # ids because the router consumes them at different points: the ids + # feed the EPLB mapping kernel, the rank feeds the append that happens + # after it. + self._last_shared_expert_rank: torch.Tensor | None = None + # Set by plan_placement() after an UltraEP L1 solve, read by + # _snapshot() on every forward until the next solve replaces them. + # None for every other L1 policy, and before the first solve. + # + # The candidate table and quota must come from the same solve: the + # quota indexes replicas by column against MLB's own candidate + # ordering, which is not guaranteed to be the same table -- same + # width, even -- as vLLM's own logical_to_physical_map, independently + # rebuilt from physical_to_logical_map through compute_logical_maps. + # So route with MLB's own table for this policy, not vLLM's. + self._ultraep_rank_quota_prefix: torch.Tensor | None = None + self._ultraep_logical_to_physical: torch.Tensor | None = None + self._ultraep_replica_counts: torch.Tensor | None = None + self._time_l2 = envs.MLB_TIME_L2 + # MLB_TIME_REFRESH= times n fast refreshes -- the placement solve and + # the weight commit separately -- with a device sync on each, then + # reports and disarms. Same shape as MLB_TIME_L2 above. + self._time_refresh = envs.MLB_TIME_REFRESH + # How UltraEP's fast refresh lands re-solved replica weight. Upstream + # UltraEP does this with a one-sided weight_sync (masters stay put, + # replica slots are refilled from the owning master); the movers below + # are that idea on vLLM's own tensors, best first: + # symm -- one-sided NVLink puts through torch symmetric memory, + # completed with pairwise signals (_commit_symm) + # direct -- pairwise send/recv over the EPLB communicator, no + # staging buffer (_commit_direct) + # generic -- vLLM's periodic-rearrangement mover (move_to_buffer), + # which also handles the one case the others do not: the + # first refresh after a checkpointed placement, when the + # masters themselves move + # MLB_ULTRAEP_MOVER (default direct). `auto` and `symm` try symmetric + # memory first, but only after every EP rank has agreed it can (see + # _symm_setup): a per-rank try/except around the collective rendezvous + # let some ranks fall back while the others waited in it, and the + # server deadlocked on its first refresh whenever the EP group spanned + # hosts without a multi-node NVLink fabric. Measured on a Blackwell + # multi-node NVLink EP16 deployment (DeepSeek-V3, one allocation, arms + # back to back) symm is within 1-2% + # of direct, so direct is the default and symm an explicit choice. + # Every mover leaves a layer whose masters move to generic. + self._ultraep_mover = envs.MLB_ULTRAEP_MOVER + self._ultraep_mover_warned = False + self._symm_hdl: Any = None + self._symm_buf: torch.Tensor | None = None + self._symm_layout: list[tuple[int, int, torch.dtype, tuple[int, ...]]] = [] + self._symm_expert_nbytes = 0 + # MLB_ULTRAEP_LAGGED_APPLY (default 0): with 1 the plan solved on this + # forward lands (routing snapshot + weights) on the layer's next + # forward, which keeps the plan's device->host copy off the critical + # path -- no per-layer GPU sync -- at the price of one step of + # staleness: the first batch of a burst runs on the old placement and + # the second on a plan fitted to that first batch. Measured on + # DeepSeek-V3 EP16 (same allocation, arms back to back): at interval 8 + # it changes nothing (Blackwell, 139.0k vs 139.0k tok/s, mean imbalance + # 1.30 vs 1.29); at interval 1 it buys +3% (HybridEP) to +8% (DeepEP + # normal) on Blackwell for +0.05 of steady imbalance, and on Hopper it + # buys nothing and costs the same balance. Off by default; set 1 for + # interval-1 deployments with Blackwell-class step times. + self._ultraep_lagged = envs.MLB_ULTRAEP_LAGGED_APPLY + self._ultraep_pending: dict[int, tuple[Any, ...]] = {} + self._ultraep_pending_host: dict[int, torch.Tensor] = {} + self._solve_us: list[float] = [] + self._commit_us: list[tuple[int, float]] = [] + self._commit_host_us: list[float] = [] + self._l2_us: list[float] = [] + + # Real, traffic-driven placement refresh with real weight transfer, + # additive to plan_placement()'s slow, wall-clock-cadence path above: + # that path still owns this model's *initial* placement (the buffers + # this refresh later writes per-layer slices into do not exist + # before its first solve -- see the None-guard in + # _ultraep_fast_refresh), this only adds much more frequent updates + # on top of it. Making memory match a placement the solve decided is + # the framework adapter's job and stays here: the solve produces an + # explicit before/after map pair and hands it to EplbState's own P2P + # mover (_commit_placement_weights). MLB supplies the placement + # algorithm and nothing else -- no policy brings its own weight + # mover, and no second component re-derives placement from raw + # routing to agree with this one by construction. + self._ultraep_expert_weights: Any = None + self._ultraep_expert_buffer: Any = None + self._ultraep_communicator: Any = None + self._ultraep_num_local_physical: int | None = None + # Layers whose MLB-side placement tables have actually been written by + # a refresh. Allocation of the quota tensor is not the same event: the + # bootstrap solve allocates it whole, layer contents arrive one refresh + # at a time. See the guard in _snapshot. + self._ultraep_committed_layers: set[int] = set() + self._ultraep_refresh_gate: Any = None + # Declared here (not just where it's actually populated, in + # _init_ultraep_fast_refresh below) because finish_pending_ + # ultraep_transfer() is called unconditionally from vLLM's own FFN + # forward for every MoE layer regardless of which L1 policy is + # active -- an L1 policy other than ultraep must see this attribute + # exist (as None) rather than raise AttributeError. + # 512 is the reference inference integration's default (SGLang's + # `moe_balance_refresh_min_tokens`). The solve reads only the batch in + # front of it and the answer is then held for a whole refresh interval, + # so re-solving on a batch too small to be representative fits the + # placement to noise and keeps it there. It bites only where batches + # are genuinely small: a prefill-shaped step is orders of magnitude + # above either threshold. + self._ultraep_min_representative_tokens = envs.MLB_ULTRAEP_REFRESH_MIN_TOKENS + # Parsed, not compared as a string. `algorithm` is the whole L2 + # expression, so an equality test here silently misses every + # composition that merely *contains* ultraep -- `ultraep+waterfill` + # among them. The failure was not a graceful degradation: the refresh + # never ran, while `_snapshot` still saw an allocated quota and routed + # through it, so tokens went to slots holding other experts. Everywhere + # else in this file already asks the parsed pipeline (see + # l2_pipeline_capabilities / l2_inapplicable_reason); this was the one + # place that did not. + if _algorithm_places_with_ultraep(algorithm) and expert_weights is not None: + self._init_ultraep_fast_refresh(expert_weights, expert_buffer, communicator) + + def _init_ultraep_fast_refresh( + self, + expert_weights: Any, + expert_buffer: Any, + communicator: Any, + ) -> None: + """One-time setup for the traffic-driven placement refresh. + + Holds on to the model's own expert weight tensors plus ``EplbState``'s + staging buffer and P2P communicator -- the same three things vLLM's own + periodic rearrangement moves weight with. This policy re-plans placement + far more often than that cadence, but it re-plans it into the same + memory using the same machinery; it brings no mover of its own. + + Weights are taken eagerly rather than on first use: vLLM hands this + runtime ``model.expert_weights`` at the same call site that constructs + it, so there is no "not loaded yet" window to defer past. + + Assumes the per-layer weight tensors are each shaped + ``[num_local_physical_experts, ...]`` -- true for the unquantized case + this was verified against. A model whose ``expert_weights`` also carries + separate quantization-scale tensors would need those moved too; not + attempted here. + """ + if expert_buffer is None or communicator is None: + logger.warning_once( + "ultraep fast refresh needs EplbState's expert buffer and " + "communicator to move weight; neither was supplied, so the " + "policy stays on the slow EplbState.rearrange() cadence only." + ) + return + + from moe_load_balancer.policies.l1 import RefreshGate + + self._ultraep_expert_weights = expert_weights + self._ultraep_expert_buffer = expert_buffer + self._ultraep_communicator = communicator + self._ultraep_num_local_physical = self.num_physical_experts // self.ep_size + + interval = envs.MLB_ULTRAEP_REFRESH_INTERVAL + self._ultraep_refresh_gate = RefreshGate(interval) + logger.info( + "UltraEP fast refresh enabled (interval=%d representative " + "batches, min_tokens=%d, num_layers=%d, " + "num_local_physical=%d, mover=%s, lagged_apply=%s)", + interval, + self._ultraep_min_representative_tokens, + len(expert_weights), + self._ultraep_num_local_physical, + self._ultraep_mover, + self._ultraep_lagged, + ) + + def finalize_step_counts(self) -> None: + """All-reduce per-layer local counts and update the stable LP input buffer. + + Call at the START of each forward pass (from EplbState.prepare_forward). + By the time this runs, the previous step's count_logical_experts results + are already in _logical_count_local (written per-layer inside the graph or + in the eager forward). + + One EP collective for all 40 layers combined replaces the 40 per-layer + collectives that MLB's _global_logical_count would otherwise issue. + Running outside the graph means NCCL is never captured -- graphs only + see the LP solve kernels reading from _logical_count_global. + """ + if self._logical_count_local is None or self._logical_count_global is None: + return + # Never run inside a CUDA graph capture stream. prepare_forward is + # normally called outside graph context, but guard explicitly. + if torch.cuda.is_current_stream_capturing(): + return + from vllm.distributed import get_ep_group + + ep_group = get_ep_group() + ep_group.all_reduce(self._logical_count_local) + self._logical_count_global.copy_(self._logical_count_local) + self._logical_count_local.zero_() + self._logical_count_ready = True + + def match_step_counts_collective(self) -> None: + """Issue the count collective without touching the count pipeline. + + A rank running a dummy batch has to take part in the collective, or the + EP group falls out of order and its peers wait in dispatch until DeepEP + times out. It must not run `finalize_step_counts`, though: that copies + the local buffer into the global one and then clears the local. Called a + second time within a step -- which is exactly what a dummy forward + beside a real one does -- it copies the freshly cleared buffer over the + counts that were just published, and the LP then solves against an + all-zero load for the rest of the run. + + So the dummy path reduces a scratch tensor of the same shape instead: + same collective, same size, no effect on what the solve reads. + """ + if self._logical_count_local is None: + return + if torch.cuda.is_current_stream_capturing(): + return + from vllm.distributed import get_ep_group + + if self._count_scratch is None: + self._count_scratch = torch.zeros_like(self._logical_count_local) + get_ep_group().all_reduce(self._count_scratch) + + def set_physical_to_logical_map(self, mapping: torch.Tensor) -> None: + self.physical_to_logical_map = mapping + + def _snapshot( + self, + layer_state: EplbLayerState, + layer_id: int, + ) -> Any: + from moe_load_balancer.adapters.vllm import ( + collapse_candidates_to_local, + to_placement_snapshot, + ) + + # Narrowing an expert's candidates to its local replica is a *routing + # decision* ("never pay for a cross-GPU hop"), not a view of the + # placement, so it belongs to the policy that asked for it rather than + # to every policy. SGLang draws the line the same way: the collapse + # lives in `logical_to_rank_dispatch_physical_map`, which it only builds + # when the pipeline reports `requires_rank_dispatch_map` -- while + # `init_by_eplb`, the path taken on every rebalance, hands the policy + # the full global list. + # + # Applying it unconditionally answers the question LPLB exists to ask: + # with a local replica always winning, a rank holding a copy has nothing + # left to solve, so the LP is skipped outright and the rest see only the + # experts they do not host. + # vLLM pads the candidate map to MAX_EXPERT_REDUNDANCY + 1 (1024) + # columns whatever the configured redundancy, and MLB answers with a + # table the same width as the candidates it was given. The kernel scans + # that table with a compile-time loop, so handing over the padded width + # would have made a share-table answer impossible to compile. Every + # column past `_max_replicas` is padding on both sides, so trimming to + # it changes no decision. + assert layer_state.logical_to_physical_map is not None + assert layer_state.logical_replica_count is not None + candidates = layer_state.logical_to_physical_map[:, : self._max_replicas] + counts = layer_state.logical_replica_count + # `requires_rank_dispatch_map` is true only for `static`, and static no + # longer calls this method at all -- its answer is resolved once per + # placement instead (see `dispatch_fixed_by_placement`). Kept rather + # than deleted: it is what a future rank-dispatch-map policy would need, + # and removing it now would be removing untested surface, not dead code. + if self.requires_rank_dispatch_map: + candidates, counts = collapse_candidates_to_local( + candidates, + counts, + ep_rank=self.ep_rank, + num_local_physical_experts=(self.num_physical_experts // self.ep_size), + ) + + defaults = self._default_replicas + quota = self._ultraep_rank_quota_prefix + # Allocated is not the same as committed. The bootstrap solve allocates + # the whole quota tensor before any refresh has run, and what it leaves + # in it describes a layout vLLM never applied -- all-zero quota, + # single-replica counts, candidates pointing at slots that hold other + # experts. Routing through that is worse than not routing at all: it + # sends tokens to the wrong expert instead of falling back. So the + # MLB-side tables are used only for layers a refresh has actually + # written; every other layer takes the framework's own candidates, which + # is the same path a non-UltraEP policy takes and is always consistent + # with physical_to_logical_map. + if quota is not None and layer_id in self._ultraep_committed_layers: + # The quota's replica columns are only meaningful against the + # candidate table MLB solved them from -- not vLLM's own + # candidates, independently rebuilt from physical_to_logical_map + # and not guaranteed to share its width, let alone its ordering. + assert self._ultraep_logical_to_physical is not None + assert self._ultraep_replica_counts is not None + candidates = self._ultraep_logical_to_physical[layer_id] + counts = self._ultraep_replica_counts[layer_id] + return to_placement_snapshot( + layer_state, + layer_id, + physical_to_logical_map=self.physical_to_logical_map[layer_id], + num_logical_experts=self.num_logical_experts, + num_physical_experts=self.num_physical_experts, + ep_size=self.ep_size, + candidates=candidates, + counts=counts, + default_physical_for_logical=( + None if defaults is None else defaults[layer_id] + ), + physical_to_rank_map=self._physical_to_rank_map, + rank_quota_prefix=None if quota is None else quota[layer_id], + ) + + def _rebuild_default_replicas(self) -> None: + """Synthesize the per-rank default replica table vLLM does not keep. + + Only `static` replica routing reads it (MLB reports that through + ``requires_rank_dispatch_map``), and it is recomputed only when a + placement is committed -- never per forward. + """ + if not self.requires_rank_dispatch_map: + self._default_replicas = None + self._fixed_map = None + self._fixed_counts = None + return + + from moe_load_balancer.adapters.vllm import nearest_replica_table + + self._default_replicas = torch.stack( + [ + nearest_replica_table( + self._logical_to_physical_map[layer], + self._logical_replica_count[layer], + ep_rank=self.ep_rank, + num_local_physical_experts=( + self.num_physical_experts // self.ep_size + ), + # Passed rather than inferred: inference reads the largest + # physical id present, which undercounts whenever the top + # ranks hold no replica, and this rank can then fall outside + # the table it just built. + ep_size=self.ep_size, + ) + for layer in range(self._logical_to_physical_map.shape[0]) + ] + ) + + if not self.dispatch_fixed_by_placement: + self._fixed_map = None + self._fixed_counts = None + return + + # [layers, num_logical] -> [layers, num_logical, 1]. One column, so the + # kernel's gather strides by 1 instead of by the candidate map's padded + # width of MAX_EXPERT_REDUNDANCY + 1, and reads a table small enough to + # stay cached. + self._fixed_map = ( + self._default_replicas.to(self._logical_to_physical_map.dtype) + .unsqueeze(-1) + .contiguous() + ) + # Every expert has exactly one candidate in that layout, so `hash % + # count` is 0 for every token and the kernel returns column 0 -- which + # is the replica this policy chose. Same answer as calling the policy + # per forward, without the call. + self._fixed_counts = torch.ones( + self._default_replicas.shape[1], + dtype=self._logical_replica_count.dtype, + device=self._default_replicas.device, + ).contiguous() + + def fixed_dispatch_maps( + self, layer_id: int + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """The placement-resolved map and counts for a layer, if there is one. + + Returned in place of calling the routing pipeline: for a policy that + decides from placement alone, the per-forward call produced this same + answer at a measured 426 us per layer against a policy-free path that + costs nothing. + """ + if self._fixed_map is None or self._fixed_counts is None: + return None + return self._fixed_map[layer_id], self._fixed_counts + + def register_logical_maps( + self, logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor + ) -> None: + """Keep the model-level maps so snapshots can be built without a layer. + + vLLM commits a rearrangement with in-place ``copy_`` into these + tensors, so holding them here observes every future placement without + re-registration. + """ + self._logical_to_physical_map = logical_to_physical_map + self._logical_replica_count = logical_replica_count + # Does this placement give any expert a second copy? If not, every + # replica policy -- static, dynamic, dynamic_random, lplb alike -- can + # only return the identity, and consulting one is pure cost. + # + # Read once here rather than per forward: `.any()` on a CUDA tensor is + # a device sync, which is cheap at a rearrangement and ruinous inside a + # forward. Kept as a Python bool for that reason. + # + # This belongs to the placement, not to a policy. LPLB happens to + # short-circuit itself (`num_red_log == 0` leaves its LP with nothing + # to solve), static does not -- it consults its nearest-replica table + # whatever the redundancy is. Leaving each policy to notice the + # degenerate case independently is how that asymmetry arose, and it is + # exactly what made a red0 row -- where by construction no policy can + # differ -- report a 2.6% spread between them. + self._placement_offers_replica_choice = bool( + (logical_replica_count > 1).any().item() + ) + self._rebuild_default_replicas() + + def announce_initial_placement(self) -> None: + """Tell MLB about the placement the model just loaded with. + + The contract is "after initial weight loading **and** after each + committed update"; SGLang does the first half in + ``ModelRunner._prepare_moe_topk``. Skipping it leaves a policy to build + its per-layer state lazily inside the first forward, from whatever + placement that forward happens to see. + """ + if not self.requires_placement_state: + return + num_layers = self._logical_to_physical_map.shape[0] + self._commit_layers(range(num_layers)) + logger.info("MLB: announced initial placement for %d layers", num_layers) + + def _commit_layers(self, layer_ids) -> None: + """Hand MLB the committed placement of the given layers.""" + from moe_load_balancer.core.types import PlacementSnapshot + + for layer_id in layer_ids: + self._mlb.on_placement_committed( + PlacementSnapshot( + layer_id=layer_id, + num_logical_experts=self.num_logical_experts, + num_physical_experts=self.num_physical_experts, + ep_size=self.ep_size, + num_local_physical_experts=( + self.num_physical_experts // self.ep_size + ), + physical_to_logical_map=self.physical_to_logical_map[layer_id], + # Same trim as `_snapshot`, and it has to be the same: the + # policy sizes its per-layer state from whichever snapshot + # reaches it first, and this one does -- it runs for every + # layer at startup. Committing the padded width here left + # the solver emitting 1024-wide tables that the mapping + # kernel then could not compile. + logical_to_physical_candidates=self._logical_to_physical_map[ + layer_id + ][:, : self._max_replicas], + logical_to_physical_count=self._logical_replica_count[layer_id], + default_physical_for_logical=( + None + if self._default_replicas is None + else self._default_replicas[layer_id] + ), + ) + ) + + def on_placement_committed( + self, changed_layer_ids: list[int] | None = None + ) -> None: + """Refresh policy state for layers whose placement just changed. + + Must be called after vLLM has moved weights and updated its live + metadata -- never from an uncommitted plan. + + ``changed_layer_ids`` matters for cost, not just tidiness. A policy + whose per-layer state layout changes has to rebuild it, and for `lplb` + that means a JIT build plus warmup -- about 14.5 s per layer measured + here. Refreshing all 40 layers stalled the worker for ~102 s, long + enough for the peer's collective to time out and kill it. SGLang + avoids this by passing the framework's own ``update_layer_ids`` + (``ModelRunner._notify_mlb_placement_committed``); vLLM does not track + them, so Glue diffs the maps before committing and passes the result. + """ + # The nearest-replica table is *this* side's derived state, not the + # policy's, so it is refreshed on every commit. Gating it on + # requires_placement_state left a policy that keeps no state of its + # own -- static, which reads this table and nothing else -- routing + # by the placement the run started with. After a rearrangement most + # defaults are no longer in their expert's candidate list, the share + # table falls back to column 0 for them, and the traffic those + # replicas exist to spread lands on one of them. + + self._rebuild_default_replicas() + + if not self.requires_placement_state: + return + if changed_layer_ids is None: + changed_layer_ids = list(range(self._logical_to_physical_map.shape[0])) + if not changed_layer_ids: + return + self._commit_layers(changed_layer_ids) + logger.info( + "MLB: refreshed placement state for %d/%d layers", + len(changed_layer_ids), + self._logical_to_physical_map.shape[0], + ) + + def resolve_routing( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + layer_state: EplbLayerState, + num_unpadded_tokens: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Ask MLB to route. + + Returns ``(replica_shares, physical_ids)``. Every replica policy + resolves its own ids now, so the first element is always ``None`` -- + kept in the return shape because the caller's fused kernel still has a + share-table mode (``HAS_REPLICA_PROB``) it did not stop supporting, + even though nothing on the MLB side asks for it any more. + """ + shares = self.replica_shares( + topk_ids, topk_weights, layer_state, num_unpadded_tokens + ) + ids = None if shares is not None else self._last_physical_ids + return shares, ids + + def resolve_shared_expert_rank( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + layer_state: EplbLayerState, + num_unpadded_tokens: torch.Tensor | None, + ) -> torch.Tensor | None: + """The per-token home rank for the shared expert, or None. + + Reads what the pipeline just decided rather than running it again: + ``resolve_routing`` has already crossed the routing boundary for this + layer, and Waterfill's answer came back in the same decision. Calling + MLB twice would double the only per-layer cost this integration has. + + None means "no policy answered", which the caller reads as "this + rank" -- the placement the replicated MLP had. + """ + del topk_ids, topk_weights, layer_state, num_unpadded_tokens + return self._last_shared_expert_rank + + def _ultraep_fast_refresh( + self, + layer_id: int, + topk_ids: torch.Tensor, + num_unpadded_tokens: torch.Tensor | None, + ) -> None: + """Real, traffic-driven placement refresh + weight transfer, one layer. + + Not CUDA-graph-safe: issues real collectives (all_gather, a + distributed weight transfer, a barrier) whenever it actually + refreshes, so it must never run under capture -- same reason + ``finalize_step_counts`` guards on the same check. + + Gated on ``_ultraep_rank_quota_prefix`` already being allocated: + those whole-model buffers are only created by ``plan_placement()``'s + slow-path solve (see ``mlb_runtime.py``'s ``plan_placement``), which + still owns this model's *initial* placement. This refresh writes + only a per-layer slice into buffers that solve already allocated; it + does nothing before that first solve has run. + + Writes its result directly into the same whole-model buffers + ``_snapshot()`` reads every forward (``physical_to_logical_map``, + ``_ultraep_logical_to_physical``, ``_ultraep_replica_counts``, + ``_ultraep_rank_quota_prefix``) rather than a separate pending store: + ``replica_shares()`` calls ``_snapshot()`` again immediately after + this returns, on its way to dispatch, so there is no window where a + separate "pending" representation could go stale relative to what + this just wrote. + + The placement solve is this file's own (MLB's L1, through + ``self._mlb.plan_placement``); physically moving weight data to match + it is ``_commit_placement_weights`` below, over ``EplbState``'s own P2P + mover. The bridge between the two is explicit and is the solve's actual + output: the before/after physical-to-logical map pair. No second + component re-derives placement from raw routing to arrive at the same + answer independently, so there is no shared-determinism assumption to + hold -- the mover cannot disagree with the solve, because it is handed + the solve. + """ + if torch.cuda.is_current_stream_capturing(): + return + if self._ultraep_rank_quota_prefix is None: + return + if self._ultraep_lagged: + # The plan solved on this layer's previous forward lands now + # (routing snapshot, then weights), before this forward's routing. + self._ultraep_apply_pending(layer_id) + if _current_stage() == "decode": + # UltraEP is explicitly a prefill-time algorithm (the paper's own + # scoping; SGLang's reference should_refresh gates the same way, + # rejecting only its provable "decode" and accepting everything + # else, including its own "mixed"). A decode-only batch must not + # count toward this layer's refresh interval at all -- it never + # reaches is_due(), the same way SGLang's gate never touches its + # own batch counter for a rejected stage. + # + # Everything past this point issues collectives, so this verdict + # has to be identical on every EP rank or the ranks that continue + # will wait in the all_gather below for ranks that returned here. + # That is what `_current_stage()` reads the DP-reduced + # `uniform_decode_across_dp` for rather than this rank's own batch + # descriptor: a rank with no requests runs `execute_dummy_batch()` + # -- `_dummy_run(uniform_decode=True)`, locally a perfect "decode" + # -- while its peers run real prefill. Gating on the local view + # splits the group, and permanently desynchronises is_due()'s + # per-layer counter as well, since a rank that returns here never + # advances it. + return + + # is_due() is pure Python counter bookkeeping -- no GPU access, never + # syncs -- and is always cheap to call. Computing `representative` + # below is not free either: the fallback source is a device tensor and + # its `.item()` is a real CPU-GPU sync (the DP source is a CPU tensor + # and is not), + # and paying it on every one of interval-1-out-of-interval calls + # where should_refresh could not possibly say yes anyway is exactly + # the class of per-forward cost this file goes to real lengths to + # avoid elsewhere (see finalize_step_counts's whole reason for + # existing). Check is_due first; only compute representative and + # call should_refresh when it says a refresh could happen this call. + if not self._ultraep_refresh_gate.is_due(layer_id): + return + + # Under DP, this rank's own num_unpadded_tokens says nothing about + # what its peers are doing this step -- a rank running a dummy batch + # has no local token count at all, and would otherwise have to guess. + # dp_counts is the same array on every rank (vLLM's own DP + # coordination gathers it before any MoE layer runs), so summing it + # gives every rank -- real or dummy -- an identical verdict. Only + # fall back to the local, per-rank count outside DP, where no dummy + # peer exists that could disagree with it. + dp_counts = _dp_token_counts() + if dp_counts is not None: + representative = ( + int(dp_counts.sum().item()) >= self._ultraep_min_representative_tokens + ) + else: + representative = ( + num_unpadded_tokens is not None + and int(num_unpadded_tokens.item()) + >= self._ultraep_min_representative_tokens + ) + if not self._ultraep_refresh_gate.should_refresh( + layer_id, representative=representative + ): + return + + from moe_load_balancer.adapters.vllm import to_placement_request + from moe_load_balancer.kernels.ops.counting import count_logical_experts + + from vllm.distributed import get_ep_group + + ep_group = get_ep_group().device_group + local_count = count_logical_experts( + topk_ids, self.num_logical_experts, dtype=torch.int32 + ) + per_rank_count = local_count.new_empty((self.ep_size, self.num_logical_experts)) + torch.distributed.all_gather_into_tensor( + per_rank_count, local_count.contiguous(), group=ep_group + ) + + request = to_placement_request( + per_rank_count[:, None, :], + num_replicas=self.num_physical_experts, + num_ranks=self.ep_size, + num_groups=1, + num_nodes=1, + algorithm="ultraep", + ep_rank=self.ep_rank, + ) + if self._time_refresh: + torch.accelerator.synchronize() + _ts = time.perf_counter() + plan = self._mlb.plan_placement(request) + if self._time_refresh: + torch.accelerator.synchronize() + self._solve_us.append((time.perf_counter() - _ts) * 1e6) + if self._ultraep_lagged: + self._ultraep_stage_pending(layer_id, plan) + return + + # The request carried one synthetic layer (dim 0 of per_rank_count), + # so the plan's own layer index is always 0 regardless of layer_id -- + # layer_id only selects where in *our* whole-model buffers this one + # layer's slice of the plan lands. + # + # Where each physical slot's weight lives now, and where the solve wants + # it to live. Those two rows *are* the transfer plan handed to the mover + # below -- nothing re-derives placement from raw routing a second time, + # so the mover cannot disagree with the solve. Stacked into one tensor + # because the mover reads them on the host: one device-to-host copy for + # the pair rather than one each, and the sync is visible here, at the + # only place in this path that pays it. + live = self.physical_to_logical_map[layer_id] + solved = plan.physical_to_logical_map[0] + placement_change = torch.stack((live, solved.to(live.dtype))).cpu().numpy() + + assert self._ultraep_logical_to_physical is not None + assert self._ultraep_replica_counts is not None + assert self._ultraep_rank_quota_prefix is not None + self.physical_to_logical_map[layer_id].copy_(solved) + self._ultraep_logical_to_physical[layer_id].copy_( + plan.logical_to_all_physical_map[0] + ) + self._ultraep_replica_counts[layer_id].copy_(plan.logical_to_physical_count[0]) + self._ultraep_rank_quota_prefix[layer_id].copy_( + plan.metadata["rank_quota_prefix"][0] + ) + # All four tables for this layer now describe the same solve, and the + # weight move below makes memory match it. Only from here is it sound + # for _snapshot to route through them. + self._ultraep_committed_layers.add(layer_id) + + self._timed_commit(layer_id, placement_change[0], placement_change[1]) + + def _timed_commit(self, layer_id: int, old_np: Any, new_np: Any) -> None: + """``_commit_placement_weights``, wrapped in the MLB_TIME_REFRESH timing + when it is armed: host issue time and device completion time of each + commit, reported once with the solve times and then disarmed. Shared by + the same-forward and the next-forward (lagged) apply paths.""" + if not self._time_refresh: + self._commit_placement_weights(layer_id, old_np, new_np) + return + torch.accelerator.synchronize() + _tc = time.perf_counter() + self._commit_placement_weights(layer_id, old_np, new_np) + _host = (time.perf_counter() - _tc) * 1e6 + torch.accelerator.synchronize() + _dt = (time.perf_counter() - _tc) * 1e6 + self._commit_host_us.append(_host) + # How much this commit actually had to move: slots whose occupant + # changed. A refresh that re-solves to the same placement moves + # nothing, and the first commit for a layer moves everything, so + # the duration alone says very little without this. + _moved = int((old_np != new_np).sum()) + self._commit_us.append((_moved, _dt)) + if len(self._commit_us) >= self._time_refresh: + import statistics + + def _q(v, f): + return sorted(v)[min(int(f * len(v)), len(v) - 1)] if v else 0.0 + + idle = [d for m, d in self._commit_us if m == 0] + moving = [(m, d) for m, d in self._commit_us if m > 0] + logger.info( + "MLB refresh cost over %d refreshes: solve median=%.0f us " + "p90=%.0f us | commits that moved nothing: n=%d median=%.0f us " + "| commits that moved slots: n=%d median slots=%.0f " + "median=%.0f us p90=%.0f us max=%.0f us", + len(self._commit_us), + statistics.median(self._solve_us), + _q(self._solve_us, 0.9), + len(idle), + statistics.median(idle) if idle else 0.0, + len(moving), + statistics.median([m for m, _ in moving]) if moving else 0.0, + statistics.median([d for _, d in moving]) if moving else 0.0, + _q([d for _, d in moving], 0.9), + max([d for _, d in moving], default=0.0), + ) + logger.info( + "MLB commit host-vs-device: host median=%.0f us | " + "device-complete median=%.0f us -- if the host figure is " + "well below the device one, the device is the binding " + "constraint and shaving host work changes nothing", + statistics.median(self._commit_host_us), + statistics.median([d for _, d in self._commit_us]), + ) + self._commit_host_us.clear() + self._solve_us.clear() + self._commit_us.clear() + self._time_refresh = 0 + + def _commit_placement_weights( + self, + layer_id: int, + old_np: Any, + new_np: Any, + ) -> None: + """Move expert weight so this layer's memory matches the placement the + solve just decided. + + Both arrays are logical-expert ids per physical slot, EP-wide -- exactly + the pair ``EplbState`` hands its own periodic rearrangement, so this is + that same P2P move driven at this policy's cadence instead of at the + rearrangement interval. ``move_to_buffer`` talks to the other ranks and + ``move_from_buffer`` lands the result locally; both run here, in order, + on the calling stream. + """ + if (old_np == new_np).all(): + # The solve landed on the placement already in memory. Nothing to + # move, and skipping is safe for the EP group precisely because + # every rank compares the same two EP-wide arrays and so reaches + # this the same way. + return + if self._ultraep_mover in ("auto", "symm") and self._commit_symm( + layer_id, old_np, new_np + ): + return + if self._ultraep_mover in ("auto", "direct") and self._commit_direct( + layer_id, old_np, new_np + ): + return + + from vllm.distributed.eplb.rebalance_execute import ( + move_from_buffer, + move_to_buffer, + ) + + assert self._ultraep_num_local_physical is not None + metadata = move_to_buffer( + num_local_experts=self._ultraep_num_local_physical, + old_indices=old_np, + new_indices=new_np, + expert_weights=self._ultraep_expert_weights[layer_id], + expert_weights_buffers=self._ultraep_expert_buffer, + stream=None, + ep_rank=self.ep_rank, + communicator=self._ultraep_communicator, + layer_idx=layer_id, + ) + move_from_buffer( + expert_weights=self._ultraep_expert_weights[layer_id], + expert_weights_buffers=self._ultraep_expert_buffer, + transfer_metadata=metadata, + new_indices=new_np, + ep_rank=self.ep_rank, + ) + + def _ultraep_stage_pending(self, layer_id: int, plan: Any) -> None: + """Park a solved plan for this layer until its next forward. + + The plan's GPU tensors are kept as-is (plan() allocates fresh ones per + solve, so nothing is overwritten underneath us) and the (live, solved) + physical->logical pair is copied to pinned host memory asynchronously + on the current stream, with an event marking completion. Nothing here + waits on the device. + """ + live = self.physical_to_logical_map[layer_id] + solved = plan.physical_to_logical_map[0].to(live.dtype) + host = self._ultraep_pending_host.get(layer_id) + if host is None: + host = torch.empty((2, live.numel()), dtype=live.dtype, pin_memory=True) + self._ultraep_pending_host[layer_id] = host + host.copy_(torch.stack((live, solved)), non_blocking=True) + ev = torch.cuda.Event() + ev.record() + self._ultraep_pending[layer_id] = ( + solved, + plan.logical_to_all_physical_map[0], + plan.logical_to_physical_count[0], + plan.metadata["rank_quota_prefix"][0], + host, + ev, + ) + + def _ultraep_apply_pending(self, layer_id: int) -> None: + """Land the plan staged on this layer's previous forward: snapshot first + (the routing that follows reads it), then the weight sync on the same + stream, so the experts kernel downstream sees matching weights. By now + the pinned copy finished a whole step ago, so the event wait is free. + """ + pending = self._ultraep_pending.pop(layer_id, None) + if pending is None: + return + solved, l2p, counts, quota, host, ev = pending + ev.synchronize() + change = host.numpy() + old_np = change[0].copy() + new_np = change[1].copy() + assert self._ultraep_logical_to_physical is not None + assert self._ultraep_replica_counts is not None + assert self._ultraep_rank_quota_prefix is not None + self.physical_to_logical_map[layer_id].copy_(solved) + self._ultraep_logical_to_physical[layer_id].copy_(l2p) + self._ultraep_replica_counts[layer_id].copy_(counts) + self._ultraep_rank_quota_prefix[layer_id].copy_(quota) + self._ultraep_committed_layers.add(layer_id) + self._timed_commit(layer_id, old_np, new_np) + + def _replica_changes(self, layer_id: int, old_np: Any, new_np: Any) -> list | None: + """Replica slots that acquired a new expert, as ``(dst_rank, k)`` pairs + with ``k`` the slot's index among that rank's replica slots -- or None + when either placement is not on the fixed-master layout. + + Both fast movers rest on one assumption: the first ``masters_per_rank`` + slots of every rank hold that rank's own logical experts, in order, on + both sides of the change. The first refresh after loading a + checkpointed (EPLB-permuted) placement breaks it -- that refresh moves + every master home -- and is left to the generic mover. + """ + ep_size = self.ep_size + num_local = self._ultraep_num_local_physical + masters_per_rank = self.num_logical_experts // ep_size + new_grid = np.asarray(new_np).reshape(ep_size, num_local) + old_grid = np.asarray(old_np).reshape(ep_size, num_local) + expected = ( + np.arange(ep_size)[:, None] * masters_per_rank + + np.arange(masters_per_rank)[None, :] + ) + if not ( + np.array_equal(new_grid[:, :masters_per_rank], expected) + and np.array_equal(old_grid[:, :masters_per_rank], expected) + ): + if not self._ultraep_mover_warned: + self._ultraep_mover_warned = True + logger.info( + "MLB weight sync: layer %d moves master slots (placement not " + "on the fixed-master layout); the generic mover handles it and " + "the fast path takes over from the next refresh", + layer_id, + ) + return None + return np.argwhere( + (new_grid[:, masters_per_rank:] != old_grid[:, masters_per_rank:]) + & (new_grid[:, masters_per_rank:] >= 0) + ).tolist() + + def _commit_direct(self, layer_id: int, old_np: Any, new_np: Any) -> bool: + """Weight sync over the EPLB communicator: masters stay put, and every + replica slot whose expert changed receives it straight from the rank + holding the master, in one send/recv group per layer -- no staging + buffer, no second local copy. Sources are master slots (never written + by a refresh) and destinations are replica slots (each written by one + recv), so stream order is the only ordering needed. + + Returns False to hand the layer to the generic mover. + """ + changed = self._replica_changes(layer_id, old_np, new_np) + if changed is None: + return False + masters_per_rank = self.num_logical_experts // self.ep_size + new_grid = np.asarray(new_np).reshape( + self.ep_size, self._ultraep_num_local_physical + ) + weights = self._ultraep_expert_weights[layer_id] + comm = self._ultraep_communicator + posted = False + for dst_rank, k in changed: + dst_row = masters_per_rank + k + expert = int(new_grid[dst_rank, dst_row]) + src_rank, src_row = divmod(expert, masters_per_rank) + if dst_rank == src_rank: + if self.ep_rank == dst_rank: + for w in weights: + w[dst_row].copy_(w[src_row], non_blocking=True) + continue + if self.ep_rank == src_rank: + comm.add_send([w[src_row] for w in weights], dst_rank, expert_id=expert) + posted = True + elif self.ep_rank == dst_rank: + comm.add_recv([w[dst_row] for w in weights], src_rank, expert_id=expert) + posted = True + if posted: + comm.execute() + return True + + def _symm_setup(self) -> bool: + """Allocate and rendezvous the per-rank staging buffer (collective on the + EP group, so every rank must reach its first commit together -- they do, + the refresh runs in lockstep). + + Layout: ``[2 sets, replica slots per rank, expert bytes]``. An expert's + tensors (weights and scales) are packed back to back as raw bytes in the + order ``expert_weights[layer]`` lists them; every layer has the same + shapes, so one layout serves all layers. Two sets alternate with the + layer index: the unpack of layer L is stream-ordered before this rank's + all-gather for layer L+1, and every peer's puts for layer L+1 are issued + after that all-gather completes, so set L%2 is never overwritten while + it is still being read. + + The decision to use symmetric memory is collective. Each rank first + judges locally whether it can take part -- the module imports, and the + EP group either sits on one host or runs on Blackwell-class devices + (taken as a multi-node NVLink fabric; ``MLB_ULTRAEP_MOVER=symm`` skips + that heuristic) -- and the verdicts are combined with an all-reduce + (min) over the EP group, so every rank either enters the rendezvous or + none does. Without that step a rank whose handle exchange failed fell + back to direct while its peers waited in the rendezvous, and the + server deadlocked on its first refresh whenever the EP group spanned + hosts without such a fabric. + + Returns False (and drops to the direct mover) when the group cannot use + symmetric memory. + """ + if self._symm_hdl is not None: + return True + from vllm.distributed import get_ep_group + from vllm.platforms import current_platform + + ep = get_ep_group() + device = self._ultraep_expert_weights[0][0].device + reason = "" + try: + import torch.distributed._symmetric_memory as symm # noqa: F401 + except Exception as exc: # pragma: no cover - build without symm + reason = f"symmetric memory unavailable ({exc})" + if not reason and self._ultraep_mover != "symm": + import socket + + hosts: list[str | None] = [None] * ep.world_size + torch.distributed.all_gather_object( + hosts, socket.gethostname(), group=ep.cpu_group + ) + capability = current_platform.get_device_capability() + multi_node_nvlink = capability is not None and capability.major >= 10 + if len(set(hosts)) > 1 and not multi_node_nvlink: + reason = "EP group spans hosts without a multi-node NVLink fabric" + verdict = torch.tensor([0 if reason else 1], device=device, dtype=torch.int32) + torch.distributed.all_reduce( + verdict, op=torch.distributed.ReduceOp.MIN, group=ep.device_group + ) + if int(verdict.item()) == 0: + logger.info( + "MLB weight sync: %s; using the direct send/recv mover on every rank", + reason or "another EP rank cannot use symmetric memory", + ) + self._ultraep_mover = "direct" + return False + try: + import torch.distributed._symmetric_memory as symm + + weights0 = self._ultraep_expert_weights[0] + layout: list[tuple[int, int, torch.dtype, tuple[int, ...]]] = [] + off = 0 + for w in weights0: + row = w[0] + nbytes = row.numel() * row.element_size() + layout.append((off, nbytes, w.dtype, tuple(row.shape))) + off += nbytes + assert self._ultraep_num_local_physical is not None + n_rep = ( + self._ultraep_num_local_physical + - self.num_logical_experts // self.ep_size + ) + buf = symm.empty( + (2, n_rep, off), dtype=torch.uint8, device=weights0[0].device + ) + hdl = symm.rendezvous(buf, group=ep.device_group) + except Exception as exc: + logger.warning( + "MLB weight sync: symmetric-memory rendezvous failed (%s); using " + "the direct send/recv mover", + exc, + ) + self._ultraep_mover = "direct" + return False + self._symm_layout = layout + self._symm_expert_nbytes = off + self._symm_buf, self._symm_hdl = buf, hdl + logger.info( + "MLB weight sync: symmetric-memory mover ready " + "(%d replica slot(s) x %d MiB, 2 sets, %d tensors per expert)", + n_rep, + off >> 20, + len(layout), + ) + return True + + def _commit_symm(self, layer_id: int, old_np: Any, new_np: Any) -> bool: + """UltraEP weight_sync semantics over torch symmetric memory. + + Every rank owning the master of an expert that some replica slot just + acquired writes that master straight into the acquiring rank's staging + slot -- a device copy into mapped peer memory, the peer does nothing -- + then signals that peer; each acquiring rank waits for exactly the peers + it received from and unpacks its staging slots into its replica rows. + No group-wide barrier, no host round trip. + + Returns False to hand the layer to the generic mover. + """ + if not self._symm_setup(): + return False + changed = self._replica_changes(layer_id, old_np, new_np) + if changed is None: + return False + masters_per_rank = self.num_logical_experts // self.ep_size + assert self._ultraep_num_local_physical is not None + n_rep = self._ultraep_num_local_physical - masters_per_rank + new_grid = np.asarray(new_np).reshape( + self.ep_size, self._ultraep_num_local_physical + ) + weights = self._ultraep_expert_weights[layer_id] + set_id = layer_id % 2 + hdl, buf = self._symm_hdl, self._symm_buf + assert buf is not None + signal_to: set[int] = set() + wait_for: set[int] = set() + # 1. puts: I own the master -> write into the acquirer's staging slot + for dst_rank, k in changed: + expert = int(new_grid[dst_rank, masters_per_rank + k]) + src_rank, src_row = divmod(expert, masters_per_rank) + if src_rank == dst_rank: + if self.ep_rank == dst_rank: + for w in weights: + w[masters_per_rank + k].copy_(w[src_row], non_blocking=True) + continue + if self.ep_rank == dst_rank: + wait_for.add(src_rank) + continue + if self.ep_rank != src_rank: + continue + remote = hdl.get_buffer( + dst_rank, (2, n_rep, self._symm_expert_nbytes), torch.uint8 + ) + stage = remote[set_id, k] + for w, (off, nbytes, _dtype, _shape) in zip(weights, self._symm_layout): + stage[off : off + nbytes].copy_( + w[src_row].contiguous().view(-1).view(torch.uint8), + non_blocking=True, + ) + signal_to.add(dst_rank) + # 2. pairwise completion, stream-ordered after the puts + for dst_rank in signal_to: + hdl.put_signal(dst_rank, channel=0) + for src_rank in wait_for: + hdl.wait_signal(src_rank, channel=0) + # 3. unpack what I acquired + for dst_rank, k in changed: + if dst_rank != self.ep_rank: + continue + expert = int(new_grid[dst_rank, masters_per_rank + k]) + if expert // masters_per_rank == self.ep_rank: + continue # local copy done above + stage = buf[set_id, k] + for w, (off, nbytes, dtype, shape) in zip(weights, self._symm_layout): + w[masters_per_rank + k].copy_( + stage[off : off + nbytes].view(dtype).view(shape), non_blocking=True + ) + return True + + def replica_shares( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + layer_state: EplbLayerState, + num_unpadded_tokens: torch.Tensor | None, + routed_scaling_factor: float = 1.0, + ) -> torch.Tensor | None: + """Ask MLB to route this layer, recording resolved ids as a side effect. + + Always returns ``None``: every replica policy resolves its own ids + (``self._last_physical_ids``) rather than answering with a share table + for this method to hand back. The resolved ids are what + ``resolve_routing`` reads after this call returns. + + ``routed_scaling_factor`` is plumbed but left at its default: MLB only + reads it when it materializes a shared-expert decision into expanded + ids, which is SGLang's contract. vLLM takes back the *rank* and builds + the extra top-k column itself, applying 1/routed_scaling_factor as that + column's weight in the layer (see shared_expert_fusion.py). + """ + from moe_load_balancer.adapters.vllm import ( + to_routing_request, + to_vllm_shared_expert_rank, + ) + from moe_load_balancer.kernels.ops.counting import count_logical_experts + + self._last_physical_ids = None + self._last_shared_expert_rank = None + layer_id = layer_state.moe_layer_idx + if layer_id is None: + raise RuntimeError( + "MLB routing requires EplbLayerState.moe_layer_idx; the layer " + "was registered by a path that does not set it." + ) + + # Accumulate this layer's LOCAL count into the stable buffer. Skipped + # during Dynamo tracing, where the counting kernel may not be handled; + # under graph capture and eager it runs normally and is captured + # alongside the copy_. + # + # This happens BEFORE the readiness check below, not after it. With the + # write on the far side, the opening step returned early without + # recording anything, the next step's reduction still saw zeros, and the + # solve after that was the first to meet real data -- a three-step + # warm-up where the design calls for two, leaving one step of every run + # solving against an all-zero load. + if self._logical_count_local is not None and not torch._dynamo.is_compiling(): + local = count_logical_experts(topk_ids, self.num_logical_experts) + self._logical_count_local[layer_id].copy_(local) + # A placement with no redundancy leaves a replica policy nothing to + # decide: every logical expert has exactly one physical copy, so the + # only answer any of them can give is the one the framework's own + # kernel already produces. Return before the pipeline runs. + # + # After the count write above, not before it -- the counts feed the + # next placement and LPLB's solve, and neither stops being wanted just + # because this placement happens to be degenerate. + # + # A shared-expert policy is exempt: it decides which rank runs the + # shared expert, a choice that exists whether or not any logical + # expert has a second copy. Returning here for `waterfill` (which + # expands to `static+waterfill`) would drop its decision on every + # red0 deployment. + if not self._placement_offers_replica_choice and not self.routes_shared_expert: + return None + + # Only policies that read the load have to wait for it. Gating every + # policy on this readiness flag disabled the ones that never asked for + # a count: their buffers are not allocated, so the flag never flips and + # they returned None for the lifetime of the run. + if self.consumes_global_logical_count and not self._logical_count_ready: + return None + + if self._ultraep_expert_weights is not None: + self._ultraep_fast_refresh(layer_id, topk_ids, num_unpadded_tokens) + + # Pass the PREVIOUS step's global count to MLB. The stable tensor + # slice has a fixed GPU address, so it is safe to use inside a + # captured CUDA graph: on replay the LP solve reads whatever value + # finalize_step_counts deposited before the graph was launched. + global_count = ( + None + if self._fresh_counts + else ( + self._logical_count_global[layer_id] + if self._logical_count_global is not None + else None + ) + ) + + if self._time_l2: + torch.accelerator.synchronize() + _t0 = time.perf_counter() + + decision = self._mlb.route_tokens( + to_routing_request( + layer_id=layer_id, + logical_topk_ids=topk_ids, + topk_weights=topk_weights, + placement=self._snapshot(layer_state, layer_id), + stage=_current_stage(), + token_count=num_unpadded_tokens, + routed_scaling_factor=routed_scaling_factor, + global_logical_count=global_count, + ) + ) + if self._time_l2: + torch.accelerator.synchronize() + self._l2_us.append((time.perf_counter() - _t0) * 1e6) + if len(self._l2_us) >= self._time_l2: + import statistics + + logger.info( + "MLB L2 solve cost: n=%d median=%.0f us p90=%.0f us " + "mean=%.0f us", + len(self._l2_us), + statistics.median(self._l2_us), + sorted(self._l2_us)[int(0.9 * len(self._l2_us))], + statistics.mean(self._l2_us), + ) + self._l2_us.clear() + self._time_l2 = 0 + + # MLB_DUMP_LP= captures what the solve was given and what it + # returned, so the achievable headroom can be computed offline instead + # of inferred from end-to-end throughput. Diagnostic only. Unconditional + # on which policy is active -- global_count/candidates/counts come from + # layer_state regardless, and lp_probability is None for anything but + # LPLB, which is exactly what "no LP ran here" should look like. + if self._dump_lp_skip > 0: + self._dump_lp_skip -= 1 + elif self._dump_lp_dir is not None and self._dump_lp_left > 0: + import os as _os + + self._dump_lp_left -= 1 + lp_probability = decision.metadata.get("lp_probability") + assert layer_state.logical_to_physical_map is not None + assert layer_state.logical_replica_count is not None + torch.save( + { + "layer_id": layer_id, + "ep_size": self.ep_size, + "num_local": self.num_physical_experts // self.ep_size, + "candidates": layer_state.logical_to_physical_map[ + :, : self._max_replicas + ].cpu(), + "counts": layer_state.logical_replica_count.cpu(), + "global_count": ( + None if global_count is None else global_count.float().cpu() + ), + "probability": ( + None if lp_probability is None else lp_probability.float().cpu() + ), + }, + _os.path.join( + self._dump_lp_dir, + f"lp_r{self.ep_rank}_l{layer_id}_{self._dump_lp_left}.pt", + ), + ) + + # Every replica policy resolves its own ids now; none answers with a + # share table for this runtime to apply. Record the ids for + # resolve_routing to pick up. + # + # The shared-expert rank is recorded before the `ids is topk_ids` + # early return below: a pipeline can decide a rank while leaving the + # routed ids untouched (bare `waterfill` does exactly that), and + # returning early there would discard the only decision it made. + self._last_shared_expert_rank = to_vllm_shared_expert_rank(decision) + + ids = decision.routed_physical_topk_ids + if ids is None or ids is topk_ids: + return None + self._last_physical_ids = ids + return None + + +class VllmMlbIntegration: + """One balancer per engine, serving both placement and routing. + + SGLang gives an engine a single MoELoadBalancer and lets L1 and L2 share + it. vLLM did not: the L1 policy built a throwaway instance on every + rebalance and L2 held a second one in a module global. Two instances cannot + see each other's state, which rules out any policy whose placement decision + depends on what routing observed -- the premise of the predictive layer -- + and a module global also rules out more than one engine in a process. + + The lookup below stays module-level because vLLM's placement policy is a + classmethod with nowhere to hang an instance. Ownership is not: the + integration is created and released by EplbState, so its lifetime is the + engine's. + """ + + def __init__(self, algorithm: str = "") -> None: + from moe_load_balancer import MoELoadBalancer + + self.algorithm = algorithm + self._balancer_kwargs: dict | None = None + self._balancer = None if algorithm else MoELoadBalancer() + self.routing: MlbRoutingRuntime | None = None + + def balancer(self): + """The single MoELoadBalancer this engine uses.""" + if self._balancer is None: + from moe_load_balancer import MoELoadBalancer + + if self._balancer_kwargs is None: + # Placement can be asked for before the routing geometry is + # known; a plain planner answers L1 and is replaced in place + # once routing supplies the topology. + self._balancer = MoELoadBalancer() + else: + self._balancer = MoELoadBalancer.from_algorithm( + self.algorithm, **self._balancer_kwargs + ) + return self._balancer + + def bind_routing(self, **kwargs) -> MlbRoutingRuntime: + """Create the routing runtime against this engine's balancer.""" + from moe_load_balancer import MoELoadBalancer + + self._balancer_kwargs = { + "ep_size": kwargs["ep_size"], + "source_rank": kwargs["ep_rank"], + "experts_per_rank": kwargs["num_physical_experts"] // kwargs["ep_size"], + "collectives": VllmRoutingCollectives(), + } + self._balancer = MoELoadBalancer.from_algorithm( + self.algorithm, **self._balancer_kwargs + ) + self.routing = MlbRoutingRuntime( + self.algorithm, balancer=self._balancer, **kwargs + ) + return self.routing + + +_integration: VllmMlbIntegration | None = None + + +def get_mlb_integration() -> VllmMlbIntegration: + """The engine's integration, created on first use.""" + global _integration + if _integration is None: + _integration = VllmMlbIntegration(mlb_l2_algorithm()) + return _integration + + +def set_mlb_integration(integration: VllmMlbIntegration | None) -> None: + global _integration, _runtime + _integration = integration + _runtime = None if integration is None else integration.routing + + +def l2_pipeline_capabilities(algorithm: str): + """Capabilities the named L2 pipeline declares, or None if unknown. + + Stateless, so it answers before any engine exists -- configuration + validation needs it. Living here rather than at each call site keeps the + balancer a detail of one module: everything else in vLLM asks this file. + """ + if not algorithm: + return None + try: + from moe_load_balancer.core.routing_pipeline import RoutingPipeline + + return RoutingPipeline.from_value(algorithm).capabilities + except Exception: + return None + + +def l2_inapplicable_reason(algorithm: str, num_redundant_experts: int) -> str | None: + """Why the named policy cannot act on this deployment, or None. + + The reason is the policy's own words. What this side supplies is what only + it knows: that the L1 plan's metadata reaches the committed snapshot + (``to_placement_snapshot`` carries it, which is how ultraep's quota gets + to its router), and whether the shared expert is dispatched at all: with + VLLM_FUSE_SHARED_EXPERTS off it is a per-rank replicated MLP, so there is + no rank for a shared-expert policy to choose and the policy says so + itself. + + Answered from the env switch rather than from a built layer because this + runs during configuration validation, before any model exists. A layer + that then turns out not to be fusible (uneven split, EP disabled) simply + never asks MLB for a shared-expert rank -- the router checks its own + geometry before it calls. + """ + if not algorithm: + return None + try: + from moe_load_balancer import ExpertDeploymentConfig + from moe_load_balancer.core.routing_pipeline import RoutingPipeline + + pipeline = RoutingPipeline.from_value(algorithm) + except Exception: + return None + from vllm.model_executor.layers.fused_moe.shared_expert_fusion import ( + shared_expert_fusion_enabled, + ) + + return pipeline.is_applicable( + ExpertDeploymentConfig( + num_redundant_experts=num_redundant_experts, + routes_shared_expert=shared_expert_fusion_enabled(), + carries_placement_metadata=True, + ) + ) + + +def plan_placement(request): + """Run L1 through the engine's balancer and hand back vLLM's one map.""" + from moe_load_balancer.adapters.vllm import to_vllm_physical_to_logical + + integration = get_mlb_integration() + plan = integration.balancer().plan_placement(request) + # UltraEP is the only L1 policy that publishes this; every other plan's + # metadata simply lacks the key, so this stays a no-op for them. Routing + # geometry can lag placement (see balancer()'s docstring), so there may be + # no runtime to hand the quota to yet -- it reads whatever the next solve + # after bind_routing() leaves here. + # + # The candidate table and replica counts come along too, not just the + # quota: they must be read from this same plan, not rebuilt from + # phy2log through vLLM's own compute_logical_maps, or the quota's column + # ordering and width silently stop matching the table L2 routes against. + # + # Stashed here rather than after the caller commits the weight move: no + # forward can observe this quota paired with the placement it belongs to + # before that commit happens, because MlbEplbPolicy requires synchronous + # (use_async=False) rearrangement -- nothing else runs on this thread + # between this return and register_logical_maps(). An async L1 path would + # need this to move to the commit step instead. + quota = plan.metadata.get("rank_quota_prefix") + if quota is not None and integration.routing is not None: + integration.routing._ultraep_rank_quota_prefix = quota + integration.routing._ultraep_logical_to_physical = ( + plan.logical_to_all_physical_map + ) + integration.routing._ultraep_replica_counts = plan.logical_to_physical_count + # All three come from the same plan whose physical_to_logical_map the + # caller commits next, so every layer's tables are consistent with the + # placement as of this moment. Say so: `_snapshot` routes a layer + # against MLB's tables only once it appears here, and leaving the set + # empty left this solve's routing unused until something else added the + # layer. Nothing was observably wrong -- RefreshGate refreshes on a + # layer's first call whatever the interval, so every layer was added on + # its first forward -- but the invariant was being maintained by a + # second component's incidental behaviour rather than by the code that + # knows the tables are ready. Stating it here is what makes the gate's + # meaning ("routable") match the condition it tests. + integration.routing._ultraep_committed_layers = set(range(quota.shape[0])) + return to_vllm_physical_to_logical(plan), plan + + +def placement_request(*args, **kwargs): + """Translate vLLM's rebalance arguments into the neutral request.""" + from moe_load_balancer.adapters.vllm import to_placement_request + + return to_placement_request(*args, **kwargs) + + +def _reject_graphs_with_rearranging_placement_state(rearranges: bool) -> None: + """Refuse the one combination that can fault the GPU. + + A policy that keeps per-layer state rebuilds it when a committed placement + changes that state's layout, and LPLB *replaces* the tensors in that case + rather than updating them in place -- its own ``prepare_layer`` docstring + notes that keeping a previously captured graph valid across such a change + "requires a future fixed-shape LPLB state representation". + + A CUDA graph captures pointers. Once the solver's buffers move, replaying + the graph reads freed memory: observed as ``Xid 43`` on two devices and a + dead worker. It is intermittent -- it needs a rearrangement that actually + changes the layout, so a short benchmark can pass and a long serving run + can fault. That is worth failing loudly for. + + Both escapes keep MLB routing available: pin the placement (leave + rearrangement off, which is how the reported LPLB numbers were measured), + or run eager. + """ + if not rearranges: + return + from vllm.config import get_current_vllm_config + + try: + compilation = get_current_vllm_config().compilation_config + except Exception: # noqa: BLE001 - no config context: leave the decision alone + return + if getattr(compilation, "cudagraph_mode", None) is None: + return + if compilation.cudagraph_mode.name == "NONE": + return + raise ValueError( + "VLLM_MLB_L2_ALGORITHM keeps per-layer solver state, and EPLB " + "rearrangement can change that state's layout. LPLB replaces its " + "tensors on such a change, which a captured CUDA graph cannot follow " + "-- replay then reads freed memory (Xid 43). Either disable " + "rearrangement (eplb_config step_interval high, or policy that does " + "not re-plan), or set enforce_eager=True. Fixing this properly needs " + "a fixed-shape solver state in moe_load_balancer." + ) + + +def init_mlb_routing( + *, + algorithm: str, + ep_size: int, + ep_rank: int, + num_logical_experts: int, + num_physical_experts: int, + physical_to_logical_map: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, + rearranges: bool = False, + expert_weights: Any | None = None, + expert_buffer: Any | None = None, + communicator: Any | None = None, +) -> MlbRoutingRuntime | None: + """Create the routing runtime for a configured L2 algorithm. + + ``algorithm`` comes from ``EPLBConfig.l2_algorithm``, which has already + resolved the environment default and cleared itself for placements no L2 + policy can act on. Passing it in rather than re-reading the environment is + what makes that decision binding. + + ``expert_weights`` is the model's own ``expert_weights`` (one entry per + MoE layer, present at this same call site); ``expert_buffer`` and + ``communicator`` are ``EplbState``'s own staging buffer and P2P backend + for moving expert weight between ranks. Only ``ultraep`` reads them -- it + is the one policy that re-plans placement often enough to need weight + moved mid-run rather than at the rearrangement cadence -- and it moves it + with this framework machinery rather than any of its own. Every other + algorithm ignores all three, so they are safe to leave unset. + """ + global _runtime + if not algorithm: + return None + integration = get_mlb_integration() + integration.algorithm = algorithm + _runtime = integration.bind_routing( + ep_size=ep_size, + ep_rank=ep_rank, + num_logical_experts=num_logical_experts, + num_physical_experts=num_physical_experts, + physical_to_logical_map=physical_to_logical_map, + expert_weights=expert_weights, + expert_buffer=expert_buffer, + communicator=communicator, + ) + _runtime.register_logical_maps(logical_to_physical_map, logical_replica_count) + # Keyed on the declared stability of the policy's state, not on whether it + # keeps state at all. A policy whose placement-derived tensors keep their + # addresses across a rearrangement is safe to capture; refusing it because + # some other policy is not would bar a combination that never faults. + if _runtime.graph_stability == "realloc_on_placement_change": + _reject_graphs_with_rearranging_placement_state(rearranges) + return _runtime + + +def get_mlb_routing() -> MlbRoutingRuntime | None: + return _runtime + + +def reset_mlb_routing() -> None: + """Test hook. Releases the engine's integration along with the runtime.""" + global _runtime, _integration + _runtime = None + _integration = None diff --git a/vllm/distributed/eplb/policy/__init__.py b/vllm/distributed/eplb/policy/__init__.py index 8e78d7bac0e3..be0d439482eb 100644 --- a/vllm/distributed/eplb/policy/__init__.py +++ b/vllm/distributed/eplb/policy/__init__.py @@ -6,8 +6,12 @@ from .abstract import AbstractEplbPolicy from .default import DefaultEplbPolicy +from .mlb import MlbEplbPolicy -EPLB_POLICIES = {"default": DefaultEplbPolicy} +EPLB_POLICIES: dict[str, type[AbstractEplbPolicy]] = { + "default": DefaultEplbPolicy, + "mlb": MlbEplbPolicy, +} # Ensure that the EPLB_POLICIES keys match the EPLBPolicyOption values assert set(EPLB_POLICIES.keys()) == set(get_args(EPLBPolicyOption)) @@ -16,4 +20,5 @@ "AbstractEplbPolicy", "DefaultEplbPolicy", "EPLB_POLICIES", + "MlbEplbPolicy", ] diff --git a/vllm/distributed/eplb/policy/mlb.py b/vllm/distributed/eplb/policy/mlb.py new file mode 100644 index 000000000000..51392bf390b3 --- /dev/null +++ b/vllm/distributed/eplb/policy/mlb.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM-owned glue for the framework-decoupled MoE Load Balancer (MLB). + +MLB (``moe-load-balancer``) implements expert-placement and token-routing +policies once, framework-neutrally, and each serving framework keeps a thin +adapter. This module is that adapter for vLLM's L1 placement slot: it +translates :class:`AbstractEplbPolicy` arguments into an MLB +``PlacementRequest``, runs the selected MLB policy, and translates the result +back into the ``physical_to_logical_map`` vLLM expects. + +Enable with ``--eplb-config '{"policy": "mlb", "use_async": false}'``. +The MLB placement algorithm is selected with ``VLLM_MLB_L1_ALGORITHM`` +(default ``auto``); see ``moe_load_balancer`` for the supported names. + +Note that ``use_async`` must be disabled: vLLM's async rearrangement path is +validated only against the built-in policy (see ``EPLBConfig``), and MLB's +planner runs synchronously on the caller's thread. + +``ultraep`` places redundant replicas by which *rank* is overloaded, not just +which logical expert is hot, so it needs every rank's own count rather than +the cross-rank sum every other algorithm here is happy with. +:meth:`wants_per_rank_weight` tells ``EplbState.rearrange`` to gather instead +of reduce and hand the per-rank breakdown back through ``per_rank_weight``. +It also needs to know which rank *this call* is running on -- its placement +kernel solves once per rank, not once globally, and validates ``rank`` is +in range -- so ``EplbState.rearrange`` passes its own ``ep_rank`` alongside +``per_rank_weight`` for the same policies. +""" + +from __future__ import annotations + +import torch + +import vllm.envs as envs +from vllm.distributed.eplb.policy.abstract import AbstractEplbPolicy +from vllm.distributed.eplb.policy.default import DefaultEplbPolicy +from vllm.logger import init_logger + +logger = init_logger(__name__) + +DEFAULT_ALGORITHM = "auto" +ALGORITHM_ENV = "VLLM_MLB_L1_ALGORITHM" + + +def _algorithm() -> str: + return envs.VLLM_MLB_L1_ALGORITHM or DEFAULT_ALGORITHM + + +class MlbEplbPolicy(AbstractEplbPolicy): + """Delegate expert placement to the MoE Load Balancer core.""" + + @classmethod + def wants_per_rank_weight(cls) -> bool: + return _algorithm() == "ultraep" + + @classmethod + def rebalance_experts( + cls, + weight: torch.Tensor, + num_replicas: int, + num_groups: int, + num_nodes: int, + num_ranks: int, + old_global_expert_indices: torch.Tensor | None = None, + per_rank_weight: torch.Tensor | None = None, + ep_rank: int | None = None, + ) -> torch.Tensor: + try: + from vllm.distributed.eplb.mlb_runtime import ( + placement_request, + plan_placement, + ) + except ImportError as exc: # pragma: no cover - depends on environment + raise ImportError( + "EPLB policy 'mlb' requires the moe-load-balancer package. " + "Install it into this environment with " + "`pip install -e /path/to/moe_load_balancer --no-deps`." + ) from exc + + algorithm = _algorithm() + + # ultraep's placement kernel derives the cross-rank sum itself from + # the per-rank breakdown, so hand that over whole instead of the + # already-summed weight every other algorithm here expects. Unlike + # every other algorithm's CPU numpy solve, ultraep's is a CUDA + # kernel -- it stays on-device rather than following weight.cpu(). + logical_count = ( + weight.float().cpu() if per_rank_weight is None else per_rank_weight.float() + ) + # vLLM passes num_groups=0 for models without expert groups; MLB + # expresses "no grouping" as a single group. + request = placement_request( + logical_count, + num_replicas=num_replicas, + num_ranks=num_ranks, + num_groups=num_groups or 1, + num_nodes=num_nodes, + algorithm=algorithm, + old_physical_to_logical_map=( + None + if old_global_expert_indices is None + else old_global_expert_indices.cpu() + ), + ep_rank=ep_rank, + ) + + # The engine's balancer, not a fresh one per rebalance: L1 and L2 must + # be the same instance for any policy whose placement depends on what + # routing observed. + phy2log, plan = plan_placement(request) + phy2log = phy2log.cpu().to(torch.int64) + + # MLB's L1 policies do not consume a previous placement, so the + # slot-preservation pass that the built-in policy performs inside + # rebalance_experts is applied here instead. It only permutes slots + # within a rank, so the placement MLB decided (which logical experts + # live on which rank, and how many replicas each gets) is unchanged -- + # only their slot positions are, which is what avoids weight copies. + if old_global_expert_indices is not None: + phy2log = torch.from_numpy( + DefaultEplbPolicy.preserve_intragpu_slots( + phy2log.numpy(), + num_ranks, + old_global_expert_indices.cpu().numpy(), + ) + ) + + logger.info_once( + "EPLB placement delegated to MLB (algorithm=%s, resolved=%s)", + algorithm, + plan.metadata.get("resolved_policy", "unknown"), + ) + return phy2log diff --git a/vllm/distributed/eplb/rebalance_execute.py b/vllm/distributed/eplb/rebalance_execute.py index 7a52dd5f56c0..2905aef44fcf 100644 --- a/vllm/distributed/eplb/rebalance_execute.py +++ b/vllm/distributed/eplb/rebalance_execute.py @@ -60,6 +60,21 @@ class AsyncEplbLayerResult: """ +def experts_being_acquired( + old_indices: np.ndarray, new_indices: np.ndarray +) -> np.ndarray: + """Experts that land on a slot they did not already occupy. + + Only these can need a transfer. If every slot now holding an expert already + held it, the ranks holding it afterwards are a subset of the ranks holding + it before, so its ranks_to_recv -- the new holders minus the old ones -- is + empty and the send loop skips it. Restricting the send-side query to this + set therefore issues exactly the same transfers, while sizing the query by + the rebalance instead of by the deployment. + """ + return np.unique(new_indices[new_indices != old_indices]) + + def get_ep_ranks_with_experts_batch( expert_ids: np.ndarray, num_local_experts: int, @@ -270,6 +285,16 @@ def move_to_buffer( if send_count > 0: experts = send_expert_ids[:send_count] srcs = send_src_rows[:send_count] + + # send_count covers every expert resident on this rank, but the loop + # below skips any that nobody is acquiring. The query costs roughly + # linearly in the experts it is asked about, so on a refresh that moves + # a handful of slots the unnarrowed form spends most of the transfer's + # cost deciding that almost nothing needs to move. + keep = np.isin(experts, experts_being_acquired(old_indices, new_indices)) + experts = experts[keep] + srcs = srcs[keep] + order = np.argsort(experts, kind="stable") experts = experts[order] srcs = srcs[order] diff --git a/vllm/envs.py b/vllm/envs.py index c8e1a8b70803..fddc576a6d9a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -214,6 +214,7 @@ ] = "relax" VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True VLLM_MOE_SKIP_PADDING: bool = True + VLLM_FUSE_SHARED_EXPERTS: bool = False VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False VLLM_KIMI_K3_GEMM_AR: bool = True @@ -285,6 +286,24 @@ VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "VLLM_OBJECT_STORAGE_SHM_BUFFER" VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024 + VLLM_DEEPEP_RDMA_BUFFER_SIZE_MB: int | None = None + VLLM_DEEPEP_HT_USE_MNNVL: bool = False + VLLM_DEEPEP_V2_RDMA_GBS: float = 0.0 + VLLM_DEEPEP_V2_NVLINK_GBS: float = 0.0 + VLLM_EPLB_DUMP_LOAD_PATH: str | None = None + VLLM_MLB_L1_ALGORITHM: str = "" + VLLM_MLB_L2_ALGORITHM: str = "" + MLB_ULTRAEP_MOVER: str = "direct" + MLB_ULTRAEP_LAGGED_APPLY: bool = False + MLB_ULTRAEP_REFRESH_INTERVAL: int = 64 + MLB_ULTRAEP_REFRESH_MIN_TOKENS: int = 512 + MLB_KEEP_ZERO_REDUNDANCY_COUNTS: bool = False + MLB_FRESH_COUNTS: bool = False + MLB_DUMP_LP: str | None = None + MLB_DUMP_LP_N: int = 120 + MLB_DUMP_LP_SKIP: int = 0 + MLB_TIME_L2: int = 0 + MLB_TIME_REFRESH: int = 0 VLLM_DEEPEP_HIGH_THROUGHPUT_FORCE_INTRA_NODE: bool = False VLLM_DEEPEP_LOW_LATENCY_USE_MNNVL: bool = False VLLM_DEEPEP_V2_ALLOW_HYBRID_MODE: bool = True @@ -1633,6 +1652,23 @@ def _resolve_rust_cli_path() -> str | None: # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that # treats topk_id == -1 as a skip sentinel "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))), + # Dispatch the shared expert(s) through the EP all-to-all instead of + # running a replicated MLP on every rank. + # + # Each EP rank gets its own expert slot for the shared expert, appended + # after that rank's routed slots, so a shared-expert token has exactly one + # owning rank -- the invariant DeepEP's `expert_id // experts_per_rank` + # rank derivation relies on. This is what gives a load balancer something + # to decide (MLB's Waterfill L2 picks that rank); with the shared expert + # replicated per rank there is no choice to make. + # + # Distinct from VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS, which fuses the + # shared expert into the *local* grouped GEMM under one globally shared + # expert id and de-duplicates with a token-level mask -- a layout DeepEP + # cannot route. + "VLLM_FUSE_SHARED_EXPERTS": lambda: bool( + int(os.getenv("VLLM_FUSE_SHARED_EXPERTS", "0")) + ), # Kimi-K3 only. Under sequence-parallel MoE the dense and shared-expert MLPs # are replicated on every rank, so each rank streams the whole weight to # serve its own token shard. Shard them across TP instead: the MLP then @@ -1998,6 +2034,78 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_DEEPEP_BUFFER_SIZE_MB": lambda: int( os.getenv("VLLM_DEEPEP_BUFFER_SIZE_MB", "1024") ), + # Size the DeepEP high-throughput RDMA buffer separately from the NVL one. + # DeepEP caps num_nvl_bytes at INT_MAX whenever num_rdma_bytes != 0, so the + # shared knob cannot give internode dispatch (EP > 8 NVL peers) more room. + # Unset: same as VLLM_DEEPEP_BUFFER_SIZE_MB. + "VLLM_DEEPEP_RDMA_BUFFER_SIZE_MB": lambda: ( + None + if os.getenv("VLLM_DEEPEP_RDMA_BUFFER_SIZE_MB") is None + else int(os.environ["VLLM_DEEPEP_RDMA_BUFFER_SIZE_MB"]) + ), + # Let DeepEP's high-throughput buffer use MNNVL (allow_mnnvl) so an EP group + # can span trays of a multi-node NVLink domain (e.g. Blackwell multi-node + # NVLink systems); the + # default IPC path is node-local and fails at startup there. + "VLLM_DEEPEP_HT_USE_MNNVL": lambda: bool( + int(os.getenv("VLLM_DEEPEP_HT_USE_MNNVL", "0")) + ), + # Link bandwidths (GB/s) handed to DeepEP v2's SM estimator. Its auto-probe + # reads rdma_gbs from an IB NIC, which is 0 on an MNNVL machine where the + # scale-out ranks ride NVLink, and the estimator then divides by zero. + # 0 keeps DeepEP's own probing. + "VLLM_DEEPEP_V2_RDMA_GBS": lambda: float(os.getenv("VLLM_DEEPEP_V2_RDMA_GBS", "0")), + "VLLM_DEEPEP_V2_NVLINK_GBS": lambda: float( + os.getenv("VLLM_DEEPEP_V2_NVLINK_GBS", "0") + ), + # If set, EplbState dumps the recorded global logical expert load (and the + # placement it was recorded under) to .rearr.pt at every + # rearrangement step, from EP rank 0. Observation only; off by default. + "VLLM_EPLB_DUMP_LOAD_PATH": lambda: os.getenv("VLLM_EPLB_DUMP_LOAD_PATH"), + # MoE Load Balancer (--eplb-config policy="mlb") integration knobs. + # L1 = expert placement algorithm run at each rearrangement ("auto" picks + # the balancer's default for the model's expert-group layout); L2 = the + # per-token replica routing expression, e.g. "static", "lplb", + # "static+waterfill", "ultraep+waterfill". Empty keeps vLLM's built-in + # behaviour. The L2 value is also the default of eplb_config.l2_algorithm. + "VLLM_MLB_L1_ALGORITHM": lambda: os.getenv("VLLM_MLB_L1_ALGORITHM", "").strip(), + "VLLM_MLB_L2_ALGORITHM": lambda: os.getenv("VLLM_MLB_L2_ALGORITHM", "").strip(), + # How UltraEP's fast refresh moves re-solved replica weights between EP + # ranks: "direct" (pairwise send/recv over the EPLB communicator), + # "symm" (one-sided puts through torch symmetric memory; needs one host + # or a multi-node NVLink fabric) or "auto" (symm when every rank agrees + # it can, otherwise direct). + "MLB_ULTRAEP_MOVER": lambda: os.getenv("MLB_ULTRAEP_MOVER", "direct").strip(), + # 1: a plan solved on one forward lands on the layer's next forward, which + # avoids a per-layer device sync at the price of one step of staleness. + "MLB_ULTRAEP_LAGGED_APPLY": lambda: ( + os.getenv("MLB_ULTRAEP_LAGGED_APPLY", "0").strip() != "0" + ), + # Re-solve placement and quota every N representative batches; a batch is + # representative when it routes at least MIN_TOKENS tokens. + "MLB_ULTRAEP_REFRESH_INTERVAL": lambda: int( + os.getenv("MLB_ULTRAEP_REFRESH_INTERVAL", "64") + ), + "MLB_ULTRAEP_REFRESH_MIN_TOKENS": lambda: int( + os.getenv("MLB_ULTRAEP_REFRESH_MIN_TOKENS", "512") + ), + # Diagnostics, off by default. KEEP_ZERO_REDUNDANCY_COUNTS keeps the + # per-step logical count collective with no redundant experts (where the + # result is discarded) so its cost can be measured; FRESH_COUNTS makes the + # balancer gather the count itself per layer instead of using the + # one-step-stale count; DUMP_LP saves the LP inputs/outputs of up to + # DUMP_LP_N calls after skipping DUMP_LP_SKIP into that directory; + # TIME_L2 / TIME_REFRESH time N routing calls / N refreshes with device + # syncs and log a summary once. + "MLB_KEEP_ZERO_REDUNDANCY_COUNTS": lambda: ( + os.getenv("MLB_KEEP_ZERO_REDUNDANCY_COUNTS", "0") == "1" + ), + "MLB_FRESH_COUNTS": lambda: os.getenv("MLB_FRESH_COUNTS", "0") == "1", + "MLB_DUMP_LP": lambda: os.getenv("MLB_DUMP_LP") or None, + "MLB_DUMP_LP_N": lambda: int(os.getenv("MLB_DUMP_LP_N", "120")), + "MLB_DUMP_LP_SKIP": lambda: int(os.getenv("MLB_DUMP_LP_SKIP", "0")), + "MLB_TIME_L2": lambda: int(os.getenv("MLB_TIME_L2", "0")), + "MLB_TIME_REFRESH": lambda: int(os.getenv("MLB_TIME_REFRESH", "0")), # Force DeepEP to use intranode kernel for inter-node communication in # high throughput mode. This is useful archive higher prefill throughput # on system supports multi-node nvlink (e.g GB200). diff --git a/vllm/forward_context.py b/vllm/forward_context.py index 9ddc8e9255ed..cdb1aa84ce08 100644 --- a/vllm/forward_context.py +++ b/vllm/forward_context.py @@ -157,6 +157,16 @@ class ForwardContext: cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE batch_descriptor: BatchDescriptor | None = None + # True only when *every* DP rank is running a single-token decode this + # step, reduced across ranks by the same all-reduce that agrees on token + # counts (see `coordinate_batch_across_dp`). `batch_descriptor.uniform` + # answers the same question for this rank alone, which is not the same + # thing under DP: a rank with no requests runs a dummy decode batch while + # its peers prefill. Anything that gates a collective on "is this a decode + # step" has to read this, not the descriptor, or the ranks that take the + # branch will wait in their collective for ranks that did not. + uniform_decode_across_dp: bool = False + ubatch_slices: UBatchSlices | None = None # Boolean mask over the token axis: True for padding rows that are not real @@ -238,6 +248,7 @@ def create_forward_context( additional_kwargs: dict[str, Any] | None = None, skip_compiled: bool = False, is_padding: torch.Tensor | None = None, + uniform_decode_across_dp: bool = False, ): if vllm_config.compilation_config.fast_moe_cold_start: all_moe_layers = vllm_config.compilation_config.static_all_moe_layers @@ -256,6 +267,7 @@ def create_forward_context( skip_compiled=skip_compiled, additional_kwargs=additional_kwargs or {}, is_padding=is_padding, + uniform_decode_across_dp=uniform_decode_across_dp, ) @@ -286,6 +298,7 @@ def set_forward_context( slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None, skip_compiled: bool = False, is_padding: torch.Tensor | None = None, + uniform_decode_across_dp: bool = False, ): """A context manager that stores the current forward context, can be attention metadata, etc. @@ -314,7 +327,7 @@ def set_forward_context( ): assert ubatch_slices is None assert num_tokens is not None - _, num_tokens_across_dp, _ = coordinate_batch_across_dp( + _, num_tokens_across_dp, _, _ = coordinate_batch_across_dp( num_tokens_unpadded=num_tokens, parallel_config=vllm_config.parallel_config, allow_microbatching=False, @@ -355,6 +368,7 @@ def set_forward_context( additional_kwargs, skip_compiled, is_padding=is_padding, + uniform_decode_across_dp=uniform_decode_across_dp, ) try: diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 11e2a5752098..c8e78708af8a 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1267,6 +1267,14 @@ class FusedMoEConfig: # Defaults to in_dtype if not specified. router_logits_dtype: torch.dtype | None = None + num_local_shared_experts: int = 0 + """How many of ``num_local_experts`` are this rank's shared-expert copies. + + They sit at the end of the local range and are outside EPLB's jurisdiction: + EPLB permutes logical experts among physical slots, and a shared-expert copy + has no logical expert to be. Zero unless VLLM_FUSE_SHARED_EXPERTS is on. + """ + # Defaults to hidden_dim if not specified. hidden_dim_unpadded: int | None = None # Defaults to intermediate_size_per_partition if not specified. diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index ff620622c1f1..3454da7b61c9 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -33,6 +33,9 @@ from vllm.model_executor.layers.fused_moe.runner.moe_runner import ( MoERunner, ) +from vllm.model_executor.layers.fused_moe.shared_expert_fusion import ( + maybe_build_shared_expert_fusion, +) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) @@ -252,6 +255,72 @@ def FusedMoEFactory( ) ) + # Shared-expert fusion has two mutually exclusive layouts (see + # shared_expert_fusion.py): AITER's one-global-id-plus-mask, counted by + # `num_fused_shared_experts` above, and the EP-dispatched one built here, + # which widens the expert space instead. Asking for both is a + # configuration error rather than something to resolve silently -- they + # disagree about what an expert id means. + shared_expert_fusion = maybe_build_shared_expert_fusion( + n_shared_experts=n_shared_experts, + num_physical_experts=global_num_experts, + ep_size=moe_parallel_config.ep_size, + ep_rank=moe_parallel_config.ep_rank, + use_ep=moe_parallel_config.use_ep, + is_act_and_mul=is_act_and_mul, + shared_expert_weight=( + (1.0 / routed_scaling_factor) + if (apply_routed_scale_to_output and routed_scaling_factor) + else 1.0 + ), + layer_name=layer_name, + ) + # The fused shared expert carries a constant weight (1/routed_scaling_factor + # above), which is only the right magnitude when the routed weights sum to + # one -- i.e. when the router renormalizes. Without that, the unfused path + # would have scaled the shared expert by sum(routed weights) and this one + # does not, so the two stop agreeing and the error shows up as a quality + # regression rather than a failure. SGLang's integration carries the same + # implicit assumption; making it explicit is cheap. + if shared_expert_fusion is not None and not renormalize: + raise ValueError( + "Shared-expert fusion assumes a renormalizing router: the fused " + "shared expert is given a constant weight, which is only correct " + "when the routed weights sum to 1. This model routes with " + "renormalize=False. Unset VLLM_FUSE_SHARED_EXPERTS (and drop " + "waterfill from VLLM_MLB_L2_ALGORITHM, which turns fusion on) to " + "keep the shared expert replicated per rank." + ) + + if shared_expert_fusion is not None and num_fused_shared_experts > 0: + raise ValueError( + "VLLM_FUSE_SHARED_EXPERTS and the AITER fused-shared-expert path " + "are both active. They place the shared expert under different " + "expert-id layouts and cannot be combined; unset one of " + "VLLM_FUSE_SHARED_EXPERTS / " + "VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS." + ) + + # The router still selects among routed experts only; the shared slot is + # appended after it. Everything that sizes a buffer per selected expert + # (dispatch, the grouped GEMM's workspace) must see the wider value. + routed_top_k = top_k + if shared_expert_fusion is not None: + global_num_experts = shared_expert_fusion.global_num_experts + top_k = routed_top_k + shared_expert_fusion.num_shared_experts + logger.info_once( + "Shared-expert fusion ON: %d routed + %d shared slots per rank " + "over %d EP ranks -> global_num_experts %d, top_k %d -> %d, " + "shared slot weight %.4f", + shared_expert_fusion.routed_slots_per_rank, + shared_expert_fusion.num_shared_experts, + shared_expert_fusion.ep_size, + global_num_experts, + routed_top_k, + top_k, + shared_expert_fusion.shared_expert_weight, + ) + # Initialize EPLB manager (or None?) eplb_state: EplbLayerState | None = None if enable_eplb: @@ -291,7 +360,7 @@ def FusedMoEFactory( # monolithic. if router is None: router = create_fused_moe_router( - top_k=top_k, + top_k=routed_top_k, global_num_experts=global_num_experts, eplb_state=eplb_state, renormalize=renormalize, @@ -331,6 +400,12 @@ def FusedMoEFactory( image_sentinel_lo=image_sentinel_lo, ) + # Attached rather than passed through every router constructor: the append + # happens in BaseRouter after the EPLB mapping, so it is one step in a + # template method that all routers already share, not a routing algorithm + # any individual one needs to know about. + router.shared_expert_fusion = shared_expert_fusion + if params_dtype is None: params_dtype = torch.get_default_dtype() @@ -352,6 +427,11 @@ def FusedMoEFactory( intermediate_pad=intermediate_pad, num_local_experts=expert_map_manager.local_num_experts, num_logical_experts=logical_num_experts, + num_local_shared_experts=( + shared_expert_fusion.num_shared_experts + if shared_expert_fusion is not None + else 0 + ), moe_parallel_config=moe_parallel_config, in_dtype=moe_in_dtype, moe_backend=vllm_config.kernel_config.moe_backend, diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 3de0aa8dab25..ab213aba0533 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -1235,8 +1235,16 @@ def _maybe_make_contiguous( and name not in NON_EXPERT_WEIGHTS ) + # Trailing shared-expert slots are excluded: EPLB reads + # `num_local_physical_experts` off this tensor's leading dimension and + # asserts it tiles the physical expert count, which counts routed + # experts only. The slice is a view of the same storage, so a + # rearrangement still writes through to the real parameter. + num_routed_local = ( + self.local_num_experts - self.moe_config.num_local_shared_experts + ) return [ - weight.view(self.local_num_experts, -1) + weight.view(self.local_num_experts, -1)[:num_routed_local] for name, weight in weights if name not in NON_EXPERT_WEIGHTS and weight.shape != torch.Size([]) diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index b51dbf90dc7e..89f3846041f1 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -6,13 +6,19 @@ import torch from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( FusedMoERouter, ) +from vllm.model_executor.layers.fused_moe.shared_expert_fusion import ( + SharedExpertFusion, +) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.v1.worker.ubatching import dbo_current_ubatch_id +logger = init_logger(__name__) + if current_platform.is_cuda_alike() or current_platform.is_xpu(): @triton.jit @@ -24,12 +30,17 @@ def _eplb_map_and_record_i32_kernel( out_ptr, record_enabled_ptr, num_unpadded_tokens_ptr, + replica_prob_ptr, + physical_ids_ptr, num_logical_experts, map_slots, out_size, numel, num_active_experts, HAS_NUM_UNPADDED: tl.constexpr, + HAS_REPLICA_PROB: tl.constexpr, + HAS_PHYSICAL_IDS: tl.constexpr, + PROB_SLOTS: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) @@ -40,25 +51,67 @@ def _eplb_map_and_record_i32_kernel( valid_expert = (expert_id >= 0) & (expert_id < num_logical_experts) safe_expert_id = tl.where(valid_expert, expert_id, 0) - # 1. Convert the logical expert ids to physical expert ids - replica_count = tl.load( - logical_replica_count_ptr + safe_expert_id, - mask=mask & valid_expert, - other=1, - ) - # Avoid invalid modulo/div by forcing at least 1. - replica_count = tl.maximum(replica_count, 1) - # floor(2^32 / phi), classic Knuth multiplicative hash multiplier. - KNUTH_MULTIPLIER = 2654435769 - token_idx = (offs // num_active_experts).to(tl.int64) - hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF - replica_idx = hashed % replica_count - map_index = safe_expert_id * map_slots + replica_idx - physical_id = tl.load( - logical_to_physical_ptr + map_index, - mask=mask & valid_expert, - other=-1, - ) + # 1. Convert the logical expert ids to physical expert ids. + # A policy that resolved the replica itself needs none of this: + # the count load and the hash exist only to pick a column, and + # leaving them in charged the direct path for a choice it had + # already made. + if not HAS_PHYSICAL_IDS: + replica_count = tl.load( + logical_replica_count_ptr + safe_expert_id, + mask=mask & valid_expert, + other=1, + ) + # Avoid invalid modulo/div by forcing at least 1. + replica_count = tl.maximum(replica_count, 1) + # floor(2^32 / phi), classic Knuth multiplicative hash multiplier. + KNUTH_MULTIPLIER = 2654435769 + token_idx = (offs // num_active_experts).to(tl.int64) + hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF + replica_idx = hashed % replica_count + + if HAS_REPLICA_PROB: + # A load-balancing policy supplied per-expert replica shares. Use + # the same hash as a uniform draw and invert the row's CDF, so the + # choice stays a pure function of the token index: every rank + # computes the same replica for the same token without exchanging + # anything, and no RNG state enters the CUDA graph. + row = safe_expert_id * PROB_SLOTS + total = tl.zeros_like(hashed).to(tl.float32) + for j in tl.static_range(PROB_SLOTS): + total += tl.load( + replica_prob_ptr + row + j, mask=mask & valid_expert, other=0.0 + ) + # 2^32; hashed is already reduced modulo it. + draw = (hashed.to(tl.float32) / 4294967296.0) * total + acc = tl.zeros_like(total) + sampled = tl.zeros_like(replica_idx) + for j in tl.static_range(PROB_SLOTS): + acc += tl.load( + replica_prob_ptr + row + j, mask=mask & valid_expert, other=0.0 + ) + # Index = how many prefix sums stay at or below the draw. + sampled += tl.where(acc <= draw, 1, 0) + sampled = tl.minimum(sampled, replica_count - 1) + # An all-zero row means "no preference"; keep the hash's uniform + # pick rather than collapsing every token onto replica 0. + replica_idx = tl.where(total > 0.0, sampled, replica_idx) + + if HAS_PHYSICAL_IDS: + # The policy resolved the replica itself. Selection and mapping are + # skipped; recording below is not, which is the whole reason this + # path exists rather than the caller writing out_ids and losing the + # load accounting that shares this kernel. + physical_id = tl.load( + physical_ids_ptr + offs, mask=mask & valid_expert, other=-1 + ).to(tl.int64) + else: + map_index = safe_expert_id * map_slots + replica_idx + physical_id = tl.load( + logical_to_physical_ptr + map_index, + mask=mask & valid_expert, + other=-1, + ) tl.store(out_ids_ptr + offs, physical_id, mask=mask) # 2. Record expert load metrics. @@ -99,6 +152,8 @@ def _eplb_map_and_record_triton( expert_load_view: torch.Tensor, record_enabled: torch.Tensor, num_unpadded_tokens: torch.Tensor | None, + replica_prob: torch.Tensor | None = None, + physical_ids: torch.Tensor | None = None, ) -> torch.Tensor: topk_ids_in = topk_ids.contiguous().to(dtype=torch.int32) numel = topk_ids_in.numel() @@ -108,6 +163,41 @@ def _eplb_map_and_record_triton( out_flat = torch.empty((numel,), device=topk_ids.device, dtype=topk_ids.dtype) grid = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),) assert expert_load_view.is_contiguous() + map_slots = logical_to_physical_map.shape[1] + # The share table is scanned with a compile-time loop, so its width has + # to be the number of replicas that can exist -- not the width of the + # candidate map, which is padded to MAX_EXPERT_REDUNDANCY + 1 (1024) + # however few redundant experts are configured. Unrolling that twice is + # 2048 loads per element and takes Triton well past vLLM's RPC timeout + # to compile. A narrower table is a prefix of the same columns, so a + # sampled index still resolves against the full-width map below. + prob_slots = map_slots + if replica_prob is not None: + assert ( + replica_prob.shape[0] == logical_to_physical_map.shape[0] + and replica_prob.shape[1] <= map_slots + ), ( + "replica shares must be a column prefix of the candidate map: " + f"{tuple(replica_prob.shape)} vs " + f"{tuple(logical_to_physical_map.shape)}" + ) + prob_slots = replica_prob.shape[1] + logger.info_once( + "MLB replica-share routing active: table width %d, map width %d", + prob_slots, + map_slots, + ) + replica_prob = replica_prob.contiguous().to(torch.float32) + if physical_ids is not None: + assert physical_ids.shape == topk_ids.shape, ( + "resolved physical ids must match the TopK shape: " + f"{tuple(physical_ids.shape)} vs {tuple(topk_ids.shape)}" + ) + physical_ids = physical_ids.contiguous().to(torch.int32) + logger.info_once( + "MLB direct-dispatch routing active: the policy resolves " + "replicas itself; load recording stays in this kernel" + ) _eplb_map_and_record_i32_kernel[grid]( topk_ids_in, logical_replica_count.contiguous(), @@ -116,12 +206,17 @@ def _eplb_map_and_record_triton( expert_load_view, record_enabled, num_unpadded_tokens, + replica_prob, + physical_ids, logical_replica_count.shape[0], - logical_to_physical_map.shape[1], + map_slots, expert_load_view.shape[0], numel, num_active_experts, HAS_NUM_UNPADDED=num_unpadded_tokens is not None, + HAS_REPLICA_PROB=replica_prob is not None, + HAS_PHYSICAL_IDS=physical_ids is not None, + PROB_SLOTS=prob_slots, BLOCK_SIZE=256, ) return out_flat.reshape(topk_ids.shape) @@ -133,7 +228,45 @@ def eplb_map_to_physical_and_record( logical_replica_count: torch.Tensor, record_enabled: torch.Tensor, num_unpadded_tokens: torch.Tensor | None = None, + *, + layer_state: EplbLayerState | None = None, + topk_weights: torch.Tensor | None = None, ) -> torch.Tensor: + # A pluggable load balancer decides how each logical expert's traffic + # should be split across its replicas; it does not need to own the + # per-token step. Ask it for that split and keep the fused kernel -- + # which means load recording, the record gate and the padding mask all + # stay exactly as they are, with nothing reimplemented alongside them. + replica_prob = None + physical_ids = None + if layer_state is not None and topk_weights is not None: + from vllm.distributed.eplb.mlb_runtime import get_mlb_routing + + routing = get_mlb_routing() + # A policy that decides from the committed placement alone has + # nothing to add per forward: its answer was already materialised + # into a one-column candidate map when the placement committed, so + # substitute that map and skip the call. This is the same shape as + # SGLang's `static`, which is one tensor index inside its own TopK + # path -- there is no external call there to pay for, and there + # should not be one here either. Measured at 426 us per layer per + # forward, or 26 ms of a ~1 s forward, which is what made `static` + # trail the built-in selection it reproduces exactly. + _layer_id = getattr(layer_state, "moe_layer_idx", None) + fixed = ( + routing.fixed_dispatch_maps(_layer_id) + if routing is not None + and getattr(routing, "dispatch_fixed_by_placement", False) + and _layer_id is not None + else None + ) + if fixed is not None: + logical_to_physical_map, logical_replica_count = fixed + elif routing is not None and routing.requires_post_topk_routing: + replica_prob, physical_ids = routing.resolve_routing( + topk_ids, topk_weights, layer_state, num_unpadded_tokens + ) + # Fused triton implementation: mapping + optional recording in one kernel. return _eplb_map_and_record_triton( topk_ids=topk_ids, @@ -142,6 +275,8 @@ def eplb_map_to_physical_and_record( expert_load_view=expert_load_view, record_enabled=record_enabled, num_unpadded_tokens=num_unpadded_tokens, + replica_prob=replica_prob, + physical_ids=physical_ids, ) else: @@ -152,6 +287,9 @@ def eplb_map_to_physical_and_record( logical_replica_count: torch.Tensor, record_enabled: torch.Tensor, num_unpadded_tokens: torch.Tensor | None = None, + *, + layer_state: EplbLayerState | None = None, + topk_weights: torch.Tensor | None = None, ) -> torch.Tensor: return topk_ids @@ -180,6 +318,11 @@ def __init__( self.top_k = top_k self.global_num_experts = global_num_experts self.capture_fn: Callable[[torch.Tensor], None] | None = None + # Set by FusedMoEFactory when the layer dispatches its shared expert. + # None -- the default -- leaves every routing path byte-identical to + # what it was, which is what makes this safe to leave in the template + # method rather than behind a separate router subclass. + self.shared_expert_fusion: SharedExpertFusion | None = None def set_capture_fn(self, capture_fn: Callable[[torch.Tensor], None] | None) -> None: """Set a capture callback for logical routed expert IDs.""" @@ -200,7 +343,11 @@ def _validate_eplb_state(self) -> None: if eplb_state.num_unpadded_tokens_tensors is None: raise ValueError("EPLB requires num_unpadded_tokens_tensors != None") - def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor: + def _apply_eplb_mapping( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor | None = None, + ) -> torch.Tensor: """Apply EPLB mapping to convert logical expert IDs to physical expert IDs.""" if self.eplb_state is not None: eplb_state = self.eplb_state @@ -218,6 +365,8 @@ def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor: num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ dbo_current_ubatch_id() ], + layer_state=eplb_state, + topk_weights=topk_weights, ) return topk_ids @@ -297,9 +446,55 @@ def _select_experts( self.capture_fn(topk_ids) # Step 3: Apply EPLB mapping - topk_ids = self._apply_eplb_mapping(topk_ids) + topk_ids = self._apply_eplb_mapping(topk_ids, topk_weights) + + # Step 3.5: Give the shared expert a home rank and a top-k column. + # + # After the EPLB mapping, not before: the shared slot is already + # physical and has no logical identity to map, and EPLB's expert-load + # accounting covers routed experts only -- recording a shared-slot id + # in it would attribute traffic to a logical expert that does not + # exist. + if self.shared_expert_fusion is not None: + topk_ids, topk_weights = self._append_shared_expert(topk_ids, topk_weights) # Step 4: Convert indices dtype topk_ids = self._convert_indices_dtype(topk_ids, topk_indices_dtype) return topk_weights, topk_ids + + def _append_shared_expert( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Re-map routed ids into the widened space and append the shared slot. + + The home rank comes from an L2 policy when one routes the shared + expert (MLB's Waterfill), and is this rank otherwise. Falling back to + this rank rather than refusing to run keeps the fusion testable on its + own: it places the shared expert exactly where the replicated MLP did, + so any output difference is the dispatch detour and not a routing + decision. + """ + fusion = self.shared_expert_fusion + assert fusion is not None + + shared_rank = None + eplb_state = self.eplb_state + if ( + eplb_state is not None + and eplb_state.num_unpadded_tokens_tensors is not None + ): + from vllm.distributed.eplb.mlb_runtime import get_mlb_routing + + routing = get_mlb_routing() + if routing is not None and routing.routes_shared_expert: + shared_rank = routing.resolve_shared_expert_rank( + topk_ids, + topk_weights, + eplb_state, + eplb_state.num_unpadded_tokens_tensors[dbo_current_ubatch_id()], + ) + + return fusion.append_shared_expert(topk_ids, topk_weights, shared_rank) diff --git a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py index def871f428c7..8e9e370b3e19 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_moe_router.py @@ -2,12 +2,18 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from abc import ABC, abstractmethod from collections.abc import Callable +from typing import TYPE_CHECKING import torch from vllm.distributed.eplb.eplb_state import EplbLayerState from vllm.model_executor.layers.fused_moe.config import RoutingMethodType +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.shared_expert_fusion import ( + SharedExpertFusion, + ) + class FusedMoERouter(ABC): """FusedMoERouter is an abstract class that provides a 'select_experts' @@ -21,6 +27,10 @@ def __init__(self, eplb_state: EplbLayerState | None = None): # bias for image sentinel tokens; attached by model code when present. self.bias_vl: torch.Tensor | None = None self.image_sentinel_lo: int = 0 + # EP-dispatched shared-expert layout, attached by FusedMoEFactory when + # VLLM_FUSE_SHARED_EXPERTS (or a policy that needs it) is on; None + # otherwise. Consumed by BaseRouter after the EPLB mapping. + self.shared_expert_fusion: SharedExpertFusion | None = None @abstractmethod def set_capture_fn( diff --git a/vllm/model_executor/layers/fused_moe/shared_expert_fusion.py b/vllm/model_executor/layers/fused_moe/shared_expert_fusion.py new file mode 100644 index 000000000000..eef9f178980f --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/shared_expert_fusion.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""EP-dispatched shared experts. + +By default vLLM runs a MoE layer's shared expert as a plain MLP replicated on +every rank: every rank computes it for its own tokens and nothing crosses the +all-to-all. That is cheap, but it means a shared-expert token has no *choice* +of rank, so a load balancer has nothing to decide -- which is why MLB's +Waterfill L2 policy (``RoutingPhase.SHARED_EXPERT``) is inapplicable on vLLM +today. + +This module gives the shared expert a home rank so it can travel through the +same dispatch as the routed experts. The layout it uses is the one property +everything downstream depends on: + + rank r owns global expert ids [r*S, (r+1)*S) where S = R + n_shared + ids [r*S, r*S + R) -- that rank's R routed physical experts + ids [r*S + R, (r+1)*S) -- that rank's *own copy* of the shared expert + +i.e. every rank gets its own distinct global id for the shared expert. That is +what keeps DeepEP's rank derivation (``expert_id // (num_experts // ep_size)``, +hard-coded in its CUDA kernels) correct, and it is why this cannot reuse +vLLM's existing ``num_fused_shared_experts`` machinery: that one gives the +shared expert a *single* global id that every rank maps to a local slot, and +de-duplicates the resulting replication with a token-level round-robin mask. +That layout suits AITER's "every rank computes the full batch locally, then +reduce" model and is unroutable under dispatch/combine -- ROCm's own MoRI +backend asserts it off for exactly this reason. + +Because the routed experts no longer occupy a contiguous ``[0, R*ep_size)`` +range, their physical ids must be re-mapped into the widened space before +dispatch. ``remap_routed_ids`` does that; ``append_shared_expert`` does both +the re-map and the extra top-k column in one go. +""" + +from dataclasses import dataclass + +import torch + +import vllm.envs as envs +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class SharedExpertFusion: + """Geometry of the widened expert layout, and the ops that use it. + + Immutable and built once per MoE layer: every field is fixed at load time, + which matters because ``num_local_experts`` sizes the weight tensors and + ``global_num_experts`` is baked into the DeepEP call. + """ + + ep_size: int + ep_rank: int + num_shared_experts: int + """Shared experts per rank -- ``config.n_shared_experts``, not times ep_size.""" + routed_slots_per_rank: int + """Physical routed experts per rank: ``(n_routed + n_redundant) // ep_size``.""" + shared_expert_weight: float + """Top-k weight for the shared slot. 1/routed_scaling_factor when the runner + scales the combined output, so the shared expert's net contribution stays + 1.0 -- the same correction the AITER path applies.""" + + @property + def slots_per_rank(self) -> int: + return self.routed_slots_per_rank + self.num_shared_experts + + @property + def global_num_experts(self) -> int: + """Widened physical expert count handed to DeepEP and ``expert_map``.""" + return self.slots_per_rank * self.ep_size + + @property + def num_routed_physical_experts(self) -> int: + """Expert count in the *un-widened* space the router still works in.""" + return self.routed_slots_per_rank * self.ep_size + + def shared_expert_global_ids(self, ep_rank: int) -> list[int]: + """Global ids of the shared-expert slots owned by ``ep_rank``.""" + base = ep_rank * self.slots_per_rank + self.routed_slots_per_rank + return [base + j for j in range(self.num_shared_experts)] + + def remap_routed_ids(self, topk_ids: torch.Tensor) -> torch.Tensor: + """Lift routed physical ids into the widened space. + + ``id -> id + (id // R) * n_shared``: rank r's block ``[r*R, r*R+R)`` + slides to ``[r*S, r*S+R)``, leaving its shared slots free at the end. + Negative (invalid) ids are passed through untouched. + """ + shifted = topk_ids + (topk_ids // self.routed_slots_per_rank) * ( + self.num_shared_experts + ) + return torch.where(topk_ids >= 0, shifted, topk_ids) + + def append_shared_expert( + self, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + shared_expert_rank: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Re-map the routed columns and add the shared-expert column(s). + + ``shared_expert_rank`` is the per-token home rank, as chosen by an L2 + policy. ``None`` means "this rank", which reproduces the replicated + MLP's placement exactly while still routing through dispatch -- that is + the configuration to A/B against, since it isolates the cost of the + detour from the effect of the balancing decision. + """ + num_tokens = topk_ids.shape[0] + ids = self.remap_routed_ids(topk_ids) + + if shared_expert_rank is None: + home = torch.full( + (num_tokens, 1), + self.ep_rank, + dtype=ids.dtype, + device=ids.device, + ) + else: + home = shared_expert_rank.reshape(num_tokens, 1).to(ids.dtype) + + offsets = torch.arange( + self.num_shared_experts, dtype=ids.dtype, device=ids.device + ) + # A policy answers -1 for a token it will not place -- a padded row, or + # one whose routed ids are all invalid. Arithmetic on that lands + # somewhere real: -1 * slots + routed + 0 is -1 only because the two + # terms happen to cancel at num_shared_experts == 1, and at 2 it yields + # -2, which torch reads as an index from the end rather than as "no + # expert". Carry the sentinel through explicitly instead of relying on + # that, and zero the weight so a row that reaches a kernel anyway + # contributes nothing. remap_routed_ids above guards the routed columns + # the same way; this is the matching guard for the shared one. + placed = home >= 0 + shared_ids = torch.where( + placed, + home * self.slots_per_rank + self.routed_slots_per_rank + offsets, + torch.full_like(home, -1), + ) + shared_weights = torch.where( + placed, + torch.full( + (num_tokens, self.num_shared_experts), + self.shared_expert_weight, + dtype=topk_weights.dtype, + device=topk_weights.device, + ), + torch.zeros( + (num_tokens, self.num_shared_experts), + dtype=topk_weights.dtype, + device=topk_weights.device, + ), + ) + return ( + torch.cat((ids, shared_ids), dim=-1), + torch.cat((topk_weights, shared_weights), dim=-1), + ) + + +def maybe_build_shared_expert_fusion( + *, + n_shared_experts: int | None, + num_physical_experts: int, + ep_size: int, + ep_rank: int, + use_ep: bool, + is_act_and_mul: bool, + shared_expert_weight: float, + layer_name: str = "", + warn_on_uneven: bool = True, +) -> SharedExpertFusion | None: + """Build the geometry when this layer can and should dispatch its shared expert. + + Returns ``None`` -- meaning "nothing changes" -- for every configuration + the widened layout cannot serve, rather than raising: a shared expert is + not required to be fusible, and the caller falls back to the replicated + MLP. The one case worth a log is an uneven split, because that one looks + like it should have worked. + """ + if not shared_expert_fusion_enabled(): + return None + if not n_shared_experts: + return None + if not is_act_and_mul: + # The shared expert has to be the same gated-MLP shape as a routed one + # to occupy an expert slot at all. + return None + if not use_ep or ep_size <= 1: + # Without EP every rank is the only rank: dispatching the shared expert + # would be a no-op detour, and there is no rank for a policy to choose. + return None + if num_physical_experts % ep_size != 0: + if not warn_on_uneven: + return None + logger.warning( + "Shared-expert fusion is on but %s has %d physical experts " + "over %d EP ranks, which does not divide evenly. Falling back to " + "the replicated shared-expert MLP.", + layer_name or "this MoE layer", + num_physical_experts, + ep_size, + ) + return None + + return SharedExpertFusion( + ep_size=ep_size, + ep_rank=ep_rank, + num_shared_experts=int(n_shared_experts), + routed_slots_per_rank=num_physical_experts // ep_size, + shared_expert_weight=shared_expert_weight, + ) + + +def _l2_pipeline_needs_dispatched_shared_expert() -> bool: + """Whether the configured MLB L2 pipeline picks the shared expert's rank.""" + from vllm.distributed.eplb.mlb_runtime import mlb_l2_algorithm + + algorithm = mlb_l2_algorithm() + if not algorithm: + return False + try: + from moe_load_balancer.core.routing_pipeline import RoutingPipeline + + caps = RoutingPipeline.from_value(algorithm).capabilities + except Exception: + return False + return bool(getattr(caps, "routes_shared_expert", False)) + + +def shared_expert_fusion_enabled() -> bool: + """Whether the shared expert is dispatched rather than replicated per rank. + + The env switch turns it on, and so does asking for an L2 pipeline that + routes the shared expert: waterfill picks a rank for it, and there is no + rank to pick unless it goes through dispatch. Selecting that pipeline and + leaving the switch off used to leave MLB's L2 quietly disabled -- the + configuration was refused, a line went by in the log, and the run carried + on doing none of what was asked for. Turning it on and saying so is how + SGLang resolves the same pair (`enable_deepep_waterfill` forces shared + expert fusion there and warns), and a loud change of configuration beats a + silent cancellation of one. + + Says nothing about whether any particular layer *can* fuse -- that is + ``maybe_build_shared_expert_fusion``'s answer, and it needs a layer's + geometry to give it. + """ + if envs.VLLM_FUSE_SHARED_EXPERTS: + return True + if _l2_pipeline_needs_dispatched_shared_expert(): + logger.warning_once( + "Enabling shared-expert fusion: the configured MLB L2 pipeline " + "(VLLM_MLB_L2_ALGORITHM) assigns the shared expert to a rank, " + "which requires it to go through expert-parallel dispatch. Set " + "VLLM_FUSE_SHARED_EXPERTS=1 to make this explicit, or drop " + "waterfill from the pipeline to keep the shared expert " + "replicated per rank." + ) + return True + return False diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index d5728e0c3365..a4a2985a2582 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -58,6 +58,10 @@ GateLinear, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.shared_expert_fusion import ( + SharedExpertFusion, + maybe_build_shared_expert_fusion, +) from vllm.model_executor.layers.fused_moe.utils import ( is_model_fused_shared_expert_compatible, resolve_layer_fused_shared_expert, @@ -288,6 +292,30 @@ def forward(self, x): return x +def _find_shared_expert_fusion(model: nn.Module) -> SharedExpertFusion | None: + """The EP shared-expert layout this model was built with, if any. + + Every MoE layer agrees on it (it is derived from parallel geometry, not + from anything per-layer), so the first one found answers for all. Returns + None on a model whose layers were built with the feature off, which is what + keeps `load_weights` on its original path. + """ + # The only caller is DeepseekV2Model.load_weights, which passes the model + # whose `layers` this is -- not the ForCausalLM wrapper around it. Accept + # either, because reading `model.model.layers` off a DeepseekV2Model + # silently returns None, and "no fusion" is exactly what the switch being + # off looks like: the shared-expert checkpoint tensors then take the + # replicated-MLP path to a module that was never built. + layers = getattr(model, "layers", None) + if layers is None: + layers = getattr(getattr(model, "model", None), "layers", None) + for layer in layers or (): + fusion = getattr(getattr(layer, "mlp", None), "shared_expert_fusion", None) + if fusion is not None: + return fusion + return None + + class DeepseekV2MoE(nn.Module): def __init__( self, @@ -305,6 +333,7 @@ def __init__( self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.ep_group = get_ep_group().device_group + self.ep_rank = get_ep_group().rank_in_group self.ep_size = self.ep_group.size() self.n_routed_experts: int = config.n_routed_experts self.n_shared_experts: int = config.n_shared_experts @@ -341,12 +370,30 @@ def __init__( self.n_local_physical_experts = self.n_physical_experts // self.ep_size self.is_rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - self.is_fused_shared_expert_enabled = False if config.n_shared_experts is not None: self.is_fused_shared_expert_enabled = resolve_layer_fused_shared_expert( quant_config, prefix ) + # EP-dispatched shared expert (VLLM_FUSE_SHARED_EXPERTS). Rebuilt here + # rather than read back off the layer: this decides whether the + # replicated MLP is constructed at all, which has to happen before the + # layer exists. The factory reaches the same verdict from the same + # inputs, and asserts below that it did. + self.shared_expert_fusion = maybe_build_shared_expert_fusion( + n_shared_experts=config.n_shared_experts, + num_physical_experts=self.n_physical_experts, + ep_size=self.ep_size, + ep_rank=self.ep_rank, + use_ep=parallel_config.enable_expert_parallel, + is_act_and_mul=config.hidden_act == "silu", + shared_expert_weight=( + (1.0 / self.routed_scaling_factor) + if (apply_routed_scale_to_output and self.routed_scaling_factor) + else 1.0 + ), + warn_on_uneven=False, + ) if ( self.is_rocm_aiter_moe_enabled @@ -356,7 +403,11 @@ def __init__( # Accumulates in fp32; avoids bf16->fp32 cast. self.gate.set_out_dtype(self.gate.weight.dtype) - if config.n_shared_experts is None or self.is_fused_shared_expert_enabled: + if ( + config.n_shared_experts is None + or self.is_fused_shared_expert_enabled + or self.shared_expert_fusion is not None + ): self.shared_experts = None else: intermediate_size = config.moe_intermediate_size * config.n_shared_experts @@ -393,12 +444,36 @@ def __init__( is_sequence_parallel=self.is_sequence_parallel, reduce_results=reduce_results, n_shared_experts=config.n_shared_experts - if self.is_fused_shared_expert_enabled + if ( + self.is_fused_shared_expert_enabled + or self.shared_expert_fusion is not None + ) else None, fuse_shared_experts=self.is_fused_shared_expert_enabled, router_logits_dtype=self.gate.out_dtype, ) + # The two verdicts must agree: this one decided whether to build the + # replicated MLP and where load_weights puts the checkpoint's shared + # expert, that one decided the expert-id space the router addresses. + # A silent disagreement writes the shared expert into a slot no token + # is ever routed to, which shows up as a quality regression and + # nothing else. + built_fusion = getattr( + getattr(self.experts, "router", None), "shared_expert_fusion", None + ) + if (built_fusion is None) != (self.shared_expert_fusion is None): + raise RuntimeError( + "Shared-expert fusion disagreement: the model decided " + f"{self.shared_expert_fusion is not None} and the MoE layer " + f"decided {built_fusion is not None}." + ) + if built_fusion is not None and built_fusion != self.shared_expert_fusion: + raise RuntimeError( + f"Shared-expert fusion geometry disagreement: model has " + f"{self.shared_expert_fusion}, layer has {built_fusion}." + ) + if ( self.is_rocm_aiter_moe_enabled and self.gate.e_score_correction_bias is not None @@ -1590,6 +1665,14 @@ def forward( return hidden_states def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # The EP-dispatched layout (VLLM_FUSE_SHARED_EXPERTS). Read off a built + # layer rather than recomputed, so the slot the checkpoint's shared + # expert lands in is by construction the one the router will send + # tokens to. + ep_shared_fusion = _find_shared_expert_fusion(self) + fuse_shared_into_experts = ( + self.is_fused_shared_expert_enabled or ep_shared_fusion is not None + ) stacked_params_mapping: list[tuple[str, str, int | str]] = [ # (param_name, shard_name, shard_id) ("gate_up_proj", "gate_proj", 0), @@ -1657,8 +1740,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ): continue # this layer has no indexer; drop its checkpoint weights - is_fusion_moe_shared_experts_layer = ( - self.is_fused_shared_expert_enabled and ("mlp.shared_experts" in name) + is_fusion_moe_shared_experts_layer = fuse_shared_into_experts and ( + "mlp.shared_experts" in name ) if _try_load_fp8_indexer_wk( @@ -1752,7 +1835,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # can route it chunk_name = name.replace( "mlp.shared_experts", - f"mlp.experts.{self.config.n_routed_experts + j}", + "mlp.experts.0" + if ep_shared_fusion is not None + else f"mlp.experts.{self.config.n_routed_experts + j}", ) # Use expert_params_mapping to locate the destination @@ -1767,6 +1852,28 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # attempted to load as other weights later is_expert_weight = True + if ep_shared_fusion is not None: + if is_fusion_moe_shared_experts_layer: + # This rank's own copy, in its own slot. Only + # this rank's id is written: every other rank + # loads its own from the same checkpoint + # tensor, and `weight_loader` rejects any id + # that is not local to it anyway. + expert_id = ep_shared_fusion.shared_expert_global_ids( + ep_shared_fusion.ep_rank + )[j] + else: + # Routed experts are no longer contiguous: + # rank r's block starts at r*S, not r*R. + expert_id = ( + expert_id + + ( + expert_id + // ep_shared_fusion.routed_slots_per_rank + ) + * ep_shared_fusion.num_shared_experts + ) + # Do not modify `name` since the loop may continue here # Instead, create a new variable name_mapped = chunk_name.replace(weight_name, param_name) diff --git a/vllm/models/deepseek_v32/nvidia/model.py b/vllm/models/deepseek_v32/nvidia/model.py index b42ba28c8df8..3bf6e6131ad2 100644 --- a/vllm/models/deepseek_v32/nvidia/model.py +++ b/vllm/models/deepseek_v32/nvidia/model.py @@ -27,6 +27,7 @@ DeepseekV2ForCausalLM, DeepseekV2MLP, DeepseekV2MoE, + _find_shared_expert_fusion, _try_load_fp8_indexer_wk, get_spec_layer_idx_from_weight_name, ) @@ -326,6 +327,16 @@ def forward( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # DSA-only: MLA (fused_qkv_a_proj) + the fused indexer wk/weights_proj + # routed experts. No MHA (qkv_proj) or ROCm shared-expert-fusion paths. + # + # The EP-dispatched shared expert (VLLM_FUSE_SHARED_EXPERTS) is handled: + # when the MoE layers were built with it, the replicated `shared_experts` + # MLP does not exist and the checkpoint's shared-expert tensors go into + # this rank's fused expert slot, while routed ids shift past the shared + # slots of every earlier rank -- the same routing as + # DeepseekV2Model.load_weights, mirrored here because this flat model + # has its own loader. Read off a built layer, so the slot the weight + # lands in is by construction the one the router sends tokens to. + ep_shared_fusion = _find_shared_expert_fusion(self) stacked_params_mapping = [ ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), @@ -363,12 +374,17 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ): continue + is_fused_shared_expert_weight = ( + ep_shared_fusion is not None and "mlp.shared_experts" in name + ) for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue # Experts are handled below; skip here before the name rewrite. if ("mlp.experts." in name) and name not in params_dict: continue + if is_fused_shared_expert_weight: + continue name_mapped = name.replace(weight_name, param_name) if ( param_name == "fused_qkv_a_proj" @@ -384,12 +400,34 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: break else: is_expert_weight = False + # The shared expert's tensor is routed like an expert weight, + # under an expert-style name the mapping recognises. + expert_name = ( + name.replace("mlp.shared_experts", "mlp.experts.0") + if is_fused_shared_expert_weight + else name + ) for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping # type: ignore[assignment] - if weight_name not in name: + if weight_name not in expert_name: continue is_expert_weight = True - name_mapped = name.replace(weight_name, param_name) + if ep_shared_fusion is not None: + if is_fused_shared_expert_weight: + # This rank's own copy, in its own slot; every rank + # loads its own from the same checkpoint tensor. + expert_id = ep_shared_fusion.shared_expert_global_ids( + ep_shared_fusion.ep_rank + )[0] + else: + # Routed experts are no longer contiguous: rank r's + # block starts at r*S, not r*R. + expert_id = ( + expert_id + + (expert_id // ep_shared_fusion.routed_slots_per_rank) + * ep_shared_fusion.num_shared_experts + ) + name_mapped = expert_name.replace(weight_name, param_name) if is_pp_missing_parameter(name_mapped, self): continue param = params_dict[name_mapped] diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 7edff8ae0202..3d6f20821131 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -729,6 +729,8 @@ def forward( ] if eplb_state.num_unpadded_tokens_tensors is not None else None, + layer_state=eplb_state, + topk_weights=topk_weights, ) shared_x_sf = None diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index c1818a444b53..28c46fbd65a3 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -480,6 +480,8 @@ def forward( ] if eplb_state.num_unpadded_tokens_tensors is not None else None, + layer_state=eplb_state, + topk_weights=topk_weights, ) prepare_megamoe_inputs( diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index 18ae2700e937..c43142fc0fb1 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/vllm/v1/spec_decode/extract_hidden_states.py @@ -210,7 +210,7 @@ def _determine_batch_execution_and_padding( # TODO(Flechman): support DBO ubatching should_ubatch, num_tokens_across_dp = False, None if self.vllm_config.parallel_config.data_parallel_size > 1: - should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = ( + should_ubatch, num_tokens_across_dp, synced_cudagraph_mode, _ = ( coordinate_batch_across_dp( num_tokens_unpadded=num_tokens, parallel_config=self.vllm_config.parallel_config, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index e4718bb3d7d1..ea0fe4de7e62 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -1797,7 +1797,7 @@ def _determine_batch_execution_and_padding( # TODO(Flechman): support DBO ubatching should_ubatch, num_tokens_across_dp = False, None if self.vllm_config.parallel_config.data_parallel_size > 1: - should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = ( + should_ubatch, num_tokens_across_dp, synced_cudagraph_mode, _ = ( coordinate_batch_across_dp( num_tokens_unpadded=num_tokens, parallel_config=self.vllm_config.parallel_config, diff --git a/vllm/v1/worker/dp_utils.py b/vllm/v1/worker/dp_utils.py index b2f2fe1506a9..34a1f58d02b6 100644 --- a/vllm/v1/worker/dp_utils.py +++ b/vllm/v1/worker/dp_utils.py @@ -59,6 +59,7 @@ def _run_ar( orig_num_tokens_per_ubatch: int, padded_num_tokens_per_ubatch: int, cudagraph_mode: int, + uniform_decode: bool, parallel_config: ParallelConfig, ) -> torch.Tensor: dp_size = parallel_config.data_parallel_size @@ -66,16 +67,34 @@ def _run_ar( device, group = _get_device_and_group(parallel_config) # Populate this rank's contribution on CPU to reduce GPU syncs. pin_memory = PIN_MEMORY and torch.device(device).type == "cuda" - tensor_cpu = torch.zeros(4, dp_size, dtype=torch.int32, pin_memory=pin_memory) + tensor_cpu = torch.zeros(5, dp_size, dtype=torch.int32, pin_memory=pin_memory) tensor_cpu[0][dp_rank] = orig_num_tokens_per_ubatch tensor_cpu[1][dp_rank] = padded_num_tokens_per_ubatch tensor_cpu[2][dp_rank] = 1 if should_ubatch else 0 tensor_cpu[3][dp_rank] = cudagraph_mode + tensor_cpu[4][dp_rank] = 1 if uniform_decode else 0 tensor = tensor_cpu.to(device, non_blocking=True) dist.all_reduce(tensor, group=group) return tensor +def _post_process_uniform_decode(tensor: torch.Tensor) -> bool: + """Whether *every* DP rank is running a single-token decode this step. + + Each rank knows only its own batch shape, but a policy that issues a + collective from inside the forward -- expert-placement refresh, say -- + has to reach the same verdict on every rank or the ranks that took the + branch wait in their collective for ranks that did not. A rank with no + requests runs `_dummy_run(uniform_decode=True)` and looks like a decode + to itself while its peers run real prefill, so the local flag is exactly + the wrong thing to gate on. Reduced with `all`, not `any`: one rank still + prefilling makes the step globally not-decode, which is the conservative + direction -- it keeps the step eligible for whatever work the local flag + would have skipped, on every rank alike. + """ + return bool(torch.all(tensor[4] == 1).item()) + + def _post_process_ubatch(tensor: torch.Tensor, num_ubatches: int) -> bool: orig_num_tokens_tensor = tensor[0, :] padded_num_tokens_tensor = tensor[1, :] @@ -124,8 +143,9 @@ def _synchronize_dp_ranks( num_tokens_padded: int, should_attempt_ubatching: bool, cudagraph_mode: int, + uniform_decode: bool, parallel_config: ParallelConfig, -) -> tuple[bool, torch.Tensor | None, int]: +) -> tuple[bool, torch.Tensor | None, int, bool]: """1. Decides if each DP rank is going to microbatch. Either all ranks run with microbatching or none of them do. @@ -135,11 +155,14 @@ def _synchronize_dp_ranks( 3. Synchronizes cudagraph_mode across ranks by taking the minimum. + 4. Synchronizes uniform_decode across ranks by taking the conjunction. + Returns: tuple[ should_ubatch: Are all DP ranks going to microbatch num_tokens_after_padding: A tensor containing the total number of tokens per-microbatch for each DP rank including any DP padding. synced_cudagraph_mode: The synchronized cudagraph mode (min across ranks) + synced_uniform_decode: Whether every DP rank is running a decode ] """ @@ -153,6 +176,7 @@ def _synchronize_dp_ranks( orig_num_tokens_per_ubatch=num_tokens_unpadded, padded_num_tokens_per_ubatch=num_tokens_padded, cudagraph_mode=cudagraph_mode, + uniform_decode=uniform_decode, parallel_config=parallel_config, ) @@ -184,7 +208,14 @@ def _synchronize_dp_ranks( # should_dp_pad is True num_tokens_after_padding = _post_process_dp_padding(tensor, should_dp_pad) - return should_ubatch, num_tokens_after_padding, synced_cudagraph_mode + synced_uniform_decode = _post_process_uniform_decode(tensor) + + return ( + should_ubatch, + num_tokens_after_padding, + synced_cudagraph_mode, + synced_uniform_decode, + ) def coordinate_batch_across_dp( @@ -194,7 +225,7 @@ def coordinate_batch_across_dp( num_tokens_padded: int | None = None, uniform_decode: bool | None = None, cudagraph_mode: int = 0, -) -> tuple[bool, torch.Tensor | None, int]: +) -> tuple[bool, torch.Tensor | None, int, bool]: """Coordinates amongst all DP ranks to determine if and how the full batch should be split into microbatches. @@ -204,8 +235,10 @@ def coordinate_batch_across_dp( parallel_config: The parallel config num_tokens_padded: Number of tokens including any non-DP padding (CUDA graphs, TP, etc) - uniform_decode: Only used if allow_microbatching is True. True if the batch - only contains single token decodes + uniform_decode: True if this rank's batch only contains single token + decodes. Used for the microbatching precondition check when + allow_microbatching is True, and -- independently of that -- + reduced across ranks into the returned synced_uniform_decode. cudagraph_mode: The cudagraph mode for this rank (0=NONE, 1=PIECEWISE, 2=FULL). DP padding is enabled when synced cudagraph mode across ranks is not NONE. @@ -216,12 +249,15 @@ def coordinate_batch_across_dp( tokens per-microbatch for each DP rank including padding. Will be padded up to the max value across all DP ranks when cudagraph is enabled. synced_cudagraph_mode: The synchronized cudagraph mode (min across ranks) + synced_uniform_decode: True only if every DP rank is running a decode. + False outside DP, where there is no peer to disagree with the caller's + own view of its batch. ] """ if parallel_config.data_parallel_size == 1: # Early exit. - return False, None, cudagraph_mode + return False, None, cudagraph_mode, bool(uniform_decode) # If the caller has explicitly enabled microbatching. should_attempt_ubatching = False @@ -248,16 +284,31 @@ def coordinate_batch_across_dp( num_tokens_padded, dtype=torch.int32, ) - return should_ubatch, num_tokens_after_padding, cudagraph_mode - - (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) = ( - _synchronize_dp_ranks( - num_tokens_unpadded, - num_tokens_padded, - should_attempt_ubatching, + # No peer to disagree with: this rank's own view of its batch stands. + return ( + should_ubatch, + num_tokens_after_padding, cudagraph_mode, - parallel_config, + bool(uniform_decode), ) + + ( + should_ubatch, + num_tokens_after_padding, + synced_cudagraph_mode, + synced_uniform_decode, + ) = _synchronize_dp_ranks( + num_tokens_unpadded, + num_tokens_padded, + should_attempt_ubatching, + cudagraph_mode, + bool(uniform_decode), + parallel_config, ) - return (should_ubatch, num_tokens_after_padding, synced_cudagraph_mode) + return ( + should_ubatch, + num_tokens_after_padding, + synced_cudagraph_mode, + synced_uniform_decode, + ) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 3dc0524babe2..3383e7e129a3 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -3935,6 +3935,7 @@ def _determine_batch_execution_and_padding( bool, torch.Tensor | None, CUDAGraphStat | None, + bool, ]: uniform_decode = self._is_uniform_decode( max_num_scheduled_tokens=max_num_scheduled_tokens, @@ -3986,16 +3987,22 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): # Extra coordination when running data-parallel since we need to coordinate # across ranks should_ubatch, num_tokens_across_dp = False, None + # Outside DP this rank is the only rank: its own view of its batch is + # by definition the global one, and no peer can disagree with it. + uniform_decode_across_dp = uniform_decode if self.vllm_config.parallel_config.data_parallel_size > 1: - should_ubatch, num_tokens_across_dp, synced_cudagraph_mode = ( - coordinate_batch_across_dp( - num_tokens_unpadded=num_tokens, - parallel_config=self.parallel_config, - allow_microbatching=allow_microbatching, - num_tokens_padded=num_tokens_padded, - uniform_decode=uniform_decode, - cudagraph_mode=cudagraph_mode.value, - ) + ( + should_ubatch, + num_tokens_across_dp, + synced_cudagraph_mode, + uniform_decode_across_dp, + ) = coordinate_batch_across_dp( + num_tokens_unpadded=num_tokens, + parallel_config=self.parallel_config, + allow_microbatching=allow_microbatching, + num_tokens_padded=num_tokens_padded, + uniform_decode=uniform_decode, + cudagraph_mode=cudagraph_mode.value, ) # Extract DP-synced values @@ -4026,6 +4033,7 @@ def dispatch_cudagraph(num_tokens, disable_full=False, valid_modes=None): should_ubatch, num_tokens_across_dp, cudagraph_stats, + uniform_decode_across_dp, ) def _register_layerwise_nvtx_hooks(self) -> None: @@ -4247,6 +4255,7 @@ def execute_model( should_ubatch, num_tokens_across_dp, cudagraph_stats, + uniform_decode_across_dp, ) = self._determine_batch_execution_and_padding( num_tokens=num_tokens_unpadded, num_reqs=num_reqs, @@ -4411,6 +4420,7 @@ def execute_model( self.vllm_config, num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, + uniform_decode_across_dp=uniform_decode_across_dp, cudagraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded, @@ -5905,29 +5915,33 @@ def _dummy_run( num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) - _cudagraph_mode, batch_desc, should_ubatch, num_tokens_across_dp, _ = ( - self._determine_batch_execution_and_padding( - num_tokens=num_tokens_unpadded, - num_reqs=num_reqs, - num_scheduled_tokens_np=num_scheduled_tokens, - max_num_scheduled_tokens=max_query_len, - use_cascade_attn=False, - allow_microbatching=allow_microbatching, - force_eager=is_profile - or (cudagraph_runtime_mode == CUDAGraphMode.NONE), - # `force_uniform_decode` is used for cudagraph capture; because for - # capturing mixed prefill-decode batches, we sometimes use - # num_tokens == num_reqs which looks like a uniform decode batch to the - # dispatcher; but we actually want to capture a piecewise cudagraph - force_uniform_decode=uniform_decode, - # `force_has_lora` is used for cudagraph capture; because LoRA is - # activated later in the context manager, but we need to know the - # LoRA state when determining the batch descriptor for capture - force_has_lora=num_active_loras > 0, - # `force_num_active_loras` is used for cudagraph capture; because we - # need to capture graphs for specific num_active_loras counts - force_num_active_loras=num_active_loras, - ) + ( + _cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + _, + uniform_decode_across_dp, + ) = self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens, + max_num_scheduled_tokens=max_query_len, + use_cascade_attn=False, + allow_microbatching=allow_microbatching, + force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE), + # `force_uniform_decode` is used for cudagraph capture; because for + # capturing mixed prefill-decode batches, we sometimes use + # num_tokens == num_reqs which looks like a uniform decode batch to the + # dispatcher; but we actually want to capture a piecewise cudagraph + force_uniform_decode=uniform_decode, + # `force_has_lora` is used for cudagraph capture; because LoRA is + # activated later in the context manager, but we need to know the + # LoRA state when determining the batch descriptor for capture + force_has_lora=num_active_loras > 0, + # `force_num_active_loras` is used for cudagraph capture; because we + # need to capture graphs for specific num_active_loras counts + force_num_active_loras=num_active_loras, ) if cudagraph_runtime_mode is None: @@ -6114,6 +6128,38 @@ def _dummy_run( if num_tokens_across_dp is not None: num_tokens_across_dp[:] = num_tokens_padded + # A routing policy that reads the EP-wide logical count issues one + # collective per forward, and prepare_forward -- where that + # collective lives on the real path -- is not called here. Under DP + # every rank must run a forward every step, so a rank with no + # requests runs this dummy batch while its peers run real ones. + # Issuing the collective on only one of those two paths leaves the + # EP group out of order: a peer's all_reduce meets this rank's next + # dispatch and both wait until DeepEP's CPU-recv timeout fires. + # Upstream already replays eplb_step here for the same reason. + # + # Only the collective is replayed. prepare_forward also drains + # async placement commits (rank-local, so asymmetry there cannot + # deadlock) and fills the unpadded-token tensors, which a dummy + # batch has no use for. + # + # Gated on skip_eplb the same way eplb_step() below is: it is + # True for _warmup_and_capture()'s dummy runs (CUDA graph capture + # warmup), where every rank captures independently -- there is no + # synchronized peer step for this replay to stay in order with, + # so issuing it would be the asymmetric call, not the missing + # one. + if not skip_eplb and self.eplb_state is not None: + from vllm.distributed.eplb.mlb_runtime import get_mlb_routing + + routing = get_mlb_routing() + if routing is not None: + # Not finalize_step_counts: that publishes the counts and + # clears the local buffer, so running it here as well as on + # the real path would overwrite the counts with zeros. This + # only matches the collective. + routing.match_step_counts_collective() + with ( self.maybe_randomize_inputs( input_ids, inputs_embeds, randomize_inputs=randomize_inputs @@ -6123,6 +6169,7 @@ def _dummy_run( self.vllm_config, num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, + uniform_decode_across_dp=uniform_decode_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, ubatch_slices=ubatch_slices_padded,