[zanes-ops]

Motivation

With an a2a backend (for example --moe-a2a-backend deepep), a server started with --init-expert-location set to a physical_to_logical_map in which one EP rank holds two replicas of the same logical expert never loads the second replica. A rank that holds no replica of that expert can still dispatch tokens to that slot, so those tokens go through an expert whose weights were never written. There is no error; the outputs just change. Reloading the weights through the same loader (for example with update_weights_from_disk) skips the same slots.

Such layouts come out of EPLB: the packing step can put two replicas of a hot expert on one GPU (deepseek-ai/EPLB#22), and SGLang's copy of the algorithm does too. Passing a saved, rebalanced layout back through --init-expert-location therefore hits this. In two layouts that we saved from Qwen3-30B-A3B servers (EP2, 2 redundant experts) after online rebalancing, 16 and 23 of the 48 layers had such a pair.

Where it happens (main @ 78185f4):

1. _compute_logical_to_all_physical_map: with an a2a backend, each rank replaces every logical expert's candidate list with _find_nearest_expert (EPLB: prefer to use physical experts in the same gpu or node #10874). A rank that holds replicas keeps only its first local one. A rank that holds none keeps the full list (on one node).
2. FusedMoE.weight_loader writes a checkpoint expert into the slots that logical_to_all_physical() returns, which is this narrowed list, and _weight_loader_physical keeps the local ones. The rank that holds both replicas writes only the first. For unquantized weights the second one stays torch.empty.
3. The static dispatch table gives a rank without a local replica a _fair_choices pick among all replicas, which can be the unloaded one. static is the default with an a2a backend; with --ep-dispatch-algorithm dynamic, the same list is sampled per token.

The narrowing has left rows unloaded before: #12866 removed it to fix NaN NVFP4 input scales, and #13162 / #13484 restored it with a separate path for those global scale parameters (require_global_experts). That relies on the other unloaded rows never being dispatched to, and a second replica on the same rank breaks it. #32962 skipped the narrowing when there is no a2a backend; the a2a path is unchanged.

Modifications

• ExpertLocationMetadata.logical_to_all_physical() (eplb/expert_location.py): for a routed expert, it now returns every physical slot that physical_to_logical_map assigns to the expert, on all ranks, instead of reading the per-rank dispatch list logical_to_all_physical_map.
  • The draft-layer fallback and the require_global_experts branch are unchanged. They only read the map's shape, and logical_to_all_physical_map is padded to the number of physical slots, so both maps give the same values.
  • For a layer in the map and without require_global_experts, the lookup rejects out-of-range logical expert ids with IndexError, including negative ids, which used to index from the end. The draft-layer and global-scale branches keep their existing conventions.
  • The dispatch fields (logical_to_all_physical_map, logical_to_rank_dispatch_physical_map) are not touched, so the nearest-replica dispatch from EPLB: prefer to use physical experts in the same gpu or node #10874 stays as it is.
• FusedMoE._local_routed_expert_ids() (layers/moe/fused_moe_triton/layer.py, new) keeps the slots in this rank's routed range. FusedMoE.weight_loader applies it to a routed expert's slots (not with require_global_experts) before _weight_loader_physical maps them to local ids.
  • This is needed because the list now also holds other ranks' slots. After an elastic EP scale-up, a new rank's routed slot can lie past the _num_global_routed that an existing rank's layer set at startup, and _map_global_expert_id_to_local_expert_id maps such a slot onto that rank's fused shared expert.
  • From reading the code and a CPU test with a stand-in layer, this already happens on main for a rank that holds no replica of the expert, because its narrowed list is the full list. The filter covers that case too.
• moe_wna16._local_expert_ids (layers/quantization/moe_wna16.py) writes qzeros outside the layer's loader, so it applies the same filter to routed experts. A fused shared expert's id, which the draft-layer fallback returns, is not filtered.
• The third caller, elastic_ep/expert_backup_client.py, already keeps only its own range and is unchanged.
• New CPU test file test/registered/unit/eplb/test_expert_location_weight_loading.py (base-a-test-cpu, 8 tests). It is a new file because it drives the real FusedMoE and moe_wna16 loaders through a stand-in layer, unlike the existing dispatch-table tests. The tests check that:
  • the lookup returns exactly each expert's slots: a same-rank pair, a two-node layout where a rank narrows to a same-node replica, random EP4 layouts, with and without an a2a backend;
  • the real FusedMoE.weight_loader and _weight_loader_physical fill every local slot with its own expert (the same layouts, random EP2 and EP4); the stand-in layer records the writes instead of copying tensors;
  • every entry of each rank's static dispatch table points at a slot that holds that expert;
  • after an EP 2 -> 4 scale-up, loading routed experts on the existing ranks keeps their fused shared slot;
  • a param with _sglang_require_global_experts still gets every slot, on every rank (the new filter does not apply to it);
  • moe_wna16 picks the right qzeros slots: a second replica, scale-up on ranks 0 and 1, a draft layer with a fused shared expert;
  • the dispatch candidates are unchanged;
  • the draft-layer fallback and the global-scale branch keep their conventions, and out-of-range and negative ids raise.

Alternatives considered:

• Keep every local replica when narrowing. This fixes loading but changes dispatch: dynamic gets more candidates, and the static table draws its random picks in a different order.
• Change only dispatch, so that no rank picks an unloaded replica. After an online rebalance the lists are full again, and a slot that keeps the same expert is not copied (update_expert_weights_single_layer treats it as unchanged), so a never-written slot can still be picked.
• Copy the first local replica into the others after loading. That has to cover every load path and every parameter, including quantization scales.

Accuracy Tests

CPU repro (script and full output below). EP2, 4 logical experts + 2 redundant, rank 1 holds both replicas of expert 3 (slots 3 and 4), rank 0 holds none.

• On main @ 78185f4, rank 1 never writes slot 4. Rank 0's static table sends expert 3 to slot 4 in 1 of 8 layers, and with dynamic, 4047 of rank 0's 8000 tokens for expert 3 go there.
• With this PR, rank 1 writes slots 3, 4 and 5. Both ranks' static picks for expert 3 are the same as on main, layer by layer; this PR does not touch the dispatch tables.

CPU tests

• The new file: 8 passed, 42 subtests.
• Together with the existing test/registered/unit/eplb/, test/registered/unit/model_loader/test_presharded_loader.py, test/registered/unit/model_loader/test_modelopt_loader.py and test/registered/unit/layers/quantization/test_moe_wna16_ep_shard.py: 123 passed, 2 skipped (both need CUDA), 88 subtests passed.
• With the fix reverted in place (the old lookup, a filter that keeps every id, the old moe_wna16 code), 6 of the 8 new tests fail. Each of these partial changes makes at least one new test fail:
  • reverting only the lookup;
  • making only the new filter keep every id;
  • reverting only the moe_wna16 change;
  • filtering shared ids in moe_wna16 too;
  • applying the filter to require_global_experts params too;
  • dropping the lower bound of the id check.
• Lint with the versions pinned in .pre-commit-config.yaml, plus check_registered_tests: clean.

GPU. We ran this change once, on an earlier main, with the two saved layouts; this was not a GPU run of 78185f4 itself. That tree also had the changes of #45341, which are inactive here because the weight cache is off.

• Model and flags: Qwen3-30B-A3B (BF16) with --tp 2 --ep-size 2 --enable-eplb --ep-num-redundant-experts 2 --random-seed 42 --moe-a2a-backend deepep --deepep-mode normal --moe-runner-backend deep_gemm --deepep-dispatcher-output-dtype bf16 --max-running-requests 128 --mem-fraction-static 0.6 --disable-radix-cache --eplb-rebalance-num-iterations 1000000000 --expert-distribution-recorder-buffer-size 1000 (no rebalance during the test), plus --init-expert-location set to one of the two saved layouts above.
• Controls, in the same run: the same model, flags, saved layouts and disk-loading path, without this change. Their trees also had inactive weight-cache changes.
• After loading, we hashed every tensor on every rank, slot by slot for the expert tensors.
• A slot counts as an expert only if its hashes equal those of an expert in a server that used the default layout and never rebalanced.
• The reference is a server that routes with the same layout and whose 6240 expert slots all hold the expert the layout puts there. It received its weights through the weight cache from a server that had rebalanced into this layout ([Fix] weight_cache: Keep the EPLB expert layout in the daemon #45341).
• Greedy outputs: 8 fixed prompts x 32 tokens.

	layout 1	layout 2
layers with two replicas of an expert on one rank (of 48)	16	23
slots that hold such a second replica (of 6240)	19	27
without this change: slots that hold no expert	19, exactly those	27, exactly those
without this change: expert slots equal to the reference	6221	6213
with this change: expert slots equal to the reference	6240	6240
greedy outputs equal to the reference, without / with this change	7/8 / 8/8	5/8 / 8/8

All non-expert tensors were byte-identical to the reference on both ranks, in all four servers.

GSM8K (5-shot, all 1319 questions, 128 at a time) was recorded for context and was not a criterion. Correct answers: 1204 and 1210 without this change, 1205 and 1193 with it. Two passes of the never-rebalanced server on the same setup differed by 7 and by 9.

Not tested on GPU:

• quantized checkpoints (the moe_wna16 path);
• elastic EP scale-up;
• dynamic dispatch;
• online rebalancing after startup, and weight updates after it;
• other models, EP > 2, more than one node;
• a repeat of this run.

Speed Tests and Profiling

Not measured. Only weight loading changes: the lookup compares one row of physical_to_logical_map per call instead of indexing the per-rank list, and a rank that holds two replicas of an expert now writes that expert twice. Dispatch and forward passes are unchanged.

Related

• EPLB: prefer to use physical experts in the same gpu or node #10874: the per-rank narrowing.
• The earlier fix for unloaded rows of global scale parameters: [Fix] Fix nan error for large scale ep #12866, Fix nan in global scaling factor for large scale nvfp4 EP #13162, Fix dpsk-r1-fp4 tp8 by reverting two commits (#13162 and #13341) #13348 (reverted Fix nan in global scaling factor for large scale nvfp4 EP #13162 after loading hangs), Fix global scaling factor loading hang #13484 (reintroduced it with a race fix).
• Fix silently wrong EPLB output with --moe-a2a-backend none (rank-invariant dispatch) #32962: no narrowing without an a2a backend.
• Open [Feat] eplb: Non-blocking expert rebalance #45282 and feat(eplb): add topology-aware expert placement #37923 also edit eplb/expert_location.py, in other functions.
• eplb algorithm can generate phy2log with duplicate hot expert ids on same rank deepseek-ai/EPLB#22: the algorithm can put duplicate replicas on one GPU. fix: prevent duplicate expert IDs on same GPU (#22) deepseek-ai/EPLB#27 and fix: prevent duplicate logical experts on same GPU deepseek-ai/EPLB#28 try to avoid that where possible; layouts that already exist are not changed.
• Open [Fix] moe: Gate expert location per FusedMoE for every draft layer #45322 adds an elif right before the logical_to_all_physical() call in FusedMoE.weight_loader. Whichever lands second may need a small rebase.
• Different from [Fix] eplb: Fix Qwen3 MoE output corruption after EPLB rebalance #39580 and [Bugfix] Fix model output corruption caused by EPLB rebalance (Eager and CUDA Graph modes) #18213, which are about wrong output after rebalancing without an a2a backend.

CPU repro script

"""CPU-only repro: with an a2a backend (e.g. --moe-a2a-backend deepep) and an
--init-expert-location physical_to_logical_map that puts two replicas of one
logical expert on the same EP rank, the second replica is never loaded, while
another rank's static dispatch table can still route tokens to it.

It calls the real SGLang code; no GPU, model or server is needed:
  compute_initial_expert_location_metadata    what each rank builds from --init-expert-location
  ExpertLocationMetadata.logical_to_all_physical  the slots FusedMoE.weight_loader writes
  topk_ids_logical_to_physical                the physical slot a token is dispatched to

Usage: PYTHONPATH=<sglang checkout>/python python repro_same_rank_replica.py [--layers N]
Exit code: 1 if a rank dispatches to a slot its owner never loads, else 0.
"""

import argparse
import json
import sys
from unittest import mock

import torch

from sglang.srt.eplb import expert_location as el
from sglang.srt.eplb.expert_location_dispatch import (
    ExpertLocationDispatchInfo,
    topk_ids_logical_to_physical,
)
from sglang.srt.runtime_context import get_context

EP_SIZE = 2
NUM_LOGICAL = 4
NUM_REDUNDANT = 2
NUM_PHYSICAL = NUM_LOGICAL + NUM_REDUNDANT  # 6 slots: rank 0 = 0..2, rank 1 = 3..5
SLOTS_PER_RANK = NUM_PHYSICAL // EP_SIZE
NUM_LAYERS = 8  # --layers
HOT = 3
# Same row in every layer: logical expert 3 has two replicas, both on rank 1
# (slots 3 and 4); rank 0 has no replica of it.
ROW = [0, 1, 2, 3, 3, 0]


def owner(slot):
    return slot // SLOTS_PER_RANK


def build_rank(rank, a2a_backend, dispatch_algorithm, num_tokens=1):
    """Return what one EP rank loads and where it dispatches, per layer."""
    shape = el.ModelConfigForExpertLocation(
        num_layers=NUM_LAYERS, num_logical_experts=NUM_LOGICAL
    )
    overrides = dict(
        device="cpu",
        ep_size=EP_SIZE,
        nnodes=1,
        moe_a2a_backend=a2a_backend,
        ep_dispatch_algorithm=dispatch_algorithm,
        ep_num_redundant_experts=NUM_REDUNDANT,
        init_expert_location=json.dumps(
            {"physical_to_logical_map": [ROW] * NUM_LAYERS}
        ),
    )
    # No model is loaded, so only the layer / expert counts come from the patch.
    with get_context().override_server_args(**overrides), mock.patch.object(
        el.ModelConfigForExpertLocation,
        "from_model_config",
        staticmethod(lambda model_config: shape),
    ):
        # ModelRunner.maybe_init_expert_location_metadata does this on every rank
        # before load_model(). With static dispatch it also builds this rank's
        # logical_to_rank_dispatch_physical_map (a seeded random pick, seed 42).
        metadata = el.compute_initial_expert_location_metadata(
            model_config=None, moe_ep_rank=rank
        )
        el.set_global_expert_location_metadata(metadata, allow_overwrite=True)
        loaded, dispatch = [], []
        for layer in range(NUM_LAYERS):
            # FusedMoE.weight_loader: for checkpoint expert e, write every slot in
            # logical_to_all_physical(layer, e) that lives on this rank.
            loaded.append(
                sorted(
                    slot
                    for e in range(NUM_LOGICAL)
                    for slot in metadata.logical_to_all_physical(layer, e)
                    if owner(slot) == rank
                )
            )
            info = ExpertLocationDispatchInfo.init_new(layer_id=layer)
            topk_ids = torch.full((num_tokens, 1), HOT, dtype=torch.int64)
            dispatch.append(topk_ids_logical_to_physical(topk_ids, info).view(-1))
        candidates = metadata.logical_to_all_physical(0, HOT)
    return loaded, dispatch, candidates


def run(a2a_backend, dispatch_algorithm, num_tokens=1):
    ranks = [
        build_rank(r, a2a_backend, dispatch_algorithm, num_tokens)
        for r in range(EP_SIZE)
    ]
    print(
        f"\n== --moe-a2a-backend {a2a_backend} "
        f"--ep-dispatch-algorithm {dispatch_algorithm} =="
    )
    for r, (loaded, _, candidates) in enumerate(ranks):
        never = sorted(
            set(range(r * SLOTS_PER_RANK, (r + 1) * SLOTS_PER_RANK)) - set(loaded[0])
        )
        print(
            f"rank {r}: weight loader writes slots {loaded[0]}, never writes {never}; "
            f"logical_to_all_physical(layer, {HOT}) = {candidates}"
        )
    bad_layers = set()
    for r, (_, dispatch, _) in enumerate(ranks):
        hits = {}  # layer -> tokens sent to a slot its owner never loaded
        for layer, slots in enumerate(dispatch):
            missing = [
                int(s) for s in slots if int(s) not in ranks[owner(int(s))][0][layer]
            ]
            if missing:
                hits[layer] = missing
        bad_layers |= set(hits)
        slots_hit = sorted({s for m in hits.values() for s in m})
        if num_tokens == 1:
            picks = [int(d[0]) for d in dispatch]
            print(
                f"rank {r}: logical {HOT} is dispatched to slot, per layer: {picks}"
                + (
                    f"\n        -> never-loaded slot {slots_hit} in layers {sorted(hits)}"
                    if hits
                    else ""
                )
            )
        else:
            n = sum(len(m) for m in hits.values())
            print(
                f"rank {r}: {n}/{num_tokens * NUM_LAYERS} tokens of logical {HOT} go to "
                f"never-loaded slot(s) {slots_hit} ({len(hits)}/{NUM_LAYERS} layers)"
            )
    return len(bad_layers)


def main():
    global NUM_LAYERS
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--layers", type=int, default=NUM_LAYERS)
    NUM_LAYERS = parser.parse_args().layers
    print(
        f"EP size {EP_SIZE}, {NUM_LOGICAL} logical experts + {NUM_REDUNDANT} redundant = "
        f"{NUM_PHYSICAL} slots ({SLOTS_PER_RANK} per rank), {NUM_LAYERS} layers"
    )
    print(
        f"--init-expert-location row in every layer: rank 0 slots 0-2 = {ROW[:3]}, "
        f"rank 1 slots 3-5 = {ROW[3:]}"
    )
    bad_static = run("deepep", "static")
    torch.manual_seed(0)
    bad_dynamic = run("deepep", "dynamic", num_tokens=1000)
    bad_control = run("none", "dynamic", num_tokens=1000)
    print()
    print(f"layers that route to a never-loaded slot (of {NUM_LAYERS}):")
    print(f"  deepep + static (default with an a2a backend): {bad_static}")
    print(f"  deepep + dynamic: {bad_dynamic}")
    print(f"  control, no a2a backend: {bad_control}")
    reproduced = bad_static > 0 or bad_dynamic > 0
    print("RESULT:", "REPRODUCED" if reproduced else "NOT_REPRODUCED")
    return 1 if reproduced else 0


if __name__ == "__main__":
    sys.exit(main())

CPU repro output, main @ 78185f4

EP size 2, 4 logical experts + 2 redundant = 6 slots (3 per rank), 8 layers
--init-expert-location row in every layer: rank 0 slots 0-2 = [0, 1, 2], rank 1 slots 3-5 = [3, 3, 0]

== --moe-a2a-backend deepep --ep-dispatch-algorithm static ==
rank 0: weight loader writes slots [0, 1, 2], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 1: weight loader writes slots [3, 5], never writes [4]; logical_to_all_physical(layer, 3) = [3]
rank 0: logical 3 is dispatched to slot, per layer: [3, 3, 4, 3, 3, 3, 3, 3]
        -> never-loaded slot [4] in layers [2]
rank 1: logical 3 is dispatched to slot, per layer: [3, 3, 3, 3, 3, 3, 3, 3]

== --moe-a2a-backend deepep --ep-dispatch-algorithm dynamic ==
rank 0: weight loader writes slots [0, 1, 2], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 1: weight loader writes slots [3, 5], never writes [4]; logical_to_all_physical(layer, 3) = [3]
rank 0: 4047/8000 tokens of logical 3 go to never-loaded slot(s) [4] (8/8 layers)
rank 1: 0/8000 tokens of logical 3 go to never-loaded slot(s) [] (0/8 layers)

== --moe-a2a-backend none --ep-dispatch-algorithm dynamic ==
rank 0: weight loader writes slots [0, 1, 2], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 1: weight loader writes slots [3, 4, 5], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 0: 0/8000 tokens of logical 3 go to never-loaded slot(s) [] (0/8 layers)
rank 1: 0/8000 tokens of logical 3 go to never-loaded slot(s) [] (0/8 layers)

layers that route to a never-loaded slot (of 8):
  deepep + static (default with an a2a backend): 1
  deepep + dynamic: 8
  control, no a2a backend: 0
RESULT: REPRODUCED

CPU repro output, this PR

EP size 2, 4 logical experts + 2 redundant = 6 slots (3 per rank), 8 layers
--init-expert-location row in every layer: rank 0 slots 0-2 = [0, 1, 2], rank 1 slots 3-5 = [3, 3, 0]

== --moe-a2a-backend deepep --ep-dispatch-algorithm static ==
rank 0: weight loader writes slots [0, 1, 2], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 1: weight loader writes slots [3, 4, 5], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 0: logical 3 is dispatched to slot, per layer: [3, 3, 4, 3, 3, 3, 3, 3]
rank 1: logical 3 is dispatched to slot, per layer: [3, 3, 3, 3, 3, 3, 3, 3]

== --moe-a2a-backend deepep --ep-dispatch-algorithm dynamic ==
rank 0: weight loader writes slots [0, 1, 2], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 1: weight loader writes slots [3, 4, 5], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 0: 0/8000 tokens of logical 3 go to never-loaded slot(s) [] (0/8 layers)
rank 1: 0/8000 tokens of logical 3 go to never-loaded slot(s) [] (0/8 layers)

== --moe-a2a-backend none --ep-dispatch-algorithm dynamic ==
rank 0: weight loader writes slots [0, 1, 2], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 1: weight loader writes slots [3, 4, 5], never writes []; logical_to_all_physical(layer, 3) = [3, 4]
rank 0: 0/8000 tokens of logical 3 go to never-loaded slot(s) [] (0/8 layers)
rank 1: 0/8000 tokens of logical 3 go to never-loaded slot(s) [] (0/8 layers)

layers that route to a never-loaded slot (of 8):
  deepep + static (default with an a2a backend): 0
  deepep + dynamic: 0
  control, no a2a backend: 0
RESULT: NOT_REPRODUCED

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ✅ Run #40305204419
Latest PR Test (Extra): ❌ Run #40305204539
Latest PR Test (AMD ROCm 10): ❌ Run #40305204466

[Shunkangz]

/tag-run-ci-label

[ShangmingCai]

LGTM

[ShangmingCai]

/tag-and-rerun-ci

[alexnails]

small comment on CPU perf

[alexnails]

this looks to cost the same as before.

[Agent check]
A CPU microbench at DSV3 shape (58 layers × 256 logical × 288 physical, 89k calls ≈ one reload's worth of w1/w2/w3 + scales) gives 3.66 µs/call here vs 3.38 µs for the old l2p[l, e].tolist() + filter, ~0.33 s vs ~0.30 s per full reload, per rank.

My suggestion:
a follow-up could build a List[List[List[int]]] from physical_to_logical_map_cpu.tolist() once in _init_raw/update (~1 ms). That makes this lookup 0.08 µs/call, so ~0.3 s → ~8 ms. The old code had the same cost, so don't really feel u need to add pr scope.

[zanes-ops]

Thanks for measuring it. Agreed, caching the inverse map once per update can be a separate follow-up; I'll keep this PR to the loading fix.

[alexnails]

rm this comment? the others are ok.

[zanes-ops]

Removed, thanks.

[ShangmingCai]

/rerun-failed-ci

[zanes-ops]

/rerun-failed-ci