[TE] perf(efa): optionally register device memory only on topology-local NICs - #3219
Draft
whn09 wants to merge 5 commits into
Draft
[TE] perf(efa): optionally register device memory only on topology-local NICs#3219whn09 wants to merge 5 commits into
whn09 wants to merge 5 commits into
Conversation
registerLocalMemoryBatch() and unregisterLocalMemoryBatch() spawn one std::async(std::launch::async) per buffer with no cap, which libstdc++ takes literally -- one fresh thread each. Kimi-K3 registers ~180 KV buffers per TP rank in a single batch, and with one TransferEngine per rank an 8-rank node peaks at ~1100 threads all inside fi_mr_reg / fi_mr_regattr. Replace both fan-outs with a fixed pool pulling from a shared index, sized by MC_MAX_CONCURRENT_REG_MR. The default of 0 means unbounded, so behavior is unchanged unless an operator sets the knob: with no cap the pool spawns count-1 threads and runs the caller as a worker, matching what std::async did. Error semantics are also unchanged -- every item is still attempted and the first non-zero return is propagated. The knob is opt-in with no built-in default because its effect depends on the order the caller passes buffers in, which this layer cannot see. Measured on p5.48xlarge (32 NICs) replaying K3's registration the way SGLang issues it -- 8 processes, one engine and one GPU each, 182 buffers of 2.5 KB to 391 MB per process, barrier-synchronized. Slowest rank, since nothing serves until all 8 finish: per-proc cap descending pool order (SGLang's) ascending unset 108/99/128 s 112 s 120 s 64 184/208 s 101 s 60 s 16 95/98/97 s 43 s 36/37 s 8 107 s -- 44 s 4 157 s -- -- A cap of 16 is worth 2.6x on the order SGLang actually uses, but only 1.1x on descending order, and the order alone swings the capped result by 2.7x. Unbounded is order-insensitive (99-128 s) because nothing queues. Largest-first being the worst order contradicts longest-processing-time scheduling and is still unexplained, so no sort is applied here yet. Two other platforms put the optimum elsewhere, which is the other reason not to compile in a default: p5 host memory (fi_mr_reg), 128 x 2 GiB of 4 KB pages, prefers cap 16 (59.1 s vs 274 s unbounded); a 2x p6-b300 K3 server run at 16 NICs/rank prefers cap 8 (20.7 s vs 138.7 s unbounded). A core-scaled default would be worse still on a large node -- where capping helps, the bottleneck is the provider lock, not CPU. So the code keeps its historical behavior and only exposes the pool size.
The bullet had grown into a multi-paragraph section with measurement tables, out of proportion to every other entry in this list. Keep the operationally relevant facts -- the cap is per process, the default is unbounded, a good value is platform- and order-dependent, do not core-scale it -- and leave the supporting measurements to the PR.
The comment and docs claimed the cap should not be core-scaled, on the grounds that the bottleneck was a shared provider lock rather than CPU. Both halves were wrong, and both came from a single-process benchmark where 48 registration threads on a 192-core node never contended. Measuring at the real topology instead -- one TransferEngine per TP rank, and varying both the rank count and the core budget -- shows the optimum is set by the CPU budget: 8 ranks/192 cores prefers cap 16, 4 ranks/192 cores prefers 32, and 8 ranks/64 cores prefers 8. All three are the same ~cores global thread count, so a good value is roughly cores/processes-per-node. The lock hypothesis is refuted separately by a 23x speedup from capping a serial baseline, which a size-proportional global lock cannot produce. Also drops the retracted single-process figures and shortens the comment. No functional change. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
documentation
…cal NICs
The EFA transport registers every single-chunk buffer on every NIC. That is
incidental to the `chunks.size() <= 1` branch rather than deliberate: the
multi-chunk paths already partition NICs, only the common case does not.
For device memory the cost is not linear in the NIC count but worse. EFA
charges a device-memory registration roughly in proportion to the device
bytes already registered on that libfabric domain, and that cost is paid
once per domain, so fan-out multiplies accumulation. Measured on
p5.48xlarge (32 NICs, 8 GPUs), 48 x 391 MB GPU buffers registered serially:
all NICs (32) 116.1 s
topology-local (4) 14.9 s -- 7.8x
An 8-rank replay of a Kimi-K3 KV registration (182 GPU buffers per rank,
14.3 GiB, one TransferEngine per rank as SGLang does it) at
MC_MAX_CONCURRENT_REG_MR=16, in the buffer order SGLang actually passes:
all NICs (32) 35.0 s wall 2981 CPU-s 85 cores busy
topology-local 5.2 s wall 394 CPU-s 77 cores busy -- 6.7x / 7.6x
This does not subsume MC_MAX_CONCURRENT_REG_MR. `local` cuts the work per
registration, not the number of them in flight: the thread count is the
buffer count, independent of the NIC set. Unbounded, a `local` batch still
peaked at 1289 threads and 182 of 192 cores busy, versus 168 threads and 77
cores with the cap. The two are orthogonal and both are needed.
Add MC_EFA_NIC_SELECTION=all|local. `local` narrows a single-chunk *device*
buffer to the NICs the topology reports as closest to its GPU.
No new PCI walk or hwloc dependency: Topology::discover() already computes
exactly this. For a "cuda:N" entry, preferred_hca holds the HCAs at minimum
PCI distance from that GPU within its NUMA node -- on p5.48xlarge the 4 EFA
devices sharing the GPU's PCIe root complex, the same rail group NIXL's
libfabric plugin derives from hwloc. buildLocalNicMap() just inverts that
into context_list_ indices. Verified on hardware: each of the 8 GPUs
resolves to its own 4 rails, and the timing above is symmetric across GPUs
(14.9 s on GPU 0, 15.1 s on GPU 5), so it is per-GPU rather than always
GPU 0's set.
Partial registration needed no new plumbing: lkey/rkey are already pushed
as 0 for unassigned NICs, and selectDevice() already retries past a NIC
whose rkey is 0.
Host memory is deliberately excluded. It has no per-domain accumulation to
amplify, it is ~4x cheaper to register anyway, and a host buffer has no
single owning device -- "cpu:N"'s preferred set is a NUMA node's NICs,
measured as 16 of 32 here, half the machine rather than a rail group.
Opt-in, defaulting to the historical behavior, because it trades away
per-buffer bandwidth: fewer NICs can serve a transfer touching that buffer.
The RDMA transport registers on every context unconditionally too, so
defaulting `local` on would also put the EFA transport out of step with the
rest of the codebase.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
whn09
force-pushed
the
perf/efa-per-gpu-nic-selection
branch
from
July 31, 2026 10:02
6e134eb to
ed960b7
Compare
30 tasks
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
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
2 participants
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.
Description
EfaTransport::registerLocalMemoryInternal()registers every single-chunk buffer on every NIC. That looks incidental to thechunks.size() <= 1branch rather than deliberate — the multi-chunk paths already partition NICs (disjoint or round-robin), only the common single-chunk case does not.NIXL's libfabric plugin does the opposite:
selectRailsForMemory()→getEfaDevicesForPci()registers aVRAM_SEGbuffer only on the rails topologically local to that GPU. On p5.48xlarge that is 4 NICs versus 32.Why the fan-out costs more than 8x
Registering device memory on EFA costs roughly
k x (device bytes already registered on that libfabric domain),k ~= 260 ms/GiBon p5.48xlarge, flat across 0-19 GiB and nearly independent of the buffer's own size. The cost is paid once per domain, so the NIC count multiplies the accumulation rather than adding to it.Measured on P5-1 (p5.48xlarge, 32 EFA NICs, 8x H100), 48 x 391 MB GPU buffers registered strictly serially — one buffer per
batch_register_memory()call:all(32 NICs), GPU 0local(4 NICs), GPU 0local(4 NICs), GPU 5all(32 NICs), GPU 57.8x. The GPU 0 / GPU 5 rows are the check that this is really per-GPU rail selection and not "always GPU 0's four NICs": both GPUs land at the same time (14.9 vs 15.1 s), which they would not if one of them were using a remote rail group.
The same thing at the deployed process count
The table above is single-process and strictly serial, which isolates the mechanism but is not how anything runs. Replaying a Kimi-K3 KV registration the way SGLang issues it — 8 ranks, one
TransferEngineand one GPU per rank, 182 GPU buffers / 14.3 GiB each, ranks released from a barrier so all 8 hit the NICs together — in the buffer order SGLang actually passes, atMC_MAX_CONCURRENT_REG_MR=16, slowest rank:all(32)local(4)6.7x wall, 7.6x CPU.
This does not replace
MC_MAX_CONCURRENT_REG_MR(#3210)Worth stating explicitly, because it is the natural assumption and it is wrong:
localreduces the work per registration, not the number of registrations in flight. The thread count inrunBoundedParallel()is the buffer count, which has nothing to do with the NIC set.allalllocallocalUnbounded,
localstill saturates the box: 1289 threads and 182 of 192 cores busy. On a real serving node those cores are needed for inference. The cap is what bounds concurrency;localis what bounds the work. Both are needed, andlocal's wall-clock win from the cap (2.1-2.3x) is the same asall's (2.3x), i.e. the two are orthogonal.What happened to #3216
#3216 sorted each batch smallest-first, on the theory that with
cost(i) = k x prior_bytes(i)the totalk x sum over pairs (i<j) of size(i)is minimized by non-decreasing sizes. Measuring it across this matrix showed no effect anywhere:localalllocalalllocalallcap=1is the strictly serial case the model assumes, and the largest NIC count is where the effect should be biggest —all/cap=1shows 0.35%. Unbounded, the ordering is reversed from the prediction.The reason the original 2.6x was wrong: that measurement called
batch_register_memory()once per buffer, soregistrationOrder()only ever saw a one-element list and sorted nothing. The 2.6x came from reordering the 48 calls in the harness, not from the code under test. With a real 182-buffer batch,runBoundedParallel()hands out indices vianext.fetch_add(1)toworkersthreads, soworkersregistrations are always in flight; a sort fixes the dispatch order while the model needs completion to be serial. #3216 is closed.No new topology code
Topology::discover()already computes precisely what NIXL derives from hwloc. For acuda:Nentry,preferred_hcaholds the HCAs at minimum PCI distance from that GPU within its NUMA node — on p5.48xlarge exactly the 4 EFA devices sharing the GPU's PCIe root complex.buildLocalNicMap()only inverts that mapping intocontext_list_indices. No PCI walk, no hwloc dependency, no new discovery path.(
avail_hcais not usable for this: it is "every HCA that is not preferred", i.e. the other 28.)Partial registration needed no new plumbing either.
lkey/rkeyare already pushed as0for NICs a buffer was not registered on, andEfaTransport::selectDevice()already skips a NIC whoserkeyis 0 and retries the next one.Host memory is deliberately excluded
localapplies to device memory only, for three separate reasons:cpu:N'spreferred_hcais a whole NUMA node's NICs — measured as 16 of 32 on P5-1, i.e. half the machine. That is not a rail group, so "local" would not mean the same thing.Why opt-in rather than the default
It trades bandwidth for registration time: fewer NICs can serve a transfer touching that buffer, so a single-buffer transfer is capped at its rail group's bandwidth instead of the node's aggregate. With
--gpu_id=-1(one buffer per GPU, the shapetransfer_engine_benchand real KV caches use) all 32 NICs are still covered collectively, and each GPU's traffic stays on its own rails — which is the configuration where I expect little or no loss. That is the measurement this draft is missing.Secondly,
RdmaTransport::registerLocalMemoryInternal()registers on every context unconditionally, in both its serial anduse_parallel_regpaths, with no NIC-subset logic at all (locationonly enters metadata, for the peer'sselectDevice()path selection). Defaultinglocalon would put the EFA transport out of step with the rest of the codebase; leaving it opt-in keeps the default behavior identical to today's for every transport.Module
mooncake-transfer-engine)mooncake-store)mooncake-ep)mooncake-pg)mooncake-integration)mooncake-p2p-store)mooncake-wheel)mooncake-common)mooncake-rl)Type of Change
How Has This Been Tested?
Test commands:
Test results:
EfaNicSelectionEnvTestcovers the default (all),local, an explicitalloverwriting a preset field, case-insensitivity, and that an unknown value or empty string leaves the field untouched rather than silently selecting a policy — mis-registering buffers on the wrong NIC set is a correctness-adjacent surprise, not a tuning knob.Manual testing is every table above, on P5-1 (p5.48xlarge, 32 EFA NICs, 8x H100): the serial per-GPU sweep, the 8-rank replay with CPU accounting, the cap x NIC-set matrix, and the ordering matrix that closed #3216.
What is not tested yet, and why:
allvslocal. My two p5.48xlarge instances are in different subnets / AZs. EFA is not routable across subnets, so no EFA transfer runs between them at all — AWS's ownfi_pingpong -p efaalso fails on this pair, independently of Mooncake. I need a same-subnet pair (ideally the same placement group) to produce these numbers. This is the one claim in the description that is currently an expectation rather than a measurement, and the reason this PR is a draft.Checklist
./scripts/code_format.shpre-commit run --all-filesand all hooks passFormatted with clang-format 20.1.8 (
--changed-linesagainstmain), matching the CI gate. DocumentedMC_EFA_NIC_SELECTIONindocs/source/design/transfer-engine/index.mdalongside the other EFA env vars. 226 LOC, so no RFC.Checked for overlap: no open PR touches
registerLocalMemoryInternal(). This stacks on #3210, which touchesregisterLocalMemoryBatch()in the same file.AI Assistance Disclosure
Claude Code was used to run the registration-timing and CPU measurements, to trace how
Topology::discover()and NIXL'sselectRailsForMemory()arrive at the same rail set, to design the matrix that disproved #3216, and to draft the code, tests, and this description. Every line has been reviewed by me and I can defend the change end-to-end.