Skip to content

[Fix] weight_cache: Keep the EPLB expert layout in the daemon - #45341

Open
zanes-ops wants to merge 5 commits into
sgl-project:mainfrom
zanes-ops:weight-cache-expert-layout-sync-pr
Open

zanes-ops wants to merge 5 commits into
sgl-project:mainfrom
zanes-ops:weight-cache-expert-layout-sync-pr

Conversation

@zanes-ops

@zanes-ops zanes-ops commented Oct 20, 2026 •

Copy link
Copy Markdown
Contributor

Motivation

With --weight-cache-mode client (or daemon), an engine's parameters are the weight cache daemon's tensors, mapped zero-copy. An EPLB rebalance moves routed experts in place (ExpertLocationUpdater), so it also moves them in the daemon's copy and under every other engine that maps that copy. Only the rebalancing engine's ExpertLocationMetadata follows. #36198 made --enable-eplb work with the weight cache, so that the daemon builds the same initial layout as the engine, but nothing records the layout after a rebalance.

An engine that starts on the same daemon later builds its metadata from the initial layout. That includes the same engine restarted in client mode, which is the fast-recovery case the weight cache exists for, and a warm standby. Such an engine reads the moved experts as the wrong ones. It starts and serves without any error or warning, but its outputs are wrong.

We measured this on an earlier main with #39580 applied, on one node of B200 GPUs: Qwen3-30B-A3B (BF16), --tp 2 --ep-size 2 --enable-eplb --ep-num-redundant-experts 2 --moe-runner-backend triton, no a2a backend, --eplb-rebalance-num-iterations 400 for the engines on the daemon. #39580 is needed for this model without an a2a backend; without it, the rebalancing engine itself keeps routing with the initial layout after a rebalance. Greedy outputs are 8 prompts x 32 tokens; GSM8K is 5-shot over all 1319 test questions.

GSM8K, correct of 1319 greedy tokens equal to the reference
reference: loaded from disk, never rebalanced (two passes) 1197, 1208 -
engine A on the daemon, two passes, 12 rebalances in total 1209, 1193 8/8 before its first rebalance, 8/8 after both passes
engine B, started on the same daemon after A stopped 0 0/8

B's non-expert tensors were byte-identical to the reference on both ranks; only the routed-expert tensors differed. Every one of B's expert slots held a real expert, together they matched the last layout A logged, and 6186 of the 6240 slots (48 layers x 130 slots) held a different expert than the initial layout puts there.

The in-tree test TestWeightCacheDaemonQwen3MoeEP enables EPLB, but its few requests never reach a rebalance, so CI does not catch this.

With this PR, the daemon stores the current expert layout of its weights (the physical_to_logical_map) with a version. Engines start from that layout, a rebalance reports each move to the daemons in two phases, and only an engine that is alone on its daemons may move experts.

Modifications

