Repository navigation
Conversation
zanes-ops
requested review from
Fridge003,
Ying1123,
alexnails,
ch-wan,
fzyzcjy,
hnyls2002,
ispobock,
liusy58,
merrymercy and
xutizhou
as code owners
October 20, 2026 06:57
Contributor
|
/tag-and-rerun-ci |
3 of 5 tasks
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
force-pushed
the
weight-cache-expert-layout-sync-pr
branch
from
October 20, 2026 09:17
5a45dde to
1197ded
Compare
This was referenced Oct 26, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
With
--weight-cache-mode client(ordaemon), 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'sExpertLocationMetadatafollows. #36198 made--enable-eplbwork 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
clientmode, 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
mainwith #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 400for 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.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
TestWeightCacheDaemonQwen3MoeEPenables 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)--enable-eplbkeeps the initial layout it already builds ([Weight Cache] Enhance test and support EPLB #36198) as its record: the fullphysical_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-eplbkeeps no record.expert_layout_attach,expert_layout_query,expert_layout_begin,expert_layout_commitandexpert_layout_abort. The existing requests keep their names;fetch_stategains an optional field, the layout checks described below, and a new refusal status.beginmarks a move of some layers pending.commitstores the new rows of those layers, increments the version and clears the pending move.abortclears 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 refusesfetch_stateandbeginuntil it is restarted. This is conservative: an engine killed afterbeginbut before it touched any weight also leaves the daemon dirty.beginreplies.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).fetch_statetakes an optionalexpert_layoutfield: 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 repliesstatus: "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_metadataruns beforeload_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:ExpertLocationMetadata.init_by_mappingfrom that map, the same way--init-expert-locationwith that map does. If the map equals the engine's own initial layout, its metadata stays as it is.--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.--ep-num-redundant-expertson 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.IpcModelLoadersends theexpert_layoutfield withfetch_stateand raises with the daemon's reason onlayout_refused. It also records whether the weights came from the daemon or from disk.Rebalance (
eplb/eplb_manager.py)clientmode, when no daemon was running), the weights belong to the engine and rebalance as before.--eplb-rebalance-layers-per-chunk) at a time:beginto 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 whosebeginsucceeded sendabort.update_expert_location_with_recovery, unchanged.commitwith 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 whosecommitfailed keeps the move pending, so it becomes dirty when this engine exits; a daemon that accepted thecommithas already recorded it.--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
beginonly 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) makesbeginreturnbusy, 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 firstbeginmust name the daemon's current version.Compatibility
expert_layout_attachgetsUnknown request type. The engine starts with its own layout and does not rebalance the shared weights (warning). The older daemon ignores the extrafetch_statefield.--enable-eplb--weight-cache-mode daemonTransport. The layout requests use the daemon's control socket and do not change the transport interface. The
fetch_statecheck runs before the transport sends anything, so it applies totorch_ipcand to any transport that maps the daemon's memory. This includes the plannedvmm_fdtransport (a stub intransport.pytoday): 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) andtest/registered/unit/eplb/test_eplb_weight_cache.py(5), registered inbase-a-test-cpu. On the daemon side they run the realWeightCacheDaemonrequest handling over a real Unix socket, without loading a model. They cover:begin/commit/abort, and malformed moves;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;--init-expert-locationand the same map, tensor by tensor, for a2a backend none and deepep, dynamic and static dispatch, and EP ranks 0 and 1;--enable-eplbcheck, the loader'sfetch_statefield and refusal message, both compatibility directions, andEPLBManagerwith the weight cache off, with weights from disk, through the daemons, and with a busy daemon.Alternatives considered
Accuracy Tests
CPU. The two new test files, the existing
test/registered/unit/model_loader/test_weight_cache_protocol.pyand all oftest/registered/unit/eplb/: 113 passed, 1 skipped (a CUDA-only case in the existingtest_megamoe_prefill_recorder). Lint with the versions pinned in.pre-commit-config.yamlandcheck_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,beginignoring other engines,fetch_stateskipping 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: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:TopKinforward_deepep()but not inforward_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.normalmode: in this version, an engine with--deepep-mode autoand--enable-eplbfails at startup, because the expert distribution recorder raisesNotImplementedErrorforauto. Soautocould not be tested.deep_gemm: withtriton, BF16DeepEPMoEraisesNotImplementedErrorin normal dispatch (from reading the code).--deepep-dispatcher-output-dtype bf16: DeepEP's dispatch output defaults to FP8 and is switched to BF16 inprocess_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-locationset to the map the daemons reported.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 atbeginwithbusy; 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 1and 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'scommitand the nextbegin. 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),dynamicexpert 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
beginand acommitrequest from each rank to its local daemon (plus anabortwhen another rank'sbeginfails), 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
Related
--enable-eplbwith the weight cache; the daemon builds the initial layout.EPLBManager.rebalance(); whichever lands second needs a rebase:begin/commitaround each layer, or it would have to stay off;weight_cache/daemon.pyorweight_cache/ipc_loader.py: [weight cache] Add disk-vs-daemon parity harness; fix persistence, zombie-PID liveness, and SM100 block-FP8 gating #36156, weight-cache: support post-load lifecycle and MXFP4 IPC #38306, [weight_cache] Addstatusrequest and CLI for daemon monitoring #39256 (astatusrequest in the same request handler), [Feat] deepseek-v4.1: Add experimental resident-weight recovery and warm standby #44241.Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #39533826792
Latest PR Test (Extra): ❌ Run #39533825949
Latest PR Test (AMD ROCm 10): ❌ Run #39533826023