[TE] perf(efa): register a batch smallest-buffer-first - #3216
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>
Registering device memory on EFA costs roughly k x (bytes of device memory already registered on that domain) -- ~260 ms/GiB on p5.48xlarge, flat from 0 to 19 GiB, and nearly independent of the buffer's own size. Measured single-threaded and strictly serial, so neither locks nor queueing are involved: the same 2.79 MB CUDA buffer costs 70 ms on an empty domain and 2505 ms once 9.17 GiB is in. A buffer's bytes are therefore charged once per buffer registered after it, which makes ascending size order a well-defined optimum rather than a heuristic. Registering the same 48 GPU buffers (9.23 GiB) serially takes 35.5 s ascending against 93.3 s descending; on Kimi-K3's real batch across 8 TP ranks it is 3.1x (101.3 s -> 32.8 s) on descending input. Sort the dispatch order only. A buffer's NIC set is derived from its own length and the NIC count, so nothing lands anywhere else, and the order of segment_desc->buffers was already whatever order the workers finished in. This is only reachable with MC_MAX_CONCURRENT_REG_MR set: unbounded, every buffer gets its own thread, nothing queues, and the dispatch order does not survive to the provider (112 s vs 120 s, i.e. noise). Host memory does not accumulate -- the same sweep on mmap'd host buffers is flat and order-insensitive -- so sorting is a no-op there rather than a regression, and is applied for all memory types instead of only VRAM. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
github-actions
Bot
added
documentation
The CI format gate runs clang-format 20, which fits both brace-init lists onto fewer lines than the version used to write them. Whitespace only; the test asserts the same thing. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
|
Closing this: I measured it properly and it has no effect. The 2.6x/3.1x figures in the description are an artifact of my benchmark, not a property of the code. The measurement error. Those numbers came from a harness that called With a real batch, the sort still does nothing. Replaying a Kimi-K3 KV registration at the deployed shape (8 ranks, one
Why it cannot work as written. What survives. The underlying mechanism is real and reproducible — a device-memory registration does cost roughly I have rebased #3219 off this branch, so it now stacks on #3210 alone and needs no changes here. Apologies for the noise, and thanks for the review time already spent on it. |
2 participants
Description
EfaTransport::registerLocalMemoryBatch()hands buffers to its workers in the order the caller supplied. On EFA that order has a large effect on total registration time, and there is a provably optimal one: smallest first.Mechanism: registration cost depends on what is already registered
Registering device memory on EFA costs roughly
k x (bytes of device memory already registered on that domain), withk ~= 260 ms/GiBon p5.48xlarge, flat across 0–19 GiB, and nearly independent of the buffer's own size.Measured single-threaded and strictly serial — one buffer per
batch_register_memory()call — so neither locks nor queueing are involved:So cost is a function of what is already in, not of the buffer's own size. The cost curve for 48 x 391 MB is a straight line: 143 ms for the 1st registration, 5022 ms for the 48th — a 35x spread for identical work.
Control — identical 48 buffers, identical 9.23 GiB total, only the order differs:
2.6x on identical work with one thread and no queue.
Why ascending is optimal rather than a heuristic
With
cost(i) = k x prior_bytes(i), each buffer's bytes are charged once per buffer registered after it. Total cost isk x sum over pairs (i<j) of size(i), which is minimized exactly when sizes are non-decreasing. Registering the large buffers last puts their bytes behind the fewest subsequent registrations.This is also why longest-processing-time scheduling intuition mispredicted the effect (#3210 recorded "largest-first is worst" as unexplained): it was never a scheduling phenomenon.
Measured effect
8-process replay of Kimi-K3's KV registration on p5.48xlarge (32 EFA NICs, one
TransferEngineper TP rank, 182 GPU buffers of 2.5 KB–391 MB each), slowest rank,MC_MAX_CONCURRENT_REG_MR=16:The sort makes the caller's order irrelevant — a descending batch lands level with one that was already ascending. Note the honest framing: 3.1x is the worst-case figure; on the order SGLang actually passes it is ~20%. For SGLang the cap in #3210 is the bigger lever (2.6x); this removes the remainder and the cliff for callers with a less lucky order.
Why it only matters with a cap set
Unbounded,
runBoundedParallelspawns one thread per buffer, nothing queues, and the dispatch order does not survive to the provider. Measured: 112.4 s in pool order vs 119.9 s ascending — i.e. noise, with ascending marginally worse.That is the one-directional coupling to #3210, and the reason this is not presented as a standalone win.
Host memory: no-op, not a regression
The accumulation is a device-memory property. Repeating the same sweep on
mmap'd host buffers:Host is flat. A control that allocates the identical sequence in both modes and varies only the registration order (2 reps) shows no order effect:
The direction flips between reps, so the difference is run-to-run spread. Sorting is therefore a no-op for host memory rather than a regression, which is why it is applied unconditionally instead of only for VRAM — a memory-type branch would add a code path for no measured benefit.
Likely cause on the device side is the CUDA/HMEM path rebuilding per-domain state on each registration, not a generic
fi_mr_regcost. Worth a libfabricefaissue, out of scope here.Scope and safety
The change is a dispatch-order permutation only. Nothing observable moves:
registerLocalMemoryInternal()a buffer's NIC set is computed from its own length, the page size, and the NIC count.pages_per_nicis function-local, andgetMaxPteEntries()is a compile-time constant with no cumulative tracking, so no buffer can migrate between the full-coverage and disjoint-partition branches because of what was registered before it.Topology::selectDevice()is unaffected — it keys on the location string plus athread_localcounter, so it is already call-to-call variable and registration order contributes nothing.EfaTransport::selectDevice()is unaffected — it scanssegment_desc->buffersby address range and returnsbuffer_idas an output, never indexing by a caller-side position.segment_desc->buffersordering was already nondeterministic — the workerspush_backconcurrently, so its order was whatever they finished in, both before and after this change.std::stable_sort), so dispatch is a stable function of the input.unregisterLocalMemoryBatch()is deliberately not touched: it takes anaddr_listwith no lengths, so there is nothing to sort by, and the accumulation effect does not apply to teardown.TransferMetadata::addLocalMemoryBuffer()and the other 15 transports are untouched. An earlier draft of this change also pushed to the segment descriptor in the caller's original order; that was dropped as a cross-transport refactor needing its own RFC, and it is not required for the win above.Question for maintainers
Should the same sort apply to the other transports? I can only measure EFA (this is the only fabric I have hardware for), so this PR is deliberately EFA-only. The mechanism is a libfabric/EFA-provider property and I have no evidence it exists for
rdma,cxi,barex, etc. — but if it does, the natural home would beTransferEngineImpl::registerLocalMemoryBatch()rather than 16 copies. Happy to move it there if someone can measure a second fabric.A larger lever this does not address
For visibility, since it dwarfs both this PR and #3210: Mooncake registers every single-chunk buffer on all NICs, while NIXL's libfabric plugin registers a
VRAM_SEGbuffer only on that GPU's topology-local rails (selectRailsForMemory()→getEfaDevicesForPci()). On p5.48xlarge that is 32 NICs vs 4. On 48 x 391 MB GPU buffers registered serially, 116.1 s on 32 NICs vs 14.9 s on 4 — 7.8x — because it multiplies the accumulation above, the per-domain cost being paid once per domain. That is a throughput trade-off rather than a free win (fewer NICs can serve a transfer touching that buffer), so it needs its own PR with bandwidth numbers. Filed as issue #3217.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:
EFARegistrationOrderTest.AscendingBySizecovers the empty batch, a descending input, an already-ascending input (unchanged), that the result is a permutation rather than a filter, and that equal lengths preserve the caller's order. It constructsBufferEntryvalues with dummy addresses and calls the ordering function directly, so it runs anywhere — no EFA, no allocation.Manual testing is the measurement above on P5-1 (p5.48xlarge, 32 EFA NICs, 8x H100): the serial single-buffer sweeps on GPU and host memory that establish the mechanism, the 8-process 3-order matrix at cap=16, the unbounded control showing the cap dependency, and the 32/4/1-NIC fan-out comparison.
Not yet verified on p6-b300 (16 NICs) — no capacity has been available. The mechanism is provider-side and should carry over, but the magnitude will differ with the NIC count.
Checklist
./scripts/code_format.shpre-commit run --all-filesand all hooks passFormatted with clang-format 20.1.8 (
--changed-linesagainstmain), matching the CI gate. No docs change: this alters no configuration surface and no documented behavior. 92 LOC, so no RFC.Checked for overlap: no open PR touches
registerLocalMemoryBatch(). This stacks on #3210, which touches the same function.AI Assistance Disclosure
Claude Code was used to design and run the measurement sweeps that established the mechanism, and to draft the code, test, and this description. Every line has been reviewed by me and I can defend the change end-to-end.