Daemon (new weight_cache/expert_layout.py, weight_cache/daemon.py)

  • A daemon started with --enable-eplb keeps the initial layout it already builds ([Weight Cache] Enhance test and support EPLB #36198) as its record: the full physical_to_logical_map (all layers x all physical slots, the same shape as the engine's metadata), a version that starts at 0, and a SHA-256 digest of the map. A daemon without --enable-eplb keeps no record.
  • New requests on the existing control socket: expert_layout_attach, expert_layout_query, expert_layout_begin, expert_layout_commit and expert_layout_abort. The existing requests keep their names; fetch_state gains an optional field, the layout checks described below, and a new refusal status.
  • Presence without PIDs. An engine's attach connection stays open for the engine's lifetime, and a background thread in the daemon watches it. The kernel closes the connection however the engine exits, including SIGKILL, so the daemon sees EOF and detaches the engine. The daemon never needs the engine's PID, which another container cannot see; this is the same idea as the socket-based liveness suggested on [Roadmap]Fast Engine Recovery: Weight Cache Daemon #33522 for the engine's daemon watchdog. That watchdog is unchanged.
  • Two-phase moves. begin marks a move of some layers pending. commit stores the new rows of those layers, increments the version and clears the pending move. abort clears it without a change. If the engine detaches while its move is pending, the daemon marks its layout as unknown ("dirty") and logs which engine, which layers and which base version. A dirty daemon refuses fetch_state and begin until it is restarted. This is conservative: an engine killed after begin but before it touched any weight also leaves the daemon dirty.
  • begin replies. busy (try again later): another engine is attached, or another engine's move is pending. refused (permanent): the daemon keeps no record or is dirty, the engine is not attached, the engine already has a move pending, the base version is not the current one, the layer list is invalid, or an engine that does not follow the layout has fetched the weights (see the next item).
  • Weight fetch check. fetch_state takes an optional expert_layout field: the engine's id, its version, and the digest of the map it routes with. The daemon checks it after the existing config check and before it calls the transport. It compares the digest (the version only appears in messages) and replies status: "layout_refused" with the reason when the digest differs, when the daemon is dirty, or when a move is pending. A fetch without a live attachment is still served if the layout matches: an older engine without the field while the version is 0, or an engine that has the field but did not attach. The daemon counts such engines as not following the layout, and from then on no engine may move experts until the daemon restarts, because the daemon cannot tell when they exit.

Engine startup (model_executor/model_runner.py, weight_cache/ipc_loader.py)

  • ModelRunner.maybe_init_expert_location_metadata runs before load_model. With the weight cache on, for models that have expert-location metadata, each rank now attaches there to its own daemon, at the socket path the IPC loader uses. Models without expert-location metadata, such as dense models, are not affected. If another engine's move is pending, the rank waits, with a 300 s timeout after which it fails and names that engine and its layers. The ranks then all-gather what they found over the world CPU group:
    • Every daemon keeps a record, all at the same version and digest: the engine builds its metadata with ExpertLocationMetadata.init_by_mapping from that map, the same way --init-expert-location with that map does. If the map equals the engine's own initial layout, its metadata stays as it is.
    • Any daemon is dirty, the daemons disagree, or a rank could not read its daemon: every rank fails at startup with the reason, so no rank is left waiting in a collective.
    • Some daemons keep no record (an older daemon, a daemon without --enable-eplb, or no daemon): the engine starts with its own initial layout if no recording daemon has moved an expert yet (version 0), and fails otherwise.
  • When a rank builds its metadata from the daemon's map, it first checks two things and fails with an explicit error otherwise: the map must have the engine's shape (the same --ep-num-redundant-experts on both), and the engine must have --enable-eplb (or --init-expert-location), because without it the engine routes as if every expert were in its initial slot.
  • IpcModelLoader sends the expert_layout field with fetch_state and raises with the daemon's reason on layout_refused. It also records whether the weights came from the daemon or from disk.

Rebalance (eplb/eplb_manager.py)

  • If every rank loaded its weights from disk (in client mode, when no daemon was running), the weights belong to the engine and rebalance as before.
  • Weights from the daemons move one chunk (--eplb-rebalance-layers-per-chunk) at a time:
    1. Every rank sends begin to its daemon, and the ranks all-gather the replies. They move only if every daemon agreed. If one is busy, this rebalance stops and the next interval tries again. If one refused, rebalancing is off for good, with a warning. Ranks whose begin succeeded send abort.
    2. The experts move with the existing update_expert_location_with_recovery, unchanged.
    3. A device synchronize, so that the moved weights have landed before a daemon can serve them to another engine.
    4. Every rank sends commit with the new rows of the moved layers. The ranks check that every daemon recorded it, at the same new version, and that each daemon's digest equals the digest of the engine's whole map. If not, the engine logs an error that says to restart the daemons before another engine uses them, keeps serving, and does not rebalance again. A daemon whose commit failed keeps the move pending, so it becomes dirty when this engine exits; a daemon that accepted the commit has already recorded it.
  • With weights from the daemons, rebalancing stays off (one warning per reason; the engine keeps its current layout) when the daemon is an older version, when it runs without --enable-eplb, when the engine did not attach at startup, or when some ranks loaded from disk and others from the daemons. It also stays off with elastic EP: this PR does not coordinate elastic EP's rebalances or its recovery of missing experts with the daemons.

One mover per daemon

  • A daemon accepts begin only from an engine that is alone on it. A second engine on the same daemon (a warm standby, or a restarted engine while the old one is still exiting) makes begin return busy, so the first engine's rebalancing pauses. The second engine therefore never holds a stale layout. When the first engine exits, the second one is alone and may rebalance; its first begin must name the daemon's current version.

Compatibility

Combination Behavior
new engine, older daemon expert_layout_attach gets Unknown request type. The engine starts with its own layout and does not rebalance the shared weights (warning). The older daemon ignores the extra fetch_state field.
older engine, new daemon Served as today while no expert has moved (version 0). From then on, no engine may move experts until the daemon restarts. Once an expert has moved, older engines are refused with a message to upgrade the engine or restart the daemon. The older loader raises on that refusal; it does not fall back to disk.
daemon without --enable-eplb No record. Engines start as today and do not rebalance the shared weights.
--weight-cache-mode daemon The same protocol. The daemons exit with the engine (parent-death signal), so the next start loads the initial layout from disk again. From reading the code; not run on GPU.

Transport. The layout requests use the daemon's control socket and do not change the transport interface. The fetch_state check runs before the transport sends anything, so it applies to torch_ipc and to any transport that maps the daemon's memory. This includes the planned vmm_fd transport (a stub in transport.py today): its engines will map the daemon's memory too, so an in-place rebalance changes every engine's weights in the same way, and the same requests and check cover it with no change to this protocol.

Tests. New CPU tests in test/registered/unit/model_loader/test_weight_cache_expert_layout.py (32) and test/registered/unit/eplb/test_eplb_weight_cache.py (5), registered in base-a-test-cpu. On the daemon side they run the real WeightCacheDaemon request handling over a real Unix socket, without loading a model. They cover:

  • versions, begin/commit/abort, and malformed moves;
  • an engine that disconnects mid-move, including a real engine-side process killed with SIGKILL after begin: the daemon notices only through the socket EOF, then refuses old and new engines, and a new engine fails at startup with the reason;
  • the one-mover rule, including a standby that takes over, and a second engine on only one rank's daemon pausing the move on every rank (the other rank aborts);
  • startup agreement: two daemons at different versions, a rank that cannot read its daemon, waiting for a pending move, a stuck move;
  • that metadata built from the daemon's map equals metadata built with --init-expert-location and the same map, tensor by tensor, for a2a backend none and deepep, dynamic and static dispatch, and EP ranks 0 and 1;
  • the shape check and the missing --enable-eplb check, the loader's fetch_state field and refusal message, both compatibility directions, and EPLBManager with the weight cache off, with weights from disk, through the daemons, and with a busy daemon.

Alternatives considered

  • Keep rebalancing off while the weight cache is on. Safe, but it removes dynamic EPLB from every weight cache deployment.
  • Give each engine a private copy of the routed experts. That copies most of an MoE model per engine, which is what the weight cache exists to avoid.
  • Let the first engine keep rebalancing and push each new layout to the other attached engines. That needs in-place metadata updates in a serving engine, and between a move and the update the other engines would compute with moved weights. Not in this PR; a second engine pauses rebalancing instead.

Accuracy Tests

CPU. The two new test files, the existing test/registered/unit/model_loader/test_weight_cache_protocol.py and all of test/registered/unit/eplb/: 113 passed, 1 skipped (a CUDA-only case in the existing test_megamoe_prefill_recorder). Lint with the versions pinned in .pre-commit-config.yaml and check_registered_tests: clean. We also broke 14 of the rules above one at a time (for example: no dirty mark when an engine leaves mid-move, begin ignoring other engines, fetch_state skipping the digest check, a commit trusting the daemon's digest); each change made at least one of the new tests fail.

GPU. This change on an earlier main, in a recent SGLang nightly image, on one node of B200 GPUs. Qwen3-30B-A3B (BF16), TP2/EP2. Every engine and daemon gets:

--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

The daemons run with --tp-size 2. Every server runs with --tp 2 --mem-fraction-static 0.6 --disable-radix-cache, and the servers on the daemon add --weight-cache-mode client --eplb-rebalance-num-iterations 400. Expert dispatch is the default for an a2a backend, static. Why these flags:

  • DeepEP: Qwen3 MoE passes the expert-location dispatch info to TopK in forward_deepep() but not in forward_normal() ([Fix] eplb: Fix Qwen3 MoE output corruption after EPLB rebalance #39580). Without an a2a backend, the rebalancing engine itself would be wrong after its first rebalance.
  • normal mode: in this version, an engine with --deepep-mode auto and --enable-eplb fails at startup, because the expert distribution recorder raises NotImplementedError for auto. So auto could not be tested.
  • deep_gemm: with triton, BF16 DeepEPMoE raises NotImplementedError in normal dispatch (from reading the code).
  • --deepep-dispatcher-output-dtype bf16: DeepEP's dispatch output defaults to FP8 and is switched to BF16 in process_weights_after_loading, which a client engine skips. This holds with or without this PR (from reading the code).
  • --max-running-requests 128: as in the in-tree DeepEP tests.

The reference R uses the same flags, loads from disk, and has a rebalance interval no run reaches. Greedy outputs are 8 fixed prompts x 32 tokens with EOS ignored. GSM8K is 5-shot over all 1319 test questions, greedy, 128 requests at a time. Manifests hold the SHA-256 of every tensor on every rank after loading, slot by slot for the expert tensors.

Case 1: restart on the same daemon (run twice, torch_ipc). The two runs ran at the same time, on different GPUs of the same node. Engine A rebalanced through two GSM8K passes, then ran a final greedy pass with no rebalance in it (repeated until one had none) and was stopped. Engine B then started on the same daemons. As a control, engine D loaded from disk with --init-expert-location set to the map the daemons reported.

run 1 run 2
R, GSM8K (two passes) 1203, 1200 1196, 1202
A, GSM8K while rebalancing (two passes) 1195, 1201 1206, 1203
daemon version on both ranks after A (= A's number of rebalances) 12 12
B logs that it routes with the daemon's layout at that version, on both ranks yes yes
B, every expert slot holds the expert the daemon's map names (48 layers) yes yes
B, greedy tokens equal to A's final pass 8/8 8/8
B vs D, all non-expert tensors and every expert slot D's loader fills byte-identical byte-identical
B, GSM8K 1204 1198

B's greedy pass ran before B's own first rebalance. We compare B's tokens with A's final pass, which used the same daemon weights and the same layout, rather than with D: with an a2a backend, D's loader does not fill the second replica of an expert when one rank holds two replicas of it, while another rank can still dispatch tokens to that slot. Here D had 24 and 19 such slots that held no expert at all, and its greedy tokens differed from B's on 2 and 3 of 8 prompts. This is a pre-existing loader issue that this PR does not change.

Case 2: a warm standby pauses rebalancing (once, torch_ipc). Engine A, alone on the daemons, rebalanced 6 times (version 6). Engine C then attached to the same daemons, adopted version 6, and gave the same greedy tokens as A (8/8). While C was attached, all 6 of A's rebalance attempts stopped at begin with busy; the daemons recorded no move and their maps did not change. After C stopped, A rebalanced again (6 more times) and the version went past 6. A's three GSM8K passes (before, during and after C): 1202, 1202, 1199.

Case 3: SIGKILL during a rebalance (once, torch_ipc). Engine A ran with --eplb-rebalance-layers-per-chunk 1 and was killed with SIGKILL as soon as the daemons had recorded its first chunk. Afterwards the daemons had no pending move and were at version 1, with layer 0 in the new layout and layers 1-47 in the old one, so the kill landed between that chunk's commit and the next begin. B adopted version 1 and matched D, a disk-loaded engine with that layout: all 483 tensors per rank byte-identical (D's loader filled every slot of this layout) and greedy tokens 8/8. The other outcome, a kill while a move is pending, is covered by the CPU test above. One kill point on GPU does not cover every kill point.

Not tested on GPU: DeepEP auto (see above), no a2a backend (needs #39580), dynamic expert dispatch, other models, quantized checkpoints, TP > 2, more than one node, --weight-cache-mode daemon, elastic EP, mixed engine and daemon versions (CPU tests only), and repeats of cases 2 and 3.

Speed Tests and Profiling

Not measured. Engines without the weight cache are unchanged. With the weight cache, in the normal case each rebalanced chunk adds a begin and a commit request from each rank to its local daemon (plus an abort when another rank's begin fails), two all-gathers of small objects over the world CPU group, and one device synchronize; the first rebalance adds one more all-gather to agree on how to move. Engine startup adds one attach and one all-gather.

Known limitations

  • Without an a2a backend, Qwen3 MoE keeps routing with the initial layout after a rebalance, with or without the weight cache ([Fix] eplb: Fix Qwen3 MoE output corruption after EPLB rebalance #39580). There, this PR keeps the daemon's record right, but engines are wrong after a rebalance until [Fix] eplb: Fix Qwen3 MoE output corruption after EPLB rebalance #39580 lands, and so is an engine that adopts a moved layout.
  • Engines without this change still rebalance shared weights without telling the daemon (today's behavior). The daemon cannot see those moves. It only stops other engines from moving experts once such an engine has fetched the weights.
  • A warm standby does not follow the active engine's moves. Rebalancing pauses while more than one engine is attached.
  • With elastic EP, daemon-shared weights are not rebalanced. Weights loaded from disk rebalance as before.
  • A child process forked without exec that inherits the attach connection and outlives the engine keeps the engine counted as attached. That keeps rebalancing paused; it cannot make an engine read a wrong layout.
  • A dirty daemon does not repair its weights. The daemon, and the engines that use it, must be restarted; the error says so.

Related

Checklist

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 #39533826792
Latest PR Test (Extra): ❌ Run #39533825949
Latest PR Test (AMD ROCm 10): ❌ Run #39533826023

@Shunkangz

Copy link
Copy Markdown
Contributor

/tag-and-rerun-ci

@github-actions github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Oct 20, 2026
With the weight cache on, an engine's routed expert weights are the
daemon's tensors. An EPLB rebalance moves experts in place, so it moves
them for the daemon and for every engine that maps it, while only the
rebalancing engine's expert location metadata follows. An engine that
starts on the same daemon afterwards, such as the same engine restarted
in client mode, routes with the initial layout and reads the wrong
experts, without any error.

The daemon now keeps the layout of its weights (the
physical_to_logical_map) with a version:

- An engine attaches to its rank's daemon before it builds its expert
  location metadata, and builds the metadata from the daemon's map, the
  same way as --init-expert-location with that map. All ranks must find
  the same layout. fetch_state carries a digest of the map the engine
  routes with, and the daemon serves no weights whose layout differs.
- EPLBManager moves each chunk of layers in two phases. begin marks the
  move pending on every rank's daemon; commit stores the new rows once
  the weights have moved. An engine that goes away while its move is
  pending leaves the daemon dirty, and a dirty daemon serves no weights
  until it is restarted.
- Only an engine that is alone on its daemons moves experts. The attach
  connection stays open and is the engine's presence, so a second engine
  (a warm standby, or a restarted engine while the old one still exits)
  pauses rebalancing instead of reading moved experts with a stale
  layout.

With a daemon that lacks these requests or runs without --enable-eplb,
the engine keeps its initial layout and does not rebalance the shared
weights. An older engine still gets the weights while no expert has
moved; once one has, the daemon refuses it instead of letting it read
the wrong experts. Weights that came from disk because no daemon was
running belong to the engine and rebalance as before.

Signed-off-by: Zane <zanes@nvidia.com>
Before each commit to the daemons, EPLBManager waits for the moved
weights with torch.get_device_module().synchronize(). The CPU CI runner
installs a CUDA build of torch and has no driver, so the real call
raised there. The test now records the sync in its call log, which also
checks that it comes before the commit.

Signed-off-by: Zane <zanes@nvidia.com>
@zanes-ops
zanes-ops force-pushed the weight-cache-expert-layout-sync-pr branch from 5a45dde to 1197ded Compare October 20, 2026 09:17

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